[Feat] Enable streaming CoT display for HunyuanImage-3.0 AR generation (#4148)

Signed-off-by: zengchuang <zengchuang3@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

examples/offline_inference/hunyuan_image3/end2end.py

@@ -74,6 +74,12 @@ def parse_args():
     parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
     parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
     parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
+    parser.add_argument(
+        "--stream",
+        action="store_true",
+        default=False,
+        help="Enable streaming CoT display; print AR text token-by-token in real time.",
+    )
     parser.add_argument("--log-stats", action="store_true", default=False)
     parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.")
     parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
@@ -222,6 +228,8 @@ def main():
             prompt_dict["modalities"] = ["text"]
         formatted_prompts.append(prompt_dict)
 
+    from vllm.sampling_params import RequestOutputKind
+
     params_list = list(omni.default_sampling_params_list)
 
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -253,6 +261,10 @@ def main():
                 sp.width = args.width
         elif hasattr(sp, "stop_token_ids"):
             sp.stop_token_ids = ar_stop_token_ids
+            # When --stream is set, request DELTA output from the AR stage
+            # so we can display CoT text token-by-token in real time.
+            if args.stream and hasattr(sp, "output_kind"):
+                sp.output_kind = RequestOutputKind.DELTA
 
     print(f"\n{'=' * 60}")
     print("HunyuanImage-3.0 Generation Configuration:")
@@ -281,22 +293,32 @@ def main():
     print(f"  Prompts: {prompts}")
     print(f"{'=' * 60}\n")
 
-    omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
+    # When --stream is set, print AR CoT text token-by-token in real time.
+    # Otherwise, collect and print the full AR text once when stage 0 finishes.
+    omni_outputs = omni.generate(
+        prompts=formatted_prompts,
+        sampling_params_list=params_list,
+        py_generator=True,
+        use_tqdm=False,
+    )
     img_idx = 0
     for req_output in omni_outputs:
         ro = getattr(req_output, "request_output", None)
-        txt = ""
-        if ro and getattr(ro, "outputs", None):
-            txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
-        if not txt:
-            ar_text = getattr(req_output, "custom_output", {}).get("ar_generated_text")
-            if isinstance(ar_text, list):
-                txt = "\n".join(text for text in ar_text if text)
-            else:
-                txt = ar_text or ""
-        if txt:
-            print(f"[Output] Text:\n{txt}")
+        stage_id = getattr(req_output, "stage_id", None)
 
+        # AR stage text — each CompletionOutput.text is already a delta when
+        # output_kind=DELTA, so we can print it directly (matching the pattern
+        # in serving_chat.py).
+        if stage_id == 0 and ro and getattr(ro, "outputs", None):
+            for o in ro.outputs:
+                text = getattr(o, "text", "") or ""
+                if text:
+                    print(text, end="", flush=True)
+            # Non-streaming: one shot with full text — emit a trailing newline.
+            if not args.stream:
+                print(flush=True)
+
+        # Collect images from diffusion stage
         images = getattr(req_output, "images", None)
         if not images and ro and hasattr(ro, "images"):
             images = ro.images
@@ -304,7 +326,7 @@ def main():
             for j, img in enumerate(images):
                 save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
                 img.save(save_path)
-                print(f"[Output] Saved image to {save_path}")
+                print(f"\n[Output] Saved image to {save_path}")
             img_idx += 1
 
 

vllm_omni/entrypoints/omni.py

@@ -25,16 +25,25 @@ logger = init_logger(__name__)
 class Omni(OmniBase):
     """Synchronous entrypoint for offline generation."""
 
-    def _set_final_only_for_llm_stages(
+    def _maybe_force_final_only_for_llm_stages(
         self,
         sampling_params_list: Sequence[OmniSamplingParams],
     ) -> list[OmniSamplingParams]:
-        """Return per-stage params with LLM stages forced to FINAL_ONLY."""
+        """Return per-stage params with LLM stages forced to FINAL_ONLY.
+
+        The caller may explicitly request ``output_kind = DELTA`` on a stage to
+        opt into streaming; such stages are left alone.  All other LLM stages
+        are forced to FINAL_ONLY.
+        """
         effective_params: list[OmniSamplingParams] = []
         for stage_id, params in enumerate(sampling_params_list):
             sp = copy.deepcopy(params)
             stage_meta = self.engine.get_stage_metadata(stage_id)
-            if stage_meta.stage_type != "diffusion" and hasattr(sp, "output_kind"):
+            if (
+                stage_meta.stage_type != "diffusion"
+                and hasattr(sp, "output_kind")
+                and sp.output_kind != RequestOutputKind.DELTA
+            ):
                 sp.output_kind = RequestOutputKind.FINAL_ONLY
             effective_params.append(sp)
         return effective_params
@@ -103,7 +112,7 @@ class Omni(OmniBase):
         use_tqdm: bool | Callable[..., tqdm] = True,
     ) -> Generator[OmniRequestOutput, None, None]:
         try:
-            sampling_params_list = self._set_final_only_for_llm_stages(sampling_params_list)
+            sampling_params_list = self._maybe_force_final_only_for_llm_stages(sampling_params_list)
 
             if isinstance(prompts, str) or not isinstance(prompts, Sequence):
                 request_prompts: list[OmniPromptType] = [prompts]