From 06329c64506b4a1d12a75252f98fc8bbdbbef695 Mon Sep 17 00:00:00 2001
From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com>
Date: Fri, 3 Jul 2026 11:04:17 +0800
Subject: [PATCH] [Bugfix] Fix request batching regression on Ming-flash-omni
 ImageGen (#4837)

Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
---
 .../models/ming_flash_omni/pipeline_ming_imagegen.py | 12 +++---------
 vllm_omni/engine/orchestrator.py                     |  2 +-
 2 files changed, 4 insertions(+), 10 deletions(-)

diff --git a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py
index a55a79987..f3e0f3357 100644
--- a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py
+++ b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py
@@ -446,7 +446,7 @@ class MingImagePipeline(ZImagePipeline):
             None if ref_latent is None else tuple(ref_latent.shape),
         )
         try:
-            outputs = super().forward(
+            outputs: DiffusionOutput = super().forward(
                 z_req,
                 prompt=None,
                 height=height,
@@ -462,13 +462,7 @@ class MingImagePipeline(ZImagePipeline):
         finally:
             set_forward_context_ref_latent(None)
 
-        output = outputs[0]
-        if hasattr(output, "output") and output.output is not None:
-            raw = output.output
-        elif hasattr(output, "images"):
-            raw = output.images
-        else:
-            raw = output
+        raw = outputs.output
         if not isinstance(raw, torch.Tensor):
             raise RuntimeError(f"ZImagePipeline returned non-tensor output: {type(raw).__name__}")
         if logger.isEnabledFor(logging.DEBUG):
@@ -478,7 +472,7 @@ class MingImagePipeline(ZImagePipeline):
                 raw.float().min().item(),
                 raw.float().max().item(),
             )
-        return DiffusionOutput(output=raw)
+        return outputs
 
 
 # ----------------------------------------------------------------------
diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py
index df05cb669..bdd678f27 100644
--- a/vllm_omni/engine/orchestrator.py
+++ b/vllm_omni/engine/orchestrator.py
@@ -1287,7 +1287,7 @@ class Orchestrator:
                             [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
                         )
                         return
-                    if already_submitted and len(diffusion_prompt) == 1:
+                    if len(diffusion_prompt) == 1:
                         diffusion_prompt = diffusion_prompt[0]
             else:
                 diffusion_prompt = req_state.prompt