diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py index 3ba69c2e2..62836f963 100644 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ b/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 diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 020b483f0..14bdc10be 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/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]