vllm-omni/tests/diffusion/models/qwen_image_21/test_batch_compatibility.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

import json
from types import SimpleNamespace

import pytest
from PIL import Image

from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import (
    get_qwen_image_21_pre_process_func,
)
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler, StepScheduler
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]


@pytest.fixture
def preprocess(tmp_path):
    vae_dir = tmp_path / "vae"
    vae_dir.mkdir()
    (vae_dir / "config.json").write_text(json.dumps({"temperal_downsample": [False] * 4}))
    return get_qwen_image_21_pre_process_func(SimpleNamespace(model=str(tmp_path)))


def make_request(request_id, sizes, color="red"):
    return OmniDiffusionRequest(
        request_id=request_id,
        prompt={
            "prompt": "A ceramic teapot",
            "multi_modal_data": {"image": [Image.new("RGB", size, color) for size in sizes]},
        },
        sampling_params=OmniDiffusionSamplingParams(
            height=1024, width=1024, num_inference_steps=2, true_cfg_scale=1.0, seed=42
        ),
    )


@pytest.mark.parametrize("scheduler_class", [RequestScheduler, StepScheduler])
@pytest.mark.parametrize(
    "first_sizes,second_sizes,compatible",
    [
        pytest.param([], [(64, 64)], False, id="text-and-edit"),
        pytest.param([(64, 64)], [], False, id="edit-and-text"),
        pytest.param([(64, 64)], [(64, 64), (64, 64)], False, id="image-count"),
        pytest.param([(64, 64)], [(128, 64)], False, id="image-aspect-ratio"),
        pytest.param([(64, 64), (128, 64)], [(128, 64), (64, 64)], False, id="slot-order"),
        pytest.param([], [], True, id="two-text-requests"),
        pytest.param([(64, 64)], [(64, 64)], True, id="different-image-content"),
        pytest.param([(64, 64)], [(128, 128)], True, id="same-resized-layout"),
        pytest.param([(64, 64), (128, 64)], [(128, 128), (256, 128)], True, id="same-multi-image-layout"),
    ],
)
def test_scheduler_groups_by_condition_layout(preprocess, scheduler_class, first_sizes, second_sizes, compatible):
    first = preprocess(make_request("first", first_sizes))
    second = preprocess(make_request("second", second_sizes, color="blue"))
    scheduler = scheduler_class()
    scheduler.initialize(SimpleNamespace(max_num_seqs=2))
    scheduler.add_request(first)
    scheduler.add_request(second)

    scheduled = scheduler.schedule()
    expected_ids = ["first", "second"] if compatible else ["first"]
    assert scheduled.scheduled_request_ids == expected_ids

    scheduler.finish_requests(expected_ids, DiffusionRequestStatus.FINISHED_COMPLETED)
    if not compatible:
        assert scheduler.schedule().scheduled_request_ids == ["second"]
        scheduler.finish_requests("second", DiffusionRequestStatus.FINISHED_COMPLETED)
    assert not scheduler.has_requests()


@pytest.mark.parametrize(
    "prompt", ["A teapot", {"prompt": "A teapot"}, {"prompt": "A teapot", "multi_modal_data": {"image": None}}]
)
def test_text_input_forms_share_compatibility_key(preprocess, prompt):
    request = make_request("text", [])
    request.prompt = prompt
    processed = preprocess(request)
    empty_list_request = preprocess(make_request("empty-list", []))
    assert processed.batch_compatibility_key is not None
    assert processed.batch_compatibility_key == empty_list_request.batch_compatibility_key


@pytest.mark.parametrize("image_type", ["numpy", "tensor"])
def test_preprocess_accepts_array_images_without_truth_testing(preprocess, image_type):
    import numpy as np
    import torch

    pixels = np.full((32, 32, 4), 255, dtype=np.uint8)
    image = torch.from_numpy(pixels).permute(2, 0, 1) if image_type == "tensor" else pixels
    request = make_request("array", [])
    request.prompt["multi_modal_data"]["image"] = image
    processed = preprocess(request)
    assert processed.batch_compatibility_key is not None
    assert processed.allow_mixed_step_phases is False


def test_late_prefill_waits_until_all_decode_requests_finish(preprocess):
    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput

    scheduler = StepScheduler()
    scheduler.initialize(SimpleNamespace(max_num_seqs=3))
    for request_id in ("a", "b"):
        scheduler.add_request(preprocess(make_request(request_id, [])))
    wave = scheduler.schedule()
    assert wave.scheduled_request_ids == ["a", "b"]
    scheduler.update_from_output(
        wave,
        BatchRunnerOutput.from_list(
            [RunnerOutput(request_id=request_id, step_index=1, finished=False) for request_id in ("a", "b")]
        ),
    )
    scheduler.add_request(preprocess(make_request("late", [])))
    assert scheduler.schedule().scheduled_request_ids == ["a", "b"]
    scheduler.finish_requests("a", DiffusionRequestStatus.FINISHED_COMPLETED)
    assert scheduler.schedule().scheduled_request_ids == ["b"]
    scheduler.finish_requests("b", DiffusionRequestStatus.FINISHED_COMPLETED)
    assert scheduler.schedule().scheduled_request_ids == ["late"]


def test_preprocess_normalizes_tuple_images_like_a_list(preprocess):
    image = Image.new("RGB", (64, 64), "red")
    tuple_request = make_request("tuple", [])
    tuple_request.prompt["multi_modal_data"]["image"] = (image,)

    processed = preprocess(tuple_request)
    reference = preprocess(make_request("list", [(64, 64)]))

    assert processed.batch_compatibility_key == reference.batch_compatibility_key
    assert processed.batch_compatibility_key[0] == "qwen_image_21"
    assert len(processed.batch_compatibility_key[1]) == 1
    info = processed.prompt["additional_information"]
    assert len(info["prompt_image"]) == 1
    assert len(info["vae_images"]) == 1
    assert len(info["input_image_sizes"]) == 1


def test_preprocess_rejects_more_images_than_the_model_supports(preprocess):
    from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import (
        MAX_QWEN_IMAGE_21_INPUT_IMAGES,
    )

    request = make_request("too-many", [(64, 64)] * (MAX_QWEN_IMAGE_21_INPUT_IMAGES + 1))
    with pytest.raises(ValueError, match="At most"):
        preprocess(request)


def test_preprocess_defaults_true_cfg_scale(preprocess):
    request = make_request("cfg", [])
    request.sampling_params.true_cfg_scale = None
    processed = preprocess(request)
    assert processed.sampling_params.true_cfg_scale == 1.0

    request = make_request("cfg-set", [])
    request.sampling_params.true_cfg_scale = 4.0
    assert preprocess(request).sampling_params.true_cfg_scale == 4.0


@pytest.mark.parametrize("size,expected_wider", [((128, 64), True), ((64, 128), False)])
def test_preprocess_derives_output_size_from_condition_aspect(preprocess, size, expected_wider):
    request = make_request("aspect", [size])
    request.sampling_params.height = None
    request.sampling_params.width = None

    processed = preprocess(request)

    height = processed.sampling_params.height
    width = processed.sampling_params.width
    assert height % 32 == 0 and width % 32 == 0
    assert (width > height) is expected_wider
    info = processed.prompt["additional_information"]
    assert info["calculated_height"] == height
    assert info["calculated_width"] == width


def test_preprocess_stages_rgba_vae_images_and_resized_prompt_images(preprocess):
    processed = preprocess(make_request("staged", [(96, 64)]))
    info = processed.prompt["additional_information"]

    (vae_image,) = info["vae_images"]
    # RGB input gains an opaque alpha channel for the RGBA VAE; the frame dim is
    # inserted for the 3D (video-capable) VAE.
    assert vae_image.shape[1] == 4
    assert vae_image.shape[2] == 1
    (prompt_image,) = info["prompt_image"]
    assert isinstance(prompt_image, Image.Image)
    (width, height) = info["input_image_sizes"][0]
    assert prompt_image.size == (width, height)
    assert processed.batch_compatibility_key == ("qwen_image_21", ((width, height),))