[BugFix]: Fix bagel single-stage img2img fallback to text2img bug (#2397)

Signed-off-by: princepride <wangzhipeng628@gmail.com>

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]