# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ 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 from typing import Any import numpy as np import pytest from PIL import Image from tests.helpers.runtime import DiffusionResponse, OmniServer, OnlineOmniClient, 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 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, ) VIDEO_NUM_FRAMES = 9 VIDEO_GEN_SAMPLING_PARAMS = OmniDiffusionSamplingParams( num_inference_steps=4, height=HEIGHT, width=WIDTH, num_frames=VIDEO_NUM_FRAMES, seed=42, ) # Online extra_body for diffusion requests IMAGE_GEN_EXTRA_BODY = { "height": HEIGHT, "width": WIDTH, "num_inference_steps": 4, "seed": 42, } # Online form_data for video generation requests (multipart /v1/videos API) VIDEO_GEN_FORM_DATA: dict[str, Any] = { "height": HEIGHT, "width": WIDTH, "num_inference_steps": 4, "num_frames": VIDEO_NUM_FRAMES, "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 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 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 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.""" _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 def _run_offline_t2i(omni: Omni, params: OmniDiffusionSamplingParams = IMAGE_GEN_SAMPLING_PARAMS): return omni.generate({"prompt": PROMPT}, params) def _run_offline_t2v(omni: Omni, params: OmniDiffusionSamplingParams = VIDEO_GEN_SAMPLING_PARAMS): return omni.generate({"prompt": PROMPT}, params) def _run_offline_i2i(omni: Omni, params: OmniDiffusionSamplingParams = IMAGE_GEN_SAMPLING_PARAMS): return omni.generate( {"prompt": PROMPT, "multi_modal_data": {"image": INPUT_IMAGE}}, 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, ) ### 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: OnlineOmniClient, 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: OnlineOmniClient, 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( content_text=PROMPT, image_data_url=image_data_url, ) request_config = { "model": server.model, "messages": messages, "extra_body": extra_body or 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_video_request(omni: Omni): """Run and validate a text to video request.""" _validate_video(_run_offline_t2v(omni)) def run_and_validate_image_to_video_request(omni: Omni, check_t2v_divergence: bool = True): """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. Pipelines that mandate an image and fail closed without one (SANA-Video I2V) pass check_t2v_divergence=False: a dropped image raises instead of silently degrading, so the T2V comparison is moot. """ i2v_video = _validate_video_frames(_get_offline_videos(_run_offline_i2v(omni)))[0] if not check_t2v_divergence: return 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_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( server: OmniServer, client: OnlineOmniClient, form_data: dict | None = None ) -> list[DiffusionResponse]: """Run a text to video request through the server's /v1/videos API.""" data = dict(form_data or VIDEO_GEN_FORM_DATA) data.setdefault("prompt", PROMPT) data.setdefault("model", server.model) return client.send_video_diffusion_request({"form_data": data}) def _run_online_i2v(server: OmniServer, client: OnlineOmniClient) -> 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 videos = responses[0].videos assert videos is not None assert len(videos) > 0 return videos ### Online task runners def run_and_validate_online_text_to_image_request(server: OmniServer, client: OnlineOmniClient): """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: OnlineOmniClient): """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_determinism(server: OmniServer, client: OnlineOmniClient, 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_multi_output(server: OmniServer, client: OnlineOmniClient, 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: OnlineOmniClient): """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: OnlineOmniClient): """Run and validate an image to video request through the server.""" _get_online_videos(_run_online_i2v(server, client))