vllm-omni/tests/model_tests/diffusion/task_runners.py

# 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):
    """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_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))