diff --git a/tests/e2e/online_serving/test_qwen_image.py b/tests/e2e/online_serving/test_qwen_image.py index 454a87aa5..c40cfbda8 100644 --- a/tests/e2e/online_serving/test_qwen_image.py +++ b/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) diff --git a/tests/helpers/tiny_model.py b/tests/helpers/tiny_model.py index b212a8e7f..7ada9f987 100644 --- a/tests/helpers/tiny_model.py +++ b/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) diff --git a/tests/model_tests/diffusion/config_types.py b/tests/model_tests/diffusion/config_types.py index 49845fa65..2872f7cec 100644 --- a/tests/model_tests/diffusion/config_types.py +++ b/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 diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py index 54e7a4016..ea69d95e5 100644 --- a/tests/model_tests/diffusion/diff_model_builders.py +++ b/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), + }, + ) diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py index c6c1cb1d9..037faf67d 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/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], + ), } diff --git a/tests/model_tests/diffusion/task_runners.py b/tests/model_tests/diffusion/task_runners.py index 23ff60eeb..2be2338c1 100644 --- a/tests/model_tests/diffusion/task_runners.py +++ b/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)) diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 29cfc6f21..e1795ddd0 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/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", diff --git a/tests/model_tests/diffusion/test_common_offline.py b/tests/model_tests/diffusion/test_common_offline.py index d5cf3aba6..b0eb26c9b 100644 --- a/tests/model_tests/diffusion/test_common_offline.py +++ b/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() diff --git a/tests/model_tests/diffusion/test_common_online.py b/tests/model_tests/diffusion/test_common_online.py index 31509c3b3..464adc13a 100644 --- a/tests/model_tests/diffusion/test_common_online.py +++ b/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)