lint

Signed-off-by: khluu <khluu000@gmail.com>

.buildkite/pipeline.yml

@@ -1,6 +1,21 @@
+# Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci
+# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs if nightly-test
+# / main NIGHTLY so upload-nightly is not skipped together with test-ready/test-merge.
+#
+# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite).
+steps:
+  - label: ":github: Resolve skip-ci & upload pipeline"
+    key: upload-ci-pipeline
+    commands:
+      - "bash .buildkite/scripts/upload_pipeline_with_skip_ci.sh"
+    agents:
+      queue: "cpu_queue_premerge"
+
+---
 steps:
   - label: ":docker: Build image"
     key: image-build
+    if: __IMAGE_BUILD_IF__
     commands:
       - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
       - "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
@@ -13,7 +28,7 @@ steps:
   - label: "Upload Ready Pipeline"
     depends_on: image-build
     key: upload-ready-pipeline
-    if: build.branch != "main" && build.pull_request.labels includes "ready"
+    if: __UPLOAD_READY_IF__
     commands:
       - buildkite-agent pipeline upload .buildkite/test-ready.yml
     agents:
@@ -23,7 +38,7 @@ steps:
   - label: "Upload Merge Pipeline"
     depends_on: image-build
     key: upload-merge-pipeline
-    if: build.branch == "main" && build.env("NIGHTLY") != "1"
+    if: __UPLOAD_MERGE_IF__
     commands:
       - buildkite-agent pipeline upload .buildkite/test-merge.yml
     agents:

.buildkite/scripts/generate-and-upload-nightly-index.sh

@@ -19,7 +19,7 @@ has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12)
 if [[ "$has_new_python" -eq 0 ]]; then
     # use new python from docker
     docker pull python:3-slim
-    PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3"
+    PYTHON="docker run --rm --user $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
 fi
 
 echo "Using python interpreter: $PYTHON"
@@ -36,7 +36,7 @@ mkdir -p "$INDICES_OUTPUT_DIR"
 
 # HACK: we do not need regex module here, but it is required by pre-commit hook
 # To avoid any external dependency, we simply replace it back to the stdlib re module
-sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
+sed -i.bak 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py && rm -f .buildkite/scripts/generate-nightly-index.py.bak
 
 # Generate indices -- the version is just the commit hash (not omni/{commit})
 # because relative paths are computed between the index and wheel directories,
@@ -73,15 +73,16 @@ echo "Pure version (without variant): $pure_version"
 
 # re-generate and copy to /omni/{version}/ only if it does not have "dev" in the version
 if [[ "$version" != *"dev"* ]]; then
-    echo "Re-generating indices for /omni/$pure_version/"
+    s3_version="v$pure_version"
+    echo "Re-generating indices for /omni/$s3_version/"
     rm -rf "${INDICES_OUTPUT_DIR:?}"
     mkdir -p "$INDICES_OUTPUT_DIR"
     # wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
     $PYTHON .buildkite/scripts/generate-nightly-index.py \
-        --version "$pure_version" \
+        --version "$s3_version" \
         --wheel-dir "$BUILDKITE_COMMIT" \
         --current-objects "$obj_json" \
         --output-dir "$INDICES_OUTPUT_DIR" \
         --comment "version $pure_version"
-    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/$pure_version/"
+    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/$s3_version/"
 fi

.buildkite/scripts/generate-nightly-index.py

@@ -4,6 +4,7 @@
 
 import argparse
 import json
+import re
 import sys
 from dataclasses import asdict, dataclass
 from datetime import datetime
@@ -11,8 +12,6 @@ from pathlib import Path
 from typing import Any
 from urllib.parse import quote
 
-import regex as re
-
 
 def normalize_package_name(name: str) -> str:
     """Normalize package name per PEP 503."""

.buildkite/scripts/upload_pipeline_with_skip_ci.sh

@@ -0,0 +1,137 @@
+#!/usr/bin/env bash
+# Evaluate docs-only skip-ci and upload continuation steps from the same `.buildkite/pipeline.yml`
+# (YAML document after the first `---`). Buildkite `if` is evaluated at upload time.
+set -euo pipefail
+
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+PIPELINE_YML="${ROOT}/.buildkite/pipeline.yml"
+
+# Prints a single digit to stdout: 1 = skip image CI, 0 = run. Logs go to stderr.
+is_docs_only_change() {
+  local file_path
+  local has_any=0
+
+  while IFS= read -r file_path; do
+    [[ -z "${file_path}" ]] && continue
+    has_any=1
+
+    if [[ "${file_path}" == docs/* ]]; then
+      continue
+    fi
+    if [[ "${file_path}" == *.md ]]; then
+      continue
+    fi
+    if [[ "${file_path}" == "mkdocs.yaml" ]]; then
+      continue
+    fi
+    return 1
+  done
+
+  [[ "${has_any}" -eq 1 ]]
+}
+
+resolve_skip_ci() {
+  local is_pr_build=0
+  local files
+  local base_branch base_ref
+
+  if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" && -n "${BUILDKITE_PULL_REQUEST:-}" ]]; then
+    is_pr_build=1
+  fi
+
+  if [[ "${is_pr_build}" -eq 1 ]]; then
+    base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
+    if ! git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
+      echo "resolve_skip_ci: origin/${base_branch} not found locally; trying fetch" >&2
+      git fetch --depth=200 origin "${base_branch}" >/dev/null 2>&1 || true
+    fi
+
+    base_ref=""
+    if git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
+      base_ref="origin/${base_branch}"
+    elif git rev-parse --verify "${base_branch}" >/dev/null 2>&1; then
+      base_ref="${base_branch}"
+    else
+      echo "resolve_skip_ci: cannot resolve PR base ${base_branch}; skip-ci=0" >&2
+      echo -n 0
+      return 0
+    fi
+
+    if ! files="$(git diff --name-only "${base_ref}...${BUILDKITE_COMMIT}" 2>/dev/null)"; then
+      echo "resolve_skip_ci: failed to compute PR changed files; skip-ci=0" >&2
+      echo -n 0
+      return 0
+    fi
+  elif [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then
+    if ! git rev-parse --verify "${BUILDKITE_COMMIT}^" >/dev/null 2>&1; then
+      echo "resolve_skip_ci: commit has no parent on main; skip-ci=0" >&2
+      echo -n 0
+      return 0
+    fi
+    if ! files="$(git diff --name-only "${BUILDKITE_COMMIT}^..${BUILDKITE_COMMIT}" 2>/dev/null)"; then
+      echo "resolve_skip_ci: failed to compute main changed files; skip-ci=0" >&2
+      echo -n 0
+      return 0
+    fi
+  else
+    echo "resolve_skip_ci: not PR/main build; skip-ci=0" >&2
+    echo -n 0
+    return 0
+  fi
+
+  if is_docs_only_change <<< "${files}"; then
+    echo "resolve_skip_ci: docs-only change detected; skip-ci=1" >&2
+    echo -n 1
+    return 0
+  fi
+
+  echo "resolve_skip_ci: non-doc changes detected; skip-ci=0" >&2
+  echo -n 0
+}
+
+SKIP_CI="$(resolve_skip_ci)"
+
+if [[ ! -f "${PIPELINE_YML}" ]]; then
+  echo "upload_pipeline_with_skip_ci: missing ${PIPELINE_YML}" >&2
+  exit 1
+fi
+
+export ROOT SKIP_CI PIPELINE_YML
+python3 <<'PY' | buildkite-agent pipeline upload
+import os
+import pathlib
+
+path = pathlib.Path(os.environ["PIPELINE_YML"])
+text = path.read_text(encoding="utf-8")
+sep = "\n---\n"
+if sep not in text:
+    raise SystemExit(
+        "upload_pipeline_with_skip_ci: .buildkite/pipeline.yml must contain a '\\n---\\n' separator "
+        "(document 1 = bootstrap, document 2 = uploaded steps)"
+    )
+_, continuation = text.split(sep, 1)
+
+skip = os.environ.get("SKIP_CI") == "1"
+# When docs-only skip-ci: skip default CI image, but still build for L4 nightly (PR label nightly-test or
+# main NIGHTLY=1), otherwise upload-nightly (depends_on image-build) would be skipped too.
+nightly_only = (
+    '(build.pull_request.labels includes "nightly-test") '
+    '|| (build.branch == "main" && build.env("NIGHTLY") == "1")'
+)
+# Placeholder in pipeline.yml is `if: __IMAGE_BUILD_IF__` (valid YAML); replace value only.
+if skip:
+    rep = f"'{nightly_only}'"
+    ready_rep = "'false'"
+    merge_rep = "'false'"
+else:
+    rep = "'true'"
+    ready_rep = "'build.branch != \"main\" && build.pull_request.labels includes \"ready\"'"
+    merge_rep = "'build.branch == \"main\" && build.env(\"NIGHTLY\") != \"1\"'"
+rendered = (
+    continuation
+    .replace("__IMAGE_BUILD_IF__", rep)
+    .replace("__UPLOAD_READY_IF__", ready_rep)
+    .replace("__UPLOAD_MERGE_IF__", merge_rep)
+)
+print(rendered, end="")
+PY

.buildkite/test-merge.yml

@@ -390,16 +390,6 @@ steps:
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
           pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
         '
-
-  - label: "CosyVoice3-TTS E2E Test"
-    timeout_in_minutes: 20
-    depends_on: upload-merge-pipeline
-    commands:
-      - |
-        timeout 20m bash -c '
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts.py -m "advanced_model" --run-level "advanced_model"
-        '
     agents:
       queue: "mithril-h100-pool"
     plugins:

.buildkite/test-nightly.yml

@@ -72,7 +72,7 @@ steps:
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
+      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -107,13 +107,13 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: ":full_moon: Diffusion Model (Wan2.2) Test with H100"
+  - label: ":full_moon: Diffusion Model (Wan) Test with H100"
     timeout_in_minutes: 90
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -m "advanced_model" --run-level "advanced_model"
+      - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py -m "advanced_model" --run-level "advanced_model"
     agents:
       queue: "mithril-h100-pool"
     plugins:

.gitignore

@@ -83,6 +83,9 @@ target/
 profile_default/
 ipython_config.py
 
+# uv
+uv.lock
+
 # pyenv
 #   For a library or package, you might want to ignore these files since the code is
 #   intended to run in multiple environments; otherwise, check them in:

docker/Dockerfile.cuda

@@ -0,0 +1,22 @@
+ARG BASE_IMAGE=vllm/vllm-openai:v0.19.0
+FROM ${BASE_IMAGE}
+
+ARG COMMON_WORKDIR=/app
+
+WORKDIR ${COMMON_WORKDIR}
+
+# Step 1: Setup - Install system dependencies
+RUN apt-get update && \
+    apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
+
+# Step 2: Copy vllm-omni code and install
+COPY . ${COMMON_WORKDIR}/vllm-omni
+RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir "."
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+ENTRYPOINT []

docs/contributing/ci/CI_5levels.md

@@ -271,7 +271,7 @@ Before entering specific testing levels, the project establishes two common spec
 
 L1 and L2 level testing form the foundation of the quality assurance system. L1 level testing focuses on verifying the internal logic correctness of code units (e.g., functions, classes), ensuring each independent component behaves as designed.
 
-L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle  .
+L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle.
 
 
 

docs/contributing/model/adding_diffusion_model.md

@@ -802,7 +802,7 @@ omni = Omni(model="your-model", enable_layerwise_offload=True)
 
 ```python
 class WanTransformer3DModel(nn.Module):
-    _layerwise_offload_blocks_attr = "blocks"  # Attribute name containing transformer blocks
+    _layerwise_offload_blocks_attrs = ["blocks"]  # Attribute name containing transformer blocks
 
     def __init__(self):
         self.blocks = nn.ModuleList([...])  # Transformer blocks

docs/design/feature/vae_parallel.md

@@ -1,14 +1,15 @@
 # VAE Patch Parallelism
 
 This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
-We use **Qwen-Image** as the reference implementation.
+We use **Qwen-Image** as the reference implementation for decode parallel, and **Wan2.2** for encode parallel.
 
 ---
 
 ## Table of Contents
 
 - [Overview](#overview)
-- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Step-by-Step Implementation (Decode)](#step-by-step-implementation-decode)
+- [Encode Parallel Implementation](#encode-parallel-implementation)
 - [Testing](#testing)
 - [Reference Implementations](#reference-implementations)
 - [Summary](#summary)
@@ -19,13 +20,13 @@ We use **Qwen-Image** as the reference implementation.
 
 ### What is Vae Patch parallel?
 
-**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is:
+**VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is:
 
 + Split into multiple spatial tiles
 
 + Distributed across multiple ranks
 
-+ Decoded in parallel
++ Encoded/Decoded in parallel
 
 + Merged to reconstruct the final output
 
@@ -35,10 +36,17 @@ This approach:
 
 + Reduces peak memory usage per device
 
-+ Accelerates decoding latency
++ Accelerates encoding/decoding latency
+
+### When to Use Encode vs Decode Parallel
+
+| Operation | Use Case | Example |
+|-----------|----------|---------|
+| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
+| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
 
 ### Architecture
-We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding.
+We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding.
 
 The executor is model-agnostic and accepts three function parameters:
 
@@ -84,7 +92,7 @@ Therefore:
 
 + Merge must perform blending to avoid seams
 
-## Step-by-Step Implementation
+## Step-by-Step Implementation (Decode)
 
 ### Step 1: Implement DistributedAutoencoderKLQwenImage
 `QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
@@ -205,14 +213,14 @@ def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid
 We need to override tiled_decode, the main logic is:
 + check distributed is enabled
 + select split/exec/merge
-+ Invoke self.distributed_decoder.execute to decode
++ Invoke self.distributed_executor.execute to decode
 ```
 def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
     if not self.is_distributed_enabled():
         return super().tiled_decode(z, return_dict=return_dict)
 
     logger.info("Decode run with distributed executor")
-    result = self.distributed_decoder.execute(
+    result = self.distributed_executor.execute(
         z,
         DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
         broadcast_result=True,
@@ -243,6 +251,166 @@ class YourModelPipeline(nn.Module):
 +       ).to(self.device)
 ```
 
+## Encode Parallel Implementation
+
+For models that require VAE encoding (e.g., Image-to-Video), you can also parallelize the encode operation. We use **Wan2.2** as the reference implementation.
+
+### Step 1: Implement encode_tile_split
+
+Similar to decode, split the input tensor into tiles. Key considerations:
+
++ **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly
++ **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x)
+
+```python
+def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+    _, _, num_frames, height, width = x.shape
+    encode_spatial_compression_ratio = self.spatial_compression_ratio
+
+    # Scale tile parameters for patchified coordinate system
+    tile_sample_min_height = self.tile_sample_min_height
+    tile_sample_min_width = self.tile_sample_min_width
+    tile_sample_stride_height = self.tile_sample_stride_height
+    tile_sample_stride_width = self.tile_sample_stride_width
+
+    if self.config.patch_size is not None:
+        # When input is patchified, scale tile parameters accordingly
+        encode_spatial_compression_ratio = self.spatial_compression_ratio // self.config.patch_size
+        tile_sample_min_height = tile_sample_min_height // self.config.patch_size
+        tile_sample_min_width = tile_sample_min_width // self.config.patch_size
+        tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
+        tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
+
+    latent_height = height // encode_spatial_compression_ratio
+    latent_width = width // encode_spatial_compression_ratio
+
+    tile_latent_min_height = tile_sample_min_height // encode_spatial_compression_ratio
+    tile_latent_min_width = tile_sample_min_width // encode_spatial_compression_ratio
+    tile_latent_stride_height = tile_sample_stride_height // encode_spatial_compression_ratio
+    tile_latent_stride_width = tile_sample_stride_width // encode_spatial_compression_ratio
+
+    blend_height = tile_latent_min_height - tile_latent_stride_height
+    blend_width = tile_latent_min_width - tile_latent_stride_width
+
+    tiletask_list = []
+    # Use temporal compression ratio from config instead of hardcoding
+    temporal_compression = self.config.scale_factor_temporal
+
+    for i in range(0, height, tile_sample_stride_height):
+        for j in range(0, width, tile_sample_stride_width):
+            time_list = []
+            frame_range = 1 + (num_frames - 1) // temporal_compression
+            for k in range(frame_range):
+                if k == 0:
+                    tile = x[:, :, :1, i : i + tile_sample_min_height, j : j + tile_sample_min_width]
+                else:
+                    tile = x[
+                        :, :,
+                        1 + temporal_compression * (k - 1) : 1 + temporal_compression * k,
+                        i : i + tile_sample_min_height,
+                        j : j + tile_sample_min_width,
+                    ]
+                time_list.append(tile)
+            tiletask_list.append(
+                TileTask(len(tiletask_list), (i // tile_sample_stride_height, j // tile_sample_stride_width),
+                         time_list, workload=time_list[0].shape[3] * time_list[0].shape[4])
+            )
+
+    grid_spec = GridSpec(
+        split_dims=(3, 4),
+        grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+        tile_spec={
+            "latent_height": latent_height, "latent_width": latent_width,
+            "blend_height": blend_height, "blend_width": blend_width,
+            "tile_latent_stride_height": tile_latent_stride_height,
+            "tile_latent_stride_width": tile_latent_stride_width,
+        },
+        output_dtype=self.dtype,
+    )
+    return tiletask_list, grid_spec
+```
+
+### Step 2: Implement encode_tile_exec
+
+```python
+def encode_tile_exec(self, task: TileTask) -> torch.Tensor:
+    """Encode a single sample tile into latent space."""
+    self.clear_cache()
+    time = []
+    for k, tile in enumerate(task.tensor):
+        self._enc_conv_idx = [0]
+        encoded = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+        encoded = self.quant_conv(encoded)
+        time.append(encoded)
+    result = torch.cat(time, dim=2)
+    self.clear_cache()
+    return result
+```
+
+### Step 3: Implement encode_tile_merge
+
+```python
+def encode_tile_merge(
+    self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec
+) -> torch.Tensor:
+    """Merge encoded tiles into a full latent tensor."""
+    grid_h, grid_w = grid_spec.grid_shape
+    result_rows = []
+    for i in range(grid_h):
+        result_row = []
+        for j in range(grid_w):
+            tile = coord_tensor_map[(i, j)]
+            if i > 0:
+                tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+            if j > 0:
+                tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+            result_row.append(tile[:, :, :,
+                : grid_spec.tile_spec["tile_latent_stride_height"],
+                : grid_spec.tile_spec["tile_latent_stride_width"]])
+        result_rows.append(torch.cat(result_row, dim=-1))
+
+    enc = torch.cat(result_rows, dim=3)[
+        :, :, :, : grid_spec.tile_spec["latent_height"], : grid_spec.tile_spec["latent_width"]
+    ]
+    return enc
+```
+
+### Step 4: Override tiled_encode method
+
+Override `tiled_encode` instead of `encode`. The parent's `_encode()` handles patchify before calling `tiled_encode()`, so input `x` is already patchified.
+
+```python
+def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+    """
+    Encode using distributed VAE executor.
+
+    Note: x is already patchified by parent's _encode() before calling this method.
+    """
+    if not self.is_distributed_enabled():
+        return super().tiled_encode(x)
+
+    self.clear_cache()
+    result = self.distributed_executor.execute(
+        x,
+        DistributedOperator(
+            split=self.encode_tile_split,
+            exec=self.encode_tile_exec,
+            merge=self.encode_tile_merge,
+        ),
+        broadcast_result=True,  # Latents needed by all ranks for diffusion
+    )
+    self.clear_cache()
+    return result
+```
+
+**Key differences from decode parallel:**
+
+| Aspect | Decode Parallel | Encode Parallel |
+|--------|-----------------|-----------------|
+| `broadcast_result` | Often `False` (only rank 0 needs output) | `True` (all ranks need latents for diffusion) |
+| Patchify | Applied in merge (unpatchify) | Handled by parent `_encode()` before `tiled_encode()` |
+| Temporal chunking | Frame-by-frame | Chunk-based (e.g., 1 + 4n frames) |
+
 ## Testing
 Verify numerical consistency between:
 + vae_patch_parallel_size = 1
@@ -272,18 +440,20 @@ When vae_patch_parallel_size is larger than the DiT world size, it will automati
 
 Complete examples in the codebase:
 
-| Model | Path | Notes |
-|-------|------|-------|
-| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
-| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
-| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage |
+| Model | Path | Decode Parallel | Encode Parallel |
+|-------|------|-----------------|-----------------|
+| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | ✅ | ❌ |
+| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | ✅ | ✅ |
+| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | ✅ | ❌ |
 
 ---
 
 ## Summary
 
-Adding Vae Patch Parallel support to diffusion model:
+Adding VAE Patch Parallel support to diffusion model:
 
-1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
-2. **Change vae model in pipeline to Distributed Vae**
-3. **Test** - Verify with `tensor_parallel_size=N` quality
+1. **Implement Distributed VAE** - Inherit from base VAE class and `DistributedVaeMixin`
+2. **Decode Parallel** - Refactor `tiled_decode` into `tile_split`/`tile_exec`/`tile_merge`
+3. **Encode Parallel** (optional) - Implement `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models
+4. **Change VAE model in pipeline** - Use the distributed version
+5. **Test** - Verify numerical consistency with `vae_patch_parallel_size=1` vs `N`

docs/getting_started/installation/README.md

@@ -6,4 +6,5 @@ vLLM-Omni supports the following hardware platforms:
     - [NVIDIA CUDA](gpu.md)
     - [AMD ROCm](gpu.md)
     - [Intel XPU](gpu.md)
+    - [MThreads MUSA](gpu.md)
 - [NPU](npu.md)

docs/getting_started/installation/gpu.md

@@ -22,6 +22,10 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:requirements"
 
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:requirements"
+
 ## Set up using Python
 
 ### Create a new Python environment
@@ -44,6 +48,10 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:pre-built-wheels"
 
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:pre-built-wheels"
+
 [](){ #build-from-source }
 
 ### Build wheel from source
@@ -60,6 +68,10 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-wheel-from-source"
 
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:build-wheel-from-source"
+
 ## Set up using Docker
 
 ### Pre-built images
@@ -76,6 +88,10 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:pre-built-images"
 
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:pre-built-images"
+
 ### Build your own docker image
 
 === "AMD ROCm"
@@ -85,3 +101,7 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 === "Intel XPU"
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-docker"
+
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:build-docker"

docs/getting_started/installation/gpu/musa.inc.md

@@ -0,0 +1,65 @@
+# --8<-- [start:requirements]
+
+- GPU: Moore Threads GPU with MUSA SDK installed (validated on MTT S5000)
+
+# --8<-- [end:requirements]
+# --8<-- [start:set-up-using-python]
+
+vLLM-Omni for MUSA requires building from source. Pre-built wheels are not currently available.
+
+!!! note
+    MUSA platform requires vLLM-MUSA to be installed first.
+
+# --8<-- [start:pre-built-wheels]
+
+# --8<-- [end:pre-built-wheels]
+
+# --8<-- [start:build-wheel-from-source]
+
+#### Prerequisites
+
+- **MUSA SDK**: Download from [MUSA SDK Download](https://developer.mthreads.com/sdk/download/musa)
+- **torchada**: CUDA→MUSA compatibility layer for PyTorch (`pip install torchada`)
+- **mthreads-ml-py**: MTML Python bindings (`pip install mthreads-ml-py`)
+- **MATE**: MUSA AI Tensor Engine ([GitHub](https://github.com/MooreThreads/mate))
+
+#### Installation of vLLM-MUSA
+
+```bash
+git clone https://github.com/MooreThreads/vllm-musa.git
+cd vllm-musa
+git checkout v0.18.0-dev
+pip install . --no-build-isolation -v
+```
+
+#### Installation of vLLM-Omni
+
+```bash
+git clone https://github.com/vllm-project/vllm-omni.git
+cd vllm-omni
+VLLM_OMNI_TARGET_DEVICE=musa pip install -e . --no-build-isolation
+```
+
+For Gradio demos:
+
+```bash
+pip install -e '.[demo]' --no-build-isolation
+```
+
+#### Environment Variables
+
+```bash
+export MUSA_VISIBLE_DEVICES=0,1
+export VLLM_WORKER_MULTIPROC_METHOD=spawn
+export VLLM_MUSA_CUSTOM_OP_USE_NATIVE=false
+```
+
+# --8<-- [end:build-wheel-from-source]
+
+# --8<-- [start:build-docker]
+
+# --8<-- [end:build-docker]
+
+# --8<-- [start:pre-built-images]
+
+# --8<-- [end:pre-built-images]

docs/models/supported_models.md

@@ -37,6 +37,7 @@ th {
 | `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
+| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
 | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -91,12 +91,19 @@ Models must define the blocks attribute name for layerwise offloading:
 
 ```python
 class WanTransformer3DModel(nn.Module):
-    _layerwise_offload_blocks_attr = "blocks"  # Attribute name containing transformer blocks
+    _layerwise_offload_blocks_attrs = ["blocks"]  # Attribute names containing transformer blocks
 
     def __init__(self):
         self.blocks = nn.ModuleList([...])  # Transformer blocks
 ```
 
+For models with multiple block types:
+
+```python
+class Flux2Transformer2DModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
+```
+
 ### Limitations
 - Cold start latency increases because of
     1) components are loaded to CPU first at the very first during initialization,
@@ -140,4 +147,4 @@ Factory function `get_offload_backend()` selects the appropriate backend based o
 
 **Notes:**
 - Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
-- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attr` pointing to transformer blocks
+- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks

docs/user_guide/diffusion_features.md

@@ -109,31 +109,32 @@ The following tables show which models support each feature:
 | **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
 | **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
 | **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
-| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ | ✅ | ❌ |
+| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
+| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
+| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
+| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ (decode) | ❌ | ❌ |
+| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
 
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
-> 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary.
+> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
 
 ### VideoGen
 
 | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
+| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
 | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
 | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 
 ### AudioGen

examples/offline_inference/bagel/end2end.py

@@ -2,7 +2,10 @@ import argparse
 import os
 
 from vllm_omni.inputs.data import OmniPromptType
-from vllm_omni.model_executor.stage_input_processors.bagel import GEN_THINK_SYSTEM_PROMPT
+from vllm_omni.model_executor.stage_input_processors.bagel import (
+    GEN_THINK_SYSTEM_PROMPT,
+    VLM_THINK_SYSTEM_PROMPT,
+)
 
 
 def parse_args():
@@ -171,7 +174,10 @@ def main():
         elif args.modality == "img2text":
             if args.image_path:
                 loaded_image = Image.open(args.image_path).convert("RGB")
-                final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                think_prefix = f"<|im_start|>system\n{VLM_THINK_SYSTEM_PROMPT}<|im_end|>\n" if args.think else ""
+                final_prompt_text = (
+                    f"{think_prefix}<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                )
                 prompt_dict = {
                     "prompt": final_prompt_text,
                     "multi_modal_data": {"image": loaded_image},
@@ -179,7 +185,8 @@ def main():
                 }
                 formatted_prompts.append(prompt_dict)
         elif args.modality == "text2text":
-            final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
+            think_prefix = f"<|im_start|>{VLM_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
+            final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|><|im_start|>"
             prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
             formatted_prompts.append(prompt_dict)
         else:
@@ -217,15 +224,11 @@ def main():
     img_idx = 0
     for req_output in omni_outputs:
         if args.think:
-            text_output = getattr(req_output, "text", None) or getattr(req_output, "outputs", None)
-            if text_output:
-                if isinstance(text_output, list) and text_output:
-                    for out in text_output:
-                        txt = getattr(out, "text", str(out))
-                        if txt:
-                            print(f"[Think] {txt}")
-                elif isinstance(text_output, str):
-                    print(f"[Think] {text_output}")
+            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(txt)
 
         images = getattr(req_output, "images", None)
 

examples/offline_inference/magi_human/README.md

@@ -0,0 +1,72 @@
+# MagiHuman Generation
+
+MagiHuman is an advanced, omni-modality model that generates both high-quality video and lip-synced audio from a text prompt.
+
+Because MagiHuman is a very large model featuring a powerful DiT MoE backbone and a ~9B parameter T5Gemma text encoder, it natively supports **Tensor Parallelism (TP)** in vLLM-Omni to run efficiently across multi-GPU setups, reducing device memory bottlenecks.
+
+## Setup
+
+### Install MagiCompiler (recommended)
+
+MagiHuman relies on [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for custom-op registration used by the DiT attention kernels. While the pipeline can fall back to stub implementations, installing MagiCompiler is **strongly recommended** for correct behaviour.
+
+```bash
+# Clone the repo
+git clone https://github.com/SandAI-org/MagiCompiler.git
+cd MagiCompiler
+
+# System dependencies (optional, for FX graph visualization; Debian/Ubuntu)
+sudo apt update && sudo apt install -y graphviz
+
+# Python dependencies
+pip install -r requirements.txt
+
+# Install MagiCompiler
+pip install .          # end users (recommended)
+# pip install -e .     # developers (editable install)
+```
+
+### Hardware requirements
+
+Ensure your hardware has enough VRAM. For a standard node with 80GB GPUs, running with `--tensor-parallel-size 4` is recommended to shard both the MoE weights and the T5Gemma text encoder across 4 GPUs, reducing the per-GPU peak VRAM overhead significantly (by roughly ~13.5GB per GPU compared to single-device inference).
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for further details on allocating memory.
+
+## Run Examples
+
+Get into the example folder:
+```bash
+cd examples/offline_inference/magi_human
+```
+
+### End-to-End Generation (Text to Video+Audio)
+
+Generate a video with synchronized speech natively generated by the model.
+
+```bash
+python end2end.py \
+    --model /proj-tango-pvc/users/zhipeng.wang/workspace/models/daVinci-MagiHuman \
+    --prompt "A young woman with long, wavy golden blonde hair..." \
+    --tensor-parallel-size 4 \
+    --output output_magihuman.mp4
+```
+
+## Common Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--model` | *(Required)* | Local model path or HuggingFace ID |
+| `--prompt` | *(built-in demo prompt)* | Highly detailed text prompt dictating visual look and dialogue text |
+| `--tensor-parallel-size` | `4` | Tensor parallelism size (Number of GPUs) |
+| `--height` | `256` | Initial resolution height |
+| `--width` | `448` | Initial resolution width |
+| `--num-inference-steps` | `8` | Denoising steps |
+| `--seed` | `52` | Random seed |
+| `--output` | `output_magihuman.mp4` | Output video with audio path |
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/magi_human/end2end.py"
+    ``````

examples/offline_inference/magi_human/end2end.py

@@ -0,0 +1,117 @@
+import argparse
+
+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
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="End-to-end inference script for MagiHuman.")
+    parser.add_argument("--model", type=str, required=True, help="Path or ID of the MagiHuman model.")
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default="",
+        help="Text prompt containing visual description, dialogue, and background sound.",
+    )
+    parser.add_argument(
+        "--tensor-parallel-size", "-tp", type=int, default=4, help="Tensor parallel size (number of GPUs)."
+    )
+    parser.add_argument(
+        "--output", type=str, default="output_magihuman.mp4", help="Path to save the generated mp4 file."
+    )
+    parser.add_argument("--height", type=int, default=256, help="Video height.")
+    parser.add_argument("--width", type=int, default=448, help="Video width.")
+    parser.add_argument("--num-inference-steps", type=int, default=8, help="Number of denoising steps.")
+    parser.add_argument("--seed", type=int, default=52, help="Random seed for generation.")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    print(f"Initializing MagiHuman pipeline with TP={args.tensor_parallel_size}...")
+    omni = Omni(
+        model=args.model,
+        init_timeout=1200,
+        tensor_parallel_size=args.tensor_parallel_size,
+        devices=list(range(args.tensor_parallel_size)),
+    )
+
+    prompt = args.prompt
+    if not prompt:
+        prompt = (
+            "A young woman with long, wavy golden blonde hair and bright blue eyes, "
+            "wearing a fitted ivory silk blouse with a delicate lace collar, sits "
+            "stationary in front of a softly lit, blurred warm-toned interior. Her "
+            "overall disposition is warm, composed, and gently confident. The camera "
+            "holds a static medium close-up, framing her from the shoulders up, "
+            "with shallow depth of field keeping her face in sharp focus. Soft "
+            "directional key light falls from the upper left, casting a gentle "
+            "highlight along her cheekbone and nose bridge. She draws a quiet breath, "
+            "the levator labii superiors relaxing as her lips part. She speaks in "
+            "clear, warm, unhurried American English: "
+            "\"The most beautiful things in life aren't things at all — "
+            "they're moments, feelings, and the people who make you feel truly alive.\" "
+            "Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
+            "shapes each vowel with precision. A faint, genuine smile engages the "
+            "zygomaticus major, lifting her lip corners fractionally. Her brows rest "
+            "in a soft, neutral arch throughout. She maintains steady, forward-facing "
+            "eye contact. Head position remains level; no torso displacement occurs.\n\n"
+            "Dialogue:\n"
+            "<Young blonde woman, American English>: "
+            "\"The most beautiful things in life aren't things at all — "
+            "they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
+            "Background Sound:\n"
+            "<Soft, warm indoor ambience with a faint distant piano melody>"
+        )
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.num_inference_steps,
+        seed=args.seed,
+        extra_args={
+            "seconds": 5,
+            "sr_height": 1080,
+            "sr_width": 1920,
+            "sr_num_inference_steps": 5,
+        },
+    )
+
+    print(f"Generating with prompt: {prompt[:80]}...")
+    outputs = omni.generate(
+        prompts=[prompt],
+        sampling_params_list=[sampling_params],
+    )
+
+    print(f"Generation complete. Output type: {type(outputs)}")
+    if outputs:
+        first = outputs[0]
+
+        if hasattr(first, "images") and first.images:
+            video_frames = first.images[0]
+            print(f"Video frames: shape={video_frames.shape}, dtype={video_frames.dtype}")
+
+            audio_waveform = None
+            if hasattr(first, "multimodal_output") and first.multimodal_output:
+                audio_waveform = first.multimodal_output.get("audio")
+                if audio_waveform is not None:
+                    print(f"Audio waveform: shape={audio_waveform.shape}, dtype={audio_waveform.dtype}")
+
+            video_bytes = mux_video_audio_bytes(
+                video_frames,
+                audio_waveform,
+                fps=25.0,
+                audio_sample_rate=44100,
+            )
+            with open(args.output, "wb") as f:
+                f.write(video_bytes)
+            print(f"Saved MP4 ({len(video_bytes)} bytes) to {args.output}")
+        print("SUCCESS: MagiHuman pipeline generation completed.")
+    else:
+        print("WARNING: No outputs returned.")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/text_to_image/README.md

@@ -33,6 +33,7 @@ This folder provides several entrypoints for experimenting with text-to-image di
 | `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) |
+| `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3)  | 160 |
 
 !!! info
 *Peak VRAM:  based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU.
@@ -90,6 +91,8 @@ python text_to_image.py \
 | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
 | `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
 | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
+| `--use-system-prompt` | str | `None` | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text. Recommended: `en_unified`. Only for HunyuanImage-3.0.|
+| `--system-prompt` | str | `None` | Custom system prompt text. Only used when `--use-system-prompt` is set to `custom`. Only for HunyuanImage-3.0.|
 
 **NextStep-1.1 specific arguments:**
 

examples/offline_inference/text_to_image/text_to_image.py

@@ -242,6 +242,19 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable logging of diffusion pipeline stats.",
     )
+    parser.add_argument(
+        "--use-system-prompt",
+        type=str,
+        default=None,
+        choices=["None", "dynamic", "en_vanilla", "en_recaption", "en_think_recaption", "en_unified", "custom"],
+        help="System prompt preset for generation. Recommended: en_unified.",
+    )
+    parser.add_argument(
+        "--system-prompt",
+        type=str,
+        default=None,
+        help=("Custom system prompt. Used when --use-system-prompt is custom. "),
+    )
     return parser.parse_args()
 
 
@@ -382,13 +395,13 @@ def main():
         )
 
     generation_start = time.perf_counter()
-
     extra_args = {
         "timesteps_shift": args.timesteps_shift,
         "cfg_schedule": args.cfg_schedule,
         "use_norm": args.use_norm,
+        "use_system_prompt": args.use_system_prompt,
+        "system_prompt": args.system_prompt,
     }
-
     if lora_request:
         extra_args["lora_request"] = lora_request
         extra_args["lora_scale"] = args.lora_scale

examples/offline_inference/x_to_video_audio/download_dreamid_omni.py

@@ -82,7 +82,6 @@ def main(output_dir: str):
 
     data = {
         "_class_name": "DreamIDOmniPipeline",
-        "fusion": "DreamID-Omni/dreamid_omni.safetensors",
     }
 
     with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
@@ -90,6 +89,12 @@ def main(output_dir: str):
 
     print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}")
 
+    transformer_dir = os.path.join(output_dir, "transformer")
+    os.makedirs(transformer_dir, exist_ok=True)
+    with open(os.path.join(transformer_dir, "config.json"), "w", encoding="utf-8") as f:
+        json.dump({"fusion": "DreamID-Omni/dreamid_omni.safetensors"}, f)
+    print(f"transformer/config.json created at {os.path.join(transformer_dir, 'config.json')}")
+
     # now we download the dependency code
     download_dependency()
 

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -24,7 +24,9 @@ dreamid_omni/
 │   ├── models_t5_umt5-xxl-enc-bf16.pth
 │   ├── Wan2.2_VAE.pth
 │
-├── model_index.json # create by download_dreamid_omni.py
+├── model_index.json
+└── transformer/
+    └── config.json   # create by download_dreamid_omni.py
 ```
 
 ### Run the Inference

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -132,8 +132,8 @@ def main() -> None:
     if not outputs:
         raise RuntimeError("No output returned from DreamID-Omni.")
     output = outputs[0].request_output
-    generated_video = output[0].images[0][0]
-    generated_audio = output[0].images[0][1]
+    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:

examples/online_serving/qwen3_tts/openai_speech_client.py

@@ -71,7 +71,7 @@ def run_tts_generation(args) -> None:
     payload = {
         "model": args.model,
         "input": args.text,
-        "speaker": args.speaker,
+        "voice": args.speaker,
         "response_format": args.response_format,
     }
 

examples/online_serving/text_to_image/README.md

@@ -231,6 +231,8 @@ count, use `size` and `n` rather than `height`, `width`, or
 | `seed`                   | int   | None    | Random seed (reproducible)     |
 | `negative_prompt`        | str   | None    | Negative prompt                |
 | `num_outputs_per_prompt` | int   | 1       | Number of images to generate   |
+| `use_system_prompt` | str | None | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text string. Only for HunyuanImage-3.0. |
+| `system_prompt` | str | None | Custom system prompt text. Only used when `use_system_prompt` is set to `custom`. Only for HunyuanImage-3.0. |
 
 ## Response Format
 

examples/online_serving/text_to_image/openai_chat_client.py

@@ -28,6 +28,8 @@ def generate_image(
     lora_name: str | None = None,
     lora_scale: float | None = None,
     lora_int_id: int | None = None,
+    use_system_prompt: str | None = None,
+    system_prompt: str | None = None,
 ) -> bytes | None:
     """Generate an image using the images generation API.
 
@@ -45,6 +47,8 @@ def generate_image(
         lora_name: LoRA name (optional, defaults to path stem)
         lora_scale: LoRA scale factor (default: 1.0)
         lora_int_id: LoRA integer ID (optional, derived from path if not provided)
+        use_system_prompt: System prompt for generation.
+        system_prompt: Custom system prompt.
 
     Returns:
         Image bytes or None if failed
@@ -70,7 +74,10 @@ def generate_image(
         payload["negative_prompt"] = negative_prompt
     if seed is not None:
         payload["seed"] = seed
-
+    if use_system_prompt is not None:
+        payload["use_system_prompt"] = use_system_prompt
+    if system_prompt is not None:
+        payload["system_prompt"] = system_prompt
     # Add LoRA if provided
     if lora_path:
         lora_body: dict = {
@@ -128,9 +135,21 @@ def main():
         default=None,
         help="LoRA integer id (cache key). If omitted, the server derives a stable id from lora_path.",
     )
-
+    parser.add_argument(
+        "--use-system-prompt",
+        type=str,
+        default=None,
+        help=(
+            "System prompt for generation. Use predefined types: 'en_unified', 'en_vanilla', 'en_recaption', 'en_think_recaption', 'dynamic', or 'None'; Or provide custom text string directly. Recommended en_unified. "
+        ),
+    )
+    parser.add_argument(
+        "--system-prompt",
+        type=str,
+        default=None,
+        help=("Custom system prompt. Used when --use-system-prompt is custom. "),
+    )
     args = parser.parse_args()
-
     print(f"Generating image for: {args.prompt}")
 
     image_bytes = generate_image(
@@ -146,6 +165,8 @@ def main():
         lora_name=args.lora_name,
         lora_scale=args.lora_scale if args.lora_path else None,
         lora_int_id=args.lora_int_id if args.lora_path else None,
+        use_system_prompt=args.use_system_prompt,
+        system_prompt=args.system_prompt,
     )
 
     if image_bytes:

requirements/common.txt

@@ -1,4 +1,5 @@
 # Common dependencies for all platforms
+av>=14.0.0
 omegaconf>=2.3.0
 librosa>=0.11.0
 resampy>=0.4.3

tests/dfx/perf/stage_configs/qwen3_tts.yaml

@@ -88,7 +88,7 @@ runtime:
         connector_get_max_wait_first_chunk: 3000
         connector_get_max_wait: 300
         codec_chunk_frames: 25
-        codec_left_context_frames: 25
+        codec_left_context_frames: 72
 
   edges:
     - from: 0

tests/diffusion/distributed/test_autoencoder_kl_wan_encode.py

@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for DistributedAutoencoderKLWan encode parallel (CPU-only)."""
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _DummyConfig:
+    def __init__(self, patch_size=None, scale_factor_temporal=4):
+        self.patch_size = patch_size
+        self.scale_factor_temporal = scale_factor_temporal
+
+
+class _DummyWanVae:
+    """Minimal mock of DistributedAutoencoderKLWan for testing encode_tile_split."""
+
+    def __init__(
+        self,
+        config=None,
+        spatial_compression_ratio=8,
+        tile_sample_min_height=256,
+        tile_sample_min_width=256,
+        tile_sample_stride_height=192,
+        tile_sample_stride_width=192,
+    ):
+        self.config = config or _DummyConfig()
+        self.spatial_compression_ratio = spatial_compression_ratio
+        self.tile_sample_min_height = tile_sample_min_height
+        self.tile_sample_min_width = tile_sample_min_width
+        self.tile_sample_stride_height = tile_sample_stride_height
+        self.tile_sample_stride_width = tile_sample_stride_width
+        self.dtype = torch.float32
+
+        # Mock caches
+        self._enc_feat_map = None
+        self._enc_conv_idx = [0]
+
+    def clear_cache(self):
+        self._enc_feat_map = None
+        self._enc_conv_idx = [0]
+
+    def encoder(self, x, feat_cache=None, feat_idx=None):  # noqa: ARG002
+        # Simple mock: just return the input
+        return x
+
+    def quant_conv(self, x):
+        return x
+
+    def blend_v(self, _a, b, _blend_extent):
+        return b
+
+    def blend_h(self, _a, b, _blend_extent):
+        return b
+
+
+def _import_encode_tile_split():
+    """Import the encode_tile_split method from the module."""
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+        DistributedAutoencoderKLWan,
+    )
+
+    return DistributedAutoencoderKLWan.encode_tile_split
+
+
+def _import_encode_tile_exec():
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+        DistributedAutoencoderKLWan,
+    )
+
+    return DistributedAutoencoderKLWan.encode_tile_exec
+
+
+def _import_encode_tile_merge():
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+        DistributedAutoencoderKLWan,
+    )
+
+    return DistributedAutoencoderKLWan.encode_tile_merge
+
+
+class TestEncodeTileSplit:
+    """Tests for encode_tile_split method."""
+
+    def test_basic_split_without_patch_size(self):
+        """Test basic tile splitting without patch_size."""
+        encode_tile_split = _import_encode_tile_split()
+
+        vae = _DummyWanVae(
+            config=_DummyConfig(patch_size=None, scale_factor_temporal=4),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=256,
+            tile_sample_min_width=256,
+            tile_sample_stride_height=192,
+            tile_sample_stride_width=192,
+        )
+
+        # Input: (B, C, T, H, W) = (1, 3, 5, 256, 256)
+        x = torch.randn(1, 3, 5, 256, 256)
+
+        tiletask_list, grid_spec = encode_tile_split(vae, x)
+
+        # With stride 192 and input size 256, we should get:
+        # Height: ceil(256/192) = 2 positions (0, 192) but 192+256 > 256, so only 1
+        # Actually for i in range(0, 256, 192): i = 0, 192 but 192 is out of bounds
+        # So we get 1x1 grid
+        assert len(tiletask_list) >= 1
+        assert grid_spec.grid_shape[0] >= 1
+        assert grid_spec.grid_shape[1] >= 1
+
+        # Check temporal chunking: 5 frames -> 1 + (5-1)//4 = 2 chunks
+        first_task = tiletask_list[0]
+        assert len(first_task.tensor) == 2  # 2 temporal chunks
+
+    def test_split_with_patch_size_scales_coordinates(self):
+        """Test that patch_size properly scales tile coordinates."""
+        encode_tile_split = _import_encode_tile_split()
+
+        # Without patch_size
+        vae_no_patch = _DummyWanVae(
+            config=_DummyConfig(patch_size=None, scale_factor_temporal=4),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=256,
+            tile_sample_min_width=256,
+            tile_sample_stride_height=128,
+            tile_sample_stride_width=128,
+        )
+
+        # With patch_size=2 (simulating patchified input)
+        vae_with_patch = _DummyWanVae(
+            config=_DummyConfig(patch_size=2, scale_factor_temporal=4),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=256,
+            tile_sample_min_width=256,
+            tile_sample_stride_height=128,
+            tile_sample_stride_width=128,
+        )
+
+        # Same patchified input size
+        x = torch.randn(1, 3, 5, 256, 256)
+
+        tasks_no_patch, _ = encode_tile_split(vae_no_patch, x)
+        tasks_with_patch, _ = encode_tile_split(vae_with_patch, x)
+
+        # With patch_size=2, stride becomes 128//2=64, so more tiles
+        assert len(tasks_with_patch) >= len(tasks_no_patch)
+
+    def test_temporal_compression_from_config(self):
+        """Test that temporal compression ratio is read from config."""
+        encode_tile_split = _import_encode_tile_split()
+
+        # temporal_compression=4 (default)
+        vae_4x = _DummyWanVae(
+            config=_DummyConfig(scale_factor_temporal=4),
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=512,
+            tile_sample_stride_width=512,
+        )
+
+        # temporal_compression=2
+        vae_2x = _DummyWanVae(
+            config=_DummyConfig(scale_factor_temporal=2),
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=512,
+            tile_sample_stride_width=512,
+        )
+
+        # 9 frames input
+        x = torch.randn(1, 3, 9, 512, 512)
+
+        tasks_4x, _ = encode_tile_split(vae_4x, x)
+        tasks_2x, _ = encode_tile_split(vae_2x, x)
+
+        # With 4x compression: 1 + (9-1)//4 = 3 chunks
+        assert len(tasks_4x[0].tensor) == 3
+
+        # With 2x compression: 1 + (9-1)//2 = 5 chunks
+        assert len(tasks_2x[0].tensor) == 5
+
+    def test_grid_spec_latent_dimensions(self):
+        """Test that grid_spec contains correct latent dimensions."""
+        encode_tile_split = _import_encode_tile_split()
+
+        vae = _DummyWanVae(
+            config=_DummyConfig(patch_size=None),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=512,
+            tile_sample_stride_width=512,
+        )
+
+        # Input: 512x512 with compression 8 -> 64x64 latent
+        x = torch.randn(1, 3, 5, 512, 512)
+
+        _, grid_spec = encode_tile_split(vae, x)
+
+        assert grid_spec.tile_spec["latent_height"] == 64
+        assert grid_spec.tile_spec["latent_width"] == 64
+
+
+class TestEncodeTileExec:
+    """Tests for encode_tile_exec method."""
+
+    def test_basic_exec(self):
+        """Test basic tile execution."""
+        encode_tile_exec = _import_encode_tile_exec()
+
+        vae = _DummyWanVae()
+
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+            TileTask,
+        )
+
+        # Create a simple task with 2 temporal chunks
+        tile1 = torch.randn(1, 3, 1, 32, 32)
+        tile2 = torch.randn(1, 3, 4, 32, 32)
+        task = TileTask(tile_id=0, grid_coord=(0, 0), tensor=[tile1, tile2])
+
+        result = encode_tile_exec(vae, task)
+
+        # Result should concatenate temporal dimension
+        assert result.shape[2] == 5  # 1 + 4 frames
+
+
+class TestEncodeTileMerge:
+    """Tests for encode_tile_merge method."""
+
+    def test_basic_merge(self):
+        """Test basic tile merging."""
+        encode_tile_merge = _import_encode_tile_merge()
+
+        vae = _DummyWanVae()
+
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+            GridSpec,
+        )
+
+        # Create 2x2 grid of tiles
+        tile_00 = torch.ones(1, 16, 2, 32, 32) * 0
+        tile_01 = torch.ones(1, 16, 2, 32, 32) * 1
+        tile_10 = torch.ones(1, 16, 2, 32, 32) * 2
+        tile_11 = torch.ones(1, 16, 2, 32, 32) * 3
+
+        coord_tensor_map = {
+            (0, 0): tile_00,
+            (0, 1): tile_01,
+            (1, 0): tile_10,
+            (1, 1): tile_11,
+        }
+
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(2, 2),
+            tile_spec={
+                "latent_height": 48,
+                "latent_width": 48,
+                "blend_height": 8,
+                "blend_width": 8,
+                "tile_latent_stride_height": 24,
+                "tile_latent_stride_width": 24,
+            },
+        )
+
+        result = encode_tile_merge(vae, coord_tensor_map, grid_spec)
+
+        # Output should be (1, 16, 2, 48, 48)
+        assert result.shape == (1, 16, 2, 48, 48)

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -59,9 +59,9 @@ class E2EOperator:
 class DummyMixin(DistributedVaeMixin):
     def __init__(self):
         self.use_tiling = True
-        self.distributed_decoder = MagicMock()
-        self.distributed_decoder.parallel_size = 2
-        self.distributed_decoder.group = None
+        self.distributed_executor = MagicMock()
+        self.distributed_executor.parallel_size = 2
+        self.distributed_executor.group = None
 
 
 @pytest.fixture(autouse=True)

tests/diffusion/lora/conftest.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared test helpers for diffusion LoRA tests."""
+
+from __future__ import annotations
+
+import torch
+from vllm.model_executor.layers.linear import LinearBase
+
+
+class FakeLinearBase(LinearBase):
+    """Minimal LinearBase stub for LoRA layer discovery."""
+
+    def __init__(self):
+        torch.nn.Module.__init__(self)
+
+
+class DummyBaseLayerWithLoRA(torch.nn.Module):
+    """Fake LoRA wrapper that records set/reset/create calls."""
+
+    def __init__(self, base_layer: torch.nn.Module):
+        super().__init__()
+        self.base_layer = base_layer
+
+        self.set_calls: list[
+            tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
+        ] = []
+        self.reset_calls: int = 0
+        self.create_calls: int = 0
+
+    def set_lora(self, index: int, lora_a, lora_b):
+        assert index == 0
+        self.set_calls.append((lora_a, lora_b))
+
+    def reset_lora(self, index: int):
+        assert index == 0
+        self.reset_calls += 1
+
+    def create_lora_weights(self, max_loras, lora_config, model_config):
+        self.create_calls += 1
+
+
+def fake_replace_submodule(
+    root: torch.nn.Module,
+    module_name: str,
+    submodule: torch.nn.Module,
+    replace_calls: list[str] | None = None,
+) -> None:
+    """Replace a submodule by traversing dotted paths correctly."""
+    if replace_calls is not None:
+        replace_calls.append(module_name)
+    parts = module_name.split(".")
+    parent = root
+    for attr in parts[:-1]:
+        parent = getattr(parent, attr)
+    setattr(parent, parts[-1], submodule)

tests/diffusion/lora/test_lora_manager.py

@@ -7,8 +7,12 @@ import pytest
 import torch
 from vllm.lora.lora_weights import LoRALayerWeights
 from vllm.lora.utils import get_supported_lora_modules
-from vllm.model_executor.layers.linear import LinearBase
 
+from tests.diffusion.lora.conftest import (
+    DummyBaseLayerWithLoRA,
+    FakeLinearBase,
+    fake_replace_submodule,
+)
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.lora.request import LoRARequest
 
@@ -33,35 +37,9 @@ class _DummyLoRALayer:
         self.reset_calls += 1
 
 
-class _FakeLinearBase(LinearBase):
-    def __init__(self):
-        torch.nn.Module.__init__(self)
-
-
-class _DummyBaseLayerWithLoRA(torch.nn.Module):
-    def __init__(self, base_layer: torch.nn.Module):
-        super().__init__()
-        self.base_layer = base_layer
-
-        self.set_calls: list[
-            tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
-        ] = []
-        self.reset_calls: int = 0
-        self.create_calls: int = 0
-
-    def set_lora(self, index: int, lora_a, lora_b):
-        assert index == 0
-        self.set_calls.append((lora_a, lora_b))
-
-    def reset_lora(self, index: int):
-        assert index == 0
-        self.reset_calls += 1
-
-    def create_lora_weights(self, max_loras, lora_config, model_config):
-        # Needs to be callable for scale test when rank changes, but not
-        # actually used since we mock everything and check everything based
-        # on set calls.
-        self.create_calls += 1
+# Aliases for backward compatibility within this file
+_FakeLinearBase = FakeLinearBase
+_DummyBaseLayerWithLoRA = DummyBaseLayerWithLoRA
 
 
 class _DummyPipeline(torch.nn.Module):
@@ -555,3 +533,45 @@ def test_lora_manager_max_rank_validation(monkeypatch, rank):
     req1 = _dummy_lora_request(1)
     with pytest.raises(ValueError):
         manager.add_adapter(req1)
+
+
+def test_lora_manager_discovers_bagel_component(monkeypatch):
+    """Verify that _replace_layers_with_lora finds layers under 'bagel'."""
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
+
+    def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
+        if isinstance(layer, _FakeLinearBase):
+            return _DummyBaseLayerWithLoRA(layer)
+        return layer
+
+    replace_calls: list[str] = []
+
+    monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
+    monkeypatch.setattr(
+        manager_mod,
+        "replace_submodule",
+        lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls),
+    )
+
+    # Pipeline with a 'bagel' component (no 'transformer')
+    pipeline = torch.nn.Module()
+    pipeline.bagel = torch.nn.Module()
+    pipeline.bagel.language_model = torch.nn.Module()
+    pipeline.bagel.language_model.qkv_proj = _FakeLinearBase()
+
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    peft_helper = type("_PH", (), {"r": 1})()
+    manager._replace_layers_with_lora(peft_helper)
+
+    assert "language_model.qkv_proj" in replace_calls
+    assert "bagel.language_model.qkv_proj" in manager._lora_modules
+    # Verify the module was actually replaced in the tree (not just recorded)
+    assert isinstance(pipeline.bagel.language_model.qkv_proj, _DummyBaseLayerWithLoRA)

tests/diffusion/models/bagel/test_bagel_lora.py

@@ -0,0 +1,248 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for BAGEL LoRA support across Stage 0 (Thinker) and Stage 1 (DiT)."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from tests.diffusion.lora.conftest import (
+    DummyBaseLayerWithLoRA,
+    FakeLinearBase,
+    fake_replace_submodule,
+)
+from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
+from vllm_omni.lora.request import LoRARequest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_FakeLinearBase = FakeLinearBase
+
+
+# ---------------------------------------------------------------------------
+# Stage 0 (Thinker / AR) -- packed_modules_mapping on the AR model class
+# ---------------------------------------------------------------------------
+
+
+class TestStage0ThinkerLoRA:
+    """Validate that OmniBagelForConditionalGeneration declares correct LoRA metadata."""
+
+    def test_omni_bagel_supports_lora(self):
+        from vllm_omni.model_executor.models.bagel.bagel import (
+            OmniBagelForConditionalGeneration,
+        )
+
+        assert getattr(OmniBagelForConditionalGeneration, "supports_lora", False) is True
+
+    def test_omni_bagel_packed_modules_mapping_complete(self):
+        from vllm_omni.model_executor.models.bagel.bagel import (
+            OmniBagelForConditionalGeneration,
+        )
+
+        mapping = OmniBagelForConditionalGeneration.packed_modules_mapping
+        # Standard Qwen2 projections
+        assert mapping["qkv_proj"] == ["q_proj", "k_proj", "v_proj"]
+        assert mapping["gate_up_proj"] == ["gate_proj", "up_proj"]
+        # MoE generation-mode projections
+        assert mapping["qkv_proj_moe_gen"] == [
+            "q_proj_moe_gen",
+            "k_proj_moe_gen",
+            "v_proj_moe_gen",
+        ]
+        assert mapping["mlp_moe_gen.gate_up_proj"] == [
+            "mlp_moe_gen.gate_proj",
+            "mlp_moe_gen.up_proj",
+        ]
+
+
+# ---------------------------------------------------------------------------
+# Stage 1 (DiT / Diffusion) -- DiffusionLoRAManager with bagel component
+# ---------------------------------------------------------------------------
+
+
+class TestStage1DiTLoRA:
+    """Validate DiffusionLoRAManager discovers BAGEL's packed modules."""
+
+    def test_diffusion_lora_manager_discovers_bagel_packed_modules(self):
+        """Manager should derive packed→sublayer mapping from stacked_params_mapping."""
+        pipeline = torch.nn.Module()
+        pipeline.bagel = torch.nn.Module()
+
+        # Simulate a submodule that exposes stacked_params_mapping
+        # (as Bagel does after load_weights())
+        language_model = torch.nn.Module()
+        language_model.stacked_params_mapping = [
+            (".qkv_proj_moe_gen", ".q_proj_moe_gen", "q"),
+            (".qkv_proj_moe_gen", ".k_proj_moe_gen", "k"),
+            (".qkv_proj_moe_gen", ".v_proj_moe_gen", "v"),
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+        pipeline.bagel.language_model = language_model
+
+        manager = DiffusionLoRAManager(
+            pipeline=pipeline,
+            device=torch.device("cpu"),
+            dtype=torch.bfloat16,
+            max_cached_adapters=1,
+        )
+
+        mapping = manager._packed_modules_mapping
+        assert mapping["qkv_proj"] == ["q_proj", "k_proj", "v_proj"]
+        assert mapping["qkv_proj_moe_gen"] == [
+            "q_proj_moe_gen",
+            "k_proj_moe_gen",
+            "v_proj_moe_gen",
+        ]
+        assert mapping["gate_up_proj"] == ["gate_proj", "up_proj"]
+
+    def test_diffusion_lora_manager_replaces_bagel_packed_layer_via_sublayer_target(self, monkeypatch):
+        """Targeting sublayer 'q_proj' should replace the fused 'qkv_proj' under bagel."""
+        import vllm_omni.diffusion.lora.manager as manager_mod
+
+        monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", DummyBaseLayerWithLoRA)
+
+        def _fake_from_layer_diffusion(*, layer, **_kwargs):
+            return DummyBaseLayerWithLoRA(layer)
+
+        replace_calls: list[str] = []
+
+        monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
+        monkeypatch.setattr(
+            manager_mod,
+            "replace_submodule",
+            lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls),
+        )
+
+        # Build pipeline with bagel component
+        pipeline = torch.nn.Module()
+        pipeline.bagel = torch.nn.Module()
+        lm = torch.nn.Module()
+        lm.stacked_params_mapping = [
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+        ]
+        lm.attn = torch.nn.Module()
+        lm.attn.qkv_proj = _FakeLinearBase()
+        pipeline.bagel.language_model = lm
+
+        manager = DiffusionLoRAManager(
+            pipeline=pipeline,
+            device=torch.device("cpu"),
+            dtype=torch.bfloat16,
+            max_cached_adapters=1,
+        )
+
+        # Treat qkv_proj as 3-slice packed layer
+        monkeypatch.setattr(manager, "_get_packed_modules_list", lambda _module: ["q", "k", "v"])
+
+        # Target sublayer "q_proj" -- manager should replace the packed "qkv_proj"
+        peft_helper = type("_PH", (), {"r": 1, "target_modules": ["q_proj"]})()
+        manager._replace_layers_with_lora(peft_helper)
+
+        assert "language_model.attn.qkv_proj" in replace_calls
+        assert "bagel.language_model.attn.qkv_proj" in manager._lora_modules
+        # Verify the module was actually replaced in the tree (not just recorded)
+        assert isinstance(pipeline.bagel.language_model.attn.qkv_proj, DummyBaseLayerWithLoRA)
+
+
+# ---------------------------------------------------------------------------
+# Round-trip: synthetic checkpoint → set_active_adapter → verify weights
+# ---------------------------------------------------------------------------
+
+
+def _write_synthetic_lora(
+    adapter_dir: Path,
+    module_name: str,
+    rank: int,
+    in_dim: int,
+    out_dim: int,
+) -> str:
+    """Write a minimal LoRA adapter (safetensors + config) to *adapter_dir*."""
+    adapter_dir.mkdir(parents=True, exist_ok=True)
+    lora_a = torch.ones((rank, in_dim), dtype=torch.float32)
+    lora_b = torch.ones((out_dim, rank), dtype=torch.float32) * 2.0
+    save_file(
+        {
+            f"base_model.model.{module_name}.lora_A.weight": lora_a,
+            f"base_model.model.{module_name}.lora_B.weight": lora_b,
+        },
+        str(adapter_dir / "adapter_model.safetensors"),
+    )
+    (adapter_dir / "adapter_config.json").write_text(
+        json.dumps({"r": rank, "lora_alpha": rank, "target_modules": [module_name]}),
+        encoding="utf-8",
+    )
+    return str(adapter_dir)
+
+
+class TestBagelLoRARoundTrip:
+    """End-to-end: synthetic checkpoint → load → activate → verify weights in fused layer."""
+
+    def test_set_active_adapter_loads_and_activates_bagel_lora(self, tmp_path, monkeypatch):
+        """Full round-trip through set_active_adapter for a bagel component module."""
+        import vllm_omni.diffusion.lora.manager as manager_mod
+
+        monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", DummyBaseLayerWithLoRA)
+
+        # Build pipeline with bagel.language_model.foo (simple non-packed layer)
+        pipeline = torch.nn.Module()
+        pipeline.bagel = torch.nn.Module()
+        lm = torch.nn.Module()
+        lm.foo = _FakeLinearBase()
+        pipeline.bagel.language_model = lm
+
+        def _fake_from_layer(*, layer, **_kwargs):
+            if isinstance(layer, FakeLinearBase):
+                return DummyBaseLayerWithLoRA(layer)
+            return layer
+
+        monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer)
+        monkeypatch.setattr(
+            manager_mod,
+            "replace_submodule",
+            lambda root, name, sub: fake_replace_submodule(root, name, sub),
+        )
+
+        manager = DiffusionLoRAManager(
+            pipeline=pipeline,
+            device=torch.device("cpu"),
+            dtype=torch.bfloat16,
+            max_cached_adapters=1,
+        )
+
+        # Write synthetic adapter targeting bagel.language_model.foo
+        module_name = "bagel.language_model.foo"
+        rank = 2
+        in_dim = 4
+        out_dim = 4
+        lora_dir = _write_synthetic_lora(tmp_path / "lora", module_name, rank, in_dim, out_dim)
+
+        lora_request = LoRARequest(
+            lora_name="test_bagel",
+            lora_int_id=42,
+            lora_path=lora_dir,
+        )
+
+        # Full round-trip: load from disk → replace layer → activate weights
+        manager.set_active_adapter(lora_request, lora_scale=0.5)
+
+        # Verify the layer was replaced and weights were set
+        replaced_layer = pipeline.bagel.language_model.foo
+        assert isinstance(replaced_layer, DummyBaseLayerWithLoRA), "Layer should be wrapped with LoRA"
+        assert len(replaced_layer.set_calls) == 1, "set_lora should have been called once"
+
+        lora_a, lora_b = replaced_layer.set_calls[0]
+        # A weights should be ones (as written)
+        assert torch.all(lora_a == 1.0), f"lora_a should be all ones, got {lora_a}"
+        # B weights should be 2.0 * scale(0.5) = 1.0
+        assert torch.allclose(lora_b, torch.ones_like(lora_b)), f"lora_b should be 2.0 * 0.5 = 1.0, got {lora_b}"

tests/diffusion/offloader/test_layerwise_backend.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Unit tests for LayerwiseOffloadHook."""
+"""Unit tests for LayerwiseOffloadHook and LayerWiseOffloadBackend utilities."""
 
 import gc
 import os
@@ -15,7 +15,7 @@ from torch import nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module
-from vllm_omni.diffusion.offloader.layerwise_backend import LayerwiseOffloadHook
+from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend, LayerwiseOffloadHook
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -127,3 +127,116 @@ class TestLayerwiseOffloadHook:
         assert current_block.weight.to_local().is_meta
         assert current_block.weight.to_local().shape == torch.Size([4])
         assert not hook.is_materialized
+
+
+class _DummyBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(10, 10))
+
+
+class _SingleBlockModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["blocks"]
+
+    def __init__(self, num_blocks: int = 3):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class _MultiBlockModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
+
+    def __init__(self, num_transformer: int = 2, num_single: int = 2):
+        super().__init__()
+        self.transformer_blocks = nn.ModuleList([_DummyBlock() for _ in range(num_transformer)])
+        self.single_transformer_blocks = nn.ModuleList([_DummyBlock() for _ in range(num_single)])
+
+
+class _EmptyBlocksModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["blocks"]
+
+    def __init__(self):
+        super().__init__()
+        self.blocks = nn.ModuleList([])
+
+
+class _InvalidAttrModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["nonexistent_blocks", "blocks"]
+
+    def __init__(self, num_blocks: int = 2):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class _DeprecatedSingleAttrModel(nn.Module):
+    _layerwise_offload_blocks_attr = "blocks"
+
+    def __init__(self, num_blocks: int = 2):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class _NoAttrsModel(nn.Module):
+    def __init__(self, num_blocks: int = 2):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class TestGetBlocksFromDit:
+    def test_get_blocks_from_dit_single_block_attr(self):
+        model = _SingleBlockModel(num_blocks=3)
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert attr_names == ["blocks"]
+        assert len(blocks) == 3
+        assert all(isinstance(b, _DummyBlock) for b in blocks)
+
+    def test_get_blocks_from_dit_multi_block_attrs(self):
+        model = _MultiBlockModel(num_transformer=2, num_single=3)
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert set(attr_names) == {"transformer_blocks", "single_transformer_blocks"}
+        assert len(blocks) == 5
+        assert all(isinstance(b, _DummyBlock) for b in blocks)
+
+    def test_get_blocks_from_dit_empty_blocks(self):
+        model = _EmptyBlocksModel()
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert attr_names == []
+        assert blocks == []
+
+    def test_get_blocks_from_dit_invalid_attr_name(self):
+        model = _InvalidAttrModel(num_blocks=2)
+        with pytest.raises(
+            AttributeError,
+            match="Attribute 'nonexistent_blocks' declared in _layerwise_offload_blocks_attrs does not exist",
+        ):
+            LayerWiseOffloadBackend.get_blocks_from_dit(model)
+
+    def test_get_blocks_from_dit_no_attrs_defined(self):
+        model = _NoAttrsModel(num_blocks=3)
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert attr_names == []
+        assert blocks == []
+
+    def test_get_blocks_from_dit_deprecated_single_attr(self):
+        model = _DeprecatedSingleAttrModel(num_blocks=2)
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert attr_names == ["blocks"]
+        assert len(blocks) == 2
+
+
+class TestGetBlocksAttrNames:
+    def test_get_blocks_attr_names_new_format(self):
+        model = _MultiBlockModel()
+        attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model)
+        assert attrs == ["transformer_blocks", "single_transformer_blocks"]
+
+    def test_get_blocks_attr_names_no_attrs(self):
+        model = _NoAttrsModel()
+        attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model)
+        assert attrs == []
+
+    def test_set_blocks_attr_names(self):
+        model = _NoAttrsModel()
+        LayerWiseOffloadBackend.set_blocks_attr_names(model, ["new_blocks"])
+        assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs")
+        assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"]

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -133,6 +133,22 @@ def test_save_async(build_adapter):
     assert task["is_finished"] is False
 
 
+def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch):
+    adapter, _ = build_adapter(stage_id=1)
+    request = _req("req-finished", RequestStatus.FINISHED_STOPPED, external_req_id="ext-finished")
+
+    adapter.custom_process_next_stage_input_func = lambda **kwargs: {"x": [1], "finished": True}
+    cleanup_calls = []
+    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
+
+    adapter._send_single_request({"pooling_output": None, "request": request, "is_finished": True})
+
+    assert len(cleanup_calls) == 1
+    args, _ = cleanup_calls[0]
+    assert args[0] == "req-finished"
+    assert args[1] == "ext-finished"
+
+
 def test_update_request_payload(build_adapter):
     adapter, _ = build_adapter()
 
@@ -409,3 +425,86 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
     args, _ = cleanup_calls[0]
     assert args[0] == "req-s1"
     assert args[1] == "ext-s1"
+
+
+def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerFixture):
+    """OmniARScheduler should enqueue save; adapter cleanup is handled in save thread."""
+    cleanup_calls = []
+    save_calls = []
+
+    adapter_mock = mocker.MagicMock()
+    adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
+    adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw))
+
+    from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+    scheduler = mocker.MagicMock()
+    scheduler.chunk_transfer_adapter = adapter_mock
+    scheduler.connector = None
+    scheduler.perf_metrics = None
+    scheduler.log_stats = False
+    scheduler.recompute_kv_load_failures = False
+    scheduler.structured_output_manager = mocker.MagicMock()
+    scheduler.structured_output_manager.should_advance.return_value = False
+    scheduler.finished_req_ids_dict = {}
+    scheduler.kv_cache_manager = mocker.MagicMock()
+    scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.kv_event_publisher = mocker.MagicMock()
+    scheduler.waiting_for_transfer_free = set()
+    scheduler.transfer_triggered_requests = set()
+    scheduler.active_kv_transfers = set()
+
+    request = _HashableRequest(
+        request_id="req-ar",
+        external_req_id="ext-ar",
+        status=RequestStatus.RUNNING,
+        is_finished=lambda: False,
+        num_computed_tokens=1,
+        num_prompt_tokens=1,
+        prompt_token_ids=[1],
+        num_output_placeholders=0,
+        sampling_params=None,
+        pooling_params=None,
+        stop_reason=None,
+        client_index=0,
+        take_events=lambda: [],
+        trace_headers=None,
+        num_cached_tokens=0,
+        num_external_computed_tokens=0,
+        num_nans_in_logits=0,
+        get_finished_reason=lambda: "stop",
+    )
+    scheduler.requests = {"req-ar": request}
+
+    scheduler._update_request_with_output = mocker.MagicMock(return_value=([], True))
+    scheduler._process_kv_transfer_trigger = mocker.MagicMock(return_value=False)
+    scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
+    scheduler._free_request = mocker.MagicMock(return_value=None)
+    scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
+    scheduler.running = [request]
+    scheduler.waiting = mocker.MagicMock()
+    scheduler.waiting.remove_requests = mocker.MagicMock()
+    scheduler.make_spec_decoding_stats = mocker.MagicMock(return_value=None)
+    scheduler.make_stats = mocker.MagicMock(return_value=None)
+
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={"req-ar": 1},
+        scheduled_spec_decode_tokens={},
+        num_invalid_spec_tokens=0,
+    )
+    model_runner_output = SimpleNamespace(
+        sampled_token_ids=[[123]],
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=None,
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        cudagraph_stats=None,
+        req_id_to_index={"req-ar": 0},
+        kv_extracted_req_ids=None,
+    )
+
+    OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
+
+    assert len(cleanup_calls) == 0
+    assert len(save_calls) == 1

tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py

@@ -2,13 +2,20 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import json
+import threading
 
+import pytest
 import zmq
 
 from vllm_omni.distributed.omni_coordinator import (
     OmniCoordClientForStage,
     StageStatus,
 )
+from vllm_omni.distributed.omni_coordinator import (
+    omni_coord_client_for_stage as stage_client_module,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def _bind_router() -> tuple[zmq.Context, zmq.Socket, str]:
@@ -19,7 +26,8 @@ def _bind_router() -> tuple[zmq.Context, zmq.Socket, str]:
     return ctx, router, endpoint
 
 
-def _recv_event(router: zmq.Socket) -> dict:
+def _recv_event(router: zmq.Socket, timeout_ms: int = 2000) -> dict:
+    assert router.poll(timeout=timeout_ms) != 0, "Timed out waiting for coordinator event"
     frames = router.recv_multipart()
     # ROUTER adds identity frame; the last frame is the payload.
     payload = frames[-1]
@@ -108,3 +116,197 @@ def test_stage_client_close_sends_down_status():
 
     router.close(0)
     ctx.term()
+
+
+def test_stage_client_reconnects_after_send_failure(mocker):
+    """Verify send failure path invokes reconnect before retrying send."""
+    ctx, router, endpoint = _bind_router()
+
+    client = OmniCoordClientForStage(
+        endpoint,
+        "tcp://stage:reconnect-in",
+        "tcp://stage:reconnect-out",
+        0,
+    )
+
+    # Discard initial registration event from the real socket.
+    _recv_event(router)
+
+    class _FlakySocket:
+        def __init__(self):
+            self.send_calls = 0
+            self.closed = False
+
+        def send(self, *_args, **_kwargs):
+            self.send_calls += 1
+            if self.send_calls == 1:
+                raise RuntimeError("simulated send failure")
+
+        def close(self, *_args, **_kwargs):
+            self.closed = True
+
+    flaky_socket = _FlakySocket()
+    client._socket = flaky_socket
+    client._reconnect = mocker.Mock(return_value=True)
+
+    client.update_info(queue_length=1)
+
+    client._reconnect.assert_called_once_with(max_retries=3)
+    assert flaky_socket.send_calls == 2
+
+    client.close()
+    router.close(0)
+    ctx.term()
+
+
+def test_stage_client_raises_when_reconnect_fails(mocker):
+    """Verify send failure is propagated when reconnect cannot recover."""
+    ctx, router, endpoint = _bind_router()
+
+    client = OmniCoordClientForStage(
+        endpoint,
+        "tcp://stage:reconnect-fail-in",
+        "tcp://stage:reconnect-fail-out",
+        0,
+    )
+
+    # Discard initial registration event from the real socket.
+    _recv_event(router)
+
+    class _AlwaysFailSocket:
+        def send(self, *_args, **_kwargs):
+            raise RuntimeError("simulated send failure")
+
+        def close(self, *_args, **_kwargs):
+            pass
+
+    client._socket = _AlwaysFailSocket()
+    client._reconnect = mocker.Mock(return_value=False)
+
+    with pytest.raises(RuntimeError, match="simulated send failure"):
+        client.update_info(queue_length=2)
+
+    client._reconnect.assert_called_once_with(max_retries=3)
+    client.close()
+    router.close(0)
+    ctx.term()
+
+
+def test_stage_client_close_handles_runtime_error_in_final_update(mocker):
+    """Verify close() still releases resources when final update raises RuntimeError."""
+    ctx, router, endpoint = _bind_router()
+
+    client = OmniCoordClientForStage(
+        endpoint,
+        "tcp://stage:close-runtime-in",
+        "tcp://stage:close-runtime-out",
+        0,
+    )
+
+    # Discard initial registration event from the real socket.
+    _recv_event(router)
+
+    client._send_event = mocker.Mock(side_effect=RuntimeError("simulated close-time failure"))
+    client.close()
+
+    assert client._closed
+    assert client._socket.closed
+
+    router.close(0)
+    ctx.term()
+
+
+def test_reconnect_respects_retry_limit(monkeypatch):
+    """Verify _reconnect stops after max_retries on repeated failures."""
+    attempts = {"connect": 0}
+
+    class _FailSocket:
+        def close(self, *_args, **_kwargs):
+            pass
+
+        def connect(self, *_args, **_kwargs):
+            attempts["connect"] += 1
+            raise zmq.ZMQError("simulated reconnect failure")
+
+    class _FailContext:
+        def socket(self, *_args, **_kwargs):
+            return _FailSocket()
+
+        def term(self):
+            pass
+
+    client = OmniCoordClientForStage.__new__(OmniCoordClientForStage)
+    client._closed = False
+    client._coord_zmq_addr = "tcp://127.0.0.1:9999"
+    client._stop_event = threading.Event()
+    client._send_lock = threading.RLock()
+    client._socket = _FailSocket()
+    client._ctx = _FailContext()
+
+    monkeypatch.setattr(stage_client_module.zmq, "Context", lambda: _FailContext())
+    monkeypatch.setattr(stage_client_module.time, "sleep", lambda *_args, **_kwargs: None)
+
+    assert client._reconnect(max_retries=3, retry_interval=5.0) is False
+    assert attempts["connect"] == 3
+
+
+def test_heartbeat_loop_retries_after_transient_send_failure():
+    """Verify heartbeat loop continues after one transient send failure."""
+
+    class _FakeStopEvent:
+        def __init__(self):
+            self.wait_calls = 0
+            self._set = False
+
+        def wait(self, timeout=None):
+            _ = timeout
+            self.wait_calls += 1
+            # Run two loop iterations, then stop.
+            return self._set or self.wait_calls >= 3
+
+        def is_set(self):
+            return self._set
+
+        def set(self):
+            self._set = True
+
+    client = OmniCoordClientForStage.__new__(OmniCoordClientForStage)
+    client._closed = False
+    client._heartbeat_interval = 0.0
+    client._stop_event = _FakeStopEvent()
+
+    calls = {"count": 0}
+
+    def _fake_send(event_type):
+        assert event_type == "heartbeat"
+        calls["count"] += 1
+        if calls["count"] == 1:
+            raise RuntimeError("transient heartbeat failure")
+
+    client._send_event = _fake_send
+
+    client._heartbeat_loop()
+
+    assert calls["count"] == 2
+
+
+def test_update_info_rejected_while_closing():
+    """Verify update_info is rejected once client enters closing state."""
+    ctx, router, endpoint = _bind_router()
+
+    client = OmniCoordClientForStage(
+        endpoint,
+        "tcp://stage:closing-in",
+        "tcp://stage:closing-out",
+        0,
+    )
+    _recv_event(router)
+
+    client._closing = True
+    with pytest.raises(RuntimeError, match="closing"):
+        client.update_info(queue_length=3)
+
+    client._closing = False
+    client.close()
+    router.close(0)
+    ctx.term()

tests/e2e/offline_inference/test_bagel_understanding.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end tests for Bagel text2text and img2text (understanding) tasks.
+
+These tests validate that the Bagel multistage pipeline correctly generates
+text output for understanding tasks, matching reference results.
+
+Equivalent to running:
+    python3 examples/offline_inference/bagel/end2end.py \
+        --modality text2text \
+        --prompts "Where is the capital of France?"
+
+    python3 examples/offline_inference/bagel/end2end.py \
+        --modality img2text \
+        --prompts "Please describe this image" \
+        --image-path 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+from pathlib import Path
+
+import pytest
+from vllm.assets.image import ImageAsset
+
+from tests.conftest import modify_stage_config
+from tests.utils import hardware_test
+from vllm_omni.entrypoints.omni import Omni
+
+MODEL_NAME = "ByteDance-Seed/BAGEL-7B-MoT"
+STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
+
+REFERENCE_TEXT_TEXT2TEXT = "The capital of France is Paris."
+
+REFERENCE_TEXT_IMG2TEXT = (
+    "This is a photo of a wooden boardwalk or pathway that leads through "
+    "tall green grass. The path appears to be in a natural setting, possibly "
+    "a wetland or marsh area. The sky above is blue with some scattered "
+    "clouds, suggesting it might be a sunny day. The overall scene looks "
+    "peaceful and serene."
+)
+
+
+def _resolve_stage_config(config_path: str, run_level: str) -> str:
+    """Strip load_format: dummy for advanced_model (real weights)."""
+    if run_level == "advanced_model":
+        return modify_stage_config(
+            config_path,
+            deletes={
+                "stage_args": {
+                    0: ["engine_args.load_format"],
+                    1: ["engine_args.load_format"],
+                }
+            },
+        )
+    return config_path
+
+
+def _extract_text(omni_outputs: list) -> str:
+    """Extract generated text from OmniRequestOutput list."""
+    for req_output in omni_outputs:
+        ro = getattr(req_output, "request_output", None)
+        if ro and getattr(ro, "outputs", None):
+            return "".join(getattr(o, "text", "") or "" for o in ro.outputs)
+    return ""
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
+def test_bagel_text2text(run_level):
+    """Test Bagel text2text produces correct text output."""
+    config_path = _resolve_stage_config(STAGE_CONFIG, run_level)
+    omni = Omni(
+        model=MODEL_NAME,
+        stage_configs_path=config_path,
+        stage_init_timeout=300,
+    )
+
+    try:
+        prompt = "<|im_start|>user\nWhere is the capital of France?<|im_end|>\n<|im_start|>assistant\n"
+        params_list = omni.default_sampling_params_list
+        omni_outputs = list(
+            omni.generate(
+                prompts=[{"prompt": prompt, "modalities": ["text"]}],
+                sampling_params_list=params_list,
+            )
+        )
+
+        assert len(omni_outputs) > 0, "No outputs returned"
+        text = _extract_text(omni_outputs)
+        assert len(text) > 0, "Generated text is empty"
+
+        if run_level == "advanced_model":
+            assert text == REFERENCE_TEXT_TEXT2TEXT, (
+                f"Text mismatch: expected {REFERENCE_TEXT_TEXT2TEXT!r}, got {text!r}"
+            )
+    finally:
+        omni.close()
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
+def test_bagel_img2text(run_level):
+    """Test Bagel img2text produces correct text output."""
+    input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
+    config_path = _resolve_stage_config(STAGE_CONFIG, run_level)
+    omni = Omni(
+        model=MODEL_NAME,
+        stage_configs_path=config_path,
+        stage_init_timeout=300,
+    )
+
+    try:
+        prompt = "<|im_start|>user\n<|image_pad|>\nPlease describe this image<|im_end|>\n<|im_start|>assistant\n"
+        params_list = omni.default_sampling_params_list
+        omni_outputs = list(
+            omni.generate(
+                prompts=[
+                    {
+                        "prompt": prompt,
+                        "multi_modal_data": {"image": input_image},
+                        "modalities": ["text"],
+                    }
+                ],
+                sampling_params_list=params_list,
+            )
+        )
+
+        assert len(omni_outputs) > 0, "No outputs returned"
+        text = _extract_text(omni_outputs)
+        assert len(text) > 0, "Generated text is empty"
+
+        if run_level == "advanced_model":
+            assert text == REFERENCE_TEXT_IMG2TEXT, f"Text mismatch: expected {REFERENCE_TEXT_IMG2TEXT!r}, got {text!r}"
+    finally:
+        omni.close()

tests/e2e/offline_inference/test_cosyvoice3.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Offline E2E smoke test for CosyVoice3 zero-shot reference inference.
+
+This test uses the official upstream zero-shot prompt text/audio pair and
+verifies a stable reference recipe:
+- config-derived top_p/top_k and token-length ratios
+- model EOS token as the stop token
+- a conservative repetition penalty to avoid degenerate loops
+"""
+
+from __future__ import annotations
+
+import functools
+import io
+import os
+import tempfile
+from pathlib import Path
+from urllib.request import urlopen
+
+import numpy as np
+import pytest
+import soundfile as sf
+import yaml
+from huggingface_hub import snapshot_download
+from vllm.sampling_params import SamplingParams
+
+from tests.conftest import OmniRunner
+from tests.utils import hardware_test
+from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
+from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
+
+MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
+MODEL_DIR_ENV = "VLLM_OMNI_COSYVOICE3_MODEL_DIR"
+
+REFERENCE_PROMPT_WAV_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
+REFERENCE_PROMPT_TEXT = "You are a helpful assistant.<|endofprompt|>希望你以后能够做的比我还好呦。"
+REFERENCE_SYNTH_TEXT = (
+    "CosyVoice is undergoing a comprehensive upgrade, providing more accurate, "
+    "stable, faster, and better voice generation capabilities."
+)
+REFERENCE_STAGE0_TEMPERATURE = 1.0
+REFERENCE_STAGE0_REPETITION_PENALTY = 2.0
+
+
+def _stage_config(name: str) -> str:
+    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
+
+
+STAGE_CONFIGS = [
+    _stage_config("cosyvoice3.yaml"),
+    _stage_config("cosyvoice3_async_chunk.yaml"),
+]
+
+
+@functools.lru_cache(maxsize=1)
+def _load_reference_prompt_wav() -> tuple[np.ndarray, int]:
+    with urlopen(REFERENCE_PROMPT_WAV_URL, timeout=30) as resp:
+        data = resp.read()
+    audio, sr = sf.read(io.BytesIO(data), dtype="float32", always_2d=False)
+    if isinstance(audio, np.ndarray) and audio.ndim > 1:
+        audio = np.mean(audio, axis=-1)
+    return np.asarray(audio, dtype=np.float32), int(sr)
+
+
+@functools.lru_cache(maxsize=1)
+def _resolve_model_dir() -> Path:
+    override = os.environ.get(MODEL_DIR_ENV)
+    if override:
+        return Path(override).expanduser().resolve()
+    return Path(snapshot_download(MODEL, allow_patterns=["*"]))
+
+
+def _reference_zero_shot_stage0_sampling(*, text: str) -> SamplingParams:
+    config = CosyVoice3Config()
+    sampling_cfg = config.llm.get("sampling", {})
+    eos_token_id = int(config.llm["eos_token_id"])
+    model_dir = _resolve_model_dir()
+    tokenizer = get_qwen_tokenizer(
+        token_path=str(model_dir / config.qwen_pretrain_path),
+        skip_special_tokens=config.skip_special_tokens,
+        version=config.version,
+    )
+    text_len = max(1, len(tokenizer.encode(text, allowed_special=config.allowed_special)))
+    return SamplingParams(
+        temperature=REFERENCE_STAGE0_TEMPERATURE,
+        top_p=float(sampling_cfg.get("top_p", 0.8)),
+        top_k=int(sampling_cfg.get("top_k", 25)),
+        repetition_penalty=REFERENCE_STAGE0_REPETITION_PENALTY,
+        stop_token_ids=[eos_token_id],
+        min_tokens=int(text_len * config.min_token_text_ratio),
+        max_tokens=int(text_len * config.max_token_text_ratio),
+    )
+
+
+def _concat_audio(audio_val) -> np.ndarray:
+    import torch
+
+    if isinstance(audio_val, list):
+        tensors = []
+        for t in audio_val:
+            if t is None:
+                continue
+            if hasattr(t, "detach"):
+                t = t.detach()
+            if hasattr(t, "cpu"):
+                t = t.cpu()
+            if hasattr(t, "float"):
+                t = t.float()
+            if isinstance(t, torch.Tensor):
+                tensors.append(t.reshape(-1))
+        if not tensors:
+            return np.zeros((0,), dtype=np.float32)
+        return torch.cat(tensors, dim=-1).numpy().astype(np.float32, copy=False)
+
+    if hasattr(audio_val, "detach"):
+        audio_val = audio_val.detach()
+    if hasattr(audio_val, "cpu"):
+        audio_val = audio_val.cpu()
+    if hasattr(audio_val, "float"):
+        audio_val = audio_val.float()
+    if hasattr(audio_val, "numpy"):
+        audio_val = audio_val.numpy()
+    audio_np = np.asarray(audio_val, dtype=np.float32)
+    return audio_np.reshape(-1)
+
+
+def _get_stage_engine_outputs(omni_runner: OmniRunner, stage_id: int):
+    stage_list = getattr(omni_runner.omni, "stage_list", None)
+    if stage_list is not None:
+        return getattr(stage_list[stage_id], "engine_outputs", None) or []
+
+    stage_clients = getattr(getattr(omni_runner.omni, "engine", None), "stage_clients", None)
+    if stage_clients is not None:
+        return getattr(stage_clients[stage_id], "engine_outputs", None) or []
+
+    raise AttributeError("Unable to locate stage outputs on Omni runner")
+
+
+def _patched_stage_config(base_stage_config: str, model_dir: Path, tmp_dir: Path) -> str:
+    cfg = yaml.safe_load(Path(base_stage_config).read_text(encoding="utf-8"))
+    tokenizer_path = str(model_dir / "CosyVoice-BlankEN")
+    for stage in cfg.get("stage_args", []):
+        engine_args = stage.setdefault("engine_args", {})
+        engine_args["tokenizer"] = tokenizer_path
+        engine_args["enforce_eager"] = True
+        engine_args["hf_overrides"] = {"architectures": ["CosyVoice3Model"]}
+    out_path = tmp_dir / Path(base_stage_config).name
+    out_path.write_text(yaml.safe_dump(cfg, sort_keys=False), encoding="utf-8")
+    return str(out_path)
+
+
+def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[str, object]]:
+    return [
+        {
+            "prompt": REFERENCE_SYNTH_TEXT,
+            "multi_modal_data": {"audio": prompt_audio},
+            "modalities": ["audio"],
+            "mm_processor_kwargs": {"prompt_text": REFERENCE_PROMPT_TEXT},
+        }
+    ]
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("base_stage_config", STAGE_CONFIGS)
+def test_cosyvoice3_offline_reference_zero_shot(base_stage_config: str) -> None:
+    """CosyVoice3 zero-shot reference inference should stop cleanly and produce sane audio."""
+    prompt_audio, prompt_sr = _load_reference_prompt_wav()
+    model_dir = _resolve_model_dir()
+    expected_stop_token = int(CosyVoice3Config().llm["eos_token_id"])
+
+    with tempfile.TemporaryDirectory(prefix="cv3-e2e-") as tmp:
+        stage_config = _patched_stage_config(base_stage_config, model_dir, Path(tmp))
+        with OmniRunner(
+            str(model_dir), seed=42, stage_configs_path=stage_config, stage_init_timeout=300
+        ) as omni_runner:
+            sampling_params_list = omni_runner.get_default_sampling_params_list()
+            sampling_params_list[0] = _reference_zero_shot_stage0_sampling(text=REFERENCE_SYNTH_TEXT)
+
+            outputs = omni_runner.omni.generate(
+                _build_reference_inputs((prompt_audio, prompt_sr)), sampling_params_list
+            )
+
+            assert outputs, "No outputs returned"
+            audio_mm = outputs[0].multimodal_output
+            assert "audio" in audio_mm, "No audio output found"
+
+            audio = _concat_audio(audio_mm["audio"])
+            assert audio.size > 0, "Generated audio is empty"
+
+            sr_val = audio_mm.get("sr", 24000)
+            if isinstance(sr_val, list) and sr_val:
+                sr_val = sr_val[-1]
+            if hasattr(sr_val, "item"):
+                sr_val = sr_val.item()
+            sr = int(sr_val)
+            assert sr == 24000, f"Unexpected sample_rate={sr}"
+
+            duration_s = audio.size / sr
+            assert 2.8 <= duration_s <= 8.8, f"Unexpected duration={duration_s:.3f}s (samples={audio.size}, sr={sr})"
+
+            stage0_outputs = _get_stage_engine_outputs(omni_runner, 0)
+            if stage0_outputs:
+                completion = stage0_outputs[0].outputs[0]
+                finish_reason = getattr(completion, "finish_reason", None)
+                stop_reason = getattr(completion, "stop_reason", None)
+                num_tokens = len(getattr(completion, "token_ids", []) or [])
+
+                assert finish_reason == "stop", f"Stage-0 finish_reason={finish_reason}, expected 'stop'"
+                assert int(stop_reason) == expected_stop_token, (
+                    f"Stage-0 stop_reason={stop_reason}, expected {expected_stop_token}"
+                )
+                assert 80 <= num_tokens <= 220, f"Stage-0 num_tokens={num_tokens}, expected sane stop-bound range"
+            else:
+                assert "async_chunk" in Path(base_stage_config).name, "Stage-0 produced no engine outputs"

tests/e2e/offline_inference/test_hunyuanimage3_text2img.py

@@ -0,0 +1,347 @@
+# ruff: noqa: E501
+from collections.abc import Generator
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from transformers import CLIPModel, CLIPProcessor
+
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+PROMPT = "A brown and white dog is running on the grass"
+MODEL_NAME = "tencent/HunyuanImage-3.0"
+LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32"
+REPO_ROOT = Path(__file__).resolve().parents[3]
+STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image_3_moe.yaml"
+
+pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion]
+
+# System prompt type. Options: None, dynamic, en_vanilla, en_recaption, en_think_recaption, en_unified
+# Below are the CLIP embedding tensors from the official HunyuanImage model (seed=1234, prompt: "A brown and white dog is running on the grass").
+# SEED_1234 denotes the output without system prompt, while the remaining entries correspond to outputs generated with different system prompts.
+# fmt: off
+SEED_1234 = torch.tensor(
+    [
+        0.027797, 0.028964, -0.005051, 0.001059, 0.017021, -0.034029, 0.021989, 0.033318, -0.000308, 0.016179, 0.010504, -0.034201, 0.050230, -0.021170, 0.083530, -0.003621,
+        0.040758, 0.039913, 0.044305, -0.019285, -0.058387, -0.001099, 0.042782, -0.036136, -0.014955, 0.002147, 0.009439, 0.012943, -0.028732, -0.018349, 0.002861, 0.013019,
+        0.014362, -0.038833, 0.029413, 0.020724, 0.002714, 0.010416, -0.020527, 0.050266, -0.081026, -0.006814, -0.007457, -0.032333, 0.008417, -0.122455, -0.006085, -0.025610,
+        0.012614, 0.025817, -0.005419, 0.038657, 0.000789, 0.067111, 0.002818, 0.028696, 0.047305, -0.009993, -0.019508, 0.038604, 0.099657, 0.026728, 0.012361, 0.013626,
+        0.023164, -0.037186, 0.007535, 0.054645, -0.009012, -0.019383, -0.005234, -0.018715, -0.000346, 0.051317, -0.028744, 0.029933, -0.006382, -0.018414, -0.033906, -0.028892,
+        -0.015301, -0.004276, 0.014626, -0.008505, 0.013717, -0.027323, -0.001332, -0.040227, 0.047021, -0.019082, -0.037260, -0.029780, -0.594026, 0.016573, -0.010523, 0.042616,
+        -0.013136, 0.030540, -0.151685, -0.005367, 0.016209, -0.034183, 0.009852, 0.038452, 0.005494, -0.017887, -0.007167, 0.017262, -0.038980, 0.011995, 0.021952, -0.031660,
+        0.020507, -0.035880, 0.035183, -0.026975, -0.050788, -0.002553, 0.037774, -0.020082, -0.015403, 0.045022, 0.072167, -0.029237, 0.003895, -0.051250, 0.008581, 0.023545,
+        -0.026827, 0.020895, 0.041780, -0.040766, -0.008146, 0.080630, 0.000404, 0.032003, -0.005279, -0.090707, -0.013813, 0.010204, -0.001513, 0.016394, -0.001321, 0.020535,
+        -0.038645, 0.024858, 0.024378, 0.018717, -0.056314, 0.024402, 0.018694, 0.029009, -0.008502, -0.014694, -0.028345, 0.005202, 0.046116, -0.032166, -0.030706, -0.038738,
+        -0.031356, -0.009683, 0.040069, 0.001596, -0.012621, 0.018590, -0.024138, 0.035330, 0.011546, 0.015791, -0.026932, 0.004531, 0.022455, -0.012871, 0.013915, -0.009567,
+        -0.010976, 0.013497, 0.042590, 0.002072, -0.052718, -0.045494, 0.013036, -0.005403, -0.005947, -0.003437, 0.016653, -0.016805, -0.040291, 0.007927, 0.001296, -0.008319,
+        0.021514, -0.001452, -0.121998, 0.015396, -0.022594, -0.006977, -0.040108, -0.035550, -0.021872, -0.014721, 0.019799, 0.036556, 0.015072, -0.057988, -0.011684, -0.045220,
+        -0.026295, 0.052647, 0.013741, -0.013428, 0.061794, 0.021431, -0.011316, -0.009963, 0.008198, 0.027746, 0.074219, -0.019499, 0.042673, 0.016028, 0.007214, -0.010650,
+        -0.019682, 0.001902, 0.038867, -0.007333, 0.031749, 0.004391, 0.018688, 0.044654, 0.030615, -0.027816, 0.031711, -0.056952, -0.033499, -0.039368, 0.025801, -0.027610,
+        -0.009329, -0.001799, 0.024061, -0.012593, -0.050266, -0.012512, 0.019528, -0.083434, 0.018238, 0.034138, -0.020120, -0.009910, -0.002280, 0.035325, 0.034440, -0.055205,
+        -0.017698, -0.000439, -0.034703, 0.013356, -0.037287, 0.048494, -0.018570, 0.028069, 0.019269, -0.007263, -0.008521, 0.000426, -0.016677, 0.056162, -0.011944, 0.017322,
+        0.022219, -0.014266, -0.009292, -0.009979, 0.014973, 0.011623, -0.017799, 0.032925, -0.024668, 0.007312, -0.025035, -0.008967, -0.026827, 0.011889, -0.138517, -0.009608,
+        -0.020592, -0.001272, 0.015676, -0.025706, 0.031775, -0.004195, 0.026876, -0.014748, -0.025966, -0.008741, 0.035437, 0.017139, -0.005140, -0.007101, -0.012510, -0.023600,
+        0.032969, -0.005510, 0.020010, 0.032567, 0.015558, 0.004265, -0.036300, 0.048210, 0.080424, -0.052820, -0.002063, -0.020875, 0.052530, -0.001638, -0.020299, -0.035202,
+        0.087818, 0.034614, -0.032735, 0.033201, -0.001751, 0.029574, 0.009926, 0.011619, -0.001267, -0.020149, -0.003826, -0.029860, 0.011437, -0.051276, 0.024344, 0.003096,
+        -0.011573, 0.038228, -0.005730, -0.052328, 0.001909, -0.025877, 0.019976, -0.010160, 0.023892, 0.049161, -0.028978, 0.018700, -0.026460, 0.001090, -0.072128, -0.008406,
+        0.010828, 0.020621, -0.005706, 0.023797, 0.036231, -0.112069, 0.017601, 0.007496, 0.045999, 0.016771, 0.021977, 0.022305, 0.018377, 0.002036, -0.029815, -0.082922,
+        -0.012710, -0.026355, 0.003790, 0.017472, -0.023148, -0.002901, -0.057854, 0.028393, 0.230866, -0.023486, 0.051094, 0.047508, 0.018957, -0.037130, 0.001054, -0.026126,
+        0.021970, -0.046915, -0.019419, -0.014077, 0.002502, -0.079454, -0.057149, -0.081701, 0.041979, -0.043074, -0.009425, -0.035776, -0.021794, -0.004826, -0.057263, -0.072940,
+        0.037651, -0.013991, -0.043863, -0.020581, 0.034319, -0.052566, -0.010355, -0.022963, 0.027144, -0.017339, 0.088930, -0.000670, -0.026547, -0.026586, -0.032531, 0.040314,
+        0.010148, 0.021104, 0.009228, -0.073227, 0.036650, -0.019337, 0.010211, -0.089620, -0.024676, -0.020729, -0.004070, 0.000784, -0.110561, 0.015390, 0.027151, -0.003228,
+        -0.066704, -0.004797, -0.026117, -0.018131, -0.090114, 0.020659, -0.007157, 0.013608, -0.022324, 0.027487, 0.018873, 0.027854, 0.045085, -0.039992, -0.017829, 0.011071,
+        -0.011393, -0.004454, -0.037189, -0.030299, 0.059668, 0.005064, 0.024655, -0.037239, 0.046882, -0.010356, -0.009690, 0.061909, -0.024736, 0.016849, 0.000784, 0.000201,
+        0.066165, 0.010234, -0.012134, -0.002823, -0.060847, 0.008953, 0.010348, 0.022292, -0.044602, -0.020981, 0.038839, 0.006616, -0.016836, -0.043995, -0.005463, -0.036413,
+        0.034895, -0.018008, -0.009543, -0.025080, -0.035243, 0.042696, -0.028911, -0.030676, -0.038542, -0.027798, -0.026607, 0.019467, 0.070629, -0.037356, -0.042648, -0.000284,
+        0.033095, 0.077781, -0.052930, 0.022515, -0.029926, -0.033821, -0.003277, -0.000038, -0.026871, 0.018223, -0.004221, 0.023454, -0.030611, -0.006396, -0.009873, -0.008402,
+    ],
+    dtype=torch.float32,
+)
+SYSTEM_PROMPT_DYNAMIC = torch.tensor(
+    [
+        0.010809, 0.021177, -0.017600, -0.016814, 0.012351, -0.024554, 0.018299, 0.039305, 0.003331, 0.030473, 0.005557, -0.040898, 0.047294, -0.016136, 0.076989, -0.002723,
+        0.017622, 0.042330, 0.058266, -0.016232, -0.029502, 0.004529, 0.033543, -0.041481, -0.017631, 0.002727, 0.018874, 0.019932, -0.030052, -0.009997, 0.004582, 0.002135,
+        -0.003720, -0.030923, 0.021174, 0.034033, -0.007096, 0.011522, -0.009518, 0.055688, -0.092351, -0.003914, 0.004589, -0.032635, 0.012479, -0.140607, -0.014141, -0.031821,
+        0.001396, 0.026780, -0.007623, 0.039957, 0.006434, 0.047516, 0.014377, 0.015237, 0.034212, 0.003576, -0.027357, 0.038888, 0.087272, 0.020248, 0.015165, 0.016002,
+        0.020781, -0.040509, -0.008929, 0.080857, -0.002642, -0.009738, -0.005683, -0.000615, -0.012801, 0.046457, -0.045004, 0.024689, 0.002498, -0.017333, -0.027366, -0.023231,
+        -0.006064, -0.021505, 0.007405, -0.021249, 0.026252, -0.018690, 0.020093, -0.036954, 0.037510, -0.032027, -0.030871, -0.011173, -0.618627, 0.021213, -0.004366, 0.029555,
+        -0.004324, 0.020221, -0.143832, -0.021386, 0.010482, -0.042113, 0.016164, 0.040350, 0.014627, -0.011778, -0.018102, 0.035380, -0.020305, 0.010590, 0.009227, -0.011415,
+        0.018623, -0.036384, 0.031003, -0.017073, -0.056456, -0.010423, 0.033029, -0.023511, -0.008717, 0.045716, 0.068273, -0.027886, 0.009665, -0.039801, 0.001465, 0.024361,
+        -0.015039, 0.022903, 0.033362, -0.022804, 0.008631, 0.076518, 0.000619, 0.022786, -0.015435, -0.095242, -0.006092, 0.015496, -0.009081, 0.015740, 0.004280, 0.013103,
+        -0.031836, 0.034241, 0.031836, 0.032636, -0.053721, 0.034370, 0.019172, 0.018383, 0.006907, -0.036039, -0.027927, 0.008646, 0.040496, -0.060314, -0.039116, -0.021488,
+        -0.031682, -0.005077, 0.034920, 0.002148, -0.008087, 0.002024, -0.008480, 0.041096, 0.011401, 0.020380, -0.025078, 0.005002, 0.022252, -0.014577, 0.008051, -0.014476,
+        -0.007078, 0.021075, 0.036965, 0.005343, -0.038671, -0.037222, 0.014052, -0.009952, -0.003958, -0.001878, 0.017848, -0.016608, -0.030813, 0.010921, 0.001068, 0.003095,
+        0.007076, -0.001936, -0.102996, 0.006838, -0.005243, -0.009140, -0.043796, -0.027227, -0.008426, -0.013177, 0.015602, 0.021036, 0.025484, -0.064836, -0.003593, -0.038036,
+        -0.023102, 0.064053, 0.007850, 0.000771, 0.039297, 0.011903, -0.015866, -0.017612, 0.006308, 0.024342, 0.086761, -0.016705, 0.039239, 0.025079, -0.006452, 0.003174,
+        -0.010146, 0.010787, 0.035932, -0.015346, 0.037191, 0.010990, 0.011573, 0.044958, 0.035560, -0.017339, 0.018878, -0.025394, -0.044339, -0.029852, 0.015951, -0.032248,
+        -0.012019, 0.013497, 0.012224, -0.001284, -0.034041, -0.015768, 0.000230, -0.086076, 0.024878, 0.031929, -0.016668, -0.019815, -0.001325, 0.007944, 0.017674, -0.036097,
+        -0.019651, -0.001272, -0.032842, 0.002056, -0.037140, 0.043191, -0.003710, 0.011767, 0.020313, -0.018396, -0.015935, 0.010228, -0.017349, 0.049363, -0.010007, 0.019533,
+        0.018076, 0.016608, -0.005523, -0.007793, 0.016868, 0.019341, -0.008236, 0.026765, -0.025324, -0.007849, -0.023648, -0.007791, -0.018508, 0.015357, -0.166499, -0.003718,
+        -0.035447, -0.005229, 0.019327, -0.014207, 0.028433, -0.002619, 0.013888, -0.033146, -0.017015, 0.004677, 0.039554, 0.003803, -0.014592, -0.018886, -0.023868, -0.022708,
+        0.033661, 0.008626, 0.015687, 0.046395, 0.014173, 0.015083, -0.025994, 0.039120, 0.076334, -0.061165, 0.001791, -0.017579, 0.067567, -0.002415, -0.032495, -0.025576,
+        0.079027, 0.036370, -0.013303, 0.030510, -0.009061, 0.019135, 0.015627, 0.024864, 0.015093, -0.017066, -0.014075, -0.021907, 0.017388, -0.033492, 0.013317, -0.000040,
+        0.003396, 0.044030, -0.009194, -0.049524, -0.005015, -0.040007, 0.009104, 0.000580, 0.005603, 0.035891, -0.038913, 0.023239, -0.017022, -0.002695, -0.095759, 0.018503,
+        0.017365, 0.011104, -0.003433, 0.024113, 0.052609, -0.085274, 0.027565, -0.005833, 0.020700, 0.015842, 0.019148, 0.020203, -0.000698, -0.005337, -0.037400, -0.060144,
+        -0.031893, -0.038396, -0.001949, 0.018901, -0.014268, -0.004721, -0.055913, 0.013814, 0.215024, -0.011357, 0.057530, 0.050092, 0.016513, -0.059254, 0.001494, -0.031472,
+        0.032190, -0.047512, -0.020501, -0.002571, 0.007844, -0.063630, -0.043938, -0.079595, 0.032820, -0.021659, -0.003738, -0.035267, -0.013794, -0.021172, -0.046356, -0.077079,
+        0.021526, -0.007447, -0.050276, -0.029743, 0.022208, -0.039137, -0.021426, -0.029825, 0.029390, -0.002943, 0.073158, -0.000435, -0.032029, -0.038524, -0.029886, 0.017473,
+        0.013513, 0.022738, 0.000632, -0.073718, 0.029219, -0.018896, 0.007302, -0.116122, -0.013324, -0.012214, -0.005960, -0.003720, -0.155869, 0.019896, 0.016919, -0.021133,
+        -0.066911, -0.000926, -0.020871, -0.015295, -0.086108, 0.014918, -0.009284, 0.001689, -0.038155, 0.039163, 0.015988, 0.014413, 0.034205, -0.053273, 0.001687, 0.012227,
+        -0.007341, -0.006123, -0.005731, -0.026863, 0.060196, 0.028929, 0.019328, -0.033709, 0.038789, -0.015624, 0.013323, 0.053821, -0.015538, -0.001610, 0.012959, -0.013897,
+        0.082010, 0.012866, -0.017269, 0.000017, -0.059458, 0.015870, 0.028455, 0.025234, -0.051163, -0.022976, 0.011866, -0.005613, -0.008738, -0.047658, -0.002155, -0.029432,
+        0.039242, -0.013491, -0.001641, -0.024210, -0.019187, 0.026716, -0.025698, -0.027591, -0.034678, -0.002473, -0.019391, 0.017597, 0.064385, -0.029104, -0.034501, -0.004955,
+        0.015008, 0.060749, -0.051693, 0.020279, -0.027170, -0.027003, 0.000254, 0.011352, -0.028116, 0.028938, -0.007224, 0.019978, -0.025379, -0.004874, -0.019361, -0.020278,
+    ],
+    dtype=torch.float32,
+)
+SYSTEM_EN_RECAPTION = torch.tensor(
+    [
+        0.007721, 0.015421, -0.019305, -0.000920, 0.016031, -0.019730, 0.029683, 0.026810, -0.010510, 0.021463, 0.008833, -0.040851, 0.043260, -0.007042, 0.057224, 0.011995,
+        0.007818, 0.046369, 0.059838, -0.028548, -0.047399, -0.000983, 0.024343, -0.052259, -0.013638, 0.006856, 0.009186, 0.014235, -0.031497, -0.008644, -0.009349, 0.018900,
+        0.002913, -0.022475, 0.039518, 0.019052, -0.007600, 0.010634, -0.011830, 0.075675, -0.071738, -0.014947, 0.004995, -0.025804, -0.002553, -0.093262, 0.002881, -0.033744,
+        -0.007234, 0.013659, 0.009897, 0.039185, -0.005366, 0.041534, -0.005924, 0.019786, 0.048566, -0.009356, -0.027360, 0.042557, 0.091286, 0.009286, 0.015410, 0.028166,
+        0.022476, -0.025162, 0.012144, 0.084603, -0.003150, -0.008549, -0.002099, -0.014987, -0.019480, 0.046843, -0.030613, 0.015557, -0.008965, -0.008798, -0.027032, -0.014112,
+        0.018703, -0.014749, -0.000928, -0.024660, 0.024004, 0.004560, 0.028156, -0.028467, 0.025444, -0.038699, -0.014927, -0.031593, -0.648498, 0.018529, 0.003378, 0.030188,
+        -0.002314, 0.014950, -0.146615, -0.009005, 0.016579, -0.037867, 0.020907, 0.033160, 0.007877, -0.026345, -0.056428, 0.031255, -0.018404, 0.013334, 0.009988, -0.022790,
+        0.020803, -0.036862, 0.036222, -0.006646, -0.058084, -0.012036, 0.044199, -0.027665, -0.015779, 0.051554, 0.059970, -0.025977, 0.003967, -0.035247, -0.000488, 0.023182,
+        0.000468, 0.019190, 0.047268, -0.032279, -0.005302, 0.078669, -0.001915, 0.024918, -0.014952, -0.078905, -0.018333, 0.001362, -0.015115, 0.005435, 0.002313, 0.018766,
+        -0.032773, 0.037344, 0.024061, 0.012143, -0.057106, 0.029490, 0.019537, 0.009099, 0.026064, -0.015927, -0.037047, 0.006002, 0.025191, -0.035318, -0.032245, -0.047822,
+        -0.023568, -0.004533, 0.025100, 0.002758, -0.002649, -0.012287, -0.012139, 0.043080, 0.003295, 0.024667, -0.021050, 0.006752, 0.025315, -0.011127, 0.009800, -0.021343,
+        -0.024866, 0.010098, 0.026954, 0.012467, -0.035866, -0.031780, 0.007479, -0.003388, -0.012619, -0.012099, 0.014974, -0.001908, -0.032700, 0.004703, 0.003238, -0.007498,
+        0.023241, 0.002715, -0.111739, 0.003317, 0.006475, -0.019792, -0.046558, -0.032593, -0.020762, -0.005059, 0.016934, 0.029195, 0.028744, -0.050633, 0.001907, -0.028791,
+        -0.016695, 0.052143, 0.010439, 0.007204, 0.028502, 0.012607, -0.012414, -0.031238, 0.007305, 0.032309, 0.087924, -0.010530, 0.029925, 0.032666, -0.002202, 0.017539,
+        -0.009091, -0.001631, 0.024906, -0.013102, 0.031772, 0.018465, 0.012035, 0.031460, 0.030193, 0.005289, 0.025859, -0.038971, -0.046577, -0.025852, 0.035235, -0.038514,
+        0.001042, 0.013012, 0.023701, -0.014630, -0.029269, -0.011981, 0.008219, -0.067347, -0.003456, 0.028198, -0.008657, -0.017773, 0.010540, 0.023964, 0.021012, -0.034465,
+        -0.023748, 0.004065, -0.021598, 0.008440, -0.031533, 0.038390, -0.007680, -0.003852, 0.016136, -0.017906, -0.008927, 0.006300, -0.001251, 0.029337, -0.008632, 0.020568,
+        0.021560, -0.007222, 0.005313, -0.013089, 0.012299, 0.031303, -0.013951, 0.016547, -0.024771, -0.008753, -0.030908, -0.014421, -0.017656, 0.014044, -0.114986, 0.000956,
+        -0.035588, 0.003756, 0.015383, -0.013358, 0.009385, -0.001359, 0.012623, -0.028724, 0.001607, 0.012809, 0.032668, 0.011834, -0.015587, -0.007170, -0.021344, -0.019664,
+        0.017690, -0.014538, 0.016511, 0.038037, 0.029919, 0.020907, -0.018565, 0.032964, 0.078548, -0.050386, -0.003012, -0.016965, 0.064131, 0.008077, -0.025879, -0.035820,
+        0.095075, 0.019901, -0.019114, 0.022832, 0.003741, 0.027148, 0.018231, 0.027741, 0.020328, 0.001700, -0.006939, -0.024154, 0.018523, -0.029819, 0.008050, -0.004477,
+        0.006087, 0.056878, -0.009083, -0.061537, -0.011531, -0.037551, 0.000434, -0.005843, 0.024739, 0.032020, -0.053119, 0.020704, -0.012385, -0.002726, -0.082489, 0.009072,
+        0.013341, 0.000316, 0.001899, 0.022868, 0.034407, -0.066857, 0.020589, 0.012195, 0.023211, -0.001520, 0.000897, 0.029670, -0.015930, 0.006509, -0.035172, -0.061215,
+        -0.014099, -0.038584, -0.012213, 0.018613, -0.012365, -0.002777, -0.055184, 0.017146, 0.214358, -0.015750, 0.052488, 0.045205, 0.025334, -0.054615, 0.002117, -0.038122,
+        0.012402, -0.053418, -0.025405, 0.007235, 0.013208, -0.092481, -0.048700, -0.085186, 0.029039, -0.036767, -0.000777, -0.017625, -0.012556, -0.004887, -0.033660, -0.082310,
+        0.013387, -0.003256, -0.062981, -0.019886, 0.017624, -0.037421, -0.020743, -0.020894, 0.041974, -0.008502, 0.088413, -0.018697, -0.029398, -0.029389, -0.043721, 0.013872,
+        0.003944, 0.030361, 0.005355, -0.081355, 0.041843, -0.016395, 0.011954, -0.060440, -0.000966, -0.019101, 0.006803, -0.011310, -0.148581, 0.020342, 0.012795, -0.016473,
+        -0.053300, -0.012340, -0.016640, -0.029834, -0.082405, 0.011859, -0.004255, -0.004396, -0.012515, 0.031962, 0.030438, 0.013792, 0.031557, -0.047200, 0.006485, 0.024815,
+        -0.019376, -0.011454, -0.034184, -0.021329, 0.050115, 0.021720, 0.002874, -0.047163, 0.044031, -0.014663, 0.020534, 0.056017, 0.007017, 0.003323, 0.005734, -0.002777,
+        0.082836, 0.012048, -0.023236, -0.007401, -0.071598, 0.016760, 0.017282, 0.028306, -0.026220, -0.008016, -0.000202, -0.020271, -0.019828, -0.046986, -0.005805, -0.039647,
+        0.042879, -0.004463, 0.007753, -0.028916, -0.020612, 0.028833, -0.039839, -0.052447, -0.013275, -0.002407, -0.018937, 0.033216, 0.075535, -0.045026, -0.009901, 0.016637,
+        -0.000322, 0.073925, -0.055701, 0.014912, -0.045671, -0.021189, 0.006761, -0.002015, -0.027410, 0.018250, -0.015916, 0.016254, -0.044964, 0.029261, -0.029319, -0.005222,
+    ],
+    dtype=torch.float32,
+)
+SYSTEM_EN_THINK_RECAPTION = torch.tensor(
+    [
+        0.011004, 0.017341, -0.019959, -0.018314, 0.016520, -0.027395, 0.017946, 0.039665, 0.000645, 0.035903, 0.002499, -0.045664, 0.039472, -0.013479, 0.081302, 0.000182,
+        0.006947, 0.042845, 0.059741, -0.010796, -0.035240, 0.004176, 0.029557, -0.043467, -0.017271, 0.006896, 0.010997, 0.022498, -0.023308, -0.013046, -0.000742, 0.016209,
+        -0.007152, -0.029868, 0.028747, 0.033743, -0.000227, 0.018419, -0.015023, 0.050376, -0.098475, -0.002375, 0.007897, -0.023936, 0.007843, -0.122463, -0.011680, -0.027267,
+        -0.007270, 0.021869, -0.011415, 0.043770, 0.000551, 0.048573, 0.003132, 0.014233, 0.037080, -0.004818, -0.028738, 0.044468, 0.073843, 0.016947, 0.014484, 0.021931,
+        0.020110, -0.032309, -0.003811, 0.095704, -0.006950, -0.007237, -0.005529, -0.020573, -0.016259, 0.041909, -0.038748, 0.018029, 0.005066, -0.021186, -0.020102, -0.019719,
+        0.006239, -0.021284, 0.004213, -0.024963, 0.032345, -0.012557, 0.037268, -0.038075, 0.040998, -0.032766, -0.023509, -0.016426, -0.627412, 0.022675, 0.000101, 0.023162,
+        -0.002081, 0.015922, -0.138671, -0.027995, 0.011579, -0.042859, 0.019935, 0.038077, 0.012640, -0.017377, -0.027456, 0.035151, -0.015756, 0.018530, 0.004646, -0.002589,
+        0.019645, -0.043736, 0.034947, -0.010166, -0.061165, -0.019195, 0.028909, -0.019415, -0.009485, 0.049566, 0.068621, -0.038644, 0.011278, -0.036133, 0.000564, 0.022611,
+        -0.013612, 0.020854, 0.030614, -0.025578, 0.005673, 0.076526, -0.004887, 0.027769, -0.022605, -0.092657, -0.013218, 0.008081, -0.015227, 0.018031, -0.005145, 0.015028,
+        -0.027193, 0.034767, 0.028710, 0.032007, -0.053175, 0.033528, 0.019437, 0.011517, 0.012107, -0.027679, -0.026937, 0.008612, 0.036909, -0.051484, -0.039971, -0.034372,
+        -0.023825, -0.003025, 0.033648, -0.001852, 0.007309, 0.000714, -0.001075, 0.038534, 0.007586, 0.016213, -0.025223, -0.001099, 0.015852, -0.011477, 0.020635, -0.010696,
+        -0.019634, 0.025613, 0.034374, 0.007169, -0.035000, -0.032268, 0.015114, -0.014217, -0.005229, -0.005495, 0.018189, -0.011360, -0.026755, 0.007036, -0.002333, -0.001174,
+        0.014729, 0.001739, -0.108591, 0.004699, 0.002048, -0.014801, -0.042855, -0.028846, -0.009609, -0.004500, 0.019466, 0.021848, 0.022140, -0.063035, -0.004272, -0.030798,
+        -0.018452, 0.055169, 0.012240, -0.003555, 0.038293, 0.008503, -0.016608, -0.021309, 0.000690, 0.027093, 0.088054, -0.008881, 0.034087, 0.030647, 0.003284, 0.005038,
+        -0.008359, 0.006311, 0.032462, -0.009699, 0.035283, 0.015261, 0.012827, 0.038169, 0.033959, -0.018048, 0.018122, -0.025259, -0.040084, -0.030879, 0.019853, -0.042558,
+        -0.011938, 0.019602, 0.016537, -0.003378, -0.027890, -0.014909, -0.005464, -0.071862, 0.012335, 0.021899, -0.017008, -0.023228, 0.003263, 0.004571, 0.016447, -0.029446,
+        -0.022645, -0.001261, -0.018573, 0.007431, -0.027587, 0.035362, -0.006785, -0.000614, 0.026044, -0.009056, -0.009843, 0.010467, -0.011929, 0.042025, -0.014068, 0.023113,
+        0.023880, 0.014948, 0.004370, -0.005262, 0.012587, 0.021608, -0.001783, 0.023697, -0.024945, -0.011533, -0.020953, -0.007205, -0.024693, 0.012961, -0.168760, 0.001767,
+        -0.041265, -0.007044, 0.015021, -0.008407, 0.029642, -0.000956, 0.008607, -0.035365, -0.012187, 0.011744, 0.032612, 0.006226, -0.015891, -0.017747, -0.022565, -0.024505,
+        0.031279, 0.004188, 0.011939, 0.038032, 0.008798, 0.012314, -0.024830, 0.034484, 0.076395, -0.060108, 0.001019, -0.016138, 0.067729, 0.003899, -0.029845, -0.019960,
+        0.086663, 0.040965, -0.010458, 0.027808, -0.006394, 0.017343, 0.014788, 0.024756, 0.016446, -0.012537, -0.008406, -0.028109, 0.013369, -0.033571, 0.012170, -0.002199,
+        0.005263, 0.052280, -0.018171, -0.047898, -0.010087, -0.038632, 0.006773, -0.000838, 0.011197, 0.038187, -0.049525, 0.021689, -0.007385, -0.005987, -0.094551, 0.019019,
+        0.012760, 0.009617, -0.002262, 0.030228, 0.047823, -0.079764, 0.023391, -0.005561, 0.018866, 0.012817, 0.020878, 0.027037, -0.013905, -0.002874, -0.035522, -0.046266,
+        -0.032448, -0.036010, -0.007776, 0.016512, -0.012279, -0.005665, -0.057974, 0.016967, 0.202836, -0.009066, 0.066093, 0.045689, 0.018319, -0.048465, 0.000242, -0.040874,
+        0.027824, -0.049045, -0.015616, -0.000307, 0.009163, -0.072975, -0.042979, -0.082254, 0.040549, -0.027049, 0.000725, -0.034118, -0.019604, -0.019097, -0.042483, -0.075446,
+        0.019387, -0.005218, -0.053573, -0.029975, 0.008195, -0.036608, -0.018920, -0.025610, 0.028426, -0.002688, 0.074996, -0.003423, -0.032505, -0.030565, -0.028142, 0.014437,
+        0.013359, 0.019376, 0.008356, -0.069731, 0.031824, -0.011103, 0.019327, -0.117090, -0.009352, -0.010290, -0.002129, -0.009198, -0.172915, 0.021232, 0.017274, -0.030060,
+        -0.061449, -0.006598, -0.013069, -0.012857, -0.081220, 0.019058, -0.004841, 0.003066, -0.037741, 0.041806, 0.018281, 0.009458, 0.036761, -0.044987, 0.003557, 0.008890,
+        -0.008011, -0.004063, -0.013474, -0.022090, 0.055398, 0.037475, 0.006991, -0.035962, 0.045503, -0.017162, 0.022391, 0.052754, -0.005924, -0.005936, 0.012673, -0.017922,
+        0.084548, 0.014695, -0.013817, 0.000421, -0.065167, 0.018269, 0.023317, 0.023523, -0.034229, -0.019588, 0.007911, -0.002426, -0.017109, -0.050870, 0.002848, -0.033077,
+        0.043451, -0.010609, -0.000375, -0.023206, -0.018155, 0.027102, -0.036006, -0.035115, -0.023922, 0.005989, -0.015372, 0.027123, 0.075210, -0.035302, -0.029799, 0.003642,
+        0.007714, 0.063498, -0.053234, 0.015699, -0.040459, -0.027354, -0.002433, 0.010923, -0.020134, 0.029292, -0.010176, 0.013508, -0.032403, 0.004323, -0.017504, -0.015237,
+    ],
+    dtype=torch.float32,
+)
+SYSTEM_EN_VANILLA = torch.tensor(
+    [
+        0.010809, 0.021177, -0.017600, -0.016814, 0.012351, -0.024554, 0.018299, 0.039305, 0.003331, 0.030473, 0.005557, -0.040898, 0.047294, -0.016136, 0.076989, -0.002723,
+        0.017622, 0.042330, 0.058266, -0.016232, -0.029502, 0.004529, 0.033543, -0.041481, -0.017631, 0.002727, 0.018874, 0.019932, -0.030052, -0.009997, 0.004582, 0.002135,
+        -0.003720, -0.030923, 0.021174, 0.034033, -0.007096, 0.011522, -0.009518, 0.055688, -0.092351, -0.003914, 0.004589, -0.032635, 0.012479, -0.140607, -0.014141, -0.031821,
+        0.001396, 0.026780, -0.007623, 0.039957, 0.006434, 0.047516, 0.014377, 0.015237, 0.034212, 0.003576, -0.027357, 0.038888, 0.087272, 0.020248, 0.015165, 0.016002,
+        0.020781, -0.040509, -0.008929, 0.080857, -0.002642, -0.009738, -0.005683, -0.000615, -0.012801, 0.046457, -0.045004, 0.024689, 0.002498, -0.017333, -0.027366, -0.023231,
+        -0.006064, -0.021505, 0.007405, -0.021249, 0.026252, -0.018690, 0.020093, -0.036954, 0.037510, -0.032027, -0.030871, -0.011173, -0.618627, 0.021213, -0.004366, 0.029555,
+        -0.004324, 0.020221, -0.143832, -0.021386, 0.010482, -0.042113, 0.016164, 0.040350, 0.014627, -0.011778, -0.018102, 0.035380, -0.020305, 0.010590, 0.009227, -0.011415,
+        0.018623, -0.036384, 0.031003, -0.017073, -0.056456, -0.010423, 0.033029, -0.023511, -0.008717, 0.045716, 0.068273, -0.027886, 0.009665, -0.039801, 0.001465, 0.024361,
+        -0.015039, 0.022903, 0.033362, -0.022804, 0.008631, 0.076518, 0.000619, 0.022786, -0.015435, -0.095242, -0.006092, 0.015496, -0.009081, 0.015740, 0.004280, 0.013103,
+        -0.031836, 0.034241, 0.031836, 0.032636, -0.053721, 0.034370, 0.019172, 0.018383, 0.006907, -0.036039, -0.027927, 0.008646, 0.040496, -0.060314, -0.039116, -0.021488,
+        -0.031682, -0.005077, 0.034920, 0.002148, -0.008087, 0.002024, -0.008480, 0.041096, 0.011401, 0.020380, -0.025078, 0.005002, 0.022252, -0.014577, 0.008051, -0.014476,
+        -0.007078, 0.021075, 0.036965, 0.005343, -0.038671, -0.037222, 0.014052, -0.009952, -0.003958, -0.001878, 0.017848, -0.016608, -0.030813, 0.010921, 0.001068, 0.003095,
+        0.007076, -0.001936, -0.102996, 0.006838, -0.005243, -0.009140, -0.043796, -0.027227, -0.008426, -0.013177, 0.015602, 0.021036, 0.025484, -0.064836, -0.003593, -0.038036,
+        -0.023102, 0.064053, 0.007850, 0.000771, 0.039297, 0.011903, -0.015866, -0.017612, 0.006308, 0.024342, 0.086761, -0.016705, 0.039239, 0.025079, -0.006452, 0.003174,
+        -0.010146, 0.010787, 0.035932, -0.015346, 0.037191, 0.010990, 0.011573, 0.044958, 0.035560, -0.017339, 0.018878, -0.025394, -0.044339, -0.029852, 0.015951, -0.032248,
+        -0.012019, 0.013497, 0.012224, -0.001284, -0.034041, -0.015768, 0.000230, -0.086076, 0.024878, 0.031929, -0.016668, -0.019815, -0.001325, 0.007944, 0.017674, -0.036097,
+        -0.019651, -0.001272, -0.032842, 0.002056, -0.037140, 0.043191, -0.003710, 0.011767, 0.020313, -0.018396, -0.015935, 0.010228, -0.017349, 0.049363, -0.010007, 0.019533,
+        0.018076, 0.016608, -0.005523, -0.007793, 0.016868, 0.019341, -0.008236, 0.026765, -0.025324, -0.007849, -0.023648, -0.007791, -0.018508, 0.015357, -0.166499, -0.003718,
+        -0.035447, -0.005229, 0.019327, -0.014207, 0.028433, -0.002619, 0.013888, -0.033146, -0.017015, 0.004677, 0.039554, 0.003803, -0.014592, -0.018886, -0.023868, -0.022708,
+        0.033661, 0.008626, 0.015687, 0.046395, 0.014173, 0.015083, -0.025994, 0.039120, 0.076334, -0.061165, 0.001791, -0.017579, 0.067567, -0.002415, -0.032495, -0.025576,
+        0.079027, 0.036370, -0.013303, 0.030510, -0.009061, 0.019135, 0.015627, 0.024864, 0.015093, -0.017066, -0.014075, -0.021907, 0.017388, -0.033492, 0.013317, -0.000040,
+        0.003396, 0.044030, -0.009194, -0.049524, -0.005015, -0.040007, 0.009104, 0.000580, 0.005603, 0.035891, -0.038913, 0.023239, -0.017022, -0.002695, -0.095759, 0.018503,
+        0.017365, 0.011104, -0.003433, 0.024113, 0.052609, -0.085274, 0.027565, -0.005833, 0.020700, 0.015842, 0.019148, 0.020203, -0.000698, -0.005337, -0.037400, -0.060144,
+        -0.031893, -0.038396, -0.001949, 0.018901, -0.014268, -0.004721, -0.055913, 0.013814, 0.215024, -0.011357, 0.057530, 0.050092, 0.016513, -0.059254, 0.001494, -0.031472,
+        0.032190, -0.047512, -0.020501, -0.002571, 0.007844, -0.063630, -0.043938, -0.079595, 0.032820, -0.021659, -0.003738, -0.035267, -0.013794, -0.021172, -0.046356, -0.077079,
+        0.021526, -0.007447, -0.050276, -0.029743, 0.022208, -0.039137, -0.021426, -0.029825, 0.029390, -0.002943, 0.073158, -0.000435, -0.032029, -0.038524, -0.029886, 0.017473,
+        0.013513, 0.022738, 0.000632, -0.073718, 0.029219, -0.018896, 0.007302, -0.116122, -0.013324, -0.012214, -0.005960, -0.003720, -0.155869, 0.019896, 0.016919, -0.021133,
+        -0.066911, -0.000926, -0.020871, -0.015295, -0.086108, 0.014918, -0.009284, 0.001689, -0.038155, 0.039163, 0.015988, 0.014413, 0.034205, -0.053273, 0.001687, 0.012227,
+        -0.007341, -0.006123, -0.005731, -0.026863, 0.060196, 0.028929, 0.019328, -0.033709, 0.038789, -0.015624, 0.013323, 0.053821, -0.015538, -0.001610, 0.012959, -0.013897,
+        0.082010, 0.012866, -0.017269, 0.000017, -0.059458, 0.015870, 0.028455, 0.025234, -0.051163, -0.022976, 0.011866, -0.005613, -0.008738, -0.047658, -0.002155, -0.029432,
+        0.039242, -0.013491, -0.001641, -0.024210, -0.019187, 0.026716, -0.025698, -0.027591, -0.034678, -0.002473, -0.019391, 0.017597, 0.064385, -0.029104, -0.034501, -0.004955,
+        0.015008, 0.060749, -0.051693, 0.020279, -0.027170, -0.027003, 0.000254, 0.011352, -0.028116, 0.028938, -0.007224, 0.019978, -0.025379, -0.004874, -0.019361, -0.020278,
+    ],
+    dtype=torch.float32,
+)
+SYSTEM_EN_UNIFIED = torch.tensor(
+    [
+        0.011409, 0.014191, -0.023163, -0.020119, 0.019190, -0.029559, 0.019616, 0.035872, 0.010434, 0.028709, 0.011616, -0.039422, 0.038369, -0.004631, 0.081177, 0.007400,
+        0.008903, 0.040408, 0.055323, -0.011950, -0.026940, 0.004916, 0.028101, -0.046200, -0.016732, 0.005115, 0.012100, 0.016136, -0.026057, -0.013827, -0.004914, 0.015261,
+        -0.010824, -0.028188, 0.022934, 0.026204, -0.003855, 0.013797, -0.014518, 0.050289, -0.100077, -0.002962, 0.009050, -0.028205, 0.016294, -0.128956, -0.012730, -0.023647,
+        -0.009306, 0.020066, 0.000033, 0.043619, 0.003250, 0.053425, 0.005889, 0.021529, 0.036032, -0.003254, -0.029715, 0.048345, 0.077978, 0.010674, 0.019296, 0.018721,
+        0.019244, -0.040115, -0.004245, 0.085214, -0.005280, -0.010746, -0.000164, -0.023405, -0.015641, 0.040193, -0.038735, 0.018966, -0.004031, -0.017879, -0.023017, -0.030379,
+        0.006468, -0.015959, 0.000532, -0.026530, 0.042640, -0.006095, 0.037899, -0.043658, 0.040965, -0.034682, -0.023729, -0.019291, -0.630840, 0.029658, 0.005462, 0.026650,
+        -0.000292, 0.013954, -0.149594, -0.019405, 0.015321, -0.045104, 0.030332, 0.031727, 0.012349, -0.009553, -0.022371, 0.034043, -0.014838, 0.015398, -0.003657, 0.000477,
+        0.021084, -0.041406, 0.029946, -0.013832, -0.057358, -0.018086, 0.031598, -0.031835, -0.006697, 0.040866, 0.068602, -0.042203, 0.007362, -0.036959, 0.003794, 0.026533,
+        -0.011873, 0.017343, 0.028333, -0.021804, 0.004007, 0.075133, 0.003340, 0.025326, -0.015068, -0.092280, -0.011514, 0.006827, -0.008254, 0.021181, -0.005035, 0.022263,
+        -0.022443, 0.043919, 0.026637, 0.028568, -0.056881, 0.036740, 0.024430, 0.015891, 0.012257, -0.031126, -0.030108, 0.007229, 0.026998, -0.051685, -0.033003, -0.031170,
+        -0.024021, 0.004235, 0.030164, 0.002674, 0.008018, 0.005532, 0.001621, 0.044790, 0.006413, 0.027160, -0.015022, 0.000911, 0.019723, -0.016244, 0.020077, -0.006847,
+        -0.014110, 0.022461, 0.031656, 0.002760, -0.039078, -0.026893, 0.006628, -0.011775, -0.000240, -0.005908, 0.014943, -0.012131, -0.021755, 0.004732, -0.005297, -0.002922,
+        0.014631, -0.002010, -0.112400, 0.000842, -0.002732, -0.014861, -0.052099, -0.034167, -0.011613, -0.006101, 0.013278, 0.018867, 0.026530, -0.068150, -0.003306, -0.032801,
+        -0.018523, 0.050875, 0.005488, -0.007241, 0.045707, 0.023119, -0.021519, -0.022683, 0.004806, 0.024827, 0.091371, -0.014424, 0.043836, 0.033094, 0.002390, 0.005450,
+        -0.004893, 0.013608, 0.031272, -0.002449, 0.031607, 0.014646, 0.014146, 0.043995, 0.028826, -0.012219, 0.021008, -0.020911, -0.036967, -0.036256, 0.013328, -0.038382,
+        -0.012084, 0.018183, 0.018782, -0.004697, -0.024284, -0.015474, -0.001463, -0.076015, 0.013923, 0.022125, -0.018765, -0.010793, 0.008409, 0.002067, 0.017961, -0.029716,
+        -0.020915, -0.001779, -0.009217, -0.001933, -0.036081, 0.042577, 0.000118, -0.013920, 0.014901, -0.016486, -0.010278, -0.000449, -0.017234, 0.042453, -0.009893, 0.021087,
+        0.017671, 0.009861, -0.004210, 0.004944, 0.015627, 0.014370, -0.001128, 0.030247, -0.019552, -0.014017, -0.020859, -0.002614, -0.024405, 0.016532, -0.173204, -0.001196,
+        -0.037415, -0.010990, 0.010449, -0.006124, 0.019211, 0.003695, 0.011679, -0.031852, -0.009764, 0.005773, 0.035793, 0.003455, -0.011772, -0.020532, -0.027434, -0.024761,
+        0.027483, -0.001554, 0.010411, 0.037888, 0.015619, 0.019186, -0.021204, 0.038158, 0.074991, -0.064521, -0.002503, -0.014499, 0.068165, 0.006145, -0.032891, -0.021540,
+        0.091385, 0.047584, -0.009590, 0.028004, -0.002962, 0.021061, 0.014854, 0.025840, 0.016068, -0.014364, -0.016418, -0.033454, 0.011734, -0.036518, 0.013015, -0.003966,
+        0.000855, 0.051373, -0.010960, -0.047078, -0.011048, -0.042015, 0.006818, 0.005483, 0.010251, 0.034951, -0.046162, 0.021258, -0.013397, -0.005259, -0.093775, 0.019974,
+        0.014992, 0.004043, -0.005931, 0.035662, 0.050723, -0.083293, 0.028047, -0.008042, 0.020763, 0.016763, 0.022913, 0.027129, -0.014314, -0.009854, -0.039019, -0.044870,
+        -0.028101, -0.038026, -0.006294, 0.018265, -0.015425, -0.007866, -0.052784, 0.010470, 0.200260, -0.007798, 0.064482, 0.046612, 0.025353, -0.059695, -0.001831, -0.039643,
+        0.025148, -0.042752, -0.014928, -0.010216, 0.014195, -0.069149, -0.041424, -0.078360, 0.036999, -0.021357, 0.011032, -0.026564, -0.016214, -0.023440, -0.044723, -0.064498,
+        0.018283, -0.007165, -0.051802, -0.026299, 0.005867, -0.034691, -0.020621, -0.030512, 0.024458, -0.011330, 0.066558, -0.004069, -0.031624, -0.030639, -0.037451, 0.013079,
+        0.015152, 0.008058, 0.009223, -0.069514, 0.030702, -0.009681, 0.014826, -0.115441, -0.005514, -0.011925, 0.001046, -0.007148, -0.164128, 0.018043, 0.017001, -0.026352,
+        -0.049691, -0.011637, -0.013045, -0.014851, -0.079469, 0.017692, -0.006575, 0.001063, -0.028299, 0.038777, 0.019930, 0.010641, 0.036955, -0.039004, -0.006477, 0.004278,
+        -0.001006, -0.002514, -0.017242, -0.023927, 0.049113, 0.038393, 0.011633, -0.031537, 0.041725, -0.012146, 0.023445, 0.049999, -0.008538, 0.001319, 0.012732, -0.021170,
+        0.082096, 0.009610, -0.025717, 0.002566, -0.060849, 0.017403, 0.032650, 0.018658, -0.030629, -0.025032, 0.005555, 0.000522, -0.009667, -0.043099, 0.005939, -0.027156,
+        0.045634, -0.011986, 0.002713, -0.032225, -0.015494, 0.028734, -0.036528, -0.033101, -0.027174, 0.009490, -0.016537, 0.029435, 0.065709, -0.037711, -0.020497, -0.005578,
+        0.011768, 0.061035, -0.044676, 0.016113, -0.042945, -0.022579, 0.002430, 0.012474, -0.018198, 0.030468, -0.016646, 0.019020, -0.035804, 0.001175, -0.018312, -0.010760,
+    ],
+    dtype=torch.float32,
+)
+# fmt: on
+SYSTEM_PROMPT_CASES = [
+    pytest.param("none", None, SEED_1234, id="none"),
+    pytest.param("dynamic", "dynamic", SYSTEM_PROMPT_DYNAMIC, id="dynamic"),
+    pytest.param("en_vanilla", "en_vanilla", SYSTEM_EN_VANILLA, id="en_vanilla"),
+    pytest.param("en_recaption", "en_recaption", SYSTEM_EN_RECAPTION, id="en_recaption"),
+    pytest.param("en_think_recaption", "en_think_recaption", SYSTEM_EN_THINK_RECAPTION, id="en_think_recaption"),
+    pytest.param("en_unified", "en_unified", SYSTEM_EN_UNIFIED, id="en_unified"),
+]
+
+
+@pytest.fixture(scope="session")
+def clip_bundle() -> tuple[CLIPModel, CLIPProcessor]:
+    try:
+        model = CLIPModel.from_pretrained(LOCAL_CLIP_PATH, local_files_only=True)
+        processor = CLIPProcessor.from_pretrained(LOCAL_CLIP_PATH, local_files_only=True)
+    except OSError as exc:
+        pytest.skip(f"Could not load CLIP model from local cache ({LOCAL_CLIP_PATH}): {exc}")
+
+    model.eval()
+    return model, processor
+
+
+@pytest.fixture(scope="module")
+def omni() -> Generator[Omni, None, None]:
+    engine = Omni(
+        model=MODEL_NAME,
+        stage_configs_path=str(STAGE_CONFIG_PATH),
+        stage_init_timeout=600,
+        init_timeout=900,
+    )
+    try:
+        yield engine
+    finally:
+        engine.close()
+
+
+def _extract_generated_image(outputs: list[object]) -> Image.Image:
+    if not outputs:
+        raise AssertionError("No outputs were returned from Omni.generate()")
+
+    first_output = outputs[0]
+    if images := getattr(first_output, "images", None):
+        return images[0]
+
+    request_output = getattr(first_output, "request_output", None)
+    if request_output is not None and (images := getattr(request_output, "images", None)):
+        return images[0]
+
+    raise AssertionError("No generated image found in Omni output")
+
+
+def extract_embedding(image: Image.Image, clip_model: CLIPModel, clip_processor: CLIPProcessor) -> torch.Tensor:
+    inputs = clip_processor(images=image.convert("RGB"), return_tensors="pt")
+    with torch.inference_mode():
+        features = clip_model.get_image_features(**inputs)
+        features = F.normalize(features, p=2, dim=-1)
+    return features.squeeze(0)
+
+
+def compare_semantic(
+    expected_embedding: torch.Tensor,
+    image: Image.Image,
+    clip_model: CLIPModel,
+    clip_processor: CLIPProcessor,
+) -> float:
+    features = extract_embedding(image, clip_model, clip_processor)
+    expected = F.normalize(expected_embedding, p=2, dim=-1)
+    return torch.dot(expected, features).item()
+
+
+def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image:
+    generator_device = current_omni_platform.device_type or "cuda"
+    sampling_params = OmniDiffusionSamplingParams(
+        seed=1234,
+        generator=torch.Generator(device=generator_device).manual_seed(1234),
+        num_outputs_per_prompt=1,
+    )
+    if use_system_prompt is not None:
+        sampling_params.extra_args = {"use_system_prompt": use_system_prompt}
+
+    outputs = omni.generate({"prompt": PROMPT}, sampling_params)
+    return _extract_generated_image(outputs)
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.")
+@pytest.mark.parametrize("system_prompt_name,use_system_prompt,expected_embedding", SYSTEM_PROMPT_CASES)
+def test_system_prompt_scores(
+    omni: Omni,
+    clip_bundle: tuple[CLIPModel, CLIPProcessor],
+    system_prompt_name: str,
+    use_system_prompt: str | None,
+    expected_embedding: torch.Tensor,
+) -> None:
+    clip_model, clip_processor = clip_bundle
+    generated_image = _generate_image(omni, use_system_prompt)
+    score = compare_semantic(expected_embedding, generated_image, clip_model, clip_processor)
+
+    print(f"{system_prompt_name}: CLIP cosine similarity = {score:.6f}")

tests/e2e/offline_inference/test_magi_human.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end tests for MagiHuman pipeline via vLLM-Omni."""
+
+import io
+
+import av
+import numpy as np
+import pytest
+
+from tests.utils import hardware_test
+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
+
+
+def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
+    """Validate that the MP4 contains meaningful video and audio tracks."""
+    container = av.open(io.BytesIO(video_bytes))
+
+    v_streams = [s for s in container.streams if s.type == "video"]
+    assert len(v_streams) >= 1, "No video stream found in MP4"
+
+    a_streams = [s for s in container.streams if s.type == "audio"]
+    assert len(a_streams) >= 1, "No audio stream found in MP4"
+
+    v_stream = v_streams[0]
+    assert v_stream.width >= 1080, f"Unexpected video width: {v_stream.width}"
+    assert v_stream.height >= 1056, f"Unexpected video height: {v_stream.height}"
+
+    frame_count = 0
+    for frame in container.decode(video=0):
+        frame_count += 1
+        if frame_count >= min_frames:
+            break
+    assert frame_count >= min_frames, f"Video has only {frame_count} frames (expected >= {min_frames})"
+
+    container.close()
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_magi_human_e2e(run_level):
+    """End-to-end test for MagiHuman generating video and audio."""
+    if run_level != "advanced_model":
+        pytest.skip("MagiHuman e2e test requires advanced_model run level with real weights.")
+
+    model_path = "SII-GAIR/daVinci-MagiHuman-Base-1080p"
+
+    omni = Omni(
+        model=model_path,
+        init_timeout=1200,
+        tensor_parallel_size=2,
+    )
+
+    prompt = (
+        "A young woman with long, wavy golden blonde hair and bright blue eyes, "
+        "wearing a fitted ivory silk blouse with a delicate lace collar, sits "
+        "stationary in front of a softly lit, blurred warm-toned interior. Her "
+        "overall disposition is warm, composed, and gently confident. The camera "
+        "holds a static medium close-up, framing her from the shoulders up, "
+        "with shallow depth of field keeping her face in sharp focus. Soft "
+        "directional key light falls from the upper left, casting a gentle "
+        "highlight along her cheekbone and nose bridge. She draws a quiet breath, "
+        "the levator labii superiors relaxing as her lips part. She speaks in "
+        "clear, warm, unhurried American English: "
+        "\"The most beautiful things in life aren't things at all — "
+        "they're moments, feelings, and the people who make you feel truly alive.\" "
+        "Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
+        "shapes each vowel with precision. A faint, genuine smile engages the "
+        "zygomaticus major, lifting her lip corners fractionally. Her brows rest "
+        "in a soft, neutral arch throughout. She maintains steady, forward-facing "
+        "eye contact. Head position remains level; no torso displacement occurs.\n\n"
+        "Dialogue:\n"
+        "<Young blonde woman, American English>: "
+        "\"The most beautiful things in life aren't things at all — "
+        "they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
+        "Background Sound:\n"
+        "<Soft, warm indoor ambience with a faint distant piano melody>"
+    )
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=256,
+        width=448,
+        num_inference_steps=8,
+        seed=52,
+        extra_args={
+            "seconds": 5,
+            "sr_height": 1080,
+            "sr_width": 1920,
+            "sr_num_inference_steps": 5,
+        },
+    )
+
+    try:
+        outputs = list(
+            omni.generate(
+                prompts=[prompt],
+                sampling_params_list=[sampling_params],
+            )
+        )
+
+        assert len(outputs) > 0, "No outputs returned"
+        first = outputs[0]
+
+        assert hasattr(first, "images") and first.images, "No video frames in output"
+        video_frames = first.images[0]
+        assert isinstance(video_frames, np.ndarray), f"Expected numpy array, got {type(video_frames)}"
+        assert video_frames.ndim == 4, f"Expected 4D array (T,H,W,3), got shape {video_frames.shape}"
+
+        audio_waveform = None
+        if hasattr(first, "multimodal_output") and first.multimodal_output:
+            audio_waveform = first.multimodal_output.get("audio")
+        assert audio_waveform is not None, "No audio waveform in multimodal_output"
+
+        video_bytes = mux_video_audio_bytes(
+            video_frames,
+            audio_waveform,
+            fps=25.0,
+            audio_sample_rate=44100,
+        )
+        assert isinstance(video_bytes, bytes), f"Expected MP4 bytes, got {type(video_bytes)}"
+        assert len(video_bytes) > 1000, f"MP4 too small ({len(video_bytes)} bytes)"
+
+        _validate_mp4(video_bytes)
+    finally:
+        omni.close()

tests/e2e/online_serving/test_cosyvoice3_tts.py

@@ -50,8 +50,18 @@ tts_server_params = [
     )
 ]
 
+tts_async_chunk_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_stage_config("cosyvoice3_async_chunk.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats"],
+        ),
+        id="cosyvoice3_async_chunk",
+    )
+]
+
 
-@pytest.mark.advanced_model
 @pytest.mark.core_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
@@ -76,16 +86,16 @@ def test_voice_clone_zh_001(omni_server, openai_client) -> None:
     openai_client.send_audio_speech_request(request_config)
 
 
-@pytest.mark.advanced_model
+@pytest.mark.core_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+@pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True)
 def test_voice_clone_zh_002(omni_server, openai_client) -> None:
     """
-    Test voice cloning TTS with Chinese text via OpenAI API.
-    Deploy Setting: default yaml
+    Test voice cloning TTS with Chinese text via async_chunk streaming.
+    Deploy Setting: cosyvoice3_async_chunk.yaml
     Input Modal: text + ref_audio + ref_text
-    Output Modal: audio
+    Output Modal: audio (streamed)
     Input Setting: stream=True
     Datasets: single request
     """
@@ -100,7 +110,7 @@ def test_voice_clone_zh_002(omni_server, openai_client) -> None:
     openai_client.send_audio_speech_request(request_config)
 
 
-@pytest.mark.advanced_model
+@pytest.mark.core_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)

tests/e2e/online_serving/test_flux2_expansion.py

@@ -1,6 +1,12 @@
 """
 Tests for Flux2 Klein; currently Dev is implemented separately,
 but ideally these models will fold together in the future.
+
+Coverage:
+- FP8 + CacheDiT + Ulysses=2 + TP=2
+- Layerwise CPU offload + Ulysses=2 + Ring=2
+- Layerwise CPU offload + TP=2
+- Layerwise CPU offload + HSDP
 """
 
 import pytest
@@ -42,6 +48,45 @@ def _get_diffusion_feature_cases(model: str):
             ),
             marks=FOUR_CARD_FEATURE_MARKS,
         ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-layerwise-offload",
+                    "--ulysses-degree",
+                    "2",
+                    "--ring",
+                    "2",
+                ],
+            ),
+            id="layerwise_ulysses2_ring2",
+            marks=FOUR_CARD_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-layerwise-offload",
+                    "--tensor-parallel-size",
+                    "2",
+                ],
+            ),
+            id="layerwise_tp2",
+            marks=FOUR_CARD_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-layerwise-offload",
+                    "--use-hsdp",
+                    "--hsdp-shard-size",
+                    "2",
+                ],
+            ),
+            id="layerwise_hsdp",
+            marks=FOUR_CARD_FEATURE_MARKS,
+        ),
     ]
 
 

tests/e2e/online_serving/test_zimage_expansion.py

@@ -3,9 +3,12 @@ Tests of common diffusion feature combinations in online serving mode
 for Z-Image.
 
 Coverage is intentionally limited to the minimal 4xL4 cases that
-exercise Z-Image's supported parallel feature combinations:
+exercise Z-Image's supported feature combinations:
 - CacheDiT + FP8 + Ring=2 + TP=2
 - TeaCache + FP8 + Ulysses=2 + Ring=2
+- Layerwise CPU offload + Ulysses=2 + Ring=2
+- Layerwise CPU offload + TP=2
+- Layerwise CPU offload + HSDP
 """
 
 import pytest
@@ -64,12 +67,39 @@ def _get_diffusion_feature_cases():
             OmniServerParams(
                 model=MODEL,
                 server_args=[
+                    "--enable-layerwise-offload",
+                    "--ulysses-degree",
+                    "2",
+                    "--ring",
+                    "2",
+                ],
+            ),
+            id="layerwise_ulysses2_ring2",
+            marks=FOUR_CARD_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    "--enable-layerwise-offload",
+                    "--tensor-parallel-size",
+                    "2",
+                ],
+            ),
+            id="layerwise_tp2",
+            marks=FOUR_CARD_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    "--enable-layerwise-offload",
                     "--use-hsdp",
                     "--hsdp-shard-size",
                     "2",
                 ],
             ),
-            id="parallel_hsdp",
+            id="layerwise_hsdp",
             marks=[*FOUR_CARD_MARKS, pytest.mark.skip(reason="issue #2435")],
         ),
     ]

tests/entrypoints/openai_api/test_serving_speech.py

@@ -658,11 +658,13 @@ class TestTTSMethods:
         mock_engine_client.tts_max_instructions_length = None
         mock_models = mocker.MagicMock()
         mock_models.is_base_model.return_value = True
-        return OmniOpenAIServingSpeech(
+        server = OmniOpenAIServingSpeech(
             engine_client=mock_engine_client,
             models=mock_models,
             request_logger=mocker.MagicMock(),
         )
+        yield server
+        server.shutdown()
 
     def test_is_tts_detection_no_stage(self, speech_server):
         """Test TTS model detection when no TTS stage exists."""
@@ -1026,6 +1028,181 @@ class TestTTSMethods:
 
         assert result is None
 
+    # ── speaker field alias ──
+
+    def test_speaker_alias_accepted_as_voice(self):
+        """The 'speaker' JSON key should be accepted as an alias for 'voice'."""
+        req = OpenAICreateSpeechRequest.model_validate({"input": "Hello", "speaker": "custom_voice"})
+        assert req.voice == "custom_voice"
+
+    def test_voice_field_still_accepted(self):
+        """The canonical 'voice' JSON key should still work."""
+        req = OpenAICreateSpeechRequest.model_validate({"input": "Hello", "voice": "custom_voice"})
+        assert req.voice == "custom_voice"
+
+    def test_speaker_alias_in_base_task_with_uploaded_voice(self, speech_server):
+        """Using 'speaker' key with an uploaded voice should work for Base task."""
+        speech_server.uploaded_speakers = {
+            "utesf": {
+                "name": "UTESF",
+                "file_path": "/tmp/voice_samples/utesf.wav",
+                "mime_type": "audio/wav",
+                "ref_text": None,
+            }
+        }
+        req = OpenAICreateSpeechRequest.model_validate({"input": "Hello", "speaker": "UTESF", "task_type": "Base"})
+        assert req.voice == "UTESF"
+        with patch("pathlib.Path.exists", return_value=True):
+            result = speech_server._validate_qwen_tts_request(req)
+        assert result is None
+
+    # ── uploaded voice with embedding ──
+
+    def test_build_tts_params_with_uploaded_voice_embedding(self, speech_server):
+        """Test _build_tts_params loads embedding for embedding-uploaded voices."""
+        speech_server.uploaded_speakers = {
+            "emb_voice": {
+                "name": "emb_voice",
+                "file_path": "/tmp/voice_samples/emb_voice.safetensors",
+                "mime_type": "application/x-safetensors",
+                "embedding_source": "direct",
+                "embedding_dim": 1024,
+                "cache_status": "ready",
+                "cache_file": "/tmp/voice_samples/emb_voice.safetensors",
+            }
+        }
+        speech_server.supported_speakers = {"ryan", "vivian", "emb_voice"}
+
+        fake_embedding = [0.1] * 1024
+        with patch.object(speech_server, "_get_uploaded_speaker_embedding") as mock_get_emb:
+            mock_get_emb.return_value = fake_embedding
+            req = OpenAICreateSpeechRequest(input="Hello", voice="emb_voice")
+            params = speech_server._build_tts_params(req)
+
+            assert "voice_clone_prompt" in params
+            assert params["voice_clone_prompt"][0]["ref_spk_embedding"] == fake_embedding
+            assert params["task_type"] == ["Base"]
+            assert params["x_vector_only_mode"] == [True]
+            assert "ref_audio" not in params
+
+    # ── regression: full flow from issue #1603 ──
+
+    def test_regression_1603_speaker_key_with_uploaded_audio_voice(self, speech_server):
+        """Regression test for #1603: upload audio voice, then invoke TTS with 'speaker' key.
+
+        Verifies the full validate → build_params pipeline works end-to-end.
+        """
+        speech_server.uploaded_speakers = {
+            "utesf": {
+                "name": "UTESF",
+                "file_path": "/tmp/voice_samples/utesf.wav",
+                "mime_type": "audio/wav",
+                "ref_text": "Hola, esta es una prueba.",
+            }
+        }
+        # Parse with 'speaker' alias (the key users actually send)
+        req = OpenAICreateSpeechRequest.model_validate(
+            {"input": "Hello world", "speaker": "UTESF", "task_type": "Base"}
+        )
+        assert req.voice == "UTESF"
+
+        # Validation should pass (file exists)
+        with patch("pathlib.Path.exists", return_value=True):
+            err = speech_server._validate_qwen_tts_request(req)
+        assert err is None, f"Validation failed: {err}"
+
+        # Build params should auto-set ref_audio from stored file
+        with patch.object(speech_server, "_get_uploaded_audio_data") as mock_audio:
+            mock_audio.return_value = "data:audio/wav;base64,ZmFrZQ=="
+            params = speech_server._build_tts_params(req)
+
+        assert params["task_type"] == ["Base"]
+        assert params["ref_audio"] == ["data:audio/wav;base64,ZmFrZQ=="]
+        assert params["ref_text"] == ["Hola, esta es una prueba."]
+        assert params["x_vector_only_mode"] == [False]
+        assert params["speaker"] == ["utesf"]
+
+    def test_regression_1603_speaker_key_with_uploaded_embedding_voice(self, speech_server):
+        """Regression test for #1603: upload embedding voice, then invoke TTS with 'speaker' key.
+
+        Verifies embedding-uploaded voices are loaded as voice_clone_prompt, not as audio.
+        """
+        speech_server.uploaded_speakers = {
+            "myvoice": {
+                "name": "myvoice",
+                "file_path": "/tmp/voice_samples/myvoice.safetensors",
+                "mime_type": "application/x-safetensors",
+                "embedding_source": "direct",
+                "embedding_dim": 1024,
+                "cache_status": "ready",
+                "cache_file": "/tmp/voice_samples/myvoice.safetensors",
+            }
+        }
+        # Parse with 'speaker' alias
+        req = OpenAICreateSpeechRequest.model_validate(
+            {"input": "Hello world", "speaker": "myvoice", "task_type": "Base"}
+        )
+        assert req.voice == "myvoice"
+
+        # Validation should pass
+        with patch("pathlib.Path.exists", return_value=True):
+            err = speech_server._validate_qwen_tts_request(req)
+        assert err is None, f"Validation failed: {err}"
+
+        # Build params should use embedding, NOT audio
+        fake_emb = [0.1] * 1024
+        with patch.object(speech_server, "_get_uploaded_speaker_embedding") as mock_emb:
+            mock_emb.return_value = fake_emb
+            params = speech_server._build_tts_params(req)
+
+        assert params["task_type"] == ["Base"]
+        assert params["x_vector_only_mode"] == [True]
+        assert "voice_clone_prompt" in params
+        assert params["voice_clone_prompt"][0]["ref_spk_embedding"] == fake_emb
+        # Must NOT have ref_audio — that would fail for safetensors files
+        assert "ref_audio" not in params
+
+    def test_validate_rejects_embedding_voice_with_pending_cache(self, speech_server):
+        """Validation should reject embedding voices whose cache is not yet ready."""
+        speech_server.uploaded_speakers = {
+            "myvoice": {
+                "name": "myvoice",
+                "file_path": "/tmp/myvoice.safetensors",
+                "mime_type": "application/x-safetensors",
+                "embedding_source": "direct",
+                "cache_status": "pending",
+                "cache_file": None,
+            }
+        }
+        req = OpenAICreateSpeechRequest.model_validate({"input": "Hello", "speaker": "myvoice", "task_type": "Base"})
+        with patch("pathlib.Path.exists", return_value=True):
+            err = speech_server._validate_qwen_tts_request(req)
+        assert err is not None
+        assert "not yet ready" in err
+
+    def test_x_vector_only_mode_not_overwritten_for_uploaded_embedding(self, speech_server):
+        """x_vector_only_mode set by uploaded embedding must not be overwritten by request field."""
+        speech_server.uploaded_speakers = {
+            "emb_voice": {
+                "name": "emb_voice",
+                "file_path": "/tmp/emb_voice.safetensors",
+                "mime_type": "application/x-safetensors",
+                "embedding_source": "direct",
+                "embedding_dim": 1024,
+                "cache_status": "ready",
+                "cache_file": "/tmp/emb_voice.safetensors",
+            }
+        }
+        fake_emb = [0.1] * 1024
+        with patch.object(speech_server, "_get_uploaded_speaker_embedding") as mock_emb:
+            mock_emb.return_value = fake_emb
+            # Client explicitly sends x_vector_only_mode=False, but embedding requires True
+            req = OpenAICreateSpeechRequest(input="Hello", voice="emb_voice", x_vector_only_mode=False)
+            params = speech_server._build_tts_params(req)
+
+        assert params["x_vector_only_mode"] == [True]
+        assert "voice_clone_prompt" in params
+
     def test_max_instructions_length_default(self, speech_server):
         """Test default max instructions length (500) when no config provided."""
         # Fixture creates server with no CLI override and no TTS stage
@@ -1639,11 +1816,13 @@ def fish_speech_server(mocker: MockerFixture):
     mock_models = mocker.MagicMock()
     mock_models.is_base_model.return_value = True
 
-    return OmniOpenAIServingSpeech(
+    server = OmniOpenAIServingSpeech(
         engine_client=mock_engine_client,
         models=mock_models,
         request_logger=mocker.MagicMock(),
     )
+    yield server
+    server.shutdown()
 
 
 class TestFishSpeechServing:
@@ -1717,7 +1896,7 @@ class TestFishSpeechServing:
             fish_speech_server._build_fish_speech_prompt(request)
 
     def test_prepare_speech_generation_overrides_fish_default_max_tokens(self, fish_speech_server):
-        fish_speech_server._build_fish_speech_prompt = MagicMock(
+        fish_speech_server._build_fish_speech_prompt_async = AsyncMock(
             return_value={
                 "prompt_token_ids": [1, 2, 3],
                 "additional_information": {},
@@ -1730,13 +1909,14 @@ class TestFishSpeechServing:
 
         assert request_id.startswith("speech-")
         assert generator == "generator"
+        fish_speech_server._build_fish_speech_prompt_async.assert_awaited_once()
         fish_speech_server.engine_client.generate.assert_called_once()
         sampling_params_list = fish_speech_server.engine_client.generate.call_args.kwargs["sampling_params_list"]
         assert sampling_params_list[0].max_tokens == 4096
         assert fish_speech_server.engine_client.default_sampling_params_list[0].max_tokens == 2048
 
     def test_prepare_speech_generation_uses_stage_default_max_tokens(self, fish_speech_server):
-        fish_speech_server._build_fish_speech_prompt = MagicMock(
+        fish_speech_server._build_fish_speech_prompt_async = AsyncMock(
             return_value={
                 "prompt_token_ids": [1, 2, 3],
                 "additional_information": {},
@@ -1985,3 +2165,115 @@ class TestCosyVoice3Serving:
         assert generator == "generator"
         assert tts_params == {}
         cosyvoice3_server._build_cosyvoice3_prompt.assert_awaited_once()
+
+
+class TestTTSAsyncOffloading:
+    """Tests for event-loop-safe offloading of blocking TTS operations."""
+
+    def test_build_voxtral_prompt_is_sync(self):
+        """_build_voxtral_prompt should be a regular function, not a coroutine."""
+        assert not asyncio.iscoroutinefunction(OmniOpenAIServingSpeech._build_voxtral_prompt)
+
+    @pytest.fixture
+    def voxtral_server(self, mocker: MockerFixture):
+        mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
+        mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+        mock_engine_client.model_config = mocker.MagicMock(model="mistralai/Voxtral")
+        mock_engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
+        mock_engine_client.tts_batch_max_items = 32
+        mock_engine_client.generate = mocker.MagicMock(return_value="generator")
+        mock_engine_client.stage_configs = [
+            SimpleNamespace(
+                engine_args=SimpleNamespace(model_stage="audio_generation"),
+                tts_args={},
+            )
+        ]
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+        server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+        yield server
+        server.shutdown()
+
+    @pytest.fixture
+    def qwen3_tts_server(self, mocker: MockerFixture):
+        mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
+        mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+        mock_engine_client.model_config = mocker.MagicMock(model="Qwen/Qwen3-TTS", hf_config=mocker.MagicMock())
+        mock_engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
+        mock_engine_client.tts_batch_max_items = 32
+        mock_engine_client.generate = mocker.MagicMock(return_value="generator")
+        mock_engine_client.tts_max_instructions_length = None
+        mock_engine_client.stage_configs = [
+            SimpleNamespace(
+                engine_args=SimpleNamespace(model_stage="qwen3_tts"),
+                tts_args={},
+            )
+        ]
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+        server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+        yield server
+        server.shutdown()
+
+    def test_prepare_speech_generation_awaits_voxtral_async(self, voxtral_server):
+        """Voxtral path in _prepare_speech_generation should call the async wrapper."""
+        voxtral_server._build_voxtral_prompt_async = AsyncMock(
+            return_value={
+                "prompt_token_ids": [1, 2, 3],
+                "additional_information": {"voice": ["test"]},
+            }
+        )
+        request = OpenAICreateSpeechRequest(input="hello", voice="test")
+        asyncio.run(voxtral_server._prepare_speech_generation(request))
+        voxtral_server._build_voxtral_prompt_async.assert_awaited_once()
+
+    def test_prepare_speech_generation_awaits_qwen3_tts_async(self, qwen3_tts_server):
+        """Qwen3 TTS path should call _estimate_prompt_len_async."""
+        qwen3_tts_server._validate_tts_request = MagicMock(return_value=None)
+        qwen3_tts_server._build_tts_params = MagicMock(
+            return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]}
+        )
+        qwen3_tts_server._estimate_prompt_len_async = AsyncMock(return_value=512)
+        request = OpenAICreateSpeechRequest(input="hello")
+        asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
+        qwen3_tts_server._build_tts_params.assert_called_once()
+        qwen3_tts_server._estimate_prompt_len_async.assert_awaited_once()
+
+    def test_shutdown_is_idempotent(self, mocker: MockerFixture):
+        """Calling shutdown() twice should not raise."""
+        mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
+        mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+        mock_engine_client.stage_configs = []
+        mock_engine_client.tts_max_instructions_length = None
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+        server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+        assert server._tts_executor is not None
+        server.shutdown()
+        assert server._tts_executor is None
+        server.shutdown()  # Should not raise
+        assert server._tts_executor is None
+
+    def test_diffusion_instance_shutdown_safe(self):
+        """Diffusion instances (created via for_diffusion) should have safe shutdown."""
+        server = OmniOpenAIServingSpeech.for_diffusion(diffusion_engine=MagicMock(), model_name="test-model")
+        assert server._tts_executor is None
+        server.shutdown()  # Should not raise

tests/entrypoints/test_utils.py

@@ -5,14 +5,17 @@ from collections import Counter
 from dataclasses import dataclass
 
 import pytest
+import torch
 from pytest_mock import MockerFixture
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.entrypoints.utils import (
     _convert_dataclasses_to_dict,
     _filter_dict_like_object,
     filter_dataclass_kwargs,
+    load_and_resolve_stage_configs,
     resolve_model_config_path,
 )
 
@@ -304,3 +307,18 @@ class TestResolveModelConfigPath:
 
         assert result is not None
         assert "glm_image.yaml" in result
+
+
+class TestLoadAndResolveStageConfigs:
+    def test_load_and_resolve_with_kwargs(self):
+        """Ensure that dtype survives default stage creation."""
+        kwargs = {"dtype": torch.float32}
+        config_path, stage_configs = load_and_resolve_stage_configs(
+            model="black-forest-labs/FLUX.2-klein-4B",
+            stage_configs_path=None,
+            kwargs=kwargs,
+            default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs),
+        )
+        assert config_path is None
+        assert len(stage_configs) == 1
+        assert "dtype" in stage_configs[0]["engine_args"]

tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Unit tests for CosyVoice3 components."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 import torch.nn as nn
@@ -247,3 +249,32 @@ class TestSDPAFallback:
 
         assert out.shape == (batch, seq_len, heads, dim)
         assert out.dtype == torch.float32
+
+
+def test_code2wav_forward_finalizes_hift_tail():
+    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav
+
+    class DummyHiFT(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.m_source = SimpleNamespace(l_linear=SimpleNamespace(weight=torch.ones(1, dtype=torch.float32)))
+            self.finalize_calls: list[bool] = []
+
+        def inference(self, speech_feat, finalize=True):
+            self.finalize_calls.append(bool(finalize))
+            return torch.zeros((speech_feat.shape[0], 1, speech_feat.shape[-1]), dtype=speech_feat.dtype), None
+
+    model = object.__new__(CosyVoice3Code2Wav)
+    nn.Module.__init__(model)
+    model.hift = DummyHiFT()
+    model._forward_mel = lambda **_: torch.ones((1, 80, 8), dtype=torch.float32)
+
+    out = model.forward(
+        token=torch.tensor([[1, 2, 3]], dtype=torch.int32),
+        prompt_token=torch.tensor([[4, 5]], dtype=torch.int32),
+        prompt_feat=torch.ones((1, 4, 80), dtype=torch.float32),
+        embedding=torch.ones((1, 192), dtype=torch.float32),
+    )
+
+    assert out.shape == (1, 1, 8)
+    assert model.hift.finalize_calls == [True]

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

@@ -0,0 +1,463 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from threading import Lock
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.logits_processor.state import LogitsProcessors
+from vllm.v1.sample.metadata import SamplingMetadata
+
+from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import CosyVoice3Model
+from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummyCode2Wav:
+    def __init__(
+        self,
+        vocab_size: int,
+        num_samples: int = 32,
+        outputs: list[tuple[torch.Tensor, dict[str, object] | None]] | None = None,
+    ):
+        self.input_embedding = SimpleNamespace(num_embeddings=vocab_size)
+        self.num_samples = num_samples
+        self.outputs = list(outputs or [])
+        self.forward_calls: list[dict[str, object]] = []
+        self.forward_streaming_calls: list[dict[str, object]] = []
+
+    def forward(self, **kwargs):
+        self.forward_calls.append(kwargs)
+        token = kwargs["token"]
+        num_samples = int(token.shape[-1])
+        return torch.linspace(-1.0, 1.0, max(num_samples, 1), dtype=torch.float32).reshape(1, 1, -1)
+
+    def forward_streaming(self, **kwargs):
+        self.forward_streaming_calls.append(kwargs)
+        if self.outputs:
+            return self.outputs.pop(0)
+
+        token = kwargs["token"]
+        num_samples = int(token.shape[-1])
+        audio = torch.linspace(-1.0, 1.0, max(num_samples, 1), dtype=torch.float32).reshape(1, 1, -1)
+        new_state = None
+        if not kwargs.get("finalize", False):
+            new_state = {
+                "mel": torch.ones((1, 80, max(num_samples, 1)), dtype=torch.float32),
+                "speech_offset": audio.shape[-1],
+            }
+        return audio, new_state
+
+
+def _make_code2wav_model(
+    *,
+    with_stride_cfg: bool = False,
+    num_samples: int = 32,
+    outputs: list[tuple[torch.Tensor, dict[str, object] | None]] | None = None,
+) -> CosyVoice3Model:
+    model = object.__new__(CosyVoice3Model)
+    nn.Module.__init__(model)
+    model.model_stage = "cosyvoice3_code2wav"
+    hift_cfg = {} if not with_stride_cfg else {"upsample_rates": [8, 5, 3], "istft_params": {"hop_len": 4}}
+    model.config = SimpleNamespace(
+        sample_rate=24000,
+        hift=hift_cfg,
+        token_frame_rate=25 if with_stride_cfg else 0,
+        token_mel_ratio=2 if with_stride_cfg else 0,
+    )
+    model.code2wav = _DummyCode2Wav(vocab_size=4, num_samples=num_samples, outputs=outputs)
+    model.source_cache_len = 4
+    model.speech_window = torch.hamming_window(8, periodic=False)
+    model._stream_audio_cache_by_req = {}
+    model._stream_audio_cache_lock = Lock()
+    model._stream_vocoder_cache_by_req = {}
+    return model
+
+
+def _make_talker_model() -> CosyVoice3Model:
+    model = object.__new__(CosyVoice3Model)
+    nn.Module.__init__(model)
+    model.model_stage = "cosyvoice3_talker"
+    model.config = SimpleNamespace(
+        llm={
+            "speech_token_size": 6561,
+            "eos_token_id": 6562,
+            "sampling": {
+                "top_p": 0.8,
+                "top_k": 25,
+                "win_size": 10,
+                "tau_r": 0.1,
+            },
+        },
+        vocab_size=151923,
+    )
+    return model
+
+
+def _make_sampling_metadata(
+    *,
+    output_token_ids: list[list[int]],
+    repetition_penalty: float = 2.0,
+) -> SamplingMetadata:
+    return SamplingMetadata(
+        temperature=torch.tensor([1.0], dtype=torch.float32),
+        all_greedy=False,
+        all_random=True,
+        top_p=torch.tensor([0.8], dtype=torch.float32),
+        top_k=torch.tensor([25], dtype=torch.int32),
+        generators={},
+        max_num_logprobs=None,
+        no_penalties=False,
+        prompt_token_ids=None,
+        frequency_penalties=torch.zeros(1, dtype=torch.float32),
+        presence_penalties=torch.zeros(1, dtype=torch.float32),
+        repetition_penalties=torch.tensor([repetition_penalty], dtype=torch.float32),
+        output_token_ids=output_token_ids,
+        allowed_token_ids_mask=None,
+        bad_words_token_ids={},
+        logitsprocs=LogitsProcessors(),
+    )
+
+
+def test_split_request_ids_uses_seq_token_counts():
+    ids = torch.tensor([10, 11, 12, 13, 14], dtype=torch.long)
+    chunks = CosyVoice3Model._split_request_ids(ids, [2, 2, 2])
+    assert [c.tolist() for c in chunks] == [[10, 11], [12, 13], [14]]
+
+
+def test_split_request_ids_honors_single_request_seq_token_counts():
+    ids = torch.tensor([10, 11, 12, 13, 14], dtype=torch.long)
+    chunks = CosyVoice3Model._split_request_ids(ids, [3])
+    assert [c.tolist() for c in chunks] == [[10, 11, 12]]
+
+
+def test_sanitize_codec_tokens_filters_out_of_range():
+    model = _make_code2wav_model()
+    raw = torch.tensor([-1, 0, 3, 4, 99], dtype=torch.long)
+    clean = model._sanitize_codec_tokens(raw)
+    assert clean.tolist() == [0, 3]
+
+
+def test_forward_prefers_token_offset_when_present():
+    model = _make_code2wav_model()
+
+    runtime_info = [
+        {
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "token_offset": 2,
+            "left_context_size": 1,
+        }
+    ]
+
+    out = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+
+    assert len(out.multimodal_outputs["audio"]) == 1
+    assert out.multimodal_outputs["audio"][0].numel() > 0
+    assert len(model.code2wav.forward_streaming_calls) == 1
+    call = model.code2wav.forward_streaming_calls[0]
+    assert call["token"].shape == (1, 3)
+    assert call["token_offset_tokens"] == 2
+    assert call["finalize"] is False
+
+
+def test_forward_falls_back_to_left_context_size_for_backward_compat():
+    model = _make_code2wav_model()
+
+    runtime_info = [
+        {
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "left_context_size": 2,
+        }
+    ]
+
+    model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+
+    assert model.code2wav.forward_streaming_calls[0]["token_offset_tokens"] == 2
+
+
+def test_forward_ignores_single_request_padded_tail_tokens():
+    model = _make_code2wav_model(with_stride_cfg=True)
+    runtime_info = [
+        {
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "token_offset": 0,
+        }
+    ]
+
+    out = model.forward(
+        input_ids=torch.tensor([0, 1, 2, 3, 3], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2, 3, 4], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+
+    # The padded tail must not contribute to code2wav length.
+    assert out.multimodal_outputs["audio"][0].numel() == 3
+    assert model.code2wav.forward_streaming_calls[0]["token"].tolist() == [[0, 1, 2]]
+
+
+def test_forward_uses_non_stream_decode_without_chunk_metadata():
+    model = _make_code2wav_model()
+
+    runtime_info = [
+        {
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "prefix_ids": [101, 102],
+            "generated_len": 3,
+        }
+    ]
+
+    out = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+
+    assert out.multimodal_outputs["audio"][0].numel() == 3
+    assert len(model.code2wav.forward_calls) == 1
+    assert len(model.code2wav.forward_streaming_calls) == 0
+    call = model.code2wav.forward_calls[0]
+    assert call["token"].tolist() == [[0, 1, 2]]
+
+
+def test_forward_reuses_streaming_cache_state_between_chunks():
+    model = _make_code2wav_model(
+        outputs=[
+            (
+                torch.arange(4, dtype=torch.float32).reshape(1, 1, -1),
+                {"mel": torch.ones((1, 80, 3), dtype=torch.float32), "speech_offset": 4},
+            ),
+            (
+                torch.full((1, 1, 2), 9.0, dtype=torch.float32),
+                {"mel": torch.ones((1, 80, 5), dtype=torch.float32), "speech_offset": 6},
+            ),
+        ]
+    )
+    runtime_info = [
+        {
+            "req_id": ["rid-stream"],
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "token_offset": 0,
+            "stream_finished": torch.tensor(False),
+        }
+    ]
+
+    out1 = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+    assert out1.multimodal_outputs["audio"][0].tolist() == [0.0, 1.0, 2.0, 3.0]
+    assert model.code2wav.forward_streaming_calls[0]["cache_state"] is None
+
+    out2 = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+    assert out2.multimodal_outputs["audio"][0].tolist() == [9.0, 9.0]
+    cache_state = model.code2wav.forward_streaming_calls[1]["cache_state"]
+    assert cache_state is not None
+    assert cache_state["speech_offset"] == 4
+    assert "rid-stream" in model._stream_vocoder_cache_by_req
+
+
+def test_forward_clears_streaming_cache_on_terminal_chunk():
+    model = _make_code2wav_model(
+        outputs=[
+            (
+                torch.arange(4, dtype=torch.float32).reshape(1, 1, -1),
+                {"mel": torch.ones((1, 80, 3), dtype=torch.float32), "speech_offset": 4},
+            ),
+            (
+                torch.full((1, 1, 1), 7.0, dtype=torch.float32),
+                None,
+            ),
+        ]
+    )
+    runtime_info = [
+        {
+            "req_id": ["rid-stream"],
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "token_offset": 0,
+            "stream_finished": torch.tensor(False),
+        }
+    ]
+
+    model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+    assert "rid-stream" in model._stream_vocoder_cache_by_req
+
+    runtime_info[0]["stream_finished"] = torch.tensor(True)
+    out = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+    assert out.multimodal_outputs["audio"][0].tolist() == [7.0]
+    assert "rid-stream" not in model._stream_vocoder_cache_by_req
+
+
+def test_sample_uses_ras_rejection_for_recent_repetition():
+    model = _make_talker_model()
+    metadata = _make_sampling_metadata(output_token_ids=[[1] * 10])
+    logits = torch.tensor([[-1e9, 10.0, 0.0]], dtype=torch.float32)
+
+    out = model.sample(logits, metadata)
+
+    assert out is not None
+    assert out.sampled_token_ids.tolist() == [[2]]
+
+
+def test_sample_tolerates_padded_rows_without_history():
+    model = _make_talker_model()
+    metadata = _make_sampling_metadata(output_token_ids=[[1] * 10])
+    logits = torch.tensor(
+        [
+            [-1e9, 10.0, 0.0],
+            [-1e9, 0.0, 10.0],
+        ],
+        dtype=torch.float32,
+    )
+
+    out = model.sample(logits, metadata)
+
+    assert out is not None
+    assert out.sampled_token_ids.shape == (2, 1)
+
+
+def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in():
+    metadata = _make_sampling_metadata(output_token_ids=[[1, 2, 3]])
+    expected = SamplerOutput(
+        sampled_token_ids=torch.tensor([[7]], dtype=torch.int32),
+        logprobs_tensors=None,
+    )
+    calls: list[torch.Tensor] = []
+
+    class _DummyInputBatch:
+        def __init__(self):
+            self.sampling_metadata = metadata
+            self.updated = False
+
+        def update_async_output_token_ids(self):
+            self.updated = True
+
+    runner = object.__new__(GPUARModelRunner)
+    runner.input_batch = _DummyInputBatch()
+    runner.model = SimpleNamespace(
+        prefer_model_sampler=True,
+        sample=lambda logits, sampling_metadata: calls.append(logits.clone()) or expected,
+    )
+    runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
+
+    out = runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
+
+    assert out is expected
+    assert runner.input_batch.updated is False
+    assert len(calls) == 1
+
+
+def test_gpu_ar_model_runner_supplies_req_output_history_to_model_sampler():
+    metadata = _make_sampling_metadata(output_token_ids=[])
+    seen_histories: list[list[list[int]]] = []
+
+    class _DummyInputBatch:
+        def __init__(self):
+            self.sampling_metadata = metadata
+            self.req_output_token_ids = [[1, 2, 3]]
+            self.req_ids = ["rid-1"]
+            self.sampled_token_ids_cpu = None
+            self.async_copy_ready_event = None
+            self.prev_req_id_to_index = None
+
+        def update_async_output_token_ids(self):
+            raise AssertionError("fallback async repair should not run for model sampler path")
+
+    runner = object.__new__(GPUARModelRunner)
+    runner.input_batch = _DummyInputBatch()
+    runner.model = SimpleNamespace(
+        prefer_model_sampler=True,
+        sample=lambda logits, sampling_metadata: seen_histories.append(
+            [list(x) for x in sampling_metadata.output_token_ids]
+        )
+        or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None),
+    )
+    runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
+
+    runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
+
+    assert seen_histories == [[[1, 2, 3]]]
+
+
+def test_gpu_ar_model_runner_repairs_async_placeholders_for_model_sampler():
+    metadata = _make_sampling_metadata(output_token_ids=[])
+    seen_histories: list[list[list[int]]] = []
+
+    class _ReadyEvent:
+        def __init__(self):
+            self.synced = False
+
+        def synchronize(self):
+            self.synced = True
+
+    class _DummyInputBatch:
+        def __init__(self):
+            self.sampling_metadata = metadata
+            self.req_output_token_ids = [[11, -1]]
+            self.req_ids = ["rid-1"]
+            self.sampled_token_ids_cpu = torch.tensor([[29]], dtype=torch.int32)
+            self.async_copy_ready_event = _ReadyEvent()
+            self.prev_req_id_to_index = {"rid-1": 0}
+
+        def update_async_output_token_ids(self):
+            raise AssertionError("fallback async repair should not run for model sampler path")
+
+    runner = object.__new__(GPUARModelRunner)
+    runner.input_batch = _DummyInputBatch()
+    runner.model = SimpleNamespace(
+        prefer_model_sampler=True,
+        sample=lambda logits, sampling_metadata: seen_histories.append(
+            [list(x) for x in sampling_metadata.output_token_ids]
+        )
+        or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None),
+    )
+    runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
+
+    runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
+
+    assert runner.input_batch.async_copy_ready_event.synced is True
+    assert seen_histories == [[[11, 29]]]

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for code predictor dtype alignment (fix for #2385).
+
+Verifies that the code predictor handles dtype mismatches between input
+tensors and model parameters without raising RuntimeError. This can happen
+when model weights are loaded in float16/bfloat16 but upstream modules
+produce float32 hidden states.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+import sys
+import types
+from unittest.mock import MagicMock, patch
+
+import torch
+
+# Direct file import to avoid vllm_omni.__init__ patch dependencies.
+_BASE = os.path.join(
+    os.path.dirname(__file__),
+    os.pardir,
+    os.pardir,
+    os.pardir,
+    os.pardir,
+    "vllm_omni",
+    "model_executor",
+    "models",
+    "qwen3_tts",
+)
+
+
+def _load_module(name: str, filename: str):
+    path = os.path.abspath(os.path.join(_BASE, filename))
+    spec = importlib.util.spec_from_file_location(name, path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def _build_mock_modules() -> dict[str, object]:
+    """Build the dict of modules to inject into sys.modules."""
+    platforms_mock = MagicMock()
+    platforms_mock.current_omni_platform.supports_torch_inductor.return_value = False
+
+    logger_mock = MagicMock()
+    logger_mock.init_logger = lambda name: MagicMock()
+
+    vllm_config_mod = MagicMock()
+    vllm_config_mod.set_current_vllm_config = lambda cfg: MagicMock(__enter__=MagicMock(), __exit__=MagicMock())
+
+    weight_utils_mock = MagicMock()
+    weight_utils_mock.default_weight_loader = lambda p, w: None
+
+    pkg = types.ModuleType("vllm_omni.model_executor.models.qwen3_tts")
+    pkg.__path__ = [os.path.abspath(_BASE)]
+
+    return {
+        "vllm_omni": MagicMock(),
+        "vllm_omni.platforms": platforms_mock,
+        "vllm.logger": logger_mock,
+        "vllm.config": MagicMock(),
+        "vllm.config.vllm": vllm_config_mod,
+        "vllm.model_executor.model_loader.weight_utils": weight_utils_mock,
+        "vllm_omni.model_executor": types.ModuleType("vllm_omni.model_executor"),
+        "vllm_omni.model_executor.models": types.ModuleType("vllm_omni.model_executor.models"),
+        "vllm_omni.model_executor.models.qwen3_tts": pkg,
+    }
+
+
+def _load_target_classes():
+    """Load config and code predictor modules with mocked dependencies.
+
+    Uses patch.dict to ensure sys.modules is always restored, even on failure.
+    """
+    mocks = _build_mock_modules()
+    with patch.dict(sys.modules, mocks):
+        config_mod = _load_module(
+            "vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts",
+            "configuration_qwen3_tts.py",
+        )
+        sys.modules["vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts"] = config_mod
+
+        cp_mod = _load_module(
+            "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm",
+            "qwen3_tts_code_predictor_vllm.py",
+        )
+
+    return config_mod, cp_mod
+
+
+_config_mod, _cp_mod = _load_target_classes()
+
+Qwen3TTSTalkerCodePredictorConfig = _config_mod.Qwen3TTSTalkerCodePredictorConfig
+Qwen3TTSTalkerConfig = _config_mod.Qwen3TTSTalkerConfig
+CodePredictorWrapper = _cp_mod.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM
+CodePredictorModel = _cp_mod.Qwen3TTSTalkerCodePredictorModelVLLM
+
+
+def _make_tiny_config() -> tuple:
+    """Create minimal configs for a tiny code predictor model."""
+    cp_config = Qwen3TTSTalkerCodePredictorConfig(
+        vocab_size=64,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=8,
+        num_code_groups=4,
+        rms_norm_eps=1e-6,
+    )
+    talker_config = Qwen3TTSTalkerConfig(
+        hidden_size=32,
+        num_code_groups=4,
+    )
+    return cp_config, talker_config
+
+
+def _make_vllm_config(max_num_seqs: int = 4) -> MagicMock:
+    """Create a mock VllmConfig with scheduler_config."""
+    vllm_config = MagicMock()
+    vllm_config.scheduler_config.max_num_seqs = max_num_seqs
+    return vllm_config
+
+
+class TestCodePredictorDtypeAlignment:
+    """Test that code predictor buffers match model parameter dtype."""
+
+    def test_ensure_buffers_uses_given_dtype(self) -> None:
+        """_ensure_buffers should create proj_buf with the given dtype."""
+        cp_config, talker_config = _make_tiny_config()
+        vllm_config = _make_vllm_config()
+
+        predictor = CodePredictorWrapper(
+            vllm_config=vllm_config,
+            config=cp_config,
+            talker_config=talker_config,
+        )
+
+        # Create buffer in float16
+        predictor._ensure_buffers(torch.device("cpu"), torch.float16)
+        assert predictor._proj_buf is not None
+        assert predictor._proj_buf.dtype == torch.float16
+
+        # Re-create buffer in float32 (different dtype triggers re-allocation)
+        predictor._ensure_buffers(torch.device("cpu"), torch.float32)
+        assert predictor._proj_buf.dtype == torch.float32
+
+    def test_warmup_aligns_buffer_to_model_params(self) -> None:
+        """_warmup_buckets should align proj_buf dtype to model parameters."""
+        cp_config, talker_config = _make_tiny_config()
+        vllm_config = _make_vllm_config(max_num_seqs=2)
+
+        predictor = CodePredictorWrapper(
+            vllm_config=vllm_config,
+            config=cp_config,
+            talker_config=talker_config,
+        )
+
+        # Cast model to float16 (simulating vLLM loading weights in half precision)
+        predictor = predictor.to(torch.float16)
+
+        # Pre-create proj_buf with WRONG dtype (float32) — simulating the bug
+        predictor._ensure_buffers(torch.device("cpu"), torch.float32)
+        assert predictor._proj_buf.dtype == torch.float32
+
+        # Simulate _setup_compile having cached model dtype and compiled forward
+        predictor._model_dtype = torch.float16
+        predictor._compiled_model_fwd = predictor.model.forward
+
+        # _warmup_buckets should fix the dtype mismatch
+        predictor._warmup_buckets()
+
+        assert predictor._proj_buf.dtype == torch.float16
+
+    def test_setup_compile_caches_model_dtype(self) -> None:
+        """_setup_compile should cache model parameter dtype."""
+        cp_config, talker_config = _make_tiny_config()
+        vllm_config = _make_vllm_config(max_num_seqs=2)
+
+        predictor = CodePredictorWrapper(
+            vllm_config=vllm_config,
+            config=cp_config,
+            talker_config=talker_config,
+        )
+        predictor = predictor.to(torch.float16)
+
+        assert predictor._model_dtype is None
+        predictor._setup_compile()
+        assert predictor._model_dtype == torch.float16
+
+    def test_forward_with_mismatched_input_dtype(self) -> None:
+        """forward() should not crash when inputs are float32 but model is float16."""
+        cp_config, talker_config = _make_tiny_config()
+        vllm_config = _make_vllm_config(max_num_seqs=2)
+
+        predictor = CodePredictorWrapper(
+            vllm_config=vllm_config,
+            config=cp_config,
+            talker_config=talker_config,
+        )
+
+        # Model in float16
+        predictor = predictor.to(torch.float16)
+
+        bsz = 1
+        num_groups = cp_config.num_code_groups
+        hidden = talker_config.hidden_size
+
+        # Inputs in float32 (simulating the dtype mismatch from #2385)
+        layer0_code = torch.zeros(bsz, dtype=torch.long)
+        layer0_embed = torch.randn(bsz, hidden, dtype=torch.float32)
+        last_talker_hidden = torch.randn(bsz, hidden, dtype=torch.float32)
+
+        # This should NOT raise RuntimeError about dtype mismatch
+        result = predictor(
+            layer0_code=layer0_code,
+            layer0_embed=layer0_embed,
+            last_talker_hidden=last_talker_hidden,
+            do_sample=False,
+        )
+
+        assert result.shape == (bsz, num_groups)
+        assert result.dtype == torch.long
+
+
+class TestCodePredictorModelDtype:
+    """Test the inner model forward with different dtypes."""
+
+    def test_model_forward_float16(self) -> None:
+        """Inner model forward should work in float16."""
+        cp_config, _ = _make_tiny_config()
+        model = CodePredictorModel(cp_config, talker_hidden_size=32).to(torch.float16)
+
+        bsz, seq_len = 1, 4
+        inputs = torch.randn(bsz, seq_len, 32, dtype=torch.float16)
+        pos_ids = torch.arange(seq_len).unsqueeze(0).expand(bsz, -1)
+
+        output = model(inputs, pos_ids)
+        assert output.dtype == torch.float16
+        assert output.shape == (bsz, seq_len, 32)
+
+    def test_model_forward_float32(self) -> None:
+        """Inner model forward should work in float32."""
+        cp_config, _ = _make_tiny_config()
+        model = CodePredictorModel(cp_config, talker_hidden_size=32).to(torch.float32)
+
+        bsz, seq_len = 1, 4
+        inputs = torch.randn(bsz, seq_len, 32, dtype=torch.float32)
+        pos_ids = torch.arange(seq_len).unsqueeze(0).expand(bsz, -1)
+
+        output = model(inputs, pos_ids)
+        assert output.dtype == torch.float32
+        assert output.shape == (bsz, seq_len, 32)

tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import Qwen3TTSCode2Wav
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeDecoder(nn.Module):
+    def __init__(self, total_upsample: int = 4):
+        super().__init__()
+        self.total_upsample = total_upsample
+
+    def chunked_decode(self, codes: torch.Tensor) -> torch.Tensor:
+        frames = codes.shape[-1]
+        wav_len = frames * self.total_upsample + 6
+        wav = torch.arange(wav_len, dtype=torch.float32)
+        return wav.view(1, 1, -1)
+
+
+def _make_model() -> Qwen3TTSCode2Wav:
+    model = Qwen3TTSCode2Wav(
+        vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(model="unused"),
+            device_config=SimpleNamespace(device=torch.device("cpu")),
+        )
+    )
+    model._decoder = _FakeDecoder()
+    model._num_quantizers = 2
+    model._output_sample_rate = 24000
+    model._total_upsample = 4
+    model._ensure_speech_tokenizer_loaded = lambda: None
+    return model
+
+
+def test_forward_trims_context_on_exact_frame_boundaries():
+    model = _make_model()
+
+    out = model.forward(
+        input_ids=torch.arange(12, dtype=torch.long),
+        runtime_additional_information=[{"left_context_size": 2}],
+    )
+
+    audio = out.multimodal_outputs["model_outputs"][0]
+    expected = torch.arange(8, 24, dtype=torch.float32)
+    torch.testing.assert_close(audio, expected)
+
+
+def test_forward_trims_trailing_padding_without_context():
+    model = _make_model()
+
+    out = model.forward(
+        input_ids=torch.arange(12, dtype=torch.long),
+        runtime_additional_information=[{"left_context_size": 0}],
+    )
+
+    audio = out.multimodal_outputs["model_outputs"][0]
+    expected = torch.arange(24, dtype=torch.float32)
+    torch.testing.assert_close(audio, expected)

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -0,0 +1,267 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import talker2code2wav_async_chunk, text2flow
+
+
+def _source_output(request_id: str, prompt_ids: list[int], out_ids: list[int], mm: dict):
+    return SimpleNamespace(
+        request_id=request_id,
+        prompt_token_ids=prompt_ids,
+        outputs=[SimpleNamespace(token_ids=out_ids, multimodal_output=mm)],
+    )
+
+
+def _transfer_manager(
+    *,
+    chunk_frames: int = 2,
+    pre_lookahead_frames: int = 0,
+    stream_scale_factor: int = 1,
+    max_chunk_frames: int | None = None,
+):
+    if max_chunk_frames is None:
+        max_chunk_frames = chunk_frames
+    return SimpleNamespace(
+        code_prompt_token_ids=defaultdict(list),
+        request_payload={},
+        connector=SimpleNamespace(
+            config={
+                "extra": {
+                    "codec_chunk_frames": chunk_frames,
+                    "codec_pre_lookahead_frames": pre_lookahead_frames,
+                    "codec_max_chunk_frames": max_chunk_frames,
+                    "codec_stream_scale_factor": stream_scale_factor,
+                    "codec_vocab_size": 6561,
+                }
+            }
+        ),
+    )
+
+
+def test_text2flow_supports_batched_source_outputs():
+    stage_list = [
+        SimpleNamespace(
+            engine_outputs=[
+                _source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[1, 2]])}),
+                _source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[3, 4]])}),
+            ]
+        )
+    ]
+
+    outputs = text2flow(stage_list=stage_list, engine_input_source=[0], prompt=None)
+
+    assert len(outputs) == 2
+    assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
+    assert outputs[1]["prompt_token_ids"] == [4, 5]
+    assert outputs[0]["additional_information"]["prefix_ids"] == [10, 11]
+    assert outputs[1]["additional_information"]["prefix_ids"] == [20, 21]
+
+
+def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
+    transfer_manager = _transfer_manager()
+    request = SimpleNamespace(
+        external_req_id="rid-0",
+        output_token_ids=[1, 2, 6562, 3],
+        additional_information={
+            "speech_token": [torch.tensor([[11, 12, 13]])],
+            "speech_feat": [torch.tensor([[[0.1, 0.2], [0.3, 0.4]]])],
+            "embedding": [torch.tensor([[0.5, 0.6]])],
+        },
+        is_finished=lambda: True,
+    )
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=True,
+    )
+
+    assert payload is not None
+    assert payload["finished"].item() is True
+    assert payload["code_predictor_codes"] == [1, 2, 3]
+    assert payload["token_offset"] == 0
+    assert payload["left_context_size"] == 0
+    assert payload["req_id"] == ["rid-0"]
+    assert payload["stream_finished"].item() is True
+    assert "speech_token" in payload
+    assert "speech_feat" in payload
+    assert "embedding" in payload
+
+
+def test_talker2code2wav_async_chunk_emits_eof_when_finished_without_valid_codes():
+    transfer_manager = _transfer_manager(chunk_frames=25)
+    request = SimpleNamespace(
+        external_req_id="rid-eof",
+        output_token_ids=[6561, 6562],  # all filtered out
+        additional_information={},
+        is_finished=lambda: True,
+    )
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=True,
+    )
+
+    assert payload is not None
+    assert payload["code_predictor_codes"] == []
+    assert payload["finished"].item() is True
+
+
+def test_talker2code2wav_async_chunk_does_not_reemit_without_new_tokens():
+    transfer_manager = _transfer_manager()
+    request = SimpleNamespace(
+        external_req_id="rid-stable",
+        output_token_ids=[1, 2],
+        additional_information={},
+        is_finished=lambda: False,
+    )
+
+    payload1 = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+    payload2 = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+
+    assert payload1 is not None
+    assert payload1["code_predictor_codes"] == [1, 2]
+    assert payload1["token_offset"] == 0
+    assert payload2 is None
+
+
+def test_talker2code2wav_async_chunk_waits_for_prelookahead_and_emits_cumulative_prefix():
+    transfer_manager = _transfer_manager(pre_lookahead_frames=1)
+    request = SimpleNamespace(
+        external_req_id="rid-pre",
+        output_token_ids=[1, 2],
+        additional_information={},
+        is_finished=lambda: False,
+    )
+
+    payload_pending = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+    request.output_token_ids = [1, 2, 3]
+    payload_ready = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+
+    assert payload_pending is None
+    assert payload_ready is not None
+    assert payload_ready["code_predictor_codes"] == [1, 2, 3]
+    assert payload_ready["token_offset"] == 0
+    assert payload_ready["finished"].item() is False
+
+
+def test_talker2code2wav_async_chunk_final_flush_uses_previous_token_offset():
+    transfer_manager = _transfer_manager(pre_lookahead_frames=1)
+    request = SimpleNamespace(
+        external_req_id="rid-tail",
+        output_token_ids=[3, 4, 5],
+        additional_information={},
+        is_finished=lambda: False,
+    )
+
+    payload_stream = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+    request.output_token_ids = [3, 4, 5, 6]
+    payload_final = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=True,
+    )
+
+    assert payload_stream is not None
+    assert payload_stream["finished"].item() is False
+    assert payload_stream["code_predictor_codes"] == [3, 4, 5]
+    assert payload_stream["token_offset"] == 0
+    assert payload_final is not None
+    assert payload_final["finished"].item() is True
+    assert payload_final["code_predictor_codes"] == [3, 4, 5, 6]
+    assert payload_final["token_offset"] == 2
+
+
+def test_talker2code2wav_async_chunk_respects_prompt_token_pad_on_first_chunk():
+    transfer_manager = _transfer_manager(pre_lookahead_frames=1)
+    request = SimpleNamespace(
+        external_req_id="rid-pad",
+        output_token_ids=[8, 9, 10],
+        additional_information={
+            "speech_token": [torch.tensor([[1, 2, 3]])],
+        },
+        is_finished=lambda: False,
+    )
+
+    payload_pending = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+    request.output_token_ids = [8, 9, 10, 11]
+    payload_ready = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+
+    assert payload_pending is None
+    assert payload_ready is not None
+    assert payload_ready["code_predictor_codes"] == [8, 9, 10, 11]
+    assert payload_ready["token_offset"] == 0
+
+
+def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio():
+    transfer_manager = _transfer_manager()
+    request = SimpleNamespace(
+        external_req_id="rid-eof-tail",
+        output_token_ids=[3, 4],
+        additional_information={},
+        is_finished=lambda: False,
+    )
+
+    payload_stream = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+    payload_final = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=True,
+    )
+
+    assert payload_stream is not None
+    assert payload_stream["finished"].item() is False
+    assert payload_stream["code_predictor_codes"] == [3, 4]
+    assert payload_final is not None
+    assert payload_final["finished"].item() is True
+    assert payload_final["code_predictor_codes"] == []

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -64,6 +64,9 @@ class OmniARScheduler(VLLMScheduler):
 
         # Track requests that have already triggered prefill transfer to avoid duplicates
         self.transfer_triggered_requests: set[str] = set()
+
+        # Cache per-request flag to avoid repeated deserialization of additional_information
+        self._omits_kv_transfer_cache: dict[str, bool] = {}
         model_config = self.vllm_config.model_config
         self.chunk_transfer_adapter = None
         if getattr(model_config, "async_chunk", False):
@@ -82,6 +85,27 @@ class OmniARScheduler(VLLMScheduler):
                 return getattr(omni_kv_config, "kv_transfer_criteria", None)
         return None
 
+    def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool:
+        """True when orchestrator will not run stage 1+ for this request (e.g. text-only).
+
+        The result is cached per request to avoid repeated deserialization of
+        additional_information on every scheduler tick.
+        """
+        rid = request.request_id
+        cached = self._omits_kv_transfer_cache.get(rid)
+        if cached is not None:
+            return cached
+
+        payload = getattr(request, "additional_information", None)
+        if payload is None:
+            result = False
+        else:
+            info = deserialize_additional_information(payload)
+            result = info.get("omni_final_stage_id") == 0
+
+        self._omits_kv_transfer_cache[rid] = result
+        return result
+
     def _process_kv_transfer_trigger(self, request: Request, new_token_ids: list[int]) -> bool:
         """
         Check triggers and process side effects (marking transfer).
@@ -91,6 +115,10 @@ class OmniARScheduler(VLLMScheduler):
         if not self.kv_transfer_criteria:
             return False
 
+        # Text-only requests finalize at stage 0; do not prefill-stop for DiT KV.
+        if self._request_omits_kv_transfer_to_next_stage(request):
+            return False
+
         if request.request_id in self.waiting_for_transfer_free:
             return False
 
@@ -316,10 +344,6 @@ class OmniARScheduler(VLLMScheduler):
                 finished = self._handle_stopped_request(request)
                 if finished:
                     kv_transfer_params = self._free_request(request)
-                    if self.chunk_transfer_adapter is not None:
-                        self.chunk_transfer_adapter.cleanup_receiver(
-                            request.request_id,
-                        )
                 if status_before_stop == RequestStatus.RUNNING:
                     stopped_running_reqs.add(request)
                 elif status_before_stop == RequestStatus.WAITING_FOR_CHUNK:
@@ -516,6 +540,8 @@ class OmniARScheduler(VLLMScheduler):
         """Mark a request as finished and free its resources."""
         assert request.is_finished()
 
+        self._omits_kv_transfer_cache.pop(request.request_id, None)
+
         # 1. Standard cleanup parts from base _free_request
         connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)
 
@@ -642,7 +668,12 @@ class OmniARScheduler(VLLMScheduler):
                 need_send = omni_kv_config.get("need_send_cache", False)
             else:
                 need_send = getattr(omni_kv_config, "need_send_cache", False)
-        return need_send
+        if not need_send:
+            return False
+        request = self.requests.get(req_id)
+        if request is not None and self._request_omits_kv_transfer_to_next_stage(request):
+            return False
+        return True
 
     def has_requests(self) -> bool:
         """Check if there are any requests to process, including KV transfers."""

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py

@@ -93,7 +93,7 @@ class DistributedAutoencoderKL_base(DistributedVaeMixin):
 
         _, _, latent_h, latent_w = z.shape
         scale = int(2 ** (len(self.config.block_out_channels) - 1))
-        max_parallel_size = self.distributed_decoder.parallel_size
+        max_parallel_size = self.distributed_executor.parallel_size
 
         root = int(math.sqrt(max_parallel_size))
         for rows in range(root, 0, -1):
@@ -187,7 +187,7 @@ class DistributedAutoencoderKL_base(DistributedVaeMixin):
         if split is not None:
             strategy = "tile" if split == self.tile_split else "patch"
             logger.info(f"Decode run with distributed executor, split strategy is {strategy}")
-            result = self.distributed_decoder.execute(
+            result = self.distributed_executor.execute(
                 z, DistributedOperator(split=split, exec=exec, merge=merge), broadcast_result=False
             )
             if not return_dict:

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py

@@ -108,8 +108,8 @@ class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMi
         if not self.is_distributed_enabled():
             return super().tiled_decode(z, return_dict=return_dict)
 
-        logger.info("Decode run with distributed executor")
-        result = self.distributed_decoder.execute(
+        logger.debug("Decode running with distributed executor")
+        result = self.distributed_executor.execute(
             z,
             DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
             broadcast_result=True,

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py

@@ -92,6 +92,119 @@ class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
         result = torch.cat(time, dim=2)
         return result
 
+    def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, num_frames, height, width = x.shape
+        encode_spatial_compression_ratio = self.spatial_compression_ratio
+        # Scale tile parameters for patchified coordinate system
+        tile_sample_min_height = self.tile_sample_min_height
+        tile_sample_min_width = self.tile_sample_min_width
+        tile_sample_stride_height = self.tile_sample_stride_height
+        tile_sample_stride_width = self.tile_sample_stride_width
+        if self.config.patch_size is not None:
+            assert encode_spatial_compression_ratio % self.config.patch_size == 0
+            encode_spatial_compression_ratio = self.spatial_compression_ratio // self.config.patch_size
+            # When input is patchified, scale tile parameters accordingly
+            tile_sample_min_height = tile_sample_min_height // self.config.patch_size
+            tile_sample_min_width = tile_sample_min_width // self.config.patch_size
+            tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
+            tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
+
+        latent_height = height // encode_spatial_compression_ratio
+        latent_width = width // encode_spatial_compression_ratio
+
+        tile_latent_min_height = tile_sample_min_height // encode_spatial_compression_ratio
+        tile_latent_min_width = tile_sample_min_width // encode_spatial_compression_ratio
+        tile_latent_stride_height = tile_sample_stride_height // encode_spatial_compression_ratio
+        tile_latent_stride_width = tile_sample_stride_width // encode_spatial_compression_ratio
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        tiletask_list = []
+        temporal_compression = self.config.scale_factor_temporal
+        for i in range(0, height, tile_sample_stride_height):
+            for j in range(0, width, tile_sample_stride_width):
+                time_list = []
+                frame_range = 1 + (num_frames - 1) // temporal_compression
+                for k in range(frame_range):
+                    if k == 0:
+                        tile = x[:, :, :1, i : i + tile_sample_min_height, j : j + tile_sample_min_width]
+                    else:
+                        tile = x[
+                            :,
+                            :,
+                            1 + temporal_compression * (k - 1) : 1 + temporal_compression * k,
+                            i : i + tile_sample_min_height,
+                            j : j + tile_sample_min_width,
+                        ]
+                    time_list.append(tile)
+                tiletask_list.append(
+                    TileTask(
+                        len(tiletask_list),
+                        (i // tile_sample_stride_height, j // tile_sample_stride_width),
+                        time_list,
+                        workload=time_list[0].shape[3] * time_list[0].shape[4],
+                    )
+                )
+
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec={
+                "latent_height": latent_height,
+                "latent_width": latent_width,
+                "blend_height": blend_height,
+                "blend_width": blend_width,
+                "tile_latent_stride_height": tile_latent_stride_height,
+                "tile_latent_stride_width": tile_latent_stride_width,
+            },
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def encode_tile_exec(self, task: TileTask) -> torch.Tensor:
+        """Encode a single sample tile into latent space."""
+        self.clear_cache()
+        time = []
+        for k, tile in enumerate(task.tensor):
+            self._enc_conv_idx = [0]
+            encoded = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+            encoded = self.quant_conv(encoded)
+            time.append(encoded)
+        result = torch.cat(time, dim=2)
+        self.clear_cache()
+        return result
+
+    def encode_tile_merge(
+        self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec
+    ) -> torch.Tensor:
+        """Merge encoded tiles into a full latent tensor."""
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+                result_row.append(
+                    tile[
+                        :,
+                        :,
+                        :,
+                        : grid_spec.tile_spec["tile_latent_stride_height"],
+                        : grid_spec.tile_spec["tile_latent_stride_width"],
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        enc = torch.cat(result_rows, dim=3)[
+            :, :, :, : grid_spec.tile_spec["latent_height"], : grid_spec.tile_spec["latent_width"]
+        ]
+        return enc
+
     def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
         """Merge decoded tiles into a full image."""
         grid_h, grid_w = grid_spec.grid_shape
@@ -130,8 +243,8 @@ class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
         if not self.is_distributed_enabled():
             return super().tiled_decode(z, return_dict=return_dict)
 
-        logger.info("Decode run with distributed executor")
-        result = self.distributed_decoder.execute(
+        logger.debug("Decode running with distributed executor")
+        result = self.distributed_executor.execute(
             z,
             DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
             broadcast_result=False,
@@ -140,3 +253,26 @@ class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
             return (result,)
 
         return DecoderOutput(sample=result)
+
+    def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode using distributed VAE executor.
+
+        Note: x is already patchified by parent's _encode() before calling this method.
+        """
+        if not self.is_distributed_enabled():
+            return super().tiled_encode(x)
+
+        logger.debug("Encode running with distributed executor")
+        self.clear_cache()
+        result = self.distributed_executor.execute(
+            x,
+            DistributedOperator(
+                split=self.encode_tile_split,
+                exec=self.encode_tile_exec,
+                merge=self.encode_tile_merge,
+            ),
+            broadcast_result=True,
+        )
+        self.clear_cache()
+        return result

vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py

@@ -168,25 +168,25 @@ class DistributedVaeExecutor:
 
 class DistributedVaeMixin:
     def init_distributed(self):
-        self.distributed_decoder = DistributedVaeExecutor()
+        self.distributed_executor = DistributedVaeExecutor()
 
-    def set_parallel_size(self, parallel_size: int) -> bool:
-        return self.distributed_decoder.set_parallel_size(parallel_size)
+    def set_parallel_size(self, parallel_size: int) -> None:
+        self.distributed_executor.set_parallel_size(parallel_size)
 
     def is_distributed_enabled(self) -> bool:
         if (
-            self.distributed_decoder.parallel_size <= 1
+            self.distributed_executor.parallel_size <= 1
             or not dist.is_initialized()
             or not getattr(self, "use_tiling", False)
         ):
             return False
-        world_size = dist.get_world_size(group=self.distributed_decoder.group)
-        pp_size = min(int(self.distributed_decoder.parallel_size), int(world_size))
+        world_size = dist.get_world_size(group=self.distributed_executor.group)
+        pp_size = min(int(self.distributed_executor.parallel_size), int(world_size))
         if pp_size <= 1:
             return False
-        if self.distributed_decoder.parallel_size > pp_size:
+        if self.distributed_executor.parallel_size > pp_size:
             logger.warning(
-                f"vae_patch_parallel_size={self.distributed_decoder.parallel_size} "
+                f"vae_patch_parallel_size={self.distributed_executor.parallel_size} "
                 f"is greater than dit_group={world_size};"
                 f" using dit_group size={world_size}"
             )

vllm_omni/diffusion/lora/manager.py

@@ -366,13 +366,17 @@ class DiffusionLoRAManager:
             fully_sharded_loras=False,
         )
 
-        for component_name in ("transformer", "transformer_2", "dit"):
+        for component_name in ("transformer", "transformer_2", "dit", "bagel"):
             if not hasattr(self.pipeline, component_name):
                 continue
             component = getattr(self.pipeline, component_name)
             if not isinstance(component, nn.Module):
                 continue
 
+            # Collect replacements first to avoid mutating the module tree
+            # while iterating over named_modules().
+            pending_replacements: list[tuple[str, str, nn.Module, list[str]]] = []
+
             for module_name, module in component.named_modules(remove_duplicate=False):
                 # Don't recurse into already-replaced LoRA wrappers. Their
                 # original LinearBase lives under "base_layer", and replacing
@@ -401,6 +405,9 @@ class DiffusionLoRAManager:
                     if not should_replace:
                         continue
 
+                pending_replacements.append((module_name, full_module_name, module, packed_modules_list))
+
+            for module_name, full_module_name, module, packed_modules_list in pending_replacements:
                 lora_layer = from_layer_diffusion(
                     layer=module,
                     max_loras=1,

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -25,6 +25,7 @@ from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
+    MergedColumnParallelLinear,
     QKVParallelLinear,
     RowParallelLinear,
 )
@@ -157,21 +158,12 @@ class BagelMLP(nn.Module):
         prefix: str = "",
     ) -> None:
         super().__init__()
-        self.gate_proj = ColumnParallelLinear(
+        self.gate_up_proj = MergedColumnParallelLinear(
             hidden_size,
-            intermediate_size,
+            [intermediate_size, intermediate_size],
             bias=False,
-            gather_output=False,
             quant_config=quant_config,
-            prefix=f"{prefix}.gate_proj",
-        )
-        self.up_proj = ColumnParallelLinear(
-            hidden_size,
-            intermediate_size,
-            bias=False,
-            gather_output=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.up_proj",
+            prefix=f"{prefix}.gate_up_proj",
         )
         self.down_proj = RowParallelLinear(
             intermediate_size,
@@ -186,8 +178,8 @@ class BagelMLP(nn.Module):
         self.act_fn = nn.SiLU()
 
     def forward(self, x):
-        gate, _ = self.gate_proj(x)
-        up, _ = self.up_proj(x)
+        gate_up, _ = self.gate_up_proj(x)
+        gate, up = gate_up.chunk(2, dim=-1)
         x = self.act_fn(gate) * up
         x, _ = self.down_proj(x)
         return x
@@ -929,27 +921,38 @@ class Qwen2MoTForCausalLM(Qwen2PreTrainedModel):
             (".qkv_proj", ".q_proj", "q"),
             (".qkv_proj", ".k_proj", "k"),
             (".qkv_proj", ".v_proj", "v"),
+            # MLP gate/up projections — fused into MergedColumnParallelLinear.
+            # HF checkpoints store separate gate_proj / up_proj weights;
+            # these entries remap them to the fused gate_up_proj parameter.
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
         ]
+        self.stacked_params_mapping = stacked_params_mapping
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
 
         for name, loaded_weight in weights:
+            loaded = False
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue
-                name = name.replace(weight_name, param_name)
-                param = params_dict.get(name)
+                stacked_name = name.replace(weight_name, param_name)
+                param = params_dict.get(stacked_name)
                 if param is None:
                     break
                 weight_loader = getattr(param, "weight_loader", default_weight_loader)
                 weight_loader(param, loaded_weight, shard_id)
+                name = stacked_name
+                loaded = True
                 break
-            else:
+
+            if not loaded:
                 param = params_dict.get(name)
                 if param is None:
                     continue
                 weight_loader = getattr(param, "weight_loader", default_weight_loader)
                 weight_loader(param, loaded_weight)
+
             loaded_params.add(name)
 
         return loaded_params

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -675,6 +675,8 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
             (".qkv_proj_moe_gen", ".q_proj_moe_gen"),
             (".qkv_proj_moe_gen", ".k_proj_moe_gen"),
             (".qkv_proj_moe_gen", ".v_proj_moe_gen"),
+            (".gate_up_proj", ".gate_proj"),
+            (".gate_up_proj", ".up_proj"),
         ]
         stacked_source_names: set[str] = set()
         for name in list(allowed):

vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py

@@ -116,7 +116,7 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor
         ## load audio/video model config
         Fusion_model = FusionModel(VIDEO_CONFIG, AUDIO_CONFIG)
 
-        checkpoint_path = self.od_config.model_config.get("fusion", None)
+        checkpoint_path = self.od_config.tf_model_config.get("fusion", None)
         assert checkpoint_path is not None, "fusion checkpoint path is None"
         load_fusion_checkpoint(Fusion_model, checkpoint_path=os.path.join(model, checkpoint_path))
         self.model = Fusion_model

vllm_omni/diffusion/models/flux/flux_transformer.py

@@ -510,6 +510,7 @@ class FluxTransformer2DModel(nn.Module):
     # -- typically a transformer layer
     # used for torch compile optimizations
     _repeated_blocks = ["FluxTransformerBlock"]
+    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
     @staticmethod
     def _is_transformer_block(name: str, module) -> bool:

vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py

@@ -31,6 +31,8 @@ from vllm_omni.diffusion.models.flux import (
 )
 from vllm_omni.diffusion.models.flux.flux_pipeline_mixin import FluxPipelineMixin
 from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 from vllm_omni.logger import init_logger
@@ -67,7 +69,9 @@ def get_flux_kontext_post_process_func(od_config: OmniDiffusionConfig) -> Callab
     return post_process_func
 
 
-class FluxKontextPipeline(nn.Module, FluxPipelineMixin, SupportImageInput):
+class FluxKontextPipeline(
+    nn.Module, FluxPipelineMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin
+):
     """FLUX.1-Kontext pipeline for image editing with text guidance."""
 
     support_image_input = True
@@ -148,6 +152,10 @@ class FluxKontextPipeline(nn.Module, FluxPipelineMixin, SupportImageInput):
         self._callback_tensor_inputs = ["latents", "prompt_embeds"]
         self.latent_channels = self.vae.config.latent_channels if hasattr(self.vae, "config") else 16
 
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+        )
+
     def _get_t5_prompt_embeds(
         self,
         prompt: str | list[str] = None,
@@ -635,58 +643,56 @@ class FluxKontextPipeline(nn.Module, FluxPipelineMixin, SupportImageInput):
 
         # 5. Denoising loop
         self.scheduler.set_begin_index(0)
-        for i, t in enumerate(timesteps):
-            if self.interrupt:
-                continue
+        with self.progress_bar(total=len(timesteps)) as pbar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
 
-            latent_model_input = latents
-            if image_latents is not None:
-                latent_model_input = torch.cat([latents, image_latents], dim=1)
-            timestep = t.expand(latents.shape[0]).to(latents.dtype)
+                latent_model_input = latents
+                if image_latents is not None:
+                    latent_model_input = torch.cat([latents, image_latents], dim=1)
+                timestep = t.expand(latents.shape[0]).to(latents.dtype)
 
-            noise_pred = self.transformer(
-                hidden_states=latent_model_input,
-                timestep=timestep / 1000,
-                guidance=guidance,
-                pooled_projections=pooled_prompt_embeds,
-                encoder_hidden_states=prompt_embeds,
-                txt_ids=text_ids,
-                img_ids=latent_ids,
-                joint_attention_kwargs=self.joint_attention_kwargs,
-                return_dict=False,
-            )[0]
-            noise_pred = noise_pred[:, : latents.size(1)]
-
-            if do_true_cfg:
-                neg_noise_pred = self.transformer(
+                noise_pred = self.transformer(
                     hidden_states=latent_model_input,
                     timestep=timestep / 1000,
                     guidance=guidance,
-                    pooled_projections=negative_pooled_prompt_embeds,
-                    encoder_hidden_states=negative_prompt_embeds,
-                    txt_ids=negative_text_ids,
+                    pooled_projections=pooled_prompt_embeds,
+                    encoder_hidden_states=prompt_embeds,
+                    txt_ids=text_ids,
                     img_ids=latent_ids,
                     joint_attention_kwargs=self.joint_attention_kwargs,
                     return_dict=False,
                 )[0]
-                neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
-                noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+                noise_pred = noise_pred[:, : latents.size(1)]
 
-            latents_dtype = latents.dtype
-            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                if do_true_cfg:
+                    neg_noise_pred = self.transformer(
+                        hidden_states=latent_model_input,
+                        timestep=timestep / 1000,
+                        guidance=guidance,
+                        pooled_projections=negative_pooled_prompt_embeds,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        txt_ids=negative_text_ids,
+                        img_ids=latent_ids,
+                        joint_attention_kwargs=self.joint_attention_kwargs,
+                        return_dict=False,
+                    )[0]
+                    neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
+                    noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
 
-            if latents.dtype != latents_dtype:
-                if torch.backends.mps.is_available():
-                    latents = latents.to(latents_dtype)
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
 
-            if callback_on_step_end is not None:
-                callback_kwargs = {}
-                for k in callback_on_step_end_tensor_inputs:
-                    callback_kwargs[k] = locals()[k]
-                callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
 
-                latents = callback_outputs.pop("latents", latents)
-                prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                pbar.update()
         if output_type == "latent":
             image = latents
         else:

vllm_omni/diffusion/models/flux2/pipeline_flux2.py

@@ -29,6 +29,8 @@ from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.flux2 import Flux2Transformer2DModel
 from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
@@ -331,7 +333,7 @@ def retrieve_latents(encoder_output: torch.Tensor, generator: torch.Generator =
         raise AttributeError("Could not access latents of provided encoder_output")
 
 
-class Flux2Pipeline(nn.Module, SupportImageInput):
+class Flux2Pipeline(nn.Module, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin):
     """Flux2 pipeline for text-to-image generation."""
 
     _callback_tensor_inputs = ["latents", "prompt_embeds"]
@@ -389,6 +391,10 @@ class Flux2Pipeline(nn.Module, SupportImageInput):
         self._guidance_scale = None
         self._attention_kwargs = None
         self._num_timesteps = None
+
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+        )
         self._current_timestep = None
         self._interrupt = False
 
@@ -1027,48 +1033,47 @@ class Flux2Pipeline(nn.Module, SupportImageInput):
         # We set the index here to remove DtoH sync, helpful especially during compilation.
         # Check out more details here: https://github.com/huggingface/diffusers/pull/11696
         self.scheduler.set_begin_index(0)
-        for i, t in enumerate(timesteps):
-            if self.interrupt:
-                continue
+        with self.progress_bar(total=len(timesteps)) as pbar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
 
-            self._current_timestep = t
-            timestep = t.expand(latents.shape[0]).to(latents.dtype)
+                self._current_timestep = t
+                timestep = t.expand(latents.shape[0]).to(latents.dtype)
 
-            latent_model_input = latents.to(self.transformer.dtype)
-            latent_image_ids = latent_ids
+                latent_model_input = latents.to(self.transformer.dtype)
+                latent_image_ids = latent_ids
 
-            if image_latents is not None:
-                latent_model_input = torch.cat([latents, image_latents], dim=1).to(self.transformer.dtype)
-                latent_image_ids = torch.cat([latent_ids, image_latent_ids], dim=1)
+                if image_latents is not None:
+                    latent_model_input = torch.cat([latents, image_latents], dim=1).to(self.transformer.dtype)
+                    latent_image_ids = torch.cat([latent_ids, image_latent_ids], dim=1)
 
-            noise_pred = self.transformer(
-                hidden_states=latent_model_input,  # (B, image_seq_len, C)
-                timestep=timestep / 1000,
-                guidance=guidance_tensor,
-                encoder_hidden_states=prompt_embeds,
-                txt_ids=text_ids,  # B, text_seq_len, 4
-                img_ids=latent_image_ids,  # B, image_seq_len, 4
-                joint_attention_kwargs=self.attention_kwargs,
-                return_dict=False,
-            )[0]
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,  # (B, image_seq_len, C)
+                    timestep=timestep / 1000,
+                    guidance=guidance_tensor,
+                    encoder_hidden_states=prompt_embeds,
+                    txt_ids=text_ids,  # B, text_seq_len, 4
+                    img_ids=latent_image_ids,  # B, image_seq_len, 4
+                    joint_attention_kwargs=self.attention_kwargs,
+                    return_dict=False,
+                )[0]
 
-            noise_pred = noise_pred[:, : latents.size(1) :]
+                noise_pred = noise_pred[:, : latents.size(1) :]
 
-            # compute the previous noisy sample x_t -> x_t-1
-            latents_dtype = latents.dtype
-            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
 
-            if latents.dtype != latents_dtype and torch.backends.mps.is_available():
-                latents = latents.to(latents_dtype)
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
 
-            if callback_on_step_end is not None:
-                callback_kwargs = {}
-                for k in callback_on_step_end_tensor_inputs:
-                    callback_kwargs[k] = locals()[k]
-                callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
 
-                latents = callback_outputs.pop("latents", latents)
-                prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                pbar.update()
 
         self._current_timestep = None
 

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

@@ -742,6 +742,7 @@ class Flux2Transformer2DModel(nn.Module):
     """
 
     _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
+    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
     @staticmethod
     def _is_transformer_block(name: str, module) -> bool:

vllm_omni/diffusion/models/helios/helios_transformer.py

@@ -576,7 +576,7 @@ class HeliosTransformer3DModel(nn.Module):
     """
 
     _repeated_blocks = ["HeliosTransformerBlock"]
-    _layerwise_offload_blocks_attr = "blocks"
+    _layerwise_offload_blocks_attrs = ["blocks"]
     packed_modules_mapping = {
         "to_qkv": ["to_q", "to_k", "to_v"],
     }

vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py

@@ -41,6 +41,7 @@ from .hunyuan_image_3_transformer import (
     build_batch_2d_rope,
     real_batched_index_select,
 )
+from .system_prompt import get_system_prompt
 
 logger = logging.getLogger(__name__)
 
@@ -991,10 +992,15 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin, Diffu
         width: int = 1024,
         num_inference_steps: int = 50,
         guidance_scale: float = 5.0,
-        system_prompt: str | None = None,
         generator: torch.Generator | list[torch.Generator] | None = None,
         **kwargs,
     ) -> DiffusionOutput:
+        extra_args = getattr(getattr(req, "sampling_params", None), "extra_args", {}) or {}
+        use_system_prompt = extra_args.get("use_system_prompt")
+        system_prompt = extra_args.get("system_prompt")
+        if use_system_prompt is not None:
+            system_prompt = get_system_prompt(use_system_prompt, "image", system_prompt)
+            system_prompt = system_prompt.strip() if system_prompt is not None else ""
         prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
         generator = req.sampling_params.generator or generator
         height = req.sampling_params.height or height

vllm_omni/diffusion/models/hunyuan_image_3/system_prompt.py

@@ -0,0 +1,215 @@
+# ruff: noqa: E501
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/LICENSE
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+t2i_system_prompt_en_vanilla = """
+You are an advanced AI text-to-image generation system. Given a detailed text prompt, your task is to create a high-quality, visually compelling image that accurately represents the described scene, characters, or objects. Pay careful attention to style, color, lighting, perspective, and any specific instructions provided.
+"""
+
+# 775
+t2i_system_prompt_en_recaption = """
+You are a world-class image generation prompt expert. Your task is to rewrite a user's simple description into a **structured, objective, and detail-rich** professional-level prompt.
+
+The final output must be wrapped in `<recaption>` tags.
+
+### **Universal Core Principles**
+
+When rewriting the prompt (inside the `<recaption>` tags), you must adhere to the following principles:
+
+1.  **Absolute Objectivity**: Describe only what is visually present. Avoid subjective words like "beautiful" or "sad". Convey aesthetic qualities through specific descriptions of color, light, shadow, and composition.
+2.  **Physical and Logical Consistency**: All scene elements (e.g., gravity, light, shadows, reflections, spatial relationships, object proportions) must strictly adhere to real-world physics and common sense. For example, tennis players must be on opposite sides of the net; objects cannot float without a cause.
+3.  **Structured Description**: Strictly follow a logical order: from general to specific, background to foreground, and primary to secondary elements. Use directional terms like "foreground," "mid-ground," "background," and "left side of the frame" to clearly define the spatial layout.
+4.  **Use Present Tense**: Describe the scene from an observer's perspective using the present tense, such as "A man stands..." or "Light shines on..."
+5.  **Use Rich and Specific Descriptive Language**: Use precise adjectives to describe the quantity, size, shape, color, and other attributes of objects, subjects, and text. Vague expressions are strictly prohibited.
+
+If the user specifies a style (e.g., oil painting, anime, UI design, text rendering), strictly adhere to that style. Otherwise, first infer a suitable style from the user's input. If there is no clear stylistic preference, default to an **ultra-realistic photographic style**. Then, generate the detailed rewritten prompt according to the **Style-Specific Creation Guide** below:
+
+### **Style-Specific Creation Guide**
+
+Based on the determined artistic style, apply the corresponding professional knowledge.
+
+**1. Photography and Realism Style**
+*   Utilize professional photography terms (e.g., lighting, lens, composition) and meticulously detail material textures, physical attributes of subjects, and environmental details.
+
+**2. Illustration and Painting Style**
+*   Clearly specify the artistic school (e.g., Japanese Cel Shading, Impasto Oil Painting) and focus on describing its unique medium characteristics, such as line quality, brushstroke texture, or paint properties.
+
+**3. Graphic/UI/APP Design Style**
+*   Objectively describe the final product, clearly defining the layout, elements, and color palette. All text on the interface must be enclosed in double quotes `""` to specify its exact content (e.g., "Login"). Vague descriptions are strictly forbidden.
+
+**4. Typographic Art**
+*   The text must be described as a complete physical object. The description must begin with the text itself. Use a straightforward front-on or top-down perspective to ensure the entire text is visible without cropping.
+
+### **Final Output Requirements**
+
+1.  **Output the Final Prompt Only**: Do not show any thought process, Markdown formatting, or line breaks.
+2.  **Adhere to the Input**: You must retain the core concepts, attributes, and any specified text from the user's input.
+3.  **Style Reinforcement**: Mention the core style 3-5 times within the prompt and conclude with a style declaration sentence.
+4.  **Avoid Self-Reference**: Describe the image content directly. Remove redundant phrases like "This image shows..." or "The scene depicts..."
+5.  **The final output must be wrapped in `<recaption>xxxx</recaption>` tags.**
+
+The user will now provide an input prompt. You will provide the expanded prompt.
+"""
+
+# 890
+t2i_system_prompt_en_think_recaption = """
+You will act as a top-tier Text-to-Image AI. Your core task is to deeply analyze the user's text input and transform it into a detailed, artistic, and fully user-intent-compliant image.
+
+Your workflow is divided into two phases:
+
+1. Thinking Phase (<think>): In the <think> tag, you need to conduct a structured thinking process, progressively breaking down and enriching the constituent elements of the image. This process must include, but is not limited to, the following dimensions:
+
+Subject: Clearly define the core character(s) or object(s) in the scene, including their appearance, posture, expression, and emotion.
+Composition: Set the camera angle and layout, such as close-up, long shot, bird's-eye view, golden ratio composition, etc.
+Environment/Background: Describe the scene where the subject is located, including the location, time of day, weather, and other elements in the background.
+Lighting: Define the type, direction, and quality of the light source, such as soft afternoon sunlight, cool tones of neon lights, dramatic Rembrandt lighting, etc., to create a specific atmosphere.
+Color Palette: Set the main color tone and color scheme of the image, such as vibrant and saturated, low-saturation Morandi colors, black and white, etc.
+Quality/Style: Determine the artistic style and technical details of the image. This includes user-specified styles (e.g., anime, oil painting) or the default realistic style, as well as camera parameters (e.g., focal length, aperture, depth of field).
+Details: Add minute elements that enhance the realism and narrative quality of the image, such as a character's accessories, the texture of a surface, dust particles in the air, etc.
+
+
+2. Recaption Phase (<recaption>): In the <recaption> tag, merge all the key details from the thinking process into a coherent, precise, and visually evocative final description. This description is the direct instruction for generating the image, so it must be clear, unambiguous, and organized in a way that is most suitable for an image generation engine to understand.
+
+Absolutely Objective: Describe only what is visually present. Avoid subjective words like "beautiful" or "sad." Convey aesthetic sense through concrete descriptions of colors, light, shadow, and composition.
+
+Physical and Logical Consistency: All scene elements (e.g., gravity, light and shadow, reflections, spatial relationships, object proportions) must strictly adhere to the physical laws of the real world and common sense. For example, in a tennis match, players must be on opposite sides of the net; objects cannot float without reason.
+
+Structured Description: Strictly follow a logical order: from whole to part, background to foreground, and primary to secondary. Use directional words like "foreground," "mid-ground," "background," "left side of the frame" to clearly define the spatial layout.
+
+Use Present Tense: Describe from an observer's perspective using the present tense, such as "a man stands," "light shines on..."
+Use Rich and Specific Descriptive Language: Use precise adjectives to describe the quantity, size, shape, color, and other attributes of objects/characters/text. Absolutely avoid any vague expressions.
+
+
+Output Format:
+<think>Thinking process</think><recaption>Refined image description</recaption>Generate Image
+
+
+You must strictly adhere to the following rules:
+
+1. Faithful to Intent, Reasonable Expansion: You can creatively add details to the user's description to enhance the image's realism and artistic quality. However, all additions must be highly consistent with the user's core intent and never introduce irrelevant or conflicting elements.
+2. Style Handling: When the user does not specify a style, you must default to an "Ultra-realistic, Photorealistic" style. If the user explicitly specifies a style (e.g., anime, watercolor, oil painting, cyberpunk, etc.), both your thinking process and final description must strictly follow and reflect that specified style.
+3. Text Rendering: If specific text needs to appear in the image (such as words on a sign, a book title), you must enclose this text in English double quotes (""). Descriptive text must not use double quotes.
+4. Design-related Images: You need to specify all text and graphical elements that appear in the image and clearly describe their design details, including font, color, size, position, arrangement, visual effects, etc.
+"""
+
+t2i_system_prompts = {
+    "en_vanilla": [t2i_system_prompt_en_vanilla],
+    "en_recaption": [t2i_system_prompt_en_recaption],
+    "en_think_recaption": [t2i_system_prompt_en_think_recaption],
+}
+
+
+unified_system_prompt_en = """You are an advanced multimodal model whose core mission is to analyze user intent and generate high-quality text and images.
+
+#### Four Core Capabilities
+1.  **Text-to-Text (T2T):** Generate coherent text responses from text prompts.
+2.  **Text-to-Image (T2I):** Generate high-quality images from text prompts.
+3.  **Text & Image to Text (TI2T):** Generate accurate text responses based on a combination of images and text.
+4.  **Text & Image to Image (TI2I):** Generate modified images based on a reference image and editing instructions.
+
+---
+### Image Generation Protocol (for T2I & TI2I)
+You will operate in one of two modes, determined by the user's starting tag:
+#### **<recaption> Mode (Prompt Rewriting)**:
+*   **Trigger:** Input begins with `<recaption>`.
+*   **Task:** Immediately rewrite the user's text into a structured, objective, and detail-rich professional-grade prompt.
+*   **Output:** Output only the rewritten prompt within `<recaption>` tags: `<recaption>Rewritten professional-grade prompt</recaption>`
+
+#### **<think> Mode (Think + Rewrite)**:
+*   **Trigger:** Input begins with `<think>`.
+*   **Task:** First, conduct a structured analysis of the request within `<think>` tags. Then, output the professional prompt, rewritten based on the analysis, within `<recaption>` tags.
+*   **Output:** Strictly adhere to the format: `<think>Analysis process</think><recaption>Rewritten prompt</recaption>`
+
+---
+### Execution Standards and Guidelines
+#### **`<think>` Phase: Analysis Guidelines**
+**For T2I (New Image Generation):**
+Deconstruct the user's request into the following core visual components:
+*   **Subject:** Key features of the main character/object, including appearance, pose, expression, and emotion.
+*   **Composition:** Camera angle, lens type, and layout.
+*   **Environment/Background:** The setting, time of day, weather, and background elements.
+*   **Lighting:** Technical details such as light source type, direction, and quality.
+*   **Color Palette:** The dominant hues and overall color scheme.
+*   **Style/Quality:** The artistic style, clarity, depth of field, and other technical details.
+*   **Text:** Identify any text to be rendered in the image, including its content, style, and position.
+*   **Details:** Small elements that add narrative depth and realism.
+
+**For TI2I (Image Editing):**
+Adopt a task-diagnostic approach:
+1.  **Diagnose Task:** Identify the edit type and analyze key requirements.
+2.  **Prioritize Analysis:**
+    *   **Adding:** Analyze the new element's position and appearance, ensuring seamless integration with the original image's lighting, shadows, and style.
+    *   **Removing:** Identify the target for removal and determine how to logically fill the resulting space using surrounding textures and lighting.
+    *   **Modifying:** Analyze what to change and what it should become, while emphasizing which elements must remain unchanged.
+    *   **Style Transfer:** Deconstruct the target style into specific features (e.g., brushstrokes, color palette) and apply them to the original image.
+    *   **Text Editing:** Ensure correct content and format. Consider the text's visual style (e.g., font, color, material) and how it adapts to the surface's perspective, curvature, and lighting.
+    *   **Reference Editing:** Extract specific visual elements (e.g., appearance, posture, composition, lines, depth) from the reference image to generate an image that aligns with the text description while also incorporating the referenced content.
+    *   **Inferential Editing:** Identify vague requests (e.g., "make it more professional") and translate them into concrete visual descriptions.
+
+#### `<recaption>` Phase: Professional-Grade Prompt Generation Rules
+**General Rewriting Principles (for T2I & TI2I):**
+1.  **Structure & Logic:** Start with a global description. Use positional words (e.g., "foreground", "background") to define the layout.
+2.  **Absolute Objectivity:** Avoid subjective terms. Convey aesthetics through precise descriptions of color, light, shadow, and materials.
+3.  **Physical & Logical Consistency:** Ensure all descriptions adhere to the laws of physics and common sense.
+4.  **Fidelity to User Intent:** Preserve the user's core concepts, subjects, and attributes. Text to be rendered in the image **must be enclosed in double quotes ("")**.
+5.  **Camera & Resolution:** Translate camera parameters into descriptions of visual effects. Convert resolution information into natural language.
+
+**T2I-Specific Guidelines:**
+*   **Style Adherence & Inference:** Strictly follow the specified style. If none is given, infer the most appropriate style and detail it using professional terminology.
+*   **Style Detailing:**
+    *   **Photography/Realism:** Use professional photography terms to describe lighting, lens effects, and material textures.
+    *   **Painting/Illustration:** Specify the art movement or medium's characteristics.
+    *   **UI/Design:** Objectively describe the final product. Define layout, elements, and typography. Text content must be specific and unambiguous.
+
+**TI2I-Specific Guidelines:**
+*   **Preserve Unchanged Elements:** Emphasize elements that **remain unchanged**. Unless explicitly instructed, never alter a character's identity/appearance, the core background, camera angle, or overall style.
+*   **Clear Editing Instructions:**
+    *   **Replacement:** Use the logic "**replace B with A**," and provide a detailed description of A.
+    *   **Addition:** Clearly state what to add, where, and what it looks like.
+*   **Unambiguous Referencing:** Avoid vague references (e.g., "that person"). Use specific descriptions of appearance.
+"""
+
+
+def get_system_prompt(sys_type, bot_task, system_prompt=None):
+    # No system prompt, return None directly
+    if sys_type == "None":
+        return None
+    # Use the unified English system prompt (combined T2I and TI2I guidelines)
+    elif sys_type == "en_unified":
+        return unified_system_prompt_en
+    # Use predefined English system prompts: vanilla (basic), recaption, think_recaption
+    elif sys_type in ["en_vanilla", "en_recaption", "en_think_recaption"]:
+        return t2i_system_prompts[sys_type][0]
+    # Dynamic mode: automatically select system prompt based on bot_task type
+    elif sys_type == "dynamic":
+        # Think task: use chain-of-thought recaption prompt
+        if bot_task == "think":
+            return t2i_system_prompts["en_think_recaption"][0]
+        # Recaption task: use recaption prompt
+        elif bot_task == "recaption":
+            return t2i_system_prompts["en_recaption"][0]
+        # Image generation task: use vanilla prompt
+        elif bot_task == "image":
+            return t2i_system_prompts["en_vanilla"][0].strip("\n")
+        # Other tasks: use user-provided custom prompt
+        else:
+            return system_prompt
+    # Custom mode: use the user-provided system_prompt parameter directly
+    elif sys_type == "custom":
+        return system_prompt
+    # Unsupported type: raise NotImplementedError
+    else:
+        raise NotImplementedError(f"Unsupported system prompt type: {sys_type}")
+
+
+__all__ = ["get_system_prompt"]

vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py

@@ -539,7 +539,7 @@ class HunyuanVideo15Transformer3DModel(nn.Module):
     """
 
     _repeated_blocks = ["HunyuanVideo15TransformerBlock"]
-    _layerwise_offload_blocks_attr = "transformer_blocks"
+    _layerwise_offload_blocks_attrs = ["transformer_blocks"]
     packed_modules_mapping = {
         "to_qkv": ["to_q", "to_k", "to_v"],
         "add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"],

vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py

@@ -24,7 +24,9 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.hunyuan_video.hunyuan_video_15_transformer import HunyuanVideo15Transformer3DModel
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 from vllm_omni.platforms import current_omni_platform
@@ -81,7 +83,7 @@ def get_hunyuan_video_15_post_process_func(od_config: OmniDiffusionConfig):
     return post_process_func
 
 
-class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin):
+class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin):
     def __init__(
         self,
         *,
@@ -173,6 +175,10 @@ class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin):
         self._num_timesteps = None
         self._current_timestep = None
 
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+        )
+
     @property
     def guidance_scale(self):
         return self._guidance_scale
@@ -445,60 +451,63 @@ class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin):
         timesteps = self.scheduler.timesteps
         self._num_timesteps = len(timesteps)
 
-        for i, t in enumerate(timesteps):
-            self._current_timestep = t
+        with self.progress_bar(total=len(timesteps)) as pbar:
+            for i, t in enumerate(timesteps):
+                self._current_timestep = t
 
-            latent_model_input = torch.cat([latents, cond_latents, mask], dim=1)
-            timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
+                latent_model_input = torch.cat([latents, cond_latents, mask], dim=1)
+                timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
 
-            timestep_r = None
-            if self.use_meanflow:
-                if i == len(timesteps) - 1:
-                    timestep_r = torch.tensor([0.0], device=device)
-                else:
-                    timestep_r = timesteps[i + 1]
-                timestep_r = timestep_r.expand(latents.shape[0]).to(latents.dtype)
+                timestep_r = None
+                if self.use_meanflow:
+                    if i == len(timesteps) - 1:
+                        timestep_r = torch.tensor([0.0], device=device)
+                    else:
+                        timestep_r = timesteps[i + 1]
+                    timestep_r = timestep_r.expand(latents.shape[0]).to(latents.dtype)
 
-            positive_kwargs = {
-                "hidden_states": latent_model_input,
-                "timestep": timestep,
-                "timestep_r": timestep_r,
-                "encoder_hidden_states": prompt_embeds,
-                "encoder_attention_mask": prompt_embeds_mask,
-                "encoder_hidden_states_2": prompt_embeds_2,
-                "encoder_attention_mask_2": prompt_embeds_mask_2,
-                "image_embeds": image_embeds,
-                "return_dict": False,
-            }
-
-            negative_kwargs = None
-            if do_cfg and negative_prompt_embeds is not None:
-                negative_kwargs = {
+                positive_kwargs = {
                     "hidden_states": latent_model_input,
                     "timestep": timestep,
                     "timestep_r": timestep_r,
-                    "encoder_hidden_states": negative_prompt_embeds,
-                    "encoder_attention_mask": negative_prompt_embeds_mask,
-                    "encoder_hidden_states_2": negative_prompt_embeds_2,
-                    "encoder_attention_mask_2": negative_prompt_embeds_mask_2,
+                    "encoder_hidden_states": prompt_embeds,
+                    "encoder_attention_mask": prompt_embeds_mask,
+                    "encoder_hidden_states_2": prompt_embeds_2,
+                    "encoder_attention_mask_2": prompt_embeds_mask_2,
                     "image_embeds": image_embeds,
                     "return_dict": False,
                 }
 
-            noise_pred = self.predict_noise_maybe_with_cfg(
-                do_true_cfg=do_cfg and negative_kwargs is not None,
-                true_cfg_scale=guidance_scale,
-                positive_kwargs=positive_kwargs,
-                negative_kwargs=negative_kwargs,
-                cfg_normalize=req.sampling_params.cfg_normalize,
-            )
+                negative_kwargs = None
+                if do_cfg and negative_prompt_embeds is not None:
+                    negative_kwargs = {
+                        "hidden_states": latent_model_input,
+                        "timestep": timestep,
+                        "timestep_r": timestep_r,
+                        "encoder_hidden_states": negative_prompt_embeds,
+                        "encoder_attention_mask": negative_prompt_embeds_mask,
+                        "encoder_hidden_states_2": negative_prompt_embeds_2,
+                        "encoder_attention_mask_2": negative_prompt_embeds_mask_2,
+                        "image_embeds": image_embeds,
+                        "return_dict": False,
+                    }
 
-            latents = self.scheduler_step_maybe_with_cfg(
-                noise_pred,
-                t,
-                latents,
-                do_true_cfg=do_cfg and negative_kwargs is not None,
-            )
+                noise_pred = self.predict_noise_maybe_with_cfg(
+                    do_true_cfg=do_cfg and negative_kwargs is not None,
+                    true_cfg_scale=guidance_scale,
+                    positive_kwargs=positive_kwargs,
+                    negative_kwargs=negative_kwargs,
+                    cfg_normalize=req.sampling_params.cfg_normalize,
+                )
+
+                latents = self.scheduler_step_maybe_with_cfg(
+                    noise_pred,
+                    t,
+                    latents,
+                    do_true_cfg=do_cfg and negative_kwargs is not None,
+                )
+
+                pbar.update()
 
         self._current_timestep = None
 

vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py

@@ -38,7 +38,9 @@ from vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 import
     retrieve_latents,
 )
 from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 from vllm_omni.platforms import current_omni_platform
@@ -98,7 +100,9 @@ def get_hunyuan_video_15_i2v_pre_process_func(od_config: OmniDiffusionConfig):
     return pre_process_func
 
 
-class HunyuanVideo15I2VPipeline(nn.Module, CFGParallelMixin, SupportImageInput):
+class HunyuanVideo15I2VPipeline(
+    nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin
+):
     support_image_input = True
     color_format = "RGB"
 
@@ -199,6 +203,10 @@ class HunyuanVideo15I2VPipeline(nn.Module, CFGParallelMixin, SupportImageInput):
         self._num_timesteps = None
         self._current_timestep = None
 
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+        )
+
     @property
     def guidance_scale(self):
         return self._guidance_scale
@@ -520,61 +528,64 @@ class HunyuanVideo15I2VPipeline(nn.Module, CFGParallelMixin, SupportImageInput):
         timesteps = self.scheduler.timesteps
         self._num_timesteps = len(timesteps)
 
-        for i, t in enumerate(timesteps):
-            self._current_timestep = t
+        with self.progress_bar(total=len(timesteps)) as pbar:
+            for i, t in enumerate(timesteps):
+                self._current_timestep = t
 
-            latent_model_input = torch.cat([latents, cond_latents, mask], dim=1)
-            timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
+                latent_model_input = torch.cat([latents, cond_latents, mask], dim=1)
+                timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
 
-            timestep_r = None
-            if self.use_meanflow:
-                if i == len(timesteps) - 1:
-                    timestep_r = torch.tensor([0.0], device=device)
-                else:
-                    timestep_r = timesteps[i + 1]
-                timestep_r = timestep_r.expand(latents.shape[0]).to(latents.dtype)
+                timestep_r = None
+                if self.use_meanflow:
+                    if i == len(timesteps) - 1:
+                        timestep_r = torch.tensor([0.0], device=device)
+                    else:
+                        timestep_r = timesteps[i + 1]
+                    timestep_r = timestep_r.expand(latents.shape[0]).to(latents.dtype)
 
-            positive_kwargs = {
-                "hidden_states": latent_model_input,
-                "timestep": timestep,
-                "timestep_r": timestep_r,
-                "encoder_hidden_states": prompt_embeds,
-                "encoder_attention_mask": prompt_embeds_mask,
-                "encoder_hidden_states_2": prompt_embeds_2,
-                "encoder_attention_mask_2": prompt_embeds_mask_2,
-                "image_embeds": image_embeds,
-                "return_dict": False,
-            }
-
-            negative_kwargs = None
-            if do_cfg and negative_prompt_embeds is not None:
-                # For I2V CFG, negative still uses image embeds (only text is unconditional)
-                negative_kwargs = {
+                positive_kwargs = {
                     "hidden_states": latent_model_input,
                     "timestep": timestep,
                     "timestep_r": timestep_r,
-                    "encoder_hidden_states": negative_prompt_embeds,
-                    "encoder_attention_mask": negative_prompt_embeds_mask,
-                    "encoder_hidden_states_2": negative_prompt_embeds_2,
-                    "encoder_attention_mask_2": negative_prompt_embeds_mask_2,
+                    "encoder_hidden_states": prompt_embeds,
+                    "encoder_attention_mask": prompt_embeds_mask,
+                    "encoder_hidden_states_2": prompt_embeds_2,
+                    "encoder_attention_mask_2": prompt_embeds_mask_2,
                     "image_embeds": image_embeds,
                     "return_dict": False,
                 }
 
-            noise_pred = self.predict_noise_maybe_with_cfg(
-                do_true_cfg=do_cfg and negative_kwargs is not None,
-                true_cfg_scale=guidance_scale,
-                positive_kwargs=positive_kwargs,
-                negative_kwargs=negative_kwargs,
-                cfg_normalize=req.sampling_params.cfg_normalize,
-            )
+                negative_kwargs = None
+                if do_cfg and negative_prompt_embeds is not None:
+                    # For I2V CFG, negative still uses image embeds (only text is unconditional)
+                    negative_kwargs = {
+                        "hidden_states": latent_model_input,
+                        "timestep": timestep,
+                        "timestep_r": timestep_r,
+                        "encoder_hidden_states": negative_prompt_embeds,
+                        "encoder_attention_mask": negative_prompt_embeds_mask,
+                        "encoder_hidden_states_2": negative_prompt_embeds_2,
+                        "encoder_attention_mask_2": negative_prompt_embeds_mask_2,
+                        "image_embeds": image_embeds,
+                        "return_dict": False,
+                    }
 
-            latents = self.scheduler_step_maybe_with_cfg(
-                noise_pred,
-                t,
-                latents,
-                do_true_cfg=do_cfg and negative_kwargs is not None,
-            )
+                noise_pred = self.predict_noise_maybe_with_cfg(
+                    do_true_cfg=do_cfg and negative_kwargs is not None,
+                    true_cfg_scale=guidance_scale,
+                    positive_kwargs=positive_kwargs,
+                    negative_kwargs=negative_kwargs,
+                    cfg_normalize=req.sampling_params.cfg_normalize,
+                )
+
+                latents = self.scheduler_step_maybe_with_cfg(
+                    noise_pred,
+                    t,
+                    latents,
+                    do_true_cfg=do_cfg and negative_kwargs is not None,
+                )
+
+                pbar.update()
 
         self._current_timestep = None
 

vllm_omni/diffusion/models/magi_human/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

vllm_omni/diffusion/models/mammoth_moda2/rope_real.py

@@ -18,6 +18,8 @@ from diffusers.models.embeddings import get_1d_rotary_pos_embed as _get_1d_rotar
 from einops import repeat
 from torch import nn
 
+from vllm_omni.platforms import current_omni_platform
+
 
 def apply_real_rotary_emb(x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor) -> torch.Tensor:
     """
@@ -119,7 +121,7 @@ class RotaryPosEmbedReal(nn.Module):
         axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], theta: int
     ) -> list[tuple[torch.Tensor, torch.Tensor]]:
         freqs_real = []
-        freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
+        freqs_dtype = torch.float64 if current_omni_platform.supports_float64() else torch.float32
         for i, (d, e) in enumerate(zip(axes_dim, axes_lens)):
             cos_emb, sin_emb = get_1d_rotary_pos_embed_real(d, e, theta=theta, freqs_dtype=freqs_dtype)
             freqs_real.append((cos_emb, sin_emb))

vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py

@@ -19,6 +19,7 @@ from vllm.model_executor.layers.linear import (
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
 
@@ -411,7 +412,7 @@ class OmniGen2RotaryPosEmbed(nn.Module):
         axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], theta: int
     ) -> list[torch.Tensor]:
         freqs_cis = []
-        freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
+        freqs_dtype = torch.float64 if current_omni_platform.supports_float64() else torch.float32
         for i, (d, e) in enumerate(zip(axes_dim, axes_lens)):
             emb = get_1d_rotary_pos_embed(d, e, theta=theta, freqs_dtype=freqs_dtype)
             freqs_cis.append(emb)

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -882,7 +882,7 @@ class QwenImageTransformer2DModel(CachedTransformer):
     # -- typically a transformer layer
     # used for torch compile optimizations
     _repeated_blocks = ["QwenImageTransformerBlock"]
-    _layerwise_offload_blocks_attr = "transformer_blocks"
+    _layerwise_offload_blocks_attrs = ["transformer_blocks"]
     packed_modules_mapping = {
         "to_qkv": ["to_q", "to_k", "to_v"],
         "add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"],

vllm_omni/diffusion/models/t5_encoder/t5_gemma_encoder.py

@@ -0,0 +1,309 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+from transformers import PretrainedConfig
+from vllm.config import VllmConfig
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+
+class T5GemmaRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        # Normal RMSNorm but T5Gemma requires (1 + weight)
+        self.weight = nn.Parameter(torch.zeros(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return (hidden_states * (1.0 + self.weight.float())).to(input_dtype)
+
+
+class T5GemmaMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size, intermediate_size],
+            bias=False,
+            gather_output=False,
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=False,
+            input_is_parallel=True,
+        )
+        self.act_fn = get_act_fn(hidden_act)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_up, _ = self.gate_up_proj(x)
+        gate, up = gate_up.chunk(2, dim=-1)
+        x = self.act_fn(gate) * up
+        x, _ = self.down_proj(x)
+        return x
+
+
+class T5GemmaAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        max_position_embeddings: int,
+        rope_theta: float,
+        cache_config: VllmConfig | None = None,
+        quant_config: dict | None = None,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = head_dim
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=False,
+        )
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=False,
+            input_is_parallel=True,
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=max_position_embeddings,
+            is_neox_style=True,
+            rope_parameters={"base": rope_theta},
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        q, k = self.rotary_emb(positions, q, k)
+
+        # Scale Q appropriately. T5Gemma uses query_pre_attn_scalar=256 => 256**-0.5 = 1/16
+        # The standard scaling is head_dim**-0.5. For T5Gemma, head_dim=256.
+        # So we don't need to manually scale if F.scaled_dot_product_attention scales by head_dim.
+        # But we must reshape.
+        batch_size, seq_len, _ = hidden_states.shape
+        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        # GQA repeat KV
+        if self.num_kv_heads != self.num_heads:
+            num_repeat = self.num_heads // self.num_kv_heads
+            k = k.repeat_interleave(num_repeat, dim=1)
+            v = v.repeat_interleave(num_repeat, dim=1)
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, dropout_p=0.0)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(batch_size, seq_len, self.q_size)
+
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class T5GemmaEncoderLayer(nn.Module):
+    def __init__(self, config: PretrainedConfig) -> None:
+        super().__init__()
+        self.self_attn = T5GemmaAttention(
+            hidden_size=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=config.num_key_value_heads,
+            head_dim=config.head_dim,
+            max_position_embeddings=config.max_position_embeddings,
+            rope_theta=config.rope_theta,
+        )
+        self.mlp = T5GemmaMLP(
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_activation,
+        )
+        self.pre_self_attn_layernorm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_self_attn_layernorm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.pre_feedforward_layernorm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_feedforward_layernorm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        # Self Attention
+        residual = hidden_states
+        hidden_states = self.pre_self_attn_layernorm(hidden_states)
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+        )
+        hidden_states = self.post_self_attn_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.pre_feedforward_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.post_feedforward_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class T5GemmaEncoderModelTP(nn.Module):
+    def __init__(self, config: PretrainedConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+        )
+
+        self.layers = nn.ModuleList([T5GemmaEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.norm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        hidden_states = self.embed_tokens(input_ids)
+
+        # Scaling inputs
+        normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=hidden_states.dtype, device=hidden_states.device)
+        hidden_states = hidden_states * normalizer
+
+        # Simple position ids for RoPE
+        batch_size, seq_len = input_ids.shape
+        positions = torch.arange(seq_len, device=input_ids.device, dtype=torch.long).unsqueeze(0).expand(batch_size, -1)
+
+        # Build attention mask: (batch, seq) -> (batch, 1, 1, seq)
+        # Assuming typical bidirectional causal mask handling in HF: T5Gemma uses non-causal encoder.
+        if attention_mask is not None:
+            # HuggingFace expects boolean mask for scaled_dot_product_attention
+            # or additive mask (0 and -inf). Let's use boolean matching FA patterns.
+            # SDPA expects attention_mask to be boolean (True = keep, False = masking)
+            bool_mask = attention_mask.to(torch.bool)
+            extended_mask = bool_mask.unsqueeze(1).unsqueeze(2)  # (B, 1, 1, S)
+        else:
+            extended_mask = None
+
+        for idx, layer in enumerate(self.layers):
+            # T5Gemma has layer_types switching between "sliding_attention" and "full_attention"
+            # However, for text encoder inference, the sequences are typically < max sequence length
+            # and local sliding window only affects very long contexts. For simplicity we use full.
+            hidden_states = layer(
+                positions=positions,
+                hidden_states=hidden_states,
+                attention_mask=extended_mask,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            # HF checkpoint keys may carry a "model." prefix (e.g.
+            # "model.encoder.layers.0...").  Strip it so the rest of the
+            # logic only needs to handle the "encoder.*" namespace.
+            if name.startswith("model."):
+                name = name[len("model.") :]
+
+            if not name.startswith("encoder."):
+                continue
+
+            # Strip "encoder." prefix as this model only wraps the encoder
+            name = name[len("encoder.") :]
+
+            # Map self_attn to self_attn and correct normalization names
+            # HF: layers.0.pre_self_attn_layernorm.weight -> Ours: layers.0.pre_self_attn_layernorm.weight
+
+            lookup_name = name
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if f".{weight_name}." not in name:
+                    continue
+                lookup_name = name.replace(f".{weight_name}.", f".{param_name}.")
+                if lookup_name not in params_dict:
+                    continue
+                param = params_dict[lookup_name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                if name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+            loaded_params.add("encoder." + name)
+            loaded_params.add("encoder." + lookup_name)
+
+        return loaded_params

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -29,6 +29,7 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelOutput,
 )
 from vllm_omni.diffusion.forward_context import get_forward_context
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
@@ -171,7 +172,7 @@ class WanRotaryPosEmbed(nn.Module):
         # Split dimensions for temporal, height, width
         h_dim = w_dim = 2 * (attention_head_dim // 6)
         t_dim = attention_head_dim - h_dim - w_dim
-        freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
+        freqs_dtype = torch.float64 if current_omni_platform.supports_float64() else torch.float32
 
         freqs_cos = []
         freqs_sin = []
@@ -725,7 +726,7 @@ class WanTransformer3DModel(nn.Module):
     """
 
     _repeated_blocks = ["WanTransformerBlock"]
-    _layerwise_offload_blocks_attr = "blocks"
+    _layerwise_offload_blocks_attrs = ["blocks"]
     packed_modules_mapping = {
         "to_qkv": ["to_q", "to_k", "to_v"],
     }

vllm_omni/diffusion/models/z_image/z_image_transformer.py

@@ -579,6 +579,7 @@ class ZImageTransformer2DModel(CachedTransformer):
     """
 
     _repeated_blocks = ["ZImageTransformerBlock"]
+    _layerwise_offload_blocks_attrs = ["layers"]
 
     @staticmethod
     def _is_transformer_block(name: str, module) -> bool:

vllm_omni/diffusion/offloader/layerwise_backend.py

@@ -312,10 +312,9 @@ class LayerWiseOffloadBackend(OffloadBackend):
             dit_name = modules.dit_names[i]
             logger.info(f"Applying hooks on {dit_name} ({dit_module.__class__.__name__})")
 
-            blocks_attr_name = LayerWiseOffloadBackend.get_blocks_attr_name(dit_module)
-            blocks = LayerWiseOffloadBackend.get_blocks_from_dit(dit_module)
+            blocks_attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(dit_module)
 
-            if not blocks_attr_name or not blocks:
+            if not blocks:
                 logger.warning(
                     "Target layers (blocks) not found. Skipping offloading on %s (%s)",
                     dit_name,
@@ -336,11 +335,20 @@ class LayerWiseOffloadBackend(OffloadBackend):
 
             # Move non-block modules to GPU (they stay resident)
             for name, m in dit_module.named_children():
-                if name == blocks_attr_name:
+                if name not in blocks_attr_names:
+                    m.to(self.device)
+                    logger.debug(f"Moved {name} to device {self.device}")
+                else:
                     logger.debug(f"Skipped blocks module {name}")
-                    continue
-                m.to(self.device)
-                logger.debug(f"Moved {name} to device {self.device}")
+
+            # Move top-level params/buffers to GPU (dit_module's own, not sub-modules)
+            for param in dit_module._parameters.values():
+                if param is not None:
+                    param.data = param.data.to(self.device, non_blocking=True)
+
+            for buffer in dit_module._buffers.values():
+                if buffer is not None:
+                    buffer.data = buffer.data.to(self.device, non_blocking=True)
 
             # Pre-fetch the first layer by manually calling the hook function on the last layer;
             # For subsequent requests, the first layer/block will be pre-fetched
@@ -395,40 +403,84 @@ class LayerWiseOffloadBackend(OffloadBackend):
         logger.info("Layer-wise offloading disabled")
 
     @staticmethod
-    def get_blocks_attr_name(model: nn.Module) -> str | None:
-        """Retrieve blocks attribute name from provided DiT model"""
-        return getattr(model.__class__, "_layerwise_offload_blocks_attr", None)
+    def get_blocks_attr_names(model: nn.Module) -> list[str]:
+        """Get block attribute names from model class."""
+        attrs: list[str] = getattr(model.__class__, "_layerwise_offload_blocks_attrs", [])
+
+        if not attrs:
+            old_attr = getattr(model.__class__, "_layerwise_offload_blocks_attr", None)
+            if old_attr is not None:
+                logger.warning(
+                    "'_layerwise_offload_blocks_attr' is deprecated, "
+                    "please use '_layerwise_offload_blocks_attrs' instead. "
+                    "Example: _layerwise_offload_blocks_attrs = ['blocks']"
+                )
+                attrs = [old_attr] if isinstance(old_attr, str) else list(old_attr)
+
+        return attrs
 
     @staticmethod
-    def set_blocks_attr_name(model: nn.Module, name: str) -> None:
-        if not hasattr(model.__class__, "_layerwise_offload_blocks_attr"):
-            setattr(model.__class__, "_layerwise_offload_blocks_attr", name)
+    def set_blocks_attr_names(model: nn.Module, names: list[str]) -> None:
+        if not hasattr(model.__class__, "_layerwise_offload_blocks_attrs"):
+            setattr(model.__class__, "_layerwise_offload_blocks_attrs", names)
 
     @staticmethod
-    def get_blocks_from_dit(model: nn.Module) -> list[nn.Module]:
+    def get_blocks_from_dit(model: nn.Module) -> tuple[list[str], list[nn.Module]]:
         """
-        Retrieve a list of blocks from provided DiT model. Blocks attribute name
-        are found by `_layerwise_offload_blocks_attr` set to DiT models. For example,
+        Retrieve blocks and attribute names from provided DiT model. Blocks attribute names
+        are found by `_layerwise_offload_blocks_attrs` set to DiT models. For example,
 
         ```
         class WanTransformer3DModel(nn.Module):
-            _layerwise_offload_blocks_attr = "blocks"
+            _layerwise_offload_blocks_attrs = ["blocks"]
         ```
+
+        Returns:
+            Tuple of (blocks_attr_names, blocks)
         """
-        blocks_attr_name = LayerWiseOffloadBackend.get_blocks_attr_name(model)
-        if blocks_attr_name is None:
+        blocks_attr_names = LayerWiseOffloadBackend.get_blocks_attr_names(model)
+        if not blocks_attr_names:
             logger.warning(
-                f"No _layerwise_offload_blocks_attr defined for {model.__class__.__name__}, "
+                f"No _layerwise_offload_blocks_attrs defined for {model.__class__.__name__}, "
                 "skipping layerwise offloading"
             )
-            return []
+            return [], []
 
-        _blocks = getattr(model, blocks_attr_name, None)
-        if _blocks is None:
+        blocks = []
+        for name in blocks_attr_names:
+            attr = getattr(model, name, None)
+            if attr is None:
+                raise AttributeError(
+                    f"Attribute '{name}' declared in _layerwise_offload_blocks_attrs "
+                    f"does not exist on model {model.__class__.__name__}"
+                )
+            try:
+                attr_iter = iter(attr)
+            except TypeError:
+                if isinstance(attr, nn.Module):
+                    logger.warning(
+                        "Attribute '%s' on %s is not iterable; treating it as one block.",
+                        name,
+                        model.__class__.__name__,
+                    )
+                    blocks.append(attr)
+                    continue
+
+                logger.warning(
+                    "Attribute '%s' on %s is not iterable (got %s); skipping it.",
+                    name,
+                    model.__class__.__name__,
+                    type(attr).__name__,
+                )
+            else:
+                blocks.extend(attr_iter)
+
+        if not blocks:
             logger.warning(
-                f"Blocks (layers) '{blocks_attr_name}' not found on {model.__class__.__name__}, "
-                "skipping layerwise offloading"
+                "No blocks found in %s for %s, skipping layerwise offloading",
+                blocks_attr_names,
+                model.__class__.__name__,
             )
-            return []
+            return [], []
 
-        return list(_blocks)
+        return blocks_attr_names, blocks

vllm_omni/diffusion/offloader/module_collector.py

@@ -21,9 +21,9 @@ class PipelineModules:
 class ModuleDiscovery:
     """Discovers pipeline components for offloading"""
 
-    DIT_ATTRS = ["transformer", "transformer_2", "dit", "language_model", "transformer_blocks"]
+    DIT_ATTRS = ["transformer", "transformer_2", "dit", "sr_dit", "language_model", "transformer_blocks"]
     ENCODER_ATTRS = ["text_encoder", "text_encoder_2", "text_encoder_3", "image_encoder"]
-    VAE_ATTRS = ["vae"]
+    VAE_ATTRS = ["vae", "audio_vae"]
 
     @staticmethod
     def discover(pipeline: nn.Module) -> PipelineModules:

vllm_omni/diffusion/registry.py

@@ -173,6 +173,11 @@ _DIFFUSION_MODELS = {
         "pipeline_hunyuan_video_1_5_i2v",
         "HunyuanVideo15I2VPipeline",
     ),
+    "MagiHumanPipeline": (
+        "magi_human",
+        "pipeline_magi_human",
+        "MagiHumanPipeline",
+    ),
     "OmniVoicePipeline": (
         "omnivoice",
         "pipeline_omnivoice",
@@ -368,6 +373,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "Flux2Pipeline": "get_flux2_post_process_func",
     "HunyuanVideo15Pipeline": "get_hunyuan_video_15_post_process_func",
     "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_post_process_func",
+    "MagiHumanPipeline": "get_magi_human_post_process_func",
     "OmniVoicePipeline": "get_omnivoice_post_process_func",
 }
 
@@ -387,6 +393,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "HeliosPipeline": "get_helios_pre_process_func",
     "HeliosPyramidPipeline": "get_helios_pre_process_func",
     "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_pre_process_func",
+    "MagiHumanPipeline": "get_magi_human_pre_process_func",
 }
 
 

vllm_omni/diffusion/utils/media_utils.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Video/audio muxing utilities using PyAV (no ffmpeg binary dependency)."""
+
+from __future__ import annotations
+
+import io
+from fractions import Fraction
+
+import av
+import numpy as np
+
+
+def mux_video_audio_bytes(
+    video_frames: np.ndarray,
+    audio_waveform: np.ndarray | None = None,
+    *,
+    fps: float = 25.0,
+    audio_sample_rate: int = 44100,
+    video_codec: str = "h264",
+    audio_codec: str = "aac",
+    crf: str = "18",
+) -> bytes:
+    """Mux video frames and optional audio waveform into MP4 bytes.
+
+    Args:
+        video_frames: uint8 array of shape ``(T, H, W, 3)`` (RGB).
+        audio_waveform: float32 array – mono ``(N,)`` or ``(N, C)`` / ``(C, N)``.
+        fps: Video frame rate.
+        audio_sample_rate: Audio sample rate in Hz.
+        video_codec: Video codec name.
+        audio_codec: Audio codec name.
+        crf: Constant rate factor for the video encoder.
+
+    Returns:
+        Raw MP4 bytes ready to be written to disk or streamed.
+    """
+    buf = io.BytesIO()
+    container = av.open(buf, mode="w", format="mp4")
+
+    v_stream = container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000))
+    v_stream.width = video_frames.shape[2]
+    v_stream.height = video_frames.shape[1]
+    v_stream.pix_fmt = "yuv420p"
+    v_stream.options = {"crf": crf}
+
+    a_stream = None
+    if audio_waveform is not None:
+        samples = audio_waveform.astype(np.float32)
+        if samples.ndim == 1:
+            samples = samples.reshape(1, -1)
+        elif samples.ndim == 2 and samples.shape[0] > samples.shape[1]:
+            samples = np.ascontiguousarray(samples.T)
+        num_channels = samples.shape[0]
+        layout = "stereo" if num_channels >= 2 else "mono"
+        a_stream = container.add_stream(audio_codec, rate=audio_sample_rate)
+        a_stream.layout = layout
+
+    for frame_data in video_frames:
+        frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24")
+        for packet in v_stream.encode(frame):
+            container.mux(packet)
+    for packet in v_stream.encode():
+        container.mux(packet)
+
+    if a_stream is not None and audio_waveform is not None:
+        audio_frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout=layout)
+        audio_frame.sample_rate = audio_sample_rate
+        for packet in a_stream.encode(audio_frame):
+            container.mux(packet)
+        for packet in a_stream.encode():
+            container.mux(packet)
+
+    container.close()
+    return buf.getvalue()

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -160,11 +160,15 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
                 new_ids = payload_data.get("code_predictor_codes", [])
                 request.prompt_token_ids = new_ids
-                # Pass additional fields (like left_context_size) to the request
-                # Only pass chunk context metadata in additional_information
-                request.additional_information = {}
-                if "left_context_size" in payload_data:
-                    request.additional_information["left_context_size"] = payload_data["left_context_size"]
+                # Preserve previously attached request metadata (e.g. prompt
+                # conditioning tensors) and update only per-chunk fields.
+                prev_info = getattr(request, "additional_information", None)
+                info = dict(prev_info) if isinstance(prev_info, dict) else {}
+                for key, value in payload_data.items():
+                    if key in {"code_predictor_codes", "finished"}:
+                        continue
+                    info[key] = value
+                request.additional_information = info
                 request.num_computed_tokens = 0
 
                 # Empty chunk with more data expected: keep polling.
@@ -240,9 +244,23 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         if success:
             self.put_req_chunk[external_req_id] += 1
             logger.debug(f"[Stage-{stage_id}] Sent {connector_put_key}")
+            finished_flag = payload_data.get("finished")
+            is_payload_finished = False
+            if isinstance(finished_flag, torch.Tensor):
+                is_payload_finished = finished_flag.numel() == 1 and bool(finished_flag.item())
+            elif finished_flag is not None:
+                is_payload_finished = bool(finished_flag)
+
+            # Reclaim per-request async state only after the terminal payload
+            # has been sent successfully. This avoids cleanup->save races.
+            if is_payload_finished:
+                self.cleanup(request.request_id, external_req_id)
 
         if is_finished:
-            self.cleanup_sender(external_req_id)
+            self.code_prompt_token_ids.pop(external_req_id, None)
+            cached_ic = getattr(self, "_cached_ic", None)
+            if cached_ic is not None:
+                cached_ic.pop(external_req_id, None)
 
     ########################################################################
     # Cleanup

vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py

@@ -45,9 +45,10 @@ class OmniCoordClientForStage:
         self._status = StageStatus.UP
         self._queue_length = 0
         self._closed = False
+        self._closing = False
         self._heartbeat_interval = 5.0
         self._stop_event = threading.Event()
-        self._send_lock = threading.Lock()
+        self._send_lock = threading.RLock()
 
         self._send_event("update")
 
@@ -57,38 +58,45 @@ class OmniCoordClientForStage:
         )
         self._heartbeat_thread.start()
 
-    def _reconnect(self) -> bool:
+    def _reconnect(self, max_retries: int = 3, retry_interval: float = 5.0) -> bool:
         """Best-effort reconnect with up to ``max_retries`` attempts.
 
-        Each attempt closes the current socket/context, sleeps 5 seconds,
-        then creates a new DEALER socket and reconnects to the coordinator.
-        Caller must hold ``_send_lock``.
+        Each attempt closes the current socket/context, sleeps ``retry_interval``
+        seconds, then creates a new DEALER socket and reconnects to the coordinator.
         Returns True on success, False if all attempts fail.
         """
-        while not self._stop_event.is_set() and not self._closed:
-            try:
-                self._socket.close(0)
-            except zmq.ZMQError:
-                pass
-            try:
-                self._ctx.term()
-            except zmq.ZMQError:
-                pass
+        if max_retries <= 0:
+            return False
 
-            time.sleep(5.0)
+        for attempt in range(1, max_retries + 1):
+            with self._send_lock:
+                if self._stop_event.is_set() or self._closed:
+                    return False
+                try:
+                    self._socket.close(0)
+                except zmq.ZMQError:
+                    pass
+                try:
+                    self._ctx.term()
+                except zmq.ZMQError:
+                    pass
 
-            try:
-                self._ctx = zmq.Context()
-                self._socket = self._ctx.socket(zmq.DEALER)
-                self._socket.connect(self._coord_zmq_addr)
-                return True
-            except zmq.ZMQError as e:
-                logger.error(
-                    "Stage client reconnect failed, will retry in 5s (coord=%s)",
-                    self._coord_zmq_addr,
-                    exc_info=e,
-                )
-                continue
+                try:
+                    self._ctx = zmq.Context()
+                    self._socket = self._ctx.socket(zmq.DEALER)
+                    self._socket.connect(self._coord_zmq_addr)
+                    return True
+                except zmq.ZMQError as e:
+                    logger.error(
+                        "Stage client reconnect failed (attempt=%d/%d, coord=%s)",
+                        attempt,
+                        max_retries,
+                        self._coord_zmq_addr,
+                        exc_info=e,
+                    )
+
+            if retry_interval > 0:
+                time.sleep(retry_interval)
         return False
 
     def _send_event(self, event_type: str) -> None:
@@ -102,20 +110,20 @@ class OmniCoordClientForStage:
         to 3 times (5s sleep each) and retries the send once after a
         successful reconnect. Raises if reconnect or the retry send fails.
         """
-        if self._closed:
-            raise RuntimeError("Client already closed")
-
-        event = InstanceEvent(
-            input_addr=self._input_addr,
-            output_addr=self._output_addr,
-            stage_id=self._stage_id,
-            event_type=event_type,
-            status=self._status,
-            queue_length=self._queue_length,
-        )
-        data = json.dumps(asdict(event)).encode("utf-8")
-
         with self._send_lock:
+            if self._closed:
+                raise RuntimeError("Client already closed")
+
+            event = InstanceEvent(
+                input_addr=self._input_addr,
+                output_addr=self._output_addr,
+                stage_id=self._stage_id,
+                event_type=event_type,
+                status=self._status,
+                queue_length=self._queue_length,
+            )
+            data = json.dumps(asdict(event)).encode("utf-8")
+
             try:
                 self._socket.send(data, flags=zmq.NOBLOCK)
                 return
@@ -124,7 +132,7 @@ class OmniCoordClientForStage:
                 return
             except (RuntimeError, zmq.ZMQError) as e:
                 # First send failed; try reconnecting a few times.
-                if not self._reconnect:
+                if not self._reconnect(max_retries=3):
                     logger.error("Failed to send event and reconnect to coordinator", exc_info=e)
                     raise
 
@@ -149,12 +157,16 @@ class OmniCoordClientForStage:
         if status is None and queue_length is None:
             raise ValueError("At least one of status or queue_length must be provided")
 
-        if status is not None:
-            self._status = status
-        if queue_length is not None:
-            self._queue_length = queue_length
+        with self._send_lock:
+            if self._closed or self._closing:
+                raise RuntimeError("Client is closing or already closed")
 
-        self._send_event("update")
+            if status is not None:
+                self._status = status
+            if queue_length is not None:
+                self._queue_length = queue_length
+
+            self._send_event("update")
 
     def _heartbeat_loop(self) -> None:
         """Periodically send heartbeat events while the client is alive."""
@@ -164,8 +176,11 @@ class OmniCoordClientForStage:
 
             try:
                 self._send_event("heartbeat")
-            except (RuntimeError, zmq.ZMQError):
-                break
+            except (RuntimeError, zmq.ZMQError) as e:
+                if self._closed or self._stop_event.is_set():
+                    break
+                logger.warning("Heartbeat send failed; will retry on next interval", exc_info=e)
+                continue
 
     def close(self) -> None:
         """Send a final down event and close the underlying socket."""
@@ -177,17 +192,23 @@ class OmniCoordClientForStage:
         if hasattr(self, "_heartbeat_thread"):
             self._heartbeat_thread.join(timeout=1.0)
 
-        # Mark status as DOWN and send one last update.
-        self._status = StageStatus.DOWN
-        try:
-            self._send_event("update")
-        except zmq.ZMQError:
-            pass  # Socket may already be broken, proceed with close
+        with self._send_lock:
+            if self._closed:
+                raise RuntimeError("Client already closed")
 
-        # Close DEALER socket and terminate this client's context.
-        self._socket.close(0)
-        try:
-            self._ctx.term()
-        except zmq.ZMQError:
-            pass
-        self._closed = True
+            self._closing = True
+
+            # Mark status as DOWN and send one last update.
+            self._status = StageStatus.DOWN
+            try:
+                self._send_event("update")
+            except (RuntimeError, zmq.ZMQError):
+                pass  # Socket may already be broken, proceed with close
+
+            # Close DEALER socket and terminate this client's context.
+            self._socket.close(0)
+            try:
+                self._ctx.term()
+            except zmq.ZMQError:
+                pass
+            self._closed = True

vllm_omni/engine/async_omni_engine.py

@@ -42,7 +42,10 @@ from vllm_omni.engine import (
 )
 from vllm_omni.engine.orchestrator import Orchestrator
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
-from vllm_omni.engine.serialization import serialize_additional_information
+from vllm_omni.engine.serialization import (
+    deserialize_additional_information,
+    serialize_additional_information,
+)
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
 from vllm_omni.engine.stage_engine_core_proc import (
     complete_stage_handshake,
@@ -170,6 +173,38 @@ def _upgrade_to_omni_request(
     )
 
 
+def _apply_omni_final_stage_metadata(
+    request: EngineCoreRequest,
+    final_stage_id: int,
+) -> EngineCoreRequest:
+    """Tag EngineCoreRequest so OmniARScheduler can skip DiT KV when final_stage_id is 0."""
+    merged: dict[str, Any] = {}
+    if isinstance(request, OmniEngineCoreRequest) and request.additional_information is not None:
+        merged = deserialize_additional_information(request.additional_information)
+    merged["omni_final_stage_id"] = final_stage_id
+    payload = serialize_additional_information(merged)
+    return OmniEngineCoreRequest(
+        request_id=request.request_id,
+        prompt_token_ids=request.prompt_token_ids,
+        mm_features=request.mm_features,
+        sampling_params=request.sampling_params,
+        pooling_params=request.pooling_params,
+        arrival_time=request.arrival_time,
+        lora_request=request.lora_request,
+        cache_salt=request.cache_salt,
+        data_parallel_rank=request.data_parallel_rank,
+        prompt_embeds=request.prompt_embeds,
+        client_index=request.client_index,
+        current_wave=request.current_wave,
+        priority=request.priority,
+        trace_headers=request.trace_headers,
+        resumable=request.resumable,
+        external_req_id=request.external_req_id,
+        reasoning_ended=request.reasoning_ended,
+        additional_information=payload,
+    )
+
+
 def _weak_shutdown_async_omni_engine(
     orchestrator_thread: threading.Thread | None,
     request_queue: janus.Queue[dict[str, Any]] | None,
@@ -713,9 +748,12 @@ class AsyncOmniEngine:
             # to match the key used in Orchestrator.request_states so that
             # output routing (output.request_id lookup) can find the req_state.
             request.external_req_id = request_id
+            request = _apply_omni_final_stage_metadata(request, final_stage_id)
 
             # Register with stage 0's output processor.
             output_prompt_text = prompt_text
+            if output_prompt_text is None and isinstance(original_prompt, dict):
+                output_prompt_text = original_prompt.get("prompt")
             self.output_processors[0].add_request(
                 request=request,
                 prompt=output_prompt_text,
@@ -882,6 +920,41 @@ class AsyncOmniEngine:
         num_devices = max(1, int(parallel_config.world_size))
         devices = ",".join(str(i) for i in range(num_devices))
 
+        stage_engine_args = {
+            "max_num_seqs": 1,
+            "parallel_config": parallel_config,
+            "model_class_name": kwargs.get("model_class_name", None),
+            "step_execution": kwargs.get("step_execution", False),
+            "vae_use_slicing": kwargs.get("vae_use_slicing", False),
+            "vae_use_tiling": kwargs.get("vae_use_tiling", False),
+            "cache_backend": cache_backend,
+            "cache_config": cache_config,
+            "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),
+            "enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
+            "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
+            "enforce_eager": kwargs.get("enforce_eager", False),
+            "diffusion_load_format": kwargs.get("diffusion_load_format", "default"),
+            "custom_pipeline_args": kwargs.get("custom_pipeline_args", None),
+            "worker_extension_cls": kwargs.get("worker_extension_cls", None),
+            "enable_sleep_mode": kwargs.get("enable_sleep_mode", False),
+            "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
+            "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
+            "quantization": kwargs.get("quantization", None),
+            "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False),
+            **(
+                {
+                    "profiler_config": asdict(kwargs["profiler_config"])
+                    if hasattr(kwargs["profiler_config"], "__dataclass_fields__")
+                    else kwargs["profiler_config"]
+                }
+                if kwargs.get("profiler_config") is not None
+                else {}
+            ),
+        }
+        # Only set dtype if it was already explicitly passed and normalized
+        if "dtype" in normalized_kwargs:
+            stage_engine_args["dtype"] = normalized_kwargs["dtype"]
+
         default_stage_cfg = [
             {
                 "stage_id": 0,
@@ -890,37 +963,7 @@ class AsyncOmniEngine:
                     "process": True,
                     "devices": devices,
                 },
-                "engine_args": {
-                    "max_num_seqs": 1,
-                    "parallel_config": parallel_config,
-                    "model_class_name": kwargs.get("model_class_name", None),
-                    "step_execution": kwargs.get("step_execution", False),
-                    "vae_use_slicing": kwargs.get("vae_use_slicing", False),
-                    "vae_use_tiling": kwargs.get("vae_use_tiling", False),
-                    "cache_backend": cache_backend,
-                    "cache_config": cache_config,
-                    "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),
-                    "enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
-                    "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
-                    "enforce_eager": kwargs.get("enforce_eager", False),
-                    "diffusion_load_format": kwargs.get("diffusion_load_format", "default"),
-                    "custom_pipeline_args": kwargs.get("custom_pipeline_args", None),
-                    "worker_extension_cls": kwargs.get("worker_extension_cls", None),
-                    "enable_sleep_mode": kwargs.get("enable_sleep_mode", False),
-                    "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
-                    "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
-                    "quantization": kwargs.get("quantization", None),
-                    "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False),
-                    **(
-                        {
-                            "profiler_config": asdict(kwargs["profiler_config"])
-                            if hasattr(kwargs["profiler_config"], "__dataclass_fields__")
-                            else kwargs["profiler_config"]
-                        }
-                        if kwargs.get("profiler_config") is not None
-                        else {}
-                    ),
-                },
+                "engine_args": stage_engine_args,
                 "final_output": True,
                 "final_output_type": "image",
             }

vllm_omni/entrypoints/openai/api_server.py

@@ -353,6 +353,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
     try:
         await shutdown_task
     finally:
+        app.state.openai_serving_speech.shutdown()
         sock.close()
 
 
@@ -1311,7 +1312,13 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
         if request.negative_prompt is not None:
             prompt["negative_prompt"] = request.negative_prompt
         gen_params = OmniDiffusionSamplingParams(num_outputs_per_prompt=request.n)
-
+        extra_args = {}
+        if request.use_system_prompt is not None:
+            extra_args["use_system_prompt"] = request.use_system_prompt
+        if request.system_prompt is not None:
+            extra_args["system_prompt"] = request.system_prompt
+        if extra_args:
+            gen_params.extra_args = extra_args
         # Parse per-request LoRA (compatible with chat's extra_body.lora shape).
         lora_request, lora_scale = _parse_lora_request(request.lora)
         _update_if_not_none(gen_params, "lora_request", lora_request)

vllm_omni/entrypoints/openai/protocol/audio.py

@@ -2,7 +2,7 @@ import math
 from typing import Literal
 
 import numpy as np
-from pydantic import BaseModel, Field, field_validator, model_validator
+from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator
 
 _MAX_EMBEDDING_DIM = 8192
 
@@ -10,8 +10,12 @@ _MAX_EMBEDDING_DIM = 8192
 class OpenAICreateSpeechRequest(BaseModel):
     input: str
     model: str | None = None
+    # Accept both "voice" (OpenAI convention) and "speaker" (model/internal
+    # convention) as input keys.  Intentionally global — all TTS backends
+    # (Qwen3-TTS, Voxtral, Fish Speech) use this field for the speaker name.
     voice: str | None = Field(
         default=None,
+        validation_alias=AliasChoices("voice", "speaker"),
         description="Speaker/voice to use. For Qwen3-TTS: vivian, ryan, aiden, etc.",
     )
     instructions: str | None = Field(

vllm_omni/entrypoints/openai/protocol/images.py

@@ -81,6 +81,24 @@ class ImageGenerationRequest(BaseModel):
 
     # vllm-omni extensions for diffusion control
     negative_prompt: str | None = Field(default=None, description="Text describing what to avoid in the image")
+    system_prompt: str | None = Field(
+        default=None, description="Custom system prompt. Used when --use_system_prompt is custom"
+    )
+    use_system_prompt: str | None = Field(
+        default=None,
+        description="System prompt type. Options: None, dynamic, en_vanilla, "
+        "en_recaption, en_think_recaption, en_unified, custom",
+    )
+
+    @field_validator("use_system_prompt")
+    @classmethod
+    def validate_use_system_prompt(cls, v):
+        """Validate system prompt type."""
+        valid_types = [None, "dynamic", "en_vanilla", "en_recaption", "en_think_recaption", "en_unified", "custom"]
+        if v not in valid_types:
+            raise ValueError(f"Invalid use_system_prompt type: {v}. Must be one of: {valid_types[1:] + [None]}")
+        return v
+
     num_inference_steps: int | None = Field(
         default=None,
         ge=1,

vllm_omni/entrypoints/openai/serving_speech.py

@@ -8,6 +8,7 @@ import re
 import struct
 import tempfile
 import time
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from typing import Any
 
@@ -22,6 +23,7 @@ from vllm.entrypoints.openai.engine.serving import OpenAIServing
 from vllm.logger import init_logger
 from vllm.multimodal.media import MediaConnector
 from vllm.utils import random_uuid
+from vllm.utils.async_utils import make_async
 
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
 from vllm_omni.entrypoints.openai.protocol.audio import (
@@ -153,6 +155,7 @@ def _validate_path_within_directory(file_path: Path, directory: Path) -> bool:
 
 class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
     _diffusion_mode: bool = False
+    _tts_executor: ThreadPoolExecutor | None = None
 
     @classmethod
     def for_diffusion(
@@ -219,6 +222,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         # Load speech tokenizer codec parameters for prompt length estimation
         self._codec_frame_rate: float | None = self._load_codec_frame_rate()
 
+        # Shared thread pool executor for blocking TTS preprocessing
+        # operations. max_workers=1 serializes tokenizer access to avoid
+        # Rust RefCell "Already borrowed" errors from concurrent use.
+        self._tts_executor = ThreadPoolExecutor(max_workers=1)
+        self._build_voxtral_prompt_async = make_async(self._build_voxtral_prompt, executor=self._tts_executor)
+        self._build_fish_speech_prompt_async = make_async(self._build_fish_speech_prompt, executor=self._tts_executor)
+        self._estimate_prompt_len_async = make_async(self._estimate_prompt_len, executor=self._tts_executor)
+
     def _load_codec_frame_rate(self) -> float | None:
         """Load codec frame rate from speech tokenizer config for prompt length estimation."""
         try:
@@ -252,6 +263,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             pass
         return None
 
+    def shutdown(self) -> None:
+        """Shut down the TTS thread pool executor."""
+        if self._tts_executor is not None:
+            self._tts_executor.shutdown(wait=False, cancel_futures=True)
+            self._tts_executor = None
+
     def _find_tts_stage(self):
         """Find and return the TTS stage config, or None if not found."""
         for stage in self.engine_client.stage_configs:
@@ -450,6 +467,48 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             logger.error(f"Could not read audio file for voice {voice_name}: {e}")
             return None
 
+    def _get_uploaded_speaker_embedding(self, voice_name: str) -> list[float] | None:
+        """Load pre-computed speaker embedding for an uploaded voice.
+
+        Returns the embedding as a list of floats, or None if the voice
+        was not uploaded with an embedding (i.e. it has audio instead).
+        """
+        voice_name_lower = voice_name.lower()
+        if voice_name_lower not in self.uploaded_speakers:
+            return None
+
+        speaker_info = self.uploaded_speakers[voice_name_lower]
+        if speaker_info.get("embedding_source") != "direct":
+            return None
+
+        cache_file = speaker_info.get("cache_file")
+        if not cache_file or not Path(cache_file).exists():
+            logger.warning("Embedding file not found for voice %s: %s", voice_name, cache_file)
+            return None
+
+        if not _validate_path_within_directory(Path(cache_file), self.uploaded_speakers_dir):
+            logger.error("Cache file path traversal detected for voice %s: %s", voice_name, cache_file)
+            return None
+
+        try:
+            from safetensors.torch import load_file
+        except ImportError:
+            logger.error(
+                "The 'safetensors' package is required to load speaker embeddings. "
+                "Install it with: pip install safetensors"
+            )
+            return None
+
+        try:
+            tensors = load_file(cache_file)
+            if "speaker_embedding" not in tensors:
+                logger.warning("Key 'speaker_embedding' not found in %s for voice %s", cache_file, voice_name)
+                return None
+            return tensors["speaker_embedding"].squeeze().tolist()
+        except Exception as e:
+            logger.error("Could not load embedding for voice %s: %s", voice_name, e)
+            return None
+
     async def upload_voice(
         self,
         audio_file: UploadFile,
@@ -841,11 +900,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 # voice is not None
                 voice_lower = request.voice.lower()
                 if voice_lower in self.uploaded_speakers:
-                    # Check if audio file exists for uploaded speaker
+                    # Check if data file exists for uploaded speaker
                     speaker_info = self.uploaded_speakers[voice_lower]
                     file_path = Path(speaker_info["file_path"])
                     if not file_path.exists():
-                        return f"Audio file for uploaded speaker '{request.voice}' not found on disk"
+                        return f"Data file for uploaded speaker '{request.voice}' not found on disk"
+                    # For embedding-uploaded voices, verify the cache is ready
+                    if speaker_info.get("embedding_source") == "direct":
+                        cache_file = speaker_info.get("cache_file")
+                        if not cache_file or not Path(cache_file).exists():
+                            status = speaker_info.get("cache_status", "unknown")
+                            return f"Speaker embedding for '{request.voice}' is not yet ready (cache_status='{status}')"
                 else:
                     # need ref_audio for built-in speaker
                     if request.ref_audio is None:
@@ -1041,9 +1106,20 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         streaming needs per-chunk delta slicing; non-streaming needs full concatenation.
         """
         mm = getattr(res, "multimodal_output", None)
+        ro = None
         if not mm:
             ro = getattr(res, "request_output", None)
             mm = getattr(ro, "multimodal_output", None) if ro else None
+        if not mm:
+            if ro is None:
+                ro = getattr(res, "request_output", None)
+            outputs = getattr(ro, "outputs", None) if ro else None
+            if outputs:
+                for completion_output in outputs:
+                    completion_mm = getattr(completion_output, "multimodal_output", None)
+                    if completion_mm:
+                        mm = completion_mm
+                        break
         if not mm:
             return None, None
         key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None)
@@ -1079,20 +1155,32 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             # Uploaded voices use task_type="Base" (CustomVoice requires built-in spk_id).
             # If ref_text was provided at upload time, use in-context cloning; otherwise x_vector only.
             if request.voice.lower() in self.uploaded_speakers and request.ref_audio is None:
-                audio_data = self._get_uploaded_audio_data(request.voice)
-                if not audio_data:
-                    raise ValueError(f"Audio file for uploaded voice '{request.voice}' is missing or corrupted")
                 speaker_info = self.uploaded_speakers[request.voice.lower()]
-                stored_ref_text = speaker_info.get("ref_text")
-                params["ref_audio"] = [audio_data]
-                params["task_type"] = ["Base"]
-                params["voice_created_at"] = [speaker_info.get("created_at", 0)]
-                if stored_ref_text:
-                    params["ref_text"] = [stored_ref_text]
-                    params["x_vector_only_mode"] = [False]
+
+                # Check if this voice was uploaded with a pre-computed embedding.
+                # Populate request.speaker_embedding so the existing code path
+                # (below) handles voice_clone_prompt and x_vector_only_mode.
+                embedding = self._get_uploaded_speaker_embedding(request.voice)
+                if embedding is not None:
+                    request.speaker_embedding = embedding
+                    params["task_type"] = ["Base"]
+                    logger.info("Auto-set speaker_embedding for uploaded voice: %s", request.voice)
                 else:
-                    params["x_vector_only_mode"] = [True]
-                logger.info("Auto-set ref_audio for uploaded voice: %s (icl=%s)", request.voice, bool(stored_ref_text))
+                    audio_data = self._get_uploaded_audio_data(request.voice)
+                    if not audio_data:
+                        raise ValueError(f"Audio file for uploaded voice '{request.voice}' is missing or corrupted")
+                    stored_ref_text = speaker_info.get("ref_text")
+                    params["ref_audio"] = [audio_data]
+                    params["task_type"] = ["Base"]
+                    params["voice_created_at"] = [speaker_info.get("created_at", 0)]
+                    if stored_ref_text:
+                        params["ref_text"] = [stored_ref_text]
+                        params["x_vector_only_mode"] = [False]
+                    else:
+                        params["x_vector_only_mode"] = [True]
+                    logger.info(
+                        "Auto-set ref_audio for uploaded voice: %s (icl=%s)", request.voice, bool(stored_ref_text)
+                    )
 
         elif params["task_type"][0] == "CustomVoice":
             params["speaker"] = ["Vivian"]  # Default for CustomVoice
@@ -1138,7 +1226,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
     # ---- Voxtral TTS helpers ----
 
-    async def _build_voxtral_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
+    def _build_voxtral_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
         """Build Voxtral TTS engine prompt from shared TTS parameters."""
         from mistral_common.protocol.speech.request import SpeechRequest
 
@@ -1278,7 +1366,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             if request.ref_audio is not None:
                 wav_list, sr = await self._resolve_ref_audio(request.ref_audio)
                 ref_audio_data = (wav_list, sr)
-            prompt = self._build_fish_speech_prompt(request, ref_audio_data=ref_audio_data)
+            prompt = await self._build_fish_speech_prompt_async(request, ref_audio_data=ref_audio_data)
             tts_params = {}
         elif self._tts_model_type == "omnivoice":
             tts_params = {}
@@ -1289,7 +1377,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 raise ValueError(validation_error)
 
             if self._tts_model_type == "voxtral_tts":
-                prompt = await self._build_voxtral_prompt(request)
+                prompt = await self._build_voxtral_prompt_async(request)
                 tts_params = {}
             elif self._tts_model_type == "cosyvoice3":
                 prompt = await self._build_cosyvoice3_prompt(request)
@@ -1306,7 +1394,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     wav_list, sr = await self._resolve_ref_audio(ref_audio_source)
                     tts_params["ref_audio"] = [[wav_list, sr]]
 
-                ph_len = self._estimate_prompt_len(tts_params)
+                ph_len = await self._estimate_prompt_len_async(tts_params)
                 prompt = {"prompt_token_ids": [1] * ph_len, "additional_information": tts_params}
         else:
             tts_params = {}
@@ -1332,6 +1420,29 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         sampling_params_list = self.engine_client.default_sampling_params_list
 
+        # CosyVoice3: set dynamic min/max tokens based on text length.
+        # The official model requires min_token_text_ratio to prevent early
+        # EOS and max_token_text_ratio to cap generation length.
+        if self._tts_model_type == "cosyvoice3" and sampling_params_list:
+            import copy
+
+            sampling_params_list = copy.deepcopy(sampling_params_list)
+            text_len = len(request.input)  # rough char-level estimate
+            # Use the model's configured ratios (defaults: min=2, max=20)
+            hf_cfg = self.model_config.hf_config
+            min_ratio = getattr(hf_cfg, "min_token_text_ratio", 2)
+            max_ratio = getattr(hf_cfg, "max_token_text_ratio", 20)
+            min_tokens = max(1, int(text_len * min_ratio))
+            max_tokens = min(2048, int(text_len * max_ratio))
+            sampling_params_list[0].min_tokens = min_tokens
+            sampling_params_list[0].max_tokens = max_tokens
+            logger.info(
+                "CosyVoice3 dynamic tokens: text_len=%d, min_tokens=%d, max_tokens=%d",
+                text_len,
+                min_tokens,
+                max_tokens,
+            )
+
         # Fish defaults come from stage_configs YAML. Only override when the caller
         # explicitly requests a different generation length.
         if self._is_fish_speech and request.max_new_tokens is not None and sampling_params_list:

vllm_omni/entrypoints/utils.py

@@ -182,6 +182,28 @@ def _convert_dataclasses_to_dict(obj: Any) -> Any:
     return obj
 
 
+def _try_resolve_omni_model_type(model: str) -> str | None:
+    """Try to resolve model_type for omni models with empty config.json.
+
+    Checks if any registered omni stage config file name matches a substring
+    in the model name (e.g. 'cosyvoice3' in 'FunAudioLLM/Fun-CosyVoice3-0.5B-2512').
+    When multiple configs match, the longest stem wins to avoid ambiguity
+    (e.g. 'bagel_single_stage' over 'bagel').
+    """
+    stage_configs_dir = PROJECT_ROOT / "vllm_omni" / "model_executor" / "stage_configs"
+    if not stage_configs_dir.exists():
+        return None
+    model_lower = model.lower().replace("-", "").replace("_", "")
+    best_match: str | None = None
+    best_len = 0
+    for config_file in sorted(stage_configs_dir.glob("*.yaml")):
+        candidate = config_file.stem.replace("-", "").replace("_", "")
+        if candidate in model_lower and len(candidate) > best_len:
+            best_match = config_file.stem
+            best_len = len(candidate)
+    return best_match
+
+
 def resolve_model_config_path(model: str) -> str:
     """Resolve the stage config file path from the model name.
 
@@ -220,7 +242,11 @@ def resolve_model_config_path(model: str) -> str:
                 if config_dict and "model_type" in config_dict:
                     model_type = config_dict["model_type"]
                 else:
-                    raise ValueError(f"config.json found but missing 'model_type' for model: {model}")
+                    # For models with empty config.json (e.g. CosyVoice3),
+                    # try matching against registered omni stage configs.
+                    model_type = _try_resolve_omni_model_type(model)
+                    if model_type is None:
+                        raise ValueError(f"config.json found but missing 'model_type' for model: {model}")
             except Exception as e:
                 raise ValueError(f"Failed to read config.json for model: {model}. Error: {e}") from e
         else:

vllm_omni/model_executor/models/bagel/bagel.py

@@ -8,7 +8,7 @@ import torch.nn as nn
 from transformers import BatchFeature
 from vllm.config import VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
-from vllm.inputs import MultiModalDataDict
+from vllm.inputs import ModalityData, MultiModalDataDict
 from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
@@ -27,7 +27,6 @@ from vllm.multimodal.inputs import (
 from vllm.multimodal.parse import (
     ImageEmbeddingItems,
     ImageProcessorItems,
-    ModalityData,
     ModalityDataItems,
     MultiModalDataItems,
     MultiModalDataParser,
@@ -408,6 +407,22 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration):
     the DiT's denoising loop.
     """
 
+    # LoRA packed→sublayer mapping for both standard Qwen2 projections
+    # and the MoE generation-mode projections added by _install_mot_modules().
+    packed_modules_mapping = {
+        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+        "qkv_proj_moe_gen": [
+            "q_proj_moe_gen",
+            "k_proj_moe_gen",
+            "v_proj_moe_gen",
+        ],
+        "mlp_moe_gen.gate_up_proj": [
+            "mlp_moe_gen.gate_proj",
+            "mlp_moe_gen.up_proj",
+        ],
+    }
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__(vllm_config=vllm_config, prefix=prefix)
         config = vllm_config.model_config.hf_config

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py

@@ -174,7 +174,7 @@ class CausalConditionalCFM(ConditionalCFM):
         super().__init__(in_channels, cfm_params, n_spks, spk_emb_dim, estimator)
 
     @torch.inference_mode()
-    def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None):
+    def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, streaming: bool = False):
         """Forward diffusion
 
         Args:
@@ -277,7 +277,9 @@ class CausalMaskedDiffWithDiT(torch.nn.Module):
         prompt_feat,
         prompt_feat_len,
         embedding,
-        finalize,
+        streaming: bool = True,
+        finalize: bool = False,
+        n_timesteps: int = 10,
     ):
         assert token.shape[0] == 1
         # xvec projection
@@ -314,7 +316,8 @@ class CausalMaskedDiffWithDiT(torch.nn.Module):
             mask=mask.unsqueeze(1),
             spks=embedding,
             cond=conds,
-            n_timesteps=10,
+            n_timesteps=max(1, int(n_timesteps)),
+            streaming=streaming,
         )
 
         feat = feat[:, :, mel_len1:]

vllm_omni/model_executor/models/cosyvoice3/config.py

@@ -7,6 +7,10 @@ class CosyVoice3Config(PretrainedConfig):
     model_type = "cosyvoice3"
 
     def __init__(self, **kwargs):
+        # Set primary speech EOS so vLLM stops generation at the right token.
+        # The official CosyVoice3 treats ALL tokens >= speech_token_size
+        # (6561-6760) as stop signals; see stop_token_ids in the YAML configs.
+        kwargs.setdefault("eos_token_id", 6562)
         super().__init__(**kwargs)
         self.sample_rate = 24000
         self.llm_input_size = 896

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py

@@ -2,14 +2,16 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import os
 from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import replace
 from functools import partial
+from threading import Lock
 
-import numpy as np
 import torch
 import torch.nn as nn
 from transformers.feature_extraction_utils import BatchFeature
 from vllm.config import VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
+from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import SupportsMultiModal
@@ -26,6 +28,9 @@ from vllm.multimodal.processing import (
     PromptUpdate,
 )
 from vllm.sequence import IntermediateTensors
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
 from vllm_omni.model_executor.models.cosyvoice3.utils import (
@@ -267,6 +272,8 @@ class CosyVoice3Model(
     supports_multimodal_raw_input_only = True
     supports_multimodal = True
     requires_raw_input_tokens = True
+    prefer_model_sampler = True
+    _sampling_eps = 1e-5
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
@@ -305,6 +312,8 @@ class CosyVoice3Model(
             self.code2wav = CosyVoice3Code2Wav(self.config)
             self.model = self.code2wav.flow_model
             self.hift = self.code2wav.hift
+            # Keep additional information synchronized for async_chunk updates.
+            self.enable_update_additional_information = True
 
             # Expose streaming parameters
             self.token_overlap_len = self.code2wav.token_overlap_len
@@ -313,6 +322,9 @@ class CosyVoice3Model(
             self.mel_cache_len = self.code2wav.mel_cache_len
             self.source_cache_len = self.code2wav.source_cache_len
             self.speech_window = self.code2wav.speech_window
+            self._stream_audio_cache_by_req: dict[str, torch.Tensor] = {}
+            self._stream_audio_cache_lock = Lock()
+            self._stream_vocoder_cache_by_req: dict[str, dict[str, torch.Tensor]] = {}
         else:
             raise ValueError(f"Model stage not supported {self.model_stage}")
 
@@ -331,19 +343,277 @@ class CosyVoice3Model(
         # Use parent's cache config - critical for PagedAttention to work correctly
         return parent_config.with_hf_config(qwen_hf_config, architectures=["Qwen2Model"])
 
+    @staticmethod
+    def _as_tensor(value: object) -> torch.Tensor | None:
+        """Extract tensor payload from runtime info fields."""
+        if isinstance(value, list):
+            if not value:
+                return None
+            value = value[0]
+        if isinstance(value, torch.Tensor):
+            return value
+        return None
+
+    @staticmethod
+    def _as_str(value: object) -> str | None:
+        """Extract string payload from runtime info fields."""
+        if isinstance(value, list):
+            if not value:
+                return None
+            value = value[0]
+        if value is None:
+            return None
+        return str(value)
+
+    @staticmethod
+    def _as_bool(value: object) -> bool:
+        """Extract boolean payload from runtime info fields."""
+        if isinstance(value, list):
+            if not value:
+                return False
+            value = value[0]
+        if isinstance(value, torch.Tensor):
+            if value.numel() == 0:
+                return False
+            return bool(value.reshape(-1)[0].item())
+        if value is None:
+            return False
+        return bool(value)
+
+    @staticmethod
+    def _cross_fade_audio(audio: torch.Tensor, prev_tail: torch.Tensor) -> torch.Tensor:
+        """Blend previous chunk tail into current chunk head using a Hamming window.
+
+        This mirrors upstream CosyVoice's `fade_in_out(...)` semantics:
+        update the current head in-place using a 2*overlap window, then
+        concatenate the unchanged remainder.
+        """
+        if audio.numel() == 0 or prev_tail.numel() == 0:
+            return audio
+        overlap = min(int(audio.numel()), int(prev_tail.numel()))
+        if overlap <= 0:
+            return audio
+        window = torch.hamming_window(2 * overlap, periodic=False, dtype=audio.dtype, device=audio.device)
+        fade_in = window[:overlap]
+        fade_out = window[overlap:]
+        blended = audio[:overlap] * fade_in + prev_tail[-overlap:].to(device=audio.device, dtype=audio.dtype) * fade_out
+        if overlap == int(audio.numel()):
+            return blended
+        return torch.cat([blended, audio[overlap:]], dim=0)
+
+    def _stitch_stream_audio(self, req_id: str | None, audio: torch.Tensor, stream_finished: bool) -> torch.Tensor:
+        """Pass-through stitching for async_chunk.
+
+        Chunk overlap is already removed in mel domain via token_offset_tokens.
+        Applying an additional waveform-domain fade/cache step introduces either
+        duplicated overlap (if no tail trim) or duration shrink (if tail trim).
+        """
+        if req_id is not None and stream_finished and hasattr(self, "_stream_audio_cache_by_req"):
+            with self._stream_audio_cache_lock:
+                self._stream_audio_cache_by_req.pop(req_id, None)
+                if hasattr(self, "_stream_vocoder_cache_by_req"):
+                    self._stream_vocoder_cache_by_req.pop(req_id, None)
+        return audio
+
+    @staticmethod
+    def _split_request_ids(ids: torch.Tensor, seq_token_counts: list[int] | None = None) -> list[torch.Tensor]:
+        """Split concatenated input_ids into per-request segments."""
+        if seq_token_counts is not None:
+            boundaries = [0]
+            for count in seq_token_counts:
+                boundaries.append(boundaries[-1] + int(count))
+            total = ids.numel()
+            return [ids[boundaries[i] : min(boundaries[i + 1], total)] for i in range(len(seq_token_counts))]
+
+        if is_forward_context_available():
+            slices = get_forward_context().ubatch_slices
+            if slices is not None and len(slices) > 1 and not any(hasattr(s, "token_slice") for s in slices):
+                boundaries = [0]
+                for s in slices:
+                    boundaries.append(boundaries[-1] + int(s))
+                return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)]
+
+        return [ids]
+
+    def _sanitize_codec_tokens(self, req_ids: torch.Tensor) -> torch.Tensor:
+        """Filter non-code tokens before feeding flow token embedding."""
+        vocab_size = int(self.code2wav.input_embedding.num_embeddings)
+        valid_mask = (req_ids >= 0) & (req_ids < vocab_size)
+        return req_ids[valid_mask]
+
+    @staticmethod
+    def _req_scalar(param: torch.Tensor | None, req_idx: int, default: float | int) -> float | int:
+        if param is None or param.numel() == 0:
+            return default
+        index = min(req_idx, int(param.numel()) - 1)
+        value = param.reshape(-1)[index].item()
+        if isinstance(default, int):
+            return int(value)
+        return float(value)
+
+    @staticmethod
+    def _multinomial_sample(probs: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor:
+        return torch.multinomial(probs, 1, replacement=True, generator=generator).reshape(())
+
+    @classmethod
+    def _nucleus_sample_one(
+        cls,
+        weighted_scores: torch.Tensor,
+        *,
+        top_p: float,
+        top_k: int,
+        generator: torch.Generator | None,
+    ) -> int:
+        probs = weighted_scores.softmax(dim=0)
+        sorted_prob, sorted_idx = probs.sort(descending=True, stable=True)
+        kept_probs: list[torch.Tensor] = []
+        kept_indices: list[torch.Tensor] = []
+        cum_prob = 0.0
+        max_keep = len(sorted_idx) if top_k <= 0 else min(int(top_k), len(sorted_idx))
+        for i in range(len(sorted_idx)):
+            if cum_prob < top_p and len(kept_probs) < max_keep:
+                cum_prob += float(sorted_prob[i].item())
+                kept_probs.append(sorted_prob[i])
+                kept_indices.append(sorted_idx[i])
+            else:
+                break
+
+        if not kept_probs:
+            return int(sorted_idx[0].item())
+
+        sample_probs = torch.stack(kept_probs)
+        sample_idx = cls._multinomial_sample(sample_probs, generator=generator)
+        return int(torch.stack(kept_indices)[int(sample_idx.item())].item())
+
+    @classmethod
+    def _ras_sample_one(
+        cls,
+        weighted_scores: torch.Tensor,
+        decoded_tokens: Sequence[int],
+        *,
+        top_p: float,
+        top_k: int,
+        win_size: int,
+        tau_r: float,
+        generator: torch.Generator | None,
+    ) -> int:
+        top_id = cls._nucleus_sample_one(
+            weighted_scores,
+            top_p=top_p,
+            top_k=top_k,
+            generator=generator,
+        )
+        if win_size > 0 and decoded_tokens:
+            recent = torch.as_tensor(
+                list(decoded_tokens[-win_size:]),
+                device=weighted_scores.device,
+                dtype=torch.long,
+            )
+            rep_num = int((recent == top_id).sum().item())
+            if rep_num >= win_size * tau_r:
+                weighted_scores = weighted_scores.clone()
+                weighted_scores[top_id] = float("-inf")
+                fallback_probs = weighted_scores.softmax(dim=0)
+                top_id = int(cls._multinomial_sample(fallback_probs, generator=generator).item())
+        return top_id
+
+    def _cosyvoice3_ras_enabled(self, sampling_metadata: SamplingMetadata) -> bool:
+        if self.model_stage != "cosyvoice3_talker":
+            return False
+        if sampling_metadata.max_num_logprobs is not None:
+            return False
+        if sampling_metadata.temperature is None:
+            return False
+        if bool(sampling_metadata.bad_words_token_ids):
+            return False
+        if torch.any(sampling_metadata.frequency_penalties != 0):
+            return False
+        if torch.any(sampling_metadata.presence_penalties != 0):
+            return False
+        return True
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> SamplerOutput | None:
+        if logits is None or logits.numel() == 0:
+            return None
+        if self.model_stage != "cosyvoice3_talker":
+            return None
+
+        sampler = getattr(self, "_talker_sampler", None)
+        if sampler is None:
+            sampler = Sampler()
+            self._talker_sampler = sampler
+
+        if not self._cosyvoice3_ras_enabled(sampling_metadata):
+            return sampler(logits=logits, sampling_metadata=sampling_metadata)
+
+        logits = logits.to(torch.float32)
+        sampling_for_processors = replace(sampling_metadata, no_penalties=True)
+        logits = sampler.apply_logits_processors(logits, sampling_for_processors, predict_bonus_token=False)
+
+        sampling_cfg = dict(self.config.llm.get("sampling", {}))
+        default_top_p = float(sampling_cfg.get("top_p", 0.8))
+        default_top_k = int(sampling_cfg.get("top_k", 25))
+        win_size = int(sampling_cfg.get("win_size", 10))
+        tau_r = float(sampling_cfg.get("tau_r", 0.1))
+
+        sampled_ids: list[int] = []
+        for req_idx in range(int(logits.shape[0])):
+            row_logits = logits[req_idx]
+
+            temperature = float(self._req_scalar(sampling_metadata.temperature, req_idx, 1.0))
+            if temperature < self._sampling_eps:
+                sampled_ids.append(int(torch.argmax(row_logits).item()))
+                continue
+
+            top_p = float(self._req_scalar(sampling_metadata.top_p, req_idx, default_top_p))
+            top_k = int(self._req_scalar(sampling_metadata.top_k, req_idx, default_top_k))
+            generator = sampling_metadata.generators.get(req_idx)
+            weighted_scores = torch.log_softmax(row_logits / max(temperature, self._sampling_eps), dim=0)
+            decoded_tokens = (
+                sampling_metadata.output_token_ids[req_idx] if req_idx < len(sampling_metadata.output_token_ids) else []
+            )
+            sampled_ids.append(
+                self._ras_sample_one(
+                    weighted_scores,
+                    decoded_tokens,
+                    top_p=top_p,
+                    top_k=top_k,
+                    win_size=win_size,
+                    tau_r=tau_r,
+                    generator=generator,
+                )
+            )
+
+        sampled = torch.tensor(sampled_ids, device=logits.device, dtype=torch.int32)
+        return SamplerOutput(sampled_token_ids=sampled.unsqueeze(-1), logprobs_tensors=None)
+
     def compute_logits(self, hidden_states: torch.Tensor | OmniOutput) -> torch.Tensor | None:
         if isinstance(hidden_states, OmniOutput):
             hidden_states = hidden_states.text_hidden_states
         if self.model_stage == "cosyvoice3_talker":
             logits = self.model.llm_decoder(hidden_states)
+            # The decoder outputs speech_token_size + 200 logits.  The official
+            # CosyVoice3 treats ALL tokens >= speech_token_size (the last 200)
+            # as stop signals.  Merge their probabilities into a single EOS
+            # token (6562) via logsumexp so that vLLM's stop_token_ids=[6562]
+            # fires with the correct aggregate stop probability.
+            speech_token_size = self.config.llm["speech_token_size"]
+            eos_idx = self.config.llm["eos_token_id"]
+            stop_logits = logits[..., speech_token_size:]  # last 200
+            merged_stop = torch.logsumexp(stop_logits, dim=-1, keepdim=True)
+            logits[..., speech_token_size:] = float("-inf")  # mask all
+            logits[..., eos_idx] = merged_stop.squeeze(-1)  # restore merged
+            # Pad to full vocab_size for vLLM token handling.
             vocab_size = self.config.vocab_size
             pad_size = vocab_size - logits.size(-1)
-            pad_shape = logits.shape[:-1] + (pad_size,)
-            pad = logits.new_full(pad_shape, float("-inf"))
-            eos_token_val = logits[..., self.config.llm["eos_token_id"]].clone()
-            logits[..., -200:] = float("-inf")
-            logits[..., self.config.llm["eos_token_id"]] = eos_token_val
-            logits = torch.cat([logits, pad], dim=-1)
+            if pad_size > 0:
+                pad_shape = logits.shape[:-1] + (pad_size,)
+                pad = logits.new_full(pad_shape, float("-inf"))
+                logits = torch.cat([logits, pad], dim=-1)
             return logits
         else:
             raise RuntimeError(f"compute_logits is only valid for {self.model_stage}.")
@@ -380,6 +650,7 @@ class CosyVoice3Model(
             hidden = int(self.config.hidden_size)
             return torch.zeros(
                 (input_ids.shape[0], hidden),
+                device=input_ids.device,
             )
         else:
             raise RuntimeError(f"embed_input_ids is not valid for {self.model_stage}.")
@@ -412,28 +683,116 @@ class CosyVoice3Model(
 
             return OmniOutput(text_hidden_states=hidden_states, multimodal_outputs=multimodal_outputs)
         elif self.model_stage == "cosyvoice3_code2wav":
-            runtime_info = kwargs.get("runtime_additional_information", [])
-            if not runtime_info:
-                length = 30 * 24000
-                audio = np.zeros((length,))
-                return OmniOutput(text_hidden_states=None, multimodal_outputs={"audio": audio})
+            runtime_info = kwargs.get("model_intermediate_buffer")
+            if runtime_info is None:
+                runtime_info = kwargs.get("runtime_additional_information", [])
+            if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs:
+                logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer")
 
-            # Remove the last eos token and add batch dimension
-            token = input_ids[..., :-1].unsqueeze(0)
+            seq_token_counts = kwargs.get("seq_token_counts")
+            flat_ids = input_ids.reshape(-1).to(dtype=torch.long)
+            request_ids_list = self._split_request_ids(flat_ids, seq_token_counts)
 
-            # Generate audio using code2wav
-            tts_speech = self.code2wav(
-                token=token,
-                prompt_token=runtime_info[0]["speech_token"][:1],
-                prompt_feat=runtime_info[0]["speech_feat"][:1],
-                embedding=runtime_info[0]["embedding"][:1],
-                n_timesteps=10,
-            )
+            num_reqs = max(1, len(request_ids_list))
+            sample_rate = torch.tensor(int(self.config.sample_rate), dtype=torch.int32)
+            empty_audio = torch.zeros((0,), dtype=torch.float32, device=input_ids.device)
+            audios: list[torch.Tensor] = [empty_audio] * num_reqs
+            srs: list[torch.Tensor] = [sample_rate] * num_reqs
+            if not isinstance(runtime_info, list):
+                runtime_info = []
 
-            return OmniOutput(
-                text_hidden_states=None,
-                multimodal_outputs={"audio": tts_speech, "sr": 22050},
-            )
+            for idx, req_ids in enumerate(request_ids_list):
+                info = runtime_info[idx] if idx < len(runtime_info) and isinstance(runtime_info[idx], dict) else {}
+                req_id = self._as_str(info.get("req_id")) if info else None
+                stream_finished = self._as_bool(info.get("stream_finished")) if info else False
+                speech_token = self._as_tensor(info.get("speech_token")) if info else None
+                speech_feat = self._as_tensor(info.get("speech_feat")) if info else None
+                embedding = self._as_tensor(info.get("embedding")) if info else None
+                if speech_token is None or speech_feat is None or embedding is None:
+                    if stream_finished and req_id is not None and hasattr(self, "_stream_vocoder_cache_by_req"):
+                        with self._stream_audio_cache_lock:
+                            self._stream_vocoder_cache_by_req.pop(req_id, None)
+                    audios[idx] = self._stitch_stream_audio(req_id, empty_audio, stream_finished)
+                    if (
+                        req_ids.numel() > 0
+                        and info
+                        and ("token_offset" in info or "left_context_size" in info or "generated_len" in info)
+                    ):
+                        info_keys = ",".join(sorted(info.keys())) if info else ""
+                        logger.warning_once(
+                            "CosyVoice3 code2wav missing prompt conditioning for non-empty codec tokens: "
+                            "raw_len=%d info_keys=%s",
+                            int(req_ids.numel()),
+                            info_keys,
+                        )
+                    continue
+
+                token = self._sanitize_codec_tokens(req_ids)
+                if token.numel() == 0:
+                    audios[idx] = self._stitch_stream_audio(req_id, empty_audio, stream_finished)
+                    if req_ids.numel() > 0:
+                        logger.warning_once(
+                            "CosyVoice3 code2wav received no valid codec tokens after filtering: "
+                            "raw_len=%d raw_range=[%d,%d] vocab_size=%d",
+                            req_ids.numel(),
+                            int(req_ids.min().item()),
+                            int(req_ids.max().item()),
+                            int(self.code2wav.input_embedding.num_embeddings),
+                        )
+                    continue
+
+                # `generated_len` is injected for many models by the generic
+                # runner, so only explicit chunk-routing fields should switch
+                # code2wav into the streaming path.
+                uses_streaming_decode = bool(info) and (
+                    "stream_finished" in info or "token_offset" in info or "left_context_size" in info
+                )
+                if uses_streaming_decode:
+                    token_offset = 0
+                    try:
+                        if info and "token_offset" in info:
+                            token_offset = max(0, int(info.get("token_offset", 0)))
+                        elif info:
+                            token_offset = max(0, int(info.get("left_context_size", 0)))
+                    except (TypeError, ValueError):
+                        token_offset = 0
+
+                    cache_state = None
+                    if req_id is not None and hasattr(self, "_stream_vocoder_cache_by_req"):
+                        with self._stream_audio_cache_lock:
+                            cache_state = self._stream_vocoder_cache_by_req.get(req_id)
+
+                    tts_speech, new_cache_state = self.code2wav.forward_streaming(
+                        token=token.unsqueeze(0),
+                        prompt_token=speech_token[:1],
+                        prompt_feat=speech_feat[:1],
+                        embedding=embedding[:1],
+                        cache_state=cache_state,
+                        n_timesteps=10,
+                        token_offset_tokens=token_offset,
+                        finalize=stream_finished,
+                    )
+
+                    if req_id is not None and hasattr(self, "_stream_vocoder_cache_by_req"):
+                        with self._stream_audio_cache_lock:
+                            if new_cache_state is None or stream_finished:
+                                self._stream_vocoder_cache_by_req.pop(req_id, None)
+                            else:
+                                self._stream_vocoder_cache_by_req[req_id] = new_cache_state
+                else:
+                    tts_speech = self.code2wav.forward(
+                        token=token.unsqueeze(0),
+                        prompt_token=speech_token[:1],
+                        prompt_feat=speech_feat[:1],
+                        embedding=embedding[:1],
+                        n_timesteps=10,
+                    )
+
+                audio = tts_speech.reshape(-1).to(dtype=torch.float32)
+
+                audios[idx] = self._stitch_stream_audio(req_id, audio, stream_finished)
+
+            return OmniOutput(text_hidden_states=None, multimodal_outputs={"audio": audios, "sr": srs})
         else:
             raise ValueError(f"Unsupported model_stage: {self.model_stage}")
 

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py

@@ -11,11 +11,12 @@ This module contains the code2wav (token-to-waveform) stage which uses:
 
 from __future__ import annotations
 
+from contextlib import nullcontext
+
 import numpy as np
 import torch
 import torch.nn as nn
 from omegaconf import DictConfig
-from torch.nn import functional as F
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT
@@ -29,7 +30,6 @@ from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
 )
 from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer
 from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
-from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask
 
 logger = init_logger(__name__)
 
@@ -150,6 +150,140 @@ class CosyVoice3Code2Wav(nn.Module):
         """Speaker embedding affine layer."""
         return self.flow_model.spk_embed_affine_layer
 
+    @torch.inference_mode()
+    def _forward_mel(
+        self,
+        token: torch.Tensor,
+        prompt_token: torch.Tensor,
+        prompt_feat: torch.Tensor,
+        embedding: torch.Tensor,
+        n_timesteps: int = 10,
+        token_offset_tokens: int = 0,
+        streaming: bool = True,
+        finalize: bool = False,
+    ) -> torch.Tensor:
+        """Generate mel features via the upstream flow-model inference path."""
+        flow_weight = next(self.flow_model.parameters())
+        device = flow_weight.device
+        dtype = flow_weight.dtype
+
+        token = token.to(device=device, dtype=torch.int32)
+        prompt_token = prompt_token.to(device=device, dtype=torch.int32)
+        prompt_feat = prompt_feat.to(device=device, dtype=dtype)
+        embedding = embedding.to(device=device, dtype=dtype)
+        token_len = torch.tensor([token.shape[1]], device=device, dtype=torch.int32)
+        prompt_token_len = torch.tensor([prompt_token.shape[1]], device=device, dtype=torch.int32)
+        prompt_feat_len = torch.tensor([prompt_feat.shape[1]], device=device, dtype=torch.int32)
+
+        with nullcontext():
+            feat, _ = self.flow_model.inference(
+                token=token,
+                token_len=token_len,
+                prompt_token=prompt_token,
+                prompt_token_len=prompt_token_len,
+                prompt_feat=prompt_feat,
+                prompt_feat_len=prompt_feat_len,
+                embedding=embedding,
+                streaming=streaming,
+                finalize=finalize,
+                n_timesteps=n_timesteps,
+            )
+
+        trim_mel = max(0, int(token_offset_tokens)) * int(self.token_mel_ratio)
+        if trim_mel > 0:
+            feat = feat[:, :, trim_mel:]
+
+        return feat
+
+    @staticmethod
+    def _fade_speech(
+        speech: torch.Tensor,
+        prev_speech: torch.Tensor,
+    ) -> torch.Tensor:
+        """Blend previous speech tail into current speech head."""
+        if speech.numel() == 0 or prev_speech.numel() == 0:
+            return speech
+        overlap = min(int(speech.shape[-1]), int(prev_speech.shape[-1]))
+        if overlap <= 0:
+            return speech
+        window = torch.hamming_window(2 * overlap, periodic=False, dtype=speech.dtype, device=speech.device)
+        fade_in = window[:overlap].view(1, -1)
+        fade_out = window[overlap:].view(1, -1)
+        blended_head = (
+            speech[:, :overlap] * fade_in
+            + prev_speech[:, -overlap:].to(device=speech.device, dtype=speech.dtype) * fade_out
+        )
+        if overlap == int(speech.shape[-1]):
+            return blended_head
+        return torch.cat([blended_head, speech[:, overlap:]], dim=-1)
+
+    @torch.inference_mode()
+    def forward_streaming(
+        self,
+        token: torch.Tensor,
+        prompt_token: torch.Tensor,
+        prompt_feat: torch.Tensor,
+        embedding: torch.Tensor,
+        *,
+        cache_state: dict[str, torch.Tensor] | None = None,
+        n_timesteps: int = 10,
+        token_offset_tokens: int = 0,
+        finalize: bool = False,
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]:
+        """Decode streaming audio using cumulative mel + emitted-speech offset.
+
+        This mirrors upstream CosyVoice3 streaming semantics more closely than
+        waveform-domain overlap-add: keep a cumulative mel history per request,
+        re-run causal HiFT on the history, and emit only the newly grown speech
+        suffix. That preserves causal look-right handling without double
+        trimming or duplicated overlap at chunk boundaries.
+        """
+        with nullcontext():
+            feat = self._forward_mel(
+                token=token,
+                prompt_token=prompt_token,
+                prompt_feat=prompt_feat,
+                embedding=embedding,
+                n_timesteps=n_timesteps,
+                token_offset_tokens=token_offset_tokens,
+                streaming=True,
+                finalize=finalize,
+            )
+        hift_weight = self.hift.m_source.l_linear.weight
+        chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype)
+
+        cached_mel = None if not cache_state else cache_state.get("mel")
+        speech_offset_obj = None if not cache_state else cache_state.get("speech_offset")
+        try:
+            speech_offset = int(speech_offset_obj) if speech_offset_obj is not None else 0
+        except (TypeError, ValueError):
+            speech_offset = 0
+
+        if isinstance(cached_mel, torch.Tensor) and cached_mel.numel() > 0:
+            cached_mel = cached_mel.to(device=chunk_mel.device, dtype=chunk_mel.dtype)
+            tts_mel = torch.cat([cached_mel, chunk_mel], dim=-1) if chunk_mel.numel() > 0 else cached_mel
+        else:
+            tts_mel = chunk_mel
+
+        if tts_mel.shape[-1] == 0:
+            tts_speech = torch.zeros((chunk_mel.shape[0], 1, 0), device=chunk_mel.device, dtype=chunk_mel.dtype)
+        else:
+            with nullcontext():
+                tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize)
+
+        tts_speech = tts_speech.reshape(tts_speech.shape[0], -1)
+        speech_offset = max(0, min(speech_offset, int(tts_speech.shape[-1])))
+        emitted_speech = tts_speech[:, speech_offset:]
+
+        if finalize:
+            return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), None
+
+        new_state = {
+            "mel": tts_mel.detach().cpu().contiguous(),
+            "speech_offset": int(tts_speech.shape[-1]),
+        }
+        return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), new_state
+
     @torch.inference_mode()
     def forward(
         self,
@@ -159,74 +293,18 @@ class CosyVoice3Code2Wav(nn.Module):
         embedding: torch.Tensor,
         n_timesteps: int = 10,
     ) -> torch.Tensor:
-        """Generate audio waveform from speech tokens.
-
-        Args:
-            token: Speech tokens from talker stage [batch, seq_len]
-            prompt_token: Prompt speech tokens [batch, prompt_len]
-            prompt_feat: Prompt mel features [batch, feat_len, mel_dim]
-            embedding: Speaker embedding [batch, spk_dim]
-            n_timesteps: Number of diffusion steps
-
-        Returns:
-            Audio waveform [batch, 1, audio_len]
-        """
-        device = token.device
-        dtype = next(self.flow_model.parameters()).dtype
-
-        # Normalize and project speaker embedding
-        embedding = embedding.to(device=device, dtype=dtype)
-        embedding = F.normalize(embedding, dim=1)
-        embedding = self.spk_embed_affine_layer(embedding)
-
-        # Prepare tokens
-        prompt_token = prompt_token.to(device=device)
-        token_len1, token_len2 = prompt_token.shape[1], token.shape[1]
-        prompt_token_len = torch.tensor([token_len1], device=device, dtype=torch.int32)
-        token_len = torch.tensor([token_len2], device=device, dtype=torch.int32)
-
-        # Concatenate prompt and target tokens
-        full_token = torch.cat([prompt_token, token], dim=1)
-        full_token_len = prompt_token_len + token_len
-
-        # Create mask
-        mask = (~make_pad_mask(full_token_len)).unsqueeze(-1).to(embedding)
-
-        # Token embedding
-        token_emb = self.input_embedding(torch.clamp(full_token, min=0)) * mask
-
-        # Pre-lookahead processing
-        h = self.pre_lookahead_layer(token_emb)
-        h = h.repeat_interleave(self.token_mel_ratio, dim=1)
-
-        # Calculate mel lengths
-        mel_len1 = prompt_feat.shape[1]
-        mel_len2 = h.shape[1] - mel_len1
-
-        # Build conditioning
-        conds = torch.zeros(
-            [1, mel_len1 + mel_len2, self.output_size],
-            device=device,
-            dtype=h.dtype,
-        )
-        conds[:, :mel_len1] = prompt_feat
-        conds = conds.transpose(1, 2)
-
-        # Create mel mask
-        mel_mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h)
-
-        # Run flow matching decoder
-        feat, _ = self.decoder(
-            mu=h.transpose(1, 2).contiguous(),
-            mask=mel_mask.unsqueeze(1),
-            spks=embedding,
-            cond=conds,
+        """Generate audio waveform from speech tokens."""
+        feat = self._forward_mel(
+            token=token,
+            prompt_token=prompt_token,
+            prompt_feat=prompt_feat,
+            embedding=embedding,
             n_timesteps=n_timesteps,
+            token_offset_tokens=0,
+            streaming=False,
+            finalize=True,
         )
 
-        # Extract generated portion (after prompt)
-        feat = feat[:, :, mel_len1:]
-
         # Run vocoder
         hift_weight = self.hift.m_source.l_linear.weight
         tts_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype)
@@ -238,7 +316,7 @@ class CosyVoice3Code2Wav(nn.Module):
                 dtype=tts_mel.dtype,
             )
         else:
-            tts_speech, _ = self.hift.inference(speech_feat=tts_mel)
+            tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=True)
 
         return tts_speech
 

vllm_omni/model_executor/models/fish_speech/dac_encoder.py

@@ -54,6 +54,9 @@ def _load_dac_codec(
     if "generator" in state_dict:
         state_dict = state_dict["generator"]
     codec.load_state_dict(state_dict, strict=False)
+    # Encoder path only uses encoder + quantizer.forward(); prune the
+    # decoder before moving to device to avoid unnecessary GPU allocation.
+    codec.decoder = None
     codec = codec.to(device=device, dtype=dtype)
     codec.eval()