diff --git a/tests/model_tests/diffusion/case_filtering.py b/tests/model_tests/diffusion/case_filtering.py index 9f53edce0..d84ba6ac5 100644 --- a/tests/model_tests/diffusion/case_filtering.py +++ b/tests/model_tests/diffusion/case_filtering.py @@ -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 ] diff --git a/tests/model_tests/diffusion/config_types.py b/tests/model_tests/diffusion/config_types.py index 8b66af191..fab480130 100644 --- a/tests/model_tests/diffusion/config_types.py +++ b/tests/model_tests/diffusion/config_types.py @@ -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. diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py index fb5376046..e2a202c89 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/tests/model_tests/diffusion/model_settings.py @@ -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], diff --git a/tests/model_tests/diffusion/test_common_offline.py b/tests/model_tests/diffusion/test_common_offline.py index b610ed013..ae2ff6901 100644 --- a/tests/model_tests/diffusion/test_common_offline.py +++ b/tests/model_tests/diffusion/test_common_offline.py @@ -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( diff --git a/tests/model_tests/diffusion/test_common_online.py b/tests/model_tests/diffusion/test_common_online.py index ef4449501..ac4393f7b 100644 --- a/tests/model_tests/diffusion/test_common_online.py +++ b/tests/model_tests/diffusion/test_common_online.py @@ -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)