[Bugfix] Support both list and generator mode & fix unnecessary error log (#599)

Signed-off-by: Gaohan123 <hgaoaf@connect.ust.hk>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/user_guide/examples/offline_inference/qwen2_5_omni.md

@@ -12,7 +12,7 @@ Get into the example folder
 ```bash
 cd examples/offline_inference/qwen2_5_omni
 ```
-Then run the command below.
+Then run the command below. Note: for processing large volume data, it uses py_generator mode, which will return a python generator from Omni class.
 ```bash
 bash run_multiple_prompts.sh
 ```

docs/user_guide/examples/offline_inference/qwen3_omni.md

@@ -12,7 +12,7 @@ Get into the example folder
 ```bash
 cd examples/offline_inference/qwen3_omni
 ```
-Then run the command below.
+Then run the command below. Note: for processing large volume data, it uses py_generator mode, which will return a python generator from Omni class.
 ```bash
 bash run_multiple_prompts.sh
 ```
@@ -92,3 +92,7 @@ sudo apt install ffmpeg
     ``````sh
     --8<-- "examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh"
     ``````
+??? abstract "text_prompts_10.txt"
+    ``````txt
+    --8<-- "examples/offline_inference/qwen3_omni/text_prompts_10.txt"
+    ``````

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -19,7 +19,7 @@ if __name__ == "__main__":
     omni = Omni(model="Qwen/Qwen-Image")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = list(outputs)[0].request_output[0].images
+    images = outputs[0].request_output[0].images
     images[0].save("coffee.png")
 ```
 

examples/offline_inference/image_to_image/image_edit.py

@@ -322,81 +322,76 @@ def main():
     print(f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}")
     print(f"{'=' * 60}\n")
 
-    try:
-        generation_start = time.perf_counter()
-        # Generate edited image
-        generate_kwargs = {
-            "prompt": args.prompt,
-            "pil_image": input_image,
-            "negative_prompt": args.negative_prompt,
-            "generator": generator,
-            "true_cfg_scale": args.cfg_scale,
-            "guidance_scale": args.guidance_scale,
-            "num_inference_steps": args.num_inference_steps,
-            "num_outputs_per_prompt": args.num_outputs_per_prompt,
-            "layers": args.layers,
-            "resolution": args.resolution,
-        }
+    generation_start = time.perf_counter()
+    # Generate edited image
+    generate_kwargs = {
+        "prompt": args.prompt,
+        "pil_image": input_image,
+        "negative_prompt": args.negative_prompt,
+        "generator": generator,
+        "true_cfg_scale": args.cfg_scale,
+        "guidance_scale": args.guidance_scale,
+        "num_inference_steps": args.num_inference_steps,
+        "num_outputs_per_prompt": args.num_outputs_per_prompt,
+        "layers": args.layers,
+        "resolution": args.resolution,
+    }
 
-        outputs = omni.generate(**generate_kwargs)
-        generation_end = time.perf_counter()
-        generation_time = generation_end - generation_start
+    outputs = omni.generate(**generate_kwargs)
+    generation_end = time.perf_counter()
+    generation_time = generation_end - generation_start
 
-        # Print profiling results
-        print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
+    # Print profiling results
+    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
-        # omni.generate() returns Generator[OmniRequestOutput, None, None], convert to list
-        outputs = list(outputs)
-        if not outputs:
-            raise ValueError("No output generated from omni.generate()")
-        logger.info("Outputs: %s", outputs)
+    if not outputs:
+        raise ValueError("No output generated from omni.generate()")
+    logger.info("Outputs: %s", outputs)
 
-        # Extract images from OmniRequestOutput
-        # omni.generate() returns Generator[OmniRequestOutput, None, None], extract images from request_output[0].images
-        first_output = outputs[0]
-        if not hasattr(first_output, "request_output") or not first_output.request_output:
-            raise ValueError("No request_output found in OmniRequestOutput")
+    # Extract images from OmniRequestOutput
+    # omni.generate() returns list[OmniRequestOutput], extract images from request_output[0].images
+    first_output = outputs[0]
+    if not hasattr(first_output, "request_output") or not first_output.request_output:
+        raise ValueError("No request_output found in OmniRequestOutput")
 
-        req_out = first_output.request_output[0]
-        if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
-            raise ValueError("Invalid request_output structure or missing 'images' key")
+    req_out = first_output.request_output[0]
+    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
+        raise ValueError("Invalid request_output structure or missing 'images' key")
 
-        images = req_out.images
-        if not images:
-            raise ValueError("No images found in request_output")
+    images = req_out.images
+    if not images:
+        raise ValueError("No images found in request_output")
 
-        # Save output image(s)
-        output_path = Path(args.output)
-        output_path.parent.mkdir(parents=True, exist_ok=True)
-        suffix = output_path.suffix or ".png"
-        stem = output_path.stem or "output_image_edit"
+    # Save output image(s)
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    suffix = output_path.suffix or ".png"
+    stem = output_path.stem or "output_image_edit"
 
-        # Handle layered output (each image may be a list of layers)
-        if args.num_outputs_per_prompt <= 1:
-            img = images[0]
+    # Handle layered output (each image may be a list of layers)
+    if args.num_outputs_per_prompt <= 1:
+        img = images[0]
+        # Check if this is a layered output (list of images)
+        if isinstance(img, list):
+            for sub_idx, sub_img in enumerate(img):
+                save_path = output_path.parent / f"{stem}_{sub_idx}{suffix}"
+                sub_img.save(save_path)
+                print(f"Saved edited image to {os.path.abspath(save_path)}")
+        else:
+            img.save(output_path)
+            print(f"Saved edited image to {os.path.abspath(output_path)}")
+    else:
+        for idx, img in enumerate(images):
             # Check if this is a layered output (list of images)
             if isinstance(img, list):
                 for sub_idx, sub_img in enumerate(img):
-                    save_path = output_path.parent / f"{stem}_{sub_idx}{suffix}"
+                    save_path = output_path.parent / f"{stem}_{idx}_{sub_idx}{suffix}"
                     sub_img.save(save_path)
                     print(f"Saved edited image to {os.path.abspath(save_path)}")
             else:
-                img.save(output_path)
-                print(f"Saved edited image to {os.path.abspath(output_path)}")
-        else:
-            for idx, img in enumerate(images):
-                # Check if this is a layered output (list of images)
-                if isinstance(img, list):
-                    for sub_idx, sub_img in enumerate(img):
-                        save_path = output_path.parent / f"{stem}_{idx}_{sub_idx}{suffix}"
-                        sub_img.save(save_path)
-                        print(f"Saved edited image to {os.path.abspath(save_path)}")
-                else:
-                    save_path = output_path.parent / f"{stem}_{idx}{suffix}"
-                    img.save(save_path)
-                    print(f"Saved edited image to {os.path.abspath(save_path)}")
-    finally:
-        omni.close()
+                save_path = output_path.parent / f"{stem}_{idx}{suffix}"
+                img.save(save_path)
+                print(f"Saved edited image to {os.path.abspath(save_path)}")
 
 
 if __name__ == "__main__":

examples/offline_inference/image_to_video/image_to_video.py

@@ -118,7 +118,6 @@ def main():
     )
 
     # Extract video frames from OmniRequestOutput
-    frames = list(frames)
     if isinstance(frames, list) and len(frames) > 0:
         first_item = frames[0]
 
@@ -151,31 +150,6 @@ def main():
     except ImportError:
         raise ImportError("diffusers is required for export_to_video.")
 
-    # Extract video frames from OmniRequestOutput
-    if isinstance(frames, list) and len(frames) > 0:
-        first_item = frames[0]
-
-        # Check if it's an OmniRequestOutput
-        if hasattr(first_item, "final_output_type"):
-            if first_item.final_output_type != "image":
-                raise ValueError(
-                    f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
-                )
-
-            # Pipeline mode: extract from nested request_output
-            if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
-                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
-                    inner_output = first_item.request_output[0]
-                    if isinstance(inner_output, dict) and "images" in inner_output:
-                        frames = inner_output["images"][0] if inner_output["images"] else None
-                        if frames is None:
-                            raise ValueError("No video frames found in output.")
-            # Diffusion mode: use direct images field
-            elif hasattr(first_item, "images") and first_item.images:
-                frames = first_item.images
-            else:
-                raise ValueError("No video frames found in OmniRequestOutput.")
-
     # frames may be np.ndarray (preferred) or torch.Tensor
     # export_to_video expects a list of frames with values in [0, 1]
     if isinstance(frames, torch.Tensor):

examples/offline_inference/qwen2_5_omni/README.md

@@ -10,7 +10,7 @@ Get into the example folder
 ```bash
 cd examples/offline_inference/qwen2_5_omni
 ```
-Then run the command below.
+Then run the command below. Note: for processing large volume data, it uses py_generator mode, which will return a python generator from Omni class.
 ```bash
 bash run_multiple_prompts.sh
 ```

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -376,7 +376,7 @@ def main(args):
         for i, prompt in enumerate(prompts):
             prompt["modalities"] = output_modalities
 
-    omni_generator = omni_llm.generate(prompts, sampling_params_list)
+    omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
     # Determine output directory: prefer --output-dir; fallback to --output-wav
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
@@ -514,6 +514,12 @@ def parse_args():
         default=None,
         help="Modalities to use for the prompts.",
     )
+    parser.add_argument(
+        "--py-generator",
+        action="store_true",
+        default=False,
+        help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
+    )
     return parser.parse_args()
 
 

examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh

@@ -1,3 +1,4 @@
 python end2end.py --output-wav output_audio \
                   --query-type text \
-                  --txt-prompts ../qwen3_omni/text_prompts_10.txt
+                  --txt-prompts ../qwen3_omni/text_prompts_10.txt \
+                  --py-generator

examples/offline_inference/qwen3_omni/README.md

@@ -10,7 +10,7 @@ Get into the example folder
 ```bash
 cd examples/offline_inference/qwen3_omni
 ```
-Then run the command below.
+Then run the command below. Note: for processing large volume data, it uses py_generator mode, which will return a python generator from Omni class.
 ```bash
 bash run_multiple_prompts.sh
 ```

examples/offline_inference/qwen3_omni/end2end.py

@@ -319,7 +319,7 @@ def main(args):
         for i, prompt in enumerate(prompts):
             prompt["modalities"] = output_modalities
 
-    omni_generator = omni_llm.generate(prompts, sampling_params_list)
+    omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
     # Determine output directory: prefer --output-dir; fallback to --output-wav
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
     os.makedirs(output_dir, exist_ok=True)
@@ -470,6 +470,12 @@ def parse_args():
         default=None,
         help="Output modalities to use for the prompts.",
     )
+    parser.add_argument(
+        "--py-generator",
+        action="store_true",
+        default=False,
+        help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
+    )
 
     return parser.parse_args()
 

examples/offline_inference/qwen3_omni/run_multiple_prompts.sh

@@ -1,3 +1,4 @@
 python end2end.py --output-wav output_audio \
                   --query-type text \
-                  --txt-prompts text_prompts_10.txt
+                  --txt-prompts text_prompts_10.txt \
+                  --py-generator

examples/offline_inference/text_to_image/README.md

@@ -16,7 +16,7 @@ if __name__ == "__main__":
     omni = Omni(model="Qwen/Qwen-Image")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = list(outputs)[0].request_output[0].images
+    images = outputs[0].request_output[0].images
     images[0].save("coffee.png")
 ```
 

examples/offline_inference/text_to_image/text_to_image.py

@@ -162,9 +162,8 @@ def main():
     print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
     # Extract images from OmniRequestOutput
-    # omni.generate() returns Generator[OmniRequestOutput, None, None], convert to list
-    outputs = list(outputs)
-    if not outputs:
+    # omni.generate() returns list[OmniRequestOutput], extract images from the first output
+    if not outputs or len(outputs) == 0:
         raise ValueError("No output generated from omni.generate()")
     logger.info(f"Outputs: {outputs}")
 

examples/offline_inference/text_to_video/text_to_video.py

@@ -54,7 +54,6 @@ def main():
         flow_shift=args.flow_shift,
     )
 
-    # omni.generate() returns Generator[OmniRequestOutput, None, None]
     frames = omni.generate(
         args.prompt,
         negative_prompt=args.negative_prompt,
@@ -68,7 +67,6 @@ def main():
     )
 
     # Extract video frames from OmniRequestOutput
-    frames = list(frames)
     if isinstance(frames, list) and len(frames) > 0:
         first_item = frames[0]
 
@@ -100,31 +98,6 @@ def main():
     except ImportError:
         raise ImportError("diffusers is required for export_to_video.")
 
-    # Extract video frames from OmniRequestOutput
-    if isinstance(frames, list) and len(frames) > 0:
-        first_item = frames[0]
-
-        # Check if it's an OmniRequestOutput
-        if hasattr(first_item, "final_output_type"):
-            if first_item.final_output_type != "image":
-                raise ValueError(
-                    f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
-                )
-
-            # Pipeline mode: extract from nested request_output
-            if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
-                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
-                    inner_output = first_item.request_output[0]
-                    if isinstance(inner_output, dict) and "images" in inner_output:
-                        frames = inner_output["images"][0] if inner_output["images"] else None
-                        if frames is None:
-                            raise ValueError("No video frames found in output.")
-            # Diffusion mode: use direct images field
-            elif hasattr(first_item, "images") and first_item.images:
-                frames = first_item.images
-            else:
-                raise ValueError("No video frames found in OmniRequestOutput.")
-
     # frames may be np.ndarray (preferred) or torch.Tensor
     # export_to_video expects a list of frames with values in [0, 1]
     if isinstance(frames, torch.Tensor):

tests/e2e/offline_inference/conftest.py

@@ -4,7 +4,6 @@
 Pytest configuration and fixtures for vllm-omni tests.
 """
 
-from collections.abc import Generator
 from typing import Any
 
 import pytest
@@ -12,6 +11,7 @@ from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 from vllm.sampling_params import SamplingParams
 
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.outputs import OmniRequestOutput
 
 PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
 PromptImageInput = list[Any] | Any | None
@@ -190,7 +190,7 @@ class OmniRunner:
         self,
         prompts: list[dict[str, Any]],
         sampling_params_list: list[SamplingParams] | None = None,
-    ) -> Generator[Any, None, None]:
+    ) -> list[OmniRequestOutput]:
         """
         Generate outputs for the given prompts.
 
@@ -206,7 +206,7 @@ class OmniRunner:
         if sampling_params_list is None:
             sampling_params_list = self.get_default_sampling_params_list()
 
-        yield from self.omni.generate(prompts, sampling_params_list)
+        return self.omni.generate(prompts, sampling_params_list)
 
     def generate_multimodal(
         self,
@@ -218,7 +218,7 @@ class OmniRunner:
         videos: PromptVideoInput = None,
         mm_processor_kwargs: dict[str, Any] | None = None,
         modalities: list[str] | None = None,
-    ) -> Generator[Any, None, None]:
+    ) -> list[OmniRequestOutput]:
         """
         Convenience method to generate with multimodal inputs.
 
@@ -243,7 +243,7 @@ class OmniRunner:
             mm_processor_kwargs=mm_processor_kwargs,
             modalities=modalities,
         )
-        yield from self.generate(omni_inputs, sampling_params_list)
+        return self.generate(omni_inputs, sampling_params_list)
 
     def generate_audio(
         self,
@@ -252,7 +252,7 @@ class OmniRunner:
         system_prompt: str | None = None,
         audios: PromptAudioInput = None,
         mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> Generator[Any, None, None]:
+    ) -> list[OmniRequestOutput]:
         """
         Convenience method to generate with multimodal inputs.
         Args:
@@ -270,7 +270,7 @@ class OmniRunner:
             audios=audios,
             mm_processor_kwargs=mm_processor_kwargs,
         )
-        yield from self.generate(omni_inputs, sampling_params_list)
+        return self.generate(omni_inputs, sampling_params_list)
 
     def generate_video(
         self,
@@ -279,7 +279,7 @@ class OmniRunner:
         system_prompt: str | None = None,
         videos: PromptVideoInput = None,
         mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> Generator[Any, None, None]:
+    ) -> list[OmniRequestOutput]:
         """
         Convenience method to generate with multimodal inputs.
         Args:
@@ -297,7 +297,7 @@ class OmniRunner:
             videos=videos,
             mm_processor_kwargs=mm_processor_kwargs,
         )
-        yield from self.generate(omni_inputs, sampling_params_list)
+        return self.generate(omni_inputs, sampling_params_list)
 
     def generate_image(
         self,
@@ -306,7 +306,7 @@ class OmniRunner:
         system_prompt: str | None = None,
         images: PromptImageInput = None,
         mm_processor_kwargs: dict[str, Any] | None = None,
-    ) -> Generator[Any, None, None]:
+    ) -> list[OmniRequestOutput]:
         """
         Convenience method to generate with multimodal inputs.
         Args:
@@ -324,7 +324,7 @@ class OmniRunner:
             images=images,
             mm_processor_kwargs=mm_processor_kwargs,
         )
-        yield from self.generate(omni_inputs, sampling_params_list)
+        return self.generate(omni_inputs, sampling_params_list)
 
     def __enter__(self):
         """Context manager entry."""

tests/e2e/offline_inference/test_cache_dit.py

@@ -62,7 +62,7 @@ def test_cache_dit(model_name: str):
         num_outputs_per_prompt=1,  # Single output for speed
     )
     # Extract images from request_output[0]['images']
-    first_output = list(outputs)[0]
+    first_output = outputs[0]
     assert first_output.final_output_type == "image"
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -133,7 +133,7 @@ def test_sequence_parallel(
             generator=torch.Generator(get_device_name()).manual_seed(seed),
             num_outputs_per_prompt=1,
         )
-        baseline_images = list(outputs)[0].request_output[0].images
+        baseline_images = outputs[0].request_output[0].images
     finally:
         baseline.close()
         if dist.is_initialized():
@@ -165,7 +165,7 @@ def test_sequence_parallel(
             generator=torch.Generator(get_device_name()).manual_seed(seed),
             num_outputs_per_prompt=1,
         )
-        sp_images = list(outputs)[0].request_output[0].images
+        sp_images = outputs[0].request_output[0].images
     finally:
         sp.close()
         if dist.is_initialized():

tests/e2e/offline_inference/test_t2i_model.py

@@ -48,7 +48,7 @@ def test_diffusion_model(model_name: str):
         num_outputs_per_prompt=2,
     )
     # Extract images from request_output[0]['images']
-    first_output = list(outputs)[0]
+    first_output = outputs[0]
     assert first_output.final_output_type == "image"
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")

tests/e2e/offline_inference/test_teacache.py

@@ -58,7 +58,7 @@ def test_teacache(model_name: str):
         num_outputs_per_prompt=1,  # Single output for speed
     )
     # Extract images from request_output[0]['images']
-    first_output = list(outputs)[0]
+    first_output = outputs[0]
     assert first_output.final_output_type == "image"
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")

tests/entrypoints/test_omni_llm.py

@@ -553,7 +553,7 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
 
     omni = Omni(model="any", init_timeout=1)
     with pytest.raises(ValueError):
-        list(omni.generate(prompts=["hi"], sampling_params_list=[]))
+        omni.generate(prompts=["hi"], sampling_params_list=[])
 
 
 def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
@@ -640,7 +640,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
     # Use dicts instead of object() for serializable sampling params
     sampling_params_list = [{"temperature": 0.7}, {"temperature": 0.8}]
     prompts = ["hi"]
-    outputs = list(omni.generate(prompts=prompts, sampling_params_list=sampling_params_list))
+    outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
 
     # Both stages have final_output=True, so should aggregate two OmniRequestOutput
     assert len(outputs) == 2
@@ -723,7 +723,7 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
     )
 
     # Use dicts instead of object() for serializable sampling params
-    outputs = list(omni.generate(prompts=["p"], sampling_params_list=[{"temperature": 0.7}, {"temperature": 0.8}]))
+    outputs = omni.generate(prompts=["p"], sampling_params_list=[{"temperature": 0.7}, {"temperature": 0.8}])
     assert outputs == []
 
 
@@ -796,7 +796,7 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
         }
     )
     # Use the default sampling params
-    list(omni.generate(prompts=["p"], sampling_params_list=None))
+    omni.generate(prompts=["p"], sampling_params_list=None)
 
 
 def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
@@ -924,7 +924,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
     # Generate should handle error gracefully (log but continue)
     # Use dict instead of object() for serializable sampling params
     sampling_params_list = [{"temperature": 0.7}]
-    outputs = list(omni.generate(prompts=["hi"], sampling_params_list=sampling_params_list))
+    outputs = omni.generate(prompts=["hi"], sampling_params_list=sampling_params_list)
     # Should return final output (error was logged but didn't stop processing)
     assert isinstance(outputs, list)
     # Since final_output=True, should have one output

vllm_omni/entrypoints/omni.py

@@ -389,7 +389,9 @@ class Omni(OmniBase):
             self._ray_pg,
         )
 
-    def generate(self, *args: Any, **kwargs: dict[str, Any]) -> Generator[OmniRequestOutput, None, None]:
+    def generate(
+        self, *args: Any, **kwargs: dict[str, Any]
+    ) -> Generator[OmniRequestOutput, None, None] | list[OmniRequestOutput]:
         """Generate outputs for the given prompts.
 
         Orchestrates the multi-stage pipeline based on YAML configuration.
@@ -413,6 +415,7 @@ class Omni(OmniBase):
         """
         prompts = args[0] if args else kwargs.get("prompts")
         sampling_params_list = args[1] if len(args) > 1 else kwargs.get("sampling_params_list")
+        py_generator = kwargs.get("py_generator", False)
         if prompts is None:
             if kwargs.get("prompt") is None:
                 raise ValueError("prompts is required for generation")
@@ -440,11 +443,32 @@ class Omni(OmniBase):
 
             sampling_params_list = per_stage_params
         try:
-            yield from self._run_generation(prompts, sampling_params_list)
+            if py_generator:
+                return self._run_generation_with_generator(prompts, sampling_params_list)
+            else:
+                outputs = list(self._run_generation(prompts, sampling_params_list))
+                self.close()
+                return outputs
+        except Exception as e:
+            logger.exception("[Orchestrator] Failed to run generation: %s", e)
+            # Always close on exception to ensure cleanup
+            self.close()
+            raise e
+
+    def _run_generation_with_generator(
+        self,
+        prompts: PromptType | Sequence[PromptType] | OmniDiffusionRequest | Sequence[OmniDiffusionRequest],
+        sampling_params_list: Any | Sequence[Any] | None,
+    ) -> Generator[OmniRequestOutput, None, None]:
+        """Run generation through all stages in the pipeline and return a generator."""
+        gen = self._run_generation(prompts, sampling_params_list)
+        try:
+            yield from gen
         except Exception as e:
             logger.exception("[Orchestrator] Failed to run generation: %s", e)
             raise e
         finally:
+            # Cleanup when generator is exhausted or closed
             self.close()
 
     def _run_generation(

vllm_omni/entrypoints/omni_stage.py

@@ -633,7 +633,7 @@ def _stage_worker(
         _recv_dequeue_ts = _time.time()
         task_type = task.get("type", OmniStageTaskType.GENERATE)
         if task_type == OmniStageTaskType.SHUTDOWN:
-            logger.error("Received shutdown signal")
+            logger.info("Received shutdown signal")
             break
 
         batch_tasks: list[dict[str, Any]] = [task]