[Tests] Further extend the tiny model testing framework and enable multiple models (#5353)

Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/e2e/online_serving/test_qwen_image.py

@@ -58,7 +58,6 @@ def _get_default_case(model: str):
     ]
 
 
-@pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
 @pytest.mark.parametrize("omni_server", _get_default_case(MODEL), indirect=True)

tests/helpers/tiny_model.py

@@ -4,11 +4,13 @@ Mirrors the component loading loop in DiffusionPipeline.from_pretrained,
 but monkeypatches from_pretrained on ModelMixin and PreTrainedModel to
 initialize with random weights instead of loading checkpoint files.
 Components with vendored configs use those; others load configs from
-the upstream HF model.
+the upstream HF model, optionally passed through a per-component transform
+function first.
 """
 
 import os
 import tempfile
+from collections.abc import Callable
 from pathlib import Path
 from unittest.mock import patch
 
@@ -16,7 +18,7 @@ import torch
 from diffusers import ModelMixin
 from diffusers.pipelines.pipeline_loading_utils import _get_pipeline_class, simple_get_class_obj
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from transformers import AutoConfig, PreTrainedModel
+from transformers import PretrainedConfig, PreTrainedModel
 
 TINY_MODEL_DIR = os.path.join(tempfile.gettempdir(), "vllm-omni-tiny-models")
 
@@ -27,39 +29,73 @@ def _get_tiny_model_path(name: str) -> str:
     return path
 
 
-def _diffusers_from_config(cls, pretrained_model_name_or_path, **kwargs):
-    """Replacement for ModelMixin.from_pretrained that initializes random weights."""
-    subfolder = kwargs.get("subfolder")
-    load_kwargs = {"subfolder": subfolder} if subfolder else {}
-    config = cls.load_config(pretrained_model_name_or_path, **load_kwargs)
-    return cls.from_config(config)
-
-
-def _transformers_from_config(cls, pretrained_model_name_or_path, **kwargs):
-    """Replacement for PreTrainedModel.from_pretrained that initializes random weights."""
-    subfolder = kwargs.get("subfolder")
-    load_kwargs = {"subfolder": subfolder} if subfolder else {}
-    config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **load_kwargs)
-    return cls(config)
-
-
-def build_tiny_from_configs(pipeline_name: str, model_id: str, configs_dir: str | Path) -> str:
-    """Build a tiny diffusion model from vendored configs with random weights.
+def build_tiny_from_configs(
+    pipeline_name: str,
+    model_id: str,
+    configs_dir: str | Path | None = None,
+    transform: dict[str, Callable[[dict], dict]] | None = None,
+) -> str:
+    """Build a tiny diffusion model with random weights.
 
     Args:
         pipeline_name: Name of the pipeline (used as output directory name).
         model_id: HuggingFace model ID for loading upstream components
-            (tokenizer, scheduler) that don't have vendored configs.
-        configs_dir: Path to the directory containing vendored config files
+            (tokenizer, scheduler, and any component without a vendored
+            config or transform) that don't have vendored configs.
+        configs_dir: Path to a directory containing vendored config files
             (model_index.json and per-component config.json).
+        transform: Maps a component (subfolder) name to a function that
+            takes that component's raw config dict (fetched from model_id)
+            and returns a shrunk one, applied only when that component has
+            no vendored config in configs_dir. Lets each pipeline encode
+            exactly the shrink it needs (layer counts, and any component-
+            internal width that isn't read by another component), including
+            keeping sibling config fields in sync (e.g. layer_types matching
+            num_hidden_layers).
 
     Returns:
         Path to the saved tiny model directory with safetensors weights.
     """
     model_dir = _get_tiny_model_path(pipeline_name)
-    config_dir = Path(configs_dir)
+    config_dir = Path(configs_dir) if configs_dir is not None else None
+    transform = transform or {}
 
-    config_dict = DiffusionPipeline.load_config(config_dir)
+    if config_dir is None and not transform:
+        raise ValueError(
+            f"Neither configs_dir nor transform given for {pipeline_name!r}. Pass a "
+            "configs_dir with vendored configs, and/or transform={{name: fn}} to shrink "
+            f"specific components fetched from the upstream model {model_id!r}."
+        )
+
+    def _diffusers_from_config(cls, pretrained_model_name_or_path, **kwargs):
+        """Replacement for ModelMixin.from_pretrained that initializes random weights."""
+        subfolder = kwargs.get("subfolder")
+        load_kwargs = {"subfolder": subfolder} if subfolder else {}
+        config = cls.load_config(pretrained_model_name_or_path, **load_kwargs)
+        if subfolder in transform:
+            config = transform[subfolder](dict(config))
+        return cls.from_config(config)
+
+    def _transformers_from_config(cls, pretrained_model_name_or_path, **kwargs):
+        """Replacement for PreTrainedModel.from_pretrained that initializes random weights.
+
+        Uses the model's own config_class (every PreTrainedModel subclass
+        declares one) so the raw dict can be transformed before constructing
+        the config object.
+        """
+        subfolder = kwargs.get("subfolder")
+        load_kwargs = {"subfolder": subfolder} if subfolder else {}
+        config_dict, _ = PretrainedConfig.get_config_dict(pretrained_model_name_or_path, **load_kwargs)
+        if subfolder in transform:
+            config_dict = transform[subfolder](config_dict)
+        config = cls.config_class.from_dict(config_dict)
+        return cls(config)
+
+    # model_index.json comes from the vendored configs_dir when it's there, otherwise
+    # from model_id. Independent of the per-component checks below, so a configs_dir
+    # can vendor just one component without also needing its own model_index.json.
+    vendored_index = config_dir is not None and (config_dir / "model_index.json").exists()
+    config_dict = DiffusionPipeline.load_config(config_dir if vendored_index else model_id)
     pipeline_cls = _get_pipeline_class(DiffusionPipeline, config=config_dict)
 
     init_dict, _, _ = pipeline_cls.extract_init_dict(config_dict)
@@ -77,7 +113,7 @@ def build_tiny_from_configs(pipeline_name: str, model_id: str, configs_dir: str
             cls = simple_get_class_obj(library_name, class_name)
 
             # Use vendored config dir if available, otherwise upstream model
-            if (config_dir / name).exists():
+            if config_dir is not None and (config_dir / name).exists():
                 init_kwargs[name] = cls.from_pretrained(config_dir / name)
             else:
                 init_kwargs[name] = cls.from_pretrained(model_id, subfolder=name)

tests/model_tests/diffusion/config_types.py

@@ -35,6 +35,7 @@ class DiffusionTasks(str, Enum):
     TEXT_TO_IMAGE = "text_to_image"
     IMAGE_TO_IMAGE = "image_to_image"
     TEXT_TO_VIDEO = "text_to_video"
+    IMAGE_TO_VIDEO = "image_to_video"
     # Text to audio, etc should be added here as needed
 
 

tests/model_tests/diffusion/diff_model_builders.py

@@ -1,3 +1,4 @@
+from functools import partial
 from pathlib import Path
 
 from tests.helpers.tiny_model import build_tiny_from_configs
@@ -5,6 +6,25 @@ from tests.helpers.tiny_model import build_tiny_from_configs
 TINY_CONFIGS_DIR = Path(__file__).parent / "tiny_configs"
 
 
+def _shrink_dit_rope_config(
+    config: dict,
+    num_layers: int = 2,
+    num_single_layers: int | None = None,
+    target_head_dim: int = 32,
+    target_heads: int = 4,
+    default_axes_dims_rope: list[int] | None = None,
+) -> dict:
+    config["num_layers"] = num_layers
+    if num_single_layers is not None:
+        config["num_single_layers"] = num_single_layers
+    axes_dims_rope = config.get("axes_dims_rope", default_axes_dims_rope)
+    factor = config["attention_head_dim"] / target_head_dim
+    config["attention_head_dim"] = target_head_dim
+    config["num_attention_heads"] = target_heads
+    config["axes_dims_rope"] = [int(d / factor) for d in axes_dims_rope]
+    return config
+
+
 def tiny_flux2_klein_builder() -> str:
     """Build a tiny Flux2Klein model from vendored configs."""
     return build_tiny_from_configs(
@@ -15,3 +35,72 @@ def tiny_flux2_klein_builder() -> str:
 def tiny_ltx2_builder() -> str:
     """Build a tiny LTX2 model from vendored configs."""
     return build_tiny_from_configs("LTX2Pipeline", "Lightricks/LTX-2", TINY_CONFIGS_DIR / "LTX2Pipeline")
+
+
+def tiny_qwen_image_builder() -> str:
+    def shrink_text_encoder(config: dict) -> dict:
+        config["num_hidden_layers"] = 2
+        config["intermediate_size"] = 64
+        config["text_config"]["num_hidden_layers"] = 2
+        config["text_config"]["intermediate_size"] = 64
+        config["text_config"]["layer_types"] = config["text_config"]["layer_types"][:2]
+        config["vision_config"]["depth"] = 2
+        config["vision_config"]["intermediate_size"] = 64
+        config["vision_config"]["fullatt_block_indexes"] = [0, 1]
+        return config
+
+    return build_tiny_from_configs(
+        "QwenImagePipeline",
+        "Qwen/Qwen-Image",
+        transform={"text_encoder": shrink_text_encoder, "transformer": _shrink_dit_rope_config},
+    )
+
+
+def tiny_flux_builder() -> str:
+    def shrink_clip_text_encoder(config: dict) -> dict:
+        config["num_hidden_layers"] = 2
+        return config
+
+    def shrink_t5_text_encoder(config: dict) -> dict:
+        config["num_layers"] = 2
+        config["num_decoder_layers"] = 2
+        config["d_ff"] = 64
+        config["num_heads"] = 4
+        return config
+
+    return build_tiny_from_configs(
+        "FluxPipeline",
+        "black-forest-labs/FLUX.1-schnell",
+        transform={
+            "text_encoder": shrink_clip_text_encoder,
+            "text_encoder_2": shrink_t5_text_encoder,
+            "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, default_axes_dims_rope=[16, 56, 56]),
+        },
+    )
+
+
+def tiny_flux2_builder() -> str:
+    def shrink_text_encoder(config: dict) -> dict:
+        # The real checkpoint ships language_model.lm_head.weight as its own
+        # tensor even though tie_word_embeddings defaults to True; keep it
+        # untied here too so transformers doesn't drop it on save.
+        config["tie_word_embeddings"] = False
+        config["text_config"]["num_hidden_layers"] = 31
+        config["text_config"]["intermediate_size"] = 64
+        config["text_config"]["num_attention_heads"] = 4
+        config["text_config"]["head_dim"] = 32
+        config["text_config"]["num_key_value_heads"] = 2
+        config["vision_config"]["num_hidden_layers"] = 2
+        config["vision_config"]["intermediate_size"] = 64
+        config["vision_config"]["num_attention_heads"] = 4
+        config["vision_config"]["head_dim"] = 32
+        return config
+
+    return build_tiny_from_configs(
+        "Flux2Pipeline",
+        "black-forest-labs/FLUX.2-dev",
+        transform={
+            "text_encoder": shrink_text_encoder,
+            "transformer": partial(_shrink_dit_rope_config, num_single_layers=2),
+        },
+    )

tests/model_tests/diffusion/model_settings.py

@@ -42,13 +42,26 @@ DIFFUSION_TEST_SETTINGS = {
     "LTX2Pipeline": DiffusionModelTestOpts(
         model="Lightricks/LTX-2",
         builder=diff_model_builders.tiny_ltx2_builder,
-        supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO],
-        check_determinism=False,
-        check_multi_output=False,
+        supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO, DiffusionTasks.IMAGE_TO_VIDEO],
         extra_test_groups=[
             [DiffusionAccs.HSDP, DiffusionAccs.CACHE_DIT],
             [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
             [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
         ],
     ),
+    "QwenImagePipeline": DiffusionModelTestOpts(
+        model="Qwen/Qwen-Image",
+        builder=diff_model_builders.tiny_qwen_image_builder,
+        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
+    ),
+    "FluxPipeline": DiffusionModelTestOpts(
+        model="black-forest-labs/FLUX.1-schnell",
+        builder=diff_model_builders.tiny_flux_builder,
+        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
+    ),
+    "Flux2Pipeline": DiffusionModelTestOpts(
+        model="black-forest-labs/FLUX.2-dev",
+        builder=diff_model_builders.tiny_flux2_builder,
+        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
+    ),
 }

tests/model_tests/diffusion/task_runners.py

@@ -8,9 +8,11 @@ import io
 from dataclasses import replace
 
 import numpy as np
+import pytest
 from PIL import Image
 
 from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.model_tests.diffusion.config_types import DiffusionTasks
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -70,7 +72,9 @@ def _validate_image_gen_determinism(images_a: list[Image.Image], images_b: list[
     """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]))
+    assert len(images_a) == len(images_b)
+    for image_a, image_b in zip(images_a, images_b, strict=True):
+        assert np.array_equal(np.array(image_a), np.array(image_b))
 
 
 ### Output extractor utils for offline / online paths respectively
@@ -88,20 +92,50 @@ def _get_online_images(responses: list[DiffusionResponse]) -> list[Image.Image]:
     return images
 
 
+def _validate_video_frames(videos: list, expected_n: int = 1, expected_frames: int = VIDEO_NUM_FRAMES):
+    """Given a list of videos (one array per generation, each holding all of
+    that generation's frames), ensure count and shape are correct."""
+    assert len(videos) == expected_n
+    for video in videos:
+        assert isinstance(video, np.ndarray)
+        assert video.ndim == 4, f"Expected 4D video array (frames, H, W, C), got shape {video.shape}"
+        assert video.shape[-4] == expected_frames, f"Expected {expected_frames} frames, got {video.shape[-4]}"
+        assert video.shape[-3] == HEIGHT, f"Expected height {HEIGHT}, got {video.shape[-3]}"
+        assert video.shape[-2] == WIDTH, f"Expected width {WIDTH}, got {video.shape[-2]}"
+        assert video.shape[-1] == 3, f"Expected 3 channels (RGB), got {video.shape[-1]}"
+    return videos
+
+
+def _get_offline_videos(outputs: list[OmniRequestOutput]) -> list:
+    """Extract the list of videos (one array per generation) from an Omni .generate() call.
+
+    Multi-output video generations come back as a single 5D array
+    (num_outputs, num_frames, H, W, C) rather than one 4D array per output;
+    split it into a flat per-output list here so callers always see one
+    4D array per generation.
+    """
+    assert len(outputs) == 1
+    videos = []
+    for video in outputs[0].images:
+        if video.ndim == 5:
+            videos.extend(video)
+        else:
+            videos.append(video)
+    return videos
+
+
 def _validate_video(outputs: list[OmniRequestOutput], expected_n: int = 1):
     """Given a set of outputs, ensure we got video frames with the expected shape."""
-    assert len(outputs) == expected_n
-    for output in outputs:
-        # Video models return numpy arrays via output.images
-        images = output.images
-        assert len(images) > 0
-        for frame_data in images:
-            assert isinstance(frame_data, np.ndarray)
-            # (num_outputs, num_frames, H, W, C) or (num_frames, H, W, C)
-            assert frame_data.ndim in (4, 5), f"Expected 4D or 5D video array, got shape {frame_data.shape}"
-            assert frame_data.shape[-3] == HEIGHT, f"Expected height {HEIGHT}, got {frame_data.shape[-3]}"
-            assert frame_data.shape[-2] == WIDTH, f"Expected width {WIDTH}, got {frame_data.shape[-2]}"
-            assert frame_data.shape[-1] == 3, f"Expected 3 channels (RGB), got {frame_data.shape[-1]}"
+    _validate_video_frames(_get_offline_videos(outputs), expected_n=expected_n)
+
+
+def _validate_video_gen_determinism(videos_a: list, videos_b: list):
+    """Ensure that two video generations are valid and that the results match."""
+    _validate_video_frames(videos_a)
+    _validate_video_frames(videos_b)
+    assert len(videos_a) == len(videos_b)
+    for video_a, video_b in zip(videos_a, videos_b, strict=True):
+        assert np.array_equal(video_a, video_b)
 
 
 ### Offline helpers
@@ -113,10 +147,17 @@ def _run_offline_t2v(omni: Omni, params: OmniDiffusionSamplingParams = VIDEO_GEN
     return omni.generate({"prompt": PROMPT}, params)
 
 
-def _run_offline_i2i(omni: Omni):
+def _run_offline_i2i(omni: Omni, params: OmniDiffusionSamplingParams = IMAGE_GEN_SAMPLING_PARAMS):
     return omni.generate(
         {"prompt": PROMPT, "multi_modal_data": {"image": INPUT_IMAGE}},
-        IMAGE_GEN_SAMPLING_PARAMS,
+        params,
+    )
+
+
+def _run_offline_i2v(omni: Omni, params: OmniDiffusionSamplingParams = VIDEO_GEN_SAMPLING_PARAMS):
+    return omni.generate(
+        {"prompt": PROMPT, "multi_modal_data": {"image": INPUT_IMAGE}},
+        params,
     )
 
 
@@ -142,7 +183,9 @@ def _run_online_t2i(
     return client.send_diffusion_request(request_config)
 
 
-def _run_online_i2i(server: OmniServer, client: OpenAIClientHandler) -> list[DiffusionResponse]:
+def _run_online_i2i(
+    server: OmniServer, client: OpenAIClientHandler, extra_body: dict | None = None
+) -> 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(
@@ -152,7 +195,7 @@ def _run_online_i2i(server: OmniServer, client: OpenAIClientHandler) -> list[Dif
     request_config = {
         "model": server.model,
         "messages": messages,
-        "extra_body": IMAGE_GEN_EXTRA_BODY,
+        "extra_body": extra_body or IMAGE_GEN_EXTRA_BODY,
     }
     return client.send_diffusion_request(request_config)
 
@@ -173,18 +216,65 @@ def run_and_validate_text_to_video_request(omni: Omni):
     _validate_video(_run_offline_t2v(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_image_to_video_request(omni: Omni):
+    """Run and validate an image to video request.
+
+    LTX2Pipeline uses a unified text/image entry and can fall back to T2V if
+    the image never reaches conditioning, so also run T2V with the same seed
+    and prompt and assert the outputs differ; a regression that silently
+    drops the input image would otherwise still pass a "did we get a valid
+    video back" check alone.
+    """
+    i2v_video = _validate_video_frames(_get_offline_videos(_run_offline_i2v(omni)))[0]
+    t2v_video = _validate_video_frames(_get_offline_videos(_run_offline_t2v(omni)))[0]
+    assert not np.array_equal(i2v_video, t2v_video), (
+        "I2V output is identical to T2V output with the same seed and prompt; "
+        "the input image may not be reaching conditioning."
     )
 
 
-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)
+def run_and_validate_determinism(omni: Omni, task_type: DiffusionTasks):
+    """Checks for determinism, dispatching by task type."""
+    if task_type == DiffusionTasks.TEXT_TO_IMAGE:
+        _validate_image_gen_determinism(
+            _get_offline_images(_run_offline_t2i(omni)),
+            _get_offline_images(_run_offline_t2i(omni)),
+        )
+    elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
+        _validate_image_gen_determinism(
+            _get_offline_images(_run_offline_i2i(omni)),
+            _get_offline_images(_run_offline_i2i(omni)),
+        )
+    elif task_type == DiffusionTasks.TEXT_TO_VIDEO:
+        _validate_video_gen_determinism(
+            _get_offline_videos(_run_offline_t2v(omni)),
+            _get_offline_videos(_run_offline_t2v(omni)),
+        )
+    elif task_type == DiffusionTasks.IMAGE_TO_VIDEO:
+        _validate_video_gen_determinism(
+            _get_offline_videos(_run_offline_i2v(omni)),
+            _get_offline_videos(_run_offline_i2v(omni)),
+        )
+    else:
+        raise ValueError(f"Task type {task_type} is not yet supported for determinism checks")
+
+
+def run_and_validate_multi_output(omni: Omni, task_type: DiffusionTasks):
+    """Checks for multi-output, dispatching by task type."""
+    if task_type == DiffusionTasks.TEXT_TO_IMAGE:
+        params = replace(IMAGE_GEN_SAMPLING_PARAMS, num_outputs_per_prompt=2)
+        _validate_images(_get_offline_images(_run_offline_t2i(omni, params)), expected_n=2)
+    elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
+        params = replace(IMAGE_GEN_SAMPLING_PARAMS, num_outputs_per_prompt=2)
+        _validate_images(_get_offline_images(_run_offline_i2i(omni, params)), expected_n=2)
+    elif task_type == DiffusionTasks.TEXT_TO_VIDEO:
+        params = replace(VIDEO_GEN_SAMPLING_PARAMS, num_outputs_per_prompt=2)
+        _validate_video_frames(_get_offline_videos(_run_offline_t2v(omni, params)), expected_n=2)
+    elif task_type == DiffusionTasks.IMAGE_TO_VIDEO:
+        params = replace(VIDEO_GEN_SAMPLING_PARAMS, num_outputs_per_prompt=2)
+        _validate_video_frames(_get_offline_videos(_run_offline_i2v(omni, params)), expected_n=2)
+    else:
+        raise ValueError(f"Task type {task_type} is not yet supported for multi-output checks")
 
 
 def _run_online_t2v(
@@ -197,6 +287,15 @@ def _run_online_t2v(
     return client.send_video_diffusion_request({"form_data": data})
 
 
+def _run_online_i2v(server: OmniServer, client: OpenAIClientHandler) -> list[DiffusionResponse]:
+    """Run an image to video request through the server's /v1/videos API."""
+    data = dict(VIDEO_GEN_FORM_DATA)
+    data.setdefault("prompt", PROMPT)
+    data.setdefault("model", server.model)
+    image_data_url = _build_online_image_data_url()
+    return client.send_video_diffusion_request({"form_data": data, "image_reference": image_data_url})
+
+
 def _get_online_videos(responses: list[DiffusionResponse]) -> list:
     """Extract the videos from a server response."""
     assert len(responses) == 1
@@ -217,23 +316,63 @@ def run_and_validate_online_image_to_image_request(server: OmniServer, client: O
     _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_determinism(server: OmniServer, client: OpenAIClientHandler, task_type: DiffusionTasks):
+    """Checks for determinism through the server, dispatching by task type.
+
+    Video is skipped online only: send_video_diffusion_request returns encoded
+    video bytes, not a frame-level array, and re-encoding the same content
+    isn't guaranteed to be byte-identical even for deterministic generation,
+    so a reliable check isn't possible with what's implemented here. The
+    offline determinism check (run_and_validate_determinism) does cover
+    video, since it has direct access to the raw frame arrays.
+    """
+    if task_type == DiffusionTasks.TEXT_TO_IMAGE:
+        _validate_image_gen_determinism(
+            _get_online_images(_run_online_t2i(server, client)),
+            _get_online_images(_run_online_t2i(server, client)),
+        )
+    elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
+        _validate_image_gen_determinism(
+            _get_online_images(_run_online_i2i(server, client)),
+            _get_online_images(_run_online_i2i(server, client)),
+        )
+    elif task_type in (DiffusionTasks.TEXT_TO_VIDEO, DiffusionTasks.IMAGE_TO_VIDEO):
+        pytest.skip(f"Online determinism isn't supported for {task_type}.")
+    else:
+        raise ValueError(f"Task type {task_type} is not yet supported for determinism checks")
 
 
-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,
-    )
+def run_and_validate_online_multi_output(server: OmniServer, client: OpenAIClientHandler, task_type: DiffusionTasks):
+    """Checks for multi-output through the server, dispatching by task type.
+
+    Video is skipped online only: send_video_diffusion_request explicitly
+    disallows more than one request/output per call. The offline multi-output
+    check (run_and_validate_multi_output) does cover video, since
+    Omni.generate() supports num_outputs_per_prompt directly.
+    """
+    if task_type == DiffusionTasks.TEXT_TO_IMAGE:
+        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,
+        )
+    elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
+        extra_body = {**IMAGE_GEN_EXTRA_BODY, "num_outputs_per_prompt": 2}
+        _validate_images(
+            _get_online_images(_run_online_i2i(server, client, extra_body=extra_body)),
+            expected_n=2,
+        )
+    elif task_type in (DiffusionTasks.TEXT_TO_VIDEO, DiffusionTasks.IMAGE_TO_VIDEO):
+        pytest.skip(f"Online multi-output isn't supported for {task_type}.")
+    else:
+        raise ValueError(f"Task type {task_type} is not yet supported for multi-output checks")
 
 
 def run_and_validate_online_text_to_video_request(server: OmniServer, client: OpenAIClientHandler):
     """Run and validate a text to video request through the server."""
     _get_online_videos(_run_online_t2v(server, client))
+
+
+def run_and_validate_online_image_to_video_request(server: OmniServer, client: OpenAIClientHandler):
+    """Run and validate an image to video request through the server."""
+    _get_online_videos(_run_online_i2v(server, client))

tests/model_tests/diffusion/test_alignment.py

@@ -21,7 +21,6 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 # 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",
@@ -50,14 +49,12 @@ EXCLUDED_MODELS = [
     "HunyuanImage3ForCausalMM",
     "ErnieImagePipeline",
     "NextStep11Pipeline",
-    "FluxPipeline",
     "FluxDMD2Pipeline",
     "Krea2Pipeline",
     "QwenImageDMD2Pipeline",
     "OmniGen2Pipeline",
     "HeliosPipeline",
     "HeliosPyramidPipeline",
-    "Flux2Pipeline",
     "DreamIDOmniPipeline",
     "SenseNovaU1Pipeline",
     "AudioXPipeline",

tests/model_tests/diffusion/test_common_offline.py

@@ -8,9 +8,10 @@ from tests.model_tests.diffusion.config_types import (
 )
 from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
 from tests.model_tests.diffusion.task_runners import (
+    run_and_validate_determinism,
     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_image_to_video_request,
+    run_and_validate_multi_output,
     run_and_validate_text_to_image_request,
     run_and_validate_text_to_video_request,
 )
@@ -56,6 +57,8 @@ def test_pipeline_on_supported_tasks(
                     run_and_validate_image_to_image_request(omni)
                 elif task_type == DiffusionTasks.TEXT_TO_VIDEO:
                     run_and_validate_text_to_video_request(omni)
+                elif task_type == DiffusionTasks.IMAGE_TO_VIDEO:
+                    run_and_validate_image_to_video_request(omni)
                 else:
                     raise ValueError(f"Task type {task_type} is not yet supported")
 
@@ -63,10 +66,12 @@ def test_pipeline_on_supported_tasks(
         # 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)
+            for task_type in supported_tasks:
+                with subtests.test(msg=f"determinism[{task_type}]"):
+                    run_and_validate_determinism(omni, task_type)
         if check_multioutput:
-            with subtests.test(msg="multi_output"):
-                run_and_validate_text_to_image_multi_output(omni)
+            for task_type in supported_tasks:
+                with subtests.test(msg=f"multi_output[{task_type}]"):
+                    run_and_validate_multi_output(omni, task_type)
     finally:
         omni.close()

tests/model_tests/diffusion/test_common_online.py

@@ -15,9 +15,10 @@ from tests.model_tests.diffusion.config_types import (
 )
 from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
 from tests.model_tests.diffusion.task_runners import (
+    run_and_validate_online_determinism,
     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_image_to_video_request,
+    run_and_validate_online_multi_output,
     run_and_validate_online_text_to_image_request,
     run_and_validate_online_text_to_video_request,
 )
@@ -64,6 +65,8 @@ def test_online_on_supported_tasks(
                     run_and_validate_online_image_to_image_request(server, client)
                 elif task_type == DiffusionTasks.TEXT_TO_VIDEO:
                     run_and_validate_online_text_to_video_request(server, client)
+                elif task_type == DiffusionTasks.IMAGE_TO_VIDEO:
+                    run_and_validate_online_image_to_video_request(server, client)
                 else:
                     raise ValueError(f"Task type {task_type} is not yet supported")
 
@@ -71,8 +74,10 @@ def test_online_on_supported_tasks(
         # 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)
+            for task_type in supported_tasks:
+                with subtests.test(msg=f"determinism[{task_type}]"):
+                    run_and_validate_online_determinism(server, client, task_type)
         if check_multioutput:
-            with subtests.test(msg="multi_output"):
-                run_and_validate_online_text_to_image_multi_output(server, client)
+            for task_type in supported_tasks:
+                with subtests.test(msg=f"multi_output[{task_type}]"):
+                    run_and_validate_online_multi_output(server, client, task_type)