From 3d6a1c4d9a0a09d8392bda3a54f8dfb890d5fb05 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Mon, 21 Sep 2026 08:36:42 +0000 Subject: [PATCH] Address review: examples policy, SPDX headers, recipe docs, test coverage - Register QwenImage21Pipeline in model_extras _EXTRA_SPECS with an image-to-image prompt builder that forwards height/width - Rename run_qwen_image_21_edit.sh to run_qwen_image_21.sh - Fix SPDX-FileCopyrightText headers - Recipe: document mixed-KV-phase admission, add head-vs-diffusers A/B comparison section (measurements pending) - Add test_registry.py and test_pipeline.py (preprocess, prepare_latents ordering, mixed-phase _assemble_kv_cache, post-process); extend test_batch_compatibility.py --- ..._image_21_edit.sh => run_qwen_image_21.sh} | 0 recipes/Qwen/Qwen-Image-2.1.md | 105 +++++-- .../qwen_image_21/test_batch_compatibility.py | 71 +++++ .../models/qwen_image_21/test_decode_graph.py | 37 ++- .../models/qwen_image_21/test_pipeline.py | 292 ++++++++++++++++++ .../models/qwen_image_21/test_registry.py | 64 ++++ tests/model_extras/test_model_extras.py | 21 ++ .../autoencoder_kl_qwenimage21.py | 2 +- .../models/qwen_image_21/__init__.py | 2 +- vllm_omni/model_extras/qwen_image_21.py | 28 ++ vllm_omni/model_extras/registry.py | 6 + 11 files changed, 597 insertions(+), 31 deletions(-) rename examples/offline_inference/image_to_image/{run_qwen_image_21_edit.sh => run_qwen_image_21.sh} (100%) create mode 100644 tests/diffusion/models/qwen_image_21/test_pipeline.py create mode 100644 tests/diffusion/models/qwen_image_21/test_registry.py create mode 100644 vllm_omni/model_extras/qwen_image_21.py diff --git a/examples/offline_inference/image_to_image/run_qwen_image_21_edit.sh b/examples/offline_inference/image_to_image/run_qwen_image_21.sh similarity index 100% rename from examples/offline_inference/image_to_image/run_qwen_image_21_edit.sh rename to examples/offline_inference/image_to_image/run_qwen_image_21.sh diff --git a/recipes/Qwen/Qwen-Image-2.1.md b/recipes/Qwen/Qwen-Image-2.1.md index 9e603c6c5..3d9295e60 100644 --- a/recipes/Qwen/Qwen-Image-2.1.md +++ b/recipes/Qwen/Qwen-Image-2.1.md @@ -65,7 +65,7 @@ python examples/offline_inference/image_to_image/image_edit.py \ or via the bundled launcher: ```bash -bash examples/offline_inference/image_to_image/run_qwen_image_21_edit.sh +bash examples/offline_inference/image_to_image/run_qwen_image_21.sh ``` Up to 4 condition images can be passed via `--image`: @@ -106,37 +106,45 @@ vllm serve Qwen/Qwen-Image-2.1 --omni \ --max-num-seqs 8 ``` -Step execution batches compatible requests at the same KV-cache phase. Once -running requests have completed their first denoising step, newly arriving -requests wait for that batch to finish before starting their own prefill. -`--max-num-seqs` controls batch capacity; mixed prefill/decode admission is not -supported yet. +Step execution batches compatible requests at the same KV-cache phase. The +pipeline tags every request with a `batch_compatibility_key` (conditioning +mode plus condition-image layout), so a newly arriving request that would +otherwise join a batch already past its first denoising step is isolated by +the scheduler and deferred until the in-flight batch drains — it no longer +fails the whole batch. `--max-num-seqs` controls batch capacity. ### CUDA Graph decode Qwen-Image-2.1 automatically uses CUDA Graph for supported fixed-shape denoising -decode steps when `enforce_eager=False` (the default). Prefill remains eager. -This works with both request execution and `--step-execution`. +decode steps when `enable_cuda_graph_decode=True` (the default) and +`enforce_eager=False`. Prefill remains eager. This works with both request +execution and `--step-execution`. -To disable graph capture: +To disable graph capture while keeping the rest of the compilation setup: ```bash -vllm serve Qwen/Qwen-Image-2.1 --omni --enforce-eager --vae-use-tiling +vllm serve Qwen/Qwen-Image-2.1 --omni --no-enable-cuda-graph-decode ``` -The same control is available as `Omni(..., enforce_eager=True)` and as -`enforce_eager: true` on a stage in the deployment YAML. It disables both graph -capture and automatic `torch.compile`. Otherwise, this transformer uses its -model-specific decode graphs instead of automatic `torch.compile`. +The same control is available as `Omni(..., enable_cuda_graph_decode=False)` +and as `enable_cuda_graph_decode: false` on a stage in the deployment YAML. +`--enforce-eager` / `enforce_eager: true` remains the broader switch: it +disables both graph capture and automatic `torch.compile`. Otherwise, this +transformer uses its model-specific decode graphs instead of automatic +`torch.compile`. -The autoregressive engine's `compilation_config.cudagraph_mode` does not control -this diffusion path; use `enforce_eager` to force eager execution. +The autoregressive engine's `compilation_config.cudagraph_mode` does not +control this diffusion path; use `enable_cuda_graph_decode` (or +`enforce_eager` for everything) to force eager decode execution. TP/SP/ring parallelism, HSDP, offload/cache hooks, compiled blocks, quantized -KV caches, dynamic LoRA, and padded text masks fall back to eager decode. -Graphs keep separate entries for different image layouts and copy request-owned -prefix K/V into static buffers before replay. Those buffers require additional -memory; graph entries are bounded by the transformer's cache limit. +KV caches, dynamic LoRA, padded text masks, and a second in-flight request +whose cache aliases an already-owned graph key fall back to eager decode. +Graphs keep separate entries for different image layouts and copy +request-owned prefix K/V into static buffers before replay; each graph key +has a single live owner at a time, so concurrent same-shape requests never +overwrite each other's prefix. Those buffers require additional memory; graph +entries are bounded by the transformer's cache limit. ### Verification @@ -311,6 +319,63 @@ piecewise-span masking and Ulysses/CFG constraints, for no memory benefit over the per-request dense tensors. The FP8 storage above captures the actual memory win at ~1% of the integration cost. +## Quality Comparison (head vs diffusers) + +The reference is the diffusers Qwen-Image 2.1 pipeline +([huggingface/diffusers#14804](https://github.com/huggingface/diffusers/pull/14804)). +Both sides must pin the same checkpoint, inputs, seed, precision (BF16), and +hardware. + +vLLM-Omni (this recipe): + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model Qwen/Qwen-Image-2.1 \ + --prompt "A ceramic teapot on a wooden table" \ + --seed 42 \ + --output qwen_image_21_vllm_omni.png \ + --num-inference-steps 50 \ + --cfg-scale 1.0 +``` + +diffusers reference: + +```python +import torch +from diffusers import QwenImage21Pipeline # diffusers PR #14804 + +pipe = QwenImage21Pipeline.from_pretrained( + "Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16 +).to("cuda") +image = pipe( + prompt="A ceramic teapot on a wooden table", + num_inference_steps=50, + true_cfg_scale=1.0, + generator=torch.Generator(device="cuda").manual_seed(42), + height=1024, + width=1024, +).images[0] +image.save("qwen_image_21_diffusers.png") +``` + +| Metric | vLLM-Omni | diffusers (#14804) | Status | +| --- | --- | --- | --- | +| Output image (1024×1024, 50 steps, seed 42, BF16) | pending measurement | pending measurement | not yet run | +| PSNR / LPIPS vs reference | pending measurement | — | not yet run | +| Latency (time/image, same GPU) | pending measurement | pending measurement | not yet run | + +The prefix KV cache changes the compute path (cached prefix vs. full recompute +at every step), so exact numerical parity is not expected; match the cache +setting on both sides or compare with PSNR/LPIPS rather than exact pixels. + +The following capabilities are implemented but the PR-body checkboxes predate +any attached measurements; treat them as unverified until the A/B above is +filled in: + +- [ ] TP=2/4 output parity vs single-card — pending measurement. +- [ ] VAE tiling vs non-tiled decode — pending measurement. +- [ ] CUDA graph decode vs eager — pending measurement. + ## Known Limitations - Cache acceleration backends (`cache_dit`, `tea_cache`) are not supported; diff --git a/tests/diffusion/models/qwen_image_21/test_batch_compatibility.py b/tests/diffusion/models/qwen_image_21/test_batch_compatibility.py index 4b4136f29..c66b27e14 100644 --- a/tests/diffusion/models/qwen_image_21/test_batch_compatibility.py +++ b/tests/diffusion/models/qwen_image_21/test_batch_compatibility.py @@ -119,3 +119,74 @@ def test_late_prefill_waits_until_all_decode_requests_finish(preprocess): 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),)) diff --git a/tests/diffusion/models/qwen_image_21/test_decode_graph.py b/tests/diffusion/models/qwen_image_21/test_decode_graph.py index ff4ea1a0a..e0ccfb17c 100644 --- a/tests/diffusion/models/qwen_image_21/test_decode_graph.py +++ b/tests/diffusion/models/qwen_image_21/test_decode_graph.py @@ -197,20 +197,32 @@ def test_capture_failure_releases_buffers_and_keeps_request_cache(graph_case, mo def test_replay_uses_active_request_and_layout(graph_case): manager, device, _ = graph_case first, second = make_cache(device, 1.0), make_cache(device, 5.0) + first_key = weakref.ref(first[0]["cond"]["key"]) with torch.inference_mode(): register(manager, first) + # A second live request aliasing the same key is refused: it keeps its + # own cache and decodes eagerly while the owner is in flight. register(manager, second) - for cache in (first, second, first): - output = decode(manager, cache) - expected = ( - torch.ones_like(output).sin() + cache[0]["cond"]["key"].mean() + cache[0]["cond"]["value"].mean() + 2 - ) - torch.testing.assert_close(output, expected) - original = decode(manager, first).clone() + assert decode(manager, second) is None + output = decode(manager, first) + expected = ( + torch.ones_like(output).sin() + first[0]["cond"]["key"].mean() + first[0]["cond"]["value"].mean() + 2 + ) + torch.testing.assert_close(output, expected) + # Once the owner's cache is released, the key is free to adopt. + del first + assert first_key() is None + register(manager, second) + output = decode(manager, second) + expected = ( + torch.ones_like(output).sin() + second[0]["cond"]["key"].mean() + second[0]["cond"]["value"].mean() + 2 + ) + torch.testing.assert_close(output, expected) + original = decode(manager, second).clone() register(manager, second, (1, 1, 4)) other = decode(manager, second, (1, 1, 4)) torch.testing.assert_close(other, torch.ones_like(other).sin() + 5 + 6 + 1) - torch.testing.assert_close(decode(manager, first), original) + torch.testing.assert_close(decode(manager, second), original) assert all(entry.captures == 1 for entry in manager.entries.values()) @@ -294,7 +306,14 @@ def test_batch_changes_and_padding_do_not_reuse_stale_prefix(graph_case): assert result is not None and result.shape == (2, 4, 4) expected = torch.ones_like(result).sin() + torch.tensor([5.0, 11.0], device=device)[:, None, None] torch.testing.assert_close(result, expected) - kwargs["kv_cache"] = [{"cond": {part: tensor.flip(0) for part, tensor in merged[0]["cond"].items()}}] + # A rebuilt cache (e.g. reordered rows) is a different live cache + # aliasing the same key: decode refuses the shared entry until the + # previous owner's tensors are released. + flipped = [{"cond": {part: tensor.flip(0) for part, tensor in merged[0]["cond"].items()}}] + kwargs["kv_cache"] = flipped + assert manager.try_decode(**kwargs) is None + del merged + register(manager, flipped) torch.testing.assert_close(manager.try_decode(**kwargs), expected.flip(0)) kwargs["encoder_hidden_states_mask"][1, 0] = False assert manager.try_decode(**kwargs) is None diff --git a/tests/diffusion/models/qwen_image_21/test_pipeline.py b/tests/diffusion/models/qwen_image_21/test_pipeline.py new file mode 100644 index 000000000..363b1ad13 --- /dev/null +++ b/tests/diffusion/models/qwen_image_21/test_pipeline.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CPU unit tests for QwenImage21Pipeline helpers: prepare_latents duplication +order, KV-cache assembly across step phases, and the post-process func.""" + +import json +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.models.qwen_image_21.pipeline_qwen_image_21 import ( + QwenImage21Pipeline, + get_qwen_image_21_post_process_func, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _toy_pipeline() -> QwenImage21Pipeline: + """A pipeline shell with the attributes prepare_latents reads, no weights.""" + pipeline = QwenImage21Pipeline.__new__(QwenImage21Pipeline) + nn.Module.__init__(pipeline) + pipeline.vae_scale_factor = 16 + pipeline.latent_channels = 4 + return pipeline + + +# --------------------------------------------------------------------------- +# prepare_latents +# --------------------------------------------------------------------------- + + +def test_prepare_latents_duplicates_condition_latents_request_major(): + """batch=2 prompts x n=2 outputs must order condition latents A,A,B,B. + + ``encode_prompt`` expands prompt embeddings request-major (A,A,B,B); the + condition latents must follow the same order or generations 2 and 3 swap + condition images. ``repeat`` would produce the wrong (A,B,A,B) tile-major + order; ``repeat_interleave`` is correct. + """ + pipeline = _toy_pipeline() + cond_a = torch.full((1, 4, 1, 8, 8), 1.0) + cond_b = torch.full((1, 4, 1, 8, 8), 2.0) + stacked = torch.cat([cond_a, cond_b], dim=0) # one row per request + + latents, image_latents = pipeline.prepare_latents( + [stacked], + batch_size=4, + num_channels_latents=4, + height=64, + width=64, + dtype=torch.float32, + device=torch.device("cpu"), + generator=torch.Generator().manual_seed(0), + ) + + assert latents.shape == (4, 16, 4) # packed: 4x4 latent tokens, 4 channels + assert image_latents.shape == (4, 64, 4) # packed: 8x8 condition tokens + row_means = image_latents.mean(dim=(1, 2)) + torch.testing.assert_close(row_means, torch.tensor([1.0, 1.0, 2.0, 2.0])) + + +def test_prepare_latents_without_duplication_when_rows_match_batch(): + pipeline = _toy_pipeline() + stacked = torch.arange(4, dtype=torch.float32).view(4, 1, 1, 1, 1).expand(4, 4, 1, 2, 2).contiguous() + + _, image_latents = pipeline.prepare_latents( + [stacked], + batch_size=4, + num_channels_latents=4, + height=32, + width=32, + dtype=torch.float32, + device=torch.device("cpu"), + generator=torch.Generator().manual_seed(0), + ) + + torch.testing.assert_close(image_latents.mean(dim=(1, 2)), torch.arange(4, dtype=torch.float32)) + + +def test_prepare_latents_rejects_non_divisible_duplication(): + pipeline = _toy_pipeline() + stacked = torch.zeros(2, 4, 1, 8, 8) + + with pytest.raises(ValueError, match="Cannot duplicate"): + pipeline.prepare_latents( + [stacked], + batch_size=3, + num_channels_latents=4, + height=64, + width=64, + dtype=torch.float32, + device=torch.device("cpu"), + generator=torch.Generator().manual_seed(0), + ) + + +def test_prepare_latents_rejects_generator_list_length_mismatch(): + pipeline = _toy_pipeline() + + with pytest.raises(ValueError, match="list of generators"): + pipeline.prepare_latents( + None, + batch_size=2, + num_channels_latents=4, + height=32, + width=32, + dtype=torch.float32, + device=torch.device("cpu"), + generator=[torch.Generator().manual_seed(0)], + ) + + +def test_prepare_latents_concatenates_multiple_condition_images(): + pipeline = _toy_pipeline() + first = torch.full((2, 4, 1, 4, 4), 1.0) + second = torch.full((2, 4, 1, 2, 2), 3.0) + + _, image_latents = pipeline.prepare_latents( + [first, second], + batch_size=2, + num_channels_latents=4, + height=32, + width=32, + dtype=torch.float32, + device=torch.device("cpu"), + generator=torch.Generator().manual_seed(0), + ) + + # Token dim concatenates the packed condition images: 16 + 4 per row. + assert image_latents.shape == (2, 20, 4) + torch.testing.assert_close(image_latents[:, :16].mean(), torch.tensor(1.0)) + torch.testing.assert_close(image_latents[:, 16:].mean(), torch.tensor(3.0)) + + +# --------------------------------------------------------------------------- +# KV-cache assembly across step phases +# --------------------------------------------------------------------------- + + +def _state(cache) -> SimpleNamespace: + return SimpleNamespace(extra={"kv_cache": cache}) + + +def _decode_cache(rows: int, value: float) -> list[dict]: + return [ + { + "cond": { + "key": torch.full((rows, 2), value), + "value": torch.full((rows, 2), value + 0.5), + } + } + ] + + +def test_kv_cache_phase_classification(): + assert QwenImage21Pipeline._kv_cache_phase(None) == "none" + assert QwenImage21Pipeline._kv_cache_phase([{}]) == "prefill" + assert QwenImage21Pipeline._kv_cache_phase(_decode_cache(1, 0.0)) == "decode" + + +def test_assemble_kv_cache_all_none(): + assert QwenImage21Pipeline._assemble_kv_cache([_state(None), _state(None)]) == (None, False) + + +def test_assemble_kv_cache_prefill_returns_fresh_owned_cache(): + states = [_state([{} for _ in range(3)]), _state([{} for _ in range(3)])] + + kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache(states) + + assert take_ownership is True + assert kv_cache == [{}, {}, {}] + # Fresh blocks, not aliases of either request's cache: the batched prefill + # writes into these and _scatter_kv_cache splits them back per request. + for block in kv_cache: + assert all(block is not request_block for state in states for request_block in state.extra["kv_cache"]) + + +def test_assemble_kv_cache_decode_concatenates_in_request_order(): + first, second = _decode_cache(2, 1.0), _decode_cache(1, 3.0) + states = [_state(first), _state(second)] + + kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache(states) + + assert take_ownership is False + assert len(kv_cache) == 1 + cond = kv_cache[0]["cond"] + torch.testing.assert_close(cond["key"], torch.tensor([[1.0, 1.0], [1.0, 1.0], [3.0, 3.0]])) + torch.testing.assert_close(cond["value"], torch.tensor([[1.5, 1.5], [1.5, 1.5], [3.5, 3.5]])) + # Decode only reads the cache: per-request tensors must stay untouched. + assert first[0]["cond"]["key"].shape == (2, 2) + assert second[0]["cond"]["key"].shape == (1, 2) + + +def test_assemble_kv_cache_decode_merges_every_part(): + first = [{"cond": {"key": torch.ones(1, 1), "value": torch.ones(1, 1), "key_scale": torch.ones(1)}}] + second = [{"cond": {"key": torch.zeros(1, 1), "value": torch.zeros(1, 1), "key_scale": torch.zeros(1)}}] + + kv_cache, take_ownership = QwenImage21Pipeline._assemble_kv_cache([_state(first), _state(second)]) + + assert take_ownership is False + assert set(kv_cache[0]["cond"]) == {"key", "value", "key_scale"} + torch.testing.assert_close(kv_cache[0]["cond"]["key_scale"], torch.tensor([1.0, 0.0])) + + +def test_assemble_kv_cache_mixed_phases_raises(): + """Defensive backstop: a late prefill must never join a decoding batch. + + The preprocessor sets ``allow_mixed_step_phases = False`` so the + StepScheduler defers late requests until the batch finishes decoding; this + raise is the last line of defense if a mixed batch is ever assembled. + """ + states = [_state([{}]), _state(_decode_cache(1, 0.0))] + + with pytest.raises(ValueError, match="mixed KV-cache phases"): + QwenImage21Pipeline._assemble_kv_cache(states) + + with pytest.raises(ValueError, match="mixed KV-cache phases"): + QwenImage21Pipeline._assemble_kv_cache([_state(None), _state(_decode_cache(1, 0.0))]) + + +def test_scatter_kv_cache_splits_rows_back_per_request(): + states = [_state([{}]), _state([{}])] + merged = [ + { + "cond": { + "key": torch.tensor([[1.0], [2.0], [3.0]]), + "value": torch.tensor([[4.0], [5.0], [6.0]]), + } + } + ] + + QwenImage21Pipeline._scatter_kv_cache(states, [1, 2], merged) + + torch.testing.assert_close(states[0].extra["kv_cache"][0]["cond"]["key"], torch.tensor([[1.0]])) + torch.testing.assert_close(states[0].extra["kv_cache"][0]["cond"]["value"], torch.tensor([[4.0]])) + torch.testing.assert_close(states[1].extra["kv_cache"][0]["cond"]["key"], torch.tensor([[2.0], [3.0]])) + assert states[1].extra["kv_cache"][0]["cond"]["key"].is_contiguous() + + +# --------------------------------------------------------------------------- +# post-process func +# --------------------------------------------------------------------------- + + +@pytest.fixture +def postprocess(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_post_process_func(SimpleNamespace(model=str(tmp_path))) + + +def test_postprocess_tensor_returns_pil_images(postprocess): + from PIL import Image + + images = torch.rand(2, 4, 16, 16) * 2 - 1 + + result = postprocess(images) + + assert len(result) == 2 + assert all(isinstance(image, Image.Image) for image in result) + + +def test_postprocess_envelope_unwraps_payload_and_keeps_metadata(postprocess): + from PIL import Image + + envelope = { + "payload": {"image": torch.rand(1, 4, 16, 16), "seed": 42}, + "metadata": {"request_id": "abc"}, + } + + result = postprocess(envelope) + + assert result["metadata"] == {"request_id": "abc"} + assert result["payload"]["seed"] == 42 + assert len(result["payload"]["image"]) == 1 + assert isinstance(result["payload"]["image"][0], Image.Image) + + +def test_postprocess_envelope_requires_image_payload(postprocess): + with pytest.raises(ValueError, match=r"payload\['image'\]"): + postprocess({"payload": {"seed": 42}, "metadata": {}}) + + +def test_postprocess_envelope_normalizes_non_dict_metadata(postprocess): + result = postprocess({"payload": {"image": torch.rand(1, 4, 16, 16)}, "metadata": "bogus"}) + + assert result["metadata"] == {} diff --git a/tests/diffusion/models/qwen_image_21/test_registry.py b/tests/diffusion/models/qwen_image_21/test_registry.py new file mode 100644 index 000000000..2a09ea79d --- /dev/null +++ b/tests/diffusion/models/qwen_image_21/test_registry.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""L1 tests for checkpoint-to-pipeline routing of Qwen-Image 2.1.""" + +import json + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_registry_entry_loads_the_pipeline_class(): + from vllm_omni.diffusion.models.qwen_image_21 import QwenImage21Pipeline + from vllm_omni.diffusion.registry import DiffusionModelRegistry + + assert DiffusionModelRegistry._try_load_model_cls("QwenImage21Pipeline") is QwenImage21Pipeline + + +def test_pipeline_declares_request_batching_and_step_execution(): + from vllm_omni.diffusion.models.qwen_image_21 import QwenImage21Pipeline + + assert QwenImage21Pipeline.supports_request_batch + assert QwenImage21Pipeline.supports_step_execution + + +def test_process_funcs_resolve_to_the_pipeline_module(): + from vllm_omni.diffusion.models.qwen_image_21 import pipeline_qwen_image_21 + from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS, _DIFFUSION_PRE_PROCESS_FUNCS + + pre_name = _DIFFUSION_PRE_PROCESS_FUNCS["QwenImage21Pipeline"] + post_name = _DIFFUSION_POST_PROCESS_FUNCS["QwenImage21Pipeline"] + assert pre_name == "get_qwen_image_21_pre_process_func" + assert post_name == "get_qwen_image_21_post_process_func" + assert callable(getattr(pipeline_qwen_image_21, pre_name)) + assert callable(getattr(pipeline_qwen_image_21, post_name)) + + +def test_multimodal_metadata_caps_condition_images_at_four(): + from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_21_MAX_INPUT_IMAGES, get_diffusion_model_metadata + + metadata = get_diffusion_model_metadata("QwenImage21Pipeline") + + assert metadata.supports_multimodal_inputs + assert metadata.max_multimodal_image_inputs == 4 + assert metadata.max_multimodal_image_inputs == QWEN_IMAGE_21_MAX_INPUT_IMAGES + + +def test_model_index_resolves_and_enriches(tmp_path): + from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name + + (tmp_path / "model_index.json").write_text( + json.dumps({"_class_name": "QwenImage21Pipeline"}), + encoding="utf-8", + ) + + assert resolve_model_class_name(str(tmp_path)) == "QwenImage21Pipeline" + + config = OmniDiffusionConfig(model=str(tmp_path)) + config.enrich_config() + + assert config.model_class_name == "QwenImage21Pipeline" + assert config.supports_multimodal_inputs + assert config.max_multimodal_image_inputs == 4 diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index a187c570c..fea96cda4 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -495,6 +495,27 @@ def test_bagel_image_to_image_prompt_builder() -> None: assert result["negative_prompt"] == "ugly" +@pytest.mark.core_model +@pytest.mark.cpu +def test_qwen_image_21_image_to_image_prompt_builder() -> None: + dummy_image = Image.new("RGBA", (64, 64)) + result = build_image_to_image_prompt( + "QwenImage21Pipeline", + prompt="make it dance", + negative_prompt="blurry", + input_image=dummy_image, + height=1024, + width=768, + ) + assert result == { + "prompt": "make it dance", + "multi_modal_data": {"image": dummy_image}, + "height": 1024, + "width": 768, + "negative_prompt": "blurry", + } + + @pytest.mark.core_model @pytest.mark.cpu def test_unknown_pipeline_preserves_canonical_text_to_image_prompt() -> None: diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py index 771c41ec1..7ef61b3b6 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage21.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from typing import Any diff --git a/vllm_omni/diffusion/models/qwen_image_21/__init__.py b/vllm_omni/diffusion/models/qwen_image_21/__init__.py index 6aa269033..8d1dcc168 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/__init__.py +++ b/vllm_omni/diffusion/models/qwen_image_21/__init__.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Qwen Image 2.1 diffusion model components. Lazy by necessity: the pipeline imports the distributed VAE diff --git a/vllm_omni/model_extras/qwen_image_21.py b/vllm_omni/model_extras/qwen_image_21.py new file mode 100644 index 000000000..b45fa87bc --- /dev/null +++ b/vllm_omni/model_extras/qwen_image_21.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from typing import Any + +from PIL import Image + + +def build_image_to_image_prompt( + prompt: str, + negative_prompt: str | None, + input_image: Image.Image | list[Image.Image], + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + img_prompt: dict[str, Any] = { + "prompt": prompt, + "multi_modal_data": {"image": input_image}, + } + if height is not None: + img_prompt["height"] = height + if width is not None: + img_prompt["width"] = width + if negative_prompt is not None: + img_prompt["negative_prompt"] = negative_prompt + return img_prompt diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 9c4a6eed6..fb05754b6 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -75,6 +75,9 @@ from vllm_omni.model_extras.ming_flash_omni import ( from vllm_omni.model_extras.ming_flash_omni import ( build_text_to_image_prompt as build_ming_flash_omni_text_to_image_prompt, ) +from vllm_omni.model_extras.qwen_image_21 import ( + build_image_to_image_prompt as build_qwen_image_21_image_to_image_prompt, +) from vllm_omni.model_extras.sana_video import SANA_VIDEO_EXTRA_BODY_PARAMS from vllm_omni.model_extras.sensenova_u1 import ( SENSENOVA_U1_EXTRA_BODY_PARAMS, @@ -299,6 +302,9 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { "text_to_image_prompt_builder": build_ming_flash_omni_text_to_image_prompt, "image_to_image_prompt_builder": build_ming_flash_omni_image_to_image_prompt, }, + "QwenImage21Pipeline": { + "image_to_image_prompt_builder": build_qwen_image_21_image_to_image_prompt, + }, } for model_class_name in ("LTX2Pipeline", "LTX2TwoStagePipeline"):