add

2026-04-01 04:35:27 -07:00 · 2026-04-01 04:26:09 -07:00 · 2026-04-01 04:20:25 -07:00
180 changed files with 2218 additions and 10516 deletions
@@ -10,7 +10,7 @@ steps:
          DOCKER_BUILDKIT: "1"
          # Buildkite will automatically replace this with the actual commit hash
          VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.19.0"
+          VLLM_VERSION: "v0.18.0"
        priority: 100
        timeout_in_minutes: 60
        soft_fail: true
@@ -56,8 +56,8 @@ steps:
    timeout_in_minutes: 20
    depends_on: upload-merge-pipeline
    commands:
-      # Single pytest session for one combined summary at end of log.
-      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
    agents:
      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
    plugins:
@@ -111,7 +111,8 @@ steps:
    timeout_in_minutes: 20
    depends_on: upload-merge-pipeline
    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py tests/diffusion/distributed/test_ulysses_uaa_perf.py
+      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+      - pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
    agents:
      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
    plugins:
@@ -173,7 +174,7 @@ steps:
                pytest -s -v tests/engine/test_async_omni_engine_abort.py
        '
    agents:
-      queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
    plugins:
      - docker#v5.2.0:
          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -192,7 +193,8 @@ steps:
    commands:
      - export VLLM_LOGGING_LEVEL=DEBUG
      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
+      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+      - pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
    agents:
      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
    plugins:
@@ -214,7 +216,7 @@ steps:
          export VLLM_LOGGING_LEVEL=DEBUG
          export VLLM_WORKER_MULTIPROC_METHOD=spawn
          export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
        '
    agents:
      queue: "gpu_1_queue"
@@ -237,7 +239,7 @@ steps:
          export VLLM_LOGGING_LEVEL=DEBUG
          export VLLM_WORKER_MULTIPROC_METHOD=spawn
          export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_base.py
        '
    agents:
      queue: "gpu_1_queue"
@@ -257,8 +259,9 @@ steps:
    depends_on: upload-merge-pipeline
    commands:
      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
+      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
+      - pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
    agents:
      queue: "mithril-h100-pool"
    plugins:
@@ -344,7 +347,8 @@ steps:
          export VLLM_TEST_CLEAN_GPU_MEMORY=1
          export VLLM_IMAGE_FETCH_TIMEOUT=60
          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory"
-          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
        '
    agents:
      queue: "mithril-h100-pool"
@@ -388,17 +392,8 @@ steps:
        timeout 20m bash -c '
          export VLLM_LOGGING_LEVEL=DEBUG
          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"
+          pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
        '
    agents:
      queue: "mithril-h100-pool"
@@ -418,11 +413,6 @@ steps:
                env:
                  - name: HF_HOME
                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
            nodeSelector:
              node.kubernetes.io/instance-type: gpu-h100-sxm
            volumes:
@@ -6,10 +6,16 @@ steps:
    commands:
      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
      - |
-        pytest -s -v \
-          tests/examples/ \
-          tests/e2e/online_serving/test_*_expansion.py \
-          -m "advanced_model and H100 and omni" --run-level "advanced_model"
+        set +e
+        pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "advanced_model" --run-level "advanced_model"
+        EXIT1=$$?
+        pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
+        EXIT2=$$?
+        pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+        EXIT3=$$?
+        pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+        EXIT4=$$?
+        exit $$((EXIT1 | EXIT2 | EXIT3 | EXIT4))
    agents:
      queue: "mithril-h100-pool"
    plugins:
@@ -51,7 +57,8 @@ steps:
    commands:
      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
      - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-      - pytest -s -v tests/examples/ tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
+      - pytest -s -v tests/examples/ -m "advanced_model and L4 and omni" --run-level "advanced_model"
+      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
    agents:
      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
    plugins:
@@ -107,7 +114,7 @@ steps:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate

-  - label: ":full_moon: Diffusion Model (Wan2.2) Test with H100"
+  - label: ":full_moon: Diffusion Model Wan22 completed Test with H100"
    timeout_in_minutes: 90
    depends_on: upload-nightly-pipeline
    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
@@ -148,7 +155,7 @@ steps:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate

-  - label: ":full_moon: Diffusion Model Test"
+  - label: ":full_moon: Diffusion Model Test with L4"
    timeout_in_minutes: 60
    depends_on: upload-nightly-pipeline
    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
@@ -170,7 +177,7 @@ steps:
            - "/fsx/hf_cache:/fsx/hf_cache"


-  - label: ":full_moon: Doc Example Code Test with H100"
+  - label: ":full_moon: Documentation Example Code Test with H100"
    timeout_in_minutes: 60
    depends_on: upload-nightly-pipeline
    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
@@ -212,7 +219,7 @@ steps:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate

-  - label: ":full_moon: Omni Model Perf Test & Testcase Statistics with H100"
+  - label: ":full_moon: Omni Model Perf Test & Test Case Statistics"
    key: nightly-omni-performance
    timeout_in_minutes: 180
    depends_on: upload-nightly-pipeline
@@ -362,7 +369,7 @@ steps:
      - kubernetes:
          podSpec:
            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:c392ce21e9cf9ea65c52b866447793db10e0261c
                resources:
                  limits:
                    nvidia.com/gpu: 2
@@ -390,9 +397,9 @@ steps:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate

-  - label: ":full_moon: Diffusion Perf Test with H100"
+  - label: ":full_moon: Qwen-Image Diffusion Perf Test with H100"
    key: nightly-qwen-image-performance
-    timeout_in_minutes: 180
+    timeout_in_minutes: 300
    depends_on: upload-nightly-pipeline
    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
    commands:
@@ -180,7 +180,7 @@ steps:
                pytest -s -v tests/engine/test_async_omni_engine_abort.py
        '
    agents:
-      queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
    plugins:
      - docker#v5.2.0:
          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -271,7 +271,7 @@ steps:
              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
                resources:
                  limits:
-                    nvidia.com/gpu: 1
+                    nvidia.com/gpu: 2
                volumeMounts:
                  - name: devshm
                    mountPath: /dev/shm
@@ -320,29 +320,6 @@ steps:
          volumes:
            - "/fsx/hf_cache:/fsx/hf_cache"

-  - label: "OmniVoice E2E Test"
-    timeout_in_minutes: 20
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 20m bash -c '
-          export VLLM_LOGGING_LEVEL=DEBUG
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/online_serving/test_omnivoice.py -m "core_model" --run-level "core_model"
-        '
-    agents:
-      queue: "gpu_1_queue"
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
  - label: "Voxtral-TTS E2E Test"
    timeout_in_minutes: 20
    depends_on: upload-ready-pipeline
@@ -351,7 +328,8 @@ steps:
        timeout 20m bash -c '
          export VLLM_LOGGING_LEVEL=DEBUG
          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "core_model" --run-level "core_model"
+          pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
        '
    agents:
      queue: "mithril-h100-pool"
@@ -548,46 +526,3 @@ steps:
                hostPath:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate
-
-  - label: "CosyVoice3-TTS E2E Test"
-    timeout_in_minutes: 20
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 20m bash -c '
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts.py -m "core_model" --run-level "core_model"
-        '
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
@@ -7,7 +7,7 @@ body:
 - type: markdown
  attributes:
    value: >
-      #### Please take a look at previous [RFCs](https://github.com/vllm-project/vllm-omni/issues?q=in%3Atitle%20RFC%20sort%3Aupdated-desc) for reference.
+      #### Please take a look at previous [RFCs](https://github.com/vllm-project/vllm-omni/issues?q=label%3ARFC+sort%3Aupdated-desc) for reference.
 - type: textarea
  attributes:
    label: Motivation.
@@ -21,7 +21,7 @@ body:
    description: >
      The proposed change of the RFC.
    value: |
-      Please provide the detailed design document of the RFC using the [template](https://docs.google.com/document/d/1jcgR3cDaUQH3VczD4ZcKaJAoYWHjCmnYzHkCNyz-9fk/edit?usp=sharing).
+      Please provide the detailed design document of the RFC using the [template](https://docs.google.com/document/d/12YxSsVeD1jvL-InClkeAEnZyWFDndz_65JmXvsamuV4/edit?tab=t.0#heading=h.4ef4szrsgspp).
  validations:
    required: true
 - type: textarea
@@ -32,9 +32,6 @@ bash run_benchmark.sh --hf-only
 # Use a different model (e.g. 1.7B)
 MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only

-# Use a Voice Clone model
-MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-Base TASK_TYPE=Base bash run_benchmark.sh --async-only
-
 # Use bs16 config for higher throughput
 STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs16.yaml bash run_benchmark.sh --async-only

@@ -23,9 +23,6 @@
 #   # Use 1.7B model:
 #   MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
 #
-#   # Use Voice Clone model
-#   MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-Base TASK_TYPE=Base bash run_benchmark.sh --async-only
-#
 #   # Use batch_size=4 config:
 #   STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
 #
@@ -38,7 +35,6 @@
 #   GPU_MEM_TALKER   - gpu_memory_utilization for talker stage (default: 0.3)
 #   GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.2)
 #   STAGE_CONFIG     - Path to stage config YAML (default: configs/qwen3_tts_bs1.yaml)
-#   TASK_TYPE        - Task type: CustomVoice, VoiceDesign, Base (default: CustomVoice)

 set -euo pipefail

@@ -57,7 +53,6 @@ NUM_WARMUPS="${NUM_WARMUPS:-3}"
 STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/configs/qwen3_tts_bs1.yaml}"
 RESULT_DIR="${SCRIPT_DIR}/results"
 TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
-TASK_TYPE="${TASK_TYPE:-CustomVoice}"

 # Parse args
 RUN_ASYNC=true
@@ -82,7 +77,6 @@ echo " Concurrency:  ${CONCURRENCY}"
 echo " Port:         ${PORT}"
 echo " Stage config: ${STAGE_CONFIG}"
 echo " Results:      ${RESULT_DIR}"
-echo " Task type:    ${TASK_TYPE}"
 echo "============================================================"

 # Prepare stage config with correct GPU device and memory settings
@@ -201,8 +195,7 @@ run_bench() {
        --max-concurrency ${conc_args} \
        --num-warmups "${NUM_WARMUPS}" \
        --config-name "${config_name}" \
-        --result-dir "${RESULT_DIR}" \
-        --task-type "${TASK_TYPE}"
+        --result-dir "${RESULT_DIR}"

    stop_server

@@ -229,8 +222,7 @@ if [ "${RUN_HF}" = true ]; then
        --num-warmups "${NUM_WARMUPS}" \
        --gpu-device "${GPU_DEVICE}" \
        --config-name "hf_transformers" \
-        --result-dir "${RESULT_DIR}" \
-        --task-type "${TASK_TYPE}"
+        --result-dir "${RESULT_DIR}"

    # Allow GPU memory to settle
    sleep 5
@@ -38,10 +38,6 @@ PROMPTS = [
    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
 ]

-REF_AUDIO = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
-REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
-INSTRUCT = "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
-

@dataclass
 class BenchmarkResult:
@@ -79,29 +75,6 @@ class BenchmarkResult:
    per_request: list = field(default_factory=list)


-def generate_audio(model, prompt: str, args):
-    if args.task_type == "Base":
-        return model.generate_voice_clone(
-            text=prompt,
-            language=args.language,
-            ref_audio=REF_AUDIO,
-            ref_text=REF_TEXT,
-        )
-
-    if args.task_type == "VoiceDesign":
-        return model.generate_voice_design(
-            text=prompt,
-            language=args.language,
-            instruct=INSTRUCT,
-        )
-
-    return model.generate_custom_voice(
-        text=prompt,
-        language=args.language,
-        speaker=args.voice,
-    )
-
-
 def run_benchmark(args):
    from qwen_tts import Qwen3TTSModel

@@ -122,7 +95,11 @@ def run_benchmark(args):
        print(f"Warming up with {args.num_warmups} requests...")
        for i in range(args.num_warmups):
            p = PROMPTS[i % len(PROMPTS)]
-            wavs, sr = generate_audio(model, p, args)
+            wavs, sr = model.generate_custom_voice(
+                text=p,
+                language=args.language,
+                speaker=args.voice,
+            )
        # Sync GPU
        torch.cuda.synchronize(device)
        print("Warmup done.")
@@ -147,7 +124,11 @@ def run_benchmark(args):
            torch.cuda.synchronize(device)
            st = time.perf_counter()

-            wavs, sr = generate_audio(model, prompt, args)
+            wavs, sr = model.generate_custom_voice(
+                text=prompt,
+                language=args.language,
+                speaker=args.voice,
+            )

            torch.cuda.synchronize(device)
            elapsed = time.perf_counter() - st
@@ -287,7 +268,6 @@ def parse_args():
    parser.add_argument("--gpu-device", type=int, default=0)
    parser.add_argument("--voice", type=str, default="Vivian")
    parser.add_argument("--language", type=str, default="English")
-    parser.add_argument("--task-type", type=str, default="CustomVoice", choices=["CustomVoice", "VoiceDesign", "Base"])
    parser.add_argument(
        "--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
    )
@@ -37,9 +37,6 @@ PROMPTS = [
    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
 ]
-REF_AUDIO = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
-REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
-INSTRUCT = "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."


@dataclass
@@ -96,39 +93,22 @@ def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width
    return num_samples / sample_rate


-def create_payload(
-    prompt: str, task_type: str = "CustomVoice", voice: str = "vivian", language: str = "English"
-) -> dict:
-    payload = {
-        "input": prompt,
-        "language": language,
-        "stream": True,
-        "response_format": "pcm",
-        "task_type": task_type,
-    }
-
-    if task_type == "Base":
-        payload["ref_audio"] = REF_AUDIO
-        payload["ref_text"] = REF_TEXT
-    elif task_type == "CustomVoice":
-        payload["voice"] = voice
-    elif task_type == "VoiceDesign":
-        payload["instructions"] = INSTRUCT
-
-    return payload
-
-
 async def send_tts_request(
    session: aiohttp.ClientSession,
    api_url: str,
    prompt: str,
-    task_type: str = "CustomVoice",
    voice: str = "vivian",
    language: str = "English",
    pbar: tqdm | None = None,
 ) -> RequestResult:
    """Send a streaming TTS request and measure latency metrics."""
-    payload = create_payload(prompt, task_type, voice, language)
+    payload = {
+        "input": prompt,
+        "voice": voice,
+        "language": language,
+        "stream": True,
+        "response_format": "pcm",
+    }

    result = RequestResult(prompt=prompt)
    st = time.perf_counter()
@@ -173,7 +153,6 @@ async def run_benchmark(
    num_prompts: int,
    max_concurrency: int,
    num_warmups: int = 3,
-    task_type: str = "CustomVoice",
    voice: str = "vivian",
    language: str = "English",
 ) -> BenchmarkResult:
@@ -196,7 +175,7 @@ async def run_benchmark(
        warmup_tasks = []
        for i in range(num_warmups):
            prompt = PROMPTS[i % len(PROMPTS)]
-            warmup_tasks.append(send_tts_request(session, api_url, prompt, task_type, voice, language))
+            warmup_tasks.append(send_tts_request(session, api_url, prompt, voice, language))
        await asyncio.gather(*warmup_tasks)
        print("  Warmup done.")

@@ -210,7 +189,7 @@ async def run_benchmark(

    async def limited_request(prompt):
        async with semaphore:
-            return await send_tts_request(session, api_url, prompt, task_type, voice, language, pbar)
+            return await send_tts_request(session, api_url, prompt, voice, language, pbar)

    start_time = time.perf_counter()
    tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
@@ -327,7 +306,6 @@ async def main(args):
            num_prompts=args.num_prompts,
            max_concurrency=concurrency,
            num_warmups=args.num_warmups,
-            task_type=args.task_type,
            voice=args.voice,
            language=args.language,
        )
@@ -356,7 +334,6 @@ def parse_args():
        "--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
    )
    parser.add_argument("--num-warmups", type=int, default=3)
-    parser.add_argument("--task-type", type=str, default="CustomVoice", choices=["CustomVoice", "VoiceDesign", "Base"])
    parser.add_argument("--voice", type=str, default="vivian")
    parser.add_argument("--language", type=str, default="English")
    parser.add_argument(
@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.19.0
+ARG VLLM_BASE_TAG=v0.18.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -11,7 +11,9 @@ RUN apt-get update && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

-RUN uv pip install --system ".[dev]"
+# Install vllm-omni into the same uv-managed Python environment used by the base image.
+# Use bash -c so that $(python3 -c ...) is expanded inside the container.
+RUN uv pip install --system --no-cache-dir ".[dev]"

 RUN ln -sf /usr/bin/python3 /usr/bin/python

@@ -1,20 +1,20 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.18.0rc1
+ARG VLLM_ASCEND_TAG=v0.17.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}

-# WORKDIR /vllm-workspace/vllm
-# RUN git fetch origin --tags && git checkout v0.18.0
+WORKDIR /vllm-workspace/vllm
+RUN git fetch origin --tags && git checkout v0.18.0

-# WORKDIR /vllm-workspace/vllm-ascend
-# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
-# # Install vllm-ascend
-# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
-# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-#     source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
-#     source /usr/local/Ascend/nnal/atb/set_env.sh && \
-#     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-#     python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
-#     python3 -m pip cache purge
+WORKDIR /vllm-workspace/vllm-ascend
+RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+# Install vllm-ascend
+# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
+    python3 -m pip cache purge

 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -1,20 +1,20 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.18.0rc1-a3
+ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}

-# WORKDIR /vllm-workspace/vllm
-# RUN git fetch origin --tags && git checkout v0.18.0
+WORKDIR /vllm-workspace/vllm
+RUN git fetch origin --tags && git checkout v0.18.0

-# WORKDIR /vllm-workspace/vllm-ascend
-# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
-# # Install vllm-ascend
-# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
-# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-#     source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
-#     source /usr/local/Ascend/nnal/atb/set_env.sh && \
-#     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-#     python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
-#     python3 -m pip cache purge
+WORKDIR /vllm-workspace/vllm-ascend
+RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+# Install vllm-ascend
+# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
+    python3 -m pip cache purge

 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.19.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.18.0
 FROM ${BASE_IMAGE} AS base

 # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -76,7 +76,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"

-ARG VLLM_VERSION=v0.19.0
+ARG VLLM_VERSION=v0.18.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm

@@ -5,6 +5,7 @@
 Main entry points for vLLM-Omni inference and serving.

 - [vllm_omni.entrypoints.async_omni.AsyncOmni][]
+- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
 - [vllm_omni.entrypoints.cfg_companion_tracker.CfgCompanionTracker][]
 - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
 - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
@@ -168,6 +168,7 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
+│   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
 │   │   └── ...                                  │   │   └── test_*.py                  ⬜
@@ -73,6 +73,7 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
+│   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
 │   │   └── ...                                  │   │   └── test_*.py                  ⬜
@@ -69,7 +69,7 @@
 [5] Orchestrator._orchestration_loop (loop)
    -> poll stage output
       - llm stage: await get_output_async()
-       - diffusion stage: get_diffusion_output_nowait()
+       - diffusion stage: get_diffusion_output_async()
    -> (llm stage) output_processors[i].process_outputs(...)
    -> _route_output(...)
    -> if finished and not final_stage and non-async-chunk:
@@ -112,7 +112,7 @@ sequenceDiagram
    ORCH->>S0: add_request_async

    loop poll route forward
-        ORCH->>S0: get_output_async / get_diffusion_output_nowait
+        ORCH->>S0: get_output_async / get_diffusion_output_async
        ORCH->>ORCH: _route_output
        alt need forward to next stage
            ORCH->>SN: add_request_async
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --8<-- [start:build-wheel-from-source]

 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library

 ```bash
-uv pip install vllm==0.19.0 --torch-backend=auto
+uv pip install vllm==0.18.0 --torch-backend=auto
 ```

-The 0.19.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.18.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.19.0
+git checkout v0.18.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.19.0/vllm-0.19.0+cu130-cp38-abi3-manylinux_2_35_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.18.0/vllm-0.18.0+cu129-cp38-abi3-manylinux_2_35_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.

 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
+uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
 ```

 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
 # --8<-- [start:build-wheel-from-source]

 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library

 ```bash
-uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
+uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
 ```

-The pre-built 0.19.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.18.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.19.0
+git checkout v0.18.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  --env "HF_TOKEN=$HF_TOKEN" \
  -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.19.0 \
+  vllm/vllm-omni-rocm:v0.18.0 \
  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```

@@ -149,7 +149,7 @@ docker run --rm -it \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  --env "HF_TOKEN=$HF_TOKEN" \
  --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.19.0
+  vllm/vllm-omni-rocm:v0.18.0
 ```

 # --8<-- [end:pre-built-images]
@@ -10,10 +10,10 @@ The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built
 ```bash
 # Update the vllm-ascend image
 # Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1-a3
+export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1
 docker run --rm \
    --name vllm-omni-npu \
    --shm-size=1g \
@@ -33,6 +33,17 @@ docker run --rm \
    -p 8000:8000 \
    -it $IMAGE bash

+cd /vllm-workspace/vllm
+git fetch origin --tags
+git checkout v0.18.0
+
+# Because vllm-ascend will release v0.18.0rc1 after vllm-omni 0.16.0,
+# we have to pin vllm-ascend at the current commit.
+cd /vllm-workspace/vllm-ascend
+git pull origin main
+git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+pip install -v -e .
+
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
 git clone -b v0.18.0 https://github.com/vllm-project/vllm-omni.git
@@ -57,10 +68,15 @@ You can also build vLLM-Omni from the latest main branch if you want to use the

 ```bash
 # Pin vLLM version to 0.18.0
-git clone -b v0.18.0 https://github.com/vllm-project/vllm.git
-VLLM_TARGET_DEVICE=empty pip install -v -e .
+cd /vllm-workspace/vllm
+git fetch origin --tags
+git checkout v0.18.0

-git clone -b v0.18.0rc1 https://github.com/vllm-project/vllm-ascend.git
+# Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
+cd /vllm-workspace/vllm-ascend
+git pull origin main
+git fetch origin --tags
+git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
 pip install -v -e .

 # Install vLLM-Omni from the latest main branch
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate

 # On CUDA
-uv pip install vllm==0.19.0 --torch-backend=auto
+uv pip install vllm==0.18.0 --torch-backend=auto

 # On ROCm
-uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
+uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700

 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -120,7 +120,6 @@ def extract_omni_serve_subparser_init():
                        "_FlexibleArgumentParser": _FlexibleArgumentParser,
                        "FlexibleArgumentParser": _FlexibleArgumentParser,
                        "make_arg_parser": lambda parser: parser,  # no-op for doc
-                        "_ensure_vllm_platform": lambda: None,  # no-op for doc
                        "VLLM_SUBCMD_PARSER_EPILOG": "",
                        "logger": logger,
                        "DummySubparsers": DummySubparsers,
@@ -31,11 +31,8 @@ th {
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
 | `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
-| `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` | ✅︎ | ✅︎ | ✅︎ | |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
 | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
@@ -118,7 +118,6 @@ Content-Type: application/json
 | `instructions` | string | "" | Voice style/emotion instructions |
 | `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
 | `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
-| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |

 **Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian

@@ -144,23 +143,9 @@ Lists available voices for the loaded model.

 ```json
 {
-    "voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian", "custom_voice_1"],
-    "uploaded_voices": [
-        {
-            "name": "custom_voice_1",
-            "consent": "user_consent_id",
-            "created_at": 1738660000,
-            "file_size": 1024000,
-            "mime_type": "audio/wav",
-            "ref_text": "The exact transcript of the audio sample.",
-            "speaker_description": "warm narrator"
-        }
-    ]
+    "voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
 }
 ```
-
-`uploaded_voices` is always present (empty list when no custom voices have been uploaded). Fields `ref_text` and `speaker_description` are omitted per-entry when not provided at upload time.
-
 ```
 POST /v1/audio/voices
 Content-Type: multipart/form-data
@@ -176,7 +161,6 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
 | `consent` | string | Yes | Consent recording ID |
 | `name` | string | Yes | Name for the new voice |
 | `ref_text` | string | No | Transcript of the audio. When provided, enables in-context voice cloning (higher quality). Without it, only the speaker embedding is extracted. |
-| `speaker_description` | string | No | Free-form description of the voice (e.g. "warm narrator", "energetic presenter"). Stored as metadata and returned in `GET /v1/audio/voices`. |

 **Response Example:**

@@ -188,15 +172,11 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
    "consent": "user_consent_id",
    "created_at": 1738660000,
    "mime_type": "audio/wav",
-    "file_size": 1024000,
-    "ref_text": "The exact transcript of the audio sample.",
-    "speaker_description": "warm narrator"
+    "file_size": 1024000
  }
 }
 ```

-Fields `ref_text` and `speaker_description` are omitted when not provided at upload time.
-
 **Usage Example:**

 ```bash
@@ -204,8 +184,7 @@ curl -X POST http://localhost:8091/v1/audio/voices \
  -F "audio_sample=@/path/to/voice_sample.wav" \
  -F "consent=user_consent_id" \
  -F "name=custom_voice_1" \
-  -F "ref_text=The exact transcript of the audio sample." \
-  -F "speaker_description=warm narrator"
+  -F "ref_text=The exact transcript of the audio sample."
 ```

 ## Streaming Text Input (WebSocket)
@@ -164,7 +164,7 @@ cache_config={

 **Performance Tips**:

- Default `Fn_compute_blocks=1` works well for most cases. Some models (e.g., FLUX.2-klein) use a larger value for `Fn_compute_blocks` for a balanced performance.
+- Default `Fn_compute_blocks=1` works well for most cases. Some models (e.g., [FLUX.2-klein](https://github.com/wtomin/vllm-omni/blob/main/vllm_omni/diffusion/cache/cache_dit_backend.py#L363)) use a larger value for `Fn_compute_blocks` for a balanced performance.
 - Increase `residual_diff_threshold` (e.g., 0.12-0.15) for faster inference with slight quality trade-off, or decrease from default 0.24 for higher quality.
 - Default `max_warmup_steps=4` is optimized for few-step models. Increase to 6-8 for more steps if needed.

@@ -1,91 +0,0 @@
-# AutoRound Quantization
-
-## Overview
-
-[AutoRound](https://github.com/intel/auto-round) is an advanced quantization toolkit designed for Large Language Models (LLMs), Vision-Language Models (VLMs), and diffusion models. It achieves high accuracy at ultra-low bit widths (2–4 bits) with minimal tuning by leveraging sign-gradient descent, while providing broad hardware compatibility with multi-datatype support.
-
-The quantization config is auto-detected from the checkpoint's `config.json` (`quantization_config.quant_method = "auto-round"`). No extra CLI flags are needed.
-
-### Supported Schemes
-
-| Scheme | Bits | Status |
-|--------|------|--------|
-| W4A16 | 4 | ✅ Supported |
-| W8A16 | 8 | Planned |
-
-W4A16 is the first supported scheme. Additional schemes will be added in future releases.
-
-## Configuration
-
-1. **Python API**: point `model` at a pre-quantized checkpoint. The quantization is detected automatically.
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(model="vllm-project-org/FLUX.1-dev-AutoRound-w4a16")
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=28),
-)
-outputs[0].save_images("output.png")
-```
-
-2. **CLI**: pass the quantized model path directly.
-
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model vllm-project-org/FLUX.1-dev-AutoRound-w4a16 \
-  --prompt "A cat sitting on a windowsill" \
-  --num-inference-steps 28 \
-  --output outputs/flux_w4a16.png
-```
-
-No `--quantization` flag is needed — the quantization method is read from the checkpoint.
-
-## How It Works
-
-The checkpoint's `config.json` contains:
-
-```json
-{
-  "quantization_config": {
-    "quant_method": "auto-round",
-    "bits": 4,
-    "group_size": 128,
-    "sym": true,
-    "packing_format": "auto_round:auto_gptq",
-    "block_name_to_quantize": "transformer_blocks,single_transformer_blocks"
-  }
-}
-```
-
-At load time:
-
-1. `TransformerConfig.from_dict()` parses the `quantization_config` section and builds a vLLM `INCConfig` via `build_quant_config("auto-round", ...)`.
-2. `OmniDiffusionConfig.set_tf_model_config()` propagates the detected config to the engine.
-3. The appropriate compute kernel (e.g. GPTQ-Marlin for W4A16) is selected automatically based on the checkpoint's bit-width and packing format.
-
-## Supported Models
-
-| Model | HF Checkpoint | Scheme | Group Size | Backend |
-|-------|--------------|--------|------------|---------|
-| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | W4A16 | 128 | GPTQ-Marlin |
-
-## Creating a Quantized Checkpoint
-
-Use [AutoRound](https://github.com/intel/auto-round) to quantize a BF16 model. The `--scheme` flag selects the quantization scheme:
-
-```bash
-# W4A16 (4-bit weight, 16-bit activation)
-auto-round \
-    --model black-forest-labs/FLUX.1-dev \
-    --scheme W4A16 \
-    --batch_size 1 \
-    --disable_opt_rtn \
-    --dataset coco2014 \
-    --iters 0
-```
-
-The output directory can be used directly as the `model` argument. See the [AutoRound documentation](https://github.com/intel/auto-round) for all available schemes and options.
@@ -11,7 +11,6 @@ vLLM-Omni provides a unified quantization framework that supports both diffusion
 | FP8 | [FP8](fp8.md) | FP8 W8A8, dynamic or static | Z-Image, Qwen-Image, Flux, Bagel | SM 89 (Ada) |
 | Int8 | [Int8](int8.md) | Int8 W8A8 | Z-Image, Qwen-Image | SM 89 (Ada) / Ascend NPU |
 | GGUF | [GGUF](gguf.md) | GGUF format, dequant+GEMM for N-D tensors | Z-Image, Flux | SM 60 |
-| AutoRound | [AutoRound](autoround.md) | W4A16 (pre-quantized) | Flux | SM 80 (Ampere) |

 ### Multi-stage Omni Models (Pre-quantized Checkpoints)

@@ -103,9 +102,6 @@ config = build_quant_config("fp8")
 # Dict with parameters
 config = build_quant_config({"method": "fp8", "activation_scheme": "static"})

-# AutoRound / INC (auto-detected from checkpoint, or explicit)
-config = build_quant_config("auto-round", bits=4, group_size=128)
-
 # Per-component dict
 config = build_quant_config({
    "transformer": {"method": "fp8"},
@@ -105,21 +105,21 @@ The following tables show which models support each feature:
 | **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
 | **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
 | **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
+| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
 | **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **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** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ |
+| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ |
+| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ |
+| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ❌ | ✅ | ✅ | ❌ |

 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
@@ -130,7 +130,6 @@ The following tables show which models support each feature:
 | 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** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
@@ -159,7 +159,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
    -H "Content-Type: application/json" \
    -d '{
        "input": "Hello, how are you?",
-        "voice": "vivian",
+        "speaker": "vivian",
        "language": "English"
    }' --output output.wav

@@ -168,7 +168,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
    -H "Content-Type: application/json" \
    -d '{
        "input": "I am so excited!",
-        "voice": "vivian",
+        "speaker": "vivian",
        "instructions": "Speak with great enthusiasm"
    }' --output excited.wav

@@ -185,7 +185,7 @@ client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")

 response = client.audio.speech.create(
    model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
-    voice="vivian",
+    speaker="vivian",
    input="Hello, how are you?",
 )

@@ -201,7 +201,7 @@ response = httpx.post(
    "http://localhost:8091/v1/audio/speech",
    json={
        "input": "Hello, how are you?",
-        "voice": "vivian",
+        "speaker": "vivian",
        "language": "English",
    },
    timeout=300.0,
@@ -237,16 +237,12 @@ List all available voices/speakers from the loaded model, including both built-i
      "consent": "user_consent_id",
      "created_at": 1738660000,
      "file_size": 1024000,
-      "mime_type": "audio/wav",
-      "ref_text": "The exact transcript of the audio sample.",
-      "speaker_description": "warm narrator"
+      "mime_type": "audio/wav"
    }
  ]
 }
 ```

-Fields `ref_text` and `speaker_description` are omitted per-entry when not provided at upload time.
-
 #### POST /v1/audio/voices

 Upload a new voice sample for voice cloning in Base task TTS requests.
@@ -256,7 +252,6 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
 - `consent` (required): Consent recording ID
 - `name` (required): Name for the new voice
 - `ref_text` (optional): Transcript of the audio. Enables in-context voice cloning (higher quality).
- `speaker_description` (optional): Free-form description of the voice (e.g. "warm narrator", "energetic presenter"). Stored as metadata.

 **Response Example:**
 ```json
@@ -267,23 +262,18 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
    "consent": "user_consent_id",
    "created_at": 1738660000,
    "mime_type": "audio/wav",
-    "file_size": 1024000,
-    "ref_text": "The exact transcript of the audio sample.",
-    "speaker_description": "warm narrator"
+    "file_size": 1024000
  }
 }
 ```

-Fields `ref_text` and `speaker_description` are omitted when not provided at upload time.
-
 **Usage Example:**
 ```bash
-curl -X POST http://localhost:8091/v1/audio/voices \
+curl -X POST http://localhost:8000/v1/audio/voices \
  -F "audio_sample=@/path/to/voice_sample.wav" \
  -F "consent=user_consent_id" \
  -F "name=custom_voice_1" \
-  -F "ref_text=The exact transcript of the audio sample." \
-  -F "speaker_description=warm narrator"
+  -F "ref_text=The exact transcript of the audio sample."
 ```

 ### Endpoint
@@ -300,7 +290,7 @@ This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/
 ```json
 {
    "input": "Text to synthesize",
-    "voice": "vivian",
+    "speaker": "vivian",
    "response_format": "wav",
    "task_type": "CustomVoice",
    "language": "Auto",
@@ -320,7 +310,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w

 ### Voice and language (summary)

- **Speaker**: Use the `voice` request field to select the speaker (e.g., `vivian`, `ryan`, `aiden`). List available speakers with `GET /v1/audio/voices`.
+- **Speaker**: Use the `speaker` request field to select the speaker (e.g., `vivian`, `ryan`, `aiden`). List available speakers with `GET /v1/audio/voices`.
 - **Language**: Use the `language` field for the codec language tag (`Auto`, `Chinese`, `English`, etc.). Default is `Auto` for automatic detection.
 - **CustomVoice**: Requires a valid `voice` from the model’s speaker set. **VoiceDesign**: Use `instructions` to describe the voice. **Base**: Use `ref_audio` and `ref_text` for voice cloning.

@@ -332,7 +322,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
 | ----------------- | ------ | -------------- | ----------------------------------------------------------- |
 | `input`           | string | **required**   | Text to synthesize                                          |
 | `model`           | string | server's model | Model to use (optional, should match server if specified)   |
-| `voice`           | string | "vivian"       | Speaker name (e.g., vivian, ryan, aiden)                    |
+| `speaker`         | string | "vivian"       | Speaker name (e.g., vivian, ryan, aiden)                    |
 | `response_format` | string | "wav"          | Audio format: wav, mp3, flac, pcm, aac, opus                |
 | `speed`           | float  | 1.0            | Playback speed (0.25-4.0, not supported with `stream=true`) |

@@ -367,7 +357,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
    -H "Content-Type: application/json" \
    -d '{
        "input": "Hello, how are you?",
-        "voice": "vivian",
+        "speaker": "vivian",
        "language": "English",
        "stream": true,
        "response_format": "pcm"
@@ -2,7 +2,6 @@ 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


 def parse_args():
@@ -66,17 +65,6 @@ def parse_args():
        help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
    )
    parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.")
-    parser.add_argument(
-        "--cfg-interval",
-        type=float,
-        nargs=2,
-        default=None,
-        help="CFG interval [start, end] (default: pipeline default)",
-    )
-    parser.add_argument(
-        "--cfg-renorm-type", type=str, default=None, help="CFG renorm type: global, text_channel, channel"
-    )
-    parser.add_argument("--cfg-renorm-min", type=float, default=None, help="CFG renorm min")
    parser.add_argument(
        "--enable-diffusion-pipeline-profiler",
        action="store_true",
@@ -88,12 +76,6 @@ def parse_args():
        default=None,
        help="Quantization method (e.g. 'fp8').",
    )
-    parser.add_argument(
-        "--think",
-        action="store_true",
-        default=False,
-        help="Enable thinking mode: AR stage decodes <think>...</think> planning tokens before image generation.",
-    )

    args = parser.parse_args()
    return args
@@ -128,12 +110,8 @@ def main():
    from vllm_omni.entrypoints.omni import Omni

    omni_kwargs = {}
-    stage_configs_path = args.stage_configs_path
-    if args.think and stage_configs_path is None:
-        stage_configs_path = "vllm_omni/model_executor/stage_configs/bagel_think.yaml"
-        print(f"[Info] Think mode enabled, using stage config: {stage_configs_path}")
-    if stage_configs_path:
-        omni_kwargs["stage_configs_path"] = stage_configs_path
+    if args.stage_configs_path:
+        omni_kwargs["stage_configs_path"] = args.stage_configs_path

    omni_kwargs.update(
        {
@@ -158,8 +136,7 @@ def main():
            if not args.image_path or not os.path.exists(args.image_path):
                raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}")
            loaded_image = Image.open(args.image_path).convert("RGB")
-            think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
-            final_prompt_text = f"{think_prefix}<|fim_middle|><|im_start|>{p}<|im_end|>"
+            final_prompt_text = f"<|fim_middle|><|im_start|>{p}<|im_end|>"
            prompt_dict = {
                "prompt": final_prompt_text,
                "multi_modal_data": {"img2img": loaded_image},
@@ -183,8 +160,7 @@ def main():
            prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
            formatted_prompts.append(prompt_dict)
        else:
-            think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
-            final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|>"
+            final_prompt_text = f"<|im_start|>{p}<|im_end|>"
            prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
            if args.negative_prompt is not None:
                prompt_dict["negative_prompt"] = args.negative_prompt
@@ -202,12 +178,6 @@ def main():
                "cfg_text_scale": args.cfg_text_scale,
                "cfg_img_scale": args.cfg_img_scale,
            }
-            if args.cfg_interval is not None:
-                extra["cfg_interval"] = tuple(args.cfg_interval)
-            if args.cfg_renorm_type is not None:
-                extra["cfg_renorm_type"] = args.cfg_renorm_type
-            if args.cfg_renorm_min is not None:
-                extra["cfg_renorm_min"] = args.cfg_renorm_min
            if args.negative_prompt is not None:
                extra["negative_prompt"] = args.negative_prompt
            diffusion_params.extra_args = extra  # type: ignore
@@ -216,17 +186,6 @@ 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}")
-
        images = getattr(req_output, "images", None)

        if not images:
@@ -235,7 +194,6 @@ def main():
        for j, img in enumerate(images):
            save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
            img.save(save_path)
-            print(f"[Output] Saved image to {save_path}")
        img_idx += 1

    print(omni_outputs)
@@ -490,6 +490,10 @@ def main():
        if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
            frames_np = frames_np[..., :3]

+        frames_np = np.clip(frames_np, 0.0, 1.0)
+        frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
+        video_tensor = torch.from_numpy(frames_u8)
+
        audio_out = None
        if audio is not None:
            if isinstance(audio, list):
@@ -503,7 +507,7 @@ def main():
                audio_out = audio_out.float().cpu()

        encode_video(
-            frames_np,
+            video_tensor,
            fps=fps,
            audio=audio_out,
            audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
@@ -1,73 +0,0 @@
-# OmniVoice
-
-This directory contains an offline demo for running OmniVoice TTS models with vLLM Omni. It generates speech from text and saves WAV files locally.
-
-## Model Overview
-
-[OmniVoice](https://huggingface.co/k2-fsa/OmniVoice) is a zero-shot multilingual TTS model supporting 600+ languages. It uses a diffusion language model (Qwen3-0.6B backbone) with iterative masked unmasking to generate speech.
-
-Three inference modes are supported:
-
- **Auto Voice**: Generate speech without any reference — the model picks a voice automatically.
- **Voice Clone**: Clone a voice from a reference audio + transcription.
- **Voice Design**: Control voice style via natural language instruction (e.g., "female, low pitch, british accent").
-
-## Setup
-
-Ensure the model is downloaded:
-
-```bash
-huggingface-cli download k2-fsa/OmniVoice
-```
-
-> **Note:** Voice cloning requires `transformers>=5.3.0` for `HiggsAudioV2TokenizerModel`. Auto voice and voice design modes work with `transformers>=4.57.0`.
-
-## Quick Start
-
-Auto voice (text only):
-
-```bash
-python end2end.py --model k2-fsa/OmniVoice --text "Hello, this is a test."
-```
-
-Voice design (with style instruction):
-
-```bash
-python end2end.py --model k2-fsa/OmniVoice \
-    --text "Hello, this is a test." \
-    --instruct "female, low pitch, british accent"
-```
-
-Voice clone (with reference audio):
-
-```bash
-python end2end.py --model k2-fsa/OmniVoice \
-    --text "Hello, this is a test." \
-    --ref-audio ref.wav \
-    --ref-text "This is the reference transcription."
-```
-
-## Language Support
-
-Specify a language for improved quality:
-
-```bash
-python end2end.py --model k2-fsa/OmniVoice \
-    --text "你好，这是一个测试。" \
-    --lang zh
-```
-
-## Architecture
-
-OmniVoice uses a two-stage pipeline:
-
- **Stage 0 (Generator)**: Qwen3-0.6B transformer with 32-step iterative unmasking and classifier-free guidance. Generates 8-codebook audio tokens from text.
- **Stage 1 (Decoder)**: HiggsAudioV2 RVQ quantizer + DAC acoustic decoder. Converts tokens to 24kHz waveform.
-
-Both stages use `GPUGenerationWorker` with `OmniGenerationScheduler`.
-
-## Notes
-
- Output audio is saved to `output.wav` by default. Use `--output` to change the path.
- The model estimates duration from text automatically via `RuleDurationEstimator`.
- Use `--stage-init-timeout` to increase the stage initialization timeout for first-time model downloads.
@@ -1,164 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""End-to-end OmniVoice TTS inference via vLLM-Omni.
-
-Supports:
- Auto voice mode: text only → generated speech
- Voice cloning mode: text + reference audio → cloned voice speech
-
-Usage:
-    # Auto voice
-    python end2end.py --model k2-fsa/OmniVoice --text "Hello world"
-
-    # Voice cloning
-    python end2end.py --model k2-fsa/OmniVoice --text "Hello" \
-        --ref-audio ref.wav --ref-text "reference transcription"
-"""
-
-import argparse
-import os
-
-import numpy as np
-import soundfile as sf
-
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-
-def run_e2e():
-    parser = argparse.ArgumentParser(description="OmniVoice E2E TTS inference")
-    parser.add_argument(
-        "--model",
-        type=str,
-        default="k2-fsa/OmniVoice",
-        help="Model name or path (HuggingFace or local)",
-    )
-    parser.add_argument(
-        "--stage-config",
-        type=str,
-        default="vllm_omni/model_executor/stage_configs/omnivoice.yaml",
-    )
-    parser.add_argument(
-        "--text",
-        type=str,
-        default="Hello, this is a test of the OmniVoice text to speech system.",
-    )
-    parser.add_argument(
-        "--ref-audio",
-        type=str,
-        default=None,
-        help="Reference audio for voice cloning (WAV file)",
-    )
-    parser.add_argument(
-        "--ref-text",
-        type=str,
-        default=None,
-        help="Transcription of reference audio",
-    )
-    parser.add_argument(
-        "--lang",
-        type=str,
-        default=None,
-        help="Language code (e.g., 'en', 'zh')",
-    )
-    parser.add_argument(
-        "--instruct",
-        type=str,
-        default=None,
-        help="Voice design instruction (e.g., 'female, low pitch, british accent')",
-    )
-    parser.add_argument(
-        "--output",
-        type=str,
-        default="output.wav",
-        help="Output audio file path",
-    )
-    parser.add_argument(
-        "--stage-init-timeout",
-        type=int,
-        default=600,
-        help="Stage initialization timeout in seconds",
-    )
-    args = parser.parse_args()
-
-    if not os.path.exists(args.stage_config):
-        raise FileNotFoundError(f"Stage config not found: {args.stage_config}")
-
-    print(f"Initializing OmniVoice with model={args.model}")
-
-    omni = Omni(
-        model=args.model,
-        stage_configs_path=args.stage_config,
-        trust_remote_code=True,
-        log_stats=True,
-    )
-
-    print("Model initialized. Preparing inputs...")
-
-    # Build prompt
-    mm_processor_kwargs = {}
-    multi_modal_data = {}
-
-    if args.ref_audio:
-        if not os.path.exists(args.ref_audio):
-            raise FileNotFoundError(f"Reference audio not found: {args.ref_audio}")
-
-        import librosa
-
-        audio_signal, sr = librosa.load(args.ref_audio, sr=None)
-        multi_modal_data["audio"] = (audio_signal.astype(np.float32), sr)
-        mm_processor_kwargs["ref_text"] = args.ref_text or ""
-        mm_processor_kwargs["sample_rate"] = sr
-
-    if args.lang:
-        mm_processor_kwargs["lang"] = args.lang
-    if args.instruct:
-        mm_processor_kwargs["instruct"] = args.instruct
-
-    prompts = {"prompt": args.text}
-    if multi_modal_data:
-        prompts["multi_modal_data"] = multi_modal_data
-    if mm_processor_kwargs:
-        prompts["mm_processor_kwargs"] = mm_processor_kwargs
-
-    sampling_params_list = [OmniDiffusionSamplingParams()]
-
-    print(f"Generating speech for: {args.text}")
-
-    outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list))
-
-    print(f"Received {len(outputs)} outputs.")
-    for i, output in enumerate(outputs):
-        try:
-            ro = output.request_output
-            if ro is None:
-                print("No request_output found.")
-                continue
-
-            mm = getattr(ro, "multimodal_output", None)
-            if not mm and ro.outputs:
-                mm = getattr(ro.outputs[0], "multimodal_output", None)
-
-            if mm:
-                print(f"Multimodal output keys: {mm.keys()}")
-                if "audio" in mm:
-                    audio_out = mm["audio"]
-                    sr = mm.get("sr", 24000)
-                    if isinstance(audio_out, np.ndarray):
-                        audio_np = audio_out
-                    else:
-                        audio_np = audio_out.cpu().numpy().squeeze()
-                    out_path = args.output if i == 0 else f"output_{i}.wav"
-                    sf.write(out_path, audio_np, sr)
-                    print(f"Saved audio to {out_path} ({sr}Hz, {len(audio_np) / sr:.2f}s)")
-            else:
-                print("No multimodal output found.")
-        except Exception as e:
-            print(f"Error inspecting output: {e}")
-
-    omni.close()
-    print("Done.")
-
-
-if __name__ == "__main__":
-    run_e2e()
@@ -56,13 +56,8 @@ def parse_args() -> argparse.Namespace:
        "Examples: Wan-AI/Wan2.2-T2V-A14B-Diffusers, "
        "hunyuanvideo-community/HunyuanVideo-1.5-480p_t2v",
    )
-    parser.add_argument(
-        "--model-class-name",
-        default=None,
-        help="Override model class name (e.g., LTX2TwoStagesVideoPipeline).",
-    )
    parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
-    parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
+    parser.add_argument("--negative-prompt", default="", help="Negative prompt (Wan2.2 only).")
    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
    parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
    parser.add_argument(
@@ -190,7 +185,6 @@ def parse_args() -> argparse.Namespace:

 def main():
    args = parse_args()
-    model_class_name = args.model_class_name

    preset = _detect_preset(args.model)
    for key, default_val in preset.items():
@@ -235,7 +229,6 @@ def main():
        enable_cpu_offload=args.enable_cpu_offload,
        parallel_config=parallel_config,
        enforce_eager=args.enforce_eager,
-        model_class_name=model_class_name,
        cache_backend=args.cache_backend,
        cache_config=cache_config,
        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
@@ -1,88 +0,0 @@
-# VACE Video Generation
-
-[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single model.
-
-| Model | Architecture | Model Weights (bf16) | HuggingFace |
-|-------|-------------|----------------------|-------------|
-| Wan2.1-VACE (1.3B) | Wan2.1 | ~10 GB | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
-| Wan2.1-VACE (14B) | Wan2.1 | ~38 GB | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
-
-## Text-to-Video (T2V)
-
-```bash
-python vace_video_generation.py \
-  --mode t2v \
-  --prompt "A sleek robot stands in a vast warehouse filled with boxes" \
-  --height 480 --width 832 --num-frames 81 \
-  --num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
-  --output t2v_output.mp4
-```
-
-## Image-to-Video (I2V)
-
-First frame is kept, remaining frames are generated:
-
-```bash
-python vace_video_generation.py \
-  --mode i2v \
-  --image astronaut.jpg \
-  --prompt "An astronaut emerging from a cracked egg on the moon" \
-  --height 480 --width 832 --num-frames 81 \
-  --output i2v_output.mp4
-```
-
-## First-Last-Frame Interpolation (FLF2V)
-
-```bash
-python vace_video_generation.py \
-  --mode flf2v \
-  --image first_frame.jpg --last-image last_frame.jpg \
-  --prompt "A bird takes off from a branch and lands on another" \
-  --height 512 --width 512 --num-frames 81 \
-  --output flf2v_output.mp4
-```
-
-## Inpainting
-
-Center vertical stripe is masked and regenerated:
-
-```bash
-python vace_video_generation.py \
-  --mode inpaint \
-  --image scene.jpg \
-  --prompt "Shrek walks out of a building" \
-  --height 480 --width 832 --num-frames 81 \
-  --output inpaint_output.mp4
-```
-
-## Reference Image-guided (R2V)
-
-```bash
-python vace_video_generation.py \
-  --mode r2v \
-  --image reference.jpg \
-  --prompt "Camera slowly zooms out from the character" \
-  --height 480 --width 832 --num-frames 81 \
-  --output r2v_output.mp4
-```
-
-## Key Arguments
-
- `--mode`: VACE task mode (`t2v`, `i2v`, `flf2v`, `inpaint`, `r2v`).
- `--model`: Model ID (default: `Wan-AI/Wan2.1-VACE-1.3B-diffusers`).
- `--image`: Input image for I2V, inpainting, and R2V modes.
- `--last-image`: Last frame image for FLF2V mode.
- `--prompt`: Text description of desired video.
- `--height/--width`: Output resolution (default 480x832). Dimensions should be multiples of 16.
- `--num-frames`: Number of frames (default 81).
- `--guidance-scale`: CFG scale (default 5.0).
- `--flow-shift`: Scheduler flow shift (default 5.0).
- `--num-inference-steps`: Number of denoising steps (default 30).
- `--fps`: Frames per second for the saved MP4 (default 16).
- `--output`: Path to save the generated video.
- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
- `--ulysses-degree`: Ulysses sequence parallelism degree for multi-GPU.
- `--cfg-parallel-size`: CFG parallel size for multi-GPU.
- `--tensor-parallel-size`: Tensor parallel size.
-
-> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options.
@@ -1,209 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""VACE video generation example.
-
-VACE (Video All-in-one Creation Engine) supports multiple video tasks:
-  - T2V:        Text-to-Video
-  - I2V:        Image-to-Video (first frame conditioning)
-  - V2LF:       Video-to-Last-Frame
-  - FLF2V:      First-Last-Frame interpolation
-  - Inpainting:  Masked region generation
-  - R2V:        Reference image-guided generation
-
-Usage examples:
-  # T2V (text-to-video)
-  python vace_video_generation.py --mode t2v --prompt "A robot in a warehouse"
-
-  # I2V (image-to-video, first frame kept)
-  python vace_video_generation.py --mode i2v --image input.jpg --prompt "..."
-
-  # FLF2V (first-last frame interpolation)
-  python vace_video_generation.py --mode flf2v --image first.jpg --last-image last.jpg
-
-  # R2V (reference image guided)
-  python vace_video_generation.py --mode r2v --image ref.jpg --prompt "..."
-"""
-
-import argparse
-import time
-from pathlib import Path
-
-import numpy as np
-import PIL.Image
-import torch
-
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.platforms import current_omni_platform
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="VACE video generation.")
-    parser.add_argument(
-        "--model",
-        default="Wan-AI/Wan2.1-VACE-14B-diffusers",
-        help="VACE model ID or local path.",
-    )
-    parser.add_argument(
-        "--mode",
-        default="t2v",
-        choices=["t2v", "i2v", "v2lf", "flf2v", "inpaint", "r2v"],
-        help="Generation mode.",
-    )
-    parser.add_argument("--prompt", default="A cat walking in a garden", help="Text prompt.")
-    parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
-    parser.add_argument("--image", type=str, default=None, help="Input image path (for I2V, R2V, FLF2V, inpaint).")
-    parser.add_argument("--last-image", type=str, default=None, help="Last frame image path (for FLF2V).")
-    parser.add_argument("--video-dir", type=str, default=None, help="Directory of video frames (for inpaint).")
-    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
-    parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG guidance scale.")
-    parser.add_argument("--height", type=int, default=480, help="Video height.")
-    parser.add_argument("--width", type=int, default=832, help="Video width.")
-    parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
-    parser.add_argument("--num-inference-steps", type=int, default=30, help="Sampling steps.")
-    parser.add_argument("--flow-shift", type=float, default=5.0, help="Scheduler flow_shift.")
-    parser.add_argument("--output", type=str, default="vace_output.mp4", help="Output video path.")
-    parser.add_argument("--fps", type=int, default=16, help="Output video FPS.")
-    parser.add_argument("--vae-use-tiling", action="store_true", default=True, help="Enable VAE tiling.")
-    parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
-    parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
-    parser.add_argument("--ring-degree", type=int, default=1, help="Ring attention degree.")
-    parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
-    return parser.parse_args()
-
-
-def build_prompts(args):
-    """Build prompt dict with multi_modal_data based on mode."""
-    h, w, nf = args.height, args.width, args.num_frames
-
-    gray = PIL.Image.new("RGB", (w, h), (128, 128, 128))
-    mask_black = PIL.Image.new("L", (w, h), 0)
-    mask_white = PIL.Image.new("L", (w, h), 255)
-
-    prompt_data = {
-        "prompt": args.prompt,
-        "negative_prompt": args.negative_prompt,
-    }
-
-    if args.mode == "t2v":
-        return prompt_data
-
-    if args.mode == "r2v":
-        assert args.image, "--image required for R2V mode"
-        ref_img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
-        prompt_data["multi_modal_data"] = {"reference_images": [ref_img]}
-        return prompt_data
-
-    if args.mode == "i2v":
-        assert args.image, "--image required for I2V mode"
-        img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
-        prompt_data["multi_modal_data"] = {
-            "video": [img] + [gray] * (nf - 1),
-            "mask": [mask_black] + [mask_white] * (nf - 1),
-        }
-        return prompt_data
-
-    if args.mode == "v2lf":
-        assert args.image, "--image required for V2LF mode"
-        img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
-        prompt_data["multi_modal_data"] = {
-            "video": [gray] * (nf - 1) + [img],
-            "mask": [mask_white] * (nf - 1) + [mask_black],
-        }
-        return prompt_data
-
-    if args.mode == "flf2v":
-        assert args.image and args.last_image, "--image and --last-image required for FLF2V"
-        first = PIL.Image.open(args.image).convert("RGB").resize((w, h))
-        last = PIL.Image.open(args.last_image).convert("RGB").resize((w, h))
-        prompt_data["multi_modal_data"] = {
-            "video": [first] + [gray] * (nf - 2) + [last],
-            "mask": [mask_black] + [mask_white] * (nf - 2) + [mask_black],
-        }
-        return prompt_data
-
-    if args.mode == "inpaint":
-        assert args.image, "--image required for inpaint mode"
-        img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
-        d = 80
-        frames, masks = [], []
-        for _ in range(nf):
-            base = np.array(img).copy()
-            mask = PIL.Image.new("L", (w, h), 0)
-            stripe = PIL.Image.new("L", (2 * d, h), 255)
-            mask.paste(stripe, (w // 2 - d, 0))
-            base[np.array(mask) > 128] = 128
-            frames.append(PIL.Image.fromarray(base))
-            masks.append(mask)
-        prompt_data["multi_modal_data"] = {"video": frames, "mask": masks}
-        return prompt_data
-
-    raise ValueError(f"Unknown mode: {args.mode}")
-
-
-def main():
-    args = parse_args()
-    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
-
-    parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree,
-        ring_degree=args.ring_degree,
-        cfg_parallel_size=args.cfg_parallel_size,
-    )
-
-    omni = Omni(
-        model=args.model,
-        vae_use_tiling=args.vae_use_tiling,
-        flow_shift=args.flow_shift,
-        enforce_eager=args.enforce_eager,
-        parallel_config=parallel_config,
-    )
-
-    prompt_data = build_prompts(args)
-
-    print(f"\n{'=' * 60}")
-    print(f"VACE {args.mode.upper()} Generation")
-    print(f"  Model: {args.model}")
-    print(f"  Size: {args.width}x{args.height}, {args.num_frames} frames, {args.num_inference_steps} steps")
-    print(f"{'=' * 60}\n")
-
-    start = time.perf_counter()
-    outputs = omni.generate(
-        prompt_data,
-        OmniDiffusionSamplingParams(
-            height=args.height,
-            width=args.width,
-            num_frames=args.num_frames,
-            num_inference_steps=args.num_inference_steps,
-            guidance_scale=args.guidance_scale,
-            generator=generator,
-        ),
-    )
-    elapsed = time.perf_counter() - start
-
-    video = outputs[0].images
-    if isinstance(video, list):
-        video = video[0]
-    if isinstance(video, torch.Tensor):
-        video = video.cpu().numpy()
-    if video.ndim == 5:
-        video = video[0]
-    print(f"Output shape: {video.shape}, Time: {elapsed:.1f}s")
-
-    output_path = Path(args.output)
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-
-    from diffusers.utils import export_to_video
-
-    if np.issubdtype(video.dtype, np.integer):
-        video = video.astype(np.float32) / 255.0
-    export_to_video(list(video), str(output_path), fps=args.fps)
-    print(f"Saved to {output_path}")
-
-    omni.close()
-
-
-if __name__ == "__main__":
-    main()
@@ -1,131 +0,0 @@
-# OmniVoice
-
-## Model Overview
-
-| Model | Description |
-|-------|-------------|
-| `k2-fsa/OmniVoice` | Zero-shot multilingual TTS (600+ languages) with diffusion language model (Qwen3-0.6B backbone) |
-
-> **Note:** Requires `transformers>=5.3.0` for voice cloning (HiggsAudioV2 tokenizer). Auto voice and voice design work with `transformers>=4.57.0`.
-
-## Launch the Server
-
-```bash
-vllm serve k2-fsa/OmniVoice \
-    --omni \
-    --port 8091 \
-    --trust-remote-code
-```
-
-Or use the convenience script:
-
-```bash
-./run_server.sh
-```
-
-## Send TTS Request
-
-### Using curl
-
-```bash
-# Basic TTS (auto voice)
-curl -X POST http://localhost:8091/v1/audio/speech \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "Hello, how are you?",
-        "voice": "default",
-        "response_format": "wav"
-    }' --output output.wav
-```
-
-### Using Python
-
-```python
-import httpx
-
-response = httpx.post(
-    "http://localhost:8091/v1/audio/speech",
-    json={
-        "input": "Hello, how are you?",
-        "voice": "default",
-        "response_format": "wav",
-    },
-    timeout=300.0,
-)
-
-with open("output.wav", "wb") as f:
-    f.write(response.content)
-```
-
-### Using OpenAI SDK
-
-```python
-from openai import OpenAI
-
-client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
-
-response = client.audio.speech.create(
-    model="k2-fsa/OmniVoice",
-    voice="default",
-    input="Hello, how are you?",
-)
-
-response.stream_to_file("output.wav")
-```
-
-### Using the CLI Client
-
-```bash
-cd examples/online_serving/omnivoice
-
-# Basic TTS
-python speech_client.py --text "Hello, how are you?"
-
-# Specify language for improved quality
-python speech_client.py --text "Bonjour, comment allez-vous?" --language French
-```
-
-The CLI client supports:
-
- `--api-base`: API base URL (default: `http://localhost:8091`)
- `--model` (or `-m`): Model name (default: `k2-fsa/OmniVoice`)
- `--text`: Text to synthesize (required)
- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav)
- `--language`: Language hint (default: Auto)
- `--output` (or `-o`): Output file path (default: `omnivoice_output.wav`)
-
-## Inference Modes
-
-OmniVoice supports three inference modes. Currently, **auto voice** is supported through the online Speech API. Voice cloning and voice design are available via offline inference (see `examples/offline_inference/omnivoice/`).
-
-| Mode | Description | Online API | Offline |
-|------|-------------|:----------:|:-------:|
-| Auto Voice | Generate speech without reference | Yes | Yes |
-| Voice Clone | Clone from reference audio + transcript | - | Yes |
-| Voice Design | Control style via natural language instruction | - | Yes |
-
-## Architecture
-
-OmniVoice uses a single-stage diffusion pipeline:
-
- **Stage 0 (Generator)**: Qwen3-0.6B transformer with 32-step iterative masked unmasking and classifier-free guidance. Generates 8-codebook audio tokens from text, then decodes to 24kHz waveform via HiggsAudioV2 RVQ quantizer + DAC acoustic decoder.
-
-## API Parameters
-
-OmniVoice uses the standard `/v1/audio/speech` endpoint. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full documentation.
-
-Key parameters:
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `input` | string | **required** | Text to synthesize |
-| `voice` | string | "default" | Voice name |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
-| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
-
-## Troubleshooting
-
-1. **TTS model did not produce audio output**: Ensure the model is fully downloaded (`huggingface-cli download k2-fsa/OmniVoice`)
-2. **Connection refused**: Make sure the server is running on the correct port
-3. **Out of memory**: Reduce `--gpu-memory-utilization` (default stage config uses 0.5)
-4. **Slow first request**: The model performs warmup on first inference; subsequent requests are faster
@@ -1,19 +0,0 @@
-#!/bin/bash
-# Launch vLLM-Omni server for OmniVoice TTS
-#
-# Usage:
-#   ./run_server.sh
-#   CUDA_VISIBLE_DEVICES=0 ./run_server.sh
-
-set -e
-
-MODEL="${MODEL:-k2-fsa/OmniVoice}"
-PORT="${PORT:-8091}"
-
-echo "Starting OmniVoice server with model: $MODEL"
-
-vllm serve "$MODEL" \
-    --host 0.0.0.0 \
-    --port "$PORT" \
-    --trust-remote-code \
-    --omni
@@ -1,84 +0,0 @@
-"""Client for OmniVoice TTS via /v1/audio/speech endpoint.
-
-Examples:
-    # Basic TTS (auto voice)
-    python speech_client.py --text "Hello, how are you?"
-
-    # Specify language
-    python speech_client.py --text "Bonjour, comment allez-vous?" --language French
-"""
-
-import argparse
-
-import httpx
-
-DEFAULT_API_BASE = "http://localhost:8091"
-DEFAULT_API_KEY = "EMPTY"
-
-
-def run_tts(args) -> None:
-    """Generate speech via /v1/audio/speech API."""
-    payload = {
-        "model": args.model,
-        "input": args.text,
-        "voice": "default",
-        "response_format": args.response_format,
-    }
-
-    if args.language:
-        payload["language"] = args.language
-
-    print(f"Model: {args.model}")
-    print(f"Text: {args.text}")
-    if args.language:
-        print(f"Language: {args.language}")
-    print("Generating audio...")
-
-    api_url = f"{args.api_base}/v1/audio/speech"
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {args.api_key}",
-    }
-
-    with httpx.Client(timeout=300.0) as client:
-        response = client.post(api_url, json=payload, headers=headers)
-
-    if response.status_code != 200:
-        print(f"Error: {response.status_code}")
-        print(response.text)
-        return
-
-    try:
-        text = response.content.decode("utf-8")
-        if text.startswith('{"error"'):
-            print(f"Error: {text}")
-            return
-    except UnicodeDecodeError:
-        pass
-
-    output_path = args.output or "omnivoice_output.wav"
-    with open(output_path, "wb") as f:
-        f.write(response.content)
-    print(f"Audio saved to: {output_path}")
-
-
-def main():
-    parser = argparse.ArgumentParser(description="OmniVoice TTS client")
-    parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
-    parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
-    parser.add_argument("--model", "-m", default="k2-fsa/OmniVoice", help="Model name")
-    parser.add_argument("--text", required=True, help="Text to synthesize")
-    parser.add_argument("--language", default=None, help="Language hint (e.g., English, Chinese, French)")
-    parser.add_argument(
-        "--response-format",
-        default="wav",
-        choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
-        help="Audio format (default: wav)",
-    )
-    parser.add_argument("--output", "-o", default=None, help="Output file path")
-    args = parser.parse_args()
-    run_tts(args)
-
-
-if __name__ == "__main__":
-    main()
@@ -36,45 +36,6 @@ cd examples/online_serving/qwen3_omni
 python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --query-type use_image --port 8091 --host "localhost"
 ```

-#### Realtime WebSocket client (`openai_realtime_client.py`)
-
-[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://<host>:<port>/v1/realtime`**, uploads a local audio file as **PCM16 mono @ 16 kHz** chunks (OpenAI-style `input_audio_buffer.append` / `commit`), and prints **streaming transcription** (`transcription.delta` / `transcription.done`).
-
-**Dependencies:**
-
-```bash
-pip install websockets librosa numpy
-```
-
-(ffmpeg may be required by `librosa` for some formats; see the FAQ below.)
-
-**From this directory** (`examples/online_serving/qwen3_omni`):
-
-```bash
-python openai_realtime_client.py \
-  --host localhost \
-  --port 8091 \
-  --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
-  --audio_path /path/to/your.wav
-```
-
-If `--audio_path` is omitted, the script uses a bundled default clip (`mary_had_lamb` via vLLM assets).
-
-**Arguments:**
-
-| Flag | Default | Description |
-|------|---------|-------------|
-| `--host` | `localhost` | API server host |
-| `--port` | `8000` | API server port (match your `vllm serve` port, e.g. `8091`) |
-| `--model` | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | Must match the served model (also sent in `session.update`) |
-| `--audio_path` | *(optional)* | Path to input audio; resampled to 16 kHz mono inside the client |
-
-Ensure the vLLM-Omni server is running with realtime support for this endpoint, for example:
-
-```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
-```
-
 The Python client supports the following command-line arguments:

 - `--query-type` (or `-q`): Query type (default: `use_video`). Options: `text`, `use_audio`, `use_image`, `use_video`
@@ -1,146 +0,0 @@
-"""
-This script demonstrates how to use the vLLM-Omni Realtime WebSocket API to perform
-audio transcription by uploading an audio file.
-
-Before running this script, you must start the vLLM-Omni server with a realtime-capable
-model, for example:
-
-    vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni
-
-Requirements:
- vllm with audio support
- websockets
- librosa
- numpy
-
-The script:
-1. Connects to the Realtime WebSocket endpoint
-2. Converts an audio file to PCM16 @ 16kHz
-3. Sends audio chunks to the server
-4. Receives and prints transcription as it streams
-"""
-
-import argparse
-import asyncio
-import base64
-import json
-
-import librosa
-import numpy as np
-import websockets
-from vllm.assets.audio import AudioAsset
-
-
-def audio_to_pcm16_base64(audio_path: str) -> str:
-    """
-    Load an audio file and convert it to base64-encoded PCM16 @ 16kHz.
-    """
-    # Load audio and resample to 16kHz mono
-    audio, _ = librosa.load(audio_path, sr=16000, mono=True)
-    # Convert to PCM16
-    pcm16 = (audio * 32767).astype(np.int16)
-    # Encode as base64
-    return base64.b64encode(pcm16.tobytes()).decode("utf-8")
-
-
-async def realtime_transcribe(audio_path: str, host: str, port: int, model: str):
-    """
-    Connect to the Realtime API and transcribe an audio file.
-    """
-    uri = f"ws://{host}:{port}/v1/realtime"
-
-    async with websockets.connect(uri) as ws:
-        # Wait for session.created
-        response = json.loads(await ws.recv())
-        if response["type"] == "session.created":
-            print(f"Session created: {response['id']}")
-        else:
-            print(f"Unexpected response: {response}")
-            return
-
-        # Validate model
-        await ws.send(json.dumps({"type": "session.update", "model": model}))
-
-        # Signal ready to start
-        await ws.send(json.dumps({"type": "input_audio_buffer.commit"}))
-
-        # Convert audio file to base64 PCM16
-        print(f"Loading audio from: {audio_path}")
-        audio_base64 = audio_to_pcm16_base64(audio_path)
-
-        # Send audio in chunks (4KB of raw audio = ~8KB base64)
-        chunk_size = 4096
-        audio_bytes = base64.b64decode(audio_base64)
-        total_chunks = (len(audio_bytes) + chunk_size - 1) // chunk_size
-
-        print(f"Sending {total_chunks} audio chunks...")
-        for i in range(0, len(audio_bytes), chunk_size):
-            chunk = audio_bytes[i : i + chunk_size]
-            await ws.send(
-                json.dumps(
-                    {
-                        "type": "input_audio_buffer.append",
-                        "audio": base64.b64encode(chunk).decode("utf-8"),
-                    }
-                )
-            )
-
-        # Signal all audio is sent
-        await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
-        print("Audio sent. Waiting for transcription...\n")
-
-        # Receive transcription
-        print("Transcription: ", end="", flush=True)
-        while True:
-            response = json.loads(await ws.recv())
-            if response["type"] == "transcription.delta":
-                print(response["delta"], end="", flush=True)
-            elif response["type"] == "transcription.done":
-                print(f"\n\nFinal transcription: {response['text']}")
-                if response.get("usage"):
-                    print(f"Usage: {response['usage']}")
-                break
-            elif response["type"] == "error":
-                print(f"\nError: {response['error']}")
-                break
-
-
-def main(args):
-    if args.audio_path:
-        audio_path = args.audio_path
-    else:
-        # Use default audio asset
-        audio_path = str(AudioAsset("mary_had_lamb").get_local_path())
-        print(f"No audio path provided, using default: {audio_path}")
-
-    asyncio.run(realtime_transcribe(audio_path, args.host, args.port, args.model))
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Realtime WebSocket Transcription Client")
-    parser.add_argument(
-        "--model",
-        type=str,
-        default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
-        help="Model that is served and should be pinged.",
-    )
-    parser.add_argument(
-        "--audio_path",
-        type=str,
-        default=None,
-        help="Path to the audio file to transcribe.",
-    )
-    parser.add_argument(
-        "--host",
-        type=str,
-        default="localhost",
-        help="vLLM-Omni server host (default: localhost)",
-    )
-    parser.add_argument(
-        "--port",
-        type=int,
-        default=8000,
-        help="vLLM-Omni server port (default: 8000)",
-    )
-    args = parser.parse_args()
-    main(args)
@@ -233,7 +233,6 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
 - `consent` (required): Consent recording ID
 - `name` (required): Name for the new voice
 - `ref_text` (optional): Transcript of the audio. Enables in-context voice cloning (higher quality).
- `speaker_description` (optional): Free-form description of the voice (e.g. "warm narrator", "energetic presenter").

 **Response Example:**
 ```json
@@ -244,23 +243,18 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
    "consent": "user_consent_id",
    "created_at": 1738660000,
    "mime_type": "audio/wav",
-    "file_size": 1024000,
-    "ref_text": "The exact transcript of the audio sample.",
-    "speaker_description": "warm narrator"
+    "file_size": 1024000
  }
 }
 ```

-Fields `ref_text` and `speaker_description` are omitted when not provided at upload time.
-
 **Usage Example:**
 ```bash
 curl -X POST http://localhost:8000/v1/audio/voices \
  -F "audio_sample=@/path/to/voice_sample.wav" \
  -F "consent=user_consent_id" \
  -F "name=custom_voice_1" \
-  -F "ref_text=The exact transcript of the audio sample." \
-  -F "speaker_description=warm narrator"
+  -F "ref_text=The exact transcript of the audio sample."
 ```

 ### Endpoint
@@ -223,4 +223,3 @@ extend-ignore-identifiers-re = [
 ue = "ue"
 semantics = "semantics"
 fullset = "fullset"
-Vai = "Vai"
@@ -1146,6 +1146,18 @@ def convert_audio_bytes_to_text(raw_bytes: bytes) -> str:
    return text


+def merge_base64_and_convert_to_text(base64_list):
+    """
+    Merge a list of base64 encoded audio data and convert to text.
+    """
+    merged_audio = _merge_base64_audio_to_segment(base64_list)
+    output_path = f"./test_{uuid.uuid4().hex}.wav"
+    merged_audio.export(output_path, format="wav")
+    print(f"audio data is saved: {output_path}")
+    text = convert_audio_file_to_text(output_path)
+    return text
+
+
 def modify_stage_config(
    yaml_path: str,
    updates: dict[str, Any] = None,
@@ -1730,7 +1742,7 @@ def _estimate_voice_gender_from_audio(audio_bytes: bytes) -> str:
        label = str(top.get("label", "")).lower()
        conf = float(top.get("score", 0.0))

-        if conf < 0.5:
+        if conf < 0.6:
            gender = "unknown"
        # Some models use non-English labels (e.g., Russian). Normalize to 'male'/'female'.
        elif ("female" in label) or ("жен" in label):
@@ -1759,34 +1771,6 @@ def _estimate_voice_gender_from_audio(audio_bytes: bytes) -> str:
        return "unknown"


-_PRESET_VOICE_GENDER_MAP: dict[str, str] = {
-    "serena": "female",
-    "uncle_fu": "male",
-    "chelsie": "female",
-    "clone": "female",
-    "ethan": "male",
-}
-
-
-def _assert_preset_voice_gender_from_audio(
-    audio_bytes: bytes | None,
-    voice_name: str | None,
-) -> None:
-    """If ``voice_name`` matches a known preset, assert classifier gender matches (skip when unknown)."""
-    if not voice_name or not audio_bytes:
-        return
-    key = str(voice_name).lower()
-    expected_gender = _PRESET_VOICE_GENDER_MAP.get(key)
-    if expected_gender is None:
-        return
-    estimated_gender = _estimate_voice_gender_from_audio(audio_bytes)
-    print(f"Preset voice gender check: preset={key!r}, estimated={estimated_gender!r}, expected={expected_gender!r}")
-    if estimated_gender != "unknown":
-        assert estimated_gender == expected_gender, (
-            f"{voice_name!r} is expected {expected_gender}, but estimated gender is {estimated_gender!r}"
-        )
-
-
 # Threshold aligned with _compute_pcm_hnr_db docstring (clean clone vs distorted).
 _MIN_PCM_SPEECH_HNR_DB = 1.0

@@ -1853,12 +1837,6 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any],
        if "audio" in modalities:
            assert response.audio_content is not None, "No audio output is generated"
            print(f"audio content is: {response.audio_content}")
-            speaker = request_config.get("speaker")
-            if speaker:
-                _assert_preset_voice_gender_from_audio(
-                    response.audio_bytes,
-                    speaker,
-                )

        if "text" in modalities:
            assert response.text_content is not None, "No text output is generated"
@@ -1871,14 +1849,12 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any],
            keywords = keywords_dict.get(word_type)
            if "text" in modalities:
                if keywords:
-                    text_lower = response.text_content.lower()
-                    assert any(str(kw).lower() in text_lower for kw in keywords), (
+                    assert any(keyword in response.text_content.lower() for keyword in keywords), (
                        "The output does not contain any of the keywords."
                    )
            else:
                if keywords:
-                    audio_lower = response.audio_content.lower()
-                    assert any(str(kw).lower() in audio_lower for kw in keywords), (
+                    assert any(keyword in response.audio_content.lower() for keyword in keywords), (
                        "The output does not contain any of the keywords."
                    )

@@ -1932,12 +1908,24 @@ def assert_audio_speech_response(
                f"Transcript doesn't match input: similarity={similarity:.2f}, transcript='{transcript}'"
            )

-        # Voice gender consistency check (preset names in ``_PRESET_VOICE_GENDER_MAP``).
+        # Voice gender consistency check:
        # When the estimator returns 'unknown', we treat it as inconclusive and do NOT fail the test.
-        _assert_preset_voice_gender_from_audio(
-            response.audio_bytes,
-            request_config.get("voice"),
-        )
+        voice = (request_config.get("voice") or "").lower()
+        if voice and response.audio_bytes:
+            estimated_gender = _estimate_voice_gender_from_audio(response.audio_bytes)
+            voice_gender_map = {
+                # adjust this mapping to your actual voice names
+                "serena": "female",
+                "uncle_fu": "male",
+                "clone": "female",
+            }
+            expected_gender = voice_gender_map.get(voice)
+            if expected_gender is not None:
+                print(f"Estimated voice gender from audio: {estimated_gender} (voice='{voice}')")
+                if estimated_gender != "unknown":
+                    assert estimated_gender == expected_gender, (
+                        f"Voice '{voice}' is expected {expected_gender}, but estimated gender is '{estimated_gender}'"
+                    )


 def assert_diffusion_response(response: DiffusionResponse, request_config: dict[str, Any], run_level: str = None):
@@ -2053,11 +2041,7 @@ class OpenAIClientHandler:

            if audio_data or text_content:
                if audio_data:
-                    merged_seg = _merge_base64_audio_to_segment(audio_data)
-                    wav_buf = BytesIO()
-                    merged_seg.export(wav_buf, format="wav")
-                    result.audio_bytes = wav_buf.getvalue()
-                    audio_content = convert_audio_bytes_to_text(result.audio_bytes)
+                    audio_content = merge_base64_and_convert_to_text(audio_data)
                if audio_content and text_content:
                    similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())

@@ -2112,8 +2096,7 @@ class OpenAIClientHandler:

            if audio_data or text_content:
                if audio_data:
-                    result.audio_bytes = base64.b64decode(audio_data)
-                    audio_content = convert_audio_bytes_to_text(result.audio_bytes)
+                    audio_content = convert_audio_to_text(audio_data)
                if audio_content and text_content:
                    similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())

@@ -2282,9 +2265,8 @@ class OpenAIClientHandler:
            request_config: Request configuration dictionary containing parameters like model, messages, stream.
                Optional ``use_audio_in_video`` (bool): when true, sets
                ``extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}`` for Qwen-Omni video+audio
-                extraction.
-                Optional top-level ``speaker`` (str): Qwen3-Omni preset TTS speaker name; sent as
-                ``extra_body["speaker"]`` to ``chat.completions.create``.
+                extraction (merged with any existing ``extra_body`` / ``mm_processor_kwargs``).
+                Optional ``extra_body`` (dict): passed through to ``chat.completions.create`` after merge.
            request_num: Number of requests, defaults to 1 (single request)

        Returns:
@@ -2296,8 +2278,9 @@ class OpenAIClientHandler:
        modalities = request_config.get("modalities", ["text", "audio"])

        extra_body: dict[str, Any] = {}
-        if "speaker" in request_config:
-            extra_body["speaker"] = request_config["speaker"]
+        raw_extra = request_config.get("extra_body")
+        if raw_extra:
+            extra_body.update(raw_extra)
        if request_config.get("use_audio_in_video"):
            mm = dict(extra_body.get("mm_processor_kwargs") or {})
            mm["use_audio_in_video"] = True
@@ -2329,15 +2312,12 @@ class OpenAIClientHandler:
            # Send concurrent requests: run create + process in worker so e2e_latency includes full round-trip.
            def _one_omni_request():
                start = time.perf_counter()
-                worker_kwargs: dict[str, Any] = {
-                    "model": request_config.get("model"),
-                    "messages": request_config.get("messages"),
-                    "modalities": modalities,
-                    "stream": stream,
-                }
-                if extra_body_arg is not None:
-                    worker_kwargs["extra_body"] = extra_body_arg
-                chat_completion = self.client.chat.completions.create(**worker_kwargs)
+                chat_completion = self.client.chat.completions.create(
+                    model=request_config.get("model"),
+                    messages=request_config.get("messages"),
+                    modalities=modalities,
+                    stream=stream,
+                )
                if stream:
                    response = self._process_stream_omni_response(chat_completion)
                else:
@@ -1,237 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for shared AdaLayerNorm layers used by FLUX and other models."""
-
-import os
-
-import pytest
-import torch
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-@pytest.fixture(autouse=True)
-def _init_distributed():
-    """Initialize the minimal distributed environment required by
-    ReplicatedLinear (tensor-parallel group must exist)."""
-    from vllm.distributed.parallel_state import (
-        cleanup_dist_env_and_memory,
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-
-    os.environ.setdefault("MASTER_ADDR", "localhost")
-    os.environ.setdefault("MASTER_PORT", "29501")
-    init_distributed_environment(
-        world_size=1,
-        rank=0,
-        local_rank=0,
-        distributed_init_method="env://",
-    )
-    initialize_model_parallel()
-    yield
-    cleanup_dist_env_and_memory()
-
-
-@pytest.fixture(autouse=True)
-def _force_default_gemm(monkeypatch):
-    """Force CPU-compatible GEMM dispatch for tests using CPU tensors.
-
-    vLLM's dispatch_unquantized_gemm() selects the backend by platform
-    (e.g. rocm_unquantized_gemm on AMD machines), not by tensor device.
-    CPU test tensors crash with NotImplementedError on ROCm.  Monkeypatch
-    the dispatcher to always return the default (torch.nn.functional.linear)
-    implementation which works on any device."""
-    from vllm.model_executor.layers.utils import default_unquantized_gemm
-
-    monkeypatch.setattr(
-        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
-    )
-
-
-def test_adalayernorm_import_from_shared_module():
-    """Verify imports work from the shared adalayernorm module."""
-    from vllm_omni.diffusion.layers.adalayernorm import (  # noqa: F401
-        AdaLayerNormContinuous,
-        AdaLayerNormZero,
-        AdaLayerNormZeroSingle,
-    )
-
-
-def test_adalayernorm_zero_forward_shape():
-    """AdaLayerNormZero produces correct output shapes (x, gate, shift, scale, gate)."""
-    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZero
-
-    dim = 64
-    batch = 2
-    seq_len = 4
-    norm = AdaLayerNormZero(dim)
-
-    x = torch.randn(batch, seq_len, dim)
-    emb = torch.randn(batch, dim)
-
-    out_x, gate_msa, shift_mlp, scale_mlp, gate_mlp = norm(x, emb)
-
-    assert out_x.shape == (batch, seq_len, dim)
-    assert gate_msa.shape == (batch, dim)
-    assert shift_mlp.shape == (batch, dim)
-    assert scale_mlp.shape == (batch, dim)
-    assert gate_mlp.shape == (batch, dim)
-
-
-def test_adalayernorm_zero_single_forward_shape():
-    """AdaLayerNormZeroSingle produces (x, gate) with correct shapes."""
-    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZeroSingle
-
-    dim = 64
-    batch = 2
-    seq_len = 4
-    norm = AdaLayerNormZeroSingle(dim)
-
-    x = torch.randn(batch, seq_len, dim)
-    emb = torch.randn(batch, dim)
-
-    out_x, gate = norm(x, emb)
-
-    assert out_x.shape == (batch, seq_len, dim)
-    assert gate.shape == (batch, dim)
-
-
-def test_adalayernorm_continuous_forward_shape():
-    """AdaLayerNormContinuous produces correct output shape."""
-    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormContinuous
-
-    dim = 64
-    cond_dim = 64
-    batch = 2
-    seq_len = 4
-    norm = AdaLayerNormContinuous(dim, cond_dim)
-
-    x = torch.randn(batch, seq_len, dim)
-    conditioning = torch.randn(batch, cond_dim)
-
-    out = norm(x, conditioning)
-
-    assert out.shape == (batch, seq_len, dim)
-
-
-def test_adalayernorm_zero_accepts_quant_config():
-    """Constructor accepts quant_config=None and prefix='test' without error."""
-    from vllm_omni.diffusion.layers.adalayernorm import (
-        AdaLayerNormContinuous,
-        AdaLayerNormZero,
-        AdaLayerNormZeroSingle,
-    )
-
-    # Should not raise with quant_config=None and prefix
-    AdaLayerNormZero(64, quant_config=None, prefix="test.norm1")
-    AdaLayerNormZeroSingle(64, quant_config=None, prefix="test.norm")
-    AdaLayerNormContinuous(64, 64, quant_config=None, prefix="test.norm_out")
-
-
-def test_adalayernorm_uses_replicated_linear():
-    """Verify .linear is a ReplicatedLinear instance (not nn.Linear)."""
-    from vllm.model_executor.layers.linear import ReplicatedLinear
-
-    from vllm_omni.diffusion.layers.adalayernorm import (
-        AdaLayerNormContinuous,
-        AdaLayerNormZero,
-        AdaLayerNormZeroSingle,
-    )
-
-    norm_zero = AdaLayerNormZero(64)
-    assert isinstance(norm_zero.linear, ReplicatedLinear)
-
-    norm_zero_single = AdaLayerNormZeroSingle(64)
-    assert isinstance(norm_zero_single.linear, ReplicatedLinear)
-
-    norm_continuous = AdaLayerNormContinuous(64, 64)
-    assert isinstance(norm_continuous.linear, ReplicatedLinear)
-
-
-# ── Numerical equivalence tests against diffusers originals ──
-
-
-def _copy_weights(src_linear, dst_replicated_linear):
-    """Copy weights from nn.Linear to ReplicatedLinear for comparison."""
-    dst_replicated_linear.weight.data.copy_(src_linear.weight.data)
-    if src_linear.bias is not None and dst_replicated_linear.bias is not None:
-        dst_replicated_linear.bias.data.copy_(src_linear.bias.data)
-
-
-def test_adalayernorm_zero_matches_diffusers():
-    """Verify AdaLayerNormZero produces identical output to diffusers original."""
-    from diffusers.models.normalization import (
-        AdaLayerNormZero as DiffusersAdaLayerNormZero,
-    )
-
-    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZero
-
-    dim = 64
-    torch.manual_seed(42)
-    ours = AdaLayerNormZero(dim)
-    ref = DiffusersAdaLayerNormZero(dim)
-
-    # Copy weights: nn.Linear -> ReplicatedLinear
-    _copy_weights(ref.linear, ours.linear)
-
-    x = torch.randn(2, 4, dim)
-    emb = torch.randn(2, dim)
-
-    out_ours = ours(x, emb)
-    out_ref = ref(x, emb=emb)
-
-    for o, r in zip(out_ours, out_ref):
-        torch.testing.assert_close(o, r, atol=1e-5, rtol=1e-5)
-
-
-def test_adalayernorm_zero_single_matches_diffusers():
-    """Verify AdaLayerNormZeroSingle produces identical output to diffusers original."""
-    from diffusers.models.normalization import (
-        AdaLayerNormZeroSingle as DiffusersAdaLayerNormZeroSingle,
-    )
-
-    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZeroSingle
-
-    dim = 64
-    torch.manual_seed(42)
-    ours = AdaLayerNormZeroSingle(dim)
-    ref = DiffusersAdaLayerNormZeroSingle(dim)
-
-    _copy_weights(ref.linear, ours.linear)
-
-    x = torch.randn(2, 4, dim)
-    emb = torch.randn(2, dim)
-
-    out_ours = ours(x, emb)
-    out_ref = ref(x, emb=emb)
-
-    for o, r in zip(out_ours, out_ref):
-        torch.testing.assert_close(o, r, atol=1e-5, rtol=1e-5)
-
-
-def test_adalayernorm_continuous_matches_diffusers():
-    """Verify AdaLayerNormContinuous produces identical output to diffusers original."""
-    from diffusers.models.normalization import (
-        AdaLayerNormContinuous as DiffusersAdaLayerNormContinuous,
-    )
-
-    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormContinuous
-
-    dim = 64
-    cond_dim = 64
-    torch.manual_seed(42)
-    # Match constructor args: diffusers defaults elementwise_affine=True, eps=1e-5
-    ours = AdaLayerNormContinuous(dim, cond_dim, elementwise_affine=False, eps=1e-6)
-    ref = DiffusersAdaLayerNormContinuous(dim, cond_dim, elementwise_affine=False, eps=1e-6)
-
-    _copy_weights(ref.linear, ours.linear)
-
-    x = torch.randn(2, 4, dim)
-    cond = torch.randn(2, cond_dim)
-
-    out_ours = ours(x, cond)
-    out_ref = ref(x, cond)
-
-    torch.testing.assert_close(out_ours, out_ref, atol=1e-5, rtol=1e-5)
@@ -1,134 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests that FLUX transformer blocks correctly propagate `quant_config` and
-`prefix` through all sub-layers.
-
-The tests instantiate blocks with a known prefix and verify that all quantization-
-aware sub-layers (AdaLayerNorm, FeedForward, Attention projections) receive the
-prefix rooted at the block prefix.  This is critical for quantized weight loading
-to match checkpoint keys to the correct model parameters.
-"""
-
-import os
-
-import pytest
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-# Standard dimensions for a minimal FLUX block
-_DIM = 64
-_HEADS = 2
-_HEAD_DIM = _DIM // _HEADS
-
-
-@pytest.fixture(autouse=True)
-def _init_distributed():
-    """Initialize the minimal distributed environment required by
-    vLLM parallel linear layers (tensor-parallel group must exist)."""
-    from vllm.distributed.parallel_state import (
-        cleanup_dist_env_and_memory,
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-
-    os.environ.setdefault("MASTER_ADDR", "localhost")
-    os.environ.setdefault("MASTER_PORT", "29502")
-    init_distributed_environment(
-        world_size=1,
-        rank=0,
-        local_rank=0,
-        distributed_init_method="env://",
-    )
-    initialize_model_parallel()
-    yield
-    cleanup_dist_env_and_memory()
-
-
-def _param_names(module) -> set[str]:
-    """Return the set of all parameter names in a module."""
-    return {name for name, _ in module.named_parameters()}
-
-
-def test_flux_transformer_block_passes_prefix():
-    """FluxTransformerBlock propagates prefix to norm1, norm1_context, attn, ff, ff_context."""
-    from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformerBlock
-
-    prefix = "transformer_blocks.0"
-    block = FluxTransformerBlock(
-        dim=_DIM,
-        num_attention_heads=_HEADS,
-        attention_head_dim=_HEAD_DIM,
-        quant_config=None,
-        prefix=prefix,
-    )
-
-    params = _param_names(block)
-
-    # norm1 and norm1_context (AdaLayerNormZero) should have linear weights
-    assert any(name.startswith("norm1.linear.") for name in params), (
-        f"norm1.linear.* not found in params: {sorted(params)}"
-    )
-    assert any(name.startswith("norm1_context.linear.") for name in params), (
-        f"norm1_context.linear.* not found in params: {sorted(params)}"
-    )
-
-    # attn should have QKV projections
-    assert any(name.startswith("attn.to_qkv.") for name in params), (
-        f"attn.to_qkv.* not found in params: {sorted(params)}"
-    )
-
-    # ff and ff_context should have net layers
-    assert any(name.startswith("ff.net.") for name in params), f"ff.net.* not found in params: {sorted(params)}"
-    assert any(name.startswith("ff_context.net.") for name in params), (
-        f"ff_context.net.* not found in params: {sorted(params)}"
-    )
-
-
-def test_flux_single_transformer_block_passes_prefix():
-    """FluxSingleTransformerBlock propagates prefix to norm, proj_mlp, attn."""
-    from vllm_omni.diffusion.models.flux.flux_transformer import FluxSingleTransformerBlock
-
-    prefix = "single_transformer_blocks.0"
-    block = FluxSingleTransformerBlock(
-        dim=_DIM,
-        num_attention_heads=_HEADS,
-        attention_head_dim=_HEAD_DIM,
-        quant_config=None,
-        prefix=prefix,
-    )
-
-    params = _param_names(block)
-
-    # norm (AdaLayerNormZeroSingle) should have linear weights
-    assert any(name.startswith("norm.linear.") for name in params), (
-        f"norm.linear.* not found in params: {sorted(params)}"
-    )
-
-    # proj_mlp (ReplicatedLinear) should have weight
-    assert any(name.startswith("proj_mlp.") for name in params), f"proj_mlp.* not found in params: {sorted(params)}"
-
-    # attn should have QKV projection
-    assert any(name.startswith("attn.to_qkv.") for name in params), (
-        f"attn.to_qkv.* not found in params: {sorted(params)}"
-    )
-
-
-def test_flux_feedforward_passes_prefix():
-    """FeedForward propagates prefix to net.0 (GELU proj) and net.2 (output proj)."""
-    from vllm_omni.diffusion.models.flux.flux_transformer import FeedForward
-
-    prefix = "transformer_blocks.0.ff"
-    ff = FeedForward(
-        dim=_DIM,
-        dim_out=_DIM,
-        quant_config=None,
-        prefix=prefix,
-    )
-
-    params = _param_names(ff)
-
-    # net.0 is ColumnParallelApproxGELU which wraps a ColumnParallelLinear
-    assert any("net.0" in name for name in params), f"net.0 not found in params: {sorted(params)}"
-
-    # net.2 is RowParallelLinear
-    assert any("net.2" in name for name in params), f"net.2 not found in params: {sorted(params)}"
@@ -1,147 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for INC/AutoRound quantization via the unified framework."""
-
-import pytest
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
-
-
-def test_build_quant_config_autoround():
-    """build_quant_config("auto-round", ...) should produce an INCConfig."""
-    from vllm.model_executor.layers.quantization.inc import INCConfig
-
-    from vllm_omni.quantization import build_quant_config
-
-    config = build_quant_config(
-        "auto-round",
-        bits=4,
-        group_size=128,
-        sym=True,
-        packing_format="auto_round:auto_gptq",
-    )
-    assert config is not None
-    assert isinstance(config, INCConfig)
-    assert config.weight_bits == 4
-    assert config.group_size == 128
-
-
-def test_build_quant_config_inc():
-    """build_quant_config("inc", ...) should also produce an INCConfig."""
-    from vllm.model_executor.layers.quantization.inc import INCConfig
-
-    from vllm_omni.quantization import build_quant_config
-
-    config = build_quant_config("inc", bits=4, group_size=128)
-    assert isinstance(config, INCConfig)
-    assert config.weight_bits == 4
-
-
-def test_build_quant_config_autoround_dict():
-    """Dict-style config with method=auto-round should work."""
-    from vllm.model_executor.layers.quantization.inc import INCConfig
-
-    from vllm_omni.quantization import build_quant_config
-
-    config = build_quant_config(
-        {
-            "method": "auto-round",
-            "bits": 4,
-            "group_size": 128,
-            "sym": True,
-            "packing_format": "auto_round:auto_gptq",
-        }
-    )
-    assert isinstance(config, INCConfig)
-    assert config.weight_bits == 4
-
-
-def test_build_quant_config_autoround_filters_metadata():
-    """Checkpoint metadata keys (autoround_version, batch_size, iters)
-    should be silently filtered out instead of causing TypeError."""
-    from vllm.model_executor.layers.quantization.inc import INCConfig
-
-    from vllm_omni.quantization import build_quant_config
-
-    config = build_quant_config(
-        "auto-round",
-        bits=4,
-        group_size=128,
-        sym=True,
-        packing_format="auto_round:auto_gptq",
-        block_name_to_quantize="transformer_blocks,single_transformer_blocks",
-        autoround_version="0.12.0",  # metadata — must be filtered
-        batch_size=1,  # metadata — must be filtered
-        iters=0,  # metadata — must be filtered
-    )
-    assert isinstance(config, INCConfig)
-    assert config.weight_bits == 4
-    assert config.group_size == 128
-
-
-def test_build_quant_config_bits_to_weight_bits_mapping():
-    """The 'bits' key from checkpoints should be mapped to 'weight_bits'."""
-    from vllm.model_executor.layers.quantization.inc import INCConfig
-
-    from vllm_omni.quantization import build_quant_config
-
-    # If weight_bits is already provided, bits should be ignored
-    config = build_quant_config("auto-round", weight_bits=4, group_size=128)
-    assert isinstance(config, INCConfig)
-    assert config.weight_bits == 4
-
-
-def test_autoround_in_supported_methods():
-    """auto-round and inc should appear in SUPPORTED_QUANTIZATION_METHODS."""
-    from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
-
-    assert "auto-round" in SUPPORTED_QUANTIZATION_METHODS
-    assert "inc" in SUPPORTED_QUANTIZATION_METHODS
-
-
-def test_integration_autoround_via_omni_diffusion_config():
-    """OmniDiffusionConfig with auto-round quantization dict should resolve."""
-    from vllm.model_executor.layers.quantization.inc import INCConfig
-
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
-
-    config = OmniDiffusionConfig(
-        model="test",
-        quantization_config={
-            "method": "auto-round",
-            "bits": 4,
-            "group_size": 128,
-            "sym": True,
-        },
-    )
-    assert isinstance(config.quantization_config, INCConfig)
-    assert config.quantization_config.weight_bits == 4
-
-
-def test_integration_autodetect_from_transformer_config():
-    """When TransformerConfig has quant_config, OmniDiffusionConfig should
-    auto-detect it even without explicit quantization_config."""
-    from vllm.model_executor.layers.quantization.inc import INCConfig
-
-    from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
-
-    tf_config = TransformerConfig.from_dict(
-        {
-            "quantization_config": {
-                "quant_method": "auto-round",
-                "bits": 4,
-                "group_size": 128,
-                "sym": True,
-                "packing_format": "auto_round:auto_gptq",
-                "autoround_version": "0.12.0",
-                "batch_size": 1,
-                "iters": 0,
-            }
-        }
-    )
-    assert tf_config.quant_method == "auto-round"
-    assert isinstance(tf_config.quant_config, INCConfig)
-
-    od_config = OmniDiffusionConfig(model="test", tf_model_config=tf_config)
-    assert isinstance(od_config.quantization_config, INCConfig)
-    assert od_config.quantization_config.weight_bits == 4
@@ -118,12 +118,8 @@ def _generate_image(omni, config: QualityTestConfig):

    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
    first = outputs[0]
-    if hasattr(first, "images") and first.images:
-        return first.images[0], peak_mem
-    inner = first.request_output
-    if inner is not None and hasattr(inner, "images") and inner.images:
-        return inner.images[0], peak_mem
-    raise ValueError("Could not extract image from output.")
+    req_out = first.request_output[0] if hasattr(first, "request_output") else first
+    return req_out.images[0], peak_mem


 def _generate_video(omni, config: QualityTestConfig):
@@ -229,24 +229,6 @@ def _wait_for_port(host: str, port: int, timeout: int = 30) -> bool:
    return False


-def _is_mooncake_master_available() -> bool:
-    """Check if mooncake_master binary is present and can actually execute."""
-    import shutil
-
-    binary = shutil.which("mooncake_master")
-    if binary is None:
-        return False
-    try:
-        result = subprocess.run(
-            [binary, "--help"],
-            capture_output=True,
-            timeout=5,
-        )
-        return result.returncode != 127
-    except (subprocess.TimeoutExpired, OSError):
-        return True
-
-
 def _cleanup_mooncake_processes(timeout_secs: int = 5) -> None:
    """Clean up any existing mooncake_master processes.

@@ -310,8 +292,6 @@ def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str:
@hardware_test(res={"cuda": "H100"})
 def test_bagel_text2img_mooncake_connector(run_level):
    """Test Bagel text2img with Mooncake connector for inter-stage communication."""
-    if not _is_mooncake_master_available():
-        pytest.skip("mooncake_master is not available or cannot execute (missing shared libraries like libibverbs)")
    MOONCAKE_HOST = "127.0.0.1"
    MOONCAKE_RPC_PORT = _find_free_port()
    MOONCAKE_HTTP_PORT = _find_free_port()
@@ -1,127 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""E2E tests for FLUX AutoRound W4A16 quantized inference.
-
-These tests require:
-  - A CUDA GPU
-  - The quantized model checkpoint (vllm-project-org/FLUX.1-dev-AutoRound-w4a16)
-"""
-
-import gc
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
-
-from tests.utils import DeviceMemoryMonitor, hardware_test
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.platforms import current_omni_platform
-
-# ruff: noqa: E402
-REPO_ROOT = Path(__file__).resolve().parents[2]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-from vllm_omni import Omni
-
-QUANTIZED_MODEL = "vllm-project-org/FLUX.1-dev-AutoRound-w4a16"
-BASELINE_MODEL = "black-forest-labs/FLUX.1-dev"
-
-# Allow overriding via environment for local testing
-import os as _os
-
-QUANTIZED_MODEL = _os.environ.get("FLUX_AUTOROUND_MODEL", QUANTIZED_MODEL)
-BASELINE_MODEL = _os.environ.get("FLUX_BASELINE_MODEL", BASELINE_MODEL)
-
-# Small resolution to keep GPU memory & time manageable
-HEIGHT = 256
-WIDTH = 256
-NUM_STEPS = 2  # minimal for smoke-test
-
-
-def _generate_image(model_name: str, **extra_kwargs) -> tuple[list, float]:
-    """Load a FLUX model, generate one image, return (images, peak_memory_mb)."""
-    gc.collect()
-    current_omni_platform.empty_cache()
-    device_index = current_omni_platform.current_device()
-    current_omni_platform.reset_peak_memory_stats()
-    monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
-    monitor.start()
-
-    m = Omni(model=model_name, enforce_eager=True, **extra_kwargs)
-
-    current_omni_platform.reset_peak_memory_stats()
-    outputs = m.generate(
-        "a photo of a cat sitting on a laptop keyboard",
-        OmniDiffusionSamplingParams(
-            height=HEIGHT,
-            width=WIDTH,
-            num_inference_steps=NUM_STEPS,
-            guidance_scale=0.0,
-            generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
-        ),
-    )
-
-    peak = monitor.peak_used_mb
-    monitor.stop()
-
-    first_output = outputs[0]
-    assert first_output.final_output_type == "image"
-    req_out = first_output.request_output
-    assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
-    images = req_out.images
-
-    del m
-    gc.collect()
-    current_omni_platform.empty_cache()
-
-    return images, peak
-
-
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "L4"})
-def test_flux_autoround_w4a16_generates_image():
-    """Load the W4A16 quantized FLUX model and verify it produces a valid image."""
-    images, _ = _generate_image(QUANTIZED_MODEL)
-
-    assert len(images) >= 1, "Expected at least one generated image"
-    img = images[0]
-    assert img.width == WIDTH, f"Expected width {WIDTH}, got {img.width}"
-    assert img.height == HEIGHT, f"Expected height {HEIGHT}, got {img.height}"
-
-    # Sanity: image should not be blank (all-zero)
-    import numpy as np
-
-    arr = np.array(img)
-    assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)"
-
-
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "L4"})
-def test_flux_autoround_w4a16_memory_savings():
-    """Compare peak GPU memory of quantized vs FP16 baseline.
-
-    The W4A16 model should use meaningfully less memory than the
-    BF16/FP16 baseline since weights are 4-bit instead of 16-bit.
-    """
-    quant_images, quant_peak = _generate_image(QUANTIZED_MODEL)
-    cleanup_dist_env_and_memory()
-    _, baseline_peak = _generate_image(BASELINE_MODEL)
-
-    print(f"Quantized (W4A16) peak memory: {quant_peak:.0f} MB")
-    print(f"Baseline (BF16) peak memory:   {baseline_peak:.0f} MB")
-    print(f"Savings:                        {baseline_peak - quant_peak:.0f} MB")
-
-    # W4A16 weights are 4x smaller than BF16/FP16.  FLUX.1-dev transformer
-    # is ~12 GB in BF16, so we expect ~9 GB savings on weights alone.
-    # Use a conservative threshold to account for activations and overhead.
-    min_savings_mb = 2000
-    assert quant_peak + min_savings_mb < baseline_peak, (
-        f"Quantized model ({quant_peak:.0f} MB) should use at least "
-        f"{min_savings_mb} MB less than baseline ({baseline_peak:.0f} MB)"
-    )
@@ -1,84 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-E2E offline tests for OmniVoice TTS model with text input and audio output.
-
-Uses GPUGenerationWorker for both stages (iterative unmasking + DAC decoder).
-"""
-
-import os
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
-
-from pathlib import Path
-
-import numpy as np
-import pytest
-
-from tests.utils import hardware_test
-
-MODEL = "k2-fsa/OmniVoice"
-
-
-def get_stage_config():
-    return str(
-        Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "omnivoice.yaml"
-    )
-
-
-@pytest.mark.advanced_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_omnivoice_text_to_audio() -> None:
-    """
-    Test OmniVoice text-to-audio generation via offline Omni runner.
-    Deploy Setting: omnivoice.yaml (enforce_eager=true)
-    Input Modal: text
-    Output Modal: audio
-    """
-    from vllm_omni.entrypoints.omni import Omni
-
-    omni = Omni(
-        model=MODEL,
-        stage_configs_path=get_stage_config(),
-        trust_remote_code=True,
-        log_stats=True,
-    )
-
-    try:
-        prompts = {"prompt": "Hello, this is a test for text to audio."}
-
-        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-        sampling_params_list = [OmniDiffusionSamplingParams()]
-
-        outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list))
-
-        assert len(outputs) > 0, "No outputs generated"
-
-        # Check final output has audio
-        final_output = outputs[-1]
-        ro = final_output.request_output
-        assert ro is not None, "No request_output"
-
-        mm = getattr(ro, "multimodal_output", None)
-        if not mm and ro.outputs:
-            mm = getattr(ro.outputs[0], "multimodal_output", None)
-
-        assert mm is not None, "No multimodal_output"
-        assert "audio" in mm, f"No 'audio' key in multimodal_output: {mm.keys()}"
-
-        audio = mm["audio"]
-        if isinstance(audio, np.ndarray):
-            audio_np = audio
-        else:
-            audio_np = audio.cpu().numpy().squeeze()
-
-        assert audio_np.size > 0, "Audio output is empty"
-        rms = np.sqrt(np.mean(audio_np**2))
-        assert rms > 0.01, f"Audio RMS too low ({rms:.4f}), likely silence"
-
-        print(f"Generated audio: {len(audio_np) / 24000:.2f}s, rms={rms:.4f}")
-    finally:
-        omni.close()
@@ -57,7 +57,7 @@ def get_question(prompt_type="mix"):
    return prompts.get(prompt_type, prompts["mix"])


-@pytest.mark.advanced_model
+@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 3})
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
@@ -88,7 +88,7 @@ def test_mix_to_audio(omni_runner, omni_runner_handler) -> None:
    omni_runner_handler.send_request(request_config)


-@pytest.mark.advanced_model
+@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 3})
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
@@ -56,7 +56,7 @@ def get_question(prompt_type="video"):
    return prompts.get(prompt_type, prompts["video"])


-@pytest.mark.advanced_model
+@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
@@ -509,7 +509,7 @@ def test_diffusion_batching_async_explicit_batch(model_name: str):
    all prompts in a single engine call and returns a single combined result.

    The list-prompt path routes through the orchestrator's
-    ``add_batch_request_async`` → ``AsyncOmni.generate_batch``
+    ``add_batch_request_async`` → ``AsyncOmniDiffusion.generate_batch``
    and yields ONE ``OmniRequestOutput`` with ALL images combined.
    """

@@ -1,124 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-E2E Online tests for CosyVoice3 TTS model with voice cloning.
-
-These tests verify the /v1/audio/speech endpoint works correctly with
-the CosyVoice3 model, which requires reference audio for voice cloning.
-"""
-
-import os
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
-
-from pathlib import Path
-
-import pytest
-
-from tests.conftest import OmniServerParams
-from tests.utils import hardware_test
-
-MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
-
-# Official CosyVoice zero-shot prompt audio and its transcript
-REF_AUDIO_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
-REF_TEXT = "希望你以后能够做的比我还好呦。"
-
-
-def get_stage_config(name: str = "cosyvoice3.yaml"):
-    """Get the stage config path from vllm_omni model_executor stage_configs."""
-    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
-
-
-def get_prompt(prompt_type="zh"):
-    prompts = {
-        "zh": "收到好友从远方寄来的生日礼物，那份意外的惊喜与深深的感动让我热泪盈眶。",
-        "en": "Hello, this is a voice cloning test with English text.",
-    }
-    return prompts.get(prompt_type, prompts["zh"])
-
-
-tts_server_params = [
-    pytest.param(
-        OmniServerParams(
-            model=MODEL,
-            stage_config_path=get_stage_config(),
-            server_args=["--trust-remote-code", "--disable-log-stats"],
-        ),
-        id="cosyvoice3",
-    )
-]
-
-
-@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)
-def test_voice_clone_zh_001(omni_server, openai_client) -> None:
-    """
-    Test voice cloning TTS with Chinese text via OpenAI API.
-    Deploy Setting: default yaml
-    Input Modal: text + ref_audio + ref_text
-    Output Modal: audio
-    Input Setting: stream=False
-    Datasets: single request
-    """
-    request_config = {
-        "model": omni_server.model,
-        "input": get_prompt("zh"),
-        "stream": False,
-        "response_format": "wav",
-        "ref_audio": REF_AUDIO_URL,
-        "ref_text": REF_TEXT,
-    }
-    openai_client.send_audio_speech_request(request_config)
-
-
-@pytest.mark.advanced_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", tts_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
-    Input Modal: text + ref_audio + ref_text
-    Output Modal: audio
-    Input Setting: stream=True
-    Datasets: single request
-    """
-    request_config = {
-        "model": omni_server.model,
-        "input": get_prompt("zh"),
-        "stream": True,
-        "response_format": "wav",
-        "ref_audio": REF_AUDIO_URL,
-        "ref_text": REF_TEXT,
-    }
-    openai_client.send_audio_speech_request(request_config)
-
-
-@pytest.mark.advanced_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_voice_clone_en_001(omni_server, openai_client) -> None:
-    """
-    Test voice cloning TTS with English text via OpenAI API.
-    Deploy Setting: default yaml
-    Input Modal: text + ref_audio + ref_text
-    Output Modal: audio
-    Input Setting: stream=False
-    Datasets: single request
-    """
-    request_config = {
-        "model": omni_server.model,
-        "input": get_prompt("en"),
-        "stream": False,
-        "response_format": "wav",
-        "ref_audio": REF_AUDIO_URL,
-        "ref_text": REF_TEXT,
-    }
-    openai_client.send_audio_speech_request(request_config)
@@ -63,27 +63,20 @@ def download_tokenizer():


 # CI stage config for H100 / MI325
-# Guard module-level setup so test collection doesn't fail in environments
-# where the model cache is read-only or models aren't available.
-try:
-    stage_configs = [get_chunk_config()]
-    tokenizer_path = download_tokenizer()
-    os.environ["MIMO_AUDIO_TOKENIZER_PATH"] = tokenizer_path
+stage_configs = [get_chunk_config()]
+tokenizer_path = download_tokenizer()
+os.environ["MIMO_AUDIO_TOKENIZER_PATH"] = tokenizer_path

-    test_params = [
-        OmniServerParams(
-            model=model,
-            stage_config_path=stage_config,
-            server_args=["--chat-template", CHAT_TEMPLATE_PATH],
-        )
-        for model in models
-        for stage_config in stage_configs
-    ]
-except Exception as exc:
-    pytest.skip(
-        f"MiMo-Audio online serving tests skipped: module setup failed ({type(exc).__name__}: {exc})",
-        allow_module_level=True,
+# Create parameter combinations for model and stage config
+test_params = [
+    OmniServerParams(
+        model=model,
+        stage_config_path=stage_config,
+        server_args=["--chat-template", CHAT_TEMPLATE_PATH],
    )
+    for model in models
+    for stage_config in stage_configs
+]


 def get_prompt(prompt_type="text_only"):
@@ -102,7 +95,7 @@ def get_max_batch_size(size_type="few"):
@pytest.mark.advanced_model
@pytest.mark.core_model
@pytest.mark.omni
-@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
    """
@@ -135,7 +128,7 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None:

@pytest.mark.advanced_model
@pytest.mark.omni
-@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_to_text_001(omni_server, openai_client) -> None:
    """
@@ -1,84 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-E2E Online tests for OmniVoice TTS model via /v1/audio/speech endpoint.
-
-Tests verify that the OmniVoice model generates valid audio when
-accessed through the standard OpenAI-compatible speech API.
-"""
-
-import os
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
-
-from pathlib import Path
-
-import httpx
-import pytest
-
-from tests.conftest import OmniServerParams
-from tests.utils import hardware_test
-
-MODEL = "k2-fsa/OmniVoice"
-
-STAGE_CONFIG = str(
-    Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "omnivoice.yaml"
-)
-EXTRA_ARGS = [
-    "--trust-remote-code",
-    "--disable-log-stats",
-]
-TEST_PARAMS = [
-    OmniServerParams(
-        model=MODEL,
-        stage_config_path=STAGE_CONFIG,
-        server_args=EXTRA_ARGS,
-    )
-]
-
-MIN_AUDIO_BYTES = 5000
-
-
-def make_speech_request(
-    host: str,
-    port: int,
-    text: str,
-    timeout: float = 180.0,
-) -> httpx.Response:
-    """Make a request to the /v1/audio/speech endpoint for OmniVoice."""
-    url = f"http://{host}:{port}/v1/audio/speech"
-    payload = {"input": text}
-
-    with httpx.Client(timeout=timeout) as client:
-        return client.post(url, json=payload)
-
-
-def verify_wav_audio(content: bytes) -> bool:
-    """Verify that content is valid WAV audio data."""
-    if len(content) < 44:
-        return False
-    return content[:4] == b"RIFF" and content[8:12] == b"WAVE"
-
-
-@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
-class TestOmniVoiceTTS:
-    """E2E tests for OmniVoice TTS model."""
-
-    @pytest.mark.core_model
-    @pytest.mark.omni
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
-    def test_speech_auto_voice(self, omni_server) -> None:
-        """Test auto voice TTS generation (text only, no reference audio)."""
-        response = make_speech_request(
-            host=omni_server.host,
-            port=omni_server.port,
-            text="Hello, this is a test of the OmniVoice text to speech system.",
-        )
-
-        assert response.status_code == 200, f"Request failed: {response.text}"
-        assert response.headers.get("content-type") == "audio/wav"
-        assert verify_wav_audio(response.content), "Response is not valid WAV audio"
-        assert len(response.content) > MIN_AUDIO_BYTES, (
-            f"Audio too small ({len(response.content)} bytes), expected > {MIN_AUDIO_BYTES}"
-        )
@@ -26,7 +26,7 @@ from tests.utils import hardware_test
 model = "Qwen/Qwen3-Omni-30B-A3B-Instruct"

 AUDIO_KEY = ["test"]
-IMAGE_KEY = ["square", "quadrate", "rectangle"]
+IMAGE_KEY = ["square", "quadrate"]
 VIDEO_KEY = ["sphere", "globe", "circle", "round", "ball"]


@@ -103,7 +103,6 @@ def get_prompt(prompt_type="text_only"):
        "text_audio": "What is in this audio? ",
        "text_audio_video": "First, what is in this audio? Then, what is in this video? ",
        "one_word": "What is the capital of UK? Answer in one word",
-        "text_chinese": "北京，中国的首都，是一座融合了长城等历史地点与现代建筑的国际化大都市，充满了独特的文化与活力。请重复这句话。",
    }
    return prompts.get(prompt_type, prompts["text_only"])

@@ -465,10 +464,20 @@ def test_audio_in_video_002(omni_server, openai_client) -> None:
        "messages": messages,
        "stream": True,
        "use_audio_in_video": True,
-        "key_words": {"video": VIDEO_KEY},
+        "key_words": {"video": VIDEO_KEY, "audio": AUDIO_KEY + ["beep", "electronic"]},
    }

-    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+    # Retry when assert_omni_response fails on key_words (see tests/conftest.py).
+    _keyword_assert_msg = "The output does not contain any of the keywords."
+    _max_retries = 3
+    for attempt in range(_max_retries):
+        try:
+            openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+            break
+        except AssertionError as e:
+            if _keyword_assert_msg not in str(e) or attempt == _max_retries - 1:
+                raise
+            print(f"Keyword assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")


@pytest.mark.advanced_model
@@ -505,120 +514,3 @@ def test_one_word_prompt_001(omni_server, openai_client) -> None:
            if _similarity_assert_msg not in str(e) or attempt == _max_retries - 1:
                raise
            print(f"Similarity assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")
-
-
-@pytest.mark.advanced_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_speaker_001(omni_server, openai_client) -> None:
-    """
-    Input Modal: text only (one-word answer constraint).
-    Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
-    Input Setting: stream=True
-    Datasets: single request
-    """
-    messages = dummy_messages_from_mix_data(
-        system_prompt=get_system_prompt(),
-        content_text=get_prompt("text"),
-    )
-
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "stream": True,
-        "speaker": "Chelsie",
-        "key_words": {"text": ["beijing"]},
-    }
-
-    openai_client.send_omni_request(request_config)
-
-
-@pytest.mark.advanced_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_speaker_002(omni_server, openai_client) -> None:
-    """
-    Input Modal: text only (one-word answer constraint).
-    Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
-    Input Setting: stream=True
-    Datasets: single request
-    """
-    messages = dummy_messages_from_mix_data(
-        system_prompt=get_system_prompt(),
-        content_text=get_prompt("text"),
-    )
-
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "stream": True,
-        "speaker": "Ethan",
-        "key_words": {"text": ["beijing"]},
-    }
-
-    # Retry only when assert_omni_response fails on preset voice gender (see tests/conftest.py).
-    _gender_assert_substr = "estimated gender"
-    _max_retries = 3
-    for attempt in range(_max_retries):
-        try:
-            openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
-            break
-        except AssertionError as e:
-            if _gender_assert_substr not in str(e) or attempt == _max_retries - 1:
-                raise
-            print(f"Gender assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")
-
-
-@pytest.mark.advanced_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_speaker_003(omni_server, openai_client) -> None:
-    """
-    Input Modal: text only (one-word answer constraint).
-    Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
-    Input Setting: stream=True
-    Datasets: single request
-    """
-    messages = dummy_messages_from_mix_data(
-        system_prompt=get_system_prompt(),
-        content_text=get_prompt("text"),
-    )
-
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "stream": True,
-        "speaker": "CHELSIE",
-        "key_words": {"text": ["beijing"]},
-    }
-
-    openai_client.send_omni_request(request_config)
-
-
-@pytest.mark.advanced_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_language_001(omni_server, openai_client) -> None:
-    """
-    Input Modal: text only (one-word answer constraint).
-    Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
-    Input Setting: stream=True
-    Datasets: single request
-    """
-    messages = dummy_messages_from_mix_data(
-        system_prompt=get_system_prompt(),
-        content_text=get_prompt("text_chinese"),
-    )
-
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "stream": True,
-        "key_words": {"text": ["北京"]},
-    }
-
-    openai_client.send_omni_request(request_config)
@@ -120,31 +120,6 @@ def test_voice_002(omni_server, openai_client) -> None:
    openai_client.send_audio_speech_request(request_config)


-@pytest.mark.advanced_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_voice_003(omni_server, openai_client) -> None:
-    """
-    Test text input processing and audio output via OpenAI API.
-    Deploy Setting: default yaml
-    Input Modal: text
-    Output Modal: audio
-    Input Setting: stream=False, language=chinese
-    Datasets: few requests
-    """
-    request_config = {
-        "model": omni_server.model,
-        "input": get_prompt(),
-        "stream": False,
-        "response_format": "wav",
-        "task_type": "CustomVoice",
-        "voice": "SERENA",
-    }
-
-    openai_client.send_audio_speech_request(request_config)
-
-
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@@ -98,18 +98,6 @@ def _get_diffusion_feature_cases(model: str):
            id="parallel_004",
            marks=PARALLEL_FEATURE_MARKS,
        ),
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=[
-                    "--use-hsdp",
-                    "--hsdp-shard-size",
-                    "2",
-                ],
-            ),
-            id="parallel_005",
-            marks=PARALLEL_FEATURE_MARKS,
-        ),
    ]


@@ -107,18 +107,6 @@ def _get_diffusion_feature_cases(model: str):
            id="vae_patch_parallel_2",
            marks=PARALLEL_FEATURE_MARKS,
        ),
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=[
-                    "--use-hsdp",
-                    "--hsdp-shard-size",
-                    "2",
-                ],
-            ),
-            id="parallel_hsdp",
-            marks=PARALLEL_FEATURE_MARKS,
-        ),
    ]


@@ -62,18 +62,6 @@ FEATURE_CASES = [
        id="cfg_parallel_001",
        marks=PARALLEL_FEATURE_MARKS,
    ),
-    pytest.param(
-        OmniServerParams(
-            model=MODEL,
-            server_args=[
-                "--use-hsdp",
-                "--hsdp-shard-size",
-                "2",
-            ],
-        ),
-        id="parallel_hsdp",
-        marks=PARALLEL_FEATURE_MARKS,
-    ),
 ]


@@ -1,161 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Comprehensive e2e tests of diffusion features for Wan2.1-VACE in online serving mode.
-
-Wan2.1-VACE supports: Cache-DiT, Ulysses-SP, Ring, CFG-Parallel, TP,
-VAE-Patch-Parallel, HSDP. TeaCache is NOT supported for this model, so
-Cache-DiT is used in place of TeaCache for single-card and CFG tests.
-
-Uses the 1.3B variant for faster CI testing.
-
-Coverage:
-  Single GPU:
-    - Cache-DiT + layerwise CPU offload
-  Two GPUs:
-    - Cache-DiT + Ulysses-SP = 2
-    - Cache-DiT + Ring = 2
-    - Cache-DiT + CFG-Parallel = 2
-    - Cache-DiT + TP = 2 + VAE-Patch-Parallel = 2
-    - Cache-DiT + HSDP = 2 + VAE-Patch-Parallel = 2
-"""
-
-import pytest
-
-from tests.conftest import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-)
-from tests.utils import hardware_marks
-
-MODEL = "Wan-AI/Wan2.1-VACE-1.3B-diffusers"
-PROMPT = "A cat walking slowly across a sunlit garden path"
-
-SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
-PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
-
-
-def _get_vace_feature_cases():
-    return [
-        # Single GPU: Cache-DiT + layerwise CPU offload
-        pytest.param(
-            OmniServerParams(
-                model=MODEL,
-                server_args=[
-                    "--cache-backend",
-                    "cache_dit",
-                    "--enable-layerwise-offload",
-                    "--vae-use-tiling",
-                ],
-            ),
-            id="single_card_001",
-            marks=SINGLE_CARD_FEATURE_MARKS,
-        ),
-        # 2 GPUs: Cache-DiT + Ulysses-SP = 2
-        pytest.param(
-            OmniServerParams(
-                model=MODEL,
-                server_args=[
-                    "--cache-backend",
-                    "cache_dit",
-                    "--ulysses-degree",
-                    "2",
-                    "--vae-use-tiling",
-                ],
-            ),
-            id="parallel_001",
-            marks=PARALLEL_FEATURE_MARKS,
-        ),
-        # 2 GPUs: Cache-DiT + Ring = 2
-        pytest.param(
-            OmniServerParams(
-                model=MODEL,
-                server_args=[
-                    "--cache-backend",
-                    "cache_dit",
-                    "--ring",
-                    "2",
-                    "--vae-use-tiling",
-                ],
-            ),
-            id="parallel_002",
-            marks=PARALLEL_FEATURE_MARKS,
-        ),
-        # 2 GPUs: Cache-DiT + CFG-Parallel = 2
-        pytest.param(
-            OmniServerParams(
-                model=MODEL,
-                server_args=[
-                    "--cache-backend",
-                    "cache_dit",
-                    "--cfg-parallel-size",
-                    "2",
-                    "--vae-use-tiling",
-                ],
-            ),
-            id="parallel_003",
-            marks=PARALLEL_FEATURE_MARKS,
-        ),
-        # 2 GPUs: Cache-DiT + TP = 2 + VAE-Patch-Parallel = 2
-        pytest.param(
-            OmniServerParams(
-                model=MODEL,
-                server_args=[
-                    "--cache-backend",
-                    "cache_dit",
-                    "--tensor-parallel-size",
-                    "2",
-                    "--vae-patch-parallel-size",
-                    "2",
-                    "--vae-use-tiling",
-                ],
-            ),
-            id="parallel_004",
-            marks=PARALLEL_FEATURE_MARKS,
-        ),
-        # 2 GPUs: Cache-DiT + HSDP = 2 + VAE-Patch-Parallel = 2
-        pytest.param(
-            OmniServerParams(
-                model=MODEL,
-                server_args=[
-                    "--cache-backend",
-                    "cache_dit",
-                    "--hsdp-shard-size",
-                    "2",
-                    "--vae-patch-parallel-size",
-                    "2",
-                    "--vae-use-tiling",
-                ],
-            ),
-            id="parallel_005",
-            marks=PARALLEL_FEATURE_MARKS,
-        ),
-    ]
-
-
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-@pytest.mark.parametrize(
-    "omni_server",
-    _get_vace_feature_cases(),
-    indirect=True,
-)
-def test_wan_2_1_vace(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Test VACE T2V generation with all supported diffusion acceleration features."""
-    openai_client.send_video_diffusion_request(
-        {
-            "model": MODEL,
-            "form_data": {
-                "prompt": PROMPT,
-                "height": 480,
-                "width": 320,
-                "num_frames": 5,
-                "fps": 8,
-                "num_inference_steps": 2,
-                "guidance_scale": 5.0,
-                "seed": 42,
-            },
-        }
-    )
@@ -60,18 +60,6 @@ def _get_diffusion_feature_cases():
            id="parallel_teacache_fp8_ulysses2_ring2",
            marks=FOUR_CARD_MARKS,
        ),
-        pytest.param(
-            OmniServerParams(
-                model=MODEL,
-                server_args=[
-                    "--use-hsdp",
-                    "--hsdp-shard-size",
-                    "2",
-                ],
-            ),
-            id="parallel_hsdp",
-            marks=[*FOUR_CARD_MARKS, pytest.mark.skip(reason="issue #2435")],
-        ),
    ]


@@ -60,7 +60,6 @@ async def generate(

@pytest.mark.core_model
@pytest.mark.omni
-@pytest.mark.real_hf_config
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
@pytest.mark.asyncio
 async def test_abort():
@@ -61,31 +61,3 @@ def test_build_add_request_message_preserves_additional_information():
    assert request.additional_information.entries["text"].list_data == ["hello world"]
    assert request.additional_information.entries["speaker"].list_data == ["vivian"]
    output_processor.add_request.assert_called_once()
-
-
-def test_build_add_request_message_with_resumable_streaming():
-    engine = object.__new__(AsyncOmniEngine)
-    params = SamplingParams(max_tokens=8)
-    engine.default_sampling_params_list = [params]
-    engine.stage_metadata = [{"stage_type": "llm"}]
-    engine.supported_tasks = ("generate",)
-
-    input_processor = Mock()
-    input_processor.process_inputs.return_value = _make_engine_core_request()
-    engine.input_processor = input_processor
-
-    output_processor = Mock()
-    engine.output_processors = [output_processor]
-
-    msg = engine._build_add_request_message(
-        request_id="req-stream",
-        prompt={"prompt_token_ids": [1, 2, 3]},
-        sampling_params_list=[params],
-        final_stage_id=0,
-        resumable=True,
-        message_type="streaming_update",
-    )
-
-    assert msg["type"] == "streaming_update"
-    input_processor.process_inputs.assert_called_once()
-    assert input_processor.process_inputs.call_args.kwargs["resumable"] is True
@@ -113,7 +113,6 @@ def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch):
        executor_class=object,
        engine_manager=object(),
        coordinator=object(),
-        proc=None,
        addresses=types.SimpleNamespace(
            inputs=["inproc://input"],
            outputs=["inproc://output"],
@@ -106,7 +106,7 @@ def test_encode_image_base64():


 class MockGenerationResult:
-    """Mock result object from AsyncOmni.generate()"""
+    """Mock result object from AsyncOmniDiffusion.generate()"""

    def __init__(self, images):
        self.images = images
@@ -233,20 +233,17 @@ def test_app(mocker: MockerFixture):
        uploaded_voices = []
        if hasattr(speech_server, "uploaded_speakers"):
            for voice_name, info in speech_server.uploaded_speakers.items():
-                voice_entry = {
-                    "name": info.get("name", voice_name),
-                    "consent": info.get("consent", ""),
-                    "created_at": info.get("created_at", 0),
-                    "file_size": info.get("file_size", 0),
-                    "mime_type": info.get("mime_type", ""),
-                    "embedding_source": info.get("embedding_source", "audio"),
-                    "embedding_dim": info.get("embedding_dim"),
-                }
-                if info.get("ref_text"):
-                    voice_entry["ref_text"] = info["ref_text"]
-                if info.get("speaker_description"):
-                    voice_entry["speaker_description"] = info["speaker_description"]
-                uploaded_voices.append(voice_entry)
+                uploaded_voices.append(
+                    {
+                        "name": info.get("name", voice_name),
+                        "consent": info.get("consent", ""),
+                        "created_at": info.get("created_at", 0),
+                        "file_size": info.get("file_size", 0),
+                        "mime_type": info.get("mime_type", ""),
+                        "embedding_source": info.get("embedding_source", "audio"),
+                        "embedding_dim": info.get("embedding_dim"),
+                    }
+                )
        return {"voices": speakers, "uploaded_voices": uploaded_voices}

    app.add_api_route("/v1/audio/voices", list_voices, methods=["GET"])
@@ -258,8 +255,7 @@ def test_app(mocker: MockerFixture):
        speaker_embedding: str | None = Form(None),
        consent: str = Form(...),
        name: str = Form(...),
-        ref_text: str | None = Form(None),
-        speaker_description: str | None = Form(None),
+        ref_text: str = Form(None),
    ):
        try:
            if speaker_embedding is not None and audio_sample is not None:
@@ -267,13 +263,7 @@ def test_app(mocker: MockerFixture):
            if speaker_embedding is not None:
                result = await speech_server.upload_voice_embedding(speaker_embedding, consent, name)
            elif audio_sample is not None:
-                result = await speech_server.upload_voice(
-                    audio_sample,
-                    consent,
-                    name,
-                    ref_text=ref_text,
-                    speaker_description=speaker_description,
-                )
+                result = await speech_server.upload_voice(audio_sample, consent, name, ref_text=ref_text)
            else:
                raise ValueError("Either 'audio_sample' or 'speaker_embedding' must be provided")
            return {"success": True, "voice": result}
@@ -407,44 +397,6 @@ class TestSpeechAPI:
        assert result["voice"].get("ref_text") == "Hello world transcript"
        response = client.delete("/v1/audio/voices/test_voice_rt")

-    def test_upload_voice_with_speaker_description(self, client, tmp_path):
-        """Test voice upload with speaker_description stores and returns the description."""
-        # Pre-cleanup in case a previous test run left this voice behind
-        client.delete("/v1/audio/voices/test_voice_vd")
-
-        audio_content = b"fake audio content" * 1000
-        files = {"audio_sample": ("test.wav", audio_content, "audio/wav")}
-        data = {"consent": "c1", "name": "test_voice_vd", "speaker_description": "  warm, energetic narrator  "}
-
-        response = client.post("/v1/audio/voices", files=files, data=data)
-        try:
-            assert response.status_code == 200
-            result = response.json()
-            assert result["success"] is True
-            assert result["voice"]["name"] == "test_voice_vd"
-            assert result["voice"].get("speaker_description") == "warm, energetic narrator"
-        finally:
-            client.delete("/v1/audio/voices/test_voice_vd")
-
-    def test_upload_voice_speaker_description_in_listing(self, client):
-        """Test that speaker_description survives the upload → list round-trip."""
-        client.delete("/v1/audio/voices/test_voice_sd_list")
-
-        audio_content = b"fake audio content" * 1000
-        files = {"audio_sample": ("test.wav", audio_content, "audio/wav")}
-        data = {"consent": "c1", "name": "test_voice_sd_list", "speaker_description": "calm female narrator"}
-
-        response = client.post("/v1/audio/voices", files=files, data=data)
-        try:
-            assert response.status_code == 200
-
-            listing = client.get("/v1/audio/voices").json()
-            uploaded = {v["name"]: v for v in listing["uploaded_voices"]}
-            assert "test_voice_sd_list" in uploaded
-            assert uploaded["test_voice_sd_list"]["speaker_description"] == "calm female narrator"
-        finally:
-            client.delete("/v1/audio/voices/test_voice_sd_list")
-
    def test_upload_voice_file_too_large(self, client):
        """Test voice upload with file exceeding size limit."""
        # Create a file larger than 10MB
@@ -898,7 +850,6 @@ class TestTTSMethods:
                "file_path": "/tmp/voice_samples/custom_voice_consent_123.wav",
                "mime_type": "audio/wav",
                "ref_text": None,
-                "created_at": 1711234567.89,
            }
        }
        speech_server.supported_speakers = {"ryan", "vivian", "custom_voice"}
@@ -911,7 +862,6 @@ class TestTTSMethods:
            assert params["ref_audio"] == ["data:audio/wav;base64,ZmFrZWF1ZGlv"]
            assert params["x_vector_only_mode"] == [True]
            assert params["task_type"] == ["Base"]
-            assert params["voice_created_at"] == [1711234567.89]
            assert "ref_text" not in params

    def test_build_tts_params_with_uploaded_voice_ref_text(self, speech_server):
@@ -922,7 +872,6 @@ class TestTTSMethods:
                "file_path": "/tmp/voice_samples/custom_voice_consent_123.wav",
                "mime_type": "audio/wav",
                "ref_text": "Hello world transcript",
-                "created_at": 1711234567.89,
            }
        }
        speech_server.supported_speakers = {"ryan", "vivian", "custom_voice"}
@@ -936,7 +885,6 @@ class TestTTSMethods:
            assert params["x_vector_only_mode"] == [False]
            assert params["task_type"] == ["Base"]
            assert params["ref_text"] == ["Hello world transcript"]
-            assert params["voice_created_at"] == [1711234567.89]

    def test_build_tts_params_without_uploaded_voice(self, speech_server):
        """Test _build_tts_params does not auto-set ref_audio for non-uploaded voices."""
@@ -1664,7 +1612,6 @@ class TestFishSpeechServing:
        assert all(allowed_special is None for _, _, allowed_special in tokenizer.calls)

    def test_build_fish_clone_prompt_normalizes_text_fields(self, fish_speech_server):
-        fish_speech_server._fish_speech_tokenizer = _FakeFishTokenizer()
        fish_speech_server._estimate_fish_prompt_len = MagicMock(return_value=123)

        request = OpenAICreateSpeechRequest(
@@ -1873,115 +1820,3 @@ class TestWAVStreaming:
        for fmt in unsupported_formats:
            response = client.post("/v1/audio/speech", json={"input": "Hello", "stream": True, "response_format": fmt})
            assert response.status_code == 422
-
-
-# ---- CosyVoice3 Serving Tests ----
-
-
-@pytest.fixture
-def cosyvoice3_server(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="FunAudioLLM/Fun-CosyVoice3-0.5B-2512")
-    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="cosyvoice3_talker"),
-            tts_args={},
-        )
-    ]
-
-    mock_models = mocker.MagicMock()
-    mock_models.is_base_model.return_value = True
-
-    return OmniOpenAIServingSpeech(
-        engine_client=mock_engine_client,
-        models=mock_models,
-        request_logger=mocker.MagicMock(),
-    )
-
-
-class TestCosyVoice3Serving:
-    def test_cosyvoice3_model_type_detection(self, cosyvoice3_server):
-        assert cosyvoice3_server._tts_model_type == "cosyvoice3"
-        assert cosyvoice3_server._is_tts is True
-        assert cosyvoice3_server._is_cosyvoice3 is True
-
-    def test_cosyvoice3_stage_registered(self):
-        from vllm_omni.entrypoints.openai.serving_speech import (
-            _COSYVOICE3_TTS_MODEL_STAGES,
-            _TTS_MODEL_STAGES,
-        )
-
-        assert "cosyvoice3_talker" in _COSYVOICE3_TTS_MODEL_STAGES
-        assert "cosyvoice3_talker" in _TTS_MODEL_STAGES
-
-    def test_validate_cosyvoice3_empty_input(self, cosyvoice3_server):
-        request = OpenAICreateSpeechRequest(input="", ref_audio="data:audio/wav;base64,abc", ref_text="hello")
-        error = cosyvoice3_server._validate_cosyvoice3_request(request)
-        assert error is not None
-        assert "empty" in error.lower()
-
-    def test_validate_cosyvoice3_missing_ref_audio(self, cosyvoice3_server):
-        request = OpenAICreateSpeechRequest(input="Hello", ref_text="hello")
-        error = cosyvoice3_server._validate_cosyvoice3_request(request)
-        assert error is not None
-        assert "ref_audio" in error.lower()
-
-    def test_validate_cosyvoice3_missing_ref_text(self, cosyvoice3_server):
-        request = OpenAICreateSpeechRequest(input="Hello", ref_audio="data:audio/wav;base64,abc")
-        error = cosyvoice3_server._validate_cosyvoice3_request(request)
-        assert error is not None
-        assert "ref_text" in error.lower()
-
-    def test_validate_cosyvoice3_invalid_ref_audio_format(self, cosyvoice3_server):
-        request = OpenAICreateSpeechRequest(input="Hello", ref_audio="/local/path.wav", ref_text="hello")
-        error = cosyvoice3_server._validate_cosyvoice3_request(request)
-        assert error is not None
-        assert "url" in error.lower() or "format" in error.lower()
-
-    def test_validate_cosyvoice3_valid_request(self, cosyvoice3_server):
-        request = OpenAICreateSpeechRequest(
-            input="Hello world",
-            ref_audio="data:audio/wav;base64,abc123",
-            ref_text="Reference transcript",
-        )
-        error = cosyvoice3_server._validate_cosyvoice3_request(request)
-        assert error is None
-
-    def test_validate_cosyvoice3_max_new_tokens_range(self, cosyvoice3_server):
-        request = OpenAICreateSpeechRequest(
-            input="Hello",
-            ref_audio="data:audio/wav;base64,abc",
-            ref_text="hello",
-            max_new_tokens=0,
-        )
-        error = cosyvoice3_server._validate_cosyvoice3_request(request)
-        assert error is not None
-        assert "max_new_tokens" in error
-
-    def test_prepare_speech_generation_cosyvoice3(self, cosyvoice3_server):
-        cosyvoice3_server._build_cosyvoice3_prompt = AsyncMock(
-            return_value={
-                "prompt": "Hello",
-                "multi_modal_data": {"audio": (np.zeros(24000), 24000)},
-                "mm_processor_kwargs": {"prompt_text": "ref text", "sample_rate": 24000},
-            }
-        )
-
-        request = OpenAICreateSpeechRequest(
-            input="Hello",
-            ref_audio="data:audio/wav;base64,abc",
-            ref_text="Reference text",
-        )
-        request_id, generator, tts_params = asyncio.run(cosyvoice3_server._prepare_speech_generation(request))
-
-        assert request_id.startswith("speech-")
-        assert generator == "generator"
-        assert tts_params == {}
-        cosyvoice3_server._build_cosyvoice3_prompt.assert_awaited_once()
@@ -13,8 +13,8 @@ def test_generate_accepts_request_after_repeated_cancellations():
        submitted_request_ids = []
        aborted_request_batches = []

-        async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id, **kwargs):
-            del prompt, sampling_params_list, final_stage_id, kwargs
+        async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id):
+            del prompt, sampling_params_list, final_stage_id
            submitted_request_ids.append(request_id)

        async def fake_abort_async(request_ids):
@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+import vllm_omni.diffusion.stage_diffusion_client as stage_diffusion_client_module
+from vllm_omni.diffusion.data import DiffusionRequestAbortedError
+from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_get_diffusion_od_config_returns_direct_config():
+    diffusion = object.__new__(AsyncOmniDiffusion)
+    diffusion.od_config = object()
+
+    assert diffusion.get_diffusion_od_config() is diffusion.od_config
+
+
+def test_async_omni_diffusion_generate_aborts_engine_on_cancel():
+    async def run_test():
+        started = threading.Event()
+        release = threading.Event()
+        abort = Mock()
+
+        def step(request):
+            del request
+            started.set()
+            release.wait(timeout=5)
+            return [SimpleNamespace(request_id="req-1")]
+
+        diffusion = object.__new__(AsyncOmniDiffusion)
+        diffusion.engine = SimpleNamespace(step=step, abort=abort)
+        diffusion._executor = ThreadPoolExecutor(max_workers=1)
+
+        task = asyncio.create_task(
+            diffusion.generate(
+                prompt="hello",
+                sampling_params=OmniDiffusionSamplingParams(),
+                request_id="req-1",
+            )
+        )
+        try:
+            assert await asyncio.to_thread(started.wait, 1)
+            task.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+        finally:
+            release.set()
+            diffusion._executor.shutdown(wait=True)
+
+        abort.assert_called_once_with("req-1")
+
+    asyncio.run(run_test())
+
+
+def test_stage_diffusion_client_abort_requests_forwards_to_engine():
+    async def run_test():
+        aborted_request_ids: list[list[str]] = []
+
+        async def abort(request_ids):
+            aborted_request_ids.append(request_ids)
+
+        client = object.__new__(StageDiffusionClient)
+        client._engine = SimpleNamespace(abort=abort)
+        client._tasks = {}
+
+        task = asyncio.create_task(asyncio.sleep(60))
+        client._tasks["req-1"] = task
+
+        await client.abort_requests_async(["req-1", "req-2"])
+
+        with pytest.raises(asyncio.CancelledError):
+            await task
+        assert client._tasks == {}
+        assert aborted_request_ids == [["req-1", "req-2"]]
+
+    asyncio.run(run_test())
+
+
+def test_stage_diffusion_client_run_treats_abort_as_normal_path(monkeypatch):
+    async def run_test():
+        async def generate(prompt, sampling_params, request_id):
+            del prompt, sampling_params
+            raise DiffusionRequestAbortedError(f"Request {request_id} aborted.")
+
+        info = Mock()
+        exception = Mock()
+        monkeypatch.setattr(stage_diffusion_client_module.logger, "info", info)
+        monkeypatch.setattr(stage_diffusion_client_module.logger, "exception", exception)
+
+        client = object.__new__(StageDiffusionClient)
+        client.stage_id = 3
+        client._engine = SimpleNamespace(generate=generate)
+        client._output_queue = asyncio.Queue()
+        client._tasks = {"req-1": object()}
+
+        await client._run("req-1", "prompt", OmniDiffusionSamplingParams())
+
+        assert client._output_queue.empty()
+        assert client._tasks == {}
+        info.assert_called_once()
+        exception.assert_not_called()
+
+    asyncio.run(run_test())
@@ -113,7 +113,6 @@ class FakeAsyncOmniEngine:
        sampling_params_list: list[Any] | None = None,
        final_stage_id: int = 0,
        arrival_time: float | None = None,
-        **kwargs: Any,
    ) -> None:
        msg = {
            "request_id": request_id,
@@ -14,6 +14,7 @@ from collections.abc import Callable
 from pathlib import Path
 from typing import Any, NamedTuple, cast

+import mistune
 import pytest
 import torch
 from safetensors.torch import save_file
@@ -61,8 +62,6 @@ class ReadmeSnippet(NamedTuple):
        readme_path: Path,
        skipif: ReadmeSnippetExtractionSkipPredicate | None = None,
    ) -> list["ReadmeSnippet"]:
-        import mistune
-
        markdown = mistune.create_markdown(renderer="ast")
        tokens = markdown(readme_path.read_text(encoding="utf-8"))
        tokens = cast(list[dict[str, Any]], tokens)  # mistune's AST renderer always produces a list, not a str
@@ -1,108 +0,0 @@
-import math
-
-import pytest
-import torch
-
-from vllm_omni.model_executor.models.fish_speech import fish_speech_slow_ar as slow_ar_module
-from vllm_omni.model_executor.models.fish_speech.fish_speech_dac_decoder import FishSpeechDACDecoder
-from vllm_omni.model_executor.models.fish_speech.fish_speech_slow_ar import (
-    FishSpeechSlowARForConditionalGeneration,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class _FakeCodec:
-    def decode(self, codes_bqf: torch.Tensor, feature_lengths: torch.Tensor):
-        del codes_bqf, feature_lengths
-        wav = torch.arange(100, dtype=torch.float32).view(1, 1, 100)
-        audio_lengths = torch.tensor([100], dtype=torch.long)
-        return wav, audio_lengths
-
-
-class _FakeTokenizer:
-    def __init__(self, mapping, unk_token_id=-1):
-        self._mapping = mapping
-        self.unk_token_id = unk_token_id
-
-    def convert_tokens_to_ids(self, token: str) -> int:
-        return self._mapping.get(token, self.unk_token_id)
-
-
-def test_dac_decoder_mixed_batch_empty_request_does_not_misalign_indices():
-    decoder = object.__new__(FishSpeechDACDecoder)
-    torch.nn.Module.__init__(decoder)
-    decoder._codec = _FakeCodec()
-    decoder._num_codebooks = 10
-    decoder._output_sample_rate = 44100
-    decoder._hop_length = 512
-    decoder._logged_codec_stats = False
-    decoder._ensure_codec_loaded = lambda: None
-    decoder._split_request_ids = lambda ids, seq_token_counts=None: [
-        torch.empty((0,), dtype=torch.long),
-        torch.arange(20, dtype=torch.long),
-    ]
-
-    out = decoder.forward(
-        input_ids=torch.arange(20, dtype=torch.long),
-        runtime_additional_information=[{}, {"left_context_size": 1}],
-    )
-
-    audios = out.multimodal_outputs["model_outputs"]
-    assert len(audios) == 2
-    assert audios[0].numel() == 0
-    # 2 total frames with 1 frame of left context => proportional trim removes half the samples.
-    assert audios[1].shape[0] == 50
-
-
-def test_structured_voice_clone_prefill_adds_full_codebooks_with_decode_scale(monkeypatch):
-    model = object.__new__(FishSpeechSlowARForConditionalGeneration)
-    torch.nn.Module.__init__(model)
-    model._num_codebooks = 2
-    model._codebook_size = 8
-    model._semantic_begin_id = 100
-    model.model_path = "unused"
-
-    hidden_size = 3
-    text_embed = torch.nn.Embedding(256, hidden_size)
-    codebook_embed = torch.nn.Embedding(model._num_codebooks * model._codebook_size, hidden_size)
-    with torch.no_grad():
-        text_embed.weight.zero_()
-        text_embed.weight[20] = torch.tensor([1.0, 2.0, 3.0])
-        text_embed.weight[21] = torch.tensor([4.0, 5.0, 6.0])
-        codebook_embed.weight.zero_()
-        codebook_embed.weight[1] = torch.tensor([10.0, 0.0, 0.0])
-        codebook_embed.weight[10] = torch.tensor([0.0, 20.0, 0.0])
-        codebook_embed.weight[3] = torch.tensor([30.0, 0.0, 0.0])
-        codebook_embed.weight[12] = torch.tensor([0.0, 40.0, 0.0])
-
-    model.embed_input_ids = lambda ids: text_embed(ids)
-    model.codebook_embeddings = codebook_embed
-    model._get_tokenizer = lambda: _FakeTokenizer({"<|audio_start|>": 10, "<|audio_end|>": 11})
-
-    monkeypatch.setattr(slow_ar_module.np, "load", lambda path: [0.0])
-    monkeypatch.setattr(slow_ar_module.os, "remove", lambda path: None)
-    monkeypatch.setattr(
-        slow_ar_module,
-        "encode_reference_audio_codes",
-        lambda *args, **kwargs: torch.tensor([[1, 2], [3, 4]], dtype=torch.long),
-    )
-    monkeypatch.setattr(
-        slow_ar_module,
-        "build_fish_voice_clone_prompt_ids",
-        lambda tokenizer, text, ref_text, semantic_token_ids: ([1, 10, 20, 21, 11, 2], None, None),
-    )
-
-    prefill = model._build_structured_voice_clone_prefill_embeds(
-        {
-            "ref_text": "ref",
-            "text": "target",
-            "ref_audio_path": "unused.npy",
-            "ref_audio_sr": 16000,
-        }
-    )
-
-    expected_0 = (torch.tensor([1.0, 2.0, 3.0]) + torch.tensor([10.0, 20.0, 0.0])) / math.sqrt(3.0)
-    expected_1 = (torch.tensor([4.0, 5.0, 6.0]) + torch.tensor([30.0, 40.0, 0.0])) / math.sqrt(3.0)
-    assert torch.allclose(prefill[2].to(dtype=torch.float32), expected_0, atol=2e-2, rtol=0)
-    assert torch.allclose(prefill[3].to(dtype=torch.float32), expected_1, atol=2e-2, rtol=0)
@@ -21,10 +21,9 @@ from vllm.config.multimodal import (
    ImageDummyOptions,
    VideoDummyOptions,
 )
-from vllm.inputs import MultiModalDataDict, MultiModalInput
-from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalDataDict
 from vllm.multimodal.cache import MultiModalProcessorOnlyCache
-from vllm.multimodal.inputs import batched_tensors_equal
+from vllm.multimodal.inputs import MultiModalInputs, batched_tensors_equal
 from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext
 from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config

@@ -303,8 +302,8 @@ def test_omni_processing_correctness(


 def _assert_inputs_equal(
-    a: MultiModalInput,
-    b: MultiModalInput,
+    a: MultiModalInputs,
+    b: MultiModalInputs,
    *,
    ignore_mm_keys: set[str] | None = None,
    msg: str = "",
@@ -1,129 +0,0 @@
-import threading
-
-import pytest
-
-from vllm_omni.utils.voice_cache import VoiceEmbeddingCache
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-@pytest.fixture
-def cache():
-    return VoiceEmbeddingCache(max_entries=4)
-
-
-class TestVoiceEmbeddingCache:
-    def test_miss_returns_none(self, cache: VoiceEmbeddingCache):
-        assert cache.get("nonexistent") is None
-        assert cache.stats()["misses"] == 1
-
-    def test_put_and_hit(self, cache: VoiceEmbeddingCache):
-        cache.put("abc", {"val": 42})
-        result = cache.get("abc")
-        assert result is not None
-        assert result["val"] == 42
-        assert cache.stats()["hits"] == 1
-
-    def test_lru_eviction(self, cache: VoiceEmbeddingCache):
-        for i in range(5):
-            cache.put(f"key{i}", {"i": i})
-        # key0 should have been evicted (oldest, max_entries=4)
-        assert cache.get("key0") is None
-        # key1..key4 should still be present
-        for i in range(1, 5):
-            assert cache.get(f"key{i}") is not None
-        assert cache.stats()["entries"] == 4
-
-    def test_lru_access_promotes(self, cache: VoiceEmbeddingCache):
-        cache.put("a", {"v": 1})
-        cache.put("b", {"v": 2})
-        cache.put("c", {"v": 3})
-        cache.put("d", {"v": 4})
-        # Access "a" to promote it to MRU
-        cache.get("a")
-        # Insert "e" -- should evict "b" (now the oldest), not "a"
-        cache.put("e", {"v": 5})
-        assert cache.get("a") is not None
-        assert cache.get("b") is None
-
-    def test_put_overwrites(self, cache: VoiceEmbeddingCache):
-        cache.put("k", {"old": True})
-        cache.put("k", {"new": True})
-        result = cache.get("k")
-        assert result is not None
-        assert "new" in result
-        assert "old" not in result
-        assert cache.stats()["entries"] == 1
-
-    def test_make_cache_key_includes_mode(self):
-        k1 = VoiceEmbeddingCache.make_cache_key("alice", xvec_only=True)
-        k2 = VoiceEmbeddingCache.make_cache_key("alice", xvec_only=False)
-        assert k1 != k2
-        assert "xvec" in k1
-        assert "icl" in k2
-
-    def test_make_cache_key_deterministic(self):
-        k1 = VoiceEmbeddingCache.make_cache_key("bob", xvec_only=True)
-        k2 = VoiceEmbeddingCache.make_cache_key("bob", xvec_only=True)
-        assert k1 == k2
-
-    def test_make_cache_key_created_at_isolation(self):
-        """Different created_at timestamps must produce different keys (stale-cache protection)."""
-        k1 = VoiceEmbeddingCache.make_cache_key("alice", xvec_only=False, created_at=1000.0)
-        k2 = VoiceEmbeddingCache.make_cache_key("alice", xvec_only=False, created_at=2000.0)
-        assert k1 != k2
-
-    def test_stale_cache_protection(self, cache: VoiceEmbeddingCache):
-        """Re-upload (new created_at) must produce a cache miss, not a stale hit."""
-        key_old = VoiceEmbeddingCache.make_cache_key("alice", xvec_only=False, created_at=1000.0)
-        key_new = VoiceEmbeddingCache.make_cache_key("alice", xvec_only=False, created_at=2000.0)
-        cache.put(key_old, {"ref_spk_embedding": "old_emb"})
-        # Re-upload produces a new created_at → different key → cold miss
-        assert cache.get(key_new) is None
-        # Old key still in cache (not yet evicted)
-        assert cache.get(key_old) is not None
-
-    def test_cache_mode_isolation(self, cache: VoiceEmbeddingCache):
-        """xvec entry must NOT be served for an icl request (same voice)."""
-        key_xvec = VoiceEmbeddingCache.make_cache_key("alice", xvec_only=True)
-        key_icl = VoiceEmbeddingCache.make_cache_key("alice", xvec_only=False)
-        cache.put(key_xvec, {"ref_code": None, "ref_spk_embedding": "emb"})
-        # icl request should miss — different key
-        assert cache.get(key_icl) is None
-        # xvec request should hit
-        assert cache.get(key_xvec) is not None
-
-    def test_stats_counters(self, cache: VoiceEmbeddingCache):
-        cache.put("x", {"v": 1})
-        cache.get("x")  # hit
-        cache.get("x")  # hit
-        cache.get("y")  # miss
-        s = cache.stats()
-        assert s["hits"] == 2
-        assert s["misses"] == 1
-        assert s["entries"] == 1
-        assert s["max_entries"] == 4
-
-    def test_thread_safety(self):
-        cache = VoiceEmbeddingCache(max_entries=32)
-        errors = []
-
-        def worker(thread_id: int):
-            try:
-                for i in range(50):
-                    key = f"t{thread_id}_k{i}"
-                    cache.put(key, {"tid": thread_id, "i": i})
-                    cache.get(key)
-                    cache.get(f"t{(thread_id + 1) % 10}_k{i}")
-            except Exception as e:
-                errors.append(e)
-
-        threads = [threading.Thread(target=worker, args=(t,)) for t in range(10)]
-        for t in threads:
-            t.start()
-        for t in threads:
-            t.join()
-
-        assert not errors, f"Thread safety errors: {errors}"
-        s = cache.stats()
-        assert s["entries"] <= 32
@@ -358,7 +358,6 @@ async def benchmark(
    lora_modules: Iterable[str] | None,
    extra_headers: dict | None,
    extra_body: dict | None,
-    lora_assignment: Literal["random", "round-robin"] = "random",
    ramp_up_strategy: Literal["linear", "exponential"] | None = None,
    ramp_up_start_rps: int | None = None,
    ramp_up_end_rps: int | None = None,
@@ -455,11 +454,8 @@ async def benchmark(
    print("Starting main benchmark run...")

    if lora_modules:
-        lora_modules_list = list(lora_modules)
-        if lora_assignment == "round-robin":
-            lora_modules = iter([lora_modules_list[i % len(lora_modules_list)] for i in range(len(input_requests))])
-        else:
-            lora_modules = iter([random.choice(lora_modules_list) for _ in range(len(input_requests))])
+        # For each input request, choose a LoRA module at random.
+        lora_modules = iter([random.choice(lora_modules) for _ in range(len(input_requests))])

    if profile:
        print("Starting profiler...")
@@ -95,54 +95,54 @@ class OmniARScheduler(VLLMScheduler):
            return False

        criteria_type = self.kv_transfer_criteria.get("type")
-        stop_decode_on_trigger = self.kv_transfer_criteria.get("stop_after_transfer", True)
+        if (
+            self.kv_transfer_criteria.get("stop_after_transfer", True)
+            and request.request_id in self.transfer_triggered_requests
+        ):
+            # For split pipelines that only need the transferred KV
+            # snapshot, stop AR decode once KV extraction has completed.
+            # This frees stage-0 resources without requiring an
+            # orchestrator-side abort.
+            if request.request_id not in self.active_kv_transfers:
+                request.status = RequestStatus.FINISHED_STOPPED
+                return True
+            return False

        if request.request_id in self.transfer_triggered_requests:
-            # Already triggered.  When stop_decode_on_trigger is True AND
-            # transfer was actually queued, the request was already stopped
-            # at trigger time (see below).  Any request that reaches this
-            # point either has stop_decode_on_trigger=False (continue
-            # decoding) or was not actually queued (should not be stopped).
            return False

        if criteria_type == "prefill_finished":
            if request.num_computed_tokens >= request.num_prompt_tokens:
+                logger.debug(f"[Omni] Request {request.request_id} triggered prefill_finished transfer (Non-Stop)")
                self.transfer_triggered_requests.add(request.request_id)
                self._mark_request_for_kv_transfer(request.request_id, request.num_computed_tokens)
-                actually_queued = request.request_id in self.requests_needing_kv_transfer
-
-                if stop_decode_on_trigger and actually_queued:
-                    # Stop immediately so the request is NOT scheduled in
-                    # the next step, freeing scheduling budget for companion
-                    # requests whose chunked-prefill boundaries must be
-                    # deterministic.  waiting_for_transfer_free keeps blocks
-                    # alive until the model runner finishes KV extraction.
-                    self.waiting_for_transfer_free.add(request.request_id)
-                    request.status = RequestStatus.FINISHED_STOPPED
-                    return True

+                # Return False means "Do NOT stop the request" -> Continue Decoding
                return False

        elif criteria_type == "special_token":
            target_token_id = self.kv_transfer_criteria.get("token_id")
            if target_token_id is not None and target_token_id in new_token_ids:
+                logger.debug(f"[Omni] Request {request.request_id} triggered special_token criteria (Non-Stop)")
+
                self.transfer_triggered_requests.add(request.request_id)

+                # Calculate precise snapshot length (trim to sentinel)
+                # Find the FIRST occurrence of the sentinel
                try:
                    idx = new_token_ids.index(target_token_id)
+                    # seq_len = tokens_before_this_step + idx + 1 (include sentinel)
+                    # request.num_computed_tokens already includes ALL new_token_ids
+                    # so we subtract (len(new_token_ids) - (idx + 1))
                    tokens_to_exclude = len(new_token_ids) - (idx + 1)
                    snapshot_len = request.num_computed_tokens - tokens_to_exclude
                except ValueError:
                    snapshot_len = request.num_computed_tokens

+                # Trigger Transfer
                self._mark_request_for_kv_transfer(request.request_id, snapshot_len)
-                actually_queued = request.request_id in self.requests_needing_kv_transfer
-
-                if stop_decode_on_trigger and actually_queued:
-                    self.waiting_for_transfer_free.add(request.request_id)
-                    request.status = RequestStatus.FINISHED_STOPPED
-                    return True

+                # Do NOT stop request
                return False

        return False
@@ -235,10 +235,7 @@ class OmniARScheduler(VLLMScheduler):
            # These blocks contain externally computed tokens that failed to
            # load. Identify affected requests and adjust their computed token
            # count to trigger recomputation of the invalid blocks.
-            failed_kv_load_req_ids = self._handle_invalid_blocks(
-                kv_connector_output.invalid_block_ids,
-                num_scheduled_tokens,
-            )
+            failed_kv_load_req_ids = self._handle_invalid_blocks(kv_connector_output.invalid_block_ids)

        # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
        # the below loop can be a performance bottleneck. We should do our best
@@ -535,12 +532,9 @@ class OmniARScheduler(VLLMScheduler):
                    # It triggered but hasn't finished yet. We MUST wait.
                    logger.debug(f"[Omni] Request {request_id} finished but transfer is still ACTIVE. Waiting.")
                    self.waiting_for_transfer_free.add(request_id)
-                    kv_xfer_params = None
+                    # We do NOT mark for transfer again, just wait.
+                    kv_xfer_params = None  # No new transfer params
                    return kv_xfer_params
-                elif request_id in self.waiting_for_transfer_free:
-                    # Stopped immediately by stop_decode_on_trigger; blocks are
-                    # held until KV extraction completes in a future step.
-                    return None
                else:
                    logger.debug(
                        f"[Omni] Request {request_id} finished and transfer no longer ACTIVE (extracted/acked). "
@@ -272,7 +272,7 @@ class OmniGenerationScheduler(VLLMScheduler):

        # KVTransfer: package metadata
        if self.connector is not None:
-            meta = self._build_kv_connector_meta(self.connector, scheduler_output)
+            meta = self.connector.build_connector_meta(scheduler_output)
            scheduler_output.kv_connector_metadata = meta
        # EC Connector: package metadata
        if self.ec_connector is not None:
@@ -368,10 +368,7 @@ class OmniGenerationScheduler(VLLMScheduler):

        failed_kv_load_req_ids = None
        if kv_connector_output and getattr(kv_connector_output, "invalid_block_ids", None):
-            failed_kv_load_req_ids = self._handle_invalid_blocks(
-                kv_connector_output.invalid_block_ids,
-                num_scheduled_tokens,
-            )
+            failed_kv_load_req_ids = self._handle_invalid_blocks(kv_connector_output.invalid_block_ids)

        # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
        # the below loop can be a performance bottleneck. We should do our best
@@ -5,9 +5,6 @@

 import torch
 import torch.nn.functional as F
-from vllm.logger import init_logger
-
-logger = init_logger(__name__)

 __all__ = ["update_out_and_lse", "flatten_varlen_lse", "unflatten_varlen_lse"]

@@ -82,9 +79,9 @@ def _update_out_and_lse(
        out = out - F.sigmoid(block_lse - lse) * (out - block_out)
        lse = lse - F.logsigmoid(lse - block_lse)
    except RuntimeError as e:
-        logger.error("_update_out_and_lse failed: %s", e)
-        logger.error("out: %s, lse: %s", out.shape, lse.shape)
-        logger.error("block_out: %s, block_lse: %s", block_out.shape, block_lse.shape)
+        print(f"ERROR in _update_out_and_lse: {e}")
+        print(f"out: {out.shape}, lse: {lse.shape}")
+        print(f"block_out: {block_out.shape}, block_lse: {block_lse.shape}")
        # raise e
        raise e

@@ -32,9 +32,7 @@ if current_omni_platform.is_rocm():
        pass
 elif current_omni_platform.is_xpu():
    try:
-        from vllm._xpu_ops import xpu_ops  # noqa: F401
-
-        flash_attn_varlen_func = xpu_ops.flash_attn_varlen_func
+        from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func  # noqa: F401
    except (ImportError, ModuleNotFoundError):
        pass
 elif current_omni_platform.is_musa():
@@ -36,7 +36,6 @@ class Attention(nn.Module):
        scatter_idx: int = 2,
        gather_idx: int = 1,
        use_sync: bool = False,
-        skip_sequence_parallel: bool = False,
    ):
        super().__init__()
        self.attn_backend = get_attn_backend(-1)
@@ -63,7 +62,6 @@ class Attention(nn.Module):
        self.gather_idx = gather_idx
        self.use_sync = use_sync
        self.causal = causal
-        self.skip_sequence_parallel = skip_sequence_parallel

        self.use_ring = False
        self.ring_pg = None
@@ -100,8 +98,6 @@ class Attention(nn.Module):
        (e.g., in noise_refiner/context_refiner before unified_prepare in Z-Image).
        This avoids unnecessary SP communication for layers not covered by _sp_plan.
        """
-        if self.skip_sequence_parallel:
-            return self._no_parallel_strategy
        if is_forward_context_available():
            ctx = get_forward_context()
            if not ctx.sp_active:
@@ -24,7 +24,6 @@ from vllm.logger import init_logger

 from vllm_omni.diffusion.cache.base import CacheBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig, OmniDiffusionConfig
-from vllm_omni.diffusion.utils.tf_utils import get_transformer_from_pipeline

 logger = init_logger(__name__)

@@ -534,7 +533,7 @@ def enable_cache_for_sd3(pipeline: Any, cache_config: Any) -> Callable[[int], No

 def enable_cache_for_ltx2(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
    """Enable cache-dit for LTX2 pipelines (audio-video transformer blocks)."""
-    transformer = get_transformer_from_pipeline(pipeline)
+    transformer = pipeline.transformer

    db_cache_config = _build_db_cache_config(cache_config)

@@ -567,12 +566,11 @@ def enable_cache_for_ltx2(pipeline: Any, cache_config: Any) -> Callable[[int], N
    )

    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        transformer = get_transformer_from_pipeline(pipeline)
        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
        else:
            cache_dit.refresh_context(
-                transformer,
+                pipeline.transformer,
                cache_config=DBCacheConfig().reset(
                    num_inference_steps=num_inference_steps,
                    steps_computation_mask=cache_dit.steps_mask(
@@ -615,9 +613,8 @@ def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], No
    )

    # Enable cache-dit on the transformer
-    transformer = get_transformer_from_pipeline(pipeline)
    cache_dit.enable_cache(
-        transformer,
+        pipeline.transformer,
        cache_config=db_cache_config,
        calibrator_config=calibrator_config,
    )
@@ -629,12 +626,11 @@ def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], No
            pipeline: The diffusion pipeline instance.
            num_inference_steps: New number of inference steps.
        """
-        transformer = get_transformer_from_pipeline(pipeline)
        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
        else:
            cache_dit.refresh_context(
-                transformer,
+                pipeline.transformer,
                cache_config=DBCacheConfig().reset(
                    num_inference_steps=num_inference_steps,
                    steps_computation_mask=cache_dit.steps_mask(
@@ -1215,8 +1211,6 @@ CUSTOM_DIT_ENABLERS.update(
        "StableDiffusion3Pipeline": enable_cache_for_sd3,
        "LTX2Pipeline": enable_cache_for_ltx2,
        "LTX2ImageToVideoPipeline": enable_cache_for_ltx2,
-        "LTX2TwoStagesPipeline": enable_cache_for_ltx2,
-        "LTX2ImageToVideoTwoStagesPipeline": enable_cache_for_ltx2,
        "BagelPipeline": enable_cache_for_bagel,
        "GlmImagePipeline": enable_cache_for_glm_image,
        "Flux2Pipeline": enable_cache_for_flux2,
@@ -193,24 +193,12 @@ class TransformerConfig:
    """Container for raw transformer configuration dictionaries."""

    params: dict[str, Any] = field(default_factory=dict)
-    quant_method: str | None = None
-    quant_config: "QuantizationConfig | None" = None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "TransformerConfig":
        if not isinstance(data, dict):
            raise TypeError(f"Expected transformer config dict, got {type(data)!r}")
-        params = dict(data)  # copy to avoid mutating caller's dict
-
-        quant_method: str | None = None
-        quant_config: QuantizationConfig | None = None
-        disk_qc = params.get("quantization_config")
-        if isinstance(disk_qc, dict) and "quant_method" in disk_qc:
-            quant_method = disk_qc["quant_method"]
-            kwargs = {k: v for k, v in disk_qc.items() if k != "quant_method"}
-            quant_config = build_quant_config(quant_method, **kwargs)
-
-        return cls(params=params, quant_method=quant_method, quant_config=quant_config)
+        return cls(params=dict(data))

    def to_dict(self) -> dict[str, Any]:
        return dict(self.params)
@@ -610,17 +598,6 @@ class OmniDiffusionConfig:
            # If it's neither dict nor DiffusionCacheConfig, convert to empty config
            self.cache_config = DiffusionCacheConfig()

-        # Auto-detect quantization from TransformerConfig if not explicitly set.
-        # This covers the case where tf_model_config is passed at construction
-        # time.  For late (post-construction) assignment, callers should use
-        # set_tf_model_config() which propagates quant_config automatically.
-        if self.quantization_config is None and self.tf_model_config.quant_config is not None:
-            self.quantization_config = self.tf_model_config.quant_config
-            logger.info(
-                "Auto-detected quantization '%s' from model config",
-                self.tf_model_config.quant_method,
-            )
-
        # Resolve quantization_config: str/dict -> QuantizationConfig via build_quant_config.
        if self.quantization_config is not None:
            if isinstance(self.quantization_config, QuantizationConfig):
@@ -640,28 +617,6 @@ class OmniDiffusionConfig:
        elif self.max_cpu_loras < 1:
            raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA")

-    def set_tf_model_config(self, tf_config: "TransformerConfig") -> None:
-        """Assign `tf_model_config` and propagate quantization if detected.
-
-        In the normal startup flow `OmniDiffusionConfig` is created
-        *before* the transformer `config.json` is loaded from disk, so
-        `__post_init__` sees an empty `TransformerConfig`.  Callers
-        that load the config later should use this method instead of bare
-        assignment so that an embedded `quant_config` is propagated to
-        `self.quantization_config` automatically.
-
-        Args:
-            tf_config: Transformer configuration, typically built via
-                `TransformerConfig.from_dict`.
-        """
-        self.tf_model_config = tf_config
-        if self.quantization_config is None and tf_config.quant_config is not None:
-            self.quantization_config = tf_config.quant_config
-            logger.info(
-                "Auto-detected quantization '%s' from model config",
-                tf_config.quant_method,
-            )
-
    def update_multimodal_support(self) -> None:
        self.supports_multimodal_inputs = self.model_class_name in {"QwenImageEditPlusPipeline"}

@@ -115,7 +115,7 @@ class DiffusionEngine:
        if output.aborted:
            raise DiffusionRequestAbortedError(output.abort_message or "Diffusion request aborted.")
        if output.error:
-            raise RuntimeError(f"{output.error}")
+            raise Exception(f"{output.error}")
        logger.info("Generation completed successfully.")

        if output.output is None:
@@ -361,8 +361,8 @@ class DiffusionEngine:
    def _dummy_run(self):
        """A dummy run to warm up the model."""
        num_inference_steps = 1
-        height = 512
-        width = 512
+        height = 1024
+        width = 1024
        if supports_image_input(self.od_config.model_class_name):
            # Provide a dummy image input if the model supports it
            color_format = image_color_format(self.od_config.model_class_name)
@@ -1,9 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from typing import Any
-
-
-def is_transformer_block_module(name: str, module: Any) -> bool:
-    """Return True for numbered modules under `transformer_blocks`."""
-    return "transformer_blocks" in name and name.split(".")[-1].isdigit()
@@ -31,12 +31,6 @@ def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]:
    import numpy as np

    tensor = tensor.detach().cpu().contiguous()
-    original_dtype = tensor.dtype
-    # NumPy does not support bfloat16; promote to float32 for the SHM
-    # transfer and record the original dtype so _tensor_from_shm can
-    # convert back.  The round-trip is lossless for bfloat16 values.
-    if original_dtype == torch.bfloat16:
-        tensor = tensor.to(torch.float32)
    arr = tensor.numpy()
    nbytes = arr.nbytes
    shm = shared_memory.SharedMemory(create=True, size=nbytes)
@@ -46,7 +40,7 @@ def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]:
        "__tensor_shm__": True,
        "name": shm.name,
        "shape": list(tensor.shape),
-        "torch_dtype": str(original_dtype),
+        "torch_dtype": str(tensor.dtype),
        "numpy_dtype": str(arr.dtype),
        "nbytes": nbytes,
    }
@@ -65,13 +59,6 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
        np_dtype = np.dtype(handle["numpy_dtype"])
        arr = np.ndarray(handle["shape"], dtype=np_dtype, buffer=shm.buf[: handle["nbytes"]])
        tensor = torch.from_numpy(arr.copy())
-        # Restore the original dtype if it differs from the numpy-compatible
-        # dtype used for the SHM transfer (e.g. bfloat16 → float32 → bfloat16).
-        torch_dtype_str = handle.get("torch_dtype", "")
-        if torch_dtype_str:
-            original_dtype = getattr(torch, torch_dtype_str.replace("torch.", ""), None)
-            if original_dtype is not None and tensor.dtype != original_dtype:
-                tensor = tensor.to(original_dtype)
    finally:
        shm.close()
        shm.unlink()
@@ -1,16 +1,11 @@
 from importlib.util import find_spec
-from typing import TYPE_CHECKING

 import torch
 import torch.nn as nn
 from vllm.logger import init_logger
-from vllm.model_executor.layers.linear import ReplicatedLinear

 from vllm_omni.diffusion.layers.custom_op import CustomOp

-if TYPE_CHECKING:
-    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
-
 logger = init_logger(__name__)

 _HAS_MINDIESD = find_spec("mindiesd") is not None
@@ -128,102 +123,3 @@ class AdaLayerNorm(CustomOp):
        shift_result, scale_result, gate_result = self.preprocess(mod_params, index)

        return self.layernorm(x) * (1 + scale_result) + shift_result, gate_result
-
-
-class AdaLayerNormZero(nn.Module):
-    def __init__(
-        self,
-        embedding_dim: int,
-        bias: bool = True,
-        quant_config: "QuantizationConfig | None" = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.emb = None
-        self.silu = nn.SiLU()
-        self.linear = ReplicatedLinear(
-            embedding_dim,
-            6 * embedding_dim,
-            bias=bias,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.linear",
-        )
-        self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        emb: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        emb = self.linear(self.silu(emb))
-        if isinstance(emb, tuple):
-            emb = emb[0]
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1)
-        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
-        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
-
-
-class AdaLayerNormZeroSingle(nn.Module):
-    def __init__(
-        self,
-        embedding_dim: int,
-        bias: bool = True,
-        quant_config: "QuantizationConfig | None" = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.silu = nn.SiLU()
-        self.linear = ReplicatedLinear(
-            embedding_dim,
-            3 * embedding_dim,
-            bias=bias,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.linear",
-        )
-        self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        emb: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        emb = self.linear(self.silu(emb))
-        if isinstance(emb, tuple):
-            emb = emb[0]
-        shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1)
-        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
-        return x, gate_msa
-
-
-class AdaLayerNormContinuous(nn.Module):
-    def __init__(
-        self,
-        embedding_dim: int,
-        conditioning_embedding_dim: int,
-        elementwise_affine: bool = False,
-        eps: float = 1e-6,
-        bias: bool = True,
-        quant_config: "QuantizationConfig | None" = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.silu = nn.SiLU()
-        self.linear = ReplicatedLinear(
-            conditioning_embedding_dim,
-            embedding_dim * 2,
-            bias=bias,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.linear",
-        )
-        self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine)
-
-    def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor:
-        emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
-        if isinstance(emb, tuple):
-            emb = emb[0]
-        scale, shift = torch.chunk(emb, 2, dim=1)
-        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
-        return x
@@ -7,7 +7,7 @@ import re
 import time
 from collections.abc import Generator, Iterable
 from pathlib import Path
-from typing import TYPE_CHECKING, cast
+from typing import cast

 import torch
 from huggingface_hub import hf_hub_download
@@ -34,9 +34,6 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
 from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
 from vllm_omni.diffusion.registry import initialize_model

-if TYPE_CHECKING:
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
-
 logger = init_logger(__name__)


@@ -335,60 +332,12 @@ class DiffusersPipelineLoader:
            weights_scale_not_loaded = {name for name in weights_not_loaded if name.endswith("weight_scale")}
            weights_not_loaded = weights_not_loaded - weights_scale_not_loaded
            if weights_not_loaded:
-                self._check_unloaded_weights(weights_not_loaded)
+                raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}")
            if weights_scale_not_loaded:
                logger.warning(
                    f"Following weight_scale weights were not initialized from checkpoint: {weights_scale_not_loaded}"
                )

-    @staticmethod
-    def _is_expected_quantized_weight(name: str) -> bool:
-        """Return True if *name* is a quantization-specific parameter.
-
-        Quantization methods (GPTQ, AWQ, FP8, GGUF, Autoround, etc.) create extra
-        parameters that have no counterpart in an unquantized checkpoint.
-        These are expected to be absent and should not trigger a load error.
-        """
-        # Weight suffixes that quantization methods register in the model but
-        # are not present in unquantized checkpoints.
-        _QUANTIZED_WEIGHT_SUFFIXES = (
-            # GPTQ / AWQ / AutoRound – g_idx is optional (not all checkpoints include it)
-            ".g_idx",
-            # FP8
-            ".weight_scale",
-            ".weight_scale_inv",
-            ".input_scale",
-            # GGUF
-            ".qweight_type",
-            # INT8  (weight_scale already covered above)
-        )
-        return name.endswith(_QUANTIZED_WEIGHT_SUFFIXES)
-
-    def _check_unloaded_weights(
-        self,
-        weights_not_loaded: set[str],
-    ) -> None:
-        """Validate unloaded weights, tolerating expected quantization artifacts.
-
-        For quantized models, weights matching known quant-specific suffixes
-        are logged as a warning.  Any *other* missing weight raises
-        ``ValueError`` regardless of quantization.
-        """
-        od_config = getattr(self, "od_config", None)
-        if od_config is None or od_config.quantization_config is None:
-            raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}")
-
-        expected_missing = {w for w in weights_not_loaded if self._is_expected_quantized_weight(w)}
-        unexpected_missing = weights_not_loaded - expected_missing
-
-        if expected_missing:
-            logger.warning(
-                "Following weights were not initialized from checkpoint (expected for quantized models): %s",
-                expected_missing,
-            )
-        if unexpected_missing:
-            raise ValueError(f"Following weights were not initialized from checkpoint: {unexpected_missing}")
-
    def _is_gguf_quantization(self, od_config: OmniDiffusionConfig) -> bool:
        quant_config = od_config.quantization_config
        if quant_config is None:
@@ -326,18 +326,11 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
        cfg_text_scale = extra_args.get("cfg_text_scale", 4.0)
        cfg_img_scale = extra_args.get("cfg_img_scale", 1.5)

-        cfg_interval = extra_args.get("cfg_interval", (0.4, 1.0))
-        cfg_renorm_type = extra_args.get("cfg_renorm_type", "global")
-        cfg_renorm_min = extra_args.get("cfg_renorm_min", 0.0)
-
        gen_params = BagelGenParams(
            num_timesteps=int(req.sampling_params.num_inference_steps or 50),
            timestep_shift=3.0,
            cfg_text_scale=cfg_text_scale,
            cfg_img_scale=cfg_img_scale,
-            cfg_interval=cfg_interval,
-            cfg_renorm_type=cfg_renorm_type,
-            cfg_renorm_min=cfg_renorm_min,
        )

        gen_context = {
@@ -1,6 +1,5 @@
 import torch
 import torch.nn as nn
-from vllm.logger import init_logger

 from vllm_omni.diffusion.attention.layer import Attention

@@ -12,8 +11,6 @@ except ImportError:
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.models.dreamid_omni.wan2_2 import WanModel, rope_apply

-logger = init_logger(__name__)
-

 class FusionModel(nn.Module):
    def __init__(self, video_config=None, audio_config=None):
@@ -25,14 +22,14 @@ class FusionModel(nn.Module):
        else:
            has_video = False
            self.video_model = None
-            logger.warning("No video model is provided!")
+            print("Warning: No video model is provided!")

        if audio_config is not None:
            self.audio_model = WanModel(**audio_config)
        else:
            has_audio = False
            self.audio_model = None
-            logger.warning("No audio model is provided!")
+            print("Warning: No audio model is provided!")

        if has_video and has_audio:
            assert len(self.video_model.blocks) == len(self.audio_model.blocks)
@@ -12,6 +12,7 @@ from diffusers.models.embeddings import (
    get_1d_rotary_pos_embed,
 )
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
 from diffusers.utils import is_torch_npu_available
 from vllm.distributed import get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather
 from vllm.logger import init_logger
@@ -31,11 +32,6 @@ if TYPE_CHECKING:

 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.layers.adalayernorm import (
-    AdaLayerNormContinuous,
-    AdaLayerNormZero,
-    AdaLayerNormZeroSingle,
-)
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk

 logger = init_logger(__name__)
@@ -50,7 +46,6 @@ class ColumnParallelApproxGELU(nn.Module):
        approximate: str,
        bias: bool = True,
        quant_config: "QuantizationConfig | None" = None,
-        prefix: str = "",
    ):
        super().__init__()
        self.proj = ColumnParallelLinear(
@@ -60,7 +55,6 @@ class ColumnParallelApproxGELU(nn.Module):
            gather_output=False,
            return_bias=False,
            quant_config=quant_config,
-            prefix=f"{prefix}.proj",
        )
        self.approximate = approximate

@@ -79,7 +73,6 @@ class FeedForward(nn.Module):
        inner_dim: int | None = None,
        bias: bool = True,
        quant_config: "QuantizationConfig | None" = None,
-        prefix: str = "",
    ) -> None:
        super().__init__()

@@ -89,9 +82,7 @@ class FeedForward(nn.Module):
        dim_out = dim_out or dim

        layers: list[nn.Module] = [
-            ColumnParallelApproxGELU(
-                dim, inner_dim, approximate="tanh", bias=bias, quant_config=quant_config, prefix=f"{prefix}.net.0"
-            ),
+            ColumnParallelApproxGELU(dim, inner_dim, approximate="tanh", bias=bias, quant_config=quant_config),
            nn.Identity(),  # placeholder for weight loading
            RowParallelLinear(
                inner_dim,
@@ -99,7 +90,6 @@ class FeedForward(nn.Module):
                input_is_parallel=True,
                return_bias=False,
                quant_config=quant_config,
-                prefix=f"{prefix}.net.2",
            ),
        ]

@@ -127,7 +117,6 @@ class FluxAttention(torch.nn.Module):
        context_pre_only: bool | None = None,
        pre_only: bool = False,
        quant_config: "QuantizationConfig | None" = None,
-        prefix: str = "",
    ):
        super().__init__()

@@ -152,7 +141,6 @@ class FluxAttention(torch.nn.Module):
            total_num_heads=self.heads,
            bias=bias,
            quant_config=quant_config,
-            prefix=f"{prefix}.to_qkv",
        )

        if not self.pre_only:
@@ -165,7 +153,6 @@ class FluxAttention(torch.nn.Module):
                        input_is_parallel=True,
                        return_bias=False,
                        quant_config=quant_config,
-                        prefix=f"{prefix}.to_out.0",
                    ),
                    nn.Dropout(dropout),
                ]
@@ -181,8 +168,8 @@ class FluxAttention(torch.nn.Module):
                total_num_heads=self.heads,
                bias=added_proj_bias,
                quant_config=quant_config,
-                prefix=f"{prefix}.add_kv_proj",
            )
+
            self.to_add_out = RowParallelLinear(
                self.inner_dim,
                query_dim,
@@ -190,7 +177,6 @@ class FluxAttention(torch.nn.Module):
                input_is_parallel=True,
                return_bias=False,
                quant_config=quant_config,
-                prefix=f"{prefix}.to_add_out",
            )

        self.rope = RotaryEmbedding(is_neox_style=False)
@@ -286,11 +272,11 @@ class FluxTransformerBlock(nn.Module):
        qk_norm: str = "rms_norm",
        eps: float = 1e-6,
        quant_config: "QuantizationConfig | None" = None,
-        prefix: str = "",
    ):
        super().__init__()
-        self.norm1 = AdaLayerNormZero(dim, quant_config=quant_config, prefix=f"{prefix}.norm1")
-        self.norm1_context = AdaLayerNormZero(dim, quant_config=quant_config, prefix=f"{prefix}.norm1_context")
+
+        self.norm1 = AdaLayerNormZero(dim)
+        self.norm1_context = AdaLayerNormZero(dim)

        self.attn = FluxAttention(
            query_dim=dim,
@@ -302,14 +288,13 @@ class FluxTransformerBlock(nn.Module):
            bias=True,
            eps=eps,
            quant_config=quant_config,
-            prefix=f"{prefix}.attn",
        )

        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.ff = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config, prefix=f"{prefix}.ff")
+        self.ff = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config)

        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.ff_context = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config, prefix=f"{prefix}.ff_context")
+        self.ff_context = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config)

    def forward(
        self,
@@ -376,28 +361,17 @@ class FluxSingleTransformerBlock(nn.Module):
        attention_head_dim: int,
        mlp_ratio: float = 4.0,
        quant_config: "QuantizationConfig | None" = None,
-        prefix: str = "",
    ):
        super().__init__()
        self.mlp_hidden_dim = int(dim * mlp_ratio)

-        self.norm = AdaLayerNormZeroSingle(dim, quant_config=quant_config, prefix=f"{prefix}.norm")
+        self.norm = AdaLayerNormZeroSingle(dim)
        self.proj_mlp = ReplicatedLinear(
-            dim,
-            self.mlp_hidden_dim,
-            bias=True,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.proj_mlp",
+            dim, self.mlp_hidden_dim, bias=True, return_bias=False, quant_config=quant_config
        )
        self.act_mlp = nn.GELU(approximate="tanh")
        self.proj_out = ReplicatedLinear(
-            dim + self.mlp_hidden_dim,
-            dim,
-            bias=True,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.proj_out",
+            dim + self.mlp_hidden_dim, dim, bias=True, return_bias=False, quant_config=quant_config
        )

        self.attn = FluxAttention(
@@ -409,7 +383,6 @@ class FluxSingleTransformerBlock(nn.Module):
            eps=1e-6,
            pre_only=True,
            quant_config=quant_config,
-            prefix=f"{prefix}.attn",
        )

    def forward(
@@ -569,9 +542,8 @@ class FluxTransformer2DModel(nn.Module):
                    num_attention_heads=num_attention_heads,
                    attention_head_dim=attention_head_dim,
                    quant_config=quant_config,
-                    prefix=f"transformer_blocks.{i}",
                )
-                for i in range(num_layers)
+                for _ in range(num_layers)
            ]
        )

@@ -582,20 +554,12 @@ class FluxTransformer2DModel(nn.Module):
                    num_attention_heads=num_attention_heads,
                    attention_head_dim=attention_head_dim,
                    quant_config=quant_config,
-                    prefix=f"single_transformer_blocks.{i}",
                )
-                for i in range(num_single_layers)
+                for _ in range(num_single_layers)
            ]
        )

-        self.norm_out = AdaLayerNormContinuous(
-            self.inner_dim,
-            self.inner_dim,
-            elementwise_affine=False,
-            eps=1e-6,
-            quant_config=quant_config,
-            prefix="norm_out",
-        )
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)

    def forward(
@@ -578,11 +578,9 @@ class Flux2Transformer2DModel(nn.Module):
        guidance_embeds: bool = True,
    ):
        super().__init__()
-        self.guidance_embeds = guidance_embeds
        self.stacked_params_mapping = None
        self.out_channels = out_channels or in_channels
        self.inner_dim = num_attention_heads * attention_head_dim
-
        self.config = SimpleNamespace(
            patch_size=patch_size,
            in_channels=in_channels,
@@ -928,7 +928,6 @@ class Flux2Pipeline(nn.Module, SupportImageInput):
        self._attention_kwargs = attention_kwargs
        self._current_timestep = None
        self._interrupt = False
-        guidance_tensor = None

        # 2. Define call parameters
        if prompt is not None and isinstance(prompt, str):
@@ -1018,11 +1017,6 @@ class Flux2Pipeline(nn.Module, SupportImageInput):
        )
        self._num_timesteps = len(timesteps)

-        # handle guidance
-        if self.transformer.guidance_embeds is not None:
-            guidance_tensor = torch.full([1], self.guidance_scale, device=device, dtype=torch.float32)
-            guidance_tensor = guidance_tensor.expand(latents.shape[0])
-
        # 7. Denoising loop
        # 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
@@ -1044,7 +1038,7 @@ class Flux2Pipeline(nn.Module, SupportImageInput):
            noise_pred = self.transformer(
                hidden_states=latent_model_input,  # (B, image_seq_len, C)
                timestep=timestep / 1000,
-                guidance=guidance_tensor,
+                guidance=None,
                encoder_hidden_states=prompt_embeds,
                txt_ids=text_ids,  # B, text_seq_len, 4
                img_ids=latent_image_ids,  # B, image_seq_len, 4
@@ -22,7 +22,6 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.cache.base import CachedTransformer
 from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module

 logger = init_logger(__name__)

@@ -725,8 +724,6 @@ class GlmImageTransformer2DModel(CachedTransformer):

    _repeated_blocks = ["GlmImageTransformerBlock"]

-    _hsdp_shard_conditions = [is_transformer_block_module]
-
    def __init__(
        self,
        od_config: OmniDiffusionConfig,
@@ -11,12 +11,9 @@ import torch
 import torch.nn.functional as F
 from diffusers.utils.outputs import BaseOutput
 from transformers import AutoTokenizer
-from vllm.logger import init_logger

 from .hunyuan_image_3_transformer import ImageInfo, JointImageInfo, default

-logger = init_logger(__name__)
-

 class TokenizerEncodeOutput(BaseOutput):
    tokens: torch.Tensor | None = None
@@ -124,7 +121,7 @@ class TokenizerWrapper:
        elif isinstance(uncond_enabled, bool):
            uncond_enabled = [uncond_enabled] * len(texts)
        if len(uncond_enabled) != len(texts):
-            logger.debug("uncond_enabled=%s, texts=%s", uncond_enabled, texts)
+            print(uncond_enabled, texts)
        assert len(uncond_enabled) == len(texts), (
            f"Length of uncond_flags should be equal to the number of texts, "
            f"but got {len(uncond_enabled)} and {len(texts)}."
@@ -2036,7 +2036,7 @@ class HunyuanImage3Model(nn.Module):
        for name, loaded_weight in weights:
            # print(f"Loading weight name: {name}, tp_rank: {tp_rank}", flush=True)
            if contains_unexpected_keyword(name, unexpected_keywords):
-                logger.warning("Skipping unexpected weight name: %s", name)
+                print(f"Skipping unexpected weight name: {name}")
                continue
            if "rotary_emb.inv_freq" in name:
                continue
Author	SHA1	Message	Date
Roger Wang	c1f3dac1a3	add	2026-04-01 04:35:27 -07:00
Roger Wang	1f8fc437b1	add	2026-04-01 04:26:09 -07:00
Roger Wang	a14336d9c8	add	2026-04-01 04:20:25 -07:00