From 6834d7f5d720adcbd603beec1e10fe7812e8ac9c Mon Sep 17 00:00:00 2001 From: princepride Date: Wed, 1 Apr 2026 03:24:24 +0000 Subject: [PATCH] fix bagel single-stage img2img bug 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 c4155a9fc8b..aa4f0a74f02 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -387,7 +387,12 @@ def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: 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]