From bf5bd0a4c00feed487f5f5e70810de84fe3d4604 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B1=AA=E5=BF=97=E9=B9=8F?= Date: Wed, 1 Apr 2026 17:29:02 +0800 Subject: [PATCH] [BugFix]: Fix bagel single-stage img2img fallback to text2img bug (#2397) Signed-off-by: princepride --- vllm_omni/diffusion/models/bagel/pipeline_bagel.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index c4155a9fc..aa4f0a74f 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -387,7 +387,12 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin): else: image_input = ( - None if isinstance(first_prompt, str) else (first_prompt.get("multi_modal_data") or {}).get("image") + None + if isinstance(first_prompt, str) + else ( + (first_prompt.get("multi_modal_data") or {}).get("image") + or (first_prompt.get("multi_modal_data") or {}).get("img2img") + ) ) if image_input and not isinstance(image_input, list): image_input = [image_input]