diff --git a/benchmarks/diffusion/backends.py b/benchmarks/diffusion/backends.py index 425667d8c..dcc86710e 100644 --- a/benchmarks/diffusion/backends.py +++ b/benchmarks/diffusion/backends.py @@ -107,6 +107,8 @@ async def async_request_image_edits( if extra_body.get("guidance_scale") is not None: form.add_field("guidance_scale", str(extra_body["guidance_scale"])) + if extra_body.get("guidance_scale_2") is not None: + form.add_field("guidance_scale_2", str(extra_body["guidance_scale_2"])) if extra_body.get("negative_prompt") is not None: form.add_field("negative_prompt", str(extra_body["negative_prompt"])) if extra_body.get("true_cfg_scale") is not None: diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py index 6bf64ec1d..ad687a4ad 100644 --- a/benchmarks/diffusion/diffusion_benchmark_serving.py +++ b/benchmarks/diffusion/diffusion_benchmark_serving.py @@ -1107,6 +1107,11 @@ def calculate_metrics( error_outputs = [o for o in outputs if not o.success] num_success = len(success_outputs) + error_log_path = os.environ.get("BENCH_ERROR_LOG") + if error_outputs and error_log_path: + with open(error_log_path, "a") as f: + for output in error_outputs: + f.write(f"ERROR latency={output.latency:.3f} detail={output.error!r}\n") latencies = [o.latency for o in success_outputs] peak_memories = [o.peak_memory_mb for o in success_outputs if o.peak_memory_mb > 0] diff --git a/recipes/Boogu/Boogu-Image.md b/recipes/Boogu/Boogu-Image.md index 6ff848539..95c87d2c7 100644 --- a/recipes/Boogu/Boogu-Image.md +++ b/recipes/Boogu/Boogu-Image.md @@ -286,7 +286,6 @@ curl -X POST http://localhost:8091/v1/images/generations \ "size": "1024x1024", "seed": 42 }' | jq -r '.data[0].b64_json' | base64 -d > output-turbo.png -``` ### Notes @@ -418,6 +417,45 @@ curl -s http://localhost:8091/v1/chat/completions \ supports Boogu-Image-Edit directly (`--model Boogu/Boogu-Image-0.1-Edit --guidance-scale 5.0`, optional `--guidance-scale-2`). +- **Request batching (`--max-num-seqs > 1`) is not bit-exact vs. serial.** + Compatible TI2I requests (same task type, same `guidance_scale_2` + explicitness) can share one forward pass. The instruction encoder batches + with `padding="longest"`, so a request's *padded batch shape* depends on + whichever partner request(s) happen to co-arrive, even though production + RoPE derives each row's own effective position count from that row's own + `attention_mask.sum()` rather than from the batch's padded width. In the + measurements below, `--max-num-seqs 1` and `--max-num-seqs > 1` do not + produce pixel-identical output for the same seed; the exact mechanism + (padding-shape-dependent kernel/reduction-order numerics vs. some other + batch-shape effect) has not been isolated, so treat this as an observed + drift under the specific configs below, not a diagnosed root cause, and + do not extrapolate the numbers below as a general upper bound on other + configs. Two distinct regimes have been measured so far, and they are + **not** interchangeable tolerances: + - **Fixed-shape co-batch** (partner instructions encode to the same + length): `Boogu-Image-0.1-Edit`, 512x512, 20 steps, guidance 5/2, one + fixed reference image and instruction pair, seeds 42-73, B=2 vs. serial + -- mean pixel MAE **1.70/255**, max per-image MAE **2.68/255**, min PSNR + **34.35 dB** (2/32 image pairs were pixel-exact). + - **Heterogeneous partner length** (co-batched instructions encode to + different lengths): MAD **2.6-6.0** in the fixed-length partner-mixing + check on this PR's test head -- larger than the fixed-shape case above + in that check. A targeted current-head A/B run observed a comparable + range: `Boogu-Image-0.1-Edit`, 512x512, 20 steps, guidance 5/2, one + fixed A request (5-token instruction) co-batched in turn with a + 2-token and a 49-token partner instruction (both also differing in + reference image and content, not length alone; `max_num_seqs=2`, + co-batching verified via request tracing, not assumed) -- + solo vs. `A`+2-token partner MAD **4.06/255** (PSNR 29.9 dB), + solo vs. `A`+49-token partner MAD **3.87/255** (PSNR 29.7 dB), + 2-token vs. 49-token partner MAD **4.05/255** (PSNR 30.2 dB). A + controlled ablation (same partner content/reference, varying only + instruction length) would be needed to isolate a pure length effect, + and has not been run. + Pin `--max-num-seqs 1` for any workload that requires strict + serial-reproducibility (e.g. golden-image regression tests); use + `--max-num-seqs > 1` only when this drift is acceptable for the throughput + gain. ## Fast image editing (Boogu-Image-0.1-Edit-Turbo) diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py index 32061f648..55ae4e7d1 100644 --- a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py +++ b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py @@ -699,6 +699,116 @@ def test_transformer_forward_ti2i_shape(): assert torch.isfinite(out).all() +def test_transformer_forward_batch_invariant_to_partner_padding(): + """Row A must not change when only row B's padded length changes. + + Both rows carry a non-empty reference image (the TI2I/editing path), so + this exercises the reference-image refiner and the joint + instruction/reference/noise sequence concatenation and position + offsets -- not just the T2I path, which skips the reference refiner + entirely and would leave that concatenation logic uncovered. Real + attention masking, RoPE, and sequence concatenation are exercised + end-to-end through the actual (tiny) weights -- the mechanism the CPU + pipeline-level fakes in ``test_pipeline_boogu_image.py`` cannot reach, + since those stop at a fake encoder/transformer boundary. A's instruction + content, reference latent, and attention mask are held byte-identical + across runs; only B's padded region grows. + """ + from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import ( + BooguImageDoubleStreamRotaryPosEmbed, + BooguImageTransformer2DModel, + ) + + torch.manual_seed(0) + model = BooguImageTransformer2DModel(od_config=_tiny_od_config()) + _randomize_parameters(model) + # _randomize_parameters' uniform(-0.02, 0.02) scale (tuned for the other + # shape/finite-value tests in this file) makes the instruction stream's + # contribution to the final output smaller than float32 rounding noise + # after 4 layers of untrained mixing with the image latents -- both the + # invariance check and the negative control below were passing/failing + # vacuously at that scale (diffs ~1e-6) until confirmed via a scaled-up + # diagnostic that masking is in fact respected; amplify locally so a real + # masking regression is distinguishable from noise. + with torch.no_grad(): + for param in model.parameters(): + param.mul_(20.0) + model.eval() + + batch_size = 2 + in_channels = 4 + latent_h = latent_w = 8 + instruction_feat_dim = 32 + a_len = 5 + + torch.manual_seed(0) + latents = torch.randn(batch_size, in_channels, latent_h, latent_w) + timestep = torch.full((batch_size,), 0.5) + freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(model.axes_dim_rope, model.axes_lens, theta=10000) + a_content = torch.randn(a_len, instruction_feat_dim) + # Non-empty, fixed per-row reference latents route both rows through the + # TI2I reference-image refiner and the joint instruction/reference/noise + # concatenation, rather than the T2I path (which skips both entirely). + a_ref = torch.randn(in_channels, 6, 10) + b_ref = torch.randn(in_channels, 6, 10) + ref_image_hidden_states = [[a_ref], [b_ref]] + + def run(b_len, a_mask_len=a_len, a_tail_value=0.0): + """``a_mask_len`` marks how many of row A's leading positions are valid. + + Real usage always leaves it at ``a_len``; the negative control below + widens it to include A's own padding tail so that tail's content + starts to matter, proving the invariance check exercises real + masking rather than an architecture that ignores trailing positions + regardless of the mask. + """ + width = max(a_len, b_len) + instruction_hidden_states = torch.zeros(batch_size, width, instruction_feat_dim) + instruction_hidden_states[0, :a_len] = a_content + if width > a_len: + instruction_hidden_states[0, a_len:] = a_tail_value + torch.manual_seed(100 + b_len) + instruction_hidden_states[1, :b_len] = torch.randn(b_len, instruction_feat_dim) + + instruction_attention_mask = torch.zeros(batch_size, width, dtype=torch.bool) + instruction_attention_mask[0, :a_mask_len] = True + instruction_attention_mask[1, :b_len] = True + + with torch.no_grad(): + return model( + latents, + timestep, + instruction_hidden_states, + freqs_real, + instruction_attention_mask, + ref_image_hidden_states=ref_image_hidden_states, + ) + + out_short_b = run(b_len=3) + out_long_b = run(b_len=8) # B now sits in a much wider padded batch than A. + + torch.testing.assert_close( + out_short_b[0], + out_long_b[0], + rtol=1e-4, + atol=1e-4, + msg="row A changed when only B's padded length grew", + ) + + # Negative control: mark A's own padding tail "valid" (mask length == + # full width) and show that what sits there now measurably changes A's + # output -- a magnitude-based outlier there did not (SDPA masking makes + # softmax weight on a correctly-excluded position exactly zero regardless + # of its magnitude), so this instead varies the tail's *content* under a + # broken mask and checks the two runs diverge. + width = max(a_len, 8) + tail_x = torch.randn(width - a_len, instruction_feat_dim) + tail_y = torch.randn(width - a_len, instruction_feat_dim) + out_broken_x = run(b_len=8, a_mask_len=width, a_tail_value=tail_x) + out_broken_y = run(b_len=8, a_mask_len=width, a_tail_value=tail_y) + assert not torch.allclose(out_broken_x[0], out_broken_y[0], rtol=1e-4, atol=1e-4) + + # --------------------------------------------------------------------------- # The three tests below are for the fused qk-norm+RoPE kernel's wiring into # this transformer: the packed-table layout conversion, the fused path against diff --git a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py index 9ad022a10..b2dea0c44 100644 --- a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py +++ b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py @@ -14,6 +14,7 @@ Two groups: """ import os +import zlib from types import SimpleNamespace import pytest @@ -520,8 +521,17 @@ def test_constructor_rejects_non_cuda_sequence_parallelism( # --------------------------------------------------------------------------- +def _content_id(text: str) -> int: + """Bounded text IDs, stable across processes unlike builtin ``hash()``. + + Seed/reference checks use separate transformers so BF16 rounding cannot + hide their smaller signals beneath the content term. + """ + return zlib.crc32(text.encode()) % 9973 + + class _RecordingProcessor: - """Fake Qwen3VLProcessor: deterministic token ids derived from the text.""" + """Fake Qwen3VLProcessor with content-derived token IDs.""" def __init__(self): self.calls = [] @@ -533,8 +543,8 @@ class _RecordingProcessor: for i, messages in enumerate(prompts): system_text = messages[0]["content"][0]["text"] user_text = messages[1]["content"][0]["text"] - input_ids[i, 0] = len(system_text) % 997 - input_ids[i, 1] = len(user_text) % 997 + input_ids[i, 0] = _content_id(system_text) + input_ids[i, 1] = _content_id(user_text) input_ids[i, 2:] = torch.arange(2, _SEQ_LEN) + i attention_mask = torch.ones(batch, _SEQ_LEN, dtype=torch.long) attention_mask[:, -1] = 0 # fake right-padding @@ -1283,7 +1293,10 @@ def test_apply_chat_template_ti2i_places_image_before_text(): class _ImageAwareRecordingProcessor: - """Records whether reference images reached the processor.""" + """Record reference-image presence and encode text with stable IDs. + + Find text by type because TI2I puts the image first. + """ def __init__(self): self.calls = [] @@ -1295,7 +1308,13 @@ class _ImageAwareRecordingProcessor: has_image.append(any(c.get("type") == "image" for c in user_content)) self.calls.append({"prompts": prompts, "kwargs": kwargs, "has_image": has_image}) batch = len(prompts) - input_ids = torch.arange(batch * _SEQ_LEN, dtype=torch.long).view(batch, _SEQ_LEN) + input_ids = torch.zeros(batch, _SEQ_LEN, dtype=torch.long) + for i, messages in enumerate(prompts): + system_text = messages[0]["content"][0]["text"] + user_text = next(c["text"] for c in messages[1]["content"] if c.get("type") == "text") + input_ids[i, 0] = _content_id(system_text) + input_ids[i, 1] = _content_id(user_text) + input_ids[i, 2:] = torch.arange(2, _SEQ_LEN) + i attention_mask = torch.ones(batch, _SEQ_LEN, dtype=torch.long) return {"input_ids": input_ids, "attention_mask": attention_mask} @@ -1558,55 +1577,10 @@ def test_forward_image_guidance_ignored_without_reference(): # --------------------------------------------------------------------------- -# Request-batch: compatibility key, generator routing, output split +# Request-batch: generator routing, output split # --------------------------------------------------------------------------- -def test_boogu_batch_compatibility_key_t2i_stable_ti2i_unique(): - from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import _boogu_batch_compatibility_key - - # t2i: request_id does not enter the key -> t2i requests batch together. - assert _boogu_batch_compatibility_key(False, "req-a") == _boogu_batch_compatibility_key(False, "req-b") - assert _boogu_batch_compatibility_key(False, "req-a")[1] == "t2i" - - # ti2i: request_id in the key -> each edit gets a unique key, never co-batched. - assert _boogu_batch_compatibility_key(True, "req-a") != _boogu_batch_compatibility_key(True, "req-b") - assert _boogu_batch_compatibility_key(True, "req-a")[1] == "ti2i" - - # t2i and ti2i never share a key. - assert _boogu_batch_compatibility_key(False, "req-a") != _boogu_batch_compatibility_key(True, "req-a") - - -def test_pre_process_key_wiring_t2i_batches_ti2i_isolated(tmp_path): - # End-to-end: real pre-process sets request.batch_compatibility_key, and the - # scheduler's key builder reads it into condition_key. Two t2i requests share - # a key (co-batchable); two edit requests get distinct keys (batch=1). - import PIL.Image - - from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import get_boogu_image_pre_process_func - from vllm_omni.diffusion.request import OmniDiffusionRequest - from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key - from vllm_omni.inputs.data import OmniDiffusionSamplingParams - - pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path)) - - def condition_key(prompt, rid): - req = OmniDiffusionRequest( - prompt=prompt, sampling_params=OmniDiffusionSamplingParams(height=512, width=512), request_id=rid - ) - pre(req) - return build_request_batch_sampling_params_key(req).condition_key - - t2i_a = condition_key({"prompt": "a cat"}, "t-a") - t2i_b = condition_key({"prompt": "a dog"}, "t-b") - assert t2i_a == t2i_b and t2i_a[1] == "t2i" - - img = PIL.Image.new("RGB", (64, 64)) - ti2i_a = condition_key({"prompt": "edit", "multi_modal_data": {"image": img}}, "e-a") - ti2i_b = condition_key({"prompt": "edit", "multi_modal_data": {"image": img}}, "e-b") - assert ti2i_a != ti2i_b and ti2i_a[1] == "ti2i" - - class _GeneratorRecordingVAE: """Fake VAE that records the generator passed to each latent ``sample()``.""" @@ -1736,78 +1710,6 @@ def test_forward_request_batch_num_outputs_slices_and_generators(): assert float(outs[1].output[0, 0, 0, 0]) == 2.0 and float(outs[1].output[1, 0, 0, 0]) == 3.0 -def test_forward_batch_isolation_partner_content_and_seed(): - """CFG-on, B=2: request A's output must not change when only the - co-batched partner's prompt content, negative prompt, or seed changes. - - ``_FakeTransformer``/``_FakeScheduler`` are content-blind (always-zero - velocity, latents passed through unchanged), so they cannot catch a - cross-request value leak. This test swaps in a transformer whose output - depends on both ``instruction_embeds`` (content, positive or negative - depending on which CFG branch called it) and ``latents`` (seed), and a - scheduler that actually applies the predicted velocity, so a batching bug - that mixes rows in either the cond or uncond prediction would change A's - result. A negative-prompt-only perturbation is required to cover the - uncond branch: varying only the positive prompt never touches - ``negative_instruction_embeds``, so an earlier version of this test - passed even with a synthetic row-mixing bug injected into the uncond - predict() call (verified via a RED-arm check before this fix). - """ - - class _ContentAwareTransformer(_FakeTransformer): - def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs): - content = instruction_embeds.mean(dim=(1, 2)).view(-1, 1, 1, 1) - return latents + content - - class _ApplyingScheduler(_FakeScheduler): - def step(self, model_output, t, latents, return_dict=False): - return (model_output,) - - def run(prompt_a, seed_a, neg_a, prompt_b, seed_b, neg_b): - pipeline = _make_forward_pipeline() - pipeline.transformer = _ContentAwareTransformer() - pipeline.scheduler = _ApplyingScheduler() - kw = dict(height=64, width=64, num_inference_steps=2, guidance_scale=4.0, output_type="latent") - req = _wrap_request_batch( - [ - ( - {"prompt": prompt_a, "negative_prompt": neg_a}, - _sampling(**kw, generator=torch.Generator().manual_seed(seed_a)), - ), - ( - {"prompt": prompt_b, "negative_prompt": neg_b}, - _sampling(**kw, generator=torch.Generator().manual_seed(seed_b)), - ), - ] - ) - return pipeline.forward(req)[0].output - - baseline = run("a cat on a mat", 1, "ugly", "a dog in a park", 2, "blurry") - assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a totally different scene", 2, "blurry")) - assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a dog in a park", 999, "blurry")) - assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a dog in a park", 2, "watermark")) - - -def test_forward_batched_ti2i_fails_closed(): - # A batched ti2i must fail closed (it is gated to batch=1). - pipeline = _make_forward_pipeline() - - def edit_prompt(): - return { - "prompt": "make it winter", - "additional_information": {"preprocessed_image": torch.zeros(1, 3, 64, 64), "prompt_image": None}, - } - - req = _wrap_request_batch( - [ - (edit_prompt(), _sampling(num_inference_steps=1, guidance_scale=1.0)), - (edit_prompt(), _sampling(num_inference_steps=1, guidance_scale=1.0)), - ] - ) - with pytest.raises(RuntimeError, match="gated to batch=1"): - pipeline.forward(req) - - def test_supports_request_batch_enabled(): from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline diff --git a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image_batching.py b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image_batching.py new file mode 100644 index 000000000..3fff22ea6 --- /dev/null +++ b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image_batching.py @@ -0,0 +1,746 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Boogu-Image request-batching tests: compatibility keys, cross-request +isolation, CFG-parallel dispatch, and variable-length padding. + +Split out from ``test_pipeline_boogu_image.py`` to keep that file's +constructor/prompt-encoding focus separate from these request-batch-specific +cases. Shared pipeline/fake builders are imported from there rather than +duplicated. +""" + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from .test_pipeline_boogu_image import ( + _content_id, + _EditForwardVAE, + _FakeScheduler, + _FakeTransformer, + _make_edit_forward_pipeline, + _make_edit_od_config, + _make_forward_pipeline, + _sampling, + _wrap_request_batch, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _RefContentVAE(_EditForwardVAE): + """Fake VAE whose ``encode`` reflects the input image's mean into the latent.""" + + def encode(self, img): + val = float(img.float().mean().item()) + dist = SimpleNamespace(sample=lambda generator=None: torch.full((1, 4, 8, 8), val)) + return SimpleNamespace(latent_dist=dist) + + +def _edit_prompt(text, neg, ref_fill): + import PIL.Image + + return { + "prompt": text, + "negative_prompt": neg, + "additional_information": { + "prompt_image": PIL.Image.new("RGB", (64, 64)), + "preprocessed_image": torch.full((1, 3, 64, 64), ref_fill), + }, + } + + +# --------------------------------------------------------------------------- +# Request-batch: compatibility key +# --------------------------------------------------------------------------- + + +def test_boogu_batch_compatibility_key_t2i_and_ti2i_stable_and_isolated(): + from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import _boogu_batch_compatibility_key + + # Each task shares a stable key; guidance compatibility is checked upstream. + assert _boogu_batch_compatibility_key(False) == ("boogu_image", "t2i") + + assert _boogu_batch_compatibility_key(True) == ("boogu_image", "ti2i") + + # T2I and TI2I use different denoise paths. + assert _boogu_batch_compatibility_key(False) != _boogu_batch_compatibility_key(True) + + +def test_pre_process_key_wiring_t2i_and_ti2i_batch(tmp_path): + # Verify preprocessing keys reach the scheduler and keep T2I/TI2I apart. + import PIL.Image + + from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import get_boogu_image_pre_process_func + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path)) + + def condition_key(prompt, rid): + req = OmniDiffusionRequest( + prompt=prompt, sampling_params=OmniDiffusionSamplingParams(height=512, width=512), request_id=rid + ) + pre(req) + return build_request_batch_sampling_params_key(req).condition_key + + t2i_a = condition_key({"prompt": "a cat"}, "t-a") + t2i_b = condition_key({"prompt": "a dog"}, "t-b") + assert t2i_a == t2i_b and t2i_a[1] == "t2i" + + img = PIL.Image.new("RGB", (64, 64)) + ti2i_a = condition_key({"prompt": "edit", "multi_modal_data": {"image": img}}, "e-a") + ti2i_b = condition_key({"prompt": "edit", "multi_modal_data": {"image": img}}, "e-b") + assert ti2i_a == ti2i_b and ti2i_a[1] == "ti2i" + assert t2i_a != ti2i_a + + +def test_ti2i_batch_key_distinguishes_explicit_image_guidance(tmp_path): + import PIL.Image + + from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import get_boogu_image_pre_process_func + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path)) + image = PIL.Image.new("RGB", (64, 64)) + + def make_request(guidance_scale_2: float | None) -> OmniDiffusionRequest: + request = OmniDiffusionRequest( + prompt={"prompt": "edit", "multi_modal_data": {"image": image}}, + sampling_params=OmniDiffusionSamplingParams( + height=512, + width=512, + seed=123, + guidance_scale=2.0, + guidance_scale_2=guidance_scale_2, + ), + request_id="ti2i", + ) + pre(request) + return request + + omitted = build_request_batch_sampling_params_key(make_request(None)) + explicit = build_request_batch_sampling_params_key(make_request(2.0)) + + # Omission auto-fills the same value but must not enable Boogu's image-guidance branch. + assert omitted.condition_key == explicit.condition_key == ("boogu_image", "ti2i") + assert omitted.guidance_scale_2 == explicit.guidance_scale_2 == 2.0 + assert omitted.guidance_scale_2_provided is False + assert explicit.guidance_scale_2_provided is True + assert omitted != explicit + + +# --------------------------------------------------------------------------- +# Request-batch: cross-request isolation +# --------------------------------------------------------------------------- + + +def test_forward_batch_isolation_partner_content_and_seed(monkeypatch): + """Keep A unchanged under B's prompt, negative-prompt, or seed changes. + + Exercise the real CFG/encoding path at B=2 with signal-sensitive fakes. + """ + + class _ContentAwareTransformer(_FakeTransformer): + def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs): + content = instruction_embeds.mean(dim=(1, 2)).view(-1, 1, 1, 1) + return latents + content + + class _ApplyingScheduler(_FakeScheduler): + def step(self, model_output, t, latents, return_dict=False): + return (model_output,) + + def run(prompt_a, seed_a, neg_a, prompt_b, seed_b, neg_b): + pipeline = _make_forward_pipeline() + pipeline.transformer = _ContentAwareTransformer() + pipeline.scheduler = _ApplyingScheduler() + kw = dict(height=64, width=64, num_inference_steps=2, guidance_scale=4.0, output_type="latent") + req = _wrap_request_batch( + [ + ( + {"prompt": prompt_a, "negative_prompt": neg_a}, + _sampling(**kw, generator=torch.Generator().manual_seed(seed_a)), + ), + ( + {"prompt": prompt_b, "negative_prompt": neg_b}, + _sampling(**kw, generator=torch.Generator().manual_seed(seed_b)), + ), + ] + ) + return pipeline.forward(req)[0].output + + baseline = run("a cat on a mat", 1, "ugly", "a dog in a park", 2, "blurry") + assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a totally different scene", 2, "blurry")) + assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a dog in a park", 999, "blurry")) + assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a dog in a park", 2, "watermark")) + # Same-length partner rewrites exercise content rather than length. + assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a fox in a cave", 2, "blurry")) + assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a dog in a park", 2, "grainy")) + + # Sensitivity control: A's own prompt and negative prompt must affect A. + assert not torch.equal(baseline, run("a totally different scene", 1, "ugly", "a dog in a park", 2, "blurry")) + assert not torch.equal(baseline, run("a cat on a mat", 1, "watermark", "a dog in a park", 2, "blurry")) + + # Test seed separately: BF16 content magnitudes can hide O(1) noise. + class _LatentsOnlyTransformer(_FakeTransformer): + def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs): + return latents + + def run_seed(seed_a, seed_b): + pipeline = _make_forward_pipeline() + pipeline.transformer = _LatentsOnlyTransformer() + pipeline.scheduler = _ApplyingScheduler() + kw = dict(height=64, width=64, num_inference_steps=2, guidance_scale=4.0, output_type="latent") + req = _wrap_request_batch( + [ + ("a cat on a mat", _sampling(**kw, generator=torch.Generator().manual_seed(seed_a))), + ("a dog in a park", _sampling(**kw, generator=torch.Generator().manual_seed(seed_b))), + ] + ) + return pipeline.forward(req)[0].output + + seed_baseline = run_seed(1, 2) + assert torch.equal(seed_baseline, run_seed(1, 999)) # B's seed must not affect A + assert not torch.equal(seed_baseline, run_seed(999, 2)) # A's own seed must affect A + + # RED: compare partner seeds under the same reversed-generator mutation. + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + + real_collate = DiffusionRequestBatch.collate_request_generators + + def reversed_collate(self, num_outputs_per_prompt, default_generator): + result = real_collate(self, num_outputs_per_prompt, default_generator) + return list(reversed(result)) if isinstance(result, list) and len(result) > 1 else result + + monkeypatch.setattr(DiffusionRequestBatch, "collate_request_generators", reversed_collate) + buggy_seed_baseline = run_seed(1, 2) + assert not torch.equal(buggy_seed_baseline, run_seed(1, 999)) + + +def test_forward_batched_ti2i_partner_reference_isolation(monkeypatch): + """Keep A independent of B's content, seed, and reference at B=2. + + Double guidance exercises all three prediction branches; separate fakes + keep seed/reference sensitivity independent of content magnitude. + """ + + class _ContentRefTransformer(_FakeTransformer): + def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs): + out = latents + instruction_embeds.mean(dim=(1, 2)).view(-1, 1, 1, 1) + # Reference latents are request-major: one optional list per row. + ref = kwargs.get("ref_image_hidden_states") + if ref is not None: + ref_means = torch.tensor( + [float(r[0].float().mean()) if r is not None else 0.0 for r in ref], + dtype=out.dtype, + device=out.device, + ).view(-1, 1, 1, 1) + out = out + ref_means + return out + + class _ApplyingScheduler(_FakeScheduler): + def step(self, model_output, t, latents, return_dict=False): + return (model_output,) + + def run(prompt_a, neg_a, seed_a, ref_fill_a, prompt_b, neg_b, seed_b, ref_fill_b): + pipeline = _make_edit_forward_pipeline() + pipeline.transformer = _ContentRefTransformer() + pipeline.scheduler = _ApplyingScheduler() + pipeline.vae = _RefContentVAE() + kw = dict( + height=64, + width=64, + num_inference_steps=2, + guidance_scale=5.0, + guidance_scale_2=2.0, + output_type="latent", + ) + req = _wrap_request_batch( + [ + ( + _edit_prompt(prompt_a, neg_a, ref_fill_a), + _sampling(**kw, generator=torch.Generator().manual_seed(seed_a)), + ), + ( + _edit_prompt(prompt_b, neg_b, ref_fill_b), + _sampling(**kw, generator=torch.Generator().manual_seed(seed_b)), + ), + ] + ) + return pipeline.forward(req)[0].output + + A = ("make it winter", "ugly", 1, 0.25) + baseline = run(*A, "a dog in a park", "blurry", 2, 0.75) + assert torch.equal(baseline, run(*A, "a totally different scene", "blurry", 2, 0.75)) + assert torch.equal(baseline, run(*A, "a dog in a park", "watermark", 2, 0.75)) + assert torch.equal(baseline, run(*A, "a dog in a park", "blurry", 999, 0.75)) + assert torch.equal(baseline, run(*A, "a dog in a park", "blurry", 2, 0.5)) + # Same-length partner rewrites exercise content rather than length. + assert torch.equal(baseline, run(*A, "a fox in a cave", "blurry", 2, 0.75)) + assert torch.equal(baseline, run(*A, "a dog in a park", "grainy", 2, 0.75)) + + # Sensitivity control: A's own prompt and negative prompt must affect A. + B = ("a dog in a park", "blurry", 2, 0.75) + assert not torch.equal(baseline, run("a totally different scene", "ugly", 1, 0.25, *B)) + assert not torch.equal(baseline, run("make it winter", "watermark", 1, 0.25, *B)) + + # Keep reference and seed signals separate from the larger BF16 content term. + class _RefOnlyTransformer(_FakeTransformer): + def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs): + ref = kwargs.get("ref_image_hidden_states") + if ref is None: + # Branches [r, r, 0] give 2r at image guidance 2. + # Returning latents would cancel r after two scheduler steps. + return torch.zeros_like(latents) + ref_means = torch.tensor( + [float(r[0].float().mean()) if r is not None else 0.0 for r in ref], + dtype=latents.dtype, + device=latents.device, + ).view(-1, 1, 1, 1) + return torch.zeros_like(latents) + ref_means + + class _LatentsOnlyTransformer(_FakeTransformer): + def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs): + return latents + + def run_isolated(transformer_cls, ref_fill_a, seed_a, ref_fill_b, seed_b): + pipeline = _make_edit_forward_pipeline() + pipeline.transformer = transformer_cls() + pipeline.scheduler = _ApplyingScheduler() + pipeline.vae = _RefContentVAE() + kw = dict( + height=64, width=64, num_inference_steps=2, guidance_scale=5.0, guidance_scale_2=2.0, output_type="latent" + ) + req = _wrap_request_batch( + [ + ( + _edit_prompt("make it winter", "ugly", ref_fill_a), + _sampling(**kw, generator=torch.Generator().manual_seed(seed_a)), + ), + ( + _edit_prompt("a dog in a park", "blurry", ref_fill_b), + _sampling(**kw, generator=torch.Generator().manual_seed(seed_b)), + ), + ] + ) + return pipeline.forward(req)[0].output + + ref_baseline = run_isolated(_RefOnlyTransformer, 0.25, 1, 0.75, 2) + assert torch.equal(ref_baseline, run_isolated(_RefOnlyTransformer, 0.25, 1, 0.5, 2)) # B's ref must not affect A + assert not torch.equal( + ref_baseline, run_isolated(_RefOnlyTransformer, 0.75, 1, 0.75, 2) + ) # A's own ref must affect A + + seed_baseline = run_isolated(_LatentsOnlyTransformer, 0.25, 1, 0.75, 2) + # Only A's seed changes A. + assert torch.equal(seed_baseline, run_isolated(_LatentsOnlyTransformer, 0.25, 1, 0.75, 999)) + assert not torch.equal(seed_baseline, run_isolated(_LatentsOnlyTransformer, 0.25, 999, 0.75, 2)) + + # RED: compare baseline and partner changes under each active mutation. + from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + + real_build_ref_latents = BooguImagePipeline._build_ref_latents + + def reversed_build_ref_latents(self, preprocessed_images, num_images_per_prompt, device, generators=None): + swapped = list(reversed(preprocessed_images)) + return real_build_ref_latents(self, swapped, num_images_per_prompt, device, generators) + + monkeypatch.setattr(BooguImagePipeline, "_build_ref_latents", reversed_build_ref_latents) + buggy_ref_baseline = run_isolated(_RefOnlyTransformer, 0.25, 1, 0.75, 2) + assert not torch.equal(buggy_ref_baseline, run_isolated(_RefOnlyTransformer, 0.25, 1, 0.5, 2)) + monkeypatch.undo() + + real_collate = DiffusionRequestBatch.collate_request_generators + + def reversed_collate(self, num_outputs_per_prompt, default_generator): + result = real_collate(self, num_outputs_per_prompt, default_generator) + return list(reversed(result)) if isinstance(result, list) and len(result) > 1 else result + + monkeypatch.setattr(DiffusionRequestBatch, "collate_request_generators", reversed_collate) + buggy_seed_baseline = run_isolated(_LatentsOnlyTransformer, 0.25, 1, 0.75, 2) + assert not torch.equal(buggy_seed_baseline, run_isolated(_LatentsOnlyTransformer, 0.25, 1, 0.75, 999)) + + +# --------------------------------------------------------------------------- +# Request-batch: CFG-parallel dispatch +# --------------------------------------------------------------------------- + + +def _patch_cfg_parallel_state(monkeypatch, *, world_size, rank, fake_group): + """Make ``CFGParallelMixin`` see an active M-rank CFG group. + + ``is_cfg_group_initialized()`` reads ``parallel_state._CFG`` through that + function's own closure, so patching ``_CFG`` flips it regardless of which + module imported the function. ``get_classifier_free_guidance_world_size``, + ``get_classifier_free_guidance_rank``, and ``get_cfg_group`` are each + rebound by name inside ``cfg_parallel.py`` (``from ...parallel_state + import ...``), so those three must be patched on both modules or + ``cfg_parallel``'s copy keeps resolving to the real (single-rank) + implementation. + """ + from vllm_omni.diffusion.distributed import cfg_parallel as cfg_parallel_mod + from vllm_omni.diffusion.distributed import parallel_state + + monkeypatch.setattr(parallel_state, "_CFG", fake_group) + for module in (parallel_state, cfg_parallel_mod): + monkeypatch.setattr(module, "get_classifier_free_guidance_world_size", lambda: world_size) + monkeypatch.setattr(module, "get_classifier_free_guidance_rank", lambda: rank) + monkeypatch.setattr(module, "get_cfg_group", lambda: fake_group) + + +class _SlotTableCfgGroup: + """Fake CFG process group backed by a precomputed per-rank slot table. + + Mirrors the real ``all_gather`` contract: called once per (slot, tensor + element) in slot-major order. Validates that the calling rank's local + prediction (``input_``) matches what that rank was supposed to compute + for this slot, and uses the *actual* ``input_`` (not the precomputed + value) for its own rank in the returned list -- a dispatch bug (wrong + branch id, wrong rank assignment) fails the assertion below instead of + being silently overwritten by the precomputed "correct" table. + """ + + def __init__(self, world_size, rank_in_group, slots_per_rank): + self.world_size = world_size + self.rank_in_group = rank_in_group + self._slots_per_rank = slots_per_rank + self.gather_calls = 0 + + def all_gather(self, input_, dim=0, separate_tensors=False): + del dim + assert separate_tensors + slot_idx = self.gather_calls + self.gather_calls += 1 + expected_local = self._slots_per_rank[self.rank_in_group][slot_idx] + torch.testing.assert_close( + input_, + expected_local, + rtol=0, + atol=0, + msg=( + f"rank {self.rank_in_group} computed the wrong local branch prediction " + f"for slot {slot_idx} (dispatch bug)" + ), + ) + return [ + input_ if rank == self.rank_in_group else self._slots_per_rank[rank][slot_idx] + for rank in range(self.world_size) + ] + + +def _make_counting_predict_noise(real_predict_noise, call_log): + def counting_predict_noise(**kwargs): + call_log.append(kwargs) + return real_predict_noise(**kwargs) + + return counting_predict_noise + + +def test_forward_batched_ti2i_double_guidance_cfg_parallel_matches_sequential(monkeypatch): + """CFG-parallel dispatch/gather must reproduce the sequential result. + + Builds a real TI2I B=2 double-guidance request batch (distinct prompts, + negatives, seeds, and references per row, via the same signal-visible + fakes as ``test_forward_batched_ti2i_partner_reference_isolation``), then + replays the exact same branch kwargs through the real + ``_dispatch_branches``/``all_gather`` path for 2- and 3-rank CFG-parallel + layouts. A wrongly dispatched branch or a row swap during gather would change + the result, since request A and B carry distinct content/reference + signals into every branch. + """ + from vllm_omni.diffusion.distributed.cfg_parallel import _dispatch_branches + from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline + + class _ContentRefTransformer(_FakeTransformer): + def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs): + # FP32 throughout: the encode path is BF16, and the content term + # (derived from crc32-scaled token ids) reaches several hundred in + # magnitude, which would round away the 0.25/0.75 reference delta + # in BF16 and make the with-ref/without-ref branches indistinguishable. + out = latents.float() + instruction_embeds.float().mean(dim=(1, 2)).view(-1, 1, 1, 1) + ref = kwargs.get("ref_image_hidden_states") + if ref is not None: + ref_means = torch.tensor( + [float(r[0].float().mean()) if r is not None else 0.0 for r in ref], + dtype=out.dtype, + device=out.device, + ).view(-1, 1, 1, 1) + out = out + ref_means + return out + + pipeline = _make_edit_forward_pipeline() + pipeline.transformer = _ContentRefTransformer() + pipeline.vae = _RefContentVAE() + kw = dict( + height=64, + width=64, + num_inference_steps=1, + guidance_scale=5.0, + guidance_scale_2=2.0, + output_type="latent", + ) + req = _wrap_request_batch( + [ + ( + _edit_prompt("make it winter", "ugly", 0.25), + _sampling(**kw, generator=torch.Generator().manual_seed(1)), + ), + ( + _edit_prompt("a dog in a park", "blurry", 0.75), + _sampling(**kw, generator=torch.Generator().manual_seed(2)), + ), + ] + ) + + # Step 1: run the real sequential (single-rank) forward once, capturing + # the exact double-guidance branch kwargs and the resulting model_pred + # for this request batch's one denoising step. + captured: dict[str, Any] = {} + real_multi_branch = BooguImagePipeline.predict_noise_with_multi_branch_cfg + + def capturing_multi_branch( + self, do_true_cfg, true_cfg_scale, branches_kwargs, cfg_normalize=False, output_slice=None + ): + result = real_multi_branch( + self, do_true_cfg, true_cfg_scale, branches_kwargs, cfg_normalize=cfg_normalize, output_slice=output_slice + ) + captured.update( + branches_kwargs=branches_kwargs, + true_cfg_scale=true_cfg_scale, + cfg_normalize=cfg_normalize, + output_slice=output_slice, + sequential_pred=result, + ) + return result + + monkeypatch.setattr(BooguImagePipeline, "predict_noise_with_multi_branch_cfg", capturing_multi_branch) + pipeline.forward(req) + monkeypatch.undo() + + branches_kwargs = captured["branches_kwargs"] + assert len(branches_kwargs) == 3 # cond+ref, neg+ref, neg+no-ref + sequential_pred = captured["sequential_pred"] + + # Step 2: replay the identical branches through the real dispatch/gather + # path for 2- and 3-rank layouts (rank 0: [0, 2] / [0]; rank 1: [1] / [1]; + # rank 2 only exists for the 3-rank layout: [2]). + n_branches = len(branches_kwargs) + # Branch predictions are a deterministic function of (pipeline, branches_kwargs) + # alone; compute once and reuse across both rank-layout replays below. + with torch.no_grad(): + branch_preds = [pipeline.predict_noise(**branch_kw) for branch_kw in branches_kwargs] + + # If any two branches ever collapsed to the same tensor, a wrongly + # dispatched or wrongly gathered branch below could go unnoticed even with exact-equality + # checks -- confirm the fakes are actually signal-visible. + for i in range(n_branches): + for j in range(i + 1, n_branches): + assert not torch.equal(branch_preds[i], branch_preds[j]), ( + f"branches {i} and {j} are numerically identical; the fakes cannot " + "distinguish a branch mix-up from a correct dispatch" + ) + + for cfg_world_size in (2, 3): + assignments = _dispatch_branches(n_branches, cfg_world_size) + max_per_rank = max(len(a) for a in assignments) + + slots_per_rank = [] + for branch_ids in assignments: + slots = [branch_preds[bid] for bid in branch_ids] + while len(slots) < max_per_rank: + slots.append(torch.zeros_like(branch_preds[0])) + slots_per_rank.append(slots) + + rank_outputs = [] + for cfg_rank in range(cfg_world_size): + fake_group = _SlotTableCfgGroup(cfg_world_size, cfg_rank, slots_per_rank) + _patch_cfg_parallel_state(monkeypatch, world_size=cfg_world_size, rank=cfg_rank, fake_group=fake_group) + + call_log: list[dict] = [] + pipeline.predict_noise = _make_counting_predict_noise(pipeline.predict_noise, call_log) + try: + with torch.no_grad(): + rank_pred = pipeline.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale=captured["true_cfg_scale"], + branches_kwargs=branches_kwargs, + cfg_normalize=captured["cfg_normalize"], + output_slice=captured["output_slice"], + ) + finally: + del pipeline.predict_noise # drop the instance override, restore the class method + monkeypatch.undo() + + # A silent fallback to the sequential path (e.g. the CFG-group + # patch not taking effect) would compute every branch locally and + # never call all_gather; catch that instead of letting it pass + # only because sequential also happens to match the baseline. + expected_calls = len(assignments[cfg_rank]) + assert len(call_log) == expected_calls, ( + f"cfg_world_size={cfg_world_size} rank={cfg_rank} called predict_noise " + f"{len(call_log)}x, expected {expected_calls}x (silent sequential fallback?)" + ) + assert fake_group.gather_calls == max_per_rank, ( + f"cfg_world_size={cfg_world_size} rank={cfg_rank} issued {fake_group.gather_calls} " + f"all_gather call(s), expected {max_per_rank} (silent sequential fallback?)" + ) + rank_outputs.append(rank_pred) + + for cfg_rank, rank_pred in enumerate(rank_outputs): + torch.testing.assert_close( + rank_pred, + sequential_pred, + rtol=0, + atol=0, + msg=( + f"cfg_world_size={cfg_world_size} rank={cfg_rank} diverged from the " + "sequential (single-rank) baseline" + ), + ) + + +# --------------------------------------------------------------------------- +# Request-batch: variable-length padding +# --------------------------------------------------------------------------- + + +def test_forward_batched_ti2i_variable_length_partner_padding(monkeypatch): + """A must be invariant to B's real padded length, not just same-length content. + + ``test_forward_batched_ti2i_partner_reference_isolation`` pins every fake + encoder call to a fixed ``_SEQ_LEN`` with an all-ones mask, so it never + exercises production's ``padding="longest"`` batching or a real 0/1 + ``instruction_attention_mask``. Here the fake processor tokenizes by + content length and right-pads to the batch max, and the fake transformer + does a masked reduction over ``instruction_attention_mask`` instead of a + blind ``.mean()`` -- so A's row can only stay correct if the real + encode/reshape/CFG-duplication plumbing threads the real per-row mask + through unchanged when B's length (or B's negative-prompt length) grows. + """ + + class _VariableLengthImageAwareProcessor: + """Tokenizes by word count and right-pads to the batch max length.""" + + def __init__(self, force_all_ones_mask=False): + self.calls = [] + self.force_all_ones_mask = force_all_ones_mask + + def apply_chat_template(self, prompts, **kwargs): + has_image = [] + rows = [] + for messages in prompts: + user_content = messages[1]["content"] + has_image.append(any(c.get("type") == "image" for c in user_content)) + system_text = messages[0]["content"][0]["text"] + user_text = next(c["text"] for c in messages[1]["content"] if c.get("type") == "text") + length = 2 + len(user_text.split()) + ids = torch.zeros(length, dtype=torch.long) + ids[0] = _content_id(system_text) + ids[1] = _content_id(user_text) + if length > 2: + ids[2:] = torch.arange(2, length) + rows.append(ids) + + batch = len(prompts) + max_len = max(len(r) for r in rows) + input_ids = torch.zeros(batch, max_len, dtype=torch.long) + attention_mask = torch.zeros(batch, max_len, dtype=torch.long) + for i, ids in enumerate(rows): + input_ids[i, : len(ids)] = ids + attention_mask[i, : len(ids)] = 1 + if self.force_all_ones_mask: + attention_mask = torch.ones_like(attention_mask) + + self.calls.append( + {"prompts": prompts, "kwargs": kwargs, "has_image": has_image, "lengths": [len(r) for r in rows]} + ) + return {"input_ids": input_ids, "attention_mask": attention_mask} + + class _MaskedContentRefTransformer(_FakeTransformer): + """Masked-mean over real padding, unlike the fixed-``_SEQ_LEN`` fakes.""" + + def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs): + mask = instruction_attention_mask.float().unsqueeze(-1) + valid_counts = mask.sum(dim=1).clamp(min=1.0) + content = ((instruction_embeds.float() * mask).sum(dim=1) / valid_counts).mean(dim=-1).view(-1, 1, 1, 1) + out = latents.float() + content + ref = kwargs.get("ref_image_hidden_states") + if ref is not None: + ref_means = torch.tensor( + [float(r[0].float().mean()) if r is not None else 0.0 for r in ref], + dtype=out.dtype, + device=out.device, + ).view(-1, 1, 1, 1) + out = out + ref_means + return out + + class _ApplyingScheduler(_FakeScheduler): + def step(self, model_output, t, latents, return_dict=False): + return (model_output,) + + A = _edit_prompt("make it winter", "ugly", 0.25) + + def run(partner=None, force_all_ones_mask=False): + pipeline = _make_edit_forward_pipeline() + pipeline.processor = _VariableLengthImageAwareProcessor(force_all_ones_mask=force_all_ones_mask) + pipeline.transformer = _MaskedContentRefTransformer() + pipeline.scheduler = _ApplyingScheduler() + pipeline.vae = _RefContentVAE() + kw = dict( + height=64, + width=64, + num_inference_steps=1, + guidance_scale=5.0, + guidance_scale_2=2.0, + output_type="latent", + ) + items = [(A, _sampling(**kw, generator=torch.Generator().manual_seed(1)))] + if partner is not None: + items.append((partner, _sampling(**kw, generator=torch.Generator().manual_seed(2)))) + req = _wrap_request_batch(items) + outs = pipeline.forward(req) + return outs[0].output, pipeline.processor.calls + + a_alone, _ = run(partner=None) + + short_b = _edit_prompt("a dog", "blurry", 0.75) + long_b = _edit_prompt("a dog running through a large green park full of tall autumn trees", "blurry", 0.75) + long_neg_b = _edit_prompt( + "a dog in a park", "extremely blurry ugly low quality watermark grainy artifact noisy", 0.75 + ) + + a_with_short_b, calls_short = run(partner=short_b) + a_with_long_b, calls_long = run(partner=long_b) + a_with_long_neg_b, calls_long_neg = run(partner=long_neg_b) + + # calls[0] is the positive-instruction batch call, calls[1] the negative. + width_short = max(calls_short[0]["lengths"]) + width_long = max(calls_long[0]["lengths"]) + assert width_long > width_short, "B's longer instruction must actually widen the padded batch" + + a_len_in_short = calls_short[0]["lengths"][0] + a_len_in_long = calls_long[0]["lengths"][0] + assert a_len_in_short == a_len_in_long # A's own content length never changes... + assert width_long > a_len_in_long # ...but A now sits inside a wider, padded row. + + neg_width_short = max(calls_short[1]["lengths"]) + neg_width_long_neg = max(calls_long_neg[1]["lengths"]) + assert neg_width_long_neg > neg_width_short, "B's longer negative prompt must widen the negative-branch batch too" + + assert torch.equal(a_alone, a_with_short_b) + assert torch.equal(a_alone, a_with_long_b) + assert torch.equal(a_alone, a_with_long_neg_b) + + # Negative control: force an all-ones mask (ignore real padding). A's row + # then sees B's padding-region positions as "real" content, so this must + # break invariance -- otherwise the checks above could pass vacuously + # even with no masking at all, exactly like the fixed-_SEQ_LEN fakes. + a_with_long_b_broken_mask, _ = run(partner=long_b, force_all_ones_mask=True) + assert not torch.equal(a_alone, a_with_long_b_broken_mask) diff --git a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py index d96d4a6c1..09c85f9ab 100644 --- a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py +++ b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py @@ -112,18 +112,15 @@ def get_boogu_image_post_process_func(od_config: OmniDiffusionConfig): return post_process_func -def _boogu_batch_compatibility_key(has_reference: bool, request_id: str) -> tuple: - """Request-batch isolation key. ``forward`` reads shared guidance/shape fields - from the batch's first request, so t2i and ti2i must not share a key. +def _boogu_batch_compatibility_key(has_reference: bool) -> tuple: + """Separate T2I and TI2I denoise paths. - ti2i is held at batch=1 (request-unique key) because - ``guidance_scale_2_provided`` is absent from ``RequestBatchSamplingParamsKey`` - while ``forward`` reads it from the first request, so mixed-``_provided`` edits - with equal numeric ``guidance_scale_2`` would co-batch into the wrong mode. + Shared shape/guidance fields, including ``guidance_scale_2_provided``, + are checked by ``RequestBatchSamplingParamsKey``. """ if not has_reference: return ("boogu_image", "t2i") - return ("boogu_image", "ti2i", request_id) + return ("boogu_image", "ti2i") def get_boogu_image_pre_process_func(od_config: OmniDiffusionConfig): @@ -157,14 +154,14 @@ def get_boogu_image_pre_process_func(od_config: OmniDiffusionConfig): prompt = request.prompt if isinstance(prompt, str): # Plain-text prompt cannot carry an image -> text-to-image. - request.batch_compatibility_key = _boogu_batch_compatibility_key(False, request.request_id) + request.batch_compatibility_key = _boogu_batch_compatibility_key(False) return request multi_modal_data = prompt.get("multi_modal_data") or {} raw_image = multi_modal_data.get("image") if not raw_image: # No reference image -> text-to-image (Base checkpoint). - request.batch_compatibility_key = _boogu_batch_compatibility_key(False, request.request_id) + request.batch_compatibility_key = _boogu_batch_compatibility_key(False) return request if isinstance(raw_image, list): @@ -199,7 +196,7 @@ def get_boogu_image_pre_process_func(od_config: OmniDiffusionConfig): prompt["additional_information"]["preprocessed_image"] = preprocessed_image prompt["additional_information"]["prompt_image"] = prompt_image request.prompt = prompt - request.batch_compatibility_key = _boogu_batch_compatibility_key(True, request.request_id) + request.batch_compatibility_key = _boogu_batch_compatibility_key(True) return request return pre_process_func @@ -961,14 +958,6 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports if self._is_turbo and req.num_reqs > 1: raise RuntimeError("BooguImageTurboPipeline does not support request batching.") - # Fail-closed: a batched ti2i must never reach here (it is gated to batch=1). - if has_reference and req.num_reqs > 1: - raise RuntimeError( - f"BooguImagePipeline received a batched TI2I (edit) request " - f"(num_reqs={req.num_reqs}); TI2I batching is gated to batch=1 " - "pending guidance-mode / compatibility-key validation." - ) - sampling_params_list = req.sampling_params_list # Shared shape/step/guidance fields are guaranteed identical across the # batch by RequestBatchSamplingParamsKey; request-local values (seeds,