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

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),))

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

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"] == {}

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