anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Bugfix] Apply function module marker per-model instead of hardcoding diffusion (#5875)

Browse Source 
Signed-off-by: khairulkabir1661 <khairulkabir1661@users.noreply.github.com>
Co-authored-by: khairulkabir1661 <khairulkabir1661@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Khairul Kabirkhairulkabir1661Claude Opus 4.6
 authored and GitHub committed 2026-09-25 14:17:19 -06:00
1 parent a8576ccb72
commit 5e36b5a86b
5 files changed
+50 -7
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/model_tests/diffusion/case_filtering.py
+19 -3
View File
@@ -20,7 +20,12 @@ import itertools
import pytest
from tests.helpers.mark import hardware_marks
from tests.model_tests.diffusion.config_types import DiffusionAccs, DiffusionModelTestOpts, get_required_device_count
from tests.model_tests.diffusion.config_types import (
DiffusionAccs,
DiffusionModelTestOpts,
ModelTypeMarker,
get_required_device_count,
)
from vllm_omni.platforms import current_omni_platform
# These tests are intended to run on L4 GPUs since they test parallelism (i.e., need GPUs),
@@ -32,7 +37,12 @@ from vllm_omni.platforms import current_omni_platform
MAX_CI_DEVICES = 4
def get_test_group_marks(model_name: str, test_group: list[DiffusionAccs] | None, model_marks: list | None) -> list:
def get_test_group_marks(
model_name: str,
test_group: list[DiffusionAccs] | None,
model_marks: list | None,
model_type_marker: ModelTypeMarker,
) -> list:
"""Build the full set of pytest marks for a test group.
For now, single device groups default to core model, and multi device groups
@@ -40,6 +50,7 @@ def get_test_group_marks(model_name: str, test_group: list[DiffusionAccs] | None
the max number of GPUs for the CI to prevent accidentally adding configurations that
would be skipped in the CI."""
marks = list(model_marks) if model_marks is not None else []
marks.append(getattr(pytest.mark, model_type_marker.value))
required_devices = get_required_device_count(test_group)
if required_devices > MAX_CI_DEVICES:
@@ -83,7 +94,12 @@ def get_model_parametrization(model_name: str, test_info: DiffusionModelTestOpts
# The divergence check compares generated outputs, which only the offline test does.
*([] if online else [test_info.check_i2v_t2v_divergence]),
id=f"{model_name}[{'+'.join(test_group)}]" if test_group else model_name,
marks=get_test_group_marks(model_name, test_group, test_info.marks),
marks=get_test_group_marks(
model_name,
test_group,
test_info.marks,
model_type_marker=test_info.model_type_marker,
),
)
for test_group in test_groups
]
tests/model_tests/diffusion/config_types.py
+12
View File
@@ -42,6 +42,14 @@ class DiffusionTasks(str, Enum):
# Text to audio, etc should be added here as needed
class ModelTypeMarker(str, Enum):
"""Model type markers for pytest parametrization."""
DIFFUSION = "diffusion"
OMNI = "omni"
TTS = "tts"
class DiffusionModelTestOpts(NamedTuple):
"""Configuration for one Diffusion model's tests."""
@@ -59,6 +67,10 @@ class DiffusionModelTestOpts(NamedTuple):
# Actual tasks which controls the tests actually run
supported_tasks: list[DiffusionTasks]
# Function module marker applied per-test-case via parametrization.
# Set to ModelTypeMarker.OMNI or ModelTypeMarker.TTS for non-diffusion models.
model_type_marker: ModelTypeMarker
# Additional acceleration groups to run beyond the base case (no acceleration).
# The base case is always run for every model in the test settings. None means
# we only run the base case.
tests/model_tests/diffusion/model_settings.py
+15
View File
@@ -7,6 +7,7 @@ from tests.model_tests.diffusion.config_types import (
DiffusionAccs,
DiffusionModelTestOpts,
DiffusionTasks,
ModelTypeMarker,
)
# This object defines the (tiny) model configurations for common tests.
@@ -37,12 +38,14 @@ DIFFUSION_TEST_SETTINGS = {
model=real_anima_model,
builder=tiny_anima_builder,
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
model_type_marker=ModelTypeMarker.DIFFUSION,
checkpoint_filename=CHECKPOINT_FILENAME,
),
"Flux2KleinPipeline": DiffusionModelTestOpts(
model="black-forest-labs/FLUX.2-klein-4B",
builder=diff_model_builders.tiny_flux2_klein_builder,
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE, DiffusionTasks.IMAGE_TO_IMAGE],
model_type_marker=ModelTypeMarker.DIFFUSION,
extra_test_groups=[
[DiffusionAccs.HSDP, DiffusionAccs.TEA_CACHE],
[DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
@@ -53,6 +56,7 @@ DIFFUSION_TEST_SETTINGS = {
model="Lightricks/LTX-2",
builder=diff_model_builders.tiny_ltx2_builder,
supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO, DiffusionTasks.IMAGE_TO_VIDEO],
model_type_marker=ModelTypeMarker.DIFFUSION,
extra_test_groups=[
[DiffusionAccs.HSDP, DiffusionAccs.CACHE_DIT],
[DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
@@ -63,6 +67,7 @@ DIFFUSION_TEST_SETTINGS = {
model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
builder=diff_model_builders.tiny_sana_video_builder,
supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO],
model_type_marker=ModelTypeMarker.DIFFUSION,
extra_test_groups=[
[DiffusionAccs.TENSOR_PARALLEL],
[DiffusionAccs.CFG_PARALLEL],
@@ -82,6 +87,7 @@ DIFFUSION_TEST_SETTINGS = {
model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
builder=diff_model_builders.tiny_sana_video_i2v_builder,
supported_tasks=[DiffusionTasks.IMAGE_TO_VIDEO],
model_type_marker=ModelTypeMarker.DIFFUSION,
extra_test_groups=[
[DiffusionAccs.CFG_PARALLEL],
[DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL],
@@ -95,11 +101,13 @@ DIFFUSION_TEST_SETTINGS = {
model="Qwen/Qwen-Image",
builder=diff_model_builders.tiny_qwen_image_builder,
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
model_type_marker=ModelTypeMarker.DIFFUSION,
),
"LongCatImagePipeline": DiffusionModelTestOpts(
model="meituan-longcat/LongCat-Image",
builder=diff_model_builders.tiny_longcat_image_builder,
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
model_type_marker=ModelTypeMarker.DIFFUSION,
extra_test_groups=[
[DiffusionAccs.TEA_CACHE],
[DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
@@ -110,6 +118,7 @@ DIFFUSION_TEST_SETTINGS = {
model="meituan-longcat/LongCat-Image-Edit",
builder=diff_model_builders.tiny_longcat_image_edit_builder,
supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
model_type_marker=ModelTypeMarker.DIFFUSION,
extra_test_groups=[
[DiffusionAccs.TEA_CACHE],
[DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
@@ -120,11 +129,13 @@ DIFFUSION_TEST_SETTINGS = {
model="black-forest-labs/FLUX.1-schnell",
builder=diff_model_builders.tiny_flux_builder,
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
model_type_marker=ModelTypeMarker.DIFFUSION,
),
"FluxKontextPipeline": DiffusionModelTestOpts(
model="black-forest-labs/FLUX.1-Kontext-dev",
builder=diff_model_builders.tiny_flux_kontext_builder,
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE, DiffusionTasks.IMAGE_TO_IMAGE],
model_type_marker=ModelTypeMarker.DIFFUSION,
extra_test_groups=[
[DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
[DiffusionAccs.CFG_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
@@ -135,6 +146,7 @@ DIFFUSION_TEST_SETTINGS = {
model="black-forest-labs/FLUX.2-dev",
builder=diff_model_builders.tiny_flux2_builder,
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
model_type_marker=ModelTypeMarker.DIFFUSION,
extra_test_groups=[
[DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
[DiffusionAccs.CFG_PARALLEL, DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
@@ -145,16 +157,19 @@ DIFFUSION_TEST_SETTINGS = {
model="Qwen/Qwen-Image-Edit",
builder=diff_model_builders.tiny_qwen_image_edit_builder,
supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
model_type_marker=ModelTypeMarker.DIFFUSION,
),
"QwenImageEditPlusPipeline": DiffusionModelTestOpts(
model="Qwen/Qwen-Image-Edit-2511",
builder=diff_model_builders.tiny_qwen_image_edit_plus_builder,
supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
model_type_marker=ModelTypeMarker.DIFFUSION,
),
"StableDiffusion3Pipeline": DiffusionModelTestOpts(
model="stabilityai/stable-diffusion-3.5-medium",
builder=diff_model_builders.tiny_sd3_builder,
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
model_type_marker=ModelTypeMarker.DIFFUSION,
extra_test_groups=[
[DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
[DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
tests/model_tests/diffusion/test_common_offline.py
+2 -2
View File
@@ -21,8 +21,8 @@ from tests.model_tests.diffusion.task_runners import (
run_and_validate_text_to_video_request,
)
# NOTE : Hardware marks are added dynamically based on test requirements
pytestmark = [pytest.mark.diffusion, pytest.mark.xdist]
# NOTE: Hardware and model type marks are added dynamically based on test requirements and model type
pytestmark = [pytest.mark.xdist]
@pytest.mark.parametrize(
tests/model_tests/diffusion/test_common_online.py
+2 -2
View File
@@ -31,8 +31,8 @@ from tests.model_tests.diffusion.task_runners import (
)
from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
# NOTE : Hardware marks are added dynamically based on test requirements
pytestmark = [pytest.mark.diffusion, pytest.mark.xdist]
# NOTE: Hardware and model type marks are added dynamically based on test requirements and model type
pytestmark = [pytest.mark.xdist]
@pytest.fixture(autouse=True)