[Tests] Tiny Diffusion Model Pattern (#4524)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/model_tests/diffusion/case_filtering.py

@@ -0,0 +1,104 @@
+"""
+Analogous to: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/vlm_utils/case_filtering.py
+
+This follows a similar pattern so that we can use `get_parametrized_options` to expand our test
+configurations out, similar to the way in which vLLM expands supported input type tests. For now,
+this mostly just means parametrizing over the test groups, which may be unique per model, but
+doing so in this way lets us cleanly define test groups that may be heterogeneous + avoid having
+excessive skip marks for features that may not be supported in every pipeline.
+"""
+
+# Engine initialization is currently quite expensive. For efficiency,
+# when we run the tests, we should try to collapse compatible accelerations,
+# because even with tiny models, the overhead of spinning up a new instance
+# stacks up fast.
+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 vllm_omni.platforms import current_omni_platform
+
+# These tests are intended to run on L4 GPUs since they test parallelism (i.e., need GPUs),
+# but are intended to be smoke tests that we can properly validate without real model
+# weights in the CI.
+#
+# As such, no DiffusionAccs may require more than 4 GPUs, since we would otherwise skip
+# them in the CI.
+MAX_CI_DEVICES = 4
+
+
+def get_test_group_marks(model_name: str, test_group: list[DiffusionAccs] | None, model_marks: list | None) -> 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
+    default to advanced_model marks. We also validate and throw if a test group exceeds
+    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 []
+
+    required_devices = get_required_device_count(test_group)
+    if required_devices > MAX_CI_DEVICES:
+        raise ValueError(
+            f"Test group {test_group} for {model_name} requires {required_devices} devices, "
+            f"but the max CI device count is {MAX_CI_DEVICES}. "
+        )
+
+    marks.extend(hardware_marks(res={"cuda": "L4"}, num_cards=required_devices))
+    # hardware_marks only adds skipif for num_cards > 1, so we currently handle the single-device case
+    # directly here. This should probably be handled in a more common way later on.
+    assert current_omni_platform is not None and current_omni_platform.device_count is not None
+    device_count = current_omni_platform.device_count()
+    if current_omni_platform.is_cuda() and device_count < required_devices:
+        marks.append(pytest.mark.skip(reason=f"Need {required_devices} devices, got {device_count}"))
+    if required_devices > 1:
+        marks.append(pytest.mark.full_model)
+    else:
+        marks.append(pytest.mark.core_model)
+    return marks
+
+
+def get_model_parametrization(model_name: str, test_info: DiffusionModelTestOpts, online: bool):
+    """Given a model & its corresponding test options, build the list of pytest params
+    to be run for this model. The base case (no accelerations) is always included. Extra
+    test groups are always appended for offline tests, but are only added to the online
+    tests if online_base_only is overridden to False to avoid redundant testing.
+    """
+    test_groups: list[list[DiffusionAccs] | None] = [None]
+    if not (online and test_info.online_base_only) and test_info.extra_test_groups:
+        test_groups.extend(test_info.extra_test_groups)
+
+    return [
+        pytest.param(
+            model_name,
+            test_group,
+            test_info.supported_tasks,
+            test_info.check_multi_output and test_group is None,
+            test_info.check_determinism and test_group is None,
+            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),
+        )
+        for test_group in test_groups
+    ]
+
+
+def get_parametrized_options(
+    test_settings: dict[str, DiffusionModelTestOpts],
+    online: bool = False,
+):
+    """Converts all the DiffusionModelTestOpts into an expanded list of parameters
+    based on which accelerations are available.
+
+    When online=True and a model has online_base_only=True, only the base case
+    (no accelerations) is included for that model's parametrization to minimize
+    redundancy between offline / online tests.
+    """
+    # Get a list per model type, where each entry contains a tuple of all of
+    # that model type's cases, then flatten them into the top level so that
+    # we can consume them in one mark.parametrize call.
+    parametrization = [
+        get_model_parametrization(model_name, test_info, online=online)
+        for model_name, test_info in test_settings.items()
+    ]
+    return list(itertools.chain(*parametrization))

tests/model_tests/diffusion/config_types.py

@@ -0,0 +1,187 @@
+"""
+Common definitions for controlling what tests run where.
+"""
+
+from collections.abc import Callable
+from enum import StrEnum, auto
+from typing import NamedTuple, TypeAlias
+
+from pytest import MarkDecorator
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+
+# All builder funcs take no params and return a path
+TinyDiffusionBuilder: TypeAlias = Callable[[], str]
+
+
+class DiffusionAccs(StrEnum):
+    """Supported acceleration types / test settings for Diffusion Models."""
+
+    HSDP = auto()
+    TEA_CACHE = auto()
+    CACHE_DIT = auto()
+    SEQUENCE_PARALLEL = auto()
+    CFG_PARALLEL = auto()
+    TENSOR_PARALLEL = auto()
+    CPU_OFFLOAD = auto()
+    LAYERWISE_OFFLOAD = auto()
+    VAE_PATCH_PARALLEL = auto()
+
+
+class DiffusionTasks(StrEnum):
+    """Supported tasks for Diffusion Models."""
+
+    TEXT_TO_IMAGE = auto()
+    IMAGE_TO_IMAGE = auto()
+    # Text to video, text to audio, etc should be added here as needed
+
+
+class DiffusionModelTestOpts(NamedTuple):
+    """Configuration for one Diffusion model's tests."""
+
+    # HF model name for real-weight tests (advanced_model / full_model level).
+    # For now, whether we use the real weights vs tiny weights in the common tests
+    # depends on the run level.
+    model: str
+
+    # Creates a tiny model for the given architecture. We should always use tiny
+    # model weights for tests that do not require us to check the model quality.
+    builder: TinyDiffusionBuilder
+
+    # Actual tasks which controls the tests actually run
+    supported_tasks: list[DiffusionTasks]
+
+    # 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.
+    extra_test_groups: list[list[DiffusionAccs]] | None = None
+
+    # Pytest Marks for this model. This may be useful for selecting which models
+    # we want to run where, similar to the way vLLM's multimodal tests mark some
+    # as core models to always run in the CI.
+    # Example: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/test_common.py#L131
+    marks: list[MarkDecorator] | None = None
+
+    # When True (default), online tests only run the base case (no accelerations).
+    # When False, online tests run all test_groups, same as offline. This should be
+    # True unless there is a good reason for it not to be, because the execution of
+    # the acceleration should be the same on both codepaths, and CLI parsing etc should
+    # be tested by tests adding the acceleration, and not per model.
+    online_base_only: bool = True
+
+    # Additional checks to run for the base case.
+    check_multi_output: bool = True  # Runs multiple generations in one request
+    check_determinism: bool = True  # Runs 2 generations with the same seed and check determinism
+
+
+### Mappings & utils for building offline Omni() instances given a list of enabled accelerations
+ACC_OMNI_KWARGS = {
+    DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_use_tiling": True},
+    DiffusionAccs.CPU_OFFLOAD: {"enable_cpu_offload": True},
+    DiffusionAccs.LAYERWISE_OFFLOAD: {"enable_layerwise_offload": True},
+    DiffusionAccs.CACHE_DIT: {"cache_backend": "cache_dit"},
+    DiffusionAccs.TEA_CACHE: {"cache_backend": "tea_cache"},
+}
+
+ACC_PARALLEL_KWARGS = {
+    DiffusionAccs.HSDP: {"use_hsdp": True, "hsdp_shard_size": 2},
+    DiffusionAccs.TENSOR_PARALLEL: {"tensor_parallel_size": 2},
+    DiffusionAccs.CFG_PARALLEL: {"cfg_parallel_size": 2},
+    DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_patch_parallel_size": 2},
+    # For SP, we don't run ring here to conserve devices, since it is easy
+    # to blow up the number of needed devices fast. Compatibility for ring
+    # and ulysses together should be tested generically outside of these tests.
+    DiffusionAccs.SEQUENCE_PARALLEL: {"ulysses_degree": 2},
+}
+
+ACC_DEVICE_COUNT_KEYS = {
+    DiffusionAccs.TENSOR_PARALLEL: "tensor_parallel_size",
+    DiffusionAccs.CFG_PARALLEL: "cfg_parallel_size",
+    DiffusionAccs.SEQUENCE_PARALLEL: "ulysses_degree",
+    DiffusionAccs.VAE_PATCH_PARALLEL: "vae_patch_parallel_size",
+    DiffusionAccs.HSDP: "hsdp_shard_size",
+}
+
+### CLI args for launching an OmniServer subprocess per acceleration.
+# Follows the same pattern as e2e online serving tests, which hand-code
+# their server_args lists in OmniServerParams.
+ACC_SERVER_ARGS: dict[DiffusionAccs, list[str]] = {
+    DiffusionAccs.HSDP: ["--use-hsdp", "--hsdp-shard-size", "2"],
+    DiffusionAccs.TEA_CACHE: ["--cache-backend", "tea_cache"],
+    DiffusionAccs.CACHE_DIT: ["--cache-backend", "cache_dit"],
+    DiffusionAccs.SEQUENCE_PARALLEL: ["--usp", "2"],
+    DiffusionAccs.CFG_PARALLEL: ["--cfg-parallel-size", "2"],
+    DiffusionAccs.TENSOR_PARALLEL: ["--tensor-parallel-size", "2"],
+    DiffusionAccs.CPU_OFFLOAD: ["--enable-cpu-offload"],
+    DiffusionAccs.LAYERWISE_OFFLOAD: ["--enable-layerwise-offload"],
+    DiffusionAccs.VAE_PATCH_PARALLEL: ["--vae-use-tiling", "--vae-patch-parallel-size", "2"],
+}
+
+## TODO ^ These should be in the same object, it's getting messy.
+
+
+def get_required_device_count(accelerations: list[DiffusionAccs] | None) -> int:
+    """Compute the minimum number of devices needed for a set of accelerations.
+    The total is the product of all parallel dimensions (defaulting to 1).
+
+    If not enough devices are available for a test group's accelerations,
+    that test will be skipped."""
+    count = 1
+    if accelerations is None:
+        return count
+
+    for acc in accelerations:
+        key = ACC_DEVICE_COUNT_KEYS.get(acc)
+        if key is not None:
+            count *= ACC_PARALLEL_KWARGS[acc][key]
+    return count
+
+
+def build_parallel_config_from_diff_accelerations(accelerations: list[DiffusionAccs]) -> DiffusionParallelConfig | None:
+    """Given a list of accelerations pertaining to the current test group,
+    build the parallel config needed for the Omni() object (if any)."""
+    config_kwargs = {}
+    for acc in accelerations:
+        update_dict = ACC_PARALLEL_KWARGS.get(acc, {})
+        config_kwargs.update(update_dict)
+    if config_kwargs:
+        return DiffusionParallelConfig(**config_kwargs)
+    return None
+
+
+### Offline Omni() object builder
+def build_omni_from_diff_accelerations(accelerations: list[DiffusionAccs] | None, **kwargs) -> Omni:
+    """Given one or more acceleration types, build the corresponding Omni() object."""
+    # Coerce to a list and build the parallel config, since that depends on the accelerations
+    if accelerations is None:
+        accelerations = []
+    parallel_config = build_parallel_config_from_diff_accelerations(accelerations)
+
+    # Then add anything else that's a top-level kwarg
+    acc_kwargs = {}
+    if parallel_config is not None:
+        acc_kwargs["parallel_config"] = parallel_config
+    for acc in accelerations:
+        update_dict = ACC_OMNI_KWARGS.get(acc, {})
+        acc_kwargs.update(update_dict)
+
+    # Keys passed through should mostly be things like enforce_eager;
+    # if there's overlap, it's probably due to a misconfiguration
+    shared_keys = acc_kwargs.keys() & kwargs.keys()
+    if shared_keys:
+        raise ValueError(f"Explicit Omni kwargs and inferred Omni kwargs for accelerations overlap: {shared_keys}")
+    omni_kwargs = {**acc_kwargs, **kwargs}
+    return Omni(**omni_kwargs)
+
+
+### Online server flag builder
+def build_server_args_from_diff_accelerations(accelerations: list[DiffusionAccs] | None) -> list[str]:
+    """Given one or more acceleration types, build the corresponding CLI args
+    for launching an OmniServer subprocess."""
+    if accelerations is None:
+        return []
+    args = []
+    for acc in accelerations:
+        args.extend(ACC_SERVER_ARGS.get(acc, []))
+    return args

tests/model_tests/diffusion/diff_model_builders.py

@@ -0,0 +1,63 @@
+import os
+import tempfile
+
+import torch
+from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
+from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel
+from diffusers.pipelines.flux2.pipeline_flux2_klein import Flux2KleinPipeline
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+from transformers import AutoTokenizer, Qwen3Config, Qwen3ForCausalLM
+
+TINY_MODEL_DIR = os.path.join(tempfile.gettempdir(), "vllm-omni-tiny-models")
+
+
+def _get_tiny_model_path(name: str) -> str:
+    path = os.path.join(TINY_MODEL_DIR, name)
+    os.makedirs(path, exist_ok=True)
+    return path
+
+
+def tiny_flux2_klein_builder() -> str:
+    """Build a tiny Flux2Klein model."""
+    model_dir = _get_tiny_model_path("Flux2KleinPipeline")
+
+    pipe = Flux2KleinPipeline(
+        scheduler=FlowMatchEulerDiscreteScheduler(),
+        vae=AutoencoderKLFlux2(
+            down_block_types=("DownEncoderBlock2D",),
+            up_block_types=("UpDecoderBlock2D",),
+            block_out_channels=(32,),
+            layers_per_block=1,
+            latent_channels=16,
+            norm_num_groups=16,
+        ),
+        # NOTE: For now we need 28 layers because of hardcoded stuff in the model :(
+        text_encoder=Qwen3ForCausalLM(
+            Qwen3Config(
+                hidden_size=32,
+                intermediate_size=64,
+                num_hidden_layers=28,
+                num_attention_heads=2,
+                num_key_value_heads=2,
+                head_dim=16,
+                vocab_size=151936,
+                max_position_embeddings=512,
+            )
+        ),
+        tokenizer=AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B"),
+        # NOTE: For now we need at least 2 layers for the transformer
+        # due to hardcoded hacks in CacheDiT for Flux2Klein specifically.
+        transformer=Flux2Transformer2DModel(
+            in_channels=64,
+            num_layers=2,
+            num_single_layers=2,
+            attention_head_dim=16,
+            num_attention_heads=2,
+            joint_attention_dim=96,
+            timestep_guidance_channels=32,
+            axes_dims_rope=(4, 4, 4, 4),
+        ),
+    )
+    # Need dtypes to be consistent; for now we just put it on bfloat16
+    pipe.to(torch.bfloat16).save_pretrained(model_dir)
+    return model_dir

tests/model_tests/diffusion/model_settings.py

@@ -0,0 +1,42 @@
+from tests.model_tests.diffusion import diff_model_builders
+from tests.model_tests.diffusion.config_types import (
+    DiffusionAccs,
+    DiffusionModelTestOpts,
+    DiffusionTasks,
+)
+
+# This object defines the (tiny) model configurations for common tests.
+#
+# TO ADD A NEW MODEL:
+# Map the pipeline name (i.e., the name of the pipeline in _DIFFUSION_MODELS)
+# to a DiffusionModelTestOpts and define the supported tasks. A base case with no
+# accelerations is always run. Use extra_test_groups to define additional
+# acceleration configurations to run a basic smoke test for (i.e., simple inference).
+#
+# The model will only be loaded once per test group, and will execute each of its
+# supported tasks as a pytest subtest. The base case also runs additional checks
+# that are independent of accelerations, e.g., for determinism and compatibility with
+# multi-output, unless explicitly disabled in the test settings.
+#
+# TO RUN A SUBSET OF THE TESTS:
+# These tests should be fast, but if you'd like to run a subset of the tests, you can do so
+# with `-k`; the IDs of the tests are the name of the pipeline joined by +.
+#
+# Example:
+# $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline]
+#   ^ Runs only the base (no accelerations) case for Flux2KleinPipeline
+#
+# $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline
+#   ^ Runs all test groups for Flux2KleinPipeline only
+DIFFUSION_TEST_SETTINGS = {
+    "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],
+        extra_test_groups=[
+            [DiffusionAccs.HSDP, DiffusionAccs.TEA_CACHE],
+            [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
+            [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
+        ],
+    )
+}

tests/model_tests/diffusion/task_runners.py

@@ -0,0 +1,172 @@
+"""
+Definitions running and validating individual tasks, e.g., text to image,
+image to image, and so on. These are called by the core test runner.
+"""
+
+import base64
+import io
+from dataclasses import replace
+
+import numpy as np
+from PIL import Image
+
+from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler, dummy_messages_from_mix_data
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+PROMPT = "Dummy prompt"
+IMAGE_DIMS = (512, 512)
+HEIGHT, WIDTH = IMAGE_DIMS
+INPUT_IMAGE = Image.new("RGB", IMAGE_DIMS)
+
+# Offline sampling params
+IMAGE_GEN_SAMPLING_PARAMS = OmniDiffusionSamplingParams(
+    num_inference_steps=4,
+    height=HEIGHT,
+    width=WIDTH,
+    seed=42,
+)
+
+# Online extra_body for diffusion requests
+IMAGE_GEN_EXTRA_BODY = {
+    "height": HEIGHT,
+    "width": WIDTH,
+    "num_inference_steps": 4,
+    "seed": 42,
+}
+
+
+### Shared validation
+def _validate_images(images: list[Image.Image], expected_n: int = 1):
+    """Given a set of images, ensure we got the expected count of images,
+    and that all provided images match the expected dimensions."""
+    assert len(images) == expected_n
+    for img in images:
+        assert isinstance(img, Image.Image)
+        assert img.size == IMAGE_DIMS
+    return images
+
+
+def _validate_image_gen_determinism(images_a: list[Image.Image], images_b: list[Image.Image]):
+    """Ensure that two image sets are valid and that the results match."""
+    _validate_images(images_a)
+    _validate_images(images_b)
+    assert np.array_equal(np.array(images_a[0]), np.array(images_b[0]))
+
+
+### Output extractor utils for offline / online paths respectively
+def _get_offline_images(outputs: list[OmniRequestOutput]) -> list[Image.Image]:
+    """Extract the images from an Omni .generate() call."""
+    assert len(outputs) == 1
+    return outputs[0].images
+
+
+def _get_online_images(responses: list[DiffusionResponse]) -> list[Image.Image]:
+    """Extract the images from a server response."""
+    assert len(responses) == 1
+    images = responses[0].images
+    assert images is not None
+    return images
+
+
+### Offline helpers
+def _run_offline_t2i(omni: Omni, params: OmniDiffusionSamplingParams = IMAGE_GEN_SAMPLING_PARAMS):
+    return omni.generate({"prompt": PROMPT}, params)
+
+
+def _run_offline_i2i(omni: Omni):
+    return omni.generate(
+        {"prompt": PROMPT, "multi_modal_data": {"image": INPUT_IMAGE}},
+        IMAGE_GEN_SAMPLING_PARAMS,
+    )
+
+
+### Online helpers
+def _build_online_image_data_url() -> str:
+    """Get a valid base 64 encoded data URL corresponding to an image."""
+    buf = io.BytesIO()
+    INPUT_IMAGE.save(buf, format="PNG")
+    b64 = base64.b64encode(buf.getvalue()).decode()
+    return f"data:image/png;base64,{b64}"
+
+
+def _run_online_t2i(
+    server: OmniServer, client: OpenAIClientHandler, extra_body: dict | None = None
+) -> list[DiffusionResponse]:
+    """Run a text to image request through the server."""
+    messages = dummy_messages_from_mix_data(content_text=PROMPT)
+    request_config = {
+        "model": server.model,
+        "messages": messages,
+        "extra_body": extra_body or IMAGE_GEN_EXTRA_BODY,
+    }
+    return client.send_diffusion_request(request_config)
+
+
+def _run_online_i2i(server: OmniServer, client: OpenAIClientHandler) -> list[DiffusionResponse]:
+    """Run an image to image request through the server."""
+    image_data_url = _build_online_image_data_url()
+    messages = dummy_messages_from_mix_data(
+        content_text=PROMPT,
+        image_data_url=image_data_url,
+    )
+    request_config = {
+        "model": server.model,
+        "messages": messages,
+        "extra_body": IMAGE_GEN_EXTRA_BODY,
+    }
+    return client.send_diffusion_request(request_config)
+
+
+### Offline task runners
+def run_and_validate_text_to_image_request(omni: Omni):
+    """Run and validate a text to image request."""
+    _validate_images(_get_offline_images(_run_offline_t2i(omni)))
+
+
+def run_and_validate_image_to_image_request(omni: Omni):
+    """Run and validate an image to image request."""
+    _validate_images(_get_offline_images(_run_offline_i2i(omni)))
+
+
+def run_and_validate_text_to_image_determinism(omni: Omni):
+    """Checks for determinism; for now we just keep this for TTI."""
+    _validate_image_gen_determinism(
+        _get_offline_images(_run_offline_t2i(omni)),
+        _get_offline_images(_run_offline_t2i(omni)),
+    )
+
+
+def run_and_validate_text_to_image_multi_output(omni: Omni):
+    """Checks for multi-output; for now we just keep this for TTI."""
+    params = replace(IMAGE_GEN_SAMPLING_PARAMS, num_outputs_per_prompt=2)
+    _validate_images(_get_offline_images(_run_offline_t2i(omni, params)), expected_n=2)
+
+
+### Online task runners
+def run_and_validate_online_text_to_image_request(server: OmniServer, client: OpenAIClientHandler):
+    """Run and validate a text to image request through the server."""
+    _validate_images(_get_online_images(_run_online_t2i(server, client)))
+
+
+def run_and_validate_online_image_to_image_request(server: OmniServer, client: OpenAIClientHandler):
+    """Run and validate an image to image request through the server."""
+    _validate_images(_get_online_images(_run_online_i2i(server, client)))
+
+
+def run_and_validate_online_text_to_image_determinism(server: OmniServer, client: OpenAIClientHandler):
+    """Checks for determinism through the server; for now we just keep this for TTI."""
+    _validate_image_gen_determinism(
+        _get_online_images(_run_online_t2i(server, client)),
+        _get_online_images(_run_online_t2i(server, client)),
+    )
+
+
+def run_and_validate_online_text_to_image_multi_output(server: OmniServer, client: OpenAIClientHandler):
+    """Checks for multi-output through the server; for now we just keep this for TTI."""
+    extra_body = {**IMAGE_GEN_EXTRA_BODY, "num_outputs_per_prompt": 2}
+    _validate_images(
+        _get_online_images(_run_online_t2i(server, client, extra_body=extra_body)),
+        expected_n=2,
+    )

tests/model_tests/diffusion/test_alignment.py

@@ -0,0 +1,98 @@
+"""
+Tests for alignment between _DIFFUSION_MODELS and DIFFUSION_TEST_SETTINGS; if
+tests in this file are failing, you are probably adding a new model, and need
+to account for it in DIFFUSION_TEST_SETTINGS
+"""
+
+import pytest
+
+from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
+from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+# These diffusion models currently do not have any added (tiny) model tests.
+# In general, pipelines should only be in the list below if they have not been
+# migrated yet or there is a compelling reason that they should not be (which for
+# now may just be CI pressure etc due to the volume of models and the cost of
+# initializing the server).
+#
+# The tests here validate that all keys in the diffusion registry are accounted for;
+# if you are adding a new model and see the tests below fail, please follow the pattern
+# for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS.
+EXCLUDED_MODELS = [
+    "QwenImagePipeline",
+    "QwenImageEditPipeline",
+    "QwenImageEditPlusPipeline",
+    "QwenImageLayeredPipeline",
+    "GlmImagePipeline",
+    "ZImagePipeline",
+    "OvisImagePipeline",
+    "WanPipeline",
+    "WanVACEPipeline",
+    "LTX2Pipeline",
+    "LTX2ImageToVideoPipeline",
+    "LTX2TwoStagesPipeline",
+    "LTX2ImageToVideoTwoStagesPipeline",
+    "LTX2T2VDMD2Pipeline",
+    "LTX2I2VDMD2Pipeline",
+    "LTX23Pipeline",
+    "LTX23ImageToVideoPipeline",
+    "StableAudioPipeline",
+    "WanImageToVideoPipeline",
+    "WanS2VPipeline",
+    "WanT2VDMD2Pipeline",
+    "WanI2VDMD2Pipeline",
+    "LongCatImagePipeline",
+    "BagelPipeline",
+    "LancePipeline",
+    "MingImagePipeline",
+    "InternVLAA1Pipeline",
+    "LongCatImageEditPipeline",
+    "StableDiffusion3Pipeline",
+    "FluxKontextPipeline",
+    "HunyuanImage3ForCausalMM",
+    "ErnieImagePipeline",
+    "NextStep11Pipeline",
+    "FluxPipeline",
+    "FluxDMD2Pipeline",
+    "QwenImageDMD2Pipeline",
+    "OmniGen2Pipeline",
+    "HeliosPipeline",
+    "HeliosPyramidPipeline",
+    "Flux2Pipeline",
+    "DreamIDOmniPipeline",
+    "SenseNovaU1Pipeline",
+    "AudioXPipeline",
+    "HunyuanVideo15Pipeline",
+    "HunyuanVideo15ImageToVideoPipeline",
+    "MagiHumanPipeline",
+    "OmniVoicePipeline",
+    "OmniVoice",
+    "Cosmos3OmniDiffusersPipeline",
+    "DiffusersAdapterPipeline",
+    "HiDreamImagePipeline",
+    "DreamZeroPipeline",
+    "StableDiffusionXLPipeline",
+    "Gr00tN1d7Pipeline",
+    "SoulXSingerPipeline",
+    "SoulXSingerSVCPipeline",
+]
+
+
+def test_all_non_excluded_pipelines_have_tests():
+    """Ensure that every pipeline in the diffusion registry either has a test
+    configuration in DIFFUSION_TEST_SETTINGS, or is explicitly excluded for a
+    known reason."""
+    non_excluded_diff_pipes = _DIFFUSION_MODELS.keys() - set(EXCLUDED_MODELS)
+    missing_tests = non_excluded_diff_pipes - DIFFUSION_TEST_SETTINGS.keys()
+    assert len(missing_tests) == 0, f"Pipelines missing test settings: {missing_tests}"
+
+
+def test_no_excluded_models_have_test_settings():
+    """Ensure that no models in the exclude list have a test configuration in
+    DIFFUSION_TEST_SETTINGS (i.e., if they have test settings, we should not
+    exclude them)."""
+    non_excluded_diff_pipes = _DIFFUSION_MODELS.keys() - set(EXCLUDED_MODELS)
+    missing_pipes = DIFFUSION_TEST_SETTINGS.keys() - set(non_excluded_diff_pipes)
+    assert len(missing_pipes) == 0

tests/model_tests/diffusion/test_common_offline.py

@@ -0,0 +1,69 @@
+import pytest
+
+from tests.model_tests.diffusion.case_filtering import get_parametrized_options
+from tests.model_tests.diffusion.config_types import (
+    DiffusionAccs,
+    DiffusionTasks,
+    build_omni_from_diff_accelerations,
+)
+from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
+from tests.model_tests.diffusion.task_runners import (
+    run_and_validate_image_to_image_request,
+    run_and_validate_text_to_image_determinism,
+    run_and_validate_text_to_image_multi_output,
+    run_and_validate_text_to_image_request,
+)
+
+# NOTE : Hardware marks are added dynamically based on test requirements
+pytestmark = [pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize(
+    "model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
+    get_parametrized_options(DIFFUSION_TEST_SETTINGS),
+)
+def test_pipeline_on_supported_tasks(
+    model_name,
+    accelerations: list[DiffusionAccs] | None,
+    supported_tasks: list[DiffusionTasks],
+    check_multioutput: bool,
+    check_determinism: bool,
+    tiny_model_paths: dict[str, str],
+    subtests,
+):
+    """Run a smoke test on all of the pipelines supported tasks using a set of enabled accelerations."""
+    assert len(supported_tasks) > 0
+    # We initialize the Omni object before running the tasks, then run each task as a pytest subtask.
+    # This lets us init the model once, but display separate failures in pytest, and avoid halting the
+    # checks on other tasks if one fails.
+    #
+    # This allows us to have some degree of test isolation without the cost of redundant initialization,
+    # since starting the server can take 10+ seconds, even for tiny models.
+    #
+    # NOTE: Be sure to install pytest-subtests if you're running on pytest < 9
+    omni = build_omni_from_diff_accelerations(
+        accelerations=accelerations,
+        model=tiny_model_paths[model_name],
+        enforce_eager=True,
+    )
+    try:
+        for task_type in supported_tasks:
+            with subtests.test(msg=task_type):
+                if task_type == DiffusionTasks.TEXT_TO_IMAGE:
+                    run_and_validate_text_to_image_request(omni)
+                elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
+                    run_and_validate_image_to_image_request(omni)
+                else:
+                    raise ValueError(f"Task type {task_type} is not yet supported")
+
+        # NOTE: For now, we only check determinism + multi output for the base case,
+        # since checking it on every extra acceleration configuration is redundant
+        # (see case_filtering).
+        if check_determinism:
+            with subtests.test(msg="determinism"):
+                run_and_validate_text_to_image_determinism(omni)
+        if check_multioutput:
+            with subtests.test(msg="multi_output"):
+                run_and_validate_text_to_image_multi_output(omni)
+    finally:
+        omni.close()

tests/model_tests/diffusion/test_common_online.py

@@ -0,0 +1,75 @@
+"""
+Analogous to test_common_offline, but for server tests. Validates the full
+online serving stack (CLI arg parsing, subprocess, API routing, response
+encoding) using tiny models.
+"""
+
+import pytest
+
+from tests.helpers.runtime import OmniServer, OpenAIClientHandler
+from tests.model_tests.diffusion.case_filtering import get_parametrized_options
+from tests.model_tests.diffusion.config_types import (
+    DiffusionAccs,
+    DiffusionTasks,
+    build_server_args_from_diff_accelerations,
+)
+from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
+from tests.model_tests.diffusion.task_runners import (
+    run_and_validate_online_image_to_image_request,
+    run_and_validate_online_text_to_image_determinism,
+    run_and_validate_online_text_to_image_multi_output,
+    run_and_validate_online_text_to_image_request,
+)
+
+# NOTE : Hardware marks are added dynamically based on test requirements
+pytestmark = [pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize(
+    "model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
+    get_parametrized_options(DIFFUSION_TEST_SETTINGS, online=True),
+)
+def test_online_on_supported_tasks(
+    model_name: str,
+    accelerations: list[DiffusionAccs] | None,
+    supported_tasks: list[DiffusionTasks],
+    check_multioutput: bool,
+    check_determinism: bool,
+    tiny_model_paths: dict[str, str],
+    run_level: str,
+    subtests,
+):
+    """Smoke test: start a tiny model server and run each supported task via the API."""
+    model_path = tiny_model_paths[model_name]
+    server_args = build_server_args_from_diff_accelerations(accelerations)
+    server_args.append("--enforce-eager")
+
+    with OmniServer(model_path, server_args) as server:
+        # TODO: We may want to revisit run_level validation here,
+        # because checks for things like image size etc should not
+        # depend on whether or not the weights are real or random
+        client = OpenAIClientHandler(
+            host=server.host,
+            port=server.port,
+            api_key="EMPTY",
+            run_level=run_level,
+            log_stats=server.log_stats,
+        )
+        for task_type in supported_tasks:
+            with subtests.test(msg=task_type):
+                if task_type == DiffusionTasks.TEXT_TO_IMAGE:
+                    run_and_validate_online_text_to_image_request(server, client)
+                elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
+                    run_and_validate_online_image_to_image_request(server, client)
+                else:
+                    raise ValueError(f"Task type {task_type} is not yet supported")
+
+        # NOTE: For now, we only check determinism + multi output for the base case,
+        # since checking it on every extra acceleration configuration is redundant
+        # (see case_filtering).
+        if check_determinism:
+            with subtests.test(msg="determinism"):
+                run_and_validate_online_text_to_image_determinism(server, client)
+        if check_multioutput:
+            with subtests.test(msg="multi_output"):
+                run_and_validate_online_text_to_image_multi_output(server, client)

tests/model_tests/diffusion/utils.py

@@ -0,0 +1,33 @@
+"""
+Utilities for resolving real models to their tiny model equivalents.
+"""
+
+import logging
+
+from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
+from vllm_omni.diffusion.data import resolve_model_class_name
+
+logger = logging.getLogger(__name__)
+
+
+def resolve_tiny_model_path(model: str) -> str:
+    """Given a real model name/path, resolve it to a tiny model path.
+
+    Raises ValueError if the pipeline class cannot be determined (invalid
+    model). Returns the original model path if no tiny builder exists yet."""
+    pipeline_class = resolve_model_class_name(model)
+    if pipeline_class is None:
+        raise ValueError(
+            f"Cannot resolve pipeline class for model: {model}. The model path may be invalid or its config unreadable."
+        )
+
+    test_opts = DIFFUSION_TEST_SETTINGS.get(pipeline_class)
+    if test_opts is None:
+        logger.warning(
+            "No tiny model builder for pipeline %s (model: %s). Using original model.",
+            pipeline_class,
+            model,
+        )
+        return model
+
+    return test_opts.builder()