[Diffusion] Enable request-level batching for Boogu-Image (T2I) (#6968)

Signed-off-by: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -520,11 +520,25 @@ def _make_forward_pipeline():
     return pipeline
 
 
-def _make_request_batch(prompt, **sampling_overrides):
-    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+def _wrap_request_batch(items):
+    """Wrap ``(prompt, sampling_params)`` pairs in a real DiffusionRequestBatch.
 
-    sampling = OmniDiffusionSamplingParams(**sampling_overrides)
-    return SimpleNamespace(prompts=[prompt], sampling_params=sampling)
+    The pipeline's request-batch ``forward`` reads ``sampling_params_list`` and
+    ``collate_request_generators``, so the fake single-namespace shim no longer
+    satisfies the contract; construct the genuine wrapper instead.
+    """
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+    from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+
+    requests = [
+        OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id=f"req-{i}")
+        for i, (prompt, sampling) in enumerate(items)
+    ]
+    return DiffusionRequestBatch(requests=requests)
+
+
+def _make_request_batch(prompt, **sampling_overrides):
+    return _wrap_request_batch([(prompt, _sampling(**sampling_overrides))])
 
 
 def test_forward_returns_diffusion_output():
@@ -533,8 +547,11 @@ def test_forward_returns_diffusion_output():
     pipeline = _make_forward_pipeline()
     req = _make_request_batch("a cat", height=64, width=64, num_inference_steps=2)
 
-    out = pipeline.forward(req)
+    outs = pipeline.forward(req)
 
+    # Request-batch forward returns one DiffusionOutput per request.
+    assert isinstance(outs, list) and len(outs) == 1
+    out = outs[0]
     assert isinstance(out, DiffusionOutput)
     assert isinstance(out.output, torch.Tensor)
     assert out.output.shape[0] == 1
@@ -937,7 +954,7 @@ def _make_edit_request(**sampling_overrides):
         },
     }
     sampling = _sampling(**sampling_overrides)
-    return SimpleNamespace(prompts=[prompt], sampling_params=sampling)
+    return _wrap_request_batch([(prompt, sampling)])
 
 
 def _sampling(**overrides):
@@ -1035,10 +1052,267 @@ def test_forward_image_guidance_ignored_without_reference():
     sampling = _sampling(height=64, width=64, num_inference_steps=num_steps, guidance_scale=5.0, guidance_scale_2=2.0)
     sampling.guidance_scale_provided = True
     sampling.guidance_scale_2_provided = True
-    req = SimpleNamespace(prompts=[{"prompt": "a cat"}], sampling_params=sampling)
+    req = _wrap_request_batch([({"prompt": "a cat"}, sampling)])
 
     pipeline.forward(req)
 
     # t2i text CFG: cond + uncond -> 2 predictions/step, no ref anywhere.
     assert _count_per_step(pipeline.transformer.calls, num_steps) == 2
     assert all(ref is None for ref in pipeline.transformer.calls)
+
+
+# ---------------------------------------------------------------------------
+# Request-batch: compatibility key, 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()``."""
+
+    dtype = torch.float32
+
+    def __init__(self):
+        self.config = SimpleNamespace(scaling_factor=1.0, shift_factor=0.0)
+        self.seen_generators = []
+
+    def encode(self, img):
+        def sample(generator=None):
+            self.seen_generators.append(generator)
+            return torch.zeros(1, 4, 2, 2)
+
+        return SimpleNamespace(latent_dist=SimpleNamespace(sample=sample))
+
+
+def test_build_ref_latents_uses_per_request_generators():
+    pipeline = _make_encode_pipeline()
+    pipeline.vae = _GeneratorRecordingVAE()
+
+    g0 = torch.Generator().manual_seed(0)
+    g1 = torch.Generator().manual_seed(1)
+    preprocessed = [torch.zeros(1, 3, 16, 16), None, torch.zeros(1, 3, 16, 16)]
+
+    pipeline._build_ref_latents(
+        preprocessed,
+        num_images_per_prompt=1,
+        device=torch.device("cpu"),
+        generators=[g0, None, g1],
+    )
+
+    # Each present reference row samples with its own request's generator; the
+    # None reference row consumes none.
+    assert pipeline.vae.seen_generators == [g0, g1]
+
+
+def test_forward_request_batch_collates_generators_and_splits_outputs():
+    from vllm_omni.diffusion.data import DiffusionOutput
+
+    pipeline = _make_forward_pipeline()
+    recorded = {}
+
+    def fake_encode_prompt(prompt, **kwargs):
+        n = len(prompt)
+        embeds = torch.zeros(n, _SEQ_LEN, _EMBED_DIM)
+        mask = torch.ones(n, _SEQ_LEN, dtype=torch.long)
+        return embeds, mask, None, None
+
+    def fake_prepare_latents(batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        recorded["generator"] = generator
+        # A distinct constant per batch row so the split can be verified.
+        rows = torch.arange(batch_size, dtype=torch.float32).view(batch_size, 1, 1, 1)
+        return rows.repeat(1, num_channels_latents, 2, 2)
+
+    pipeline.encode_prompt = fake_encode_prompt
+    pipeline.prepare_latents = fake_prepare_latents
+
+    # T2I, CFG off, 1 step, latent output -> avoids VAE decode and CFG branches.
+    s0 = _sampling(height=64, width=64, num_inference_steps=1, guidance_scale=1.0, output_type="latent")
+    s1 = _sampling(height=64, width=64, num_inference_steps=1, guidance_scale=1.0, output_type="latent")
+    req = _wrap_request_batch([("a cat", s0), ("a dog", s1)])
+    g0 = torch.Generator().manual_seed(0)
+    g1 = torch.Generator().manual_seed(1)
+    req.requests[0].sampling_params.generator = g0
+    req.requests[1].sampling_params.generator = g1
+
+    outs = pipeline.forward(req)
+
+    # forward() collates per-request generators rather than reusing req0's.
+    assert recorded["generator"] == [g0, g1]
+    # One DiffusionOutput per request, in scheduler order.
+    assert len(outs) == 2
+    assert all(isinstance(o, DiffusionOutput) for o in outs)
+    assert float(outs[0].output[0, 0, 0, 0]) == 0.0
+    assert float(outs[1].output[0, 0, 0, 0]) == 1.0
+
+
+def test_reshape_mask_is_request_major_with_distinct_masks():
+    # B=2, N=2 with per-request-distinct masks: reshaped rows must be
+    # request-major [p0, p0, p1, p1] (repeat_interleave) to match the embed
+    # layout, not tiled [p0, p1, p0, p1] (a plain repeat).
+    pipeline = _make_encode_pipeline()
+    embeds = torch.zeros(2, _SEQ_LEN, _EMBED_DIM)
+    mask = torch.tensor([[1, 1, 1, 1, 1, 0], [1, 1, 1, 0, 0, 0]], dtype=torch.long)
+    _, _, _, reshaped_mask = pipeline._reshape_embeds_and_mask(embeds, mask, 2)
+    assert reshaped_mask.shape == (4, _SEQ_LEN)
+    assert torch.equal(reshaped_mask[0], mask[0])
+    assert torch.equal(reshaped_mask[1], mask[0])
+    assert torch.equal(reshaped_mask[2], mask[1])
+    assert torch.equal(reshaped_mask[3], mask[1])
+
+
+def test_forward_request_batch_num_outputs_slices_and_generators():
+    pipeline = _make_forward_pipeline()
+    recorded = {}
+
+    def fake_encode_prompt(prompt, **kwargs):
+        n = len(prompt)
+        return torch.zeros(n, _SEQ_LEN, _EMBED_DIM), torch.ones(n, _SEQ_LEN, dtype=torch.long), None, None
+
+    def fake_prepare_latents(batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        recorded["generator"] = generator
+        rows = torch.arange(batch_size, dtype=torch.float32).view(batch_size, 1, 1, 1)
+        return rows.repeat(1, num_channels_latents, 2, 2)
+
+    pipeline.encode_prompt = fake_encode_prompt
+    pipeline.prepare_latents = fake_prepare_latents
+
+    kw = dict(
+        height=64, width=64, num_inference_steps=1, guidance_scale=1.0, output_type="latent", num_outputs_per_prompt=2
+    )
+    req = _wrap_request_batch([("a cat", _sampling(**kw)), ("a dog", _sampling(**kw))])
+    g0 = torch.Generator().manual_seed(0)
+    g1 = torch.Generator().manual_seed(1)
+    req.requests[0].sampling_params.generator = g0
+    req.requests[1].sampling_params.generator = g1
+
+    outs = pipeline.forward(req)
+
+    # collated generator is request-major, output-minor.
+    assert recorded["generator"] == [g0, g0, g1, g1]
+    # each request gets its own 2-output slice: req0 rows [0,1], req1 rows [2,3].
+    assert len(outs) == 2
+    assert outs[0].output.shape[0] == 2 and outs[1].output.shape[0] == 2
+    assert float(outs[0].output[0, 0, 0, 0]) == 0.0 and float(outs[0].output[1, 0, 0, 0]) == 1.0
+    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
+
+    assert BooguImagePipeline.supports_request_batch is True

vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py

@@ -52,7 +52,7 @@ from vllm_omni.diffusion.models.boogu_image.scheduling_flow_match_euler_discrete
 from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 
 logger = init_logger(__name__)
@@ -102,6 +102,20 @@ 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.
+
+    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.
+    """
+    if not has_reference:
+        return ("boogu_image", "t2i")
+    return ("boogu_image", "ti2i", request_id)
+
+
 def get_boogu_image_pre_process_func(od_config: OmniDiffusionConfig):
     """Build the pre-process callable for Boogu-Image reference (edit) input.
 
@@ -132,13 +146,15 @@ def get_boogu_image_pre_process_func(od_config: OmniDiffusionConfig):
     def pre_process_func(request: OmniDiffusionRequest):
         prompt = request.prompt
         if isinstance(prompt, str):
-            # Plain-text prompt cannot carry an image -> text-to-image, no-op.
+            # Plain-text prompt cannot carry an image -> text-to-image.
+            request.batch_compatibility_key = _boogu_batch_compatibility_key(False, request.request_id)
             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, no-op (Base checkpoint).
+            # No reference image -> text-to-image (Base checkpoint).
+            request.batch_compatibility_key = _boogu_batch_compatibility_key(False, request.request_id)
             return request
 
         if isinstance(raw_image, list):
@@ -173,6 +189,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)
         return request
 
     return pre_process_func
@@ -200,7 +217,7 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
     transformer's reference-image refiner path.
     """
 
-    supports_request_batch = False
+    supports_request_batch = True
 
     support_image_input: ClassVar[bool] = True
     color_format: ClassVar[str] = "RGB"
@@ -424,8 +441,11 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
             embeds = embeds.repeat(1, num_images_per_prompt, 1)
             reshaped_embeds = embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
 
-        mask = mask.repeat(num_images_per_prompt, 1)
-        reshaped_mask = mask.view(batch_size * num_images_per_prompt, -1)
+        # repeat_interleave (not repeat/tile) so mask rows stay request-major
+        # [p0, p0, p1, p1] to match the reshaped embeds above; a plain repeat
+        # tiles to [p0, p1, p0, p1] and mismatches when batch_size > 1 and
+        # num_images_per_prompt > 1.
+        reshaped_mask = mask.repeat_interleave(num_images_per_prompt, dim=0)
 
         return batch_size, seq_len, reshaped_embeds, reshaped_mask
 
@@ -676,7 +696,7 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
         preprocessed_images: list[torch.Tensor | None],
         num_images_per_prompt: int,
         device: torch.device,
-        generator=None,
+        generators: list[torch.Generator | None] | None = None,
     ) -> list[list[torch.Tensor] | None]:
         """VAE-encode per-sample reference images into the transformer's format.
 
@@ -685,17 +705,18 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
         ``None`` (no reference / text-to-image) or a list of ``[C, H, W]``
         reference latents (one per reference image). Boogu editing uses a single
         reference image, so each non-empty entry is a one-element list.
-        """
-        # ``latent_dist.sample`` accepts only a single generator; a per-output
-        # generator list (num_outputs_per_prompt > 1) falls back to unseeded.
-        vae_generator = generator if isinstance(generator, torch.Generator) else None
 
+        ``generators`` holds one generator (or ``None``) per request, so each
+        reference latent is sampled with its own request's seed and batched
+        edits stay reproducible against the single-request path.
+        """
         ref_latents: list[list[torch.Tensor] | None] = []
-        for image in preprocessed_images:
+        for idx, image in enumerate(preprocessed_images):
             if image is None:
                 sample_latents: list[torch.Tensor] | None = None
             else:
-                latent = self._encode_vae_image(image.to(device=device), generator=vae_generator).squeeze(0)
+                generator = generators[idx] if generators is not None else None
+                latent = self._encode_vae_image(image.to(device=device), generator=generator).squeeze(0)
                 sample_latents = [latent]
             for _ in range(num_images_per_prompt):
                 ref_latents.append(sample_latents)
@@ -723,7 +744,7 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
             preprocessed_images.append(ai.get("preprocessed_image"))
         return prompt_images, preprocessed_images
 
-    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
         # Prompt / negative-prompt extraction (mirrors the Ovis pattern; the
         # online API sometimes passes ``{"negative_prompt": None}``).
         prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
@@ -737,7 +758,19 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
         has_reference = any(img is not None for img in preprocessed_images)
         task_type = "ti2i" if has_reference else "t2i"
 
-        sp = req.sampling_params
+        # 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,
+        # reference generators) are read per request below.
+        sp = sampling_params_list[0]
         device = self._execution_device
 
         height = sp.height or self.default_sample_size * self.vae_scale_factor
@@ -752,7 +785,15 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
         if not has_reference:
             image_guidance_scale = 1.0
         num_images_per_prompt = sp.num_outputs_per_prompt if sp.num_outputs_per_prompt > 0 else 1
-        generator = sp.generator
+        # Per-request noise generators, collated into one flat list of length
+        # batch_size * num_images_per_prompt (request-major, output-minor).
+        generator = req.collate_request_generators(num_images_per_prompt, None)
+        # One generator per request for reference-latent sampling; a per-output
+        # generator list is not usable by ``latent_dist.sample`` and falls back
+        # to unseeded, matching the single-request path.
+        ref_generators = [
+            s.generator if isinstance(s.generator, torch.Generator) else None for s in sampling_params_list
+        ]
         max_sequence_length = sp.max_sequence_length or 1280
         output_type = sp.output_type or "pil"
         cfg_range = (0.0, 1.0)
@@ -793,7 +834,7 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
         latent_channels = self.transformer.in_channels
         ref_latents = None
         if has_reference:
-            ref_latents = self._build_ref_latents(preprocessed_images, num_images_per_prompt, device, generator)
+            ref_latents = self._build_ref_latents(preprocessed_images, num_images_per_prompt, device, ref_generators)
 
         latents = self.prepare_latents(
             batch_size * num_images_per_prompt,
@@ -935,4 +976,8 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
             if (ori_height, ori_width) != (height, width):
                 image = F.interpolate(image, size=(ori_height, ori_width), mode="bilinear")
 
-        return DiffusionOutput(output=image)
+        return split_diffusion_output_by_request(
+            DiffusionOutput(output=image),
+            req,
+            num_outputs_per_prompt=num_images_per_prompt,
+        )