Isotr0py
3d6a1c4d9a
- 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
293 lines
10 KiB
Python
293 lines
10 KiB
Python
# 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"] == {}
|