anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Diffusion] Enable request-level batching for Boogu-Image (TI2I) (#7227)

Browse Source 
Signed-off-by: Shenglei Fu <sfu@confluent.io>
Signed-off-by: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Shenglei FuandClaude authored and GitHub committed 2026-09-26 22:53:35 +08:00
1 parent e75278dccc
commit eb67caaaa9
7 files changed
+935 -143
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py
+8 -19
View File
@@ -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,