Merge branch 'main' into feat/moss-tts-nano

Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-04-21 17:50:21 -04:00
commit 3bd77d3f8f
@@ -173,8 +173,6 @@ Example configuration for TP=2 on GPUs 0 and 1:

 | Parameter             | Value   | Description                      |
 | :-------------------- | :------ | :------------------------------- |
-| `window_size`         | `-1`    | Window size (-1 means unlimited) |
-| `max_inflight`        | `1`     | Maximum inflight requests        |
 | `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB)   |

 ## Using Mooncake Connector
@@ -58,7 +58,14 @@ Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
    - Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features.
    - Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform.

-Stage wiring is configured in `vllm_omni/model_executor/stage_configs/cosyvoice3.yaml`.
+Pipeline topology lives in `vllm_omni/model_executor/models/cosyvoice3/pipeline.py`;
+runtime tunables (batch size, memory limits, sampling) live in
+`vllm_omni/deploy/cosyvoice3.yaml`. The deploy config auto-loads by
+HF `model_type` and defaults to `async_chunk: true` (shared-memory
+streaming). Pass `--no-async-chunk` on `vllm serve` to switch to the
+legacy sync path where stage 1 runs `text2flow` over the full
+speech-token sequence.

 - Stage 0 emits latent speech tokens.
- Stage 1 consumes them via `custom_process_input_func` and outputs audio.
+- Stage 1 consumes them via `sync_process_input_func` (sync mode) or the
+  shared-memory connector (async-chunk mode) and outputs audio.
@@ -24,7 +24,13 @@ def run_e2e():
        required=True,
        help="Path to CosyVoice3 model directory (e.g., pretrained_models/Fun-CosyVoice3-0.5B/).",
    )
-    parser.add_argument("--stage-config", type=str, default="vllm_omni/model_executor/stage_configs/cosyvoice3.yaml")
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        default=None,
+        help="Override the deploy config path. If unset, auto-loads "
+        "vllm_omni/deploy/cosyvoice3.yaml based on the HF model_type.",
+    )
    parser.add_argument("--prompt", type=str, default="Hello, this is a test of the CosyVoice system capability.")
    parser.add_argument(
        "--prompt-text",
@@ -43,18 +49,14 @@ def run_e2e():
    if not os.path.exists(args.tokenizer):
        raise FileNotFoundError(f"{args.tokenizer} does not exist!")

-    # Ensure stage config exists
-    if not os.path.exists(args.stage_config):
-        raise FileNotFoundError(f"{args.stage_config} does not exist!")
+    if args.deploy_config is not None and not os.path.exists(args.deploy_config):
+        raise FileNotFoundError(f"{args.deploy_config} does not exist!")

    print(f"Initializing cosyvoice E2E with model={args.model}")

-    # Initialize Omni
-    # This spins up the engine(s) based on the stage config
-    # We pass trust_remote_code=True same as Qwen examples
    omni = Omni(
        model=args.model,
-        stage_configs_path=args.stage_config,
+        deploy_config=args.deploy_config,
        trust_remote_code=True,
        tokenizer=args.tokenizer,
        log_stats=True,
@@ -57,22 +57,22 @@ GLM_IMAGE_EOS_TOKEN_ID = 16385  # eos_token_id from generation_config.json
 GLM_IMAGE_VISION_VOCAB_SIZE = 16512  # top_k should be vision_vocab_size


-def compute_max_tokens(height: int, width: int, factor: int = 32) -> int:
+def compute_max_tokens(height: int, width: int, factor: int = 32, is_i2i: bool = False) -> int:
    """
    Compute max_new_tokens for GLM-Image AR generation.

-    GLM-Image generates tokens in this order for text-to-image:
-    1. Small preview image (half resolution in each dimension)
-    2. Large target image (full resolution)
-    3. EOS token
+    GLM-Image generation differs by mode:
+    - text-to-image (t2i): small preview + large target + EOS
+    - image-to-image (i2i): large target + EOS

    Args:
        height: Target image height in pixels
        width: Target image width in pixels
        factor: Downsampling factor (32 for GLM-Image AR output)
+        is_i2i: Whether the request is image-to-image mode

    Returns:
-        Total number of tokens to generate (small + large + EOS)
+        Total number of tokens to generate for the specified mode
    """
    # Large image tokens (target resolution)
    token_h = height // factor
@@ -80,11 +80,15 @@ def compute_max_tokens(height: int, width: int, factor: int = 32) -> int:
    large_tokens = token_h * token_w

    # Small preview tokens (half resolution in each dimension)
-    small_h = token_h // 2
-    small_w = token_w // 2
-    small_tokens = small_h * small_w
+    import math

-    # Total: small + large + EOS
+    ratio = token_h / token_w if token_w > 0 else 1.0
+    small_token_h = max(1, int(math.sqrt(ratio) * (factor // 2)))
+    small_token_w = max(1, int(math.sqrt(1 / ratio) * (factor // 2)))
+    small_tokens = small_token_h * small_token_w
+
+    if is_i2i:
+        return large_tokens + 1
    return small_tokens + large_tokens + 1


@@ -282,14 +286,18 @@ def main(args: argparse.Namespace) -> None:
    # Compute max_tokens dynamically based on target image size
    target_height = prompt_dict.get("height", 1024)
    target_width = prompt_dict.get("width", 1024)
-    calculated_max_tokens = compute_max_tokens(target_height, target_width)
+    is_i2i = source_image is not None
+    calculated_max_tokens = compute_max_tokens(target_height, target_width, is_i2i=is_i2i)

    # Use calculated value unless user explicitly specified a different value
    # Default args.max_tokens is 16384 (very large), so prefer calculated value
    effective_max_tokens = calculated_max_tokens if args.max_tokens == 16384 else args.max_tokens

    if args.verbose:
-        print(f"AR max_tokens: {effective_max_tokens} (calculated: {calculated_max_tokens}, arg: {args.max_tokens})")
+        print(
+            f"AR max_tokens: {effective_max_tokens} "
+            f"(calculated: {calculated_max_tokens}, arg: {args.max_tokens}, mode: {'i2i' if is_i2i else 't2i'})"
+        )

    # IMPORTANT: GLM-Image AR model requires these exact sampling parameters
    # from generation_config.json for proper image token generation.
@@ -303,6 +311,12 @@ def main(args: argparse.Namespace) -> None:
        stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID],  # 16385, CRITICAL for stopping
        seed=args.seed,
        detokenize=False,
+        # Keep target size available in runner/model for deterministic M-RoPE
+        # decode grids in t2i (no mm_features available in this path).
+        extra_args={
+            "target_h": int(target_height),
+            "target_w": int(target_width),
+        },
    )

    # For diffusion stage, sampling_params contains diffusion-specific parameters
@@ -1,25 +1,161 @@
-# HunyuanImage-3.0 Image-to-Text Inference
+# HunyuanImage-3.0-Instruct

-This example demonstrates how to run HunyuanImage-3.0 Image-to-Text with the vLLM-Omni.
+## Set up

-## Local CLI Usage
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.

-Download the example image:
+## Run examples
+
+**Note**: These examples work with the default configuration on **8x NVIDIA L40S (48GB)**. For different GPU setups, modify the stage configuration to adjust device allocation and memory utilization.
+
+Get into the hunyuan_image3 folder:

 ```bash
-wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
+cd examples/offline_inference/hunyuan_image3
 ```

-Run example:
+### Modality Control
+
+HunyuanImage-3.0-Instruct supports multiple modality modes. You can control the mode using the `--modality` argument:
+
+#### Text to Image (text2img)
+
+- **Pipeline**: Text → AR (CoT + latent tokens) → DiT (denoise) → VAE Decode → Image
+- **Stages Used**: Stage 0 (AR) + Stage 1 (DiT)
+- **KV Transfer**: AR sends KV cache to DiT for conditioned generation
+- **Default Config**: `hunyuan_image3_t2i.yaml`

 ```bash
-python image_to_text.py \
-  --image cherry_blossom.jpg \
-  --prompt "<|startoftext|>You are an assistant that understands images and outputs text.<img>Describe the content of the picture."
+python end2end.py --model tencent/HunyuanImage-3.0-Instruct \
+                  --modality text2img \
+                  --prompts "A cute cat sitting on a windowsill watching the sunset"
 ```

-Key arguments:
+#### Image to Image (img2img)

- `--model`: Model used. Default is: tencent/HunyuanImage-3.0-Instruct (Optional).
- `--image`: Path to input image (required).
- `--prompt`: Text description used to guide image understanding (required).
+- **Pipeline**: Image + Text → AR (CoT + recaption + latent) → DiT → Edited Image
+- **Stages Used**: Stage 0 (AR) + Stage 1 (DiT)
+- **KV Transfer**: AR sends KV cache to DiT
+- **Default Config**: `hunyuan_image3_it2i.yaml`
+
+```bash
+python end2end.py --model tencent/HunyuanImage-3.0-Instruct \
+                  --modality img2img \
+                  --image-path /path/to/image.png \
+                  --prompts "Make the petals neon pink"
+```
+
+#### Image to Text (img2text)
+
+- **Pipeline**: Image + Question → AR → Text description
+- **Stages Used**: Stage 0 (AR) only
+- **Default Config**: `hunyuan_image3_i2t.yaml`
+
+```bash
+python end2end.py --model tencent/HunyuanImage-3.0-Instruct \
+                  --modality img2text \
+                  --image-path /path/to/image.jpg \
+                  --prompts "Describe the content of the picture."
+```
+
+#### Text to Text (text2text)
+
+- **Pipeline**: Text → AR → Text
+- **Stages Used**: Stage 0 (AR) only
+- **Default Config**: `hunyuan_image3_t2t.yaml`
+
+```bash
+python end2end.py --model tencent/HunyuanImage-3.0-Instruct \
+                  --modality text2text \
+                  --prompts "What is the capital of France?"
+```
+
+### Inference Steps & Guidance
+
+Control generation quality for image modalities:
+
+```bash
+python end2end.py --modality text2img \
+                  --steps 50 \
+                  --guidance-scale 5.0 \
+                  --height 1024 --width 1024 \
+                  --prompts "A photo-realistic sunset over the ocean"
+```
+
+### Key Arguments
+
+#### 📌 Command Line Arguments (end2end.py)
+
+| Argument               | Type   | Default                              | Description                                                  |
+| :--------------------- | :----- | :----------------------------------- | :----------------------------------------------------------- |
+| `--model`              | string | `tencent/HunyuanImage-3.0-Instruct` | Model path or name                                           |
+| `--modality`           | choice | `text2img`                           | Modality: `text2img`, `img2img`, `img2text`, `text2text`     |
+| `--prompts`            | list   | `None`                               | Input text prompts                                           |
+| `--image-path`         | string | `None`                               | Input image path (for `img2img`/`img2text`)                  |
+| `--output`             | string | `.`                                  | Output directory for saved images                            |
+| `--steps`              | int    | `50`                                 | Number of inference steps                                    |
+| `--guidance-scale`     | float  | `5.0`                                | Classifier-free guidance scale                               |
+| `--seed`               | int    | `42`                                 | Random seed                                                  |
+| `--height`             | int    | `1024`                               | Output image height                                          |
+| `--width`              | int    | `1024`                               | Output image width                                           |
+| `--bot-task`           | string | auto                                 | Override prompt task (e.g. `it2i_think`, `t2i_recaption`)    |
+| `--sys-type`           | string | auto                                 | Override system prompt type (e.g. `en_unified`, `en_vanilla`) |
+| `--stage-configs-path` | string | auto                                 | Custom stage config YAML path                                |
+| `--enforce-eager`      | flag   | `False`                              | Disable torch.compile                                        |
+| `--init-timeout`       | int    | `300`                                | Initialization timeout (seconds)                             |
+
+------
+
+#### ⚙️ Stage Configurations
+
+| Config YAML                         | Modality  | Stages | GPUs   | Description                           |
+| :---------------------------------- | :-------- | :----- | :----- | :------------------------------------ |
+| `hunyuan_image3_t2i.yaml`           | text2img  | 2      | 8      | T2I with AR→DiT, 4 GPU each          |
+| `hunyuan_image3_it2i.yaml`          | img2img   | 2      | 8      | IT2I with AR→DiT, 4 GPU each         |
+| `hunyuan_image3_i2t.yaml`           | img2text  | 1      | 4      | I2T (AR only)                         |
+| `hunyuan_image3_t2t.yaml`           | text2text | 1      | 4      | T2T (AR only)                         |
+| `hunyuan_image3_t2i_2gpu.yaml`      | text2img  | 2      | 2      | T2I for 2-GPU setups                  |
+| `hunyuan_image3_moe.yaml`           | text2img  | 2      | 8      | T2I with MoE AR→DiT KV reuse          |
+| `hunyuan_image3_moe_dit_2gpu_fp8.yaml` | text2img | 2   | 2      | T2I with FP8 quantization             |
+
+------
+
+## Using MoE Config
+
+The `hunyuan_image3_moe.yaml` config enables AR→DiT KV cache reuse with 8 GPUs (4 for AR + 4 for DiT).
+
+```bash
+python end2end.py --model tencent/HunyuanImage-3.0-Instruct \
+                  --modality text2img \
+                  --stage-configs-path hunyuan_image3_moe.yaml \
+                  --prompts "A cute cat"
+```
+
+------
+
+## Prompt Format
+
+HunyuanImage-3.0 uses a pretrain template format:
+
+```
+<|startoftext|>{system_prompt}{<img>}{trigger_tag}{user_prompt}
+```
+
+- `<img>`: Placeholder for each input image (auto-inserted by `prompt_utils.py`)
+- Trigger tags: `<think>` (CoT), `<recaption>` (recaptioning)
+- System prompt: Auto-selected based on task
+
+The `prompt_utils.build_prompt()` handles this formatting automatically.
+
+------
+
+## FAQ
+
+- **OOM errors**: Decrease `gpu_memory_utilization` in the YAML stage config, or use a smaller `max_num_batched_tokens`.
+- **Custom image sizes**: Use `--height` and `--width` flags (multiples of 16 recommended).
+
+| Stage             | VRAM (approx)        |
+| :---------------- | :------------------- |
+| Stage 0 (AR)      | ~15 GiB + KV Cache   |
+| Stage 1 (DiT)     | ~30 GiB              |
+| Total (8-GPU)     | ~45 GiB + KV Cache   |
@@ -0,0 +1,262 @@
+"""
+HunyuanImage-3.0-Instruct unified end-to-end inference script.
+
+Supports all modalities through a single entry point:
+  - text2img:  Text → AR → DiT → Image
+  - img2img:   Text+Image → AR → DiT → Edited Image (IT2I)
+  - img2text:  Image+Text → AR → Text description (I2T)
+  - text2text: Text → AR → Text (comprehension, no image)
+
+Usage:
+    python end2end.py --modality text2img --prompts "A cute cat"
+    python end2end.py --modality img2img --image-path input.png --prompts "Make it snowy"
+    python end2end.py --modality img2text --image-path input.png --prompts "Describe this image"
+"""
+
+import argparse
+import os
+
+from vllm_omni.diffusion.models.hunyuan_image3.system_prompt import (
+    get_system_prompt,
+)
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniPromptType
+
+# task → (sys_type, bot_task, trigger_tag)
+_TASK_PRESETS: dict[str, tuple[str, str | None, str | None]] = {
+    "t2t": ("en_unified", None, None),
+    "i2t": ("en_unified", None, None),
+    "it2i_think": ("en_unified", "think", "<think>"),
+    "it2i_recaption": ("en_unified", "recaption", "<recaption>"),
+    "t2i_think": ("en_unified", "think", "<think>"),
+    "t2i_recaption": ("en_unified", "recaption", "<recaption>"),
+    "t2i_vanilla": ("en_vanilla", "image", None),
+}
+
+# Modality → prompt_utils task mapping
+_MODALITY_TASK_MAP = {
+    "text2img": "t2i_think",
+    "img2img": "it2i_think",
+    "img2text": "i2t",
+    "text2text": "t2t",
+}
+
+
+def build_prompt(
+    user_prompt: str,
+    task: str = "it2i_think",
+    sys_type: str | None = None,
+    custom_system_prompt: str | None = None,
+) -> str:
+    """Build a HunyuanImage-3.0 prompt using pretrain template format."""
+    if task not in _TASK_PRESETS:
+        raise ValueError(f"Unknown task {task!r}. Choose from: {sorted(_TASK_PRESETS)}")
+
+    preset_sys_type, preset_bot_task, trigger_tag = _TASK_PRESETS[task]
+    effective_sys_type = sys_type or preset_sys_type
+
+    system_prompt = get_system_prompt(effective_sys_type, preset_bot_task, custom_system_prompt)
+    sys_text = system_prompt.strip() if system_prompt else ""
+
+    has_image_input = task.startswith("i2t") or task.startswith("it2i")
+
+    parts = ["<|startoftext|>"]
+    if sys_text:
+        parts.append(sys_text)
+    if has_image_input:
+        parts.append("<img>")
+    if trigger_tag:
+        parts.append(trigger_tag)
+    parts.append(user_prompt)
+
+    return "".join(parts)
+
+
+# Modality → default stage config
+_MODALITY_DEFAULT_CONFIG = {
+    "text2img": "hunyuan_image3_t2i.yaml",
+    "img2img": "hunyuan_image3_it2i.yaml",
+    "img2text": "hunyuan_image3_i2t.yaml",
+    "text2text": "hunyuan_image3_t2t.yaml",
+}
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="HunyuanImage-3.0-Instruct end-to-end inference.")
+    parser.add_argument(
+        "--model",
+        default="tencent/HunyuanImage-3.0-Instruct",
+        help="Model name or local path.",
+    )
+    parser.add_argument(
+        "--modality",
+        default="text2img",
+        choices=["text2img", "img2img", "img2text", "text2text"],
+        help="Modality mode to control stage execution.",
+    )
+    parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
+    parser.add_argument(
+        "--image-path",
+        type=str,
+        default=None,
+        help="Path to input image (for img2img/img2text).",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default=".",
+        help="Output directory to save results.",
+    )
+
+    # Generation parameters
+    parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
+    parser.add_argument("--guidance-scale", type=float, default=5.0, help="Classifier-free guidance scale.")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
+    parser.add_argument("--height", type=int, default=1024, help="Output image height.")
+    parser.add_argument("--width", type=int, default=1024, help="Output image width.")
+
+    # Prompt configuration
+    parser.add_argument(
+        "--bot-task",
+        type=str,
+        default=None,
+        help="Override prompt task (e.g. it2i_think, t2i_recaption). Default: auto from modality.",
+    )
+    parser.add_argument(
+        "--sys-type",
+        type=str,
+        default=None,
+        help="Override system prompt type (e.g. en_unified, en_vanilla).",
+    )
+
+    # Omni init args
+    parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom stage config YAML path.")
+    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.")
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    os.makedirs(args.output, exist_ok=True)
+
+    # Determine task for prompt formatting
+    task = args.bot_task or _MODALITY_TASK_MAP[args.modality]
+
+    # Determine stage config
+    stage_configs_path = args.stage_configs_path or _MODALITY_DEFAULT_CONFIG[args.modality]
+
+    # Build Omni
+    omni_kwargs = {
+        "model": args.model,
+        "stage_configs_path": stage_configs_path,
+        "log_stats": args.log_stats,
+        "init_timeout": args.init_timeout,
+        "enforce_eager": args.enforce_eager,
+    }
+    if args.modality in ("text2img", "img2img"):
+        omni_kwargs["mode"] = "text-to-image"
+
+    omni = Omni(**omni_kwargs)
+
+    # Prepare prompts
+    prompts = args.prompts or ["A cute cat"]
+    if not prompts:
+        print("[Info] No prompts provided, using default.")
+        prompts = ["A cute cat"]
+
+    # Load image if needed
+    input_image = None
+    if args.modality in ("img2img", "img2text"):
+        if not args.image_path or not os.path.exists(args.image_path):
+            raise ValueError(f"--image-path required for {args.modality}, got: {args.image_path}")
+        from PIL import Image
+
+        input_image = Image.open(args.image_path).convert("RGB")
+
+    # Format prompts
+    formatted_prompts: list[OmniPromptType] = []
+    for p in prompts:
+        formatted_text = build_prompt(p, task=task, sys_type=args.sys_type)
+
+        prompt_dict: dict = {"prompt": formatted_text}
+
+        if args.modality == "text2img":
+            prompt_dict["modalities"] = ["image"]
+        elif args.modality == "img2img":
+            prompt_dict["modalities"] = ["image"]
+            prompt_dict["multi_modal_data"] = {"image": input_image}
+            prompt_dict["height"] = input_image.height
+            prompt_dict["width"] = input_image.width
+        elif args.modality == "img2text":
+            prompt_dict["modalities"] = ["text"]
+            prompt_dict["multi_modal_data"] = {"image": input_image}
+        elif args.modality == "text2text":
+            prompt_dict["modalities"] = ["text"]
+
+        formatted_prompts.append(prompt_dict)
+
+    # Build sampling params from defaults
+    params_list = list(omni.default_sampling_params_list)
+
+    # Override diffusion params if applicable
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    for i, sp in enumerate(params_list):
+        if isinstance(sp, OmniDiffusionSamplingParams):
+            sp.num_inference_steps = args.steps
+            sp.guidance_scale = args.guidance_scale
+            if args.seed is not None:
+                sp.seed = args.seed
+            if args.modality in ("text2img",):
+                sp.height = args.height
+                sp.width = args.width
+
+    # Print configuration
+    print(f"\n{'=' * 60}")
+    print("HunyuanImage-3.0 Generation Configuration:")
+    print(f"  Model: {args.model}")
+    print(f"  Modality: {args.modality}")
+    print(f"  Stage config: {stage_configs_path}")
+    print(f"  Num stages: {omni.num_stages}")
+    if args.modality in ("text2img", "img2img"):
+        print(f"  Inference steps: {args.steps}")
+        print(f"  Guidance scale: {args.guidance_scale}")
+        print(f"  Seed: {args.seed}")
+    if args.modality == "text2img":
+        print(f"  Output size: {args.width}x{args.height}")
+    if args.image_path:
+        print(f"  Input image: {args.image_path}")
+    print(f"  Prompts: {prompts}")
+    print(f"{'=' * 60}\n")
+
+    # Generate
+    omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
+
+    # Process outputs
+    img_idx = 0
+    for req_output in omni_outputs:
+        # Text output (AR stage or text-only)
+        ro = getattr(req_output, "request_output", None)
+        if ro and getattr(ro, "outputs", None):
+            txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
+            if txt:
+                print(f"[Output] Text:\n{txt}")
+
+        # Image output (DiT stage)
+        images = getattr(req_output, "images", None)
+        if not images and ro and hasattr(ro, "images"):
+            images = ro.images
+
+        if images:
+            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}")
+            img_idx += 1
+
+
+if __name__ == "__main__":
+    main()
@@ -1,92 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import argparse
-import os
-
-from PIL import Image
-
-from vllm_omni.entrypoints.omni import Omni
-
-"""
-The tencent/HunyuanImage-3.0-Instruct base model uses the tencent/Hunyuan-A13B-Instruct backbone. It utilizes two tokenizer delimiter templates:
-
-1) Pretrained template (default for gen_text mode), which concatenates system, image
-   tokens, and user question WITHOUT role delimiters:
-"<|startoftext|>{system_prompt}{image_tokens}{user_question}"
-
-   Example (before image token expansion):
-"<|startoftext|>You are an assistant that understands images and outputs text.<img>Describe the content of the picture."
-
-2) Instruct template, which uses explicit role prefixes and separators.
-"""
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Generate text from image using HunyuanImage-3.0-Instruct.")
-    parser.add_argument(
-        "--model",
-        default="tencent/HunyuanImage-3.0-Instruct",
-        help="Model name or local path.",
-    )
-    parser.add_argument(
-        "--image",
-        type=str,
-        required=True,
-        help="Path to input image file (PNG, JPG, etc.).",
-    )
-    parser.add_argument(
-        "--prompt",
-        type=str,
-        required=True,
-        help="Pretrain template prompt: <|startoftext|>{system}<img>{question}",
-    )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
-    return parser.parse_args()
-
-
-def load_image(image_path: str) -> Image.Image:
-    """Load an image from file path."""
-    if not os.path.exists(image_path):
-        raise FileNotFoundError(f"Image file not found: {image_path}")
-    return Image.open(image_path).convert("RGB")
-
-
-def main(args: argparse.Namespace) -> None:
-    omni = Omni(
-        model=args.model,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
-        mode="image-to-text",
-    )
-
-    prompt = "<|startoftext|>You are an assistant that understands images and outputs text.<img>" + args.prompt
-
-    prompt_dict = {
-        "prompt": prompt,
-        "modalities": ["text"],
-    }
-
-    # Add image input if provided
-    if args.image:
-        if not os.path.exists(args.image):
-            raise FileNotFoundError(f"Input image not found: {args.image}")
-
-        input_image = load_image(args.image)
-        prompt_dict["multi_modal_data"] = {"image": input_image}
-
-    prompts = [prompt_dict]
-    omni_outputs = omni.generate(prompts=prompts)
-
-    prompt_text = omni_outputs[0].request_output.prompt
-    generated_text = omni_outputs[0].request_output.outputs[0].text
-    print(f"Prompt: {prompt_text}")
-    print(f"Text: {generated_text}")
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    main(args)
@@ -39,7 +39,6 @@ Run a single sample for basic TTS:

 ```bash
 python3 -u end2end.py \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
  --query-type tts_sft
 ```
@@ -48,7 +47,6 @@ Run batch samples for basic TTS:

 ```bash
 python3 -u end2end.py \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
  --query-type tts_sft \
  --num-prompts {batch_size}
@@ -66,7 +64,6 @@ Generate speech from text input:

 ```bash
 python3 -u end2end.py \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
  --query-type tts_sft \
  --text "The weather is so nice today."
@@ -78,7 +75,6 @@ Generate speech with explicit voice style instructions:

 ```bash
 python3 -u end2end.py \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
  --query-type tts_sft_with_instruct \
  --text "The weather is so nice today." \
@@ -91,7 +87,6 @@ Generate speech using an audio reference for voice cloning:

 ```bash
 python3 -u end2end.py \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
  --query-type tts_sft_with_audio \
  --text "The weather is so nice today." \
@@ -104,7 +99,6 @@ Generate speech from text containing natural voice descriptions:

 ```bash
 python3 -u end2end.py \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
  --query-type tts_sft_with_natural_instruction \
  --text "In a panting young male voice, he said: I can't run anymore, wait for me!"
@@ -116,7 +110,6 @@ Transcribe audio to text:

 ```bash
 python3 -u end2end.py \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
  --query-type audio_trancribing_sft \
  --audio-path "./spoken_dialogue_assistant_turn_1.wav"
@@ -128,7 +121,6 @@ Understand and analyze audio content with text queries:

 ```bash
 python3 -u end2end.py \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
  --query-type audio_understanding_sft \
  --text "Summarize the audio." \
@@ -141,7 +133,6 @@ Audio understanding with reasoning chain:

 ```bash
 python3 -u end2end.py \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
  --query-type audio_understanding_sft_with_thinking \
  --text "Summarize the audio." \
@@ -154,7 +145,6 @@ Multi-turn dialogue with audio input and output:

 ```bash
 python3 -u end2end.py \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
  --query-type spoken_dialogue_sft_multiturn \
  --audio-path "./prompt_speech_zh_m.wav"
@@ -168,7 +158,6 @@ Multi-turn dialogue converting speech to text:

 ```bash
 python3 -u end2end.py \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
  --query-type speech2text_dialogue_sft_multiturn
 ```
@@ -181,7 +170,6 @@ Multi-turn text-only dialogue:

 ```bash
 python3 -u end2end.py \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
  --query-type text_dialogue_sft_multiturn
 ```
@@ -182,7 +182,7 @@ def main(args):

    omni = Omni(
        model=model_name,
-        stage_configs_path=args.stage_configs_path,
+        deploy_config=args.deploy_config,
        log_stats=args.enable_stats,
        log_file=("omni_pipeline.log" if args.enable_stats else None),
        init_sleep_seconds=args.init_sleep_seconds,
@@ -428,10 +428,11 @@ def parse_args():
        help="Sampling rate for audio.",
    )
    parser.add_argument(
-        "--stage-configs-path",
+        "--deploy-config",
        type=str,
-        default="../../../model_executor/stage_configs/mimo_audio.yaml",
-        help="Path to a stage configs file.",
+        default=None,
+        help="Override the deploy config path. If unset, auto-loads "
+        "vllm_omni/deploy/mimo_audio.yaml based on the HF model_type.",
    )

    return parser.parse_args()
@@ -0,0 +1,76 @@
+# Ming-flash-omni 2.0
+
+[Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with outputs in text, image, and audio. For now, Ming-flash-omni-2.0 in vLLM-Omni is supported with thinker stage (multi-modal understanding).
+
+## Setup
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+
+## Run examples
+
+### Text-only
+```bash
+python examples/offline_inference/ming_flash_omni/end2end.py --query-type text
+```
+
+#### Reasoning (Thinking Mode)
+
+Reasoning (Thinking) mode is enabled via applying "detailed thinking on" when building the system prompt template (in `apply_chat_template`).
+
+In the end2end example, a default problem for thinking mode is provided, as referred to the example usage of Ming's cookbook;
+To utilize it, you have to download the example figure from https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/3_0.png
+
+```bash
+python examples/offline_inference/ming_flash_omni/end2end.py -q reasoning --image-path ./3_0.png
+```
+
+### Image understanding
+```bash
+python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image
+
+# With a local image
+python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image --image-path /path/to/image.jpg
+```
+
+### Audio understanding
+```bash
+python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio
+
+# With a local audio file
+python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio --audio-path /path/to/audio.wav
+```
+
+### Video understanding
+```bash
+python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video
+
+# With a local video and custom frame count
+python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video --video-path /path/to/video.mp4 --num-frames 16
+```
+
+### Mixed modalities (image + audio)
+```bash
+python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_mixed_modalities \
+    --image-path /path/to/image.jpg \
+    --audio-path /path/to/audio.wav
+```
+
+If media file paths are not provided, the script uses built-in default assets.
+
+### Modality control
+To control output modalities (e.g. text-only output):
+```bash
+python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio --modalities text
+```
+
+*For now, only text output is supported*
+
+### Custom stage config
+```bash
+python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image \
+    --stage-configs-path /path/to/your_config.yaml
+```
+
+## Online serving
+
+For online serving via the OpenAI-compatible API, see [examples/online_serving/ming_flash_omni/README.md](../../online_serving/ming_flash_omni/README.md).
@@ -0,0 +1,485 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Partial example cases are referred from
+# https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/cookbook.ipynb
+import os
+import time
+from typing import NamedTuple
+
+import numpy as np
+import vllm
+from PIL import Image
+from transformers import AutoProcessor
+from vllm import SamplingParams
+from vllm.assets.audio import AudioAsset
+from vllm.assets.image import ImageAsset
+from vllm.assets.video import VideoAsset, video_to_ndarrays
+from vllm.multimodal.image import convert_image_mode
+from vllm.multimodal.media.audio import load_audio
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+import vllm_omni
+from vllm_omni.entrypoints.omni import Omni
+
+# Imports the processor also registers itself
+from vllm_omni.transformers_utils.processors.ming import MingFlashOmniProcessor  # noqa: F401
+
+SEED = 42
+MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0"
+
+
+class QueryResult(NamedTuple):
+    inputs: dict
+    limit_mm_per_prompt: dict[str, int]
+
+
+def get_text_query(processor: MingFlashOmniProcessor, question: str | None = None) -> QueryResult:
+    if question is None:
+        question = "请详细介绍鹦鹉的生活习性。"
+    conversation = [{"role": "HUMAN", "content": question}]
+    prompt = processor.apply_chat_template(conversation, tokenize=False)
+    return QueryResult(
+        inputs={"prompt": prompt},
+        limit_mm_per_prompt={},
+    )
+
+
+def get_image_query(
+    processor: MingFlashOmniProcessor,
+    question: str | None = None,
+    image_path: str | None = None,
+) -> QueryResult:
+    if question is None:
+        question = "Describe this image in detail."
+
+    if image_path:
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+        image_data = convert_image_mode(Image.open(image_path), "RGB")
+    else:
+        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
+
+    conversation = [
+        {
+            "role": "HUMAN",
+            "content": [
+                {"type": "image", "image": image_data},
+                {"type": "text", "text": question},
+            ],
+        }
+    ]
+    prompt = processor.apply_chat_template(conversation, tokenize=False)
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {"image": image_data},
+        },
+        limit_mm_per_prompt={"image": 1},
+    )
+
+
+def get_audio_query(
+    processor: MingFlashOmniProcessor,
+    question: str | None = None,
+    audio_path: str | None = None,
+    sampling_rate: int = 16000,
+) -> QueryResult:
+    if question is None:
+        question = "Please recognize the language of this speech and transcribe it. Format: oral."
+
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = load_audio(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+
+    # Use a string for "audio" so the processor counts it as 1 audio input
+    conversation = [
+        {
+            "role": "HUMAN",
+            "content": [
+                {"type": "audio", "audio": "input"},
+                {"type": "text", "text": question},
+            ],
+        }
+    ]
+    prompt = processor.apply_chat_template(conversation, tokenize=False)
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {"audio": audio_data},
+        },
+        limit_mm_per_prompt={"audio": 1},
+    )
+
+
+def get_video_query(
+    processor: MingFlashOmniProcessor,
+    question: str | None = None,
+    video_path: str | None = None,
+    num_frames: int = 16,
+) -> QueryResult:
+    if question is None:
+        question = "Describe what is happening in this video."
+
+    if video_path:
+        if not os.path.exists(video_path):
+            raise FileNotFoundError(f"Video file not found: {video_path}")
+        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
+    else:
+        video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
+
+    conversation = [
+        {
+            "role": "HUMAN",
+            "content": [
+                {"type": "video"},
+                {"type": "text", "text": question},
+            ],
+        }
+    ]
+    prompt = processor.apply_chat_template(conversation, tokenize=False)
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {"video": video_frames},
+        },
+        limit_mm_per_prompt={"video": 1},
+    )
+
+
+def get_mixed_modalities_query(
+    processor: MingFlashOmniProcessor,
+    image_path: str | None = None,
+    audio_path: str | None = None,
+    sampling_rate: int = 16000,
+) -> QueryResult:
+    """Mixed image + audio understanding."""
+    question = "Describe the image, and recognize the language of this speech and transcribe it. Format: oral"
+
+    if image_path:
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+        image_data = convert_image_mode(Image.open(image_path), "RGB")
+    else:
+        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
+
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        sig, sr = load_audio(audio_path, sr=sampling_rate)
+        audio_data = (sig.astype(np.float32), sr)
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+
+    conversation = [
+        {
+            "role": "HUMAN",
+            "content": [
+                {"type": "image", "image": image_data},
+                {"type": "audio", "audio": "input"},
+                {"type": "text", "text": question},
+            ],
+        }
+    ]
+    prompt = processor.apply_chat_template(conversation, tokenize=False)
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {"image": image_data, "audio": audio_data},
+        },
+        limit_mm_per_prompt={"image": 1, "audio": 1},
+    )
+
+
+def get_reasoning_query(
+    processor: MingFlashOmniProcessor,
+    question: str | None = None,
+    image_path: str | None = None,
+) -> QueryResult:
+    if question is None:
+        # NOTE: To use the following default question, input with example figure provided by Ming
+        # https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/3_0.png
+        # E.g.,
+        # python examples/offline_inference/ming_flash_omni/end2end.py -q reasoning --image-path ./3_0.png
+        # Otherwise, the problem solving might be false.
+        question = (
+            "Based on the following rules:\n•\tYou control the smiley face character\n"
+            "•\tYou can move up, down, left, and right, and only a single square at a time\n"
+            "•\tWalls are dark grey and cannot be moved into\n•\tThe brown square is a box\n•"
+            "\tThe box can be pushed by moving into it (i.e., if you are in the square "
+            "adjacent to the box to the left, and move onto the square with the box, "
+            "the box will move one square to the right).\n"
+            "•\tThe box cannot be pushed into walls\n"
+            "•\tThe blue door at the bottom is locked and cannot be passed through, "
+            "unless the box is placed on the blue square\n"
+            "•\tThe square beneath the blue door is the exit\n"
+            "•\tMoving from one square to another\n\n"
+            "Let's assume a coordinate system where the smiley face is "
+            "on the top left at (1,1) and the square below it is (1,2). "
+            "The smiley face performs the following moves: {down, right, right, right}, "
+            "such that the smiley face is at square (4,2) and the box is in square (5,2). "
+            "What are the next sequence of moves that must be done to move the box down to (5,3)? "
+            "Give your answer as a comma separated list."
+        )
+
+    if image_path:
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+        image_data = convert_image_mode(Image.open(image_path), "RGB")
+        conversation = [
+            {
+                "role": "HUMAN",
+                "content": [
+                    {"type": "image", "image": image_data},
+                    {"type": "text", "text": question},
+                ],
+            }
+        ]
+        prompt = processor.apply_chat_template(conversation, tokenize=False, use_cot_system_prompt=True)
+        return QueryResult(
+            inputs={
+                "prompt": prompt,
+                "multi_modal_data": {"image": image_data},
+            },
+            limit_mm_per_prompt={"image": 1},
+        )
+
+    conversation = [{"role": "HUMAN", "content": question}]
+    prompt = processor.apply_chat_template(conversation, tokenize=False, use_cot_system_prompt=True)
+    return QueryResult(
+        inputs={"prompt": prompt},
+        limit_mm_per_prompt={},
+    )
+
+
+query_map = {
+    "text": get_text_query,
+    "use_audio": get_audio_query,
+    "use_image": get_image_query,
+    "use_video": get_video_query,
+    "use_mixed_modalities": get_mixed_modalities_query,
+    "reasoning": get_reasoning_query,
+}
+
+
+def main(args):
+    print(
+        "=" * 20,
+        "\n",
+        f"vllm version: {vllm.__version__}\n",
+        f"vllm-omni version: {vllm_omni.__version__}\n",
+        "=" * 20,
+        sep="",
+    )
+
+    processor = AutoProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    assert isinstance(processor, MingFlashOmniProcessor), f"Wrong processor type being used: {type(processor)}"
+
+    query_func = query_map[args.query_type]
+    if args.query_type == "use_image":
+        query_result = query_func(processor, image_path=args.image_path)
+    elif args.query_type == "use_audio":
+        query_result = query_func(processor, audio_path=args.audio_path, sampling_rate=args.sampling_rate)
+    elif args.query_type == "use_video":
+        query_result = query_func(processor, video_path=args.video_path, num_frames=args.num_frames)
+    elif args.query_type == "use_mixed_modalities":
+        query_result = query_func(
+            processor,
+            image_path=args.image_path,
+            audio_path=args.audio_path,
+            sampling_rate=args.sampling_rate,
+        )
+    elif args.query_type == "reasoning":
+        query_result = query_func(processor, image_path=args.image_path)
+    else:
+        query_result = query_func(processor)
+
+    # Initialize Omni (with thinker-only stage config)
+    omni = Omni(
+        model=MODEL_NAME,
+        stage_configs_path=args.stage_configs_path,
+        log_stats=args.log_stats,
+        init_timeout=args.init_timeout,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    # Thinker sampling params
+    thinker_sampling_params = SamplingParams(
+        temperature=0.4,
+        top_p=0.9,
+        max_tokens=args.max_tokens,
+        repetition_penalty=1.05,
+        seed=SEED,
+        detokenize=True,
+    )
+    sampling_params_list = [thinker_sampling_params]
+
+    prompts = [query_result.inputs for _ in range(args.num_prompts)]
+
+    if args.modalities is not None:
+        output_modalities = args.modalities.split(",")
+        for prompt in prompts:
+            prompt["modalities"] = output_modalities
+
+    total_requests = len(prompts)
+    processed_count = 0
+    print(f"Query type: {args.query_type}")
+    print(f"Number of prompts: {total_requests}")
+
+    output_dir = args.output_dir
+    os.makedirs(output_dir, exist_ok=True)
+
+    profiler_enabled = args.enable_profiler
+    if profiler_enabled:
+        omni.start_profile(stages=args.profiler_stages)
+
+    for stage_outputs in omni.generate(prompts, sampling_params_list):
+        output = stage_outputs.request_output
+        if stage_outputs.final_output_type == "text":
+            request_id = output.request_id
+            text_output = output.outputs[0].text
+            lines = []
+            lines.append("Prompt:\n")
+            lines.append(str(output.prompt) + "\n")
+            lines.append("Text Output:\n")
+            lines.append(str(text_output).strip() + "\n")
+            print(*lines, sep="")
+
+            # Save to file
+            out_txt = os.path.join(output_dir, f"{request_id}.txt")
+            try:
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    f.writelines(lines)
+                print(f"Request ID: {request_id}, text saved to {out_txt}")
+            except Exception as e:
+                print(f"Failed to write output file {out_txt}: {e}")
+
+        elif stage_outputs.final_output_type == "audio":
+            raise NotImplementedError("Add audio example after talker supported.")
+
+        processed_count += 1
+        if profiler_enabled and processed_count >= total_requests:
+            print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
+            # Stop the profiler while workers are still alive
+            omni.stop_profile(stages=args.profiler_stages)
+
+            print("[Info] Waiting 30s for workers to write trace files to disk...")
+            time.sleep(30)
+            print("[Info] Trace export wait time finished.")
+
+    omni.close()
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="text",
+        choices=query_map.keys(),
+        help="Query type.",
+    )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="Path to a stage configs YAML file.",
+    )
+    parser.add_argument(
+        "--log-stats",
+        action="store_true",
+        default=False,
+        help="Enable detailed statistics logging.",
+    )
+    parser.add_argument("--init-timeout", type=int, default=2000, help="Timeout for initializing in seconds.")
+    parser.add_argument(
+        "--stage-init-timeout",
+        type=int,
+        default=2000,
+        help="Timeout for initializing a single stage in seconds.",
+    )
+    parser.add_argument(
+        "--enable-profiler",
+        action="store_true",
+        default=False,
+        help="Enables profiling when set.",
+    )
+    parser.add_argument(
+        "--profiler-stages",
+        type=int,
+        nargs="*",
+        default=[0],
+        help="List of stage IDs to profile. If not set, profiles all stages.",
+    )
+    parser.add_argument(
+        "--image-path",
+        "-i",
+        type=str,
+        default=None,
+        help="Path to local image file. Uses default asset if not provided.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file. Uses default asset if not provided.",
+    )
+    parser.add_argument(
+        "--video-path",
+        "-v",
+        type=str,
+        default=None,
+        help="Path to local video file. Uses default asset if not provided.",
+    )
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=16,
+        help="Number of frames to extract from video.",
+    )
+    parser.add_argument(
+        "--sampling-rate",
+        type=int,
+        default=16000,
+        help="Sampling rate for audio loading.",
+    )
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=16384,
+        help="Maximum tokens to generate.",
+    )
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Number of prompts to generate.",
+    )
+    parser.add_argument(
+        "--modalities",
+        type=str,
+        default=None,
+        help="Output modalities (comma-separated).",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="output_ming",
+        help="Output directory for results.",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)
@@ -320,14 +320,7 @@ def main(args):
        query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
    else:
        query_result = query_func()
-    omni = Omni(
-        model=model_name,
-        log_stats=args.log_stats,
-        stage_init_timeout=args.stage_init_timeout,
-        batch_timeout=args.batch_timeout,
-        init_timeout=args.init_timeout,
-        shm_threshold_bytes=args.shm_threshold_bytes,
-    )
+    omni = Omni.from_cli_args(args, model=model_name)
    thinker_sampling_params = SamplingParams(
        temperature=0.0,  # Deterministic - no randomness
        top_p=1.0,  # Disable nucleus sampling
@@ -70,8 +70,8 @@ For true stage-level concurrency -- where downstream stages (Talker, Code2Wav)
 start **before** the upstream stage (Thinker) finishes -- use the async_chunk
 example. This requires:

-1. A stage config YAML with ``async_chunk: true`` (e.g.
-   ``qwen3_omni_moe_async_chunk.yaml``).
+1. A deploy config YAML with ``async_chunk: true`` (e.g.
+   ``qwen3_omni_moe.yaml``).
 2. Hardware that matches the config (e.g. 2x H100 for the default 3-stage
   config).

@@ -101,7 +101,7 @@ python end2end_async_chunk.py --query-type text --modalities text
 ```bash
 python end2end_async_chunk.py \
    --query-type use_audio \
-    --stage-configs-path /path/to/your_async_chunk.yaml
+    --deploy-config /path/to/your_deploy_config.yaml
 ```

 > **Note**: The synchronous ``end2end.py`` (using ``Omni``) is still the
@@ -294,14 +294,7 @@ def main(args):
    else:
        query_result = query_func()

-    omni = Omni(
-        model=model_name,
-        dtype=args.dtype,
-        stage_configs_path=args.stage_configs_path,
-        log_stats=args.log_stats,
-        stage_init_timeout=args.stage_init_timeout,
-        init_timeout=args.init_timeout,
-    )
+    omni = Omni.from_cli_args(args, model=model_name)

    thinker_sampling_params = SamplingParams(
        temperature=0.9,
@@ -14,7 +14,7 @@ actively processing while stage-0 is still generating.
 Usage
 -----
    python end2end_async_chunk.py --query-type use_audio \
-        --stage-configs-path <path-to-async-chunk-yaml>
+        --deploy-config <path-to-deploy-config-yaml>

 See ``--help`` for all options.
 """
@@ -179,20 +179,26 @@ def clone_prompt_for_request(template: dict) -> dict:
    return cloned


-def _default_async_chunk_stage_configs_path() -> str | None:
-    """Best-effort default stage config for running Qwen3-Omni with async_chunk.
+def _default_deploy_config_path() -> str | None:
+    """Best-effort default deploy config for running Qwen3-Omni with async_chunk.

-    When this example is executed from within the repository, we resolve the
-    default YAML path relative to this file. When installed elsewhere, the
-    file may not exist and callers should pass --stage-configs-path explicitly.
+    The default ``vllm_omni/deploy/qwen3_omni_moe.yaml`` ships with
+    ``async_chunk: true`` at the top level, so loading it is enough to
+    enable async-chunk semantics. To disable it, copy the YAML and set
+    ``async_chunk: false`` (or pass ``--deploy-config`` to a YAML that
+    overrides the flag).
+
+    When this example is executed from within the repository, we resolve
+    the default YAML path relative to this file. When installed elsewhere,
+    the file may not exist and callers should pass ``--deploy-config``
+    explicitly.
    """
    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
    candidate = os.path.join(
        repo_root,
        "vllm_omni",
-        "model_executor",
-        "stage_configs",
-        "qwen3_omni_moe_async_chunk.yaml",
+        "deploy",
+        "qwen3_omni_moe.yaml",
    )
    return candidate if os.path.exists(candidate) else None

@@ -374,15 +380,16 @@ async def run_all(args):
            prompt["modalities"] = output_modalities

    # Create AsyncOmni
-    print(f"[Info] Creating AsyncOmni with stage_configs_path={args.stage_configs_path}")
+    print(f"[Info] Creating AsyncOmni with deploy_config={args.deploy_config}")
    async_omni = None
    try:
-        async_omni = AsyncOmni(
-            model=args.model,
-            stage_configs_path=args.stage_configs_path,
-            log_stats=args.log_stats,
-            stage_init_timeout=args.stage_init_timeout,
-        )
+        # ``from_cli_args`` expands vars(args) into kwargs and auto-captures
+        # ``_cli_explicit_keys`` from ``sys.argv[1:]`` so argparse defaults
+        # do not silently override deploy YAML values. Mirrors the
+        # ``EngineArgs.from_cli_args`` pattern used throughout vllm /
+        # vllm-omni. ``deploy_config=None`` (the default) falls through to
+        # the bundled ``vllm_omni/deploy/qwen3_omni_moe.yaml``.
+        async_omni = AsyncOmni.from_cli_args(args)

        # Use default sampling params from stage config (they are pre-configured
        # in the YAML for each stage).
@@ -470,11 +477,11 @@ def parse_args():
        help="Query type.",
    )
    parser.add_argument(
-        "--stage-configs-path",
+        "--deploy-config",
        type=str,
-        default=_default_async_chunk_stage_configs_path(),
+        default=_default_deploy_config_path(),
        help=(
-            "Path to an async_chunk stage config YAML. "
+            "Path to a deploy config YAML. "
            "If not set, uses the model's default config "
            "(make sure it has async_chunk: true)."
        ),
@@ -17,7 +17,7 @@ REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
 python "${SCRIPT_DIR}/end2end_async_chunk.py" \
    --query-type text \
    --txt-prompts "${SCRIPT_DIR}/text_prompts_10.txt" \
-    --stage-configs-path "${REPO_ROOT}/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml" \
+    --deploy-config "${REPO_ROOT}/vllm_omni/deploy/qwen3_omni_moe.yaml" \
    --output-dir output_audio_async_chunk \
    --max-in-flight 2 \
    "$@"
@@ -6,13 +6,13 @@
 # achieving true stage-level concurrency via chunk-level streaming.
 #
 # Prerequisites:
-#   - An async_chunk stage config YAML (e.g. qwen3_omni_moe_async_chunk.yaml)
+#   - A deploy config YAML (e.g. qwen3_omni_moe.yaml)
 #   - Hardware matching the config (e.g. 2x H100 for the default 3-stage config)
 #
 # Usage:
 #   bash run_single_prompt_async_chunk.sh
 #   bash run_single_prompt_async_chunk.sh --query-type text --modalities text
-#   bash run_single_prompt_async_chunk.sh --stage-configs-path /path/to/custom.yaml
+#   bash run_single_prompt_async_chunk.sh --deploy-config /path/to/custom.yaml

 set -euo pipefail

@@ -21,6 +21,6 @@ REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"

 python "${SCRIPT_DIR}/end2end_async_chunk.py" \
    --query-type use_audio \
-    --stage-configs-path "${REPO_ROOT}/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml" \
+    --deploy-config "${REPO_ROOT}/vllm_omni/deploy/qwen3_omni_moe.yaml" \
    --output-dir output_audio_async_chunk \
    "$@"
@@ -15,11 +15,11 @@ Please refer to the [stage configuration documentation](https://docs.vllm.ai/pro

 ### ROCm Dependencies

-You will need to install these two dependencies `onnxruntime-rocm` and `sox`.
+You will need to install the dependency `onnxruntime-rocm`.

 ```
 pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
-pip install onnxruntime-rocm sox
+pip install onnxruntime-rocm
 ```

 ## Quick Start
@@ -104,13 +104,13 @@ completes. This demonstrates that audio data is available progressively rather t

 ## Batched Decoding

-The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_num_seqs > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
+The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, set `max_num_seqs > 1` on both stages via `--stage-overrides` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.

 ```
 python end2end.py --query-type CustomVoice \
    --txt-prompts benchmark_prompts.txt \
    --batch-size 4 \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
+    --stage-overrides '{"0":{"max_num_seqs":4,"gpu_memory_utilization":0.2},"1":{"max_num_seqs":4,"gpu_memory_utilization":0.2}}'
 ```

 **Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_num_seqs >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_num_seqs`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
@@ -366,12 +366,7 @@ def main(args):
    output_dir = args.output_dir
    os.makedirs(output_dir, exist_ok=True)

-    omni = Omni(
-        model=model_name,
-        stage_configs_path=args.stage_configs_path,
-        log_stats=args.log_stats,
-        stage_init_timeout=args.stage_init_timeout,
-    )
+    omni = Omni.from_cli_args(args, model=model_name)

    batch_size = args.batch_size
    for batch_start in range(0, len(inputs), batch_size):
@@ -387,12 +382,7 @@ async def main_streaming(args):
    output_dir = args.output_dir
    os.makedirs(output_dir, exist_ok=True)

-    omni = AsyncOmni(
-        model=model_name,
-        stage_configs_path=args.stage_configs_path,
-        log_stats=args.log_stats,
-        stage_init_timeout=args.stage_init_timeout,
-    )
+    omni = AsyncOmni.from_cli_args(args, model=model_name)

    for i, prompt in enumerate(inputs):
        request_id = str(i)
@@ -27,6 +27,22 @@ python text_to_audio.py \
  --output stable_audio_output.wav
 ```

+To reduce per-GPU memory for multi-GPU inference, launch with HSDP:
+
+```bash
+python text_to_audio.py \
+  --model stabilityai/stable-audio-open-1.0 \
+  --prompt "The sound of a hammer hitting a wooden surface" \
+  --negative-prompt "Low quality" \
+  --seed 42 \
+  --guidance-scale 7.0 \
+  --audio-length 10.0 \
+  --num-inference-steps 100 \
+  --use-hsdp \
+  --hsdp-shard-size 2 \
+  --output stable_audio_output.wav
+```
+
 Key arguments:

 - `--prompt`: text description (string).
@@ -35,5 +51,8 @@ Key arguments:
 - `--guidance-scale`: classifier-free guidance scale.
 - `--audio-length`: audio duration in seconds.
 - `--num-inference-steps`: diffusion sampling steps.(more steps = higher quality, slower).
+- `--use-hsdp`: enable HSDP weight sharding for the Stable Audio DiT.
+- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
+- `--hsdp-replicate-size`: number of HSDP replica groups.
 - `--cache-backend`: cache acceleration backend. Stable Audio currently supports `tea_cache`.
 - `--output`: path to save the generated WAV file.
@@ -21,6 +21,7 @@ from pathlib import Path
 import numpy as np
 import torch

+from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -113,6 +114,23 @@ def parse_args() -> argparse.Namespace:
        action="store_true",
        help="Enable diffusion pipeline profiler to display stage durations.",
    )
+    parser.add_argument(
+        "--use-hsdp",
+        action="store_true",
+        help="Enable HSDP for Stable Audio DiT weight sharding.",
+    )
+    parser.add_argument(
+        "--hsdp-shard-size",
+        type=int,
+        default=1,
+        help="Number of GPUs to shard Stable Audio DiT weights across when HSDP is enabled.",
+    )
+    parser.add_argument(
+        "--hsdp-replicate-size",
+        type=int,
+        default=1,
+        help="Number of HSDP replica groups. Default 1 means pure sharding.",
+    )
    return parser.parse_args()


@@ -158,12 +176,23 @@ def main():
    print(f"  Inference steps: {args.num_inference_steps}")
    print(f"  Guidance scale: {args.guidance_scale}")
    print(f"  Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}")
+    if args.use_hsdp:
+        print(f"  HSDP: enabled (shard_size={args.hsdp_shard_size}, replicate_size={args.hsdp_replicate_size})")
+    else:
+        print("  HSDP: disabled")
    print(f"  Seed: {args.seed}")
    print(f"{'=' * 60}\n")

+    parallel_config = DiffusionParallelConfig(
+        use_hsdp=args.use_hsdp,
+        hsdp_shard_size=args.hsdp_shard_size,
+        hsdp_replicate_size=args.hsdp_replicate_size,
+    )
+
    # Initialize Omni with Stable Audio model
    omni = Omni(
        model=args.model,
+        parallel_config=parallel_config,
        cache_backend=args.cache_backend,
        cache_config=cache_config,
        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
@@ -29,7 +29,8 @@ This folder provides several entrypoints for experimenting with text-to-image di
 | `AIDC-AI/Ovis-Image-7B` | 1024 x 1024 | 71.8 | 17.1 |
 | `OmniGen2/OmniGen2` |  1024 x 1024 | 20.1 | 14.7 |
 | `stabilityai/stable-diffusion-3.5-medium` | 1024 x 1024 | 20.1 | 15.6 |
-| `black-forest-labs/FLUX.1-dev` | 1024 x 1024 | 77.6 | 31.4 |
+| `black-forest-labs/FLUX.1-dev` | 1024 x 1024 | 33.9 | 31.4 |
+| `black-forest-labs/FLUX.1-schnell` | 1024 x 1024 | 33.9 | 31.4 |
 | `black-forest-labs/FLUX.2-klein-4B` | 1024 x 1024 | 72.7 | 14.9 |
 | `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 |
 | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
@@ -0,0 +1,123 @@
+# VoxCPM Offline Example
+
+This directory contains the minimal offline VoxCPM example for vLLM Omni.
+
+`end2end.py` is intentionally small and only covers:
+
+- single text-to-speech
+- single voice cloning with `ref_audio` + `ref_text`
+- non-streaming with `vllm_omni/model_executor/stage_configs/voxcpm.yaml`
+- streaming with `vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml`
+
+Advanced workflows were moved out of the getting-started example:
+
+- `benchmarks/voxcpm/vllm_omni/bench_tts_offline.py`: warmup, batch prompts, profiler, offline TTFP / RTF
+- `benchmarks/voxcpm/vllm_omni/run_offline_matrix.py`: fixed offline smoke matrix
+- `benchmarks/voxcpm/`: benchmark scripts and benchmark docs
+
+## Prerequisites
+
+Install VoxCPM in one of these ways:
+
+```bash
+pip install voxcpm
+```
+
+or point vLLM Omni to the local VoxCPM source tree:
+
+```bash
+export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
+```
+
+The example writes WAV files with `soundfile`:
+
+```bash
+pip install soundfile
+```
+
+## Model Path
+
+Pass the native VoxCPM model directory directly:
+
+```bash
+export VOXCPM_MODEL=/path/to/voxcpm-model
+```
+
+If the native VoxCPM `config.json` does not contain HuggingFace metadata such as
+`model_type`, prepare a persistent HF-compatible config directory and point the
+stage configs to it with `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH`:
+
+```bash
+export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
+mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
+cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
+cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
+python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
+```
+
+If the model directory itself already has `model_type`, this extra directory is
+not required.
+
+## Quick Start
+
+Single text-to-speech, non-streaming:
+
+```bash
+python examples/offline_inference/voxcpm/end2end.py \
+  --model "$VOXCPM_MODEL" \
+  --text "This is a split-stage VoxCPM synthesis example running on vLLM Omni."
+```
+
+Single voice cloning, non-streaming:
+
+```bash
+python examples/offline_inference/voxcpm/end2end.py \
+  --model "$VOXCPM_MODEL" \
+  --text "This sentence is synthesized with a cloned voice." \
+  --ref-audio /path/to/reference.wav \
+  --ref-text "The exact transcript spoken in reference.wav."
+```
+
+Streaming:
+
+```bash
+python examples/offline_inference/voxcpm/end2end.py \
+  --model "$VOXCPM_MODEL" \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
+  --text "This is a split-stage VoxCPM streaming example running on vLLM Omni."
+```
+
+By default, `end2end.py` writes to `output_audio/` for non-streaming and
+`output_audio_streaming/` for streaming.
+
+## Advanced Workflows
+
+Use `benchmarks/voxcpm/vllm_omni/bench_tts_offline.py` when you need:
+
+- warmup runs
+- prompt files
+- batch JSONL inputs
+- profiler injection
+- offline TTFP / RTF emission
+
+Use `benchmarks/voxcpm/vllm_omni/run_offline_matrix.py` when you need the fixed offline smoke matrix that previously lived in `test.py`.
+
+Full matrix benchmark example:
+
+```bash
+python benchmarks/voxcpm/vllm_omni/run_offline_matrix.py \
+  --model "$VOXCPM_MODEL" \
+  --ref-audio /path/to/reference.wav \
+  --ref-text "The exact transcript spoken in reference.wav."
+```
+
+For online serving examples, see [examples/online_serving/voxcpm](../../online_serving/voxcpm/README.md).
+
+For benchmark reporting, see [benchmarks/voxcpm](../../../benchmarks/voxcpm/README.md).
+
+## Notes
+
+- `voxcpm.yaml` is the default non-streaming stage config.
+- `voxcpm_async_chunk.yaml` is the streaming stage config.
+- Streaming is currently single-request oriented; the fixed smoke matrix now lives in `benchmarks/voxcpm/vllm_omni/run_offline_matrix.py`.
+- `ref_text` must be the real transcript of the reference audio. Mismatched text usually causes obvious quality degradation.
@@ -0,0 +1,206 @@
+"""Minimal offline VoxCPM example for vLLM Omni."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from pathlib import Path
+from typing import Any
+
+import soundfile as sf
+import torch
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni import AsyncOmni, Omni
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+DEFAULT_SYNC_STAGE_CONFIG = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml"
+
+
+def _build_prompt(args) -> dict[str, Any]:
+    additional_information: dict[str, list[Any]] = {
+        "text": [args.text],
+        "cfg_value": [args.cfg_value],
+        "inference_timesteps": [args.inference_timesteps],
+        "min_len": [args.min_len],
+        "max_new_tokens": [args.max_new_tokens],
+    }
+    if args.streaming_prefix_len is not None:
+        additional_information["streaming_prefix_len"] = [args.streaming_prefix_len]
+    if args.ref_audio is not None:
+        additional_information["ref_audio"] = [args.ref_audio]
+    if args.ref_text is not None:
+        additional_information["ref_text"] = [args.ref_text]
+    return {
+        "prompt_token_ids": [1],
+        "additional_information": additional_information,
+    }
+
+
+def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor:
+    audio = mm.get("audio", mm.get("model_outputs"))
+    if audio is None:
+        raise ValueError("No audio output found in multimodal output.")
+    if isinstance(audio, list):
+        parts = [torch.as_tensor(item).float().cpu().reshape(-1) for item in audio]
+        audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0)
+    if not isinstance(audio, torch.Tensor):
+        audio = torch.as_tensor(audio)
+    return audio.float().cpu().reshape(-1)
+
+
+def _extract_sample_rate(mm: dict[str, Any]) -> int:
+    sr_raw = mm.get("sr", 24000)
+    if isinstance(sr_raw, list) and sr_raw:
+        sr_raw = sr_raw[-1]
+    if hasattr(sr_raw, "item"):
+        return int(sr_raw.item())
+    return int(sr_raw)
+
+
+def _is_streaming_stage_config(stage_config_path: str) -> bool:
+    return "async_chunk" in Path(stage_config_path).stem
+
+
+def _save_audio(audio: torch.Tensor, sample_rate: int, output_dir: Path, request_id: str) -> Path:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    output_path = output_dir / f"output_{request_id}.wav"
+    sf.write(
+        output_path,
+        audio.float().cpu().clamp(-1.0, 1.0).numpy(),
+        sample_rate,
+        format="WAV",
+        subtype="PCM_16",
+    )
+    return output_path
+
+
+async def _run_streaming(args) -> Path:
+    prompt = _build_prompt(args)
+    output_dir = Path(args.output_dir) if args.output_dir is not None else Path("output_audio_streaming")
+    request_id = "streaming_example"
+    sample_rate = 24000
+    buffered_samples = 0
+    chunks: list[torch.Tensor] = []
+    started = time.perf_counter()
+    omni = AsyncOmni(
+        model=args.model,
+        stage_configs_path=args.stage_configs_path,
+        log_stats=args.log_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+    try:
+        async for stage_output in omni.generate(prompt, request_id=request_id):
+            mm = getattr(stage_output, "multimodal_output", None)
+            if not isinstance(mm, dict):
+                request_output = getattr(stage_output, "request_output", None)
+                if request_output is None:
+                    continue
+                mm = getattr(request_output, "multimodal_output", None)
+                if not isinstance(mm, dict) and getattr(request_output, "outputs", None):
+                    mm = getattr(request_output.outputs[0], "multimodal_output", None)
+            if not isinstance(mm, dict):
+                continue
+            audio = _extract_audio_tensor(mm)
+            if audio.numel() == 0:
+                continue
+            sample_rate = _extract_sample_rate(mm)
+            if audio.numel() > buffered_samples:
+                delta = audio[buffered_samples:]
+                buffered_samples = int(audio.numel())
+            else:
+                delta = audio
+                buffered_samples += int(delta.numel())
+            if delta.numel() > 0:
+                chunks.append(delta)
+        if not chunks:
+            raise RuntimeError("No streaming audio chunks received from VoxCPM.")
+        output_audio = torch.cat(chunks, dim=0)
+        output_path = _save_audio(output_audio, sample_rate, output_dir, request_id)
+        print(f"Saved streaming audio to: {output_path} ({time.perf_counter() - started:.2f}s)")
+        return output_path
+    finally:
+        omni.shutdown()
+
+
+def _run_sync(args) -> Path:
+    prompt = _build_prompt(args)
+    output_dir = Path(args.output_dir) if args.output_dir is not None else Path("output_audio")
+    request_id = "sync_example"
+    started = time.perf_counter()
+    last_mm: dict[str, Any] | None = None
+    omni = Omni(
+        model=args.model,
+        stage_configs_path=args.stage_configs_path,
+        log_stats=args.log_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+    for stage_outputs in omni.generate(prompt):
+        request_output = getattr(stage_outputs, "request_output", None)
+        if request_output is None:
+            continue
+        outputs = getattr(request_output, "outputs", None)
+        if outputs:
+            for output in outputs:
+                mm = getattr(output, "multimodal_output", None)
+                if isinstance(mm, dict):
+                    last_mm = mm
+        mm = getattr(request_output, "multimodal_output", None)
+        if isinstance(mm, dict):
+            last_mm = mm
+    if last_mm is None:
+        raise RuntimeError("No audio output received from VoxCPM.")
+    output_path = _save_audio(
+        _extract_audio_tensor(last_mm),
+        _extract_sample_rate(last_mm),
+        output_dir,
+        request_id,
+    )
+    print(f"Saved audio to: {output_path} ({time.perf_counter() - started:.2f}s)")
+    return output_path
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="Minimal offline VoxCPM example for vLLM Omni.")
+    parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.")
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=str(DEFAULT_SYNC_STAGE_CONFIG),
+        help=("Stage config path. Use voxcpm.yaml for non-streaming or voxcpm_async_chunk.yaml for streaming."),
+    )
+    parser.add_argument("--text", type=str, required=True, help="Input text for synthesis.")
+    parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.")
+    parser.add_argument("--ref-text", type=str, default=None, help="Transcript of the reference audio.")
+    parser.add_argument("--output-dir", type=str, default=None, help="Output directory for generated wav files.")
+    parser.add_argument("--cfg-value", type=float, default=2.0, help="Guidance value passed to VoxCPM.")
+    parser.add_argument("--inference-timesteps", type=int, default=10, help="Number of diffusion timesteps.")
+    parser.add_argument("--min-len", type=int, default=2, help="Minimum latent length.")
+    parser.add_argument("--max-new-tokens", type=int, default=4096, help="Maximum latent length.")
+    parser.add_argument(
+        "--streaming-prefix-len",
+        type=int,
+        default=3,
+        help="Streaming prefix length used by voxcpm_async_chunk.yaml.",
+    )
+    parser.add_argument("--stage-init-timeout", type=int, default=600, help="Stage initialization timeout in seconds.")
+    parser.add_argument("--log-stats", action="store_true", help="Enable vLLM Omni stats logging.")
+    args = parser.parse_args()
+    if (args.ref_audio is None) != (args.ref_text is None):
+        raise ValueError("Voice cloning requires --ref-audio and --ref-text together.")
+    return args
+
+
+def main(args) -> None:
+    route = "streaming" if _is_streaming_stage_config(args.stage_configs_path) else "sync"
+    print(f"Model: {args.model}")
+    print(f"Stage config: {args.stage_configs_path}")
+    print(f"Route: {route}")
+    if route == "streaming":
+        asyncio.run(_run_streaming(args))
+    else:
+        _run_sync(args)
+
+
+if __name__ == "__main__":
+    main(parse_args())
@@ -65,6 +65,12 @@ def parse_args():
        default=None,
        help="Text matching --prompt-audio for continuation mode.",
    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Optional transcript of --reference-audio (enables ref_continuation mode).",
+    )
    return parser.parse_args()


@@ -74,16 +80,20 @@ def extract_audio(multimodal_output: dict) -> torch.Tensor:
    The output processor concatenates per-step delta tensors under
    ``model_outputs``.  Falls back to ``audio`` for backwards compat.
    """
-    audio = multimodal_output.get("model_outputs") or multimodal_output.get("audio")
+    audio = multimodal_output.get("model_outputs")
+    if audio is None:
+        audio = multimodal_output.get("audio")
    if audio is None:
        raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}")

    if isinstance(audio, list):
-        # Take the last valid tensor (most complete audio)
+        # Defensive: usually the output processor consolidates into a single
+        # tensor at request completion, but concatenate here too in case the
+        # caller consumes intermediate (pre-consolidation) outputs.
        valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
        if not valid:
            raise ValueError("Audio list is empty or all elements are None.")
-        return valid[-1]
+        return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]

    return torch.as_tensor(audio).float().cpu().reshape(-1)

@@ -99,24 +109,40 @@ def main():
        stage_configs_path=args.stage_configs_path,
    )

-    additional: dict = {}
-    if args.reference_audio:
-        additional["reference_audio"] = args.reference_audio
-    if args.prompt_audio and args.prompt_text:
-        additional["prompt_audio"] = args.prompt_audio
-        additional["prompt_text"] = args.prompt_text
+    from transformers import AutoTokenizer

-    prompt: dict = {"prompt": args.text}
-    if additional:
-        prompt["additional_information"] = additional
+    from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
+        build_cjk_split_map,
+        build_voxcpm2_prompt,
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
+    split_map = build_cjk_split_map(tokenizer)
+    hf_config = engine.engine.stage_vllm_configs[0].model_config.hf_config
+
+    ref_audio_arg = args.reference_audio or args.prompt_audio
+    ref_text_arg = args.ref_text or args.prompt_text
+    ref_wav, ref_sr = (None, None)
+    if ref_audio_arg:
+        ref_wav_arr, ref_sr = sf.read(ref_audio_arg)
+        ref_wav = ref_wav_arr.mean(axis=-1).tolist() if ref_wav_arr.ndim > 1 else ref_wav_arr.tolist()
+
+    prompt = build_voxcpm2_prompt(
+        hf_config=hf_config,
+        tokenizer=tokenizer,
+        split_map=split_map,
+        text=args.text,
+        ref_audio=ref_wav,
+        ref_sr=ref_sr,
+        ref_text=ref_text_arg,
+    )

    print(f"Model       : {args.model}")
    print(f"Text        : {args.text}")
-    if args.reference_audio:
-        print(f"Ref audio   : {args.reference_audio}")
-    if args.prompt_audio:
-        print(f"Prompt audio: {args.prompt_audio}")
-        print(f"Prompt text : {args.prompt_text}")
+    if ref_audio_arg:
+        print(f"Ref audio   : {ref_audio_arg}")
+    if ref_text_arg:
+        print(f"Ref text    : {ref_text_arg}")
    print(f"Output dir  : {output_dir}")

    t_start = time.perf_counter()
@@ -10,28 +10,24 @@ When `mistral_common` has `SpeechRequest` support, prompt token IDs are built vi
 ```bash
 # Basic single-prompt with cheerful_female voice preset
 python3 examples/offline_inference/voxtral_tts/end2end.py \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
    --write-audio --voice cheerful_female \
    --model mistralai/Voxtral-4B-TTS-2603 \
    --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"

 # 32 replicate prompts with cheerful_female voice preset
 python3 examples/offline_inference/voxtral_tts/end2end.py \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
    --num-prompts 32 --write-audio --voice cheerful_female \
    --model mistralai/Voxtral-4B-TTS-2603 \
    --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"

 # Streaming with neutral_female voice preset
 python3 examples/offline_inference/voxtral_tts/end2end.py \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
    --streaming --write-audio --voice neutral_female \
    --model mistralai/Voxtral-4B-TTS-2603 \
    --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"

 # 32 prompts, 8 concurrent requests per wave, streaming with neutral_female voice
 python3 examples/offline_inference/voxtral_tts/end2end.py \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
    --num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \
    --model mistralai/Voxtral-4B-TTS-2603 \
    --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
@@ -39,7 +35,6 @@ python3 examples/offline_inference/voxtral_tts/end2end.py \
 # Short debug prompt with reference audio
 # Note: Reference audio capability is not yet released.
 python3 examples/offline_inference/voxtral_tts/end2end.py \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
    --write-audio \
    --model mistralai/Voxtral-4B-TTS-2603 \
    --text "This is a test message." \
@@ -54,7 +49,7 @@ python3 examples/offline_inference/voxtral_tts/end2end.py \
 | `--text TEXT` | Text to synthesize (default: `"This is a test message."`) |
 | `--audio-path PATH` | Path to reference audio file for voice cloning |
 | `--output-dir DIR` | Directory to write output WAV files (default: `output_audio`) |
-| `--stage-configs-path PATH` | Path to stage configs YAML (currently it must be set for VoxtralTTS) |
+| `--deploy-config PATH` | Override the deploy config path. If unset, auto-loads `vllm_omni/deploy/voxtral_tts.yaml` from the HF `model_type`. |
 | `--num-prompts N` | Number of replicate prompts to run for measuring performance (default: 1) |
 | `--streaming` | Use streaming generation via `AsyncOmni` (default: blocking `Omni`) |
 | `--concurrency N` | Max concurrent requests per wave (must be used with `--streaming`, must evenly divide `--num-prompts`) |
@@ -39,7 +39,7 @@ logger = logging.getLogger(__name__)
 async def run_streaming(inputs, sampling_params_list, model_name, args, output_dir):
    async_omni = AsyncOmni(
        model=model_name,
-        stage_configs_path=args.stage_configs_path,
+        deploy_config=args.deploy_config,
        log_stats=args.log_stats,
    )

@@ -192,7 +192,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
    llm = Omni(
        model=model_name,
        log_stats=args.log_stats,
-        stage_configs_path=args.stage_configs_path,
+        deploy_config=args.deploy_config,
    )

    if args.profiling_mode:
@@ -253,10 +253,11 @@ def parse_args() -> Namespace:
        help="Directory to write output wav files.",
    )
    parser.add_argument(
-        "--stage-configs-path",
+        "--deploy-config",
        type=str,
        default=None,
-        help="Path to stage configs YAML. Auto-resolved from model if not set.",
+        help="Override the deploy config path. If unset, auto-loads "
+        "vllm_omni/deploy/voxtral_tts.yaml based on the HF model_type.",
    )
    parser.add_argument(
        "--num-prompts", type=int, default=1, help="Number of replicate prompts to run for measuring performance"
@@ -30,9 +30,9 @@ dreamid_omni/
 ```

 ### Run the Inference
-```
+```python
 python x_to_video_audio.py \
-  --model /xx/dreamid_omni \
+  --model /path/to/dreamid_omni \
  --prompt "Two people walking together and singing happily" \
  --image-path ./example0.png ./example1.png \
  --audio-path ./example0.wav ./example1.wav \
@@ -42,11 +42,33 @@ python x_to_video_audio.py \
  --num-inference-steps 45 \
  --height 704 \
  --width 1280 \
-  --output dreamid_omni.mp4
+  --output out_dreamid_omni_twoip.mp4
 ```
 In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
 The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.

+
+You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
+
+For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing.
+
+```python
+# Example usage for oneip, ref media from the official repo DreamID-Omni
+python x_to_video_audio.py \
+  --model /path/to/dreamid_omni \
+  --prompt "<img1>: In the frame, a woman with black long hair is identified as <sub1>.\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: <sub1> is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: <sub1> tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n<sub1> maintains eye contact, swallows as if choosing her words, and says, <S>I keep telling myself I’m fine,but some nights it feels like I’m just performing calm.<E>" \
+  --image-path 9.png \
+  --audio-path 9.wav \
+  --video-negative-prompt "jitter, bad hands, blur, distortion" \
+  --audio-negative-prompt "robotic, muffled, echo, distorted" \
+  --cfg-parallel-size 2 \
+  --num-inference-steps 45 \
+  --height 704 \
+  --width 1280 \
+  --output out_dreamid_omni_oneip.mp4
+```
+
+
 Key arguments:
 - `--prompt`: text description (string).
 - `--model`: path to the model local directory.
@@ -5,10 +5,12 @@ import argparse
 import re
 import time

+import numpy as np
 from PIL import Image
 from vllm.multimodal.media.audio import load_audio

 from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams

@@ -56,6 +58,11 @@ def parse_args() -> argparse.Namespace:
        default=False,
        help="Enable CPU offloading for diffusion models.",
    )
+    parser.add_argument(
+        "--enable-layerwise-offload",
+        action="store_true",
+        help="Enable layerwise (blockwise) offloading on DiT modules.",
+    )
    return parser.parse_args()


@@ -124,6 +131,7 @@ def main() -> None:
        parallel_config=parallel_config,
        model_type=args.model_type,
        enable_cpu_offload=args.enable_cpu_offload,
+        enable_layerwise_offload=args.enable_layerwise_offload,
    )
    start = time.perf_counter()
    outputs = omni.generate(prompt, sampling_params)
@@ -131,15 +139,35 @@ def main() -> None:

    if not outputs:
        raise RuntimeError("No output returned from DreamID-Omni.")
-    output = outputs[0].request_output
-    generated_video = output.images[0][0]
-    generated_audio = output.images[0][1]
-    try:
-        from dreamid_omni.utils.io_utils import save_video
-    except Exception as e:
-        raise RuntimeError(f"Failed to extract video and audio from DreamID-Omni output. Error: {e}")
+    result = outputs[0]
+    if not result.images:
+        raise RuntimeError("No video frames found in DreamID-Omni output.")
+    generated_video = result.images[0]
+    mm = result.multimodal_output or {}
+    generated_audio = mm.get("audio")
+    fps = int(mm.get("fps", 24))
+    sample_rate = int(mm.get("audio_sample_rate", 16000))
+
+    # DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
+    # mux_video_audio_bytes expects (F, H, W, C) uint8.
+    if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
+        raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
+    frames = generated_video.transpose(1, 2, 3, 0)
+    frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
+
+    audio_np = None
+    if generated_audio is not None:
+        audio_np = np.squeeze(np.asarray(generated_audio)).astype(np.float32)
+
    output_path = args.output
-    save_video(output_path, generated_video, generated_audio, fps=24, sample_rate=16000)
+    video_bytes = mux_video_audio_bytes(
+        frames,
+        audio_np,
+        fps=float(fps),
+        audio_sample_rate=sample_rate,
+    )
+    with open(output_path, "wb") as f:
+        f.write(video_bytes)
    print(f"Saved generated video to {output_path}")
    print(f"Total time: {elapsed:.2f}s")