From 0e903e9695a68817b9817885e834cd7c02148aae Mon Sep 17 00:00:00 2001 From: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com> Date: Fri, 4 Sep 2026 01:27:51 -0700 Subject: [PATCH] [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) --- .../boogu_image/test_pipeline_boogu_image.py | 288 +++++++++++++++++- .../boogu_image/pipeline_boogu_image.py | 81 +++-- 2 files changed, 344 insertions(+), 25 deletions(-) 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 fc2b5d362..fc1e0298e 100644 --- a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py +++ b/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 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 8b6fc3a45..cc03b75df 100644 --- a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py +++ b/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, + )