Merge branch 'main' into dev/rebase-v0170

Signed-off-by: Ricardo Noriega De Soto <rnoriega@redhat.com>

.buildkite/pipeline.yml

@@ -3,7 +3,7 @@ steps:
     key: image-build
     commands:
       - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-      - "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
       - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
       - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
     agents:

.buildkite/release-pipeline.yaml

@@ -0,0 +1,38 @@
+steps:
+  - label: ":docker: Build image"
+    key: image-build
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L2 Test
+  - label: "Upload Ready Pipeline"
+    depends_on: image-build
+    key: upload-ready-pipeline
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-ready.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L3 Test
+  - label: "Upload Merge Pipeline"
+    depends_on: image-build
+    key: upload-merge-pipeline
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-merge.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L4 Test
+  - label: "Upload Nightly Pipeline"
+    depends_on: image-build
+    key: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-nightly.yml
+    agents:
+      queue: "cpu_queue_premerge"

.buildkite/test-amd.yaml

@@ -33,6 +33,7 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+    - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
 
 - label: "Diffusion Cache Backend Test"
   timeout_in_minutes: 15
@@ -88,6 +89,7 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    - pytest -s -v tests/engine/test_async_omni_engine_abort.py
 
 - label: "Omni Model Test Qwen3-Omni"
   timeout_in_minutes: 15
@@ -101,7 +103,7 @@ steps:
     - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
     - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
-    - pytest -s -v tests/engine/test_async_omni_engine_abort.py
+
 
 - label: "Diffusion Image Edit Test"
   timeout_in_minutes: 15

.buildkite/test-merge.yml

@@ -24,7 +24,7 @@ steps:
       queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
       - docker#v5.2.0:
-          image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
           always-pull: true
           propagate-environment: true
           environment:
@@ -155,41 +155,31 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Benchmark & Engine Test with H100"
-    timeout_in_minutes: 15
+  - label: "Benchmark & Engine Test"
     depends_on: upload-merge-pipeline
     commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/benchmarks/test_serve_cli.py
-      - pytest -s -v tests/engine/test_async_omni_engine_abort.py
+      - |
+        timeout 15m bash -c '
+                export VLLM_WORKER_MULTIPROC_METHOD=spawn
+                set +e
+                pytest -s -v tests/benchmarks/test_serve_cli.py
+                EXIT1=$$?
+                pytest -s -v tests/engine/test_async_omni_engine_abort.py
+                EXIT2=$$?
+                exit $$((EXIT1 | EXIT2))
+        '
     agents:
-      queue: "mithril-h100-pool"
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     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: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-            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
+      - 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: "Omni Model Test"
     timeout_in_minutes: 15

.buildkite/test-nightly.yml

@@ -1,19 +1,24 @@
 steps:
-  - label: "Omni Model Test with H100"
+  - label: ":full_moon: Omni Model Test with H100"
     timeout_in_minutes: 90
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "advanced_model" --run-level "advanced_model"
-      - pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --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=$$?
+        exit $$((EXIT1 | EXIT2))
     agents:
       queue: "mithril-h100-pool"
     plugins:
       - kubernetes:
           podSpec:
             containers:
-              - image: public.ecr.aws/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:$BUILDKITE_COMMIT
                 resources:
                   limits:
                     nvidia.com/gpu: 2
@@ -36,7 +41,7 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: "Omni Model Test"
+  - label: ":full_moon: Omni Model Test"
     timeout_in_minutes: 60
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1"
@@ -51,25 +56,51 @@ steps:
           always-pull: true
           shm-size: "8gb"
           propagate-environment: true
+          shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Omni Model Perf Test"
-    timeout_in_minutes: 120
+  - label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
+    timeout_in_minutes: 30
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - |
+        huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
+        export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk
+    agents:
+      queue: "gpu_4_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: ":full_moon: Omni Model Perf Test with H100"
+    key: nightly-performance
+    timeout_in_minutes: 180
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export BENCHMARK_DIR=tests
       - pytest -s -v tests/perf/scripts/run_benchmark.py
+      - buildkite-agent artifact upload "tests/*.json"
     agents:
       queue: "mithril-h100-pool"
     plugins:
       - kubernetes:
           podSpec:
             containers:
-              - image: public.ecr.aws/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:$BUILDKITE_COMMIT
                 resources:
                   limits:
                     nvidia.com/gpu: 2
@@ -91,3 +122,18 @@ steps:
                 hostPath:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
+
+  - label: ":email: Nightly Perf Collection & Email"
+    key: nightly-perf-distribution
+    depends_on: nightly-performance
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - pip install openpyxl
+      - export DEFAULT_INPUT_DIR=tests
+      - export DEFAULT_OUTPUT_DIR=tests
+      - buildkite-agent artifact download "tests/*.json" . --step nightly-performance
+      - python tools/nightly/generate_nightly_perf_excel.py
+      - python tools/nightly/send_nightly_perf_email.py
+      - buildkite-agent artifact upload "tests/*.xlsx"
+    agents:
+      queue: "cpu_queue_premerge"

.buildkite/test-ready.yml

@@ -2,7 +2,7 @@ steps:
   - label: "Simple Unit Test"
     depends_on: upload-ready-pipeline
     commands:
-      - "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+      - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
     agents:
       queue: "gpu_1_queue"
     plugins:
@@ -16,12 +16,11 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Model Test"
-    timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
     commands:
-      - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
     agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+      queue: "gpu_1_queue"
     plugins:
       - docker#v5.2.0:
           image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -33,11 +32,17 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Model CPU offloading Test"
-    timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
     commands:
-      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
-      - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+      - |
+        timeout 20m bash -c '
+          set +e
+          pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+          EXIT1=$$?
+          pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+          EXIT2=$$?
+          exit $$((EXIT1 | EXIT2))
+        '
     agents:
       queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
@@ -51,10 +56,9 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Audio Generation Model Test"
-    timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
     commands:
-      - pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
     agents:
       queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
@@ -68,10 +72,9 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Cache Backend Test"
-    timeout_in_minutes: 15
     depends_on: upload-ready-pipeline
     commands:
-      - pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
+      - timeout 15m pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
     agents:
       queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
@@ -85,10 +88,9 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Sequence Parallelism Test"
-    timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
     commands:
-      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
@@ -103,10 +105,9 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion GPU Worker Test"
-    timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
     commands:
-      - pytest -s -v tests/diffusion/test_diffusion_worker.py
+      - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
@@ -120,49 +121,43 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  # - label: "Benchmark & Engine Test with H100"
-  #   timeout_in_minutes: 15
-  #   depends_on: upload-ready-pipeline
-  #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - pytest -s -v tests/benchmarks/test_serve_cli.py
-  #     - pytest -s -v tests/engine/test_async_omni_engine_abort.py
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #               resources:
-  #                 limits:
-  #                   nvidia.com/gpu: 2
-  #               volumeMounts:
-  #                 - name: devshm
-  #                   mountPath: /dev/shm
-  #                 - name: hf-cache
-  #                   mountPath: /root/.cache/huggingface
-  #               env:
-  #                 - name: HF_HOME
-  #                   value: /root/.cache/huggingface
-  #           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
 
-  - label: "Omni Model Test"
-    timeout_in_minutes: 17
+  - label: "Benchmark & Engine Test"
     depends_on: upload-ready-pipeline
     commands:
-      - export VLLM_LOGGING_LEVEL=DEBUG
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+      - |
+        timeout 15m bash -c '
+                export VLLM_WORKER_MULTIPROC_METHOD=spawn
+                set +e
+                pytest -s -v tests/benchmarks/test_serve_cli.py
+                EXIT1=$$?
+                pytest -s -v tests/engine/test_async_omni_engine_abort.py
+                EXIT2=$$?
+                exit $$((EXIT1 | EXIT2))
+        '
+    agents:
+      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
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+
+  - label: "Omni Model Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 17m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+        '
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
@@ -176,20 +171,22 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   # - label: "Omni Model Test with H100"
-  #   timeout_in_minutes: 20
   #   depends_on: upload-ready-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
-  #     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
+  #     - |
+  #       timeout 20m bash -c '
+  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #         export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+  #         # - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+  #         pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
+  #       '
   #   agents:
   #     queue: "mithril-h100-pool"
   #   plugins:
   #     - kubernetes:
   #         podSpec:
   #           containers:
-  #             - image: public.ecr.aws/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:$BUILDKITE_COMMIT
   #               resources:
   #                 limits:
   #                   nvidia.com/gpu: 2
@@ -213,12 +210,15 @@ steps:
   #                 type: DirectoryOrCreate
 
   - label: "Qwen3-TTS E2E Test"
-    timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
     commands:
-      - export VLLM_LOGGING_LEVEL=DEBUG
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
+      - |
+        timeout 20m bash -c '
+          huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py -k "not NoAsyncChunk"
+        '
     agents:
       queue: "gpu_4_queue"
     plugins:
@@ -233,18 +233,20 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   # - label: "Diffusion Image Edit Test with H100 (1 GPU)"
-  #   timeout_in_minutes: 20
   #   depends_on: upload-ready-pipeline
   #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+  #     - |
+  #       timeout 20m bash -c '
+  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #         pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+  #       '
   #   agents:
   #     queue: "mithril-h100-pool"
   #   plugins:
   #     - kubernetes:
   #         podSpec:
   #           containers:
-  #             - image: public.ecr.aws/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:$BUILDKITE_COMMIT
   #               resources:
   #                 limits:
   #                   nvidia.com/gpu: 1
@@ -268,18 +270,20 @@ steps:
   #                 type: DirectoryOrCreate
 
   # - label: "Bagel Text2Img Model Test with H100"
-  #   timeout_in_minutes: 30
   #   depends_on: upload-ready-pipeline
   #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+  #     - |
+  #       timeout 30m bash -c '
+  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #         pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+  #       '
   #   agents:
   #     queue: "mithril-h100-pool"
   #   plugins:
   #     - kubernetes:
   #         podSpec:
   #           containers:
-  #             - image: public.ecr.aws/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:$BUILDKITE_COMMIT
   #               resources:
   #                 limits:
   #                   nvidia.com/gpu: 1

.github/PULL_REQUEST_TEMPLATE.md

@@ -12,10 +12,10 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
 <summary> Essential Elements of an Effective PR Description Checklist </summary>
 
 - [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
-- [ ] The test plan. Please providing the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
-- [ ] The test results. Please pasting the results comparison before and after, or e2e results.
+- [ ] The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
+- [ ] The test results. Please paste the results comparison before and after, or the e2e results.
 - [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.**
-- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft.
+- [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
 </details>
 
 **BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)

.github/workflows/comfyui-publish.yml

@@ -1,27 +0,0 @@
-name: ComfyUI Integration -- Publish to Comfy registry
-on:
-  workflow_dispatch:
-  push:
-    tags:
-      - "*"
-    paths:
-      - "apps/ComfyUI-vLLM-Omni/**"
-
-permissions:
-  issues: write
-
-defaults:
-  run:
-    working-directory: apps/ComfyUI-vLLM-Omni
-
-jobs:
-  publish-node:
-    name: Publish Custom Node to registry
-    runs-on: ubuntu-latest
-    steps:
-      - name: ♻️ Check out code
-        uses: actions/checkout@v4
-      - name: 📦 Publish Custom Node
-        uses: Comfy-Org/publish-node-action@main
-        with:
-          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.github/workflows/comfyui-validate.yml

@@ -1,40 +0,0 @@
-name: ComfyUI Integration -- Validate backwards compatibility
-
-on:
-  pull_request:
-    branches:
-      - master
-      - main
-    paths:
-      - "apps/ComfyUI-vLLM-Omni/**"
-
-defaults:
-  run:
-    working-directory: apps/ComfyUI-vLLM-Omni
-
-jobs:
-  check-base-path:
-    runs-on: ubuntu-latest
-    outputs:
-      exists: ${{ steps.check.outputs.exists }}
-    defaults:
-      run:
-        working-directory: .
-    steps:
-      - uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event.pull_request.base.sha }}
-          fetch-depth: 1
-      - id: check
-        run: |
-          if [ -d "apps/ComfyUI-vLLM-Omni" ]; then
-            echo "exists=true" >> "$GITHUB_OUTPUT"
-          else
-            echo "exists=false" >> "$GITHUB_OUTPUT"
-          fi
-  validate:
-    needs: check-base-path
-    if: needs.check-base-path.outputs.exists == 'true'
-    runs-on: ubuntu-latest
-    steps:
-      - uses: comfy-org/node-diff@main

.gitignore

@@ -243,5 +243,7 @@ Dockerfile.dev
 discussion
 tmp_test
 
+# Auto-generated version file (created by setuptools_scm during build)
+vllm_omni/_version.py
 # output files
 *.wav

.pre-commit-config.yaml

@@ -8,6 +8,7 @@ repos:
       # list of supported hooks: https://pre-commit.com/hooks.html
       - id: check-yaml
         args: ["--unsafe"]
+        exclude: "examples/online_serving/chart-helm/templates/"
       - id: debug-statements
       - id: end-of-file-fixer
       - id: mixed-line-ending

README.md

@@ -1,6 +1,6 @@
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
+    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
     <img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
   </picture>
 </p>
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://drive.google.com/file/d/12k_DOo2xoEROvd4xQxYNuEi6C2GECoRx/view?usp=drive_link"><b>Slides</b></a> |
 </p>
 
 
@@ -17,6 +17,8 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 *Latest News* 🔥
 
+- [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
+- [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
 - [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
 - [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.

apps/ComfyUI-vLLM-Omni/README.md

@@ -47,12 +47,12 @@ If no, check your shell running the ComfyUI process. There may be some error mes
 This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
 
 - **Generate Image** for text-to-image and image-to-image tasks
-- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
+- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
 - **TTS** and **TTS Voice Clone** for TTS tasks
 
 This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
 
-> [!INFO]
+> [!NOTE]
 > The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support.
 
 To build a simple workflow yourself,
@@ -65,28 +65,16 @@ To build a simple workflow yourself,
     - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
     - For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages.
 
-**The following features are tested**:
-
-- Single-node workflows for
-    - Multimodal Comprehension (e.g., Qwen Omni, BAGEL)
-    - Text-to-Image Generation (e.g., Qwen-Image)
-    - Image-to-Image Generation (e.g., Qwen-Image-Edit)
-    - TTS (e.g., Qwen TTS, including VoiceDesign, VoiceClone, CustomVoice)
-
-**The following features are not currently tested**. They may work or break. You are welcomed to test it out and offer comments.
-
-- Multi-node workflow that connects multiple model services together.
-
 ## Screenshots and Examples
 
-### Multimodal comprehension (e.g., Qwen Omni series, BAGEL)
+### Multimodal understanding (e.g., Qwen Omni series, BAGEL)
 
-(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Annotated Example**)
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Multimodal Understanding**)
 
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg">
-    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg" width=55%>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-understanding.jpg">
+    <img alt="vLLM-Omni multimodal understanding" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-understanding.jpg" width=55%>
   </picture>
 </p>
 
@@ -98,7 +86,7 @@ You can configure per-stage sampling parameters for multi-stage models.
 <p align="center">
   <picture>
     <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg">
-    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg" width=55%>
+    <img alt="vLLM-Omni multiple stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg" width=55%>
   </picture>
 </p>
 
@@ -109,7 +97,7 @@ You can configure per-stage sampling parameters for multi-stage models.
 <p align="center">
   <picture>
     <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg">
-    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg" width=55%>
+    <img alt="vLLM-Omni image generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg" width=55%>
   </picture>
 </p>
 
@@ -123,13 +111,24 @@ You can configure per-stage sampling parameters for multi-stage models.
 <p align="center">
   <picture>
     <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg">
-    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg" width=55%>
+    <img alt="vLLM-Omni TTS" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg" width=55%>
   </picture>
 </p>
 
 > [!TIP]
 > There is a dedicated node for VoiceClone tasks with reference audio input. Other simple text-to-speech tasks should use the regular TTS node.
 
+### Chaining multiple model services
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Chaining Services**)
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-chaining-services.jpg">
+    <img alt="vLLM-Omni TTS" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-chaining-services.jpg" width=55%>
+  </picture>
+</p>
+
 ## Develop
 
 Follow the [development convention and rules of vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/).

apps/ComfyUI-vLLM-Omni/__init__.py

@@ -6,18 +6,18 @@ __all__ = [
     "WEB_DIRECTORY",
 ]
 
-__author__ = """Zeyu Huang"""
-__email__ = "11222265+fhfuih@users.noreply.github.com"
+__author__ = """vLLM-Omni Team"""
+__email__ = "vllm-omni@vllm.ai"
 __version__ = "0.0.1"
 
 from .comfyui_vllm_omni.nodes import (
     VLLMOmniARSampling,
-    VLLMOmniComprehension,
     VLLMOmniDiffusionSampling,
     VLLMOmniGenerateImage,
     VLLMOmniQwenTTSParams,
     VLLMOmniSamplingParamsList,
     VLLMOmniTTS,
+    VLLMOmniUnderstanding,
     VLLMOmniVoiceClone,
 )
 
@@ -25,7 +25,7 @@ from .comfyui_vllm_omni.nodes import (
 NODE_CLASS_MAPPINGS = {
     # === Generation ===
     "VLLMOmniGenerateImage": VLLMOmniGenerateImage,
-    "VLLMOmniComprehension": VLLMOmniComprehension,
+    "VLLMOmniUnderstanding": VLLMOmniUnderstanding,
     "VLLMOmniTTS": VLLMOmniTTS,
     "VLLMOmniVoiceClone": VLLMOmniVoiceClone,
     # === Params ===
@@ -39,7 +39,7 @@ NODE_CLASS_MAPPINGS = {
 NODE_DISPLAY_NAME_MAPPINGS = {
     # === Generation ===
     "VLLMOmniGenerateImage": "Generate Image",
-    "VLLMOmniComprehension": "Multimodality Comprehension",
+    "VLLMOmniUnderstanding": "Multimodality Understanding",
     "VLLMOmniTTS": "TTS (Text to Speech)",
     "VLLMOmniVoiceClone": "TTS Voice Cloning",
     # === Params ===

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py

@@ -1,4 +1,4 @@
-from typing import Literal, cast
+from typing import Literal
 
 import torch
 from comfy_api.input import AudioInput, VideoInput
@@ -6,7 +6,12 @@ from comfy_api.input import AudioInput, VideoInput
 from .utils.api_client import VLLMOmniClient
 from .utils.logger import get_logger
 from .utils.models import lookup_model_spec
-from .utils.types import AudioFormat
+from .utils.types import (
+    AudioFormat,
+    AutoregressionSamplingParams,
+    DiffusionSamplingParams,
+    QwenTTSModelSpecificParams,
+)
 from .utils.validators import (
     add_sampling_parameters_to_stage,
     validate_model_and_sampling_params_types,
@@ -86,7 +91,10 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
 
         # Prefer DALL-E compatible API for simple (one-stage) diffusion models
         if (spec is None or spec["stages"] == ["diffusion"]) and not is_bagel:
-            sampling_params = cast(dict | None, sampling_params)
+            # The number of sampling parameter groups should have been validated.
+            # Now, simply convert single-item list to dict.
+            if isinstance(sampling_params, list):
+                sampling_params = sampling_params[0]
             if audio is None and image is None and video is None:
                 # No multimodal input --- use DALL-E image generation
                 logger.info("Using DALL-E image generation endpoint")
@@ -133,7 +141,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
         return (output,)
 
 
-class VLLMOmniComprehension(_VLLMOmniGenerateBase):
+class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
     @classmethod
     def INPUT_TYPES(cls):
         return {
@@ -197,7 +205,7 @@ class VLLMOmniComprehension(_VLLMOmniGenerateBase):
             (
                 text_response,
                 _,
-            ) = await client.generate_comprehension_chat_completion(
+            ) = await client.generate_understanding_chat_completion(
                 model=model,
                 prompt=prompt,
                 image=image,
@@ -221,7 +229,7 @@ class VLLMOmniComprehension(_VLLMOmniGenerateBase):
             (
                 text_response,
                 audio,
-            ) = await client.generate_comprehension_chat_completion(
+            ) = await client.generate_understanding_chat_completion(
                 model=model,
                 prompt=prompt,
                 image=image,
@@ -287,15 +295,13 @@ class VLLMOmniTTS(_VLLMOmniGenerateBase):
         logger.info("Got extra kwargs in TTS: %s", kwargs)
 
         is_qwen_tts = "qwen3-tts" in model.lower()
-        extra_params_type = None if model_specific_params is None else model_specific_params["type"]
-        if not is_qwen_tts and extra_params_type == "qwen-tts":
+        if not is_qwen_tts and isinstance(model_specific_params, QwenTTSModelSpecificParams):
             raise ValueError(
                 "You have provided Qwen-specific TTS params."
                 "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
             )
 
         combined_params = {**kwargs, **(model_specific_params or {})}
-        combined_params.pop("type", None)  # Internal fields in model_specific_params
 
         client = VLLMOmniClient(url)
 
@@ -352,8 +358,7 @@ class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
         **kwargs,
     ):
         is_qwen_tts = "qwen3-tts" in model.lower()
-        extra_params_type = None if model_specific_params is None else model_specific_params["type"]
-        if not is_qwen_tts and extra_params_type == "qwen-tts":
+        if not is_qwen_tts and isinstance(model_specific_params, QwenTTSModelSpecificParams):
             raise ValueError(
                 "You have provided Qwen-specific TTS params."
                 "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
@@ -366,7 +371,6 @@ class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
             **kwargs,
             **(model_specific_params or {}),
         }
-        combined_params.pop("type", None)  # Internal fields in model_specific_params
 
         client = VLLMOmniClient(url)
 
@@ -419,10 +423,7 @@ class VLLMOmniARSampling:
     CATEGORY = "vLLM-Omni/Sampling Params"
 
     def get_params(self, seed, **kwargs):
-        params = {
-            "type": "autoregression",  # for internal use, removed before sending the request
-            **kwargs,
-        }
+        params = AutoregressionSamplingParams(kwargs)
         if seed >= 0:
             params["seed"] = seed
         return (params,)
@@ -479,6 +480,13 @@ class VLLMOmniDiffusionSampling:
                         "tooltip": "Enable VAE slicing for reduced memory usage (slight quality trade-off)",
                     },
                 ),
+                "vae_use_tiling": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": "Enable VAE tiling for reduced memory usage (slight quality trade-off)",
+                    },
+                ),
                 # === Put seed at last. ===
                 # Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate"
                 "seed": (
@@ -499,10 +507,7 @@ class VLLMOmniDiffusionSampling:
     CATEGORY = "vLLM-Omni/Sampling Params"
 
     def get_params(self, seed, **kwargs):
-        params = {
-            "type": "diffusion",  # for internal use, removed before sending the request
-            **kwargs,
-        }
+        params = DiffusionSamplingParams(kwargs)
         if seed >= 0:
             params["seed"] = seed
         return (params,)
@@ -566,4 +571,4 @@ class VLLMOmniQwenTTSParams:
     CATEGORY = "vLLM-Omni/TTS Params"
 
     def get_params(self, **kwargs):
-        return ({"type": "qwen-tts", **kwargs},)
+        return (QwenTTSModelSpecificParams(kwargs),)

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py

@@ -57,16 +57,7 @@ class VLLMOmniClient:
         if negative_prompt:
             payload["negative_prompt"] = negative_prompt
         if sampling_params is not None:
-            # Only select specific sampling params
-            for k in (
-                "n",
-                "num_inference_steps",
-                "guidance_scale",
-                "true_cfg_scale",
-                "vae_use_slicing",
-            ):
-                if k in sampling_params and sampling_params[k] is not None:
-                    payload[k] = sampling_params[k]
+            payload.update(sampling_params)
         logger.debug("img gen payload: %s", payload)
 
         url = self.base_url + "/images/generations"
@@ -138,10 +129,8 @@ class VLLMOmniClient:
         if negative_prompt:
             form.add_field("negative_prompt", negative_prompt)
         if sampling_params is not None:
-            # Only select specific sampling params
-            for k in ("n", "num_inference_steps", "guidance_scale", "true_cfg_scale"):
-                if k in sampling_params and sampling_params[k] is not None:
-                    form.add_field(k, str(sampling_params[k]))
+            for k, v in sampling_params.items():
+                form.add_field(k, str(v))
         if mask is not None:
             mask_filename = "mask.png"
             form.add_field(
@@ -217,7 +206,7 @@ class VLLMOmniClient:
 
         return torch.stack(image_tensors, dim=0)
 
-    async def generate_comprehension_chat_completion(
+    async def generate_understanding_chat_completion(
         self,
         *,
         model: str,
@@ -427,30 +416,28 @@ class VLLMOmniClient:
             message_content.append({"type": "video_url", "video_url": {"url": video_to_base64(video)}})
         messages = [{"role": "user", "content": message_content}]
 
+        payload: dict[str, Any] = {"messages": messages, "model": model}
+        if modalities:
+            payload["modalities"] = modalities
+
         combined_extra_body: dict[str, Any] = {}
         if sampling_params is not None:
             spec, _ = lookup_model_spec(model)
             is_single_sampling_param = isinstance(sampling_params, dict) or len(sampling_params) == 1
 
-            # Exclude internal key
-            if isinstance(sampling_params, dict):
-                sampling_params = {k: v for k, v in sampling_params.items() if k != "type"}
-            else:
-                sampling_params = [{k: v for k, v in sp.items() if k != "type"} for sp in sampling_params]
-
             if (spec is None and is_single_sampling_param) or (spec is not None and spec["stages"] == ["diffusion"]):
                 # Diffusion format: extra_body directly contains sampling params.
-                # Validation should have taken care of matching sampling params' types.
+                # Validation has already taken care of matching sampling params' types and length. Safe to take [0].
                 # * Use this mode if the model is a simple one-stage diffusion model.
                 # * Fallback to this mode if model is not registered and a single sampling param is provided.
                 sampling_params = sampling_params if isinstance(sampling_params, dict) else sampling_params[0]
-                combined_extra_body: dict[str, Any] = {**sampling_params}
+                combined_extra_body: dict[str, Any] = sampling_params.copy()
                 if "n" in combined_extra_body:
-                    combined_extra_body["num_outputs_per_prompt"] = combined_extra_body["n"]
-                    del combined_extra_body["n"]
+                    combined_extra_body["num_outputs_per_prompt"] = combined_extra_body.pop("n")
             else:
-                # Use AR style payload, extra_body has a sampling_params_list field
-                combined_extra_body: dict[str, Any] = {"sampling_params_list": sampling_params}
+                # AR format: the payload has a sampling_params_list field, containing a list.
+                sampling_params_list = sampling_params if isinstance(sampling_params, list) else [sampling_params]
+                payload["sampling_params_list"] = sampling_params_list
 
         if negative_prompt:
             combined_extra_body["negative_prompt"] = negative_prompt
@@ -458,12 +445,11 @@ class VLLMOmniClient:
         if extra_body:
             combined_extra_body.update(extra_body)
 
-        payload: dict[str, Any] = {"messages": messages, "model": model}
+        # Add extra_body only if it has any content.
         if combined_extra_body:
             payload["extra_body"] = combined_extra_body
-        if modalities:
-            payload["modalities"] = modalities
 
+        # Place to inject any model-specific payload adjustment
         spec, _ = lookup_model_spec(model)
         if spec:
             preprocessor = spec.get("payload_preprocessor", None)

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/models.py

@@ -11,8 +11,7 @@ def _bagel_payload_preprocessor(payload: dict) -> dict:
                     content["text"] = "<|im_start|>" + content["text"] + "<|im_end|>"
     except (KeyError, TypeError):
         raise RuntimeError("Internal Error: malformatted BAGEL payload")
-    extra_body = payload.pop("extra_body", {})
-    return {**payload, **extra_body}
+    return payload
 
 
 def _qwen25_payload_preprocessor(payload: dict) -> dict:
@@ -37,7 +36,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
         ],
         "modes": [
             {
-                "mode": ModelMode.COMPREHENSION,
+                "mode": ModelMode.UNDERSTANDING,
                 "input_modalities": [Modality.TEXT, Modality.IMAGE],
             }
         ],
@@ -48,7 +47,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
         "payload_preprocessor": _qwen25_payload_preprocessor,
         "modes": [
             {
-                "mode": ModelMode.COMPREHENSION,
+                "mode": ModelMode.UNDERSTANDING,
                 "input_modalities": [
                     Modality.TEXT,
                     Modality.IMAGE,
@@ -62,7 +61,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
         "stages": ["autoregression", "autoregression", "autoregression"],
         "modes": [
             {
-                "mode": ModelMode.COMPREHENSION,
+                "mode": ModelMode.UNDERSTANDING,
                 "input_modalities": [
                     Modality.TEXT,
                     Modality.IMAGE,

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py

@@ -11,11 +11,23 @@ from typing import (
 AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"]
 
 
+class AutoregressionSamplingParams(dict):
+    pass
+
+
+class DiffusionSamplingParams(dict):
+    pass
+
+
+class QwenTTSModelSpecificParams(dict):
+    pass
+
+
 class ModelMode(Enum):
     IMAGE_GENERATION = auto()
     VIDEO_GENERATION = auto()
     AUDIO_GENERATION = auto()
-    COMPREHENSION = auto()
+    UNDERSTANDING = auto()
 
 
 class Modality(Enum):

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py

@@ -1,5 +1,6 @@
 from .logger import get_logger
 from .models import lookup_model_spec
+from .types import AutoregressionSamplingParams, DiffusionSamplingParams
 
 logger = get_logger(__name__)
 
@@ -31,24 +32,20 @@ def validate_model_and_sampling_params_types(
             )
         # Check that each stage's type match
         for i, sp in enumerate(sampling_param_list):
-            if "type" not in sp:
-                raise RuntimeError("Internal error: unknown sampling parameter type")
-            if sp["type"] != stages[i]:
+            if not _check_sampling_param_matches_stage(sp, stages[i]):
                 raise ValueError(
-                    f"Sampling parameter type ({sp['type']}) does not match "
+                    f"Sampling parameter type ({sp.__class__.__name__}) does not match "
                     f"stage type ({stages[i]}) at index {i} for model {model_name}."
                 )
     elif isinstance(sampling_param_list, dict):
-        if "type" not in sampling_param_list:
-            raise RuntimeError("Internal error: unknown sampling parameter type")
         # Check that the provided single sampling param matches all stages
-        elif any(stage != sampling_param_list["type"] for stage in stages):
-            raise ValueError(
-                f"When passing a single sampling parameter node, all stages of the model must match "
-                f"the provided sampling parameter's type. "
-                f"However, the stages of model {model_name} are: {stages}. "
-                f"The provided sampling parameter is {sampling_param_list['type']}"
-            )
+        for i, stage in enumerate(stages):
+            if not _check_sampling_param_matches_stage(sampling_param_list, stage):
+                raise ValueError(
+                    f"Provided single sampling parameter type ({sampling_param_list.__class__.__name__}) must match "
+                    f"the types of all stages of the model. "
+                    f"However, stage {i} of model {model_name} is of type {stage}."
+                )
 
 
 def add_sampling_parameters_to_stage(
@@ -82,3 +79,11 @@ def add_sampling_parameters_to_stage(
                 sampling_param_list[i].update(params_to_add)
 
     return sampling_param_list
+
+
+def _check_sampling_param_matches_stage(sampling_param: dict, stage_type: str) -> bool:
+    if stage_type == "autoregression":
+        return isinstance(sampling_param, AutoregressionSamplingParams)
+    if stage_type == "diffusion":
+        return isinstance(sampling_param, DiffusionSamplingParams)
+    raise RuntimeError(f"Internal error: unknown stage type {stage_type}.")

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Annotated Example.json

@@ -1 +0,0 @@
-{"id":"1c99f525-0a37-45ba-a28a-7df7c3af66b4","revision":0,"last_node_id":12,"last_link_id":14,"nodes":[{"id":1,"type":"VLLMOmniComprehension","pos":[1191.2177053682556,144.66829928181377],"size":[400,268],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":1},{"localized_name":"video","name":"video","shape":7,"type":"VIDEO","link":2},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":3},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":14},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"output_text","name":"output_text","type":"BOOLEAN","widget":{"name":"output_text"},"link":null},{"localized_name":"output_audio","name":"output_audio","type":"BOOLEAN","widget":{"name":"output_audio"},"link":null},{"localized_name":"use_audio_in_video","name":"use_audio_in_video","type":"BOOLEAN","widget":{"name":"use_audio_in_video"},"link":null}],"outputs":[{"localized_name":"text_response","name":"text_response","type":"STRING","links":[8]},{"localized_name":"audio_response","name":"audio_response","type":"AUDIO","links":[9]}],"properties":{"Node name for S&R":"VLLMOmniComprehension"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen2.5-Omni-7B","",true,true,true]},{"id":3,"type":"LoadVideo","pos":[729.5984141255855,-198.631920454299],"size":[282.798828125,233.0743408203125],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"file","name":"file","type":"COMBO","widget":{"name":"file"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[2]}],"properties":{"Node name for S&R":"LoadVideo"},"widgets_values":["draw.mp4","image"]},{"id":4,"type":"LoadAudio","pos":[729.8037086965753,99.86963519703949],"size":[282.798828125,136],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":5,"type":"VLLMOmniARSampling","pos":[510.3517536642828,658.073751009259],"size":[270,178],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[12]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":7,"type":"VLLMOmniARSampling","pos":[503.33235181647115,419.34158016181806],"size":[270,178],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[5,13]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":8,"type":"VLLMOmniSamplingParamsList","pos":[820.6056617389042,426.38372037182273],"size":[263.066015625,66],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"param1","name":"param1","type":"SAMPLING_PARAMS","link":5},{"localized_name":"param2","name":"param2","shape":7,"type":"SAMPLING_PARAMS","link":12},{"localized_name":"param3","name":"param3","shape":7,"type":"SAMPLING_PARAMS","link":13}],"outputs":[{"localized_name":"param list","name":"param list"Line truncated

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Chaining Services.json

@@ -0,0 +1,552 @@
+{
+    "id": "6643e5fd-fa2a-4f25-935a-483173a8097c",
+    "revision": 0,
+    "last_node_id": 8,
+    "last_link_id": 6,
+    "nodes": [
+        {
+            "id": 1,
+            "type": "PreviewImage",
+            "pos": [
+                1446.2005205859512,
+                -316.7359686049902
+            ],
+            "size": [
+                305.6628685610174,
+                307.0336884172169
+            ],
+            "flags": {},
+            "order": 5,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "images",
+                    "name": "images",
+                    "type": "IMAGE",
+                    "link": 1
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewImage"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 8,
+            "type": "PreviewImage",
+            "pos": [
+                1010.1848219273943,
+                -311.54216706114363
+            ],
+            "size": [
+                305.6628685610174,
+                307.0336884172169
+            ],
+            "flags": {},
+            "order": 4,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "images",
+                    "name": "images",
+                    "type": "IMAGE",
+                    "link": 6
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewImage"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 7,
+            "type": "VLLMOmniDiffusionSampling",
+            "pos": [
+                234.61958293819296,
+                9.405681270043619
+            ],
+            "size": [
+                284.205078125,
+                226
+            ],
+            "flags": {},
+            "order": 0,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "n",
+                    "name": "n",
+                    "type": "INT",
+                    "widget": {
+                        "name": "n"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "num_inference_steps",
+                    "name": "num_inference_steps",
+                    "type": "INT",
+                    "widget": {
+                        "name": "num_inference_steps"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "guidance_scale",
+                    "name": "guidance_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "guidance_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "true_cfg_scale",
+                    "name": "true_cfg_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "true_cfg_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_slicing",
+                    "name": "vae_use_slicing",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_slicing"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_tiling",
+                    "name": "vae_use_tiling",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_tiling"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "diffusion sampling params",
+                    "name": "diffusion sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        5
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniDiffusionSampling"
+            },
+            "widgets_values": [
+                1,
+                50,
+                1,
+                1,
+                false,
+                false,
+                1525,
+                "randomize"
+            ]
+        },
+        {
+            "id": 4,
+            "type": "VLLMOmniDiffusionSampling",
+            "pos": [
+                666.8380026154548,
+                268.86271068330126
+            ],
+            "size": [
+                284.205078125,
+                226
+            ],
+            "flags": {},
+            "order": 1,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "n",
+                    "name": "n",
+                    "type": "INT",
+                    "widget": {
+                        "name": "n"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "num_inference_steps",
+                    "name": "num_inference_steps",
+                    "type": "INT",
+                    "widget": {
+                        "name": "num_inference_steps"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "guidance_scale",
+                    "name": "guidance_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "guidance_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "true_cfg_scale",
+                    "name": "true_cfg_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "true_cfg_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_slicing",
+                    "name": "vae_use_slicing",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_slicing"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_tiling",
+                    "name": "vae_use_tiling",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_tiling"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "diffusion sampling params",
+                    "name": "diffusion sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        2
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniDiffusionSampling"
+            },
+            "widgets_values": [
+                4,
+                50,
+                7,
+                1,
+                false,
+                false,
+                42,
+                "fixed"
+            ]
+        },
+        {
+            "id": 5,
+            "type": "VLLMOmniGenerateImage",
+            "pos": [
+                984.723613585788,
+                63.376900027553276
+            ],
+            "size": [
+                416.56628685610167,
+                372.1662621294205
+            ],
+            "flags": {},
+            "order": 3,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": 4
+                },
+                {
+                    "localized_name": "mask",
+                    "name": "mask",
+                    "shape": 7,
+                    "type": "MASK",
+                    "link": null
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 2
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "negative_prompt",
+                    "name": "negative_prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "negative_prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "width",
+                    "name": "width",
+                    "type": "INT",
+                    "widget": {
+                        "name": "width"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "height",
+                    "name": "height",
+                    "type": "INT",
+                    "widget": {
+                        "name": "height"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "type": "IMAGE",
+                    "links": [
+                        1
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniGenerateImage"
+            },
+            "widgets_values": [
+                "http://localhost:8001/v1",
+                "/home/models/Qwen/Qwen-Image-Edit",
+                "A high-quality, high contrast, stylized portrait of the object in the uploaded reference image. Pop art and doodle style, with abundant scribbling patterns, such as a teal crown, orange lightning bolts, colorful handwritten scripts.",
+                "Realistic",
+                800,
+                800
+            ]
+        },
+        {
+            "id": 6,
+            "type": "VLLMOmniGenerateImage",
+            "pos": [
+                541.2702717429818,
+                -160.70392744708548
+            ],
+            "size": [
+                416.56628685610167,
+                372.1662621294205
+            ],
+            "flags": {},
+            "order": 2,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": null
+                },
+                {
+                    "localized_name": "mask",
+                    "name": "mask",
+                    "shape": 7,
+                    "type": "MASK",
+                    "link": null
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 5
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "negative_prompt",
+                    "name": "negative_prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "negative_prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "width",
+                    "name": "width",
+                    "type": "INT",
+                    "widget": {
+                        "name": "width"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "height",
+                    "name": "height",
+                    "type": "INT",
+                    "widget": {
+                        "name": "height"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "type": "IMAGE",
+                    "links": [
+                        4,
+                        6
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniGenerateImage"
+            },
+            "widgets_values": [
+                "http://localhost:8000/v1",
+                "/home/models/Tongyi-MAI/Z-Image-Turbo",
+                "A headshot of a cute Siamese kitty. Blurred background due to wide aperture. Close-up look. Realistic.",
+                "Cartoonish.",
+                800,
+                800
+            ]
+        }
+    ],
+    "links": [
+        [
+            1,
+            5,
+            0,
+            1,
+            0,
+            "IMAGE"
+        ],
+        [
+            2,
+            4,
+            0,
+            5,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            4,
+            6,
+            0,
+            5,
+            0,
+            "IMAGE"
+        ],
+        [
+            5,
+            7,
+            0,
+            6,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            6,
+            6,
+            0,
+            8,
+            0,
+            "IMAGE"
+        ]
+    ],
+    "groups": [],
+    "config": {},
+    "extra": {
+        "workflowRendererVersion": "LG",
+        "ds": {
+            "scale": 1.0426432563169903,
+            "offset": [
+                -90.4071374799346,
+                575.2948920069742
+            ]
+        }
+    },
+    "version": 0.4
+}

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Image Generation.json

@@ -1 +1 @@
-{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,202],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,42,"fixed"]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[496.31859627609606,-229.71277089860084],"size":[270,314],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Input","bounding":[213.1706010087147,-313.20095750667554,560.8246471437027,407.4246532472182],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486894,"offset":[276.8427259980847,549.4885226647582]}},"version":0.4}
+{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"图像","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[496.31859627609606,-229.71277089860084],"size":[270,314],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"图像","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"选择文件上传","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"图像","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"遮罩","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,226],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"vae_use_tiling","name":"vae_use_tiling","type":"BOOLEAN","widget":{"name":"vae_use_tiling"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,false,42,"randomize"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Input","bounding":[213.1706010087147,-313.20095750667554,560.8246471437027,407.4246532472182],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486894,"offset":[-55.18047513099182,220.2553505195962]}},"version":0.4}

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Multimodal Understanding.json

@@ -0,0 +1,761 @@
+{
+    "id": "1c99f525-0a37-45ba-a28a-7df7c3af66b4",
+    "revision": 0,
+    "last_node_id": 12,
+    "last_link_id": 14,
+    "nodes": [
+        {
+            "id": 1,
+            "type": "VLLMOmniUnderstanding",
+            "pos": [
+                1191.2177053682556,
+                144.66829928181377
+            ],
+            "size": [
+                400,
+                268
+            ],
+            "flags": {},
+            "order": 8,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": 1
+                },
+                {
+                    "localized_name": "video",
+                    "name": "video",
+                    "shape": 7,
+                    "type": "VIDEO",
+                    "link": 2
+                },
+                {
+                    "localized_name": "audio",
+                    "name": "audio",
+                    "shape": 7,
+                    "type": "AUDIO",
+                    "link": 3
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 14
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "output_text",
+                    "name": "output_text",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "output_text"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "output_audio",
+                    "name": "output_audio",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "output_audio"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "use_audio_in_video",
+                    "name": "use_audio_in_video",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "use_audio_in_video"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "text_response",
+                    "name": "text_response",
+                    "type": "STRING",
+                    "links": [
+                        8
+                    ]
+                },
+                {
+                    "localized_name": "audio_response",
+                    "name": "audio_response",
+                    "type": "AUDIO",
+                    "links": [
+                        9
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniUnderstanding"
+            },
+            "widgets_values": [
+                "http://localhost:8000/v1",
+                "Qwen/Qwen2.5-Omni-7B",
+                "",
+                true,
+                true,
+                true
+            ]
+        },
+        {
+            "id": 3,
+            "type": "LoadVideo",
+            "pos": [
+                729.5984141255855,
+                -198.631920454299
+            ],
+            "size": [
+                282.798828125,
+                233.0743408203125
+            ],
+            "flags": {},
+            "order": 0,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "file",
+                    "name": "file",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "file"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "choose file to upload",
+                    "name": "upload",
+                    "type": "IMAGEUPLOAD",
+                    "widget": {
+                        "name": "upload"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "VIDEO",
+                    "name": "VIDEO",
+                    "type": "VIDEO",
+                    "links": [
+                        2
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadVideo"
+            },
+            "widgets_values": [
+                "draw.mp4",
+                "image"
+            ]
+        },
+        {
+            "id": 4,
+            "type": "LoadAudio",
+            "pos": [
+                729.8037086965753,
+                99.86963519703949
+            ],
+            "size": [
+                282.798828125,
+                136
+            ],
+            "flags": {},
+            "order": 1,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "audio",
+                    "name": "audio",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "audio"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "audioUI",
+                    "name": "audioUI",
+                    "type": "AUDIO_UI",
+                    "widget": {
+                        "name": "audioUI"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "choose file to upload",
+                    "name": "upload",
+                    "type": "AUDIOUPLOAD",
+                    "widget": {
+                        "name": "upload"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "AUDIO",
+                    "name": "AUDIO",
+                    "type": "AUDIO",
+                    "links": [
+                        3
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadAudio"
+            },
+            "widgets_values": [
+                "Megan-Fox.mp3",
+                null,
+                null
+            ]
+        },
+        {
+            "id": 5,
+            "type": "VLLMOmniARSampling",
+            "pos": [
+                510.3517536642828,
+                658.073751009259
+            ],
+            "size": [
+                270,
+                178
+            ],
+            "flags": {},
+            "order": 2,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "max_tokens",
+                    "name": "max_tokens",
+                    "type": "INT",
+                    "widget": {
+                        "name": "max_tokens"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "temperature",
+                    "name": "temperature",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "temperature"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "top_p",
+                    "name": "top_p",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "top_p"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "repetition_penalty",
+                    "name": "repetition_penalty",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "repetition_penalty"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "AR sampling params",
+                    "name": "AR sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        12
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniARSampling"
+            },
+            "widgets_values": [
+                100,
+                1,
+                1,
+                1,
+                -1,
+                "randomize"
+            ]
+        },
+        {
+            "id": 7,
+            "type": "VLLMOmniARSampling",
+            "pos": [
+                503.33235181647115,
+                419.34158016181806
+            ],
+            "size": [
+                270,
+                178
+            ],
+            "flags": {},
+            "order": 3,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "max_tokens",
+                    "name": "max_tokens",
+                    "type": "INT",
+                    "widget": {
+                        "name": "max_tokens"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "temperature",
+                    "name": "temperature",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "temperature"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "top_p",
+                    "name": "top_p",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "top_p"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "repetition_penalty",
+                    "name": "repetition_penalty",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "repetition_penalty"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "AR sampling params",
+                    "name": "AR sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        5,
+                        13
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniARSampling"
+            },
+            "widgets_values": [
+                100,
+                1,
+                1,
+                1,
+                -1,
+                "randomize"
+            ]
+        },
+        {
+            "id": 8,
+            "type": "VLLMOmniSamplingParamsList",
+            "pos": [
+                820.6056617389042,
+                426.38372037182273
+            ],
+            "size": [
+                263.066015625,
+                66
+            ],
+            "flags": {},
+            "order": 7,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "param1",
+                    "name": "param1",
+                    "type": "SAMPLING_PARAMS",
+                    "link": 5
+                },
+                {
+                    "localized_name": "param2",
+                    "name": "param2",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 12
+                },
+                {
+                    "localized_name": "param3",
+                    "name": "param3",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 13
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "param list",
+                    "name": "param list",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        14
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniSamplingParamsList"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 11,
+            "type": "MarkdownNote",
+            "pos": [
+                826.2328280438272,
+                569.1890318701705
+            ],
+            "size": [
+                333.8220435590464,
+                261.63596728060656
+            ],
+            "flags": {},
+            "order": 4,
+            "mode": 0,
+            "inputs": [],
+            "outputs": [],
+            "title": "Note: Sampling Parameters",
+            "properties": {},
+            "widgets_values": [
+                "## Sampling Parameter Types\n\nThere are two types of sampling parameters: one for autoregression and one for diffusion.\nYou should ensure that you have chosen the correct type of sampling parameters for the model you request.\n\n## Stages & Shorthand\n\nFor multi-stage models such as Qwen Omni, you can either\n- connect one sampling parameter node, which is applied to all stages.\n- connect exactly the same number of sampling parameter nodes to a \"Multi-Stage Sampling Parameter List\" node, then connect this node to the primary request node.\n\nNote that this shorthand is intended to stay consistent with the [online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/qwen2_5_omni/)"
+            ],
+            "color": "#432",
+            "bgcolor": "#000"
+        },
+        {
+            "id": 12,
+            "type": "MarkdownNote",
+            "pos": [
+                378.9866207003777,
+                152.59550252215752
+            ],
+            "size": [
+                319.7287574247016,
+                107.15904081906785
+            ],
+            "flags": {},
+            "order": 6,
+            "mode": 0,
+            "inputs": [],
+            "outputs": [],
+            "title": "Note: Input",
+            "properties": {},
+            "widgets_values": [
+                "Note that not all models support every modality as input. For example, `ByteDance-Seed/BAGEL-7B-MoT` in Multimodality Understanding mode only support text and image input.\n\nYou should ensure that the input are supported by the model. You can check the corresponding [online serving documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/) for confirmation."
+            ],
+            "color": "#432",
+            "bgcolor": "#000"
+        },
+        {
+            "id": 2,
+            "type": "LoadImage",
+            "pos": [
+                394.4674804308822,
+                -207.6987397548834
+            ],
+            "size": [
+                282.798828125,
+                314
+            ],
+            "flags": {},
+            "order": 5,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "image"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "choose file to upload",
+                    "name": "upload",
+                    "type": "IMAGEUPLOAD",
+                    "widget": {
+                        "name": "upload"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "IMAGE",
+                    "name": "IMAGE",
+                    "type": "IMAGE",
+                    "links": [
+                        1
+                    ]
+                },
+                {
+                    "localized_name": "MASK",
+                    "name": "MASK",
+                    "type": "MASK",
+                    "links": null
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadImage"
+            },
+            "widgets_values": [
+                "example.png",
+                "image"
+            ]
+        },
+        {
+            "id": 10,
+            "type": "PreviewAudio",
+            "pos": [
+                1664.548345556043,
+                297.5921292054054
+            ],
+            "size": [
+                270,
+                88
+            ],
+            "flags": {},
+            "order": 10,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "audio",
+                    "name": "audio",
+                    "type": "AUDIO",
+                    "link": 9
+                },
+                {
+                    "localized_name": "audioUI",
+                    "name": "audioUI",
+                    "type": "AUDIO_UI",
+                    "widget": {
+                        "name": "audioUI"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewAudio"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 9,
+            "type": "ShowText|pysssss",
+            "pos": [
+                1649.2506875091847,
+                66.22823888292349
+            ],
+            "size": [
+                318.7188464232943,
+                173.38502269972975
+            ],
+            "flags": {},
+            "order": 9,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "text",
+                    "name": "text",
+                    "type": "STRING",
+                    "link": 8
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "STRING",
+                    "name": "STRING",
+                    "shape": 6,
+                    "type": "STRING",
+                    "links": null
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "ShowText|pysssss"
+            },
+            "widgets_values": []
+        }
+    ],
+    "links": [
+        [
+            1,
+            2,
+            0,
+            1,
+            0,
+            "IMAGE"
+        ],
+        [
+            2,
+            3,
+            0,
+            1,
+            1,
+            "VIDEO"
+        ],
+        [
+            3,
+            4,
+            0,
+            1,
+            2,
+            "AUDIO"
+        ],
+        [
+            5,
+            7,
+            0,
+            8,
+            0,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            8,
+            1,
+            0,
+            9,
+            0,
+            "STRING"
+        ],
+        [
+            9,
+            1,
+            1,
+            10,
+            0,
+            "AUDIO"
+        ],
+        [
+            12,
+            5,
+            0,
+            8,
+            1,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            13,
+            7,
+            0,
+            8,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            14,
+            8,
+            0,
+            1,
+            3,
+            "SAMPLING_PARAMS"
+        ]
+    ],
+    "groups": [
+        {
+            "id": 1,
+            "title": "Sampling Parameters",
+            "bounding": [
+                480.1649301181556,
+                341.08402513937995,
+                692.4113568972277,
+                510.48648853403665
+            ],
+            "color": "#3f789e",
+            "font_size": 24,
+            "flags": {}
+        },
+        {
+            "id": 2,
+            "title": "Input",
+            "bounding": [
+                344.2364817813528,
+                -287.5850484183313,
+                704.8832238768634,
+                559.5124009832894
+            ],
+            "color": "#3f789e",
+            "font_size": 24,
+            "flags": {}
+        }
+    ],
+    "config": {},
+    "extra": {
+        "workflowRendererVersion": "LG",
+        "ds": {
+            "scale": 0.9478575057427204,
+            "offset": [
+                33.81037136029557,
+                307.1974296197726
+            ]
+        }
+    },
+    "version": 0.4
+}

benchmarks/qwen3-tts/vllm_omni/bench_async_chunk.py

@@ -0,0 +1,301 @@
+"""Benchmark comparing async_chunk on vs off for Qwen3-TTS.
+
+Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across
+concurrency levels for both async_chunk modes. Saves results as JSON.
+
+Usage:
+    # Run against a server already serving with a given config:
+    python bench_async_chunk.py \
+        --host 127.0.0.1 --port 8000 \
+        --config-name async_chunk_on \
+        --num-prompts 50 \
+        --max-concurrency 1 10 \
+        --result-dir results/
+"""
+
+import argparse
+import asyncio
+import json
+import time
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import aiohttp
+import numpy as np
+from tqdm.asyncio import tqdm
+
+PROMPTS = [
+    "Hello, welcome to the voice synthesis benchmark test.",
+    "She said she would be here by noon, but nobody showed up.",
+    "The quick brown fox jumps over the lazy dog near the riverbank.",
+    "I can't believe how beautiful the sunset looks from up here on the mountain.",
+    "Please remember to bring your identification documents to the appointment tomorrow morning.",
+    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
+    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
+    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
+    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
+    "The train leaves at half past seven, so we need to arrive at the station before then.",
+    "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.",
+]
+
+
+@dataclass
+class RequestResult:
+    success: bool = False
+    ttfp: float = 0.0
+    e2e: float = 0.0
+    audio_bytes: int = 0
+    audio_duration: float = 0.0
+    rtf: float = 0.0
+    prompt: str = ""
+    error: str = ""
+
+
+@dataclass
+class BenchmarkResult:
+    config_name: str = ""
+    concurrency: int = 0
+    num_prompts: int = 0
+    completed: int = 0
+    failed: int = 0
+    duration_s: float = 0.0
+    mean_ttfp_ms: float = 0.0
+    median_ttfp_ms: float = 0.0
+    std_ttfp_ms: float = 0.0
+    p90_ttfp_ms: float = 0.0
+    p95_ttfp_ms: float = 0.0
+    p99_ttfp_ms: float = 0.0
+    mean_e2e_ms: float = 0.0
+    median_e2e_ms: float = 0.0
+    std_e2e_ms: float = 0.0
+    p90_e2e_ms: float = 0.0
+    p95_e2e_ms: float = 0.0
+    p99_e2e_ms: float = 0.0
+    mean_rtf: float = 0.0
+    median_rtf: float = 0.0
+    std_rtf: float = 0.0
+    mean_audio_duration_s: float = 0.0
+    total_audio_duration_s: float = 0.0
+    audio_throughput: float = 0.0
+    request_throughput: float = 0.0
+    per_request: list = field(default_factory=list)
+
+
+def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
+    return num_bytes / sample_width / sample_rate
+
+
+async def send_tts_request(
+    session: aiohttp.ClientSession,
+    api_url: str,
+    prompt: str,
+    voice: str = "vivian",
+    language: str = "English",
+    stream: bool = True,
+    pbar: tqdm | None = None,
+) -> RequestResult:
+    payload = {
+        "input": prompt,
+        "voice": voice,
+        "language": language,
+        "stream": stream,
+        "response_format": "pcm",
+    }
+
+    result = RequestResult(prompt=prompt)
+    st = time.perf_counter()
+
+    try:
+        async with session.post(api_url, json=payload) as response:
+            if response.status != 200:
+                result.error = f"HTTP {response.status}: {await response.text()}"
+                return result
+
+            first_chunk = True
+            total_bytes = 0
+
+            async for chunk in response.content.iter_any():
+                if first_chunk and len(chunk) > 0:
+                    result.ttfp = time.perf_counter() - st
+                    first_chunk = False
+                total_bytes += len(chunk)
+
+            result.e2e = time.perf_counter() - st
+            result.audio_bytes = total_bytes
+            result.audio_duration = pcm_bytes_to_duration(total_bytes)
+            if result.audio_duration > 0:
+                result.rtf = result.e2e / result.audio_duration
+            result.success = True
+
+    except Exception as e:
+        result.error = str(e)
+        result.e2e = time.perf_counter() - st
+
+    if pbar:
+        pbar.update(1)
+    return result
+
+
+async def run_benchmark(
+    host: str,
+    port: int,
+    num_prompts: int,
+    max_concurrency: int,
+    num_warmups: int = 3,
+    voice: str = "vivian",
+    language: str = "English",
+    stream: bool = True,
+) -> BenchmarkResult:
+    api_url = f"http://{host}:{port}/v1/audio/speech"
+
+    connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
+    session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600))
+
+    if num_warmups > 0:
+        print(f"  Warming up with {num_warmups} requests...")
+        warmup_tasks = [
+            send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream)
+            for i in range(num_warmups)
+        ]
+        await asyncio.gather(*warmup_tasks)
+        print("  Warmup done.")
+
+    request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
+
+    print(f"  Running {num_prompts} requests with concurrency={max_concurrency}...")
+    semaphore = asyncio.Semaphore(max_concurrency)
+    pbar = tqdm(total=num_prompts, desc=f"  concurrency={max_concurrency}")
+
+    async def limited_request(prompt):
+        async with semaphore:
+            return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar)
+
+    start_time = time.perf_counter()
+    tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
+    results: list[RequestResult] = await asyncio.gather(*tasks)
+    duration = time.perf_counter() - start_time
+    pbar.close()
+
+    await session.close()
+
+    successful = [r for r in results if r.success]
+    failed = [r for r in results if not r.success]
+
+    bench = BenchmarkResult(
+        concurrency=max_concurrency,
+        num_prompts=num_prompts,
+        completed=len(successful),
+        failed=len(failed),
+        duration_s=duration,
+    )
+
+    if successful:
+        ttfps = [r.ttfp * 1000 for r in successful]
+        e2es = [r.e2e * 1000 for r in successful]
+        rtfs = [r.rtf for r in successful]
+        audio_durs = [r.audio_duration for r in successful]
+
+        bench.mean_ttfp_ms = float(np.mean(ttfps))
+        bench.median_ttfp_ms = float(np.median(ttfps))
+        bench.std_ttfp_ms = float(np.std(ttfps))
+        bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
+        bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
+        bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
+
+        bench.mean_e2e_ms = float(np.mean(e2es))
+        bench.median_e2e_ms = float(np.median(e2es))
+        bench.std_e2e_ms = float(np.std(e2es))
+        bench.p90_e2e_ms = float(np.percentile(e2es, 90))
+        bench.p95_e2e_ms = float(np.percentile(e2es, 95))
+        bench.p99_e2e_ms = float(np.percentile(e2es, 99))
+
+        bench.mean_rtf = float(np.mean(rtfs))
+        bench.median_rtf = float(np.median(rtfs))
+        bench.std_rtf = float(np.std(rtfs))
+
+        bench.mean_audio_duration_s = float(np.mean(audio_durs))
+        bench.total_audio_duration_s = float(np.sum(audio_durs))
+        bench.audio_throughput = bench.total_audio_duration_s / duration
+        bench.request_throughput = len(successful) / duration
+
+        bench.per_request = [
+            {
+                "ttfp_ms": r.ttfp * 1000,
+                "e2e_ms": r.e2e * 1000,
+                "rtf": r.rtf,
+                "audio_duration_s": r.audio_duration,
+                "prompt": r.prompt,
+            }
+            for r in successful
+        ]
+
+    print(f"\n{'=' * 60}")
+    print(f"  Concurrency: {max_concurrency}  |  Completed: {bench.completed}  |  Failed: {bench.failed}")
+    print(f"  Duration: {duration:.2f}s  |  Throughput: {bench.request_throughput:.2f} req/s")
+    print(
+        f"  TTFP (ms):  mean={bench.mean_ttfp_ms:.1f}  median={bench.median_ttfp_ms:.1f}"
+        f"  p90={bench.p90_ttfp_ms:.1f}  p99={bench.p99_ttfp_ms:.1f}"
+    )
+    print(
+        f"  E2E (ms):   mean={bench.mean_e2e_ms:.1f}  median={bench.median_e2e_ms:.1f}"
+        f"  p90={bench.p90_e2e_ms:.1f}  p99={bench.p99_e2e_ms:.1f}"
+    )
+    print(f"  RTF:        mean={bench.mean_rtf:.3f}  median={bench.median_rtf:.3f}")
+    print(f"  Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec")
+    print(f"{'=' * 60}\n")
+
+    if failed:
+        for r in failed[:3]:
+            print(f"  [ERROR] {r.error[:200]}")
+
+    return bench
+
+
+async def main(args):
+    all_results = []
+
+    for concurrency in args.max_concurrency:
+        result = await run_benchmark(
+            host=args.host,
+            port=args.port,
+            num_prompts=args.num_prompts,
+            max_concurrency=concurrency,
+            num_warmups=args.num_warmups,
+            voice=args.voice,
+            language=args.language,
+            stream=args.stream,
+        )
+        result.config_name = args.config_name
+        all_results.append(asdict(result))
+
+    result_dir = Path(args.result_dir)
+    result_dir.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
+
+    with open(result_file, "w") as f:
+        json.dump(all_results, f, indent=2)
+    print(f"Results saved to {result_file}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client")
+    parser.add_argument("--host", type=str, default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--num-prompts", type=int, default=50)
+    parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10])
+    parser.add_argument("--num-warmups", type=int, default=3)
+    parser.add_argument("--voice", type=str, default="vivian")
+    parser.add_argument("--language", type=str, default="English")
+    parser.add_argument("--stream", action="store_true", default=True)
+    parser.add_argument("--no-stream", dest="stream", action="store_false")
+    parser.add_argument("--config-name", type=str, default="async_chunk_on")
+    parser.add_argument("--result-dir", type=str, default="results")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    asyncio.run(main(args))

benchmarks/qwen3-tts/vllm_omni/plot_async_chunk.py

@@ -0,0 +1,249 @@
+"""Plot TTFP comparison: async_chunk off vs on.
+
+Generates a bar chart with improvement arrows, matching the Qwen3-Omni
+async_chunk benchmark figure style.
+
+Usage:
+    python plot_async_chunk.py \
+        --off results/bench_async_chunk_off_*.json \
+        --on results/bench_async_chunk_on_*.json \
+        --output results/qwen3_tts_async_chunk_ttfp.png
+
+    # Also supports E2E and RTF metrics:
+    python plot_async_chunk.py \
+        --off results/bench_async_chunk_off_*.json \
+        --on results/bench_async_chunk_on_*.json \
+        --metric e2e \
+        --output results/qwen3_tts_async_chunk_e2e.png
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+METRIC_CONFIG = {
+    "ttfp": {
+        "key": "mean_ttfp_ms",
+        "ylabel": "TTFP (s)",
+        "title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency",
+        "to_seconds": True,
+    },
+    "e2e": {
+        "key": "mean_e2e_ms",
+        "ylabel": "E2E (s)",
+        "title": "E2E Latency - Qwen3-TTS, by concurrency",
+        "to_seconds": True,
+    },
+    "rtf": {
+        "key": "mean_rtf",
+        "ylabel": "RTF",
+        "title": "Real-Time Factor - Qwen3-TTS, by concurrency",
+        "to_seconds": False,
+    },
+}
+
+
+def load_results(path: str) -> list[dict]:
+    with open(path) as f:
+        return json.load(f)
+
+
+def plot_ttfp_comparison(
+    off_results: list[dict],
+    on_results: list[dict],
+    metric: str,
+    output_path: str,
+    title_override: str | None = None,
+):
+    cfg = METRIC_CONFIG[metric]
+    key = cfg["key"]
+    to_seconds = cfg["to_seconds"]
+
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    off_vals = []
+    on_vals = []
+    for c in concurrencies:
+        v_off = off_map[c][key]
+        v_on = on_map[c][key]
+        if to_seconds:
+            v_off /= 1000.0
+            v_on /= 1000.0
+        off_vals.append(v_off)
+        on_vals.append(v_on)
+
+    fig, ax = plt.subplots(figsize=(8, 6))
+
+    x = np.arange(len(concurrencies))
+    width = 0.3
+
+    ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none")
+    ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
+
+    # Draw improvement arrows and labels
+    for i in range(len(concurrencies)):
+        v_off = off_vals[i]
+        v_on = on_vals[i]
+        if v_on > 0:
+            improvement = v_off / v_on
+        else:
+            improvement = float("inf")
+
+        # Arrow from top of off-bar to top of on-bar
+        arrow_start_x = x[i] - width / 2
+        arrow_start_y = v_off * 0.95
+        arrow_end_x = x[i] + width / 2
+        arrow_end_y = v_on * 1.05
+
+        ax.annotate(
+            "",
+            xy=(arrow_end_x, arrow_end_y),
+            xytext=(arrow_start_x, arrow_start_y),
+            arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
+        )
+
+        # Improvement label
+        label_x = (arrow_start_x + arrow_end_x) / 2
+        label_y = arrow_start_y + (v_off - v_on) * 0.15
+        ax.text(
+            label_x,
+            label_y,
+            f"{improvement:.1f}x improvement",
+            ha="center",
+            va="bottom",
+            fontsize=10,
+            color="red",
+            fontweight="bold",
+        )
+
+    title = title_override or cfg["title"]
+    ax.set_title(title, fontsize=13, fontweight="bold")
+    ax.set_ylabel(cfg["ylabel"], fontsize=12)
+    ax.set_xlabel("Max concurrency", fontsize=12)
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(c) for c in concurrencies])
+    ax.set_yscale("log")
+    ax.legend(loc="upper left", fontsize=11)
+    ax.grid(axis="y", alpha=0.3, linestyle="--")
+    ax.set_axisbelow(True)
+
+    plt.tight_layout()
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str):
+    """Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons."""
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
+    fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold")
+
+    for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]):
+        cfg = METRIC_CONFIG[metric]
+        key = cfg["key"]
+        to_seconds = cfg["to_seconds"]
+
+        off_vals = []
+        on_vals = []
+        for c in concurrencies:
+            v_off = off_map[c][key]
+            v_on = on_map[c][key]
+            if to_seconds:
+                v_off /= 1000.0
+                v_on /= 1000.0
+            off_vals.append(v_off)
+            on_vals.append(v_on)
+
+        x = np.arange(len(concurrencies))
+        width = 0.3
+        ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB")
+        ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
+
+        for i in range(len(concurrencies)):
+            if on_vals[i] > 0:
+                improvement = off_vals[i] / on_vals[i]
+                ax.annotate(
+                    "",
+                    xy=(x[i] + width / 2, on_vals[i] * 1.05),
+                    xytext=(x[i] - width / 2, off_vals[i] * 0.95),
+                    arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
+                )
+                label_y = off_vals[i] * 0.85
+                ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold")
+
+        ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold")
+        ax.set_ylabel(cfg["ylabel"], fontsize=11)
+        ax.set_xlabel("Max concurrency", fontsize=11)
+        ax.set_xticks(x)
+        ax.set_xticklabels([str(c) for c in concurrencies])
+        if metric != "rtf":
+            ax.set_yscale("log")
+        ax.legend(fontsize=9)
+        ax.grid(axis="y", alpha=0.3, linestyle="--")
+        ax.set_axisbelow(True)
+
+    plt.tight_layout()
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def print_table(off_results: list[dict], on_results: list[dict]):
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    print("\n## Benchmark Results: async_chunk off vs on\n")
+    print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |")
+    print("| --- | --- | --- | --- | --- |")
+
+    for name, key, fmt in [
+        ("TTFP (ms)", "mean_ttfp_ms", ".1f"),
+        ("E2E (ms)", "mean_e2e_ms", ".1f"),
+        ("RTF", "mean_rtf", ".3f"),
+        ("Throughput", "audio_throughput", ".2f"),
+    ]:
+        for c in concurrencies:
+            v_off = off_map[c].get(key, 0)
+            v_on = on_map[c].get(key, 0)
+            if v_on > 0 and key != "audio_throughput":
+                ratio = f"{v_off / v_on:.1f}x"
+            elif v_off > 0 and key == "audio_throughput":
+                ratio = f"{v_on / v_off:.1f}x"
+            else:
+                ratio = "N/A"
+            print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS")
+    parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off")
+    parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on")
+    parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"])
+    parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png")
+    parser.add_argument("--title", type=str, default=None, help="Custom title override")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    off_results = load_results(args.off)
+    on_results = load_results(args.on)
+
+    print_table(off_results, on_results)
+
+    if args.metric == "all":
+        plot_all_metrics(off_results, on_results, args.output)
+    else:
+        plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title)

benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh

@@ -0,0 +1,167 @@
+#!/bin/bash
+# Qwen3-TTS async_chunk on vs off Benchmark
+#
+# Starts two servers (async_chunk on and off), benchmarks both,
+# and generates comparison plots.
+#
+# Usage:
+#   bash run_async_chunk_benchmark.sh
+#
+# Environment variables:
+#   GPU_DEVICE       - GPU index (default: 0)
+#   MODEL            - Model path (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
+#   NUM_PROMPTS      - Prompts per concurrency level (default: 50)
+#   CONCURRENCY      - Space-separated concurrency levels (default: "1 10")
+#   PORT_ON          - Port for async_chunk on server (default: 8000)
+#   PORT_OFF         - Port for async_chunk off server (default: 8001)
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+cd "$PROJECT_ROOT"
+
+GPU_DEVICE="${GPU_DEVICE:-0}"
+MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
+NUM_PROMPTS="${NUM_PROMPTS:-50}"
+CONCURRENCY="${CONCURRENCY:-1 10}"
+NUM_WARMUPS="${NUM_WARMUPS:-3}"
+PORT_ON="${PORT_ON:-8000}"
+PORT_OFF="${PORT_OFF:-8001}"
+RESULT_DIR="${SCRIPT_DIR}/results"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+
+STAGE_CONFIG_ON="vllm_omni/model_executor/stage_configs/qwen3_tts.yaml"
+STAGE_CONFIG_OFF="vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml"
+
+mkdir -p "${RESULT_DIR}"
+
+echo "============================================================"
+echo " Qwen3-TTS async_chunk Benchmark"
+echo "============================================================"
+echo " GPU:            ${GPU_DEVICE}"
+echo " Model:          ${MODEL}"
+echo " Prompts:        ${NUM_PROMPTS}"
+echo " Concurrency:    ${CONCURRENCY}"
+echo " Port (on/off):  ${PORT_ON} / ${PORT_OFF}"
+echo " Results:        ${RESULT_DIR}"
+echo "============================================================"
+
+cleanup() {
+    echo "Cleaning up servers..."
+    kill "$PID_ON" 2>/dev/null || true
+    kill "$PID_OFF" 2>/dev/null || true
+    wait "$PID_ON" 2>/dev/null || true
+    wait "$PID_OFF" 2>/dev/null || true
+}
+trap cleanup EXIT
+
+wait_for_server() {
+    local port=$1
+    local name=$2
+    local max_wait=300
+    local elapsed=0
+    echo "Waiting for ${name} server on port ${port}..."
+    while ! curl -s "http://localhost:${port}/health" >/dev/null 2>&1; do
+        sleep 5
+        elapsed=$((elapsed + 5))
+        if [ $elapsed -ge $max_wait ]; then
+            echo "ERROR: ${name} server failed to start within ${max_wait}s"
+            exit 1
+        fi
+    done
+    echo "${name} server ready (${elapsed}s)"
+}
+
+# ---- Phase 1: Start async_chunk ON server ----
+echo ""
+echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..."
+CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
+    --stage-configs-path "${STAGE_CONFIG_ON}" \
+    --host 0.0.0.0 --port "${PORT_ON}" \
+    --trust-remote-code --enforce-eager --omni \
+    > "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 &
+PID_ON=$!
+
+wait_for_server "${PORT_ON}" "async_chunk_on"
+
+echo "[Phase 1] Benchmarking async_chunk ON..."
+# shellcheck disable=SC2086
+python "${SCRIPT_DIR}/bench_async_chunk.py" \
+    --host 127.0.0.1 --port "${PORT_ON}" \
+    --config-name "async_chunk_on" \
+    --num-prompts "${NUM_PROMPTS}" \
+    --max-concurrency ${CONCURRENCY} \
+    --num-warmups "${NUM_WARMUPS}" \
+    --result-dir "${RESULT_DIR}"
+
+echo "[Phase 1] Stopping async_chunk ON server..."
+kill "$PID_ON" 2>/dev/null || true
+wait "$PID_ON" 2>/dev/null || true
+sleep 5
+
+# ---- Phase 2: Start async_chunk OFF server ----
+echo ""
+echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..."
+CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
+    --stage-configs-path "${STAGE_CONFIG_OFF}" \
+    --host 0.0.0.0 --port "${PORT_OFF}" \
+    --trust-remote-code --enforce-eager --omni \
+    > "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 &
+PID_OFF=$!
+
+wait_for_server "${PORT_OFF}" "async_chunk_off"
+
+echo "[Phase 2] Benchmarking async_chunk OFF (non-streaming)..."
+# shellcheck disable=SC2086
+python "${SCRIPT_DIR}/bench_async_chunk.py" \
+    --host 127.0.0.1 --port "${PORT_OFF}" \
+    --config-name "async_chunk_off" \
+    --num-prompts "${NUM_PROMPTS}" \
+    --max-concurrency ${CONCURRENCY} \
+    --num-warmups "${NUM_WARMUPS}" \
+    --no-stream \
+    --result-dir "${RESULT_DIR}"
+
+echo "[Phase 2] Stopping async_chunk OFF server..."
+kill "$PID_OFF" 2>/dev/null || true
+wait "$PID_OFF" 2>/dev/null || true
+
+# ---- Phase 3: Plot results ----
+echo ""
+echo "[Phase 3] Generating plots..."
+
+# Find the latest result files
+RESULT_ON=$(ls -t "${RESULT_DIR}"/bench_async_chunk_on_*.json 2>/dev/null | head -1)
+RESULT_OFF=$(ls -t "${RESULT_DIR}"/bench_async_chunk_off_*.json 2>/dev/null | head -1)
+
+if [ -z "$RESULT_ON" ] || [ -z "$RESULT_OFF" ]; then
+    echo "ERROR: Could not find result files. Check logs in ${RESULT_DIR}/"
+    exit 1
+fi
+
+echo "  ON results:  ${RESULT_ON}"
+echo "  OFF results: ${RESULT_OFF}"
+
+# TTFP comparison (main figure)
+python "${SCRIPT_DIR}/plot_async_chunk.py" \
+    --off "${RESULT_OFF}" \
+    --on "${RESULT_ON}" \
+    --metric ttfp \
+    --output "${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
+
+# All metrics comparison
+python "${SCRIPT_DIR}/plot_async_chunk.py" \
+    --off "${RESULT_OFF}" \
+    --on "${RESULT_ON}" \
+    --metric all \
+    --output "${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
+
+echo ""
+echo "============================================================"
+echo " Benchmark complete!"
+echo " Results: ${RESULT_DIR}/"
+echo " Plots:"
+echo "   - ${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
+echo "   - ${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
+echo "============================================================"

docker/Dockerfile.ci

@@ -1,14 +1,13 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.16.0
+ARG VLLM_BASE_TAG=v0.17.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
-
 COPY . .
 
 # Install system dependencies
 RUN apt-get update && \
-    apt-get install -y ffmpeg sox libsox-fmt-all jq && \
+    apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 

docker/Dockerfile.npu

@@ -2,13 +2,15 @@ ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
 ARG VLLM_ASCEND_TAG=v0.14.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
+WORKDIR /vllm-workspace/vllm-ascend
+RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+RUN pip install -v -e .
+
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 
 COPY . .
 
-RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
- && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
 RUN pip install -v -e .
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn

docker/Dockerfile.npu.a3

@@ -2,13 +2,15 @@ ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
 ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
+WORKDIR /vllm-workspace/vllm-ascend
+RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+RUN pip install -v -e .
+
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 
 COPY . .
 
-RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
- && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
 RUN pip install -v -e .
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.16.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.17.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -19,7 +19,7 @@ WORKDIR ${COMMON_WORKDIR}
 
 # Step 1: Setup - Install system dependencies
 RUN apt-get update && \
-    apt-get install -y ffmpeg && \
+    apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 

docker/Dockerfile.xpu

@@ -9,12 +9,12 @@ ARG PYTHON_VERSION=3.12
 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
 
 RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
-    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
-    add-apt-repository -y ppa:kobuk-team/intel-graphics
+    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
 
 RUN apt clean && apt-get update -y && \
     apt-get install -y --no-install-recommends --fix-missing \
     curl \
+    espeak-ng \
     ffmpeg \
     git \
     libsndfile1 \
@@ -31,9 +31,22 @@ RUN apt clean && apt-get update -y && \
     python3-pip
 
 RUN apt update && apt upgrade -y && \
-    apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing intel-ocloc && \
     apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
 
+# Install UMD with fixed version
+RUN mkdir neo && \
+    cd neo && \
+    wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
+    wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
+    dpkg -i *.deb && \
+    cd .. && \
+    rm -rf neo
+
 ENV PATH="/root/.local/bin:$PATH"
 ENV VIRTUAL_ENV="/opt/venv"
 ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python

docs/.nav.yml

@@ -13,9 +13,10 @@ nav:
     - examples/README.md
     - Offline Inference:
       - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
+      - GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
       - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
-      - LoRA Inference(Diffusion): user_guide/examples/offline_inference/lora_inference.md
+      - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
       - Qwen3-TTS: user_guide/examples/offline_inference/qwen3_tts.md
@@ -24,12 +25,15 @@ nav:
       - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
     - Online Serving:
       - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
+      - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
-      - LoRA Inference(Diffusion): user_guide/examples/online_serving/lora_inference.md
+      - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
+      - Online serving Example of vLLM-Omni for MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
       - Qwen3-TTS: user_guide/examples/online_serving/qwen3_tts.md
       - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
+      - Text-To-Video: user_guide/examples/online_serving/text_to_video.md
   - General:
     - usage/*
   - Configuration:
@@ -46,18 +50,23 @@ nav:
       - Quantization:
         - Overview: user_guide/diffusion/quantization/overview.md
         - FP8: user_guide/diffusion/quantization/fp8.md
+        - GGUF: user_guide/diffusion/quantization/gguf.md
       - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
       - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
       - LoRA: user_guide/diffusion/lora.md
+      - Hybrid Sharded Data Parallel: design/feature/hsdp.md
+      - Custom Pipeline: features/custom_pipeline.md
     - ComfyUI: features/comfyui.md
 - Developer Guide:
   - General:
     - contributing/README.md
+    - pr_reviewer.md
     - glob: contributing/*
       flatten_single_child_sections: true
   - Model Implementation:
     - contributing/model/README.md
     - contributing/model/adding_omni_model.md
+    - contributing/model/adding_tts_model.md
     - contributing/model/adding_diffusion_model.md
   - CI: contributing/ci
   - Design Documents:
@@ -72,6 +81,7 @@ nav:
       - design/feature/tensor_parallel.md
       - design/feature/cache_dit.md
       - design/feature/teacache.md
+      - design/feature/async_chunk_design.md
     - Module Design:
       - design/module/ar_module.md
       - design/module/dit_module.md

docs/api/README.md

@@ -7,7 +7,6 @@ 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.async_omni_llm.AsyncOmniLLM][]
-- [vllm_omni.entrypoints.chat_utils.extract_audio_from_video_async][]
 - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
 - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
 - [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
@@ -19,6 +18,7 @@ Main entry points for vLLM-Omni inference and serving.
 - [vllm_omni.entrypoints.omni_llm.OmniLLM][]
 - [vllm_omni.entrypoints.omni_stage.OmniStage][]
 - [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
+- [vllm_omni.entrypoints.zmq_utils.ZmqQueue][]
 
 ## Inputs
 
@@ -36,6 +36,12 @@ Input data structures for multi-modal inputs.
 Engine classes for offline and online inference.
 
 - [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.BufferAllocator][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.ManagedBuffer][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.MooncakeAgentMetadata][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.MooncakeTransferEngineConnector][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.QueryRequest][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.QueryResponse][]
 - [vllm_omni.engine.AdditionalInformationEntry][]
 - [vllm_omni.engine.AdditionalInformationPayload][]
 - [vllm_omni.engine.OmniEngineCoreOutput][]
@@ -72,6 +78,8 @@ Configuration classes.
 - [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
 - [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
 - [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
+- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioConfig][]
+- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioTokenizerConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSSpeakerEncoderConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerCodePredictorConfig][]
@@ -89,8 +97,10 @@ Configuration classes.
 Worker classes and model runners for distributed inference.
 
 - [vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner][]
+- [vllm_omni.diffusion.worker.diffusion_worker.CustomPipelineWorkerExtension][]
 - [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
 - [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
+- [vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase][]
 - [vllm_omni.platforms.npu.worker.npu_ar_model_runner.ExecuteModelState][]
 - [vllm_omni.platforms.npu.worker.npu_ar_model_runner.NPUARModelRunner][]
 - [vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker][]

docs/configuration/README.md

@@ -2,7 +2,7 @@
 
 This section lists the most common options for running vLLM-Omni.
 
-For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html)
+For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.16.0/configuration/index.html)
 
 Currently, the main options are maintained by stage configs for each model.
 
@@ -19,3 +19,4 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
 - **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
 - **[Cache-DiT Configuration](../user_guide/diffusion/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
 - **[Parallelism Configuration](../user_guide/diffusion/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models
+- **[CPU Offloading](../user_guide/diffusion/cpu_offload_diffusion.md)** - Enable CPU offloading (model-level and layerwise) for for DiT models

docs/configuration/stage_configs.md

@@ -135,6 +135,14 @@ Each stage in the `stage_args` list contains the following configuration options
 
 A unique identifier for each stage in the multi-stage pipeline. Stages are numbered sequentially starting from 0, and this ID is used to reference stages in inter-stage dependencies (e.g., `engine_input_source`).
 
+### `prompt_expand_func` (Optional)
+
+A custom Python function hook for the LLM stage (Stage 0) that expands a single incoming prompt object into multiple prompts. This is primarily used for multi-modal Classifier-Free Guidance (CFG), where it generates the necessary companion requests (like a negative text prompt) and tags them with internal roles (e.g., `cfg_text`). This ensures the upstream LLM generates the needed contextual hidden states for both the conditional and unconditional generations simultaneously.
+
+### `cfg_kv_collect_func` (Optional)
+
+A custom Python function hook for downstream diffusion stages (Stage 1+) to collect, map, and process the KV caches transferred from the companion requests fired by `prompt_expand_func`. It aggregates the hidden condition states cleanly (e.g., binding them as `cfg_text_past_key_values` and `cfg_text_kv_metadata`), allowing the diffusion runtime to perform CFG smoothly without redundantly evaluating text paths on the DiT workers.
+
 ### `runtime`
 
 Configuration for disaggregated execution of the stage, controlling how the stage is deployed and executed.

docs/contributing/README.md

@@ -71,14 +71,7 @@ For additional features and advanced configurations, refer to the:
 ### Testing
 
 vLLM-Omni uses `pytest` to test the codebase.
-
-```bash
-# Run all tests
-pytest tests/
-
-# Run tests for a single test file with detailed output
-pytest -s -v tests/test_omni_llm.py
-```
+Please refer to the [test instructions](./ci/test_guide.md) for detailed testing information.
 
 !!! warning
     Currently, not all unit tests pass when run on CPU platforms. If you don't have access to a GPU platform to run unit tests locally, rely on the continuous integration system to run the tests for now.
@@ -123,6 +116,9 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
 !!! note
     If the PR spans more than one category, please include all relevant prefixes.
 
+### Local Test
+Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
+
 ### Code Quality
 
 The PR needs to meet the following code quality standards:

docs/contributing/ci/test_guide.md

@@ -0,0 +1,75 @@
+# Test Guide
+## Setting Up the Test Environment
+### Creating a Container
+vLLM-Omni provides an official Docker image for deployment. These images are built upon vLLM Docker images and are available on [Docker Hub](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which vLLM release it is based on.
+For a local test environment, you can follow the steps below to create a container:
+## Installing Dependencies
+### vLLM & vLLM-Omni
+vLLM-Omni is built based on vLLM. You can follow [install guide](../../getting_started/installation/README.md) to build your local environment.
+
+### Test Case Dependencies
+When running test cases, you may need to install the following dependencies:
+
+```bash
+uv pip install ".[dev]"
+apt-get install -y ffmpeg
+```
+
+## Running Tests
+Our test scripts use the pytest framework. First, please use `git clone https://github.com/vllm-project/vllm-omni.git` to download the vllm-omni source code. Then, in the root directory of vllm-omni, you can run the following commands in your local test environment to execute the corresponding test cases.
+
+=== "L1 level"
+
+    ```bash
+    cd tests
+    pytest -s -v -m "core_model and cpu"
+    ```
+    The latest test command is available in the "Simple Unit Test" step of this [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
+
+=== "L2 level"
+
+    ```bash
+    cd tests
+    pytest -s -v -m "core_model and not cpu" --run-level=core_model
+    ```
+    If you only want to run a specific test case, you can use:
+    ```bash
+    pytest -s -v test_xxxx.py --run-level=core_model
+    ```
+    If you only want to run specific test cases on a particular platform, you can use:
+    ```bash
+    pytest -s -v -m "core_model and distributed_cuda and L4"  --run-level=core_model
+    ```
+    The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
+
+=== "L3 level & L4 level"
+
+    ```bash
+    cd tests
+    pytest -s -v -m "advanced_model" --run-level=advanced_model
+    ```
+    If you only want to run L3 test case, you can use:
+    ```bash
+    pytest -s -v e2e/ --ignore-glob='*expansion.py' -m "advanced_model" --run-level=advanced_model
+    ```
+    If you only want to run a specific test case, you can use:
+    ```bash
+    pytest -s -v test_xxxx.py --run-level=advanced_model
+    ```
+    If you only want to run specific test cases on a particular platform, you can use:
+    ```bash
+    pytest -s -v -m "core_model and distributed_cuda and L4"  --run-level=core_model
+    ```
+    Note: To run performance tests, use:
+    ```bash
+    pytest -s -v perf/scripts/run_benchmark.py
+    ```
+
+    The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
+
+    The latest L4 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml).
+
+You can find more information about markers in the documentation: [marker doc](./tests_markers.md)
+
+## Adding New Test Cases
+Please refer to the [L5 Layering Specification document](./CI_5levels.md).

docs/contributing/model/adding_omni_model.md

@@ -460,7 +460,7 @@ def forward(self, ...):
 These keys are then accessible in your stage transition function:
 ```python
 # In stage_input_processors/qwen3_omni.py
-thinker_embeddings = output.multimodal_output["0"]  # Access by key
+thinker_prefill_embeddings = output.multimodal_output["0"]  # Access by key
 thinker_hidden_states = output.multimodal_output["24"]
 ```
 
@@ -513,11 +513,11 @@ def thinker2talker(
     for thinker_output in thinker_outputs:
         output = thinker_output.outputs[0]
         # Extract thinker embeddings and hidden states
-        thinker_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
+        thinker_prefill_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
         thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda()
 
         info = {
-            "thinker_embeddings": thinker_embeddings,
+            "thinker_prefill_embeddings": thinker_prefill_embeddings,
             "thinker_hidden_states": thinker_hidden_states,
             "thinker_sequences": thinker_output.prompt_token_ids + output.token_ids,
             "thinker_input_ids": thinker_output.prompt_token_ids,

docs/contributing/model/adding_tts_model.md

@@ -0,0 +1,584 @@
+# Adding a TTS Model
+
+This guide walks through adding a new TTS model to vLLM-Omni, using **Qwen3-TTS**
+as a reference. Qwen3-TTS demonstrates the standard two-stage TTS pipeline and the
+key optimizations all TTS models in this repo should follow.
+
+## Table of Contents
+
+1. [Overview](#overview)
+2. [Directory Structure](#directory-structure)
+3. [Step-by-Step Implementation](#step-by-step-implementation)
+4. [Key Components](#key-components)
+5. [Model Registration](#model-registration)
+6. [Stage Configuration](#stage-configuration)
+7. [Stage Input Processors](#stage-input-processors)
+8. [Testing](#testing)
+9. [Summary](#summary)
+
+## Overview
+
+vLLM-Omni supports TTS models as multi-stage pipelines where each stage runs independently
+and can be placed on different devices. Qwen3-TTS has two stages:
+
+| Stage | Name | Input | Output |
+|-------|------|-------|--------|
+| 0 | Code Predictor (AR) | Text tokens | Discrete RVQ codec codes |
+| 1 | Code2Wav (Decoder) | RVQ codec codes | Audio waveform |
+
+Each stage is a separate model class configured independently via YAML. The two stages
+are connected by the `async_chunk` framework, which enables inter-stage streaming for
+low first-packet latency (see [Async Chunk Design](../../design/feature/async_chunk_design.md)).
+
+### Without async_chunk (batch mode)
+
+Stage 0 runs to completion before Stage 1 starts, resulting in long first-packet latency:
+
+```mermaid
+flowchart TB
+    subgraph stage0["Stage 0: AR Code Predictor"]
+        direction LR
+        P[Prefill] --> D1[Decode 1]
+        D1 --> D2[Decode 2]
+        D2 --> Dots1["..."]
+        Dots1 --> DN[Decode N]
+    end
+
+    subgraph stage1["Stage 1: Code2Wav"]
+        direction LR
+        DEC[Decode all codes at once]
+    end
+
+    stage0 -- "all N codes" --> stage1
+    stage1 --> FPL["First Packet Latency = Stage 0 + Stage 1"]
+
+    style stage0 fill:#dae8fc,stroke:#6c8ebf
+    style stage1 fill:#f8d7c8,stroke:#d4856a
+    style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px
+```
+
+### With async_chunk (streaming mode)
+
+Stage 0 sends codec codes to Stage 1 every `chunk_size=25` tokens. Stage 1 begins decoding
+immediately, reducing first-packet latency from the full AR time to just the first chunk:
+
+```mermaid
+flowchart TB
+    subgraph stage0["Stage 0: Code Predictor (AR)"]
+        direction LR
+        P[Prefill] --> D1["Decode 1-25"]
+        D1 --> D2["Decode 26-50"]
+        D2 --> Dots1["..."]
+        Dots1 --> DN["Decode N"]
+    end
+
+    subgraph stage1["Stage 1: Code2Wav"]
+        direction LR
+        C1["Chunk 1\n(25 frames)"] --> C2["Chunk 2\n(context + 25)"]
+        C2 --> Dots2["..."]
+        Dots2 --> CN["Final chunk"]
+    end
+
+    D1 -. "chunk 1 (25 codes)" .-> C1
+    D2 -. "chunk 2 (context + 25)" .-> C2
+    DN -. "final" .-> CN
+
+    stage0 --> FPL["⏱ First Packet Latency = Prefill + 25 decode steps only"]
+
+    style stage0 fill:#dae8fc,stroke:#6c8ebf
+    style stage1 fill:#e8d4f8,stroke:#8a6cad
+    style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px
+```
+
+Key parameters: `chunk_size=25`, `left_context_size=25` (validated defaults from Qwen3-TTS
+and Qwen3-Omni).
+
+## Directory Structure
+
+When adding a new TTS model, create the following structure:
+
+```
+vllm_omni/model_executor/models/
+  your_model_name/
+    __init__.py
+    your_model.py                    # Unified class (stage dispatch)
+    your_model_ar_stage.py           # Stage 0: AR stage
+    your_model_decoder.py            # Stage 1: audio decoder
+
+vllm_omni/model_executor/stage_input_processors/
+  your_model_name.py                 # Stage 0 -> Stage 1 transition
+
+vllm_omni/model_executor/stage_configs/
+  your_model_name.yaml               # Batch mode config
+  your_model_name_async_chunk.yaml   # Streaming mode config
+```
+
+**Qwen3-TTS reference files:**
+
+| File | Purpose |
+|------|---------|
+| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
+| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | Stage 0 - optimized AR |
+| `models/qwen3_tts/qwen3_tts_code2wav.py` | Stage 1 - decoder |
+| `stage_configs/qwen3_tts.yaml` | Stage config (async_chunk enabled) |
+| `stage_configs/qwen3_tts_batch.yaml` | Batch mode config |
+| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
+
+## Step-by-Step Implementation
+
+### Step 1: Implement Stage 0 - AR Stage
+
+Stage 0 is the autoregressive stage that generates intermediate audio representations.
+**It must use vLLM's native decoder layers with fused ops and PagedAttention** for the LLM
+backbone - this is the primary source of speedup over HuggingFace inference.
+
+#### 1.1 Use vLLM Decoder Layers Directly
+
+Build your transformer layers from the corresponding vLLM decoder layer class (e.g.
+`Qwen3DecoderLayer` for Qwen3-based backbones, or the equivalent for LLaMA, Qwen2, etc.).
+Do not wrap the HuggingFace model directly - that bypasses PagedAttention and fused kernels.
+
+```python
+# your_model_ar_stage.py
+
+from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer
+
+class YourTTSARStage(nn.Module):
+
+    def __init__(self, config, vllm_config, prefix):
+        self.layers = nn.ModuleList([
+            Qwen3DecoderLayer(
+                config, vllm_config=vllm_config, prefix=f"{prefix}.layers.{i}"
+            )
+            for i in range(config.num_hidden_layers)
+        ])
+        self.lm_head = ParallelLMHead(config.codec_size, config.hidden_size)
+```
+
+See `qwen3_tts_code_predictor_vllm.py` for the full implementation.
+
+#### 1.2 Forward Pass
+
+Implement `forward()` to return an `OmniOutput` with intermediate data for Stage 1:
+
+```python
+def forward(self, input_ids, positions, intermediate_tensors=None,
+            inputs_embeds=None, **kwargs) -> OmniOutput:
+    hidden_states = self.run_layers(input_ids, positions, intermediate_tensors, inputs_embeds)
+    logits = self.lm_head(hidden_states)
+
+    return OmniOutput(
+        text_hidden_states=hidden_states,
+        multimodal_outputs={
+            "audio_codes": self.extract_codes(logits),
+        },
+    )
+```
+
+The keys in `multimodal_outputs` are what your stage input processor will read to build
+Stage 1 inputs.
+
+#### 1.3 Weight Loading with Fused QKV
+
+When using vLLM's fused `QKVParallelLinear`, pack the HF `q_proj`/`k_proj`/`v_proj` weights
+into `qkv_proj` using `stacked_params_mapping`. See the `load_weights()` method in
+`qwen3_tts_code_predictor_vllm.py` for the standard pattern - it can be reused as-is
+for any Qwen-family backbone.
+
+#### 1.4 Custom Stop Condition (if needed)
+
+Some TTS models use a learned stop head rather than an EOS token. If your model does this,
+implement it inside `sample()`:
+
+```python
+def sample(self, logits, sampling_metadata) -> SamplerOutput | None:
+    output = self.sampler(logits, sampling_metadata)
+    if self._stop_head_fired():
+        output = mark_as_finished(output)
+    return output
+```
+
+### Step 2: Implement Stage 1 - Decoder
+
+Stage 1 decodes Stage 0 output into audio. It runs outside the scheduler (no PagedAttention
+needed). Implement `chunked_decode_streaming()` to support async_chunk streaming:
+
+```python
+# your_model_decoder.py
+
+class YourTTSDecoder(nn.Module):
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        # Initialize your audio decoder (SpeechTokenizer, HiFiGAN, etc.)
+
+    def forward(self, codes: torch.Tensor, **kwargs) -> torch.Tensor:
+        return self.decoder(codes)
+
+    def chunked_decode_streaming(self, codes, chunk_size=25,
+                                  left_context_size=25) -> torch.Tensor:
+        """Decode with a sliding context window for smooth chunk boundaries."""
+        end_index = codes.shape[-1]
+        context_size = 0 if end_index <= chunk_size else left_context_size
+        wav_chunk = self(codes)
+        # Trim left context to avoid duplicate audio
+        return wav_chunk[..., context_size * self.total_upsample:]
+```
+
+### Step 3: Implement the Unified Model Class
+
+The unified class dispatches to the correct stage based on `model_stage` in the config:
+
+```python
+# your_model.py
+
+class YourTTSModelForConditionalGeneration(nn.Module, SupportsPP):
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.model_stage = vllm_config.model_config.model_stage
+
+        if self.model_stage == "ar_stage":
+            ar_vllm_config = vllm_config.with_hf_config(
+                vllm_config.model_config.hf_config.ar_config,
+                architectures=["YourTTSARStageForConditionalGeneration"],
+            )
+            self.ar_stage = init_vllm_registered_model(
+                vllm_config=ar_vllm_config,
+                prefix=maybe_prefix(prefix, "ar"),
+                hf_config=ar_vllm_config.model_config.hf_config,
+                architectures=["YourTTSARStageForConditionalGeneration"],
+            )
+            self.model = self.ar_stage
+
+        elif self.model_stage == "decoder":
+            self.decoder = YourTTSDecoder(vllm_config=vllm_config, prefix=prefix)
+            self.model = self.decoder
+```
+
+### Step 4: Create `__init__.py`
+
+```python
+# vllm_omni/model_executor/models/your_model_name/__init__.py
+from .your_model import YourTTSModelForConditionalGeneration
+
+__all__ = ["YourTTSModelForConditionalGeneration"]
+```
+
+## Key Components
+
+### Model Interfaces
+
+Your unified model class should implement the appropriate interfaces:
+
+- **`SupportsPP`**: Required for pipeline parallelism support (all models should implement this)
+- **`SupportsMultiModal`**: Only if your model accepts multimodal inputs (e.g. reference audio for voice cloning)
+
+### Output Format
+
+Use `OmniOutput` so the orchestrator can route intermediate data between stages:
+
+```python
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+return OmniOutput(
+    text_hidden_states=hidden_states,
+    multimodal_outputs={
+        "audio_codes": codec_codes,
+    },
+)
+```
+
+### Weight Loading from a Single Checkpoint
+
+If both stages load from one checkpoint, separate them by prefix in the unified class:
+
+```python
+def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+    ar_weights, decoder_weights = [], []
+    for name, tensor in weights:
+        if name.startswith("decoder."):
+            decoder_weights.append((name, tensor))
+        else:
+            ar_weights.append((name, tensor))
+
+    if self.model_stage == "ar_stage":
+        return self.ar_stage.load_weights(ar_weights)
+    elif self.model_stage == "decoder":
+        return self.decoder.load_weights(decoder_weights)
+```
+
+## Model Registration
+
+Register all stage classes in `vllm_omni/model_executor/models/registry.py`:
+
+```python
+_OMNI_MODELS = {
+    # (package_name, module_name, class_name)
+    "YourTTSModelForConditionalGeneration": (
+        "your_model_name", "your_model",
+        "YourTTSModelForConditionalGeneration",
+    ),
+    "YourTTSARStageForConditionalGeneration": (
+        "your_model_name", "your_model_ar_stage",
+        "YourTTSARStageForConditionalGeneration",
+    ),
+    "YourTTSDecoder": (
+        "your_model_name", "your_model_decoder",
+        "YourTTSDecoder",
+    ),
+}
+```
+
+The registry uses lazy loading - model classes are only imported when needed.
+
+## Stage Configuration
+
+Each stage has a `worker_type` that determines how it is scheduled:
+
+- `worker_type: ar` - autoregressive stage, uses `OmniARScheduler` with PagedAttention
+- `worker_type: generation` - non-AR stage (e.g. decoder), uses `OmniGenerationScheduler`
+
+Key configuration fields:
+
+| Field | Description |
+|-------|-------------|
+| `model_stage` | Which stage to initialize (`ar_stage`, `decoder`, etc.) |
+| `model_arch` | Architecture name, must match `registry.py` |
+| `engine_input_source` | List of upstream stage IDs that provide input (e.g. `[0]`) |
+| `engine_output_type` | Output type: `latent` for intermediate, `audio` for final |
+| `custom_process_next_stage_input_func` | Async chunk processor function path (streaming only) |
+| `final_output` | Whether this stage produces the final user-facing output |
+| `final_output_type` | Type of final output (`audio`, `text`, etc.) |
+
+### Batch mode
+
+```yaml
+# stage_configs/your_model_name.yaml
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 64
+    engine_args:
+      model_stage: ar_stage
+      model_arch: YourTTSModelForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      engine_output_type: latent
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 2048
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+    engine_args:
+      model_stage: decoder
+      model_arch: YourTTSModelForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: audio
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: audio
+```
+
+### Streaming mode (async_chunk)
+
+Add `async_chunk: true` at the top level and specify `custom_process_next_stage_input_func`
+on Stage 0 to define how intermediate outputs are chunked and forwarded:
+
+```yaml
+# stage_configs/your_model_name_async_chunk.yaml
+
+async_chunk: true
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 64
+    engine_args:
+      model_stage: ar_stage
+      model_arch: YourTTSModelForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      engine_output_type: latent
+      custom_process_next_stage_input_func: >
+        vllm_omni.model_executor.stage_input_processors.your_model_name.ar2decoder_async_chunk
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 2048
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+    engine_args:
+      model_stage: decoder
+      model_arch: YourTTSModelForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: audio
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: audio
+```
+
+## Stage Input Processors
+
+Stage input processors convert Stage 0 outputs into Stage 1 inputs. Create yours in
+`vllm_omni/model_executor/stage_input_processors/your_model_name.py`.
+
+See `stage_input_processors/qwen3_tts.py` for the full reference implementation.
+
+### Data structures
+
+Understanding what's available in stage outputs:
+
+- `stage_list[source_id].engine_outputs` - list of `EngineCoreOutput` objects
+- Each `EngineCoreOutput` has `outputs` - list of `RequestOutput` objects
+- Each `RequestOutput` has:
+  - `token_ids` - generated token IDs
+  - `multimodal_output` - dict with keys matching your model's `OmniOutput.multimodal_outputs`
+  - `prompt_token_ids` - original prompt token IDs
+
+### Batch mode (non-streaming)
+
+Collects all Stage 0 outputs and forwards them to Stage 1 in one shot:
+
+```python
+def ar2decoder(
+    stage_list: list[Any],
+    engine_input_source: list[int],
+    prompt: OmniTokensPrompt | TextPrompt | None = None,
+    requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    source_id = engine_input_source[0]
+    decoder_inputs = []
+
+    for output in stage_list[source_id].engine_outputs:
+        result = output.outputs[0]
+        codes = result.multimodal_output["audio_codes"].cpu()
+        decoder_inputs.append(
+            OmniTokensPrompt(prompt_token_ids=codes.reshape(-1).tolist())
+        )
+
+    return decoder_inputs
+```
+
+### Streaming mode (async_chunk)
+
+Buffers Stage 0 outputs and forwards a chunk to Stage 1 once `chunk_size` frames
+have accumulated. The function signature follows the `OmniChunkTransferAdapter` protocol:
+
+```python
+def ar2decoder_async_chunk(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any] | None,
+    request: Any,
+    is_finished: bool = False,
+) -> dict[str, Any] | None:
+    """Forward chunks of AR output to the decoder stage."""
+    request_id = request.external_req_id
+    finished = bool(is_finished or request.is_finished())
+
+    # Extract and buffer the latest frame
+    if isinstance(pooling_output, dict):
+        frame = extract_frame(pooling_output)
+        if frame is not None:
+            transfer_manager.code_prompt_token_ids[request_id].append(
+                frame.cpu().tolist()
+            )
+    elif not finished:
+        return None
+
+    # Read chunk config from connector
+    chunk_size = 25
+    left_context_size = 25
+
+    length = len(transfer_manager.code_prompt_token_ids[request_id])
+    if length <= 0:
+        if finished:
+            return {"codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
+        return None
+
+    # Wait until a full chunk is ready (or request is finished)
+    chunk_length = length % chunk_size
+    if chunk_length != 0 and not finished:
+        return None
+
+    # Build context window: left_context + chunk
+    context_length = chunk_length if chunk_length != 0 else chunk_size
+    end_index = min(length, left_context_size + context_length)
+    window = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
+
+    return {
+        "codes": torch.tensor(window).transpose(0, 1).reshape(-1).tolist(),
+        "left_context_size": max(0, int(end_index - context_length)),
+        "finished": torch.tensor(finished, dtype=torch.bool),
+    }
+```
+
+Key points:
+- `transfer_manager` is the `OmniChunkTransferAdapter` that owns the chunk lifecycle
+- Each call appends one AR decode step's output; a chunk is emitted every `chunk_size` steps
+- The final (possibly partial) chunk is flushed when `is_finished` is true
+- `left_context_size` frames of overlap are included for smooth audio boundaries
+
+## Testing
+
+For general testing conventions, see [tests_style.md](../ci/tests_style.md).
+
+Recommended test cases for a new TTS model:
+
+1. **Single request** - verify waveform output shape and sample rate
+2. **Batched requests** - verify each request in the batch finishes independently
+3. **async_chunk streaming** - verify audio chunks arrive incrementally and decode correctly
+4. **Speaker conditioning** (if applicable) - verify different speaker inputs produce different outputs
+
+Reference test: `tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py`
+
+## Adding a Model Recipe
+
+After implementing and testing your model, add a model recipe to the
+[vllm-project/recipes](https://github.com/vllm-project/recipes) repository so users can
+get started quickly. See [Adding an Omni-Modality Model](./adding_omni_model.md#adding-a-model-recipe)
+for the expected format.
+
+## Summary
+
+Adding a TTS model to vLLM-Omni involves:
+
+1. **Create model directory** with AR stage, decoder stage, and unified class
+2. **AR stage** - use vLLM's native decoder layers with fused QKV; do not wrap HF directly
+3. **Decoder stage** - thin wrapper around your audio decoder; implement `chunked_decode_streaming()`
+4. **Unified class** - dispatches on `model_stage`; same structure as `Qwen3TTSModelForGeneration`
+5. **Register** all stage classes in `registry.py`
+6. **YAML configs** - provide both batch and `async_chunk` variants
+7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25
+8. **Tests** - cover single request, batching, and async_chunk streaming
+9. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
+
+### Qwen3-TTS Reference Files
+
+| File | Purpose |
+|------|---------|
+| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
+| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | AR stage with vLLM fused ops |
+| `models/qwen3_tts/qwen3_tts_code2wav.py` | Decoder stage with `chunked_decode_streaming()` |
+| `stage_configs/qwen3_tts.yaml` | Stage configuration |
+| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
+
+For more information, see:
+
+- [Architecture Overview](../../design/architecture_overview.md)
+- [Async Chunk Design](../../design/feature/async_chunk_design.md)
+- [Stage Configuration Guide](../../configuration/stage_configs.md)

docs/contributing/profiling.md

@@ -131,9 +131,6 @@ python image_to_video.py \
 
 2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**:   [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
 
-> **Note:**
-As of now, asynchronous (online) profiling is not fully supported in vLLM-Omni. While start_profile() and stop_profile() methods exist, they are only reliable in offline inference scripts (e.g., the provided end2end.py examples). Do not use them in server-mode or streaming scenarios—traces may be incomplete or fail to flush.
-
 ### 4. Analyzing Omni Traces
 
 Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.

docs/design/architecture_overview.md

@@ -92,6 +92,13 @@ The framework achieves high performance through several optimization techniques:
     * **Quantization:** Supports various quantization implementations including FP8 and AWQ.
     * **FusedOps:** Allows for custom and third-party integration.
 
+### Classifier-Free Guidance (CFG) Companion Flow
+
+vLLM-Omni natively models Classifier-Free Guidance (CFG) across disaggregated multi-stage setups via a "companion request" paradigm, eliminating redundant textual/multimodal context computation boundaries:
+1. **Prompt Expansion:** In the initial autoregressive (AR) stage, a customized `prompt_expand_func` hook intercepts incoming generation prompts and pairs them directly with negative companion prompts (e.g., a default negative prompt) on the fly, tagging the secondary prompt with a specific internal role (`cfg_text`).
+2. **Synchronized KV Cache Transfer:** The AR stage evaluates both the primary and companion sequence batches concurrently. The `OmniConnector` captures these specific structural dependencies and reliably passes the positive and negative outcome KV caches seamlessly across stage boundaries via shared memory or network protocols.
+3. **KV Cache Collection & Injection:** Upon reaching the downstream Diffusion (DiT) Engine, an assigned `cfg_kv_collect_func` automatically intercepts the mapped companion caches (`cfg_text_past_key_values`). These auxiliary dependencies are natively gathered and seamlessly bound to the primary generation sequence variables, enabling the DiT Engine to cleanly implement cross-attention CFG guidance over accurate conditioning and unconditioning structures in parallel.
+
 ### Flexibility and Usability
 
 vLLM-Omni is designed to be flexible and straightforward for users:

docs/design/feature/async_chunk_design.md

@@ -0,0 +1,209 @@
+# Async Chunk Design
+
+## Table of Contents
+
+1. [Overview](#overview)
+2. [Performance](#performance)
+3. [Architecture](#architecture)
+4. [Configuration](#configuration)
+5. [Related Files](#related-files)
+
+## Overview
+
+The `async_chunk` feature enables asynchronous, chunked processing of data across multiple stages in a multi-stage pipeline (e.g., Qwen3-Omni with Thinker → Talker → Code2Wav stages). Instead of waiting for a complete stage output before forwarding to the next stage, this feature allows stages to process and forward data in chunks as it becomes available, significantly reducing latency and improving throughput.
+
+**Chunk Size Definition**
+
+- **Prefill Phase**: `chunk_size = num_scheduled_tokens` for chunked prefill processing
+- **Decode Phase**: `chunk_size = num_scheduled_tokens = 1 ` for per-token streaming
+
+For qwen3-omni:
+- **Thinker → Talker**: Per decode step (typically chunk_size=1)
+- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. Set `initial_codec_chunk_frames` to emit smaller chunks during the initial phase for reduced TTFA
+- **Code2Wav**: Streaming decode with code2wav chunk_size
+
+With `async_chunk`:
+- Stages can start processing as soon as chunks are available
+- Overlapping execution across stages
+- Reduced latency and improved throughput
+- Better resource utilization
+- Async scheduling: Chunk IO (get/put) overlaps with compute via background threads so the scheduler is not blocked waiting for chunks
+
+## Performance
+1. **Reduced Latency**: Next stage can start processing immediately
+2. **Streaming Support**: Enables streaming for audio generation
+3. **IO-Compute Overlap**: Chunk retrieval happens asynchronously while other requests compute
+4. **Non-blocking Scheduler**: Requests waiting for chunks don't block the entire scheduler
+5. **Code2Wav Batch Inference**: Supports batched processing in code2wav stage
+
+| Input     | Output           | Async_chunk enabled | Code2Wav batch size | Max_Concurrency | Prompts | Mean E2E   | Mean TTFT  | Mean TPOT | Mean TTFP   | Mean RTF | Mean ITL |
+|-----------|------------------|---------------------|---------------------|----------------|---------|------------|------------|-----------|-------------|----------|----------|
+| text 100  | text 100+audio   | False               | 1                   | 1              | 50      | 6581.80    | 43.22      | 8.31      | 6459.34     | 0.24     | 8.22     |
+| text 100  | text 100+audio   | False               | 1                   | 4              | 50      | 7398.63    | 67.57      | 9.14      | 7285.35     | 0.27     | 9.05     |
+| text 100  | text 100+audio   | False               | 1                   | 10             | 50      | 13522.99   | 131.82     | 12.72     | 13410.44    | 0.49     | 12.60    |
+| text 100  | text 100+audio   | False               | 64                  | 1              | 50      | 6505.13    | 43.14      | 8.52      | 6395.40     | 0.24     | 8.44     |
+| text 100  | text 100+audio   | False               | 64                  | 4              | 50      | 7668.15    | 51.15      | 9.36      | 7562.37     | 0.28     | 9.27     |
+| text 100  | text 100+audio   | False               | 64                  | 10             | 50      | 9516.18    | 138.06     | 14.75     | 9409.26     | 0.34     | 14.60    |
+| text 100  | text 100+audio   | True                | 1                   | 1              | 50      | 6179.79    | 44.58      | 8.69      | 522.99      | 0.22     | 8.60     |
+| text 100  | text 100+audio   | True                | 1                   | 4              | 50      | 7692.69    | 103.96     | 10.22     | 785.85      | 0.29     | 10.12    |
+| text 100  | text 100+audio   | True                | 1                   | 10             | 50      | 11152.71   | 685.60     | 17.64     | 1628.88     | 0.41     | 17.62    |
+
+
+Performance data collected on H800 GPUs through comprehensive benchmarking with cudagraph enabled. text input uses random dataset.
+
+Enabling **async_chunk** (False→True) sharply reduces time-to-first-audio (TTFP)—e.g. ~92% at concurrency 1 (6.5s→0.52s)—and improves E2E latency (e.g. ~6% at conc 1, ~17% at conc 10). RTF (Real Time Factor) also improves with async_chunk on (e.g. ~8% at conc 1: 0.24→0.22, ~16% at conc 10: 0.49→0.41). Enabling **Code2Wav batch size 64** (vs 1) improves E2E and TTFP at higher concurrency when async_chunk is off (e.g. ~30% at conc 10: 13.5s→9.5s E2E, 13.4s→9.4s TTFP).
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_ttfp_performance.png">
+    <img alt="TTFP Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_ttfp_performance.png" width=100%>
+  </picture>
+</p>
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_rtf_performance.png">
+    <img alt="RTF Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_rtf_performance.png" width=100%>
+  </picture>
+</p>
+
+## Architecture
+### Data Flow
+
+#### Sequential Flow
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png">
+    <img alt="Data Flow between stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png" width=100%>
+  </picture>
+</p>
+
+#### Async Chunk Flow
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png">
+    <img alt="Data Flow between stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png" width=100%>
+  </picture>
+</p>
+
+### Async Chunk architecture
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/async-chunk-architecture.png">
+    <img alt="Async Chunk Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/async-chunk-architecture.png" width=100%>
+  </picture>
+</p>
+
+
+### Key Components
+
+1. **OmniConnector**: Inter-stage data transport only
+   - Shared memory or other IPC mechanisms
+   - **Transport-only API**: `put(from_stage, to_stage, put_key, data)` and `get(from_stage, to_stage, get_key)` (optionally with timeout)
+   - **No request-specific state**: Connector does not track put_requests, get_requests, request_payload, finished_requests, or other request-bound metadata; it only performs put/get operations
+   - Chunk keys and request/chunk lifecycle are managed by **OmniChunkTransferAdapter**
+
+2. **Transfer Adapter Layer**: Extensible abstraction for managing data transfer via connectors
+   - **OmniTransferAdapterBase**: Base class with background **recv_loop** and **save_loop** threads;
+   - **OmniChunkTransferAdapter**: Chunk-specific implementation that owns the full chunk lifecycle when async_chunk is enabled
+     - **Chunk ID and key construction**: Builds keys like `{req_id}_{stage_id}_{chunk_id}` for put/get
+     - **Async get**: `load_async(request)` enqueues the request; background **recv_loop** polls the connector (non-blocking); when data is available, updates the request and marks it in `_finished_load_reqs`; scheduler calls `get_finished_requests()` to learn which requests have chunks ready
+     - **Async put**: `save_async(pooling_output, request)` invokes `custom_process_next_stage_input_func` in the main thread to build the payload, then enqueues a save task; background **save_loop** performs `connector.put()`; payload processing and chunk accumulation (e.g. code2wav chunk_size) remain in the main thread
+
+3. **Stage Input Processors**: Custom functions that process stage outputs into chunks for different models
+   - Receive **transfer_manager** (OmniChunkTransferAdapter)
+   - Qwen3-omni reference: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
+
+4. **Schedulers**: Modified to handle chunk-based scheduling with async IO-compute overlap
+   - `OmniARScheduler`: For autoregressive stages
+   - `OmniGenerationScheduler`: For generation stages
+   - Both schedulers use **OmniChunkTransferAdapter** and **before/after** hooks around `super().schedule()`:
+     - **Before** `super().schedule()`: `process_pending_chunks(waiting, running)` moves requests waiting for chunks to `WAITING_FOR_CHUNK`, enqueues load tasks for background polling
+     - **After** `super().schedule()`: `restore_queues(waiting, running)` restores requests with ready chunks back to waiting/running, `postprocess_scheduler_output(scheduler_output)` attaches cached additional_information, clears chunk-ready flags
+   - **put_chunk** `save_async(pooler_output, request)`; **get_chunk** / **get_chunk_for_generation** `load_async(request)`
+
+5. **Model Runners**: Handle chunk processing
+   - `OmniGPUModelRunner`: Processes chunks in AR stages
+   - `GPUGenerationModelRunner`: Processes chunks in generation stages
+     - Uses `ubatch_slices` from `get_forward_context()` to track per-request sequence lengths in batched inference
+     - Reuses `ubatch_slices_padded` for code2wav batching to properly split batch outputs
+     - Handles list-type multimodal outputs: iterates through requests and assigns corresponding tensor to each
+     - Improved request state management: removes unscheduled and finished requests from input batch
+
+6. **Model Implementation**: Model-specific chunk handling
+   - `Qwen3OmniMoeForConditionalGeneration`: Main model with async_chunk support
+     - **Code2Wav stage batching**: Uses `ubatch_slices` to construct batched codec codes tensor `[batch_size, 16, max_seq_len]`
+     - **Batch output handling**: `generate_audio()` returns `list[torch.Tensor]`, one audio tensor per request
+     - **Multimodal outputs**: Returns list of audio tensors for batch processing instead of single concatenated tensor
+   - `Qwen3OmniCode2WavDecoder`: Audio generation model
+     - `chunked_decode()` and `chunked_decode_streaming()`: Return `list[torch.Tensor]` (one per request)
+     - Uses `ubatch_slices` to split batched waveform output into per-request audio chunks
+     - Each request gets correctly sized audio based on its code sequence length: `waveform[:, :, :code_seq_len * total_upsample]`
+
+7. **Request status**: `RequestStatus.WAITING_FOR_CHUNK` is added via patch (e.g. in `vllm_omni/patch.py`) so requests waiting for a chunk are not scheduled by the base vLLM scheduler until the chunk is ready.
+
+## Configuration
+
+Enable async_chunk in stage configuration YAML:
+
+```yaml
+async_chunk: true
+stage_args:
+  - stage_id: 0
+    engine_args:
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk
+  - stage_id: 1
+    engine_args:
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk
+```
+
+### Stage Configuration
+
+- `async_chunk: bool`: Enable/disable async chunk mode
+- `custom_process_next_stage_input_func: str`: Path to custom chunk processing function; receives `(transfer_manager, pooling_output, request)`. For qwen3-omni: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
+- `stage_connector_config: dict`: Connector configuration
+- `worker_type: str`: Model type, e.g. `"ar"` or `"generation"` (used by OmniChunkTransferAdapter for mode-specific payload handling)
+- `max_batch_size: int`: Maximum batch size for the stage
+
+
+### Connector Configuration
+
+```yaml
+connectors:
+  - from_stage: 0
+    to_stage: 1
+    spec:
+      name: SharedMemoryConnector
+      extra:
+        stage_id: 0
+```
+
+### Code2Wav Batch Configuration
+
+For optimal performance with async_chunk, the code2wav stage should be configured with batching:
+
+```yaml
+stage_args:
+  - stage_id: 2  # code2wav stage
+    runtime:
+      devices: "1"
+      max_batch_size: 64  # Enables batched audio generation
+    engine_args:
+      model_stage: code2wav
+```
+
+## Related Files
+
+- `vllm_omni/model_executor/stage_input_processors/qwen3_omni.py`: Chunk processing functions (receive `transfer_manager` as first param)
+- `vllm_omni/distributed/omni_connectors/transfer_adapter/base.py`: OmniTransferAdapterBase (recv_loop, save_loop, load_async, save_async)
+- `vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py`: OmniChunkTransferAdapter (process_pending_chunks, restore_queues, postprocess_scheduler_output)
+- `vllm_omni/distributed/omni_connectors/connectors/shm_connector.py`: SharedMemoryConnector (transport-only put/get)
+- `vllm_omni/core/sched/omni_ar_scheduler.py`: AR scheduler with chunk_transfer_adapter
+- `vllm_omni/core/sched/omni_generation_scheduler.py`: Generation scheduler with same async chunk pattern
+- `vllm_omni/worker/gpu_model_runner.py`: Model runner with chunk handling
+- `vllm_omni/worker/gpu_generation_model_runner.py`: Generation model runner with batch output handling and ubatch_slices support
+- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py`: Model implementation with code2wav batching
+- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py`: Code2wav decoder with batch support
+- `vllm_omni/engine/arg_utils.py`: Configuration definitions (async_chunk, worker_type)
+- `vllm_omni/config/model.py`: Model config with async_chunk field

docs/design/feature/hsdp.md

@@ -0,0 +1,146 @@
+# HSDP
+
+This section describes how to add HSDP (Hybrid Sharded Data Parallel) support to a diffusion transformer model. We use the Wan2.2 transformer as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Reference Implementations](#reference-implementations)
+
+---
+
+## Overview
+
+### What is HSDP?
+
+HSDP (Hybrid Sharded Data Parallel) is a memory optimization technique that **shards model weights** across multiple GPUs using PyTorch's FSDP2. Unlike Tensor Parallelism which splits computation, HSDP:
+
+- Shards weights across GPUs to reduce per-GPU memory usage
+- Gathers weights on-demand during forward passes
+- Can work standalone or combined with other parallelism (e.g., Sequence Parallel)
+
+This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
+
+**Important constraints:**
+- HSDP cannot be used with Tensor Parallelism
+- For standalone HSDP (no other parallelism), `hsdp_shard_size` must be specified explicitly
+
+### Architecture
+
+HSDP implementation relies on:
+
+1. **`_hsdp_shard_conditions`**: Model attribute specifying which modules to shard
+2. **`apply_hsdp_to_model`**: Function that applies FSDP2 sharding
+3. **`HSDPInferenceConfig`**: Runtime configuration for HSDP
+
+---
+
+## Step-by-Step Implementation
+
+### Step 1: Identify Modules to Shard
+
+Determine which modules in your transformer should be sharded. Typically, these are:
+
+- Transformer blocks (e.g., `blocks.0`, `blocks.1`, ...)
+- Large submodules with significant weight memory
+
+**Key questions:**
+- Which modules have the largest weights?
+- Which modules are repeated (like transformer blocks)?
+
+### Step 2: Define Shard Conditions
+
+Add `_hsdp_shard_conditions` to your model class. This is a list of functions that return `True` for modules that should be sharded.
+
+**Example (Transformer Blocks):**
+
+```python
+class MyTransformerModel(nn.Module):
+
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        """Match transformer blocks for HSDP sharding.
+
+        Args:
+            name: Module name from named_modules() (e.g., "blocks.0", "blocks.0.attn")
+            module: The module instance
+
+        Returns:
+            True if this module should be sharded
+        """
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block]
+```
+
+**Multiple Conditions Example:**
+
+```python
+class MyModel(nn.Module):
+
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    @staticmethod
+    def _is_moe_expert(name: str, module) -> bool:
+        # Also shard MoE expert layers
+        return "experts" in name and name.split(".")[-1].isdigit()
+
+    # Module is sharded if ANY condition returns True
+    _hsdp_shard_conditions = [_is_transformer_block, _is_moe_expert]
+```
+
+---
+
+## Testing
+
+After adding HSDP support, test with:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+parallel_config = DiffusionParallelConfig(
+    use_hsdp=True,
+    hsdp_shard_size=8,  # Shard across 8 GPUs
+)
+omni = Omni(model="your-model-name", parallel_config=parallel_config)
+
+output = omni.generate(
+    "a cup of coffee on the table",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Or via command line:**
+
+```bash
+vllm serve Your-org/your-model --omni --port 8091 --use-hsdp
+```
+
+**Verify:**
+
+1. Check logs for "HSDP Inference: replicate_size=..., shard_size=..."
+2. Check logs for "Sharded N modules + root"
+3. Verify memory usage is reduced proportionally
+4. Compare generated output quality with HSDP disabled
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Notes |
+|-------|------|-------|
+| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Reference implementation |
+| **HSDP Core** | `vllm_omni/diffusion/distributed/hsdp.py` | `apply_hsdp_to_model`, `shard_model` |
+| **HSDP Tests** | `tests/diffusion/distributed/test_hsdp.py` | Unit tests |
+
+---

docs/design/feature/sequence_parallel.md

@@ -240,7 +240,8 @@ class TransformerWithRoPE(nn.Module):
 
 ## Approach 2: Intrusive Modification (For Complex Cases)
 
-For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls.
+For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls. Importantly, when taking this approach, be careful to ensure that you correctly manage the `_sp_shard_depth`; if the sequence parallel shard depth is 0, Ulysses will not be used.
+
 
 **When to use:**
 - Dynamic/conditional sharding logic
@@ -252,15 +253,18 @@ from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
 
 def forward(self, hidden_states, ...):
     if self.parallel_config.sequence_parallel_size > 1:
+        # <Increment the _sp_shard_depth on the fwd context>
         hidden_states = sp_shard(hidden_states, dim=1)
 
     # ... computation ...
 
     if self.parallel_config.sequence_parallel_size > 1:
         output = sp_gather(output, dim=1)
+        # <Decrement the _sp_shard_depth on the fwd context>
 
     return output
 ```
+Note that currently, `sp_shard` / `sp_gather` do *not* automatically manage the `_sp_shard_depth`; you need to be careful to manage it yourself.
 
 ---
 

docs/design/feature/vae_parallel.md

@@ -0,0 +1,289 @@
+# VAE Patch Parallelism
+
+This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
+We use **Qwen-Image** as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is Vae Patch parallel?
+
+**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is:
+
++ Split into multiple spatial tiles
+
++ Distributed across multiple ranks
+
++ Decoded in parallel
+
++ Merged to reconstruct the final output
+
+This approach:
+
++ Distributes computation across multiple devices
+
++ Reduces peak memory usage per device
+
++ Accelerates decoding latency
+
+### Architecture
+We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding.
+
+The executor is model-agnostic and accepts three function parameters:
+
++ split – Partition the latent into tiles
+
++ exec – Decode a single tile
+
++ merge – Combine decoded tiles into the final output
+
+#### Execution Flow
+
++ Call split(z) to generate a list of TileTask and a GridSpec
+
++ Dispatch tasks across ranks using workload-based balancing
+
++ Each rank executes exec(task) on its assigned tiles
+
++ Gather decoded tile results to rank 0
+
++ Rank 0 performs merge(...)
+
++ (Optional) Broadcast final result to all ranks
+
+This design separates:
+
++ Distributed execution logic
+
++ Model-specific tiling and merging logic
+
+#### Why split / exec / merge is necessary?
+
+The latent tensor cannot be arbitrarily partitioned.
+
+During decoding:
+
++ Each output pixel may depend on neighboring pixels
+
++ The receptive field is model-dependent
+
+Therefore:
+
++ Tiles must include overlap
+
++ Merge must perform blending to avoid seams
+
+## Step-by-Step Implementation
+
+### Step 1: Implement DistributedAutoencoderKLQwenImage
+`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
+
+
+```
+class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+```
+**Key points**:
++ Inherit both AutoencoderKLQwenImage and DistributedVaeMixin
++ Call init_distributed() after loading weights
+
+### Step 2: Implement split/exec/merge
+Reuse `AutoencoderKLQwenImage.tiled_decode` logic and divide it into three stages. And we need return tiles with `GridSpec` and `TileTask`:
+```
+class GridSpec:
+    split_dims: tuple[int, ...]  # Tensor dimensions being split (e.g., (2, 3) for (B, C, H, W))
+    grid_shape: tuple[int, ...]  # Tile grid layout (num_rows, num_cols)
+    tile_spec: dict = field(default_factory=dict) # Metadata required for merging
+    output_dtype: torch.dtype | None = None # Final output dtype
+```
+```
+class TileTask:
+    tile_id: int # task id
+    grid_coord: tuple[int, ...]  # Tile position in grid
+    tensor: torch.Tensor | list[torch.Tensor]  # The tile tensor
+    workload: int | float = 1 # Used for load balancing (e.g., tile area)
+```
+And tiled base split/exec/merge as follow:
+```
+def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+    # mostly copy from AutoencoderKL
+    _, _, num_frames, height, width = z.shape
+    sample_height = height * self.spatial_compression_ratio
+    sample_width = width * self.spatial_compression_ratio
+
+    tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+    tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+    tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+    tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+    blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+    blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+    # Split z into overlapping tiles and decode them separately.
+    # The tiles have an overlap to avoid seams between tiles.
+    tiletask_list = []
+    for i in range(0, height, tile_latent_stride_height):
+        for j in range(0, width, tile_latent_stride_width):
+            time_list = []
+            for k in range(num_frames):
+                self._conv_idx = [0]
+                tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                time_list.append(tile)
+            tiletask_list.append(
+                TileTask(
+                    len(tiletask_list),
+                    (i // tile_latent_stride_height, j // tile_latent_stride_width),
+                    time_list,
+                    workload=time_list[0].shape[3] * time_list[0].shape[4],
+                )
+            )
+    tile_spec = {
+        "sample_height": sample_height,
+        "sample_width": sample_width,
+        "blend_height": blend_height,
+        "blend_width": blend_width,
+    }
+    grid_spec = GridSpec(
+        split_dims=(3, 4),
+        grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+        tile_spec=tile_spec,
+        output_dtype=self.dtype,
+    )
+    return tiletask_list, grid_spec
+
+def tile_exec(self, task: TileTask) -> torch.Tensor:
+    """Decode a single latent tile into RGB space."""
+    self.clear_cache()
+    time = []
+    for k in range(len(task.tensor)):
+        self._conv_idx = [0]
+        tile = self.post_quant_conv(task.tensor[k])
+        decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx)
+        time.append(decoded)
+    result = torch.cat(time, dim=2)
+    return result
+
+def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+    """Merge decoded tiles into a full image."""
+    grid_h, grid_w = grid_spec.grid_shape
+    result_rows = []
+    self.clear_cache()
+
+    result_rows = []
+    for i in range(grid_h):
+        result_row = []
+        for j in range(grid_w):
+            tile = coord_tensor_map[(i, j)]
+            if i > 0:
+                tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+            if j > 0:
+                tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+            result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width])
+        result_rows.append(torch.cat(result_row, dim=-1))
+    dec = torch.cat(result_rows, dim=3)[
+        :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
+    ]
+    return dec
+```
+
+### Step 3: Override tiled_decode
+We need to override tiled_decode, the main logic is:
++ check distributed is enabled
++ select split/exec/merge
++ Invoke self.distributed_decoder.execute to decode
+```
+def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
+    if not self.is_distributed_enabled():
+        return super().tiled_decode(z, return_dict=return_dict)
+
+    logger.info("Decode run with distributed executor")
+    result = self.distributed_decoder.execute(
+        z,
+        DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
+        broadcast_result=True,
+    )
+    if not return_dict:
+        return (result,)
+
+    return DecoderOutput(sample=result)
+```
+`broadcast_result` is set to True or False depending on the model; when enabled, the result will be used even on ranks other than 0.
+
+### Step 4: Modify Pipeline
+Change vae model from AutoencoderKLQwenImage to DistributedAutoencoderKLQwenImage
+```
+class YourModelPipeline(nn.Module):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        ...
+-       self.vae = AutoencoderKL.from_pretrained(
+-           model, subfolder="vae", local_files_only=local_files_only).to(self.device)
++       self.vae = DistributedAutoencoderKL.from_pretrained(
++           model, subfolder="vae", local_files_only=local_files_only
++       ).to(self.device)
+```
+
+## Testing
+Verify numerical consistency between:
++ vae_patch_parallel_size = 1
+
++ vae_patch_parallel_size = N
+
+Example:
+torch.allclose(output_1, output_n, atol=1e-5)
+
+Testing requirements:
++ Fix random seed
++ Use identical tiling strategy
+
+```python
+m = Omni(
+        model=model_name,
+        vae_use_tiling=True,
+        parallel_config=DiffusionParallelConfig(
+            tensor_parallel_size=2,
+            vae_patch_parallel_size=1, # or 2
+        ),
+    )
+```
+When vae_patch_parallel_size is larger than the DiT world size, it will automatically fall back to using the DiT world size instead.
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Notes |
+|-------|------|-------|
+| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
+| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
+| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage |
+
+---
+
+## Summary
+
+Adding Vae Patch Parallel support to diffusion model:
+
+1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
+2. **Change vae model in pipeline to Distributed Vae**
+3. **Test** - Verify with `tensor_parallel_size=N` quality

docs/features/custom_pipeline.md

@@ -115,7 +115,7 @@ python image_edit.py \
     --image cherry_blossom.jpg \
     --prompt "Let this mascot dance under the moon, surrounded by floating stars" \
     --output output_image_edit.png \
-    --num_inference_steps 10
+    --num-inference-steps 10
 ```
 
 ## Advanced Usage

docs/getting_started/installation/gpu.md

@@ -30,7 +30,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

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

@@ -20,20 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-# vllm 0.16.0 is still under prerelease
-uv pip install --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
-
-# vllm 0.16.0 may have some bugs for cuda 12.9, here is how we solve them:
-export FLASHINFER_CUDA_TAG="$(python3 -c 'import torch; print((torch.version.cuda or "12.4").replace(".", ""))')"
-
-uv pip install --upgrade --force-reinstall \
-  "flashinfer-python==0.6.3" \
-  "flashinfer-cubin==0.6.3" \
-  "flashinfer-jit-cache==0.6.3" \
-  --extra-index-url "https://flashinfer.ai/whl/cu${FLASHINFER_CUDA_TAG}"
-
-uv pip install --upgrade --force-reinstall "nvidia-cublas-cu12==12.9.1.4"
-uv pip install --upgrade --force-reinstall "numpy==2.2.6"
+uv pip install vllm --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -47,13 +34,13 @@ uv pip install vllm-omni
 # --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.16.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
 
 ```bash
-uv pip install vllm==0.16.0 --torch-backend=auto
+uv pip install vllm==0.17.0 --torch-backend=auto
 ```
 
-The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
+The release 0.17.0 of vLLM is based on PyTorch 2.10.0 which requires CUDA 12.9 environment.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -69,11 +56,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.16.0
+git checkout v0.17.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
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.16.0/vllm-0.16.0-cp38-abi3-manylinux_2_31_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -104,7 +91,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.14.0 \
+    vllm/vllm-omni:v0.16.0 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

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

@@ -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.16.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
+uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.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.16.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
 
 ```bash
-uv pip install vllm==0.16.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
+uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
 ```
 
-The release 0.16.0 of vLLM requires ROCm 7.0 environment.
+The release 0.17.0 of vLLM requires ROCm 7.0 environment.
 
 #### 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.16.0
+git checkout v0.17.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.14.0 \
+  vllm/vllm-omni-rocm:v0.16.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.14.0
+  vllm/vllm-omni-rocm:v0.16.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/installation/npu/npu.inc.md

@@ -33,13 +33,17 @@ docker run --rm \
     -p 8000:8000 \
     -it $IMAGE bash
 
+# Because vllm-ascend will release v0.16.0rc1 after vllm-omni 0.16.0,
+# we have to pin vllm-ascend at the current commit.
+cd /vllm-workspace/vllm-ascend
+git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+pip install -v -e .
+
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.16.0 https://github.com/vllm-project/vllm-omni.git
 
 cd vllm-omni
-sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
- && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
 pip install -v -e .
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
 ```
@@ -48,13 +52,13 @@ The default workdir is `/workspace`, with vLLM, vLLM-Ascend and vLLM-Omni code p
 
 For other installation methods (pip installation, building from source, custom Docker builds), please refer to the [vllm-ascend installation guide](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
 
-We are keeping [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there.
+We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there.
 
 # --8<-- [end:installation-release]
 
 # --8<-- [start:installation-main]
 
-You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) for the status of the latest commit of vLLM-Omni main branch on NPU.)
+You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
 
 ```bash
 # Pin vLLM version to 0.16.0
@@ -62,12 +66,14 @@ cd /vllm-workspace/vllm
 git pull origin main
 git fetch origin --tags
 git checkout v0.16.0
+VLLM_TARGET_DEVICE=empty pip install -v -e .
 
 # 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-ascend
 git pull origin main
 git fetch origin --tags
 git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+pip install -v -e .
 
 # Install vLLM-Omni from the latest main branch
 cd ../vllm-omni
@@ -94,10 +100,10 @@ Here's an example deployment command that has been verified on 4 x NPUs:
 
 ```bash
 # Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
+# export IMAGE=quay.io/ascend/vllm-omni:v0.16.0
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-omni:v0.14.0-a3
-export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
+# export IMAGE=quay.io/ascend/vllm-omni:v0.16.0-a3
+export IMAGE=quay.io/ascend/vllm-omni:v0.16.0
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \

docs/getting_started/quickstart.md

@@ -19,21 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-# vllm 0.16.0 is still under prerelease
-uv pip install --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
-# vllm 0.16.0 may have some bugs for cuda 12.9, here is how we solve them:
-export FLASHINFER_CUDA_TAG="$(python3 -c 'import torch; print((torch.version.cuda or "12.4").replace(".", ""))')"
-uv pip install --upgrade --force-reinstall \
-  "flashinfer-python==0.6.3" \
-  "flashinfer-cubin==0.6.3" \
-  "flashinfer-jit-cache==0.6.3" \
-  --extra-index-url "https://flashinfer.ai/whl/cu${FLASHINFER_CUDA_TAG}"
-uv pip install --upgrade --force-reinstall "nvidia-cublas-cu12==12.9.1.4"
-uv pip install --upgrade --force-reinstall "numpy==2.2.6"
-
+uv pip install vllm==0.17.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.16.0 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
+uv pip install vllm==0.17.0 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni

docs/mkdocs/hooks/generate_argparse.py

@@ -124,6 +124,7 @@ def extract_omni_serve_subparser_init():
                         "logger": logger,
                         "DummySubparsers": DummySubparsers,
                         "argparse": __import__("argparse"),
+                        "json": __import__("json"),
                         "DESCRIPTION": DESCRIPTION,
                     }
                     exec(code, exec_globals, local_vars)

docs/models/supported_models.md

@@ -30,10 +30,15 @@ th {
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
 | `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `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` |
+| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` |
+| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` |
+| `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` |
 |`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
 |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
+|`CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` |
+|`MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` |
 |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
 |`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
 |`OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` |
@@ -41,7 +46,9 @@ th {
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
+|`GlmImageForConditionalGeneration` | GLM-Image | `zai-org/GLM-Image` |
 |`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` |
+|`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` |
 
 
 ## List of Supported Models for NPU
@@ -64,7 +71,10 @@ th {
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` |
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
+| `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `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` |
 |`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
+|`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
 |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |

docs/serving/audio_generate_api.md

@@ -0,0 +1,338 @@
+# Audio Generate API
+
+vLLM-Omni provides an API for text-to-audio generation using diffusion-based models such as Stable Audio.
+
+Unlike the [Speech API](speech_api.md) which targets text-to-speech synthesis, the Audio Generate API is designed for general-purpose audio generation from text descriptions (sound effects, music, ambient soundscapes, etc.).
+
+Each server instance runs a single model (specified at startup via `vllm-omni serve <model> --omni`).
+
+## Quick Start
+
+### Start the Server
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
+```
+
+### Generate Audio
+
+**Using curl:**
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0
+    }' --output cat.wav
+```
+
+**Using Python:**
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8000/v1/audio/generate",
+    json={
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0,
+    },
+    timeout=300.0,
+)
+
+with open("cat.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/generate
+Content-Type: application/json
+```
+
+### Request Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text prompt describing the audio to generate |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
+
+#### Diffusion Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `audio_length` | float | null | Audio duration in seconds (default value is the max ~47s for `stable-audio-open-1.0`) |
+| `audio_start` | float | 0.0 | Audio start time in seconds |
+| `negative_prompt` | string | null | Text describing what to avoid in generation |
+| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
+| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
+| `seed` | int | null | Random seed for reproducible generation |
+
+### Response Format
+
+Returns binary audio data with the appropriate `Content-Type` header:
+
+| `response_format` | Content-Type |
+|--------------------|--------------|
+| `wav` | `audio/wav` |
+| `mp3` | `audio/mpeg` |
+| `flac` | `audio/flac` |
+| `pcm` | `audio/pcm` |
+| `aac` | `audio/aac` |
+| `opus` | `audio/opus` |
+
+## Examples
+
+### Basic Generation
+
+Generate audio with only a text prompt (model defaults for all other parameters):
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of ocean waves crashing on a beach"
+    }' --output ocean.wav
+```
+
+### Custom Duration
+
+Specify an explicit audio length in seconds:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A dog barking",
+        "audio_length": 5.0
+    }' --output dog_5s.wav
+```
+
+### High Quality with Negative Prompt
+
+Use a negative prompt to steer generation away from undesired characteristics, and increase inference steps for higher quality:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A piano playing a gentle melody",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted, noisy",
+        "guidance_scale": 8.0,
+        "num_inference_steps": 150
+    }' --output piano_hq.wav
+```
+
+### Reproducible Generation
+
+Set a `seed` to get deterministic results across runs:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "seed": 42
+    }' --output thunder.wav
+```
+
+### Full Control
+
+Combine all parameters for precise control over generation:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "negative_prompt": "Low quality",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+        "seed": 42
+    }' --output thunder_rain.wav
+```
+
+### Quick Generation (Fewer Steps)
+
+For faster generation with slightly lower quality:
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Birds chirping in a forest",
+        "audio_length": 8.0,
+        "num_inference_steps": 50
+    }' --output birds_quick.wav
+```
+
+### Python Client
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8000/v1/audio/generate",
+    json={
+        "input": "Thunder and rain",
+        "audio_length": 15.0,
+        "negative_prompt": "Low quality",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+        "seed": 42,
+        "response_format": "wav",
+    },
+    timeout=300.0,
+)
+
+with open("thunder.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Parameter Tuning Guide
+
+### `guidance_scale`
+
+Controls how closely the generated audio follows the text prompt.
+
+| Range | Behaviour |
+|-------|-----------|
+| 3 - 5 | More creative / varied output |
+| 7 (default) | Balanced adherence |
+| 10+ | Strict adherence to the prompt |
+
+### `num_inference_steps`
+
+Controls the number of denoising steps in the diffusion process.
+
+| Steps | Quality | Speed | Use Case |
+|-------|---------|-------|----------|
+| 50 | Good | Fast | Quick previews |
+| 100 | Very Good | Medium | General purpose |
+| 150+ | Excellent | Slow | Final / critical audio |
+
+### `audio_length`
+
+Duration of the generated audio clip. For `stable-audio-open-1.0`, the maximum is approximately 47 seconds. If omitted, the model uses its own default length.
+
+### `negative_prompt`
+
+Describes characteristics to avoid. Common negative prompts include:
+
+- `"Low quality, distorted, noisy"`
+- `"Silence, static"`
+- `"Music"` (when generating sound effects only)
+
+## Supported Models
+
+| Model | Description |
+|-------|-------------|
+| `stabilityai/stable-audio-open-1.0` | Open-source audio generation model, up to ~47 seconds, 44.1 kHz stereo |
+
+## Error Responses
+
+### 400 Bad Request
+
+Invalid or missing parameters:
+
+```json
+{
+    "error": {
+        "message": "Audio generation model did not produce audio output.",
+        "type": "BadRequestError",
+        "param": null,
+        "code": 400
+    }
+}
+```
+
+### 404 Not Found
+
+Model mismatch:
+
+```json
+{
+    "error": {
+        "message": "The model `xxx` does not exist.",
+        "type": "NotFoundError",
+        "param": "model",
+        "code": 404
+    }
+}
+```
+
+### 422 Unprocessable Entity
+
+Pydantic validation failure (e.g. invalid `response_format`, `speed` out of range):
+
+```json
+{
+    "detail": [
+        {
+            "type": "literal_error",
+            "msg": "Input should be 'wav', 'pcm', 'flac', 'mp3', 'aac' or 'opus'",
+            ...
+        }
+    ]
+}
+```
+
+## Troubleshooting
+
+### "Audio generation model did not produce audio output"
+
+The model finished but returned no audio data. Verify the server started successfully and the model loaded without errors.
+
+### Server Not Responding
+
+```bash
+# Check if the server is healthy
+curl http://localhost:8000/health
+```
+
+### Audio Quality Issues
+
+- Increase `num_inference_steps` (e.g. 150).
+- Add a negative prompt: `"Low quality, distorted, noisy"`.
+- Increase `guidance_scale` for stronger prompt adherence.
+
+### Generation Timeout
+
+- Reduce `num_inference_steps`.
+- Reduce `audio_length`.
+- Check GPU memory with `nvidia-smi`.
+
+### Out of Memory
+
+- Lower `--gpu-memory-utilization` (e.g. 0.8).
+- Reduce `audio_length`.
+
+## Development
+
+Enable debug logging:
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni \
+    --uvicorn-log-level debug
+```

docs/serving/speech_api.md

@@ -96,6 +96,7 @@ Content-Type: application/json
 | `language` | string | "Auto" | Language (see supported languages below) |
 | `instructions` | string | "" | Voice style/emotion instructions |
 | `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
+| `initial_codec_chunk_frames` | integer | null | Initial chunk size for reduced TTFA (overrides stage config) |
 
 **Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
 
@@ -103,7 +104,7 @@ Content-Type: application/json
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
-| `ref_audio` | string | null | Reference audio (URL or base64 data URL) |
+| `ref_audio` | string | null | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`) |
 | `ref_text` | string | null | Transcript of reference audio |
 | `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) |
 
@@ -124,6 +125,81 @@ Lists available voices for the loaded model.
     "voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
 }
 ```
+```
+POST /v1/audio/voices
+Content-Type: multipart/form-data
+```
+
+Upload a new voice sample for voice cloning in Base task TTS requests.
+
+**Form Parameters:**
+
+| Parameter | Type | Required | Description |
+|-----------|------|----------|-------------|
+| `audio_sample` | file | Yes | Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4) |
+| `consent` | string | Yes | Consent recording ID |
+| `name` | string | Yes | Name for the new voice |
+
+**Response Example:**
+
+```json
+{
+  "success": true,
+  "voice": {
+    "name": "custom_voice_1",
+    "consent": "user_consent_id",
+    "created_at": 1738660000,
+    "mime_type": "audio/wav",
+    "file_size": 1024000
+  }
+}
+```
+
+**Usage Example:**
+
+```bash
+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"
+```
+
+
+```bash
+DELETE /v1/audio/voices/{name}
+```
+
+Delete an uploaded voice sample.
+
+**Path Parameters:**
+
+| Parameter | Type | Required | Description |
+|-----------|------|----------|-------------|
+| `name` | string | Yes | Name of the voice to delete |
+
+**Response Example:**
+
+```json
+{
+  "success": true,
+  "message": "Voice 'custom_voice_1' deleted successfully"
+}
+```
+
+**Error Response (404 Not Found):**
+
+```json
+{
+  "success": false,
+  "error": "Voice 'unknown_voice' not found"
+}
+```
+
+**Usage Example:**
+
+```bash
+curl -X DELETE http://localhost:8091/v1/audio/voices/custom_voice_1
+```
 
 ## Examples
 
@@ -184,6 +260,25 @@ curl -X POST http://localhost:8091/v1/audio/speech \
     }' --output cloned.wav
 ```
 
+upload voice
+```bash
+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"
+```
+
+use upload voice
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, this is a cloned voice",
+        "task_type": "Base",
+        "voice": "custom_voice_1"
+    }' --output cloned.wav
+```
+
 ## Supported Models
 
 | Model | Task Type | Description |

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -16,6 +16,8 @@ The following parallelism methods are currently supported in vLLM-Omni:
 
 5. [VAE Patch Parallelism](#vae-patch-parallelism): VAE patch parallelism shards VAE decode spatially across ranks. This can reduce the peak memory of VAE decode and (depending on resolution and communication overhead) speed up VAE decode.
 
+6. [HSDP](#hsdp): Hybrid Sharded Data Parallel shards model weights across GPUs using PyTorch FSDP2. This reduces per-GPU memory usage, enabling inference of large models on GPUs with limited memory.
+
 The following table shows which models are currently supported by parallelism method:
 
 ### ImageGen
@@ -25,12 +27,12 @@ The following table shows which models are currently supported by parallelism me
 | **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |
 | **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |
 | **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |         ❌          |
-| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
+| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |         ✅          |
 | **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
 | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
 | **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
 | **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |         ✅          |
-| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |
+| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ✅        |         ✅          |
 | **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |
 | **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ✅       |        ✅        |         ❌          |
 
@@ -49,9 +51,10 @@ The following table shows which models are currently supported by parallelism me
 
 ### VideoGen
 
-| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel |
-|-------|------------------|------------|---------|--------------------------|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ |
+| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | VAE-Patch-Parallel |
+|-------|------------------|:----------:|:--------------:|:---------------:|:----:| :----:|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **LTX-2** | `Lightricks/LTX-2` | ✅ | ✅ | ✅ | ❌ | ❌ |
 
 ### Tensor Parallelism
 
@@ -83,7 +86,7 @@ outputs = omni.generate(
 VAE patch parallelism distributes the VAE decode workload across multiple ranks by splitting the latent spatially. It is configured via `DiffusionParallelConfig.vae_patch_parallel_size` and can be combined with other parallelism methods (e.g., TP).
 
 !!! note "Enablement and feature gate"
-    - VAE patch parallelism is currently **enabled only for validated pipelines** (currently: `Tongyi-MAI/Z-Image-Turbo`).
+    - VAE patch parallelism is currently **enabled only for validated pipelines** (check [ImageGen](#imagegen) and [VideoGen](#videogen) for more information).
     - If `vae_patch_parallel_size > 1` is set for a validated pipeline, vLLM-Omni will automatically enable `vae_use_tiling` as a safety gate. (We use `vae_use_tiling` because it indicates the VAE supports diffusers tiling parameters like `tile_latent_min_size` and `tile_overlap_factor`.)
 
 #### Offline Inference
@@ -330,3 +333,84 @@ You can enable CFG-Parallel in online serving for diffusion models via `--cfg-pa
 ```bash
 vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2
 ```
+
+### HSDP
+
+HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
+
+Unlike Tensor Parallelism which splits computation, HSDP uses PyTorch's FSDP2 to shard and redistribute weights at runtime. Each GPU only holds a fraction of the model weights, and weights are gathered on-demand during forward passes.
+
+#### Configuration
+
+HSDP is configured via `DiffusionParallelConfig`:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `use_hsdp` | bool | False | Enable HSDP |
+| `hsdp_shard_size` | int | -1 | Number of GPUs to shard weights across. -1 = auto (requires other parallelism > 1) |
+| `hsdp_replicate_size` | int | 1 | Number of replica groups. Each group holds a full sharded copy |
+
+**Constraints:**
+
+- `hsdp_replicate_size × hsdp_shard_size == world_size`
+- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
+
+#### Operating Modes
+
+HSDP can work in two modes:
+
+- **Standalone Mode**: HSDP alone without other parallelism. Must specify `hsdp_shard_size` explicitly.
+- **Combined Mode**: HSDP overlays on top of other parallelism (Ulysses Sequence Parallel, CFG Parallel). HSDP dimensions must match world_size.
+
+#### Offline Inference
+
+**Standalone HSDP** (shard across 4 GPUs, no other parallelism):
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        use_hsdp=True,
+        hsdp_shard_size=4,  # Shard across 4 GPUs
+    ),
+)
+
+outputs = omni.generate(
+    "A cat playing piano",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Combined HSDP + Sequence Parallel**:
+
+```python
+omni = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        ulysses_degree=4,  # Sequence parallel
+        use_hsdp=True,     # HSDP overlays on SP
+    ),
+)
+```
+
+#### Online Serving
+
+**Standalone HSDP** (shard model across 4 GPUs):
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --hsdp-shard-size 4
+```
+
+**Combined with Sequence Parallel**:
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --usp 4
+```
+
+#### Adding HSDP Support to New Models
+
+For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../design/feature/hsdp.md).

docs/user_guide/diffusion/quantization/fp8.md

@@ -62,6 +62,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
 |-------|-----------|---------------|------------------|
 | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
 | Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Skip sensitive layers | `img_mlp` |
+| Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
 
 ## Combining with Other Features
 

docs/user_guide/diffusion/quantization/gguf.md

@@ -0,0 +1,185 @@
+# GGUF Quantization
+
+## Goals
+1. Reuse vLLM quantization configs and weight loaders as much as possible.
+2. Add native GGUF support to diffusion transformers without changing model definitions.
+3. Keep user-facing knobs minimal and consistent across offline and online flows.
+
+## Scope
+1. Models: Z-Image, and Flux2-klein.
+2. Components: diffusion transformer weights, loader paths, and quantization configs.
+3. Modes: native GGUF (transformer-only weights).
+
+## Architecture Overview
+1. `OmniDiffusionConfig` accepts `quantization` or `quantization_config`.
+2. Diffusion quantization wrapper (`DiffusionGgufConfig`) produces vLLM `QuantizationConfig` objects for linear layers.
+3. `DiffusersPipelineLoader` branches on quantization method and loads either HF weights or GGUF weights for the transformer.
+4. GGUF transformer loading is routed through model-specific adapters (e.g., Flux2Klein).
+5. vLLM GGUF path uses `GGUFConfig` and `GGUFLinearMethod` for matmul.
+
+## Call Chain (Offline)
+```
+CLI (examples/offline_inference/text_to_image/text_to_image.py)
+  |
+  v
+Omni (vllm_omni/entrypoints/omni.py)
+  |
+  v
+OmniStage (diffusion)
+  |
+  v
+DiffusionWorker
+  |
+  v
+DiffusionModelRunner
+  |
+  v
+DiffusersPipelineLoader
+  |
+  v
+Pipeline.forward (Flux2/Qwen/Z-Image)
+  |
+  v
+DiffusionEngine
+  |
+  v
+OmniRequestOutput
+  |
+  v
+Client (saved PNG)
+```
+
+## Call Chain (Online)
+```
+Client
+  |
+  | POST /v1/images/generations
+  v
+APIServer (vllm_omni/entrypoints/openai/api_server.py)
+  |
+  v
+_generate_with_async_omni
+  |
+  v
+AsyncOmni
+  |
+  v
+DiffusionEngine
+  |
+  v
+OmniRequestOutput
+  |
+  v
+encode_image_base64
+  |
+  v
+ImageGenerationResponse
+  |
+  v
+Client
+```
+
+## Call Chain (GGUF Operator Path)
+```
+Pipeline.forward (Flux2/Qwen/Z-Image)
+  |
+  v
+Transformer blocks
+  |
+  v
+QKVParallelLinear / ColumnParallelLinear / RowParallelLinear
+  |
+  v
+LinearBase.forward
+  |
+  v
+QuantMethod.apply (GGUFLinearMethod.apply)
+  |
+  v
+fused_mul_mat_gguf
+  |
+  v
+_fused_mul_mat_gguf (custom op)
+  |
+  v
+ops.ggml_dequantize
+  |
+  v
+x @ weight.T
+```
+
+## GGUF Weight Loading Path (Transformer-Only)
+1. `DiffusersPipelineLoader.load_model` detects `quantization_config.method == "gguf"`.
+2. `gguf_model` is resolved as one of: local file, `repo/file.gguf`, or `repo:quant_type`.
+3. GGUF weights are routed through adapters in `vllm_omni/diffusion/model_loader/gguf_adapters/`.
+4. Name mapping is applied per-architecture (Z-Image, Flux2Klein).
+5. GGUF weights are loaded into transformer modules, remaining non-transformer weights come from the HF checkpoint.
+
+## GGUF Adapter Design
+1. `GGUFAdapter` is an abstract base class for model-specific adapters.
+2. `Flux2KleinGGUFAdapter` implements Flux2-Klein remapping + qkv split + adaLN swap.
+3. `ZImageGGUFAdapter` implements Z-Image qkv + ffn shard handling and linear qweight routing.
+4. `get_gguf_adapter(...)` strictly selects by model class/config; unsupported models raise an error (no fallback adapter).
+
+Adapter paths:
+- Base: `vllm_omni/diffusion/model_loader/gguf_adapters/base.py`
+- Z-Image: `vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py`
+- Flux2-Klein: `vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py`
+
+## User Usage (Offline)
+
+### Baseline BF16
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model /workspace/models/black-forest-labs/FLUX.2-klein-4B \
+  --prompt "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture" \
+  --height 768 \
+  --width 1360 \
+  --seed 42 \
+  --cfg-scale 4.0 \
+  --num-images-per-prompt 1 \
+  --num-inference-steps 4 \
+  --output outputs/flux2_klein_4b.png
+```
+
+### Native GGUF (Transformer Only)
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model /workspace/models/black-forest-labs/FLUX.2-klein-4B \
+  --gguf-model "/workspace/models/unsloth/FLUX.2-klein-4B-GGUF/flux-2-klein-4b-Q8_0.gguf" \
+  --quantization gguf \
+  --prompt "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture" \
+  --height 768 \
+  --width 1360 \
+  --seed 42 \
+  --cfg-scale 4.0 \
+  --num-images-per-prompt 1 \
+  --num-inference-steps 4 \
+  --output outputs/flux2_klein_4b_gguf.png
+```
+
+Notes for GGUF:
+1. Many GGUF repos do not ship `model_index.json` and configs. Use the base repo for `--model` and only pass the GGUF file via `--gguf-model`.
+2. `gguf_model` supports local path, `repo/file.gguf`, or `repo:quant_type`.
+
+## User Usage (Online)
+
+### Start Server (Native GGUF via CLI)
+```bash
+vllm serve /workspace/models/black-forest-labs/FLUX.2-klein-4B \
+  --omni \
+  --port 8000 \
+  --quantization-config '{"method":"gguf","gguf_model":"/workspace/models/unsloth/FLUX.2-klein-4B-GGUF/flux-2-klein-4b-Q8_0.gguf"}'
+```
+
+### Online Request (Images API)
+```bash
+curl -X POST http://localhost:8000/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "prompt": "a dragon laying over the spine of the Green Mountains of Vermont",
+    "size": "1024x1024",
+    "seed": 42,
+    "num_inference_steps": 4
+  }'
+```

docs/user_guide/diffusion/quantization/overview.md

@@ -7,6 +7,7 @@ vLLM-Omni supports quantization of DiT linear layers to reduce memory usage and
 | Method | Guide |
 |--------|-------|
 | FP8 | [FP8](fp8.md) |
+| GGUF | [GGUF](gguf.md) |
 
 ## Device Compatibility
 

docs/user_guide/diffusion_acceleration.md

@@ -30,6 +30,8 @@ vLLM-Omni also supports parallelism methods for diffusion models, including:
 
 5. [VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism) - shards VAE decode/encode spatially across ranks to reduce VAE peak memory (and can speed up VAE decode).
 
+6. [HSDP](diffusion/parallelism_acceleration.md#hsdp) - Hybrid Sharded Data Parallel shards model weights across GPUs to reduce per-GPU memory usage, enabling inference of large models on limited GPU memory.
+
 ## Quick Comparison
 
 ### Cache Methods
@@ -56,8 +58,8 @@ The following table shows which models are currently supported by each accelerat
 | **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
 | **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
 | **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
-| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
 | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
 | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
@@ -65,13 +67,15 @@ The following table shows which models are currently supported by each accelerat
 | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
 | **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
-| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ |
+| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
+| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
 
 ### VideoGen
 
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel |
-|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:----------------:|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ |
+| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | VAE-Patch-Parallel |
+|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 
 ### Quantization
 
@@ -265,6 +269,30 @@ outputs = omni.generate(
 )
 ```
 
+### Using HSDP
+
+HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        use_hsdp=True,
+        hsdp_replicate_size=1,  # Number of replica groups
+        hsdp_shard_size=8,      # Number of GPUs to shard across (or -1 for auto)
+    ),
+)
+
+outputs = omni.generate(
+    "A cat playing piano",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
 ### Using CFG-Parallel
 
 Run image-to-image:
@@ -319,3 +347,4 @@ For detailed information on each acceleration method:
 - **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
 - **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.
 - **[VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism)** - Guidance on how to reduce VAE memory via patch/tile parallelism.
+- **[HSDP](diffusion/parallelism_acceleration.md#hsdp)** - Hybrid Sharded Data Parallel for memory-efficient inference of large models.

docs/user_guide/examples/offline_inference/bagel.md

@@ -154,6 +154,24 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
 
 ------
 
+#### Tensor Parallelism (TP)
+
+For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) by modifying the stage configuration (e.g., [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)).
+
+1. **Set `tensor_parallel_size`**: Increase this value (e.g., to `2` or `4`).
+2. **Set `devices`**: Specify the comma-separated GPU IDs to be used for the stage (e.g., `"0,1"`).
+
+Example configuration for TP=2 on GPUs 0 and 1:
+```yaml
+    engine_args:
+      tensor_parallel_size: 2
+      ...
+    runtime:
+      devices: "0,1"
+```
+
+------
+
 #### 🔗 Runtime Configuration
 
 | Parameter             | Value   | Description                      |
@@ -162,6 +180,74 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
 | `max_inflight`        | `1`     | Maximum inflight requests        |
 | `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB)   |
 
+## Using Mooncake Connector
+
+[Mooncake](https://github.com/kvcache-ai/Mooncake) is a high-performance distributed KV cache transfer engine that enables efficient cross-node data movement via TCP or RDMA, making it ideal for multi-node disaggregated inference.
+
+By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can switch to the Mooncake connector for better performance on multi-GPU setups and to enable multi-node deployment.
+
+### Prerequisites
+
+Install the Mooncake transfer engine:
+
+```bash
+# For CUDA-enabled systems (recommended)
+pip install mooncake-transfer-engine
+
+# For non-CUDA systems
+pip install mooncake-transfer-engine-non-cuda
+```
+
+### Step 1: Start the Mooncake Master
+
+On the **primary node**, start the Mooncake master service (run in a separate terminal or background with `&`):
+
+```bash
+# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
+# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
+mkdir -p ./mc_storage
+
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003 \
+  --root_fs_dir=./mc_storage/ \
+  --cluster_id=mc-local-1 &
+```
+
+### Step 2: Run Offline Inference with Mooncake
+
+Use the provided Mooncake stage config [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml). Before launching, update the `metadata_server` and `master` addresses in the YAML to match your Mooncake master node's IP (use `127.0.0.1` for single-node testing).
+
+```bash
+cd examples/offline_inference/bagel
+
+# Text to Image with Mooncake
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2img \
+                  --prompts "A cute cat" \
+                  --stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+
+# Image to Text with Mooncake
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality img2text \
+                  --image-path /path/to/image.jpg \
+                  --prompts "Describe this image" \
+                  --stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+
+# Text to Text with Mooncake
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2text \
+                  --prompts "What is the capital of France?" \
+                  --stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+```
+
+For more details on the Mooncake connector and multi-node setup, see the [Mooncake Store Connector documentation](https://github.com/vllm-project/vllm-omni/tree/main/docs/design/feature/omni_connectors/mooncake_store_connector.md).
+
+------
+
 ## FAQ
 
 - If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.

docs/user_guide/examples/offline_inference/glm_image.md

@@ -0,0 +1,156 @@
+# GLM-Image Multistage End-to-End Inference
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/glm_image>.
+
+
+This example demonstrates how to run GLM-Image with the vLLM-Omni multistage architecture.
+
+## Architecture
+
+GLM-Image uses a 2-stage pipeline:
+
+```
+┌─────────────────────────────────────────────────────────────┐
+│                     GLM-Image Pipeline                       │
+├─────────────────────────────────────────────────────────────┤
+│                                                             │
+│  Stage 0 (AR Model)              Stage 1 (Diffusion)        │
+│  ┌─────────────────┐            ┌─────────────────────┐     │
+│  │ vLLM-optimized  │            │  GlmImagePipeline   │     │
+│  │ GlmImageFor     │  prior     │  ┌───────────────┐  │     │
+│  │ Conditional     │──tokens───►│  │ DiT Denoiser  │  │     │
+│  │ Generation      │            │  └───────────────┘  │     │
+│  │ (9B AR model)   │            │         │          │     │
+│  └─────────────────┘            │         ▼          │     │
+│         ▲                       │  ┌───────────────┐  │     │
+│         │                       │  │  VAE Decode   │──┼──► Image
+│    Text/Image                   │  └───────────────┘  │     │
+│      Input                      └─────────────────────┘     │
+│                                                             │
+└─────────────────────────────────────────────────────────────┘
+```
+
+## Features
+
+- **vLLM-optimized AR**: Uses PagedAttention and tensor parallelism for faster prior token generation
+- **Flexible deployment**: AR and Diffusion stages can run on different GPUs
+- **Text-to-Image**: Generate images from text descriptions
+- **Image-to-Image**: Edit existing images with text prompts
+
+## Usage
+
+### Text-to-Image
+
+```bash
+python end2end.py \
+    --model-path /path/to/glm-image \
+    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "A beautiful sunset over the ocean with sailing boats" \
+    --height 1024 \
+    --width 1024 \
+    --output output_t2i.png
+```
+
+### Image-to-Image (Image Editing)
+
+```bash
+python end2end.py \
+    --model-path /path/to/glm-image \
+    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "Transform this scene into a winter wonderland" \
+    --image input.png \
+    --output output_i2i.png
+```
+
+### With Custom Parameters
+
+```bash
+python end2end.py \
+    --model-path /path/to/glm-image \
+    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "A photorealistic cat sitting on a window sill" \
+    --height 1024 \
+    --width 1024 \
+    --num-inference-steps 50 \
+    --guidance-scale 1.5 \
+    --seed 42 \
+    --output output.png
+```
+
+## Shell Scripts
+
+### Run Text-to-Image
+
+```bash
+./run_t2i.sh
+```
+
+### Run Image-to-Image
+
+```bash
+./run_i2i.sh --image /path/to/input.png
+```
+
+## Stage Configuration
+
+The stage config (`glm_image.yaml`) defines:
+
+- **Stage 0 (AR)**: Uses `GPUARWorker` with vLLM engine
+
+  - Model: `GlmImageForConditionalGeneration`
+  - Output: `token_ids` (prior tokens)
+
+- **Stage 1 (Diffusion)**: Uses diffusion engine
+  - Model: `GlmImagePipeline`
+  - Output: Generated image
+
+See `vllm_omni/model_executor/stage_configs/glm_image.yaml` for full configuration.
+
+## Comparison with Single-Stage
+
+| Aspect      | Single-Stage (transformers) | Multistage (vLLM)   |
+| ----------- | --------------------------- | ------------------- |
+| AR Model    | transformers native         | vLLM PagedAttention |
+| Memory      | Higher (no KV cache opt)    | Lower (optimized)   |
+| Throughput  | Lower                       | Higher              |
+| Flexibility | Single GPU                  | Multi-GPU support   |
+
+## Troubleshooting
+
+### OOM Error
+
+Try reducing memory usage:
+
+```bash
+# In glm_image.yaml, adjust:
+gpu_memory_utilization: 0.5  # Reduce from 0.6
+```
+
+### Slow Initialization
+
+The first run loads model weights. Subsequent runs are faster:
+
+```bash
+--stage-init-timeout 900  # Increase timeout for slow storage
+```
+
+## Requirements
+
+- vLLM-Omni with GLM-Image support
+- CUDA-capable GPU (recommended: H100/A100 with 80GB)
+- GLM-Image model weights
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/glm_image/end2end.py"
+    ``````
+??? abstract "run_i2i.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/glm_image/run_i2i.sh"
+    ``````
+??? abstract "run_t2i.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/glm_image/run_t2i.sh"
+    ``````

docs/user_guide/examples/offline_inference/image_to_video.md

@@ -69,6 +69,11 @@ Key arguments:
 - `--vae-use-tiling`: Enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+- `--use-hsdp`: Enable Hybrid Sharded Data Parallel to shard model weights across GPUs.
+- `--hsdp-shard-size`: Number of GPUs to shard model weights across within each replica group. -1 (default) auto-calculates as world_size / replicate_size.
+- `--hsdp-replicate-size`: Number of replica groups for HSDP. Each replica holds a full sharded copy. Default 1 means pure sharding (no replication).
+
+
 
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
 

docs/user_guide/examples/offline_inference/qwen3_tts.md

@@ -11,7 +11,7 @@ Qwen3 TTS provides multiple task variants for speech generation:
 
 - **CustomVoice**: Generate speech with a known speaker identity (speaker ID) and optional instruction.
 - **VoiceDesign**: Generate speech from text plus a descriptive instruction that designs a new voice.
-- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection.
+- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection. The `ref_audio` field accepts a local file path, HTTP URL, or base64 data URL.
 
 ## Setup
 Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
@@ -90,13 +90,44 @@ Examples:
 python end2end.py --query-type Base --mode-tag icl
 ```
 
+## Streaming Mode
+
+Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
+
+```bash
+python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_stream
+```
+
+Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames`
+and `initial_codec_chunk_frames` in the stage config). The final WAV file is written once generation
+completes. This demonstrates that audio data is available progressively rather than only at the end.
+
+> **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.
+
+## Batched Decoding
+
+The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
+
+```
+python end2end.py --query-type CustomVoice \
+    --txt-prompts benchmark_prompts.txt \
+    --batch-size 4 \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
+```
+
+**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_batch_size >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_batch_size`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
+
 ## Notes
 
 - The script uses the model paths embedded in `end2end.py`. Update them if your local cache path differs.
-- Use `--output-dir` (preferred) or `--output-wav` to change the output folder.
+- Use `--output-dir` to change the output folder.
 
 ## Example materials
 
+??? abstract "benchmark_prompts.txt"
+    ``````txt
+    --8<-- "examples/offline_inference/qwen3_tts/benchmark_prompts.txt"
+    ``````
 ??? abstract "end2end.py"
     ``````py
     --8<-- "examples/offline_inference/qwen3_tts/end2end.py"

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -3,7 +3,7 @@
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_image>.
 
 
-This folder provides several entrypoints for experimenting with `Qwen/Qwen-Image` `Qwen/Qwen-Image-2512` `Tongyi-MAI/Z-Image-Turbo` using vLLM-Omni:
+This folder provides several entrypoints for experimenting with `Qwen/Qwen-Image` `Qwen/Qwen-Image-2512` `Tongyi-MAI/Z-Image-Turbo` `stepfun-ai/NextStep-1.1` using vLLM-Omni, note that NextStep-1.1 has different architecture so we treat it differently regarding running arguments and pipeline.
 
 - `text_to_image.py`: command-line script for single image generation with advanced options.
 - `web_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
@@ -77,6 +77,8 @@ if __name__ == "__main__":
 
 ## Local CLI Usage
 
+### Qwen/Tongyi Models
+
 ```bash
 python text_to_image.py \
   --model Tongyi-MAI/Z-Image-Turbo \
@@ -90,7 +92,26 @@ python text_to_image.py \
   --output outputs/coffee.png
 ```
 
-Key arguments:
+### NextStep Models
+
+NextStep-1.1 can have extra arguments
+```bash
+python text_to_image.py \
+  --model stepfun-ai/NextStep-1.1 \
+  --prompt "A baby panda wearing an Iron Man mask, holding a board with 'NextStep-1' written on it" \
+  --height 512 \
+  --width 512 \
+  --num-inference-steps 28 \
+  --guidance-scale 7.5 \
+  --guidance-scale-2 1.0 \
+  --cfg-schedule constant \
+  --output nextstep_output.png \
+  --seed 42
+```
+
+### Key Arguments
+
+**Common arguments:**
 
 - `--prompt`: text description (string).
 - `--seed`: integer seed for deterministic sampling.
@@ -101,8 +122,15 @@ Key arguments:
 - `--output`: path to save the generated PNG.
 - `--vae-use-slicing`: enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: enable VAE tiling for memory optimization.
-- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion_acceleration.md#using-cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+- `--guidance-scale`: classifier-free guidance scale.
+
+**NextStep-1.1 specific:**
+- `--guidance-scale-2`: secondary guidance scale, e.g. image-level CFG (default: 1.0).
+- `--timesteps-shift`: timesteps shift parameter for sampling (default: 1.0).
+- `--cfg-schedule`: CFG schedule type, "constant" or "linear" (default: "constant").
+- `--use-norm`: apply layer normalization to sampled tokens.
 
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
 

docs/user_guide/examples/online_serving/bagel.md

@@ -25,12 +25,152 @@ cd /workspace/vllm-omni/examples/online_serving/bagel
 bash run_server.sh
 ```
 
-If you have a custom stage configs file, launch the server with the command below:
-
 ```bash
 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
 
+#### 🚀 Tensor Parallelism (TP)
+
+For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) for the server.
+
+1. **Modify Stage Config**: Create or modify a stage configuration yaml (e.g., [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)). Set `tensor_parallel_size` to `2` (or more) and update `devices` to include multiple GPU IDs (e.g., `"0,1"`).
+
+```yaml
+    engine_args:
+      tensor_parallel_size: 2
+      ...
+    runtime:
+      devices: "0,1"
+```
+
+2. **Launch Server**:
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/your/custom_bagel.yaml
+```
+
+#### Using Mooncake Connector
+
+By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can use the [Mooncake](https://github.com/kvcache-ai/Mooncake) connector to transfer KV cache between stages, which also enables multi-node deployment.
+
+**1. Install Mooncake**
+
+```bash
+# For CUDA-enabled systems (recommended)
+pip install mooncake-transfer-engine
+
+# For non-CUDA systems
+pip install mooncake-transfer-engine-non-cuda
+```
+
+**2. Start Mooncake Master** on the primary node:
+
+```bash
+# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
+# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
+mkdir -p ./mc_storage
+
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003 \
+  --root_fs_dir=./mc_storage/ \
+  --cluster_id=mc-local-1 &
+```
+
+**3. Launch the server** with the Mooncake stage config:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+```
+
+> **Note**: Before launching, edit [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. For single-node testing, `127.0.0.1` works.
+
+The client-side usage is identical to the default setup -- the Mooncake connector is transparent to the API. See the requests section below.
+
+For more details on the Mooncake connector configuration, see the [Mooncake Store Connector documentation](https://github.com/vllm-project/vllm-omni/tree/main/docs/design/feature/omni_connectors/mooncake_store_connector.md).
+
+#### Multi-Node Deployment
+
+You can deploy each stage on a **separate node** for better resource utilization. In this example, the orchestrator (Stage 0 / Thinker) and Stage 1 (DiT) run on different machines, connected via Mooncake.
+
+Replace `<ORCHESTRATOR_IP>` below with the actual IP address of your orchestrator node (e.g., `10.244.227.244`).
+
+> [!WARNING]
+> **Before launching**, edit [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. Mismatched addresses will cause silent connection failures.
+
+**1. Start Mooncake Master** (on the orchestrator node):
+
+```bash
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=<ORCHESTRATOR_IP> \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003
+```
+
+**2. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
+    --port 8000 \ # API server port for client requests
+    --stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
+    --stage-id 0 \
+    -oma <ORCHESTRATOR_IP> \
+    -omp 8091
+```
+
+**3. Launch Stage 1 (DiT)** on the remote node in headless mode:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
+    --stage-id 1 \
+    --headless \
+    -oma <ORCHESTRATOR_IP> \
+    -omp 8091
+```
+
+**Mooncake Master arguments:**
+
+| Argument | Description |
+| :------- | :---------- |
+| `--rpc_port` | Mooncake RPC port for control-plane coordination between stages |
+| `--enable_http_metadata_server` | Enable the HTTP metadata server for service discovery |
+| `--http_metadata_server_host` | IP address to bind the metadata server (use the orchestrator node's IP) |
+| `--http_metadata_server_port` | Port for the HTTP metadata server |
+| `--metrics_port` | Port for Prometheus-compatible metrics endpoint |
+
+**vllm serve arguments:**
+
+| Argument | Description |
+| :------- | :---------- |
+| `--stage-id` | Which stage this process runs (0 = Thinker, 1 = DiT) |
+| `--headless` | Run without the API server (worker-only mode) |
+| `-oma` | Orchestrator master address |
+| `-omp` | Orchestrator master port for Stage 1 to connect to Stage 0 for task coordination |
+
+> [!IMPORTANT]
+> **Startup Order**: Stage 0 (orchestrator) must be launched **before** Stage 1 (headless).
+> Stage 0 will appear to hang on startup until Stage 1 (worker) connects — this is expected behavior.
+
+**Network Requirements**
+
+All nodes must have network connectivity to each other. Ensure the following ports are open **between all participating nodes**:
+
+| Port | Protocol | Service | Direction |
+| :--- | :------- | :------ | :-------- |
+| 50051 | TCP | Mooncake Master RPC | Worker → Orchestrator |
+| 8080 | TCP | Mooncake HTTP Metadata Server | Worker → Orchestrator |
+| 8091 | TCP | Orchestrator Master (`-omp`) | Worker → Orchestrator |
+| 8000 | TCP | API Server (`--port`) | Client → Orchestrator |
+| 9003 | TCP | Metrics (optional) | Monitoring → Orchestrator |
+
+> **Tip**: If nodes are behind a firewall or in different VPCs/security groups, make sure the above ports are allowed in ingress/egress rules. All nodes should be reachable via their IP addresses (no NAT). Using nodes on the same subnet or VPC is recommended to minimize latency for Mooncake KV cache transfers.
+
 ### Send Multi-modal Request
 
 Get into the bagel folder:
@@ -242,3 +382,7 @@ sudo apt install ffmpeg
     ``````sh
     --8<-- "examples/online_serving/bagel/run_server.sh"
     ``````
+??? abstract "run_server_stage_cli.sh"
+    ``````sh
+    --8<-- "examples/online_serving/bagel/run_server_stage_cli.sh"
+    ``````

docs/user_guide/examples/online_serving/glm_image.md

@@ -0,0 +1,269 @@
+# GLM-Image Online Serving
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/glm_image>.
+
+
+This example demonstrates how to deploy GLM-Image for online image generation using vLLM-Omni.
+
+## 🛠️ Installation
+
+Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
+
+## Run examples (GLM-Image)
+
+**Note**: These examples work with the default configuration on **2× NVIDIA A100 (80GB)** or equivalent. Stage 0 (AR) and Stage 1 (Diffusion) each use one GPU by default. For single-GPU setups, modify the stage configuration to share the same device.
+
+### Launch the Server
+
+```bash
+# Use default configuration
+vllm serve zai-org/GLM-Image --omni --port 8091
+```
+
+Or use the convenience script:
+
+```bash
+cd examples/online_serving/glm_image
+bash run_server.sh
+```
+
+If you have a custom stage configs file:
+
+```bash
+vllm serve zai-org/GLM-Image --omni --port 8091 --stage-configs-path /path/to/glm_image.yaml
+```
+
+### Send Requests
+
+Get into the glm_image folder:
+
+```bash
+cd examples/online_serving/glm_image
+```
+
+Send request via Python:
+
+```bash
+python openai_chat_client.py --prompt "A cute cat sitting on a window sill"
+```
+
+The Python client supports the following command-line arguments:
+
+- `--prompt` (or `-p`): Text prompt for generation (default: `A beautiful sunset over the ocean with sailing boats`)
+- `--output` (or `-o`): Output file path (default: `glm_image_output.png`)
+- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
+- `--image` (or `-i`): Input image path (for image-to-image editing)
+- `--height`: Image height in pixels (default: 1024)
+- `--width`: Image width in pixels (default: 1024)
+- `--steps`: Number of inference steps (default: 50)
+- `--guidance-scale`: Classifier-free guidance scale (default: 1.5)
+- `--seed`: Random seed (default: 42)
+- `--negative`: Negative prompt
+
+## Modality Control
+
+GLM-Image supports **text-to-image** and **image-to-image** modes.
+
+The default yaml configuration deploys AR on GPU 0 and DiT on GPU 1. You can use the default configuration file: [`glm_image.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/glm_image.yaml)
+
+| Mode           | Input        | Output | Description                        |
+| -------------- | ------------ | ------ | ---------------------------------- |
+| Text-to-Image  | Text         | Image  | Generate images from text prompts  |
+| Image-to-Image | Image + Text | Image  | Edit images with text instructions |
+
+### Text-to-Image
+
+Generate images from text prompts:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "A photorealistic mountain landscape at sunset" \
+    --height 1024 \
+    --width 1024 \
+    --output landscape.png
+```
+
+**Using curl**
+
+```bash
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      {"role": "user", "content": "A beautiful sunset over the ocean with sailing boats"}
+    ],
+    "extra_body": {
+      "height": 1024,
+      "width": 1024,
+      "num_inference_steps": 50,
+      "guidance_scale": 1.5,
+      "seed": 42
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
+```
+
+Or use the script:
+
+```bash
+bash run_curl_text_to_image.sh "A futuristic city skyline at night"
+```
+
+### Image-to-Image (Image Editing)
+
+Edit images with text instructions:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "Convert this image to watercolor style" \
+    --image input.png \
+    --output watercolor.png
+```
+
+**Using curl**
+
+```bash
+IMG_B64=$(base64 < input.png | tr -d '\n')
+
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d @- <<EOF | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
+{
+  "messages": [{
+    "role": "user",
+    "content": [
+      {"type": "text", "text": "Convert this image to watercolor style"},
+      {"type": "image_url", "image_url": {"url": "data:image/png;base64,'$IMG_B64'"}}
+    ]
+  }],
+  "extra_body": {
+    "height": 1024,
+    "width": 1024,
+    "num_inference_steps": 50,
+    "guidance_scale": 1.5,
+    "seed": 42
+  }
+}
+EOF
+```
+
+Or use the script:
+
+```bash
+bash run_curl_image_edit.sh input.png "Convert to watercolor style"
+```
+
+## Generation Parameters (extra_body)
+
+| Parameter             | Type  | Default | Description                         |
+| --------------------- | ----- | ------- | ----------------------------------- |
+| `height`              | int   | 1024    | Image height in pixels              |
+| `width`               | int   | 1024    | Image width in pixels               |
+| `num_inference_steps` | int   | 50      | Number of diffusion denoising steps |
+| `guidance_scale`      | float | 1.5     | Classifier-free guidance scale      |
+| `seed`                | int   | 42      | Random seed for reproducibility     |
+| `negative_prompt`     | str   | None    | Negative prompt                     |
+
+## Response Format
+
+```json
+{
+  "id": "chatcmpl-xxx",
+  "created": 1234567890,
+  "model": "zai-org/GLM-Image",
+  "choices": [
+    {
+      "index": 0,
+      "message": {
+        "role": "assistant",
+        "content": [
+          {
+            "type": "image_url",
+            "image_url": {
+              "url": "data:image/png;base64,..."
+            }
+          }
+        ]
+      },
+      "finish_reason": "stop"
+    }
+  ],
+  "usage": {}
+}
+```
+
+## Extract Image
+
+```bash
+# From a saved JSON response
+cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
+```
+
+## Architecture
+
+GLM-Image uses a 2-stage multistage pipeline:
+
+```
+Stage 0 (AR Model)                Stage 1 (Diffusion)
+┌───────────────────┐            ┌─────────────────────┐
+│ vLLM-optimized    │  prior     │  GlmImagePipeline   │
+│ GlmImageFor       │──tokens──►│  ┌───────────────┐  │
+│ Conditional       │            │  │ DiT Denoiser  │  │
+│ Generation        │            │  └───────┬───────┘  │
+│ (9B AR model)     │            │          ▼          │
+└───────────────────┘            │  ┌───────────────┐  │
+        ▲                        │  │  VAE Decode   │──┼──► Image
+        │                        │  └───────────────┘  │
+   Text / Image                  └─────────────────────┘
+     Input
+```
+
+## VRAM Requirements
+
+| Stage             | VRAM                   |
+| :---------------- | :--------------------- |
+| Stage-0 (AR)      | **~18 GiB + KV Cache** |
+| Stage-1 (DiT+VAE) | **~20 GiB**            |
+| Total             | **~38 GiB + KV Cache** |
+
+## File Description
+
+| File                        | Description                           |
+| --------------------------- | ------------------------------------- |
+| `run_server.sh`             | Server startup script                 |
+| `run_curl_text_to_image.sh` | Text-to-image curl example            |
+| `run_curl_image_edit.sh`    | Image-to-image (editing) curl example |
+| `openai_chat_client.py`     | Python client (t2i + i2i)             |
+
+## FAQ
+
+- If you encounter OOM errors, adjust `gpu_memory_utilization` in the stage config:
+
+```yaml
+# In glm_image.yaml, reduce from default 0.6:
+gpu_memory_utilization: 0.5
+```
+
+- The first request may be slow due to model warmup. Subsequent requests will be faster.
+
+## Example materials
+
+??? abstract "openai_chat_client.py"
+    ``````py
+    --8<-- "examples/online_serving/glm_image/openai_chat_client.py"
+    ``````
+??? abstract "run_curl_image_edit.sh"
+    ``````sh
+    --8<-- "examples/online_serving/glm_image/run_curl_image_edit.sh"
+    ``````
+??? abstract "run_curl_text_to_image.sh"
+    ``````sh
+    --8<-- "examples/online_serving/glm_image/run_curl_text_to_image.sh"
+    ``````
+??? abstract "run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/glm_image/run_server.sh"
+    ``````

docs/user_guide/examples/online_serving/image_to_video.md

@@ -0,0 +1,97 @@
+# Image-To-Video
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/image_to_video>.
+
+
+This example demonstrates how to deploy the Wan2.2 image-to-video model for online video generation using vLLM-Omni.
+
+## Start Server
+
+### Basic Start
+
+```bash
+vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --port 8091
+```
+
+### Start with Parameters
+
+Or use the startup script:
+
+```bash
+bash run_server.sh
+```
+
+The script allows overriding:
+- `MODEL` (default: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`)
+- `PORT` (default: `8091`)
+- `BOUNDARY_RATIO` (default: `0.875`)
+- `FLOW_SHIFT` (default: `12.0`)
+- `CACHE_BACKEND` (default: `none`)
+- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
+
+## API Calls
+
+### Method 1: Using curl
+
+```bash
+# Basic image-to-video generation
+bash run_curl_image_to_video.sh
+
+# Or execute directly (OpenAI-style multipart)
+curl -X POST http://localhost:8091/v1/videos \
+  -H "Accept: application/json" \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "input_reference=@/path/to/qwen-bear.png" \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_frames=33" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=1.0" \
+  -F "guidance_scale_2=1.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=12.0" \
+  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_i2v_output.mp4
+```
+
+## Request Format
+
+### Required Fields
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "input_reference=@/path/to/qwen-bear.png"
+```
+
+### Generation with Parameters
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "input_reference=@/path/to/qwen-bear.png" \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_frames=33" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=1.0" \
+  -F "guidance_scale_2=1.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=12.0" \
+  -F "seed=42"
+```
+
+## Example materials
+
+??? abstract "run_curl_image_to_video.sh"
+    ``````sh
+    --8<-- "examples/online_serving/image_to_video/run_curl_image_to_video.sh"
+    ``````
+??? abstract "run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/image_to_video/run_server.sh"
+    ``````

docs/user_guide/examples/online_serving/qwen2_5_omni.md

@@ -30,7 +30,7 @@ cd examples/online_serving/qwen2_5_omni
 #### Send request via python
 
 ```bash
-python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
+python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities --port 8091 --host "localhost"
 ```
 
 The Python client supports the following command-line arguments:

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -36,7 +36,7 @@ cd examples/online_serving/qwen3_omni
 ####  Send request via python
 
 ```bash
-python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
+python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image --port 8091 --host "localhost"
 ```
 
 The Python client supports the following command-line arguments:

docs/user_guide/examples/online_serving/qwen3_tts.md

@@ -9,13 +9,13 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
 
 ## Supported Models
 
-| Model | Task Type | Description |
-|-------|-----------|-------------|
+| Model                                  | Task Type   | Description                                           |
+| -------------------------------------- | ----------- | ----------------------------------------------------- |
 | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
-| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
-| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
-| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
-| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description              |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-Base`        | Base        | Voice cloning from reference audio                    |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant                                |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-Base`        | Base        | Smaller/faster variant for voice cloning              |
 
 ## Run examples (Qwen3-TTS)
 
@@ -112,7 +112,7 @@ The Python client supports the following command-line arguments:
 - `--voice`: Speaker/voice name (default: `vivian`). Options: `vivian`, `ryan`, `aiden`, etc.
 - `--language`: Language. Options: `Auto`, `Chinese`, `English`, `Japanese`, `Korean`, `German`, `French`, `Russian`, `Portuguese`, `Spanish`, `Italian`
 - `--instructions`: Voice style/emotion instructions
-- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task)
+- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task). Local paths are automatically base64-encoded by the client before sending to the server.
 - `--ref-text`: Reference audio transcript for voice cloning (Base task)
 - `--response-format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`, `aac`, `opus`
 - `--output` (or `-o`): Output audio file path (default: `tts_output.wav`)
@@ -220,7 +220,7 @@ Lists available voices for the loaded model:
     "task_type": "CustomVoice",
     "language": "Auto",
     "instructions": "Optional style instructions",
-    "ref_audio": "URL or base64 for voice cloning",
+    "ref_audio": "HTTP URL, base64 data URL, or file:// URI for voice cloning",
     "ref_text": "Reference audio transcript",
     "x_vector_only_mode": false,
     "max_new_tokens": 2048
@@ -237,36 +237,58 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
 
 ### OpenAI Standard Parameters
 
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `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) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
-| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
+| Parameter         | Type   | Default        | Description                                                 |
+| ----------------- | ------ | -------------- | ----------------------------------------------------------- |
+| `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)                    |
+| `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`) |
 
 ### vLLM-Omni Extension Parameters
 
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `task_type` | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base |
-| `language` | string | "Auto" | Language (see supported languages below) |
-| `instructions` | string | "" | Voice style/emotion instructions |
-| `max_new_tokens` | int | 2048 | Maximum tokens to generate |
+| Parameter        | Type   | Default       | Description                                                                  |
+| ---------------- | ------ | ------------- | ---------------------------------------------------------------------------- |
+| `task_type`      | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base                                      |
+| `language`       | string | "Auto"        | Language (see supported languages below)                                     |
+| `instructions`   | string | ""            | Voice style/emotion instructions                                             |
+| `max_new_tokens` | int    | 2048          | Maximum tokens to generate                                                   |
+| `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
 
 ### Voice Clone Parameters (Base task)
 
-| Parameter | Type | Required | Description |
-|-----------|------|----------|-------------|
-| `ref_audio` | string | **Yes** | Reference audio (URL or base64 data URL) |
-| `ref_text` | string | No | Transcript of reference audio (for ICL mode) |
-| `x_vector_only_mode` | bool | No | Use speaker embedding only (no ICL) |
+| Parameter            | Type   | Required | Description                                  |
+| -------------------- | ------ | -------- | -------------------------------------------- |
+| `ref_audio`          | string | **Yes**  | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`) |
+| `ref_text`           | string | No       | Transcript of reference audio (for ICL mode) |
+| `x_vector_only_mode` | bool   | No       | Use speaker embedding only (no ICL)          |
+
+## Streaming
+
+Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded
+(one chunk per Code2Wav window, default 25 frames; configurable in the stage config):
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, how are you?",
+        "voice": "vivian",
+        "language": "English",
+        "stream": true,
+        "response_format": "pcm"
+    }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
+```
+
+**Constraints:**
+- `stream=true` requires `response_format="pcm"` (raw 16-bit signed PCM, 24 kHz mono).
+- `speed` adjustment is not supported when streaming.
+- Requires the server stage config to have `async_chunk: true` (default in `qwen3_tts.yaml`).
 
 ## Limitations
 
-- **No streaming**: Audio is generated completely before being returned. Streaming will be supported after the pipeline is disaggregated (see RFC #938).
 - **Single request**: Batch processing is not yet optimized for online serving.
 
 ## Troubleshooting

docs/user_guide/examples/online_serving/text_to_audio.md

@@ -0,0 +1,193 @@
+# Text-To-Audio
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/stable_audio>.
+
+This example demonstrates how to deploy Stable Audio models for online text-to-audio generation using vLLM-Omni.
+
+## Supported Models
+
+| Model | Description |
+|-------|-------------|
+| `stabilityai/stable-audio-open-1.0` | Open-source audio generation, up to ~47 seconds, 44.1 kHz stereo |
+
+## Start Server
+
+### Basic Start
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
+```
+
+## API Calls
+
+### Method 1: Using curl
+
+```bash
+# Run all curl examples
+bash curl_examples.sh
+
+# Or execute directly
+curl -X POST http://localhost:8000/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0
+    }' --output cat.wav
+```
+
+### Method 2: Using Python Client
+
+```bash
+cd examples/online_serving/stable_audio
+
+# Simple generation
+python stable_audio_client.py \
+    --text "The sound of a cat purring"
+
+# With custom duration
+python stable_audio_client.py \
+    --text "A dog barking" \
+    --audio_length 5.0
+
+# With all parameters
+python stable_audio_client.py \
+    --text "Thunder and rain" \
+    --audio_length 15.0 \
+    --negative_prompt "Low quality" \
+    --guidance_scale 7.0 \
+    --num_inference_steps 100 \
+    --seed 42 \
+    --output thunder.wav
+```
+
+The Python client supports the following command-line arguments:
+
+- `--api_url`: API endpoint URL (default: `http://localhost:8000/v1/audio/generate`)
+- `--text`: Text prompt for audio generation (default: `"The sound of a cat purring"`)
+- `--audio_length`: Audio length in seconds (default: `10.0`, max ~47s for `stable-audio-open-1.0`)
+- `--audio_start`: Audio start time in seconds (default: `0.0`)
+- `--negative_prompt`: Negative prompt for classifier-free guidance (default: `"Low quality"`)
+- `--guidance_scale`: Guidance scale for diffusion (default: `7.0`)
+- `--num_inference_steps`: Number of inference steps (default: `100`)
+- `--seed`: Random seed for reproducibility (default: `None`)
+- `--response_format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`
+- `--output`: Output file path (default: `stable_audio_output.wav`)
+
+### Method 3: Using Python httpx
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8000/v1/audio/generate",
+    json={
+        "input": "The sound of ocean waves crashing on a beach",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+    },
+    timeout=300.0,
+)
+
+with open("ocean.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Request Format
+
+### Simple Generation
+
+```json
+{
+    "input": "The sound of ocean waves"
+}
+```
+
+### Generation with Parameters
+
+```json
+{
+    "input": "A piano playing a gentle melody",
+    "audio_length": 10.0,
+    "negative_prompt": "Low quality, distorted, noisy",
+    "guidance_scale": 8.0,
+    "num_inference_steps": 150,
+    "seed": 42,
+    "response_format": "wav"
+}
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/generate
+Content-Type: application/json
+```
+
+### Generation Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text prompt describing the audio to generate |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
+| `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) |
+| `audio_start` | float | 0.0 | Audio start time in seconds |
+| `negative_prompt` | string | null | Text describing what to avoid in generation |
+| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
+| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
+| `seed` | int | null | Random seed for reproducible generation |
+
+### Response Format
+
+Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
+
+## Tuning Tips
+
+1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`.
+2. **Quality vs Speed**:
+   - 50 steps: Fast, decent quality (quick previews)
+   - 100 steps: Good balance (general purpose)
+   - 150+ steps: High quality, slower (final / critical audio)
+3. **Guidance Scale**:
+   - Lower (3 - 5): More creative / varied output
+   - Default (7): Good balance
+   - Higher (10+): Strict adherence to the prompt
+4. **Negative Prompts**: Use to avoid unwanted characteristics such as `"Low quality"`, `"distorted"`, `"noisy"`.
+5. **Seeds**: Set a fixed seed to get deterministic, reproducible results.
+
+## File Description
+
+| File | Description |
+|------|-------------|
+| `curl_examples.sh` | Curl examples covering common use cases |
+| `stable_audio_client.py` | Python client with full CLI argument support |
+
+## Troubleshooting
+
+1. **Audio generation model did not produce audio output**: Verify the server started successfully and the model loaded without errors.
+2. **Connection refused**: Make sure the server is running on the correct port.
+3. **Generation timeout**: Reduce `num_inference_steps` or `audio_length`, and check GPU memory with `nvidia-smi`.
+4. **Out of memory**: Lower `--gpu-memory-utilization` or reduce `audio_length`.
+5. **Audio quality issues**: Increase `num_inference_steps`, add a negative prompt, or raise `guidance_scale`.
+
+## Example materials
+
+??? abstract "stable_audio_client.py"
+    ``````py
+    --8<-- "examples/online_serving/stable_audio/stable_audio_client.py"
+    ``````
+??? abstract "curl_examples.sh"
+    ``````sh
+    --8<-- "examples/online_serving/stable_audio/curl_examples.sh"
+    ``````

docs/user_guide/examples/online_serving/text_to_image.md

@@ -23,6 +23,29 @@ Or use the startup script:
 bash run_server.sh
 ```
 
+### Start with Parallelism Acceleration
+
+Enable Tensor Parallelism and VAE Patch Parallelism for faster inference:
+
+```bash
+# With Tensor Parallelism (requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --tensor-parallel-size 2
+
+# With Tensor Parallelism and VAE Patch Parallelism (requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --tensor-parallel-size 2 --vae-patch-parallel-size 2 --vae-use-tiling
+
+# With Sequence Parallelism (Ulysses-SP, requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
+
+# With Ring-Attention (requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2
+
+# Combined: Ulysses + Ring (requires >= 4 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
+```
+
+For more details on parallelism acceleration, see the [Parallelism Acceleration Guide](../../diffusion/parallelism_acceleration.md).
+
 ## API Calls
 
 ### Method 1: Using curl

docs/user_guide/examples/online_serving/text_to_video.md

@@ -0,0 +1,130 @@
+# Text-To-Video
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/text_to_video>.
+
+
+This example demonstrates how to deploy the Wan2.2 text-to-video model for online video generation using vLLM-Omni.
+
+## Start Server
+
+### Basic Start
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091
+```
+
+### Start with Parameters
+
+Or use the startup script:
+
+```bash
+bash run_server.sh
+```
+
+The script allows overriding:
+- `MODEL` (default: `Wan-AI/Wan2.2-T2V-A14B-Diffusers`)
+- `PORT` (default: `8091`)
+- `BOUNDARY_RATIO` (default: `0.875`)
+- `FLOW_SHIFT` (default: `5.0`)
+- `CACHE_BACKEND` (default: `none`)
+- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
+
+## API Calls
+
+### Method 1: Using curl
+
+```bash
+# Basic text-to-video generation
+bash run_curl_text_to_video.sh
+
+# Or execute directly (OpenAI-style multipart)
+curl -s http://localhost:8091/v1/videos \
+  -H "Accept: application/json" \
+  -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_frames=33" \
+  -F "negative_prompt=色调艳丽 ，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=4.0" \
+  -F "guidance_scale_2=4.0" \
+  -F "boundary_ratio=0.875" \
+  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
+```
+
+## Request Format
+
+### Simple Text-to-Video Generation
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A cinematic view of a futuristic city at sunset"
+```
+
+### Generation with Parameters
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A cinematic view of a futuristic city at sunset" \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_frames=33" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=4.0" \
+  -F "guidance_scale_2=4.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=5.0" \
+  -F "seed=42"
+```
+
+## Generation Parameters
+
+| Parameter             | Type   | Default | Description                                      |
+| --------------------- | ------ | ------- | ------------------------------------------------ |
+| `prompt`              | str    | -       | Text description of the desired video            |
+| `negative_prompt`     | str    | None    | Negative prompt                                  |
+| `n`                   | int    | 1       | Number of videos to generate                     |
+| `width`               | int    | None    | Video width in pixels                            |
+| `height`              | int    | None    | Video height in pixels                           |
+| `num_frames`          | int    | None    | Number of frames to generate                     |
+| `fps`                 | int    | None    | Frames per second for output video               |
+| `num_inference_steps` | int    | None    | Number of denoising steps                        |
+| `guidance_scale`      | float  | None    | CFG guidance scale (low-noise stage)             |
+| `guidance_scale_2`    | float  | None    | CFG guidance scale (high-noise stage, Wan2.2)     |
+| `boundary_ratio`      | float  | None    | Boundary split ratio for low/high DiT (Wan2.2)   |
+| `flow_shift`          | float  | None    | Scheduler flow shift (Wan2.2)                    |
+| `seed`                | int    | None    | Random seed (reproducible)                       |
+| `lora`                | object | None    | LoRA configuration                               |
+| `extra_body`          | object | None    | Model-specific extra parameters                  |
+
+## Response Format
+
+```json
+{
+  "created": 1234567890,
+  "data": [
+    { "b64_json": "<base64-mp4>" }
+  ]
+}
+```
+
+## Extract Video
+
+```bash
+# Extract base64 from response and decode to video
+cat response.json | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
+```
+
+## Example materials
+
+??? abstract "run_curl_text_to_video.sh"
+    ``````sh
+    --8<-- "examples/online_serving/text_to_video/run_curl_text_to_video.sh"
+    ``````
+??? abstract "run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/text_to_video/run_server.sh"
+    ``````

examples/offline_inference/bagel/end2end.py

@@ -1,8 +1,7 @@
 import argparse
 import os
-from typing import cast
 
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
+from vllm_omni.inputs.data import OmniPromptType
 
 
 def parse_args():
@@ -49,8 +48,16 @@ def parse_args():
     parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (default: 4.0)")
     parser.add_argument("--cfg-img-scale", type=float, default=1.5, help="Image CFG scale (default: 1.5)")
     parser.add_argument(
-        "--negative-prompt", type=str, default=None, help="Negative prompt (not yet supported, reserved for future)"
+        "--negative-prompt", type=str, default=None, help="Negative prompt for CFG (default: empty prompt)"
     )
+    parser.add_argument(
+        "--cfg-parallel-size",
+        type=int,
+        default=1,
+        choices=[1, 2, 3],
+        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.")
 
     args = parser.parse_args()
     return args
@@ -81,101 +88,80 @@ def main():
 
     from PIL import Image
 
-    if args.modality == "img2img":
-        from PIL import Image
+    from vllm_omni.entrypoints.omni import Omni
 
-        from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
+    omni_kwargs = {}
+    if args.stage_configs_path:
+        omni_kwargs["stage_configs_path"] = args.stage_configs_path
 
-        print("[Info] Running in img2img mode (Stage 1 only)")
-        client = OmniDiffusion(model=model_name)
+    omni_kwargs.update(
+        {
+            "log_stats": args.log_stats,
+            "init_sleep_seconds": args.init_sleep_seconds,
+            "batch_timeout": args.batch_timeout,
+            "init_timeout": args.init_timeout,
+            "shm_threshold_bytes": args.shm_threshold_bytes,
+            "worker_backend": args.worker_backend,
+            "ray_address": args.ray_address,
+        }
+    )
 
-        if args.image_path:
-            if os.path.exists(args.image_path):
-                loaded_image = Image.open(args.image_path).convert("RGB")
-                prompts = [
-                    {
-                        "prompt": cast(str, p),
-                        "multi_modal_data": {"image": loaded_image},
-                    }
-                    for p in prompts
-                ]
-            else:
-                print(f"[Warning] Image path {args.image_path} does not exist.")
+    omni = Omni(model=model_name, **omni_kwargs)
 
-        result = client.generate(
-            prompts,
-            OmniDiffusionSamplingParams(
-                seed=52,
-                need_kv_receive=False,
-                num_inference_steps=args.steps,
-                extra_args={
-                    "cfg_text_scale": args.cfg_text_scale,
-                    "cfg_img_scale": args.cfg_img_scale,
-                },
-            ),
-        )
-
-        # Ensure result is a list for iteration
-        if not isinstance(result, list):
-            omni_outputs = [result]
-        else:
-            omni_outputs = result
-
-    else:
-        from vllm_omni.entrypoints.omni import Omni
-
-        omni_kwargs = {}
-        if args.stage_configs_path:
-            omni_kwargs["stage_configs_path"] = args.stage_configs_path
-
-        omni_kwargs.update(
-            {
-                "log_stats": args.log_stats,
-                "init_sleep_seconds": args.init_sleep_seconds,
-                "batch_timeout": args.batch_timeout,
-                "init_timeout": args.init_timeout,
-                "shm_threshold_bytes": args.shm_threshold_bytes,
-                "worker_backend": args.worker_backend,
-                "ray_address": args.ray_address,
+    formatted_prompts = []
+    for p in prompts:
+        if args.modality == "img2img":
+            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")
+            final_prompt_text = f"<|fim_middle|><|im_start|>{p}<|im_end|>"
+            prompt_dict = {
+                "prompt": final_prompt_text,
+                "multi_modal_data": {"img2img": loaded_image},
+                "modalities": ["img2img"],
             }
-        )
-
-        omni = Omni(model=model_name, **omni_kwargs)
-
-        formatted_prompts = []
-        for p in args.prompts:
-            if args.modality == "img2text":
-                if args.image_path:
-                    loaded_image = Image.open(args.image_path).convert("RGB")
-                    final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
-                    prompt_dict = {
-                        "prompt": final_prompt_text,
-                        "multi_modal_data": {"image": loaded_image},
-                        "modalities": ["text"],
-                    }
-                    formatted_prompts.append(prompt_dict)
-            elif args.modality == "text2text":
-                final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
-                prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
-                formatted_prompts.append(prompt_dict)
-            else:
-                # text2img
-                final_prompt_text = f"<|im_start|>{p}<|im_end|>"
-                prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
-                formatted_prompts.append(prompt_dict)
-
-        params_list = omni.default_sampling_params_list
-        if args.modality == "text2img":
-            params_list[0].max_tokens = 1  # type: ignore # The first stage is a SamplingParam (vllm)
-            if len(params_list) > 1:
-                diffusion_params = params_list[1]
-                diffusion_params.num_inference_steps = args.steps  # type: ignore
-                diffusion_params.extra_args = {  # type: ignore
-                    "cfg_text_scale": args.cfg_text_scale,
-                    "cfg_img_scale": args.cfg_img_scale,
+            if args.negative_prompt is not None:
+                prompt_dict["negative_prompt"] = args.negative_prompt
+            formatted_prompts.append(prompt_dict)
+        elif args.modality == "img2text":
+            if args.image_path:
+                loaded_image = Image.open(args.image_path).convert("RGB")
+                final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                prompt_dict = {
+                    "prompt": final_prompt_text,
+                    "multi_modal_data": {"image": loaded_image},
+                    "modalities": ["text"],
                 }
+                formatted_prompts.append(prompt_dict)
+        elif args.modality == "text2text":
+            final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
+            prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
+            formatted_prompts.append(prompt_dict)
+        else:
+            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
+            formatted_prompts.append(prompt_dict)
 
-        omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
+    params_list = omni.default_sampling_params_list
+    if args.modality in ("text2img", "img2img"):
+        params_list[0].max_tokens = 1  # type: ignore
+        if len(params_list) > 1:
+            diffusion_params = params_list[1]
+            diffusion_params.num_inference_steps = args.steps  # type: ignore
+            diffusion_params.cfg_parallel_size = args.cfg_parallel_size  # type: ignore
+            if args.seed is not None:
+                diffusion_params.seed = args.seed  # type: ignore
+            extra = {
+                "cfg_text_scale": args.cfg_text_scale,
+                "cfg_img_scale": args.cfg_img_scale,
+            }
+            if args.negative_prompt is not None:
+                extra["negative_prompt"] = args.negative_prompt
+            diffusion_params.extra_args = extra  # type: ignore
+
+    omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
 
     for i, req_output in enumerate(omni_outputs):
         images = getattr(req_output, "images", None)

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

@@ -29,8 +29,8 @@ Accelerated with cache_dit:
     python image_edit.py \
         --image input.png \
         --prompt "Edit description" \
-        --cache_backend cache_dit \
-        --cache_dit_enable_taylorseer
+        --cache-backend cache_dit \
+        --cache-dit-enable-taylorseer
 
 Layered RGBA output:
     python image_edit.py \
@@ -66,38 +66,38 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--model", default="Qwen/Qwen-Image-Edit", help="Model name or local path.")
     parser.add_argument("--image", type=str, nargs="+", required=True, help="Input image file(s).")
     parser.add_argument("--prompt", type=str, required=True, help="Edit description prompt.")
-    parser.add_argument("--negative_prompt", type=str, default=None)
+    parser.add_argument("--negative-prompt", type=str, default=None)
     parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument("--cfg_scale", type=float, default=4.0)
-    parser.add_argument("--guidance_scale", type=float, default=1.0)
+    parser.add_argument("--cfg-scale", type=float, default=4.0)
+    parser.add_argument("--guidance-scale", type=float, default=1.0)
     parser.add_argument("--output", type=str, default="output.png")
-    parser.add_argument("--num_outputs_per_prompt", type=int, default=1)
-    parser.add_argument("--num_inference_steps", type=int, default=50)
-    parser.add_argument("--cache_backend", type=str, default=None, choices=["cache_dit", "tea_cache"])
-    parser.add_argument("--ulysses_degree", type=int, default=1)
-    parser.add_argument("--ring_degree", type=int, default=1)
-    parser.add_argument("--tensor_parallel_size", type=int, default=1)
+    parser.add_argument("--num-outputs-per-prompt", type=int, default=1)
+    parser.add_argument("--num-inference-steps", type=int, default=50)
+    parser.add_argument("--cache-backend", type=str, default=None, choices=["cache_dit", "tea_cache"])
+    parser.add_argument("--ulysses-degree", type=int, default=1)
+    parser.add_argument("--ring-degree", type=int, default=1)
+    parser.add_argument("--tensor-parallel-size", type=int, default=1)
     parser.add_argument("--layers", type=int, default=4)
     parser.add_argument("--resolution", type=int, default=640)
     parser.add_argument("--color-format", type=str, default="RGB")
 
     # Acceleration + Optimization Options
-    parser.add_argument("--cache_dit_fn_compute_blocks", type=int, default=1)
-    parser.add_argument("--cache_dit_bn_compute_blocks", type=int, default=0)
-    parser.add_argument("--cache_dit_max_warmup_steps", type=int, default=4)
-    parser.add_argument("--cache_dit_residual_diff_threshold", type=float, default=0.24)
-    parser.add_argument("--cache_dit_max_continuous_cached_steps", type=int, default=3)
-    parser.add_argument("--cache_dit_enable_taylorseer", action="store_true", default=False)
-    parser.add_argument("--cache_dit_taylorseer_order", type=int, default=1)
+    parser.add_argument("--cache-dit-fn-compute-blocks", type=int, default=1)
+    parser.add_argument("--cache-dit-bn-compute-blocks", type=int, default=0)
+    parser.add_argument("--cache-dit-max-warmup-steps", type=int, default=4)
+    parser.add_argument("--cache-dit-residual-diff-threshold", type=float, default=0.24)
+    parser.add_argument("--cache-dit-max-continuous-cached-steps", type=int, default=3)
+    parser.add_argument("--cache-dit-enable-taylorseer", action="store_true", default=False)
+    parser.add_argument("--cache-dit-taylorseer-order", type=int, default=1)
     parser.add_argument(
-        "--cache_dit_scm_steps_mask_policy", type=str, default=None, choices=[None, "slow", "medium", "fast", "ultra"]
+        "--cache-dit-scm-steps-mask-policy", type=str, default=None, choices=[None, "slow", "medium", "fast", "ultra"]
     )
-    parser.add_argument("--cache_dit_scm_steps_policy", type=str, default="dynamic", choices=["dynamic", "static"])
-    parser.add_argument("--tea_cache_rel_l1_thresh", type=float, default=0.2)
-    parser.add_argument("--cfg_parallel_size", type=int, default=1, choices=[1, 2])
-    parser.add_argument("--enforce_eager", action="store_true")
-    parser.add_argument("--vae_use_slicing", action="store_true")
-    parser.add_argument("--vae_use_tiling", action="store_true")
+    parser.add_argument("--cache-dit-scm-steps-policy", type=str, default="dynamic", choices=["dynamic", "static"])
+    parser.add_argument("--tea-cache-rel-l1-thresh", type=float, default=0.2)
+    parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2])
+    parser.add_argument("--enforce-eager", action="store_true")
+    parser.add_argument("--vae-use-slicing", action="store_true")
+    parser.add_argument("--vae-use-tiling", action="store_true")
     parser.add_argument("--enable-cpu-offload", action="store_true")
 
     return parser.parse_args()

examples/offline_inference/custom_pipeline/image_to_image/run.sh

@@ -3,4 +3,4 @@ python image_edit.py \
         --image cherry_blossom.jpg \
         --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
         --output output_image_edit.png \
-        --num_inference_steps 10
+        --num-inference-steps 10

examples/offline_inference/glm_image/README.md

@@ -0,0 +1,136 @@
+# GLM-Image Multistage End-to-End Inference
+
+This example demonstrates how to run GLM-Image with the vLLM-Omni multistage architecture.
+
+## Architecture
+
+GLM-Image uses a 2-stage pipeline:
+
+```
+┌─────────────────────────────────────────────────────────────┐
+│                     GLM-Image Pipeline                       │
+├─────────────────────────────────────────────────────────────┤
+│                                                             │
+│  Stage 0 (AR Model)              Stage 1 (Diffusion)        │
+│  ┌─────────────────┐            ┌─────────────────────┐     │
+│  │ vLLM-optimized  │            │  GlmImagePipeline   │     │
+│  │ GlmImageFor     │  prior     │  ┌───────────────┐  │     │
+│  │ Conditional     │──tokens───►│  │ DiT Denoiser  │  │     │
+│  │ Generation      │            │  └───────────────┘  │     │
+│  │ (9B AR model)   │            │         │          │     │
+│  └─────────────────┘            │         ▼          │     │
+│         ▲                       │  ┌───────────────┐  │     │
+│         │                       │  │  VAE Decode   │──┼──► Image
+│    Text/Image                   │  └───────────────┘  │     │
+│      Input                      └─────────────────────┘     │
+│                                                             │
+└─────────────────────────────────────────────────────────────┘
+```
+
+## Features
+
+- **vLLM-optimized AR**: Uses PagedAttention and tensor parallelism for faster prior token generation
+- **Flexible deployment**: AR and Diffusion stages can run on different GPUs
+- **Text-to-Image**: Generate images from text descriptions
+- **Image-to-Image**: Edit existing images with text prompts
+
+## Usage
+
+### Text-to-Image
+
+```bash
+python end2end.py \
+    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "A beautiful sunset over the ocean with sailing boats" \
+    --height 1024 \
+    --width 1024 \
+    --output output_t2i.png
+```
+
+### Image-to-Image (Image Editing)
+
+```bash
+python end2end.py \
+    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "Transform this scene into a winter wonderland" \
+    --image input.png \
+    --output output_i2i.png
+```
+
+### With Custom Parameters
+
+```bash
+python end2end.py \
+    --model-path /path/to/glm-image \
+    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "A photorealistic cat sitting on a window sill" \
+    --height 1024 \
+    --width 1024 \
+    --num-inference-steps 50 \
+    --guidance-scale 1.5 \
+    --seed 42 \
+    --output output.png
+```
+
+## Shell Scripts
+
+### Run Text-to-Image
+
+```bash
+./run_t2i.sh
+```
+
+### Run Image-to-Image
+
+```bash
+./run_i2i.sh --image /path/to/input.png
+```
+
+## Stage Configuration
+
+The stage config (`glm_image.yaml`) defines:
+
+- **Stage 0 (AR)**: Uses `GPUARWorker` with vLLM engine
+
+  - Model: `GlmImageForConditionalGeneration`
+  - Output: `token_ids` (prior tokens)
+
+- **Stage 1 (Diffusion)**: Uses diffusion engine
+  - Model: `GlmImagePipeline`
+  - Output: Generated image
+
+See `vllm_omni/model_executor/stage_configs/glm_image.yaml` for full configuration.
+
+## Comparison with Single-Stage
+
+| Aspect      | Single-Stage (transformers) | Multistage (vLLM)   |
+| ----------- | --------------------------- | ------------------- |
+| AR Model    | transformers native         | vLLM PagedAttention |
+| Memory      | Higher (no KV cache opt)    | Lower (optimized)   |
+| Throughput  | Lower                       | Higher              |
+| Flexibility | Single GPU                  | Multi-GPU support   |
+
+## Troubleshooting
+
+### OOM Error
+
+Try reducing memory usage:
+
+```bash
+# In glm_image.yaml, adjust:
+gpu_memory_utilization: 0.5  # Reduce from 0.6
+```
+
+### Slow Initialization
+
+The first run loads model weights. Subsequent runs are faster:
+
+```bash
+--stage-init-timeout 900  # Increase timeout for slow storage
+```
+
+## Requirements
+
+- vLLM-Omni with GLM-Image support
+- CUDA-capable GPU (recommended: H100/A100 with 80GB)
+- GLM-Image model weights

examples/offline_inference/glm_image/end2end.py

@@ -0,0 +1,473 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+End-to-end offline inference example for GLM-Image with multistage architecture.
+
+This script tests the multistage pipeline where:
+- Stage 0 (AR): vLLM-optimized GlmImageForConditionalGeneration generates prior_token_ids
+- Stage 1 (Diffusion): GlmImagePipeline performs DiT denoising + VAE decode
+
+Usage (text-to-image):
+    python end2end.py \
+        --model-path /path/to/glm-image \
+        --config-path /path/to/glm_image.yaml \
+        --prompt "A beautiful sunset over the ocean" \
+        --output output_t2i.png
+
+Usage (image-to-image / image edit):
+    python end2end.py \
+        --model-path /path/to/glm-image \
+        --config-path /path/to/glm_image.yaml \
+        --prompt "Make it look like winter" \
+        --image input.png \
+        --output output_i2i.png
+
+Usage (with custom parameters):
+    python end2end.py \
+        --model-path /path/to/glm-image \
+        --config-path /path/to/glm_image.yaml \
+        --prompt "A cat sitting on a window sill" \
+        --height 1024 \
+        --width 1024 \
+        --num-inference-steps 50 \
+        --guidance-scale 1.5 \
+        --seed 42
+
+For more options, run:
+    python end2end.py --help
+"""
+
+import argparse
+import os
+import time
+from pathlib import Path
+
+from PIL import Image
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+# Default stage config path (relative to vllm_omni package)
+DEFAULT_CONFIG_PATH = "vllm_omni/model_executor/stage_configs/glm_image.yaml"
+
+SEED = 42
+
+# GLM-Image special tokens
+GLM_IMAGE_EOS_TOKEN_ID = 16385  # eos_token_id from generation_config.json
+GLM_IMAGE_VISION_VOCAB_SIZE = 16512  # top_k should be vision_vocab_size
+
+
+def compute_max_tokens(height: int, width: int, factor: int = 32) -> int:
+    """
+    Compute max_new_tokens for GLM-Image AR generation.
+
+    GLM-Image generates tokens in this order for text-to-image:
+    1. Small preview image (half resolution in each dimension)
+    2. Large target image (full resolution)
+    3. EOS token
+
+    Args:
+        height: Target image height in pixels
+        width: Target image width in pixels
+        factor: Downsampling factor (32 for GLM-Image AR output)
+
+    Returns:
+        Total number of tokens to generate (small + large + EOS)
+    """
+    # Large image tokens (target resolution)
+    token_h = height // factor
+    token_w = width // factor
+    large_tokens = token_h * token_w
+
+    # Small preview tokens (half resolution in each dimension)
+    small_h = token_h // 2
+    small_w = token_w // 2
+    small_tokens = small_h * small_w
+
+    # Total: small + large + EOS
+    return small_tokens + large_tokens + 1
+
+
+def load_image(image_path: str) -> Image.Image:
+    """Load an image from file path."""
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+    return Image.open(image_path).convert("RGB")
+
+
+def save_image(image: Image.Image, output_path: str) -> None:
+    """Save an image to file path."""
+    output_dir = os.path.dirname(output_path)
+    if output_dir:
+        os.makedirs(output_dir, exist_ok=True)
+    image.save(output_path)
+    print(f"Image saved to: {output_path}")
+
+
+def build_prompt_for_t2i(
+    prompt: str,
+    height: int = 1024,
+    width: int = 1024,
+) -> dict:
+    """
+    Build prompt dict for text-to-image generation.
+
+    Args:
+        prompt: Text description for image generation
+        height: Target image height
+        width: Target image width
+
+    Returns:
+        Dict containing prompt and generation parameters
+    """
+    return {
+        "prompt": prompt,
+        "height": height,
+        "width": width,
+        # Pass target dimensions to AR processor for proper grid token generation
+        "mm_processor_kwargs": {
+            "target_h": height,
+            "target_w": width,
+        },
+    }
+
+
+def build_prompt_for_i2i(
+    prompt: str,
+    image: Image.Image,
+    height: int | None = None,
+    width: int | None = None,
+) -> dict:
+    """
+    Build prompt dict for image-to-image generation.
+
+    Args:
+        prompt: Text description for image editing
+        image: Source image for editing
+        height: Target image height (default: use source image size)
+        width: Target image width (default: use source image size)
+
+    Returns:
+        Dict containing prompt, image, and generation parameters
+    """
+    # Use source image dimensions if not specified
+    if height is None:
+        height = image.height
+    if width is None:
+        width = image.width
+
+    return {
+        "prompt": prompt,
+        "multi_modal_data": {
+            "image": image,
+        },
+        "height": height,
+        "width": width,
+        # Pass target dimensions to AR processor for proper grid token generation
+        "mm_processor_kwargs": {
+            "target_h": height,
+            "target_w": width,
+        },
+    }
+
+
+def main(args: argparse.Namespace) -> None:
+    """Main entry point for GLM-Image end-to-end inference."""
+    print("=" * 60)
+    print("GLM-Image Multistage End-to-End Inference")
+    print("=" * 60)
+
+    # Validate arguments
+    if not args.prompt:
+        raise ValueError("--prompt is required")
+
+    # Determine config path
+    config_path = args.config_path
+    if config_path is None:
+        # Try to find default config
+        if os.path.exists(DEFAULT_CONFIG_PATH):
+            config_path = DEFAULT_CONFIG_PATH
+        else:
+            # Try relative to script location
+            script_dir = Path(__file__).parent.parent.parent.parent
+            config_path = script_dir / "vllm_omni/model_executor/stage_configs/glm_image.yaml"
+            if not config_path.exists():
+                raise FileNotFoundError(
+                    f"Stage config not found. Please specify --config-path. Tried: {DEFAULT_CONFIG_PATH}"
+                )
+            config_path = str(config_path)
+
+    print(f"Model path: {args.model_path}")
+    print(f"Config path: {config_path}")
+    print(f"Prompt: {args.prompt}")
+
+    # Load source image for image-to-image mode
+    source_image = None
+    if args.image:
+        print(f"Source image: {args.image}")
+        source_image = load_image(args.image)
+        print(f"  Image size: {source_image.width}x{source_image.height}")
+
+    # Build prompt based on mode
+    if source_image is not None:
+        # Image-to-image mode
+        prompt_dict = build_prompt_for_i2i(
+            prompt=args.prompt,
+            image=source_image,
+            height=args.height,
+            width=args.width,
+        )
+        mode = "image-to-image"
+    else:
+        # Text-to-image mode
+        prompt_dict = build_prompt_for_t2i(
+            prompt=args.prompt,
+            height=args.height or 1024,
+            width=args.width or 1024,
+        )
+        mode = "text-to-image"
+
+    print(f"Mode: {mode}")
+    print(f"Target size: {prompt_dict.get('height', 1024)}x{prompt_dict.get('width', 1024)}")
+
+    # Add generation parameters to prompt
+    prompt_dict["seed"] = args.seed
+    prompt_dict["num_inference_steps"] = args.num_inference_steps
+    prompt_dict["guidance_scale"] = args.guidance_scale
+
+    if args.negative_prompt:
+        prompt_dict["negative_prompt"] = args.negative_prompt
+
+    # Initialize Omni with multistage config
+    print("\nInitializing Omni with multistage pipeline...")
+    start_time = time.time()
+
+    omni = Omni(
+        model=args.model_path,
+        stage_configs_path=config_path,
+        log_stats=args.enable_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    init_time = time.time() - start_time
+    print(f"Initialization completed in {init_time:.2f}s")
+
+    # Prepare prompts (support batch generation)
+    prompts = [prompt_dict for _ in range(args.num_prompts)]
+
+    # No explicit sampling_params for diffusion - parameters are in prompt_dict
+    # For multistage, the AR stage may need sampling params
+    from vllm import SamplingParams
+
+    # Compute max_tokens dynamically based on target image size
+    target_height = prompt_dict.get("height", 1024)
+    target_width = prompt_dict.get("width", 1024)
+    calculated_max_tokens = compute_max_tokens(target_height, target_width)
+
+    # Use calculated value unless user explicitly specified a different value
+    # Default args.max_tokens is 16384 (very large), so prefer calculated value
+    effective_max_tokens = calculated_max_tokens if args.max_tokens == 16384 else args.max_tokens
+
+    if args.verbose:
+        print(f"AR max_tokens: {effective_max_tokens} (calculated: {calculated_max_tokens}, arg: {args.max_tokens})")
+
+    # IMPORTANT: GLM-Image AR model requires these exact sampling parameters
+    # from generation_config.json for proper image token generation.
+    # - temperature=0.9, top_p=0.75, top_k=16512 (vision_vocab_size)
+    # - stop_token_ids=[16385] (eos_token_id) is CRITICAL to stop generation
+    ar_sampling_params = SamplingParams(
+        temperature=0.9,  # From generation_config.json
+        top_p=0.75,  # From generation_config.json
+        top_k=GLM_IMAGE_VISION_VOCAB_SIZE,  # 16512, vision vocabulary size
+        max_tokens=effective_max_tokens,
+        stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID],  # 16385, CRITICAL for stopping
+        seed=args.seed,
+        detokenize=False,
+    )
+
+    # For diffusion stage, sampling_params contains diffusion-specific parameters
+    # These are passed as kwargs to the diffusion engine
+    diffusion_sampling_params = OmniDiffusionSamplingParams(
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        height=prompt_dict.get("height", 1024),
+        width=prompt_dict.get("width", 1024),
+        seed=args.seed,
+    )
+
+    # For multistage, we need sampling_params for each stage
+    # Stage 0 (AR): SamplingParams for vLLM
+    # Stage 1 (Diffusion): dict with diffusion kwargs
+    sampling_params_list = [ar_sampling_params, diffusion_sampling_params]
+
+    # Run generation
+    print(f"\nGenerating {args.num_prompts} image(s)...")
+    gen_start_time = time.time()
+
+    output_dir = os.path.dirname(args.output) if args.output else "outputs"
+    if output_dir:
+        os.makedirs(output_dir, exist_ok=True)
+
+    output_count = 0
+    for stage_outputs in omni.generate(prompts, sampling_params_list, py_generator=True):
+        if stage_outputs.final_output_type == "image":
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
+
+                # Get generated images
+                images = output.images if hasattr(output, "images") else []
+                if not images and hasattr(output, "multimodal_output"):
+                    images = output.multimodal_output.get("images", [])
+
+                # Save each generated image
+                for idx, img in enumerate(images):
+                    if args.num_prompts == 1 and len(images) == 1:
+                        output_path = args.output
+                    else:
+                        base, ext = os.path.splitext(args.output)
+                        output_path = f"{base}_{request_id}_{idx}{ext}"
+
+                    if isinstance(img, Image.Image):
+                        save_image(img, output_path)
+                    else:
+                        print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
+
+                    output_count += 1
+
+        elif stage_outputs.final_output_type == "text":
+            # AR stage output (intermediate, for debugging)
+            if args.verbose:
+                for output in stage_outputs.request_output:
+                    print(f"AR output for request {output.request_id}:")
+                    print(f"  Token count: {len(output.outputs[0].token_ids)}")
+
+    gen_time = time.time() - gen_start_time
+    print(f"\nGeneration completed in {gen_time:.2f}s")
+    print(f"Generated {output_count} image(s)")
+
+    # Cleanup
+    omni.close()
+    print("\nDone!")
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="GLM-Image Multistage End-to-End Inference",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+
+    # Required arguments
+    parser.add_argument(
+        "--model-path",
+        type=str,
+        default="zai-org/GLM-Image",
+        help="Path to GLM-Image model directory or HuggingFace model ID",
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        required=True,
+        help="Text prompt for image generation",
+    )
+
+    # Optional arguments
+    parser.add_argument(
+        "--config-path",
+        type=str,
+        default=None,
+        help="Path to stage config YAML file (default: auto-detect)",
+    )
+    parser.add_argument(
+        "--image",
+        type=str,
+        default=None,
+        help="Path to source image for image-to-image mode",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="output_glm_image.png",
+        help="Output image path (default: output_glm_image.png)",
+    )
+    parser.add_argument(
+        "--negative-prompt",
+        type=str,
+        default=None,
+        help="Negative prompt for classifier-free guidance",
+    )
+
+    # Generation parameters
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=None,
+        help="Output image height (default: 1024 for t2i, source size for i2i)",
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=None,
+        help="Output image width (default: 1024 for t2i, source size for i2i)",
+    )
+    parser.add_argument(
+        "--num-inference-steps",
+        type=int,
+        default=50,
+        help="Number of diffusion denoising steps (default: 50)",
+    )
+    parser.add_argument(
+        "--guidance-scale",
+        type=float,
+        default=1.5,
+        help="Classifier-free guidance scale (default: 1.5)",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=SEED,
+        help=f"Random seed for reproducibility (default: {SEED})",
+    )
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=16384,
+        help="Maximum tokens for AR generation (default: 16384)",
+    )
+
+    # Batch processing
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Number of images to generate (default: 1)",
+    )
+
+    # Runtime options
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=False,
+        help="Enable statistics logging",
+    )
+    parser.add_argument(
+        "--stage-init-timeout",
+        type=int,
+        default=600,
+        help="Timeout for stage initialization in seconds (default: 600)",
+    )
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="store_true",
+        default=False,
+        help="Enable verbose output",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

examples/offline_inference/glm_image/run_i2i.sh

@@ -0,0 +1,93 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# Run GLM-Image image-to-image (editing) with multistage pipeline
+
+set -e
+
+# Default values
+MODEL_PATH="${MODEL_PATH:-/path/to/glm-image}"
+CONFIG_PATH="${CONFIG_PATH:-vllm_omni/model_executor/stage_configs/glm_image.yaml}"
+PROMPT="${PROMPT:-Transform this image into an oil painting style}"
+INPUT_IMAGE=""
+OUTPUT="${OUTPUT:-output_i2i.png}"
+NUM_STEPS="${NUM_STEPS:-50}"
+GUIDANCE="${GUIDANCE:-1.5}"
+SEED="${SEED:-42}"
+
+# Parse command line arguments
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --model-path)
+            MODEL_PATH="$2"
+            shift 2
+            ;;
+        --config-path)
+            CONFIG_PATH="$2"
+            shift 2
+            ;;
+        --prompt)
+            PROMPT="$2"
+            shift 2
+            ;;
+        --image)
+            INPUT_IMAGE="$2"
+            shift 2
+            ;;
+        --output)
+            OUTPUT="$2"
+            shift 2
+            ;;
+        --num-steps)
+            NUM_STEPS="$2"
+            shift 2
+            ;;
+        --guidance)
+            GUIDANCE="$2"
+            shift 2
+            ;;
+        --seed)
+            SEED="$2"
+            shift 2
+            ;;
+        *)
+            echo "Unknown option: $1"
+            exit 1
+            ;;
+    esac
+done
+
+# Check if input image is provided
+if [ -z "${INPUT_IMAGE}" ]; then
+    echo "Error: --image is required for image-to-image mode"
+    echo "Usage: ./run_i2i.sh --image /path/to/input.png [--prompt \"edit instruction\"]"
+    exit 1
+fi
+
+if [ ! -f "${INPUT_IMAGE}" ]; then
+    echo "Error: Input image not found: ${INPUT_IMAGE}"
+    exit 1
+fi
+
+echo "=============================================="
+echo "GLM-Image Image-to-Image Generation"
+echo "=============================================="
+echo "Model: ${MODEL_PATH}"
+echo "Config: ${CONFIG_PATH}"
+echo "Input: ${INPUT_IMAGE}"
+echo "Prompt: ${PROMPT}"
+echo "Output: ${OUTPUT}"
+echo "Steps: ${NUM_STEPS}"
+echo "Guidance: ${GUIDANCE}"
+echo "Seed: ${SEED}"
+echo "=============================================="
+
+python end2end.py \
+    --model-path "${MODEL_PATH}" \
+    --config-path "${CONFIG_PATH}" \
+    --prompt "${PROMPT}" \
+    --image "${INPUT_IMAGE}" \
+    --output "${OUTPUT}" \
+    --num-inference-steps "${NUM_STEPS}" \
+    --guidance-scale "${GUIDANCE}" \
+    --seed "${SEED}" \
+    --verbose

examples/offline_inference/glm_image/run_t2i.sh

@@ -0,0 +1,87 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# Run GLM-Image text-to-image generation with multistage pipeline
+
+set -e
+
+# Default values
+MODEL_PATH="${MODEL_PATH:-/path/to/glm-image}"
+CONFIG_PATH="${CONFIG_PATH:-vllm_omni/model_executor/stage_configs/glm_image.yaml}"
+PROMPT="${PROMPT:-A beautiful mountain landscape with snow-capped peaks and a clear blue lake}"
+OUTPUT="${OUTPUT:-output_t2i.png}"
+HEIGHT="${HEIGHT:-1024}"
+WIDTH="${WIDTH:-1024}"
+NUM_STEPS="${NUM_STEPS:-50}"
+GUIDANCE="${GUIDANCE:-1.5}"
+SEED="${SEED:-42}"
+
+# Parse command line arguments
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --model-path)
+            MODEL_PATH="$2"
+            shift 2
+            ;;
+        --config-path)
+            CONFIG_PATH="$2"
+            shift 2
+            ;;
+        --prompt)
+            PROMPT="$2"
+            shift 2
+            ;;
+        --output)
+            OUTPUT="$2"
+            shift 2
+            ;;
+        --height)
+            HEIGHT="$2"
+            shift 2
+            ;;
+        --width)
+            WIDTH="$2"
+            shift 2
+            ;;
+        --num-steps)
+            NUM_STEPS="$2"
+            shift 2
+            ;;
+        --guidance)
+            GUIDANCE="$2"
+            shift 2
+            ;;
+        --seed)
+            SEED="$2"
+            shift 2
+            ;;
+        *)
+            echo "Unknown option: $1"
+            exit 1
+            ;;
+    esac
+done
+
+echo "=============================================="
+echo "GLM-Image Text-to-Image Generation"
+echo "=============================================="
+echo "Model: ${MODEL_PATH}"
+echo "Config: ${CONFIG_PATH}"
+echo "Prompt: ${PROMPT}"
+echo "Output: ${OUTPUT}"
+echo "Size: ${WIDTH}x${HEIGHT}"
+echo "Steps: ${NUM_STEPS}"
+echo "Guidance: ${GUIDANCE}"
+echo "Seed: ${SEED}"
+echo "=============================================="
+
+python end2end.py \
+    --model-path "${MODEL_PATH}" \
+    --config-path "${CONFIG_PATH}" \
+    --prompt "${PROMPT}" \
+    --output "${OUTPUT}" \
+    --height "${HEIGHT}" \
+    --width "${WIDTH}" \
+    --num-inference-steps "${NUM_STEPS}" \
+    --guidance-scale "${GUIDANCE}" \
+    --seed "${SEED}" \
+    --verbose

examples/offline_inference/helios/README.md

@@ -0,0 +1,193 @@
+# Helios Video Generation
+
+Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants:
+
+| Variant | Description | Key Features |
+|---------|-------------|--------------|
+| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` |
+| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support |
+| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` |
+
+## Setup
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+
+## Run Examples
+
+Get into the example folder:
+```bash
+cd examples/offline_inference/helios
+```
+
+### Text-to-Video (T2V)
+
+**Helios-Base** (Stage 1 only):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Base \
+    --sample-type t2v \
+    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
+    --guidance-scale 5.0 \
+    --output helios_t2v_base.mp4
+```
+
+**Helios-Mid** (Stage 2 + CFG-Zero*):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Mid \
+    --sample-type t2v \
+    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
+    --guidance-scale 5.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 20 20 20 \
+    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
+    --output helios_t2v_mid.mp4
+```
+
+**Helios-Distilled** (Stage 2 pyramid + DMD):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Distilled \
+    --sample-type t2v \
+    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
+    --num-frames 240 \
+    --guidance-scale 1.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 2 2 2 \
+    --is-amplify-first-chunk \
+    --output helios_t2v_distilled.mp4
+```
+
+### Image-to-Video (I2V)
+
+**Helios-Base**:
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Base \
+    --sample-type i2v \
+    --image-path /path/to/image.jpg \
+    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
+    --guidance-scale 5.0 \
+    --output helios_i2v_base.mp4
+```
+
+**Helios-Mid** (Stage 2 + CFG-Zero*):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Mid \
+    --sample-type i2v \
+    --image-path /path/to/image.jpg \
+    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
+    --guidance-scale 5.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 20 20 20 \
+    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
+    --output helios_i2v_mid.mp4
+```
+
+**Helios-Distilled**:
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Distilled \
+    --sample-type i2v \
+    --image-path /path/to/image.jpg \
+    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
+    --num-frames 240 \
+    --guidance-scale 1.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 2 2 2 \
+    --is-amplify-first-chunk \
+    --output helios_i2v_distilled.mp4
+```
+
+### Video-to-Video (V2V)
+
+**Helios-Base**:
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Base \
+    --sample-type v2v \
+    --video-path /path/to/video.mp4 \
+    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
+    --guidance-scale 5.0 \
+    --output helios_v2v_base.mp4
+```
+
+**Helios-Mid** (Stage 2 + CFG-Zero*):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Mid \
+    --sample-type v2v \
+    --video-path /path/to/video.mp4 \
+    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
+    --guidance-scale 5.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 20 20 20 \
+    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
+    --output helios_v2v_mid.mp4
+```
+
+**Helios-Distilled**:
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Distilled \
+    --sample-type v2v \
+    --video-path /path/to/video.mp4 \
+    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
+    --num-frames 240 \
+    --guidance-scale 1.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 2 2 2 \
+    --is-amplify-first-chunk \
+    --output helios_v2v_distilled.mp4
+```
+
+## Common Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path |
+| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` |
+| `--prompt` | — | Text prompt describing the video |
+| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) |
+| `--image-path` | — | Input image (required for `i2v`) |
+| `--video-path` | — | Input video (required for `v2v`) |
+| `--height` | `384` | Video height in pixels |
+| `--width` | `640` | Video width in pixels |
+| `--num-frames` | `99` | Number of output frames |
+| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) |
+| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) |
+| `--seed` | `42` | Random seed |
+| `--fps` | `16` | Output video frame rate |
+| `--output` | `helios_output.mp4` | Output file path |
+
+### Stage 2 / Pyramid Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising |
+| `--pyramid-num-stages` | `3` | Number of pyramid stages |
+| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage |
+| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) |
+
+### CFG-Zero* Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) |
+| `--use-zero-init` | off | Zero init for first steps |
+| `--zero-steps` | `1` | Number of zero-init steps |
+
+### Memory & Parallelism
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--vae-use-slicing` | off | Enable VAE slicing |
+| `--vae-use-tiling` | off | Enable VAE tiling |
+| `--enforce-eager` | off | Disable torch.compile |
+| `--enable-cpu-offload` | off | CPU offloading |
+| `--enable-layerwise-offload` | off | Layerwise offloading |
+| `--ulysses-degree` | `1` | Ulysses SP degree |
+| `--ring-degree` | `1` | Ring SP degree |
+| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) |
+| `--tensor-parallel-size` | `1` | Tensor parallelism size |

examples/offline_inference/helios/end2end.py

@@ -0,0 +1,343 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Helios video generation example supporting T2V, I2V, and V2V.
+
+Usage (T2V, Helios-Base, Stage 1 only):
+    python end2end.py \
+        --model /path/to/Helios-Base --sample-type t2v \
+        --prompt "A serene lakeside sunrise with mist over the water." \
+        --height 384 --width 640 --num-frames 99 \
+        --num-inference-steps 50 --guidance-scale 5.0
+
+Usage (I2V, Helios-Base):
+    python end2end.py \
+        --model /path/to/Helios-Base --sample-type i2v \
+        --image-path /path/to/image.jpg \
+        --prompt "Description of desired animation." \
+        --guidance-scale 5.0
+
+Usage (V2V, Helios-Base):
+    python end2end.py \
+        --model /path/to/Helios-Base --sample-type v2v \
+        --video-path /path/to/video.mp4 \
+        --prompt "Description of desired transformation." \
+        --guidance-scale 5.0
+
+Usage (Helios-Mid, Stage 2 + CFG-Zero*):
+    python end2end.py \
+        --model /path/to/Helios-Mid --sample-type t2v \
+        --prompt "A serene lakeside sunrise with mist over the water." \
+        --guidance-scale 5.0 \
+        --is-enable-stage2 \
+        --pyramid-num-inference-steps-list 20 20 20 \
+        --use-cfg-zero-star --use-zero-init --zero-steps 1
+
+Usage (Helios-Distilled, Stage 2 pyramid + DMD):
+    python end2end.py \
+        --model /path/to/Helios-Distilled --sample-type t2v \
+        --prompt "A serene lakeside sunrise with mist over the water." \
+        --guidance-scale 1.0 \
+        --is-enable-stage2 \
+        --pyramid-num-inference-steps-list 2 2 2 \
+        --is-amplify-first-chunk
+"""
+
+import argparse
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+
+def load_image_as_tensor(image_path: str, height: int, width: int) -> torch.Tensor:
+    """Load an image and return a (1, C, H, W) tensor normalized to [-1, 1]."""
+    from PIL import Image
+    from torchvision import transforms
+
+    image = Image.open(image_path).convert("RGB").resize((width, height))
+    transform = transforms.Compose(
+        [
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ]
+    )
+    return transform(image).unsqueeze(0)
+
+
+def load_video_as_tensor(video_path: str, height: int, width: int) -> torch.Tensor:
+    """Load a video and return a (1, C, T, H, W) tensor normalized to [-1, 1]."""
+    from torchvision import transforms
+    from torchvision.io import read_video
+
+    video_frames, _, _ = read_video(video_path, output_format="TCHW")
+    video_frames = video_frames.float() / 255.0
+
+    transform = transforms.Compose(
+        [
+            transforms.Resize((height, width), antialias=True),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ]
+    )
+    video_frames = torch.stack([transform(f) for f in video_frames])
+    return video_frames.permute(1, 0, 2, 3).unsqueeze(0)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate a video with Helios (T2V / I2V / V2V).")
+    parser.add_argument(
+        "--model",
+        default="BestWishYsh/Helios-Base",
+        help="Helios model ID or local path (e.g. Helios-Base, Helios-Mid, Helios-Distilled).",
+    )
+    parser.add_argument(
+        "--sample-type",
+        choices=["t2v", "i2v", "v2v"],
+        default="t2v",
+        help="Generation mode: t2v (text-to-video), i2v (image-to-video), v2v (video-to-video).",
+    )
+    parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
+    parser.add_argument(
+        "--negative-prompt",
+        default=(
+            "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, "
+            "images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, "
+            "incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, "
+            "misshapen limbs, fused fingers, still picture, messy background, three legs, many people "
+            "in the background, walking backwards"
+        ),
+        help="Negative prompt.",
+    )
+    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
+    parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.")
+    parser.add_argument("--height", type=int, default=384, help="Video height.")
+    parser.add_argument("--width", type=int, default=640, help="Video width.")
+    parser.add_argument("--num-frames", type=int, default=99, help="Number of video frames.")
+    parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps (Stage 1 only).")
+    parser.add_argument("--output", type=str, default="helios_output.mp4", help="Output video path.")
+    parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
+
+    # I2V / V2V
+    parser.add_argument("--image-path", type=str, default=None, help="Input image path for I2V mode.")
+    parser.add_argument("--video-path", type=str, default=None, help="Input video path for V2V mode.")
+
+    # Stage 2 (pyramid multi-stage denoising)
+    parser.add_argument(
+        "--is-enable-stage2",
+        action="store_true",
+        help="Enable pyramid multi-stage denoising (Stage 2). Required for Helios-Distilled.",
+    )
+    parser.add_argument(
+        "--pyramid-num-stages",
+        type=int,
+        default=3,
+        help="Number of pyramid stages for Stage 2.",
+    )
+    parser.add_argument(
+        "--pyramid-num-inference-steps-list",
+        type=int,
+        nargs="+",
+        default=[10, 10, 10],
+        help="Inference steps per pyramid stage.",
+    )
+
+    # DMD
+    parser.add_argument(
+        "--is-amplify-first-chunk",
+        action="store_true",
+        help="Enable DMD amplification for the first chunk (Helios-Distilled).",
+    )
+
+    # CFG Zero Star
+    parser.add_argument(
+        "--use-cfg-zero-star",
+        action="store_true",
+        help="Enable CFG Zero Star guidance (recommended for Helios-Mid).",
+    )
+    parser.add_argument(
+        "--use-zero-init",
+        action="store_true",
+        help="Use zero initialization for the first denoising steps with CFG-Zero*.",
+    )
+    parser.add_argument(
+        "--zero-steps",
+        type=int,
+        default=1,
+        help="Number of initial denoising steps using zero prediction (default: 1).",
+    )
+
+    # Memory & parallelism
+    parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
+    parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
+    parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
+    parser.add_argument("--enable-cpu-offload", action="store_true", help="Enable CPU offloading.")
+    parser.add_argument("--enable-layerwise-offload", action="store_true", help="Enable layerwise offloading.")
+    parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
+    parser.add_argument("--ring-degree", type=int, default=1, help="Ring SP degree.")
+    parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
+    parser.add_argument("--tensor-parallel-size", type=int, default=1, help="Tensor parallelism size.")
+
+    return parser.parse_args()
+
+
+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,
+        tensor_parallel_size=args.tensor_parallel_size,
+    )
+
+    omni = Omni(
+        model=args.model,
+        enable_layerwise_offload=args.enable_layerwise_offload,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
+        enable_cpu_offload=args.enable_cpu_offload,
+        parallel_config=parallel_config,
+        enforce_eager=args.enforce_eager,
+    )
+
+    # Validate I2V / V2V arguments
+    if args.sample_type == "i2v" and not args.image_path:
+        raise ValueError("--image-path is required for I2V mode.")
+    if args.sample_type == "v2v" and not args.video_path:
+        raise ValueError("--video-path is required for V2V mode.")
+
+    # Build extra_args for Helios-specific parameters
+    extra_args = {}
+    if args.is_enable_stage2:
+        extra_args["is_enable_stage2"] = True
+        extra_args["pyramid_num_stages"] = args.pyramid_num_stages
+        extra_args["pyramid_num_inference_steps_list"] = args.pyramid_num_inference_steps_list
+    if args.is_amplify_first_chunk:
+        extra_args["is_amplify_first_chunk"] = True
+    if args.use_cfg_zero_star:
+        extra_args["use_cfg_zero_star"] = True
+    if args.use_zero_init:
+        extra_args["use_zero_init"] = True
+        extra_args["zero_steps"] = args.zero_steps
+
+    if args.sample_type == "i2v":
+        image_tensor = load_image_as_tensor(args.image_path, args.height, args.width)
+        extra_args["image"] = image_tensor
+    elif args.sample_type == "v2v":
+        video_tensor = load_video_as_tensor(args.video_path, args.height, args.width)
+        extra_args["video"] = video_tensor
+
+    # Print generation configuration
+    print(f"\n{'=' * 60}")
+    print("Helios Generation Configuration:")
+    print(f"  Model: {args.model}")
+    print(f"  Sample type: {args.sample_type.upper()}")
+    if args.sample_type == "i2v":
+        print(f"  Image: {args.image_path}")
+    elif args.sample_type == "v2v":
+        print(f"  Video: {args.video_path}")
+    print(f"  Prompt: {args.prompt[:80]}{'...' if len(args.prompt) > 80 else ''}")
+    print(f"  Video size: {args.width}x{args.height}, {args.num_frames} frames")
+    print(f"  Inference steps: {args.num_inference_steps}")
+    print(f"  Guidance scale: {args.guidance_scale}")
+    if args.is_enable_stage2:
+        print(f"  Stage 2: enabled (stages={args.pyramid_num_stages}, steps={args.pyramid_num_inference_steps_list})")
+        if args.is_amplify_first_chunk:
+            print("  DMD amplify first chunk: enabled")
+        if args.use_cfg_zero_star:
+            print(f"  CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})")
+    else:
+        if args.use_cfg_zero_star:
+            print(f"  CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})")
+        print("  Stage 2: disabled (Stage 1 only)")
+    print(f"{'=' * 60}\n")
+
+    generation_start = time.perf_counter()
+    frames = omni.generate(
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+        },
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_frames=args.num_frames,
+            extra_args=extra_args,
+        ),
+    )
+    generation_end = time.perf_counter()
+    generation_time = generation_end - generation_start
+
+    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
+
+    # Extract video frames from OmniRequestOutput
+    if isinstance(frames, list) and len(frames) > 0:
+        first_item = frames[0]
+
+        if hasattr(first_item, "final_output_type"):
+            if first_item.final_output_type != "image":
+                raise ValueError(
+                    f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
+                )
+
+            if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
+                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
+                    inner_output = first_item.request_output[0]
+                    if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
+                        frames = inner_output.images[0] if inner_output.images else None
+                        if frames is None:
+                            raise ValueError("No video frames found in output.")
+            elif hasattr(first_item, "images") and first_item.images:
+                frames = first_item.images
+            else:
+                raise ValueError("No video frames found in OmniRequestOutput.")
+
+    # Unwrap batch list from postprocess_video: [numpy(T,H,W,C)] -> numpy(T,H,W,C)
+    if isinstance(frames, list) and len(frames) > 0 and isinstance(frames[0], np.ndarray):
+        frames = frames[0]
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        from diffusers.utils import export_to_video
+    except ImportError:
+        raise ImportError("diffusers is required for export_to_video.")
+
+    if isinstance(frames, torch.Tensor):
+        video_tensor = frames.detach().cpu()
+        if video_tensor.dim() == 5:
+            if video_tensor.shape[1] in (3, 4):
+                video_tensor = video_tensor[0].permute(1, 2, 3, 0)
+            else:
+                video_tensor = video_tensor[0]
+        elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
+            video_tensor = video_tensor.permute(1, 2, 3, 0)
+        if video_tensor.is_floating_point():
+            video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
+        video_array = video_tensor.float().numpy()
+    elif isinstance(frames, np.ndarray):
+        video_array = frames
+    else:
+        video_array = frames
+
+    if isinstance(video_array, np.ndarray) and video_array.ndim == 5:
+        video_array = video_array[0]
+
+    export_to_video(video_array, str(output_path), fps=args.fps)
+    print(f"Saved generated video to {output_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/hunyuan_image3/README.md

@@ -0,0 +1,25 @@
+# HunyuanImage-3.0 Image-to-Text Inference
+
+This example demonstrates how to run HunyuanImage-3.0 Image-to-Text with the vLLM-Omni.
+
+## Local CLI Usage
+
+Download the example image:
+
+```bash
+wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
+```
+
+Run example:
+
+```bash
+python image_to_text.py \
+  --image cherry_blossom.jpg \
+  --prompt "<|startoftext|>You are an assistant that understands images and outputs text.<img>Describe the content of the picture."
+```
+
+Key arguments:
+
+- `--model`: Model used. Default is: tencent/HunyuanImage-3.0-Instruct (Optional).
+- `--image`: Path to input image (required).
+- `--prompt`: Text description used to guide image understanding (required).

examples/offline_inference/hunyuan_image3/image_to_text.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import os
+
+from PIL import Image
+
+from vllm_omni.entrypoints.omni import Omni
+
+"""
+The tencent/HunyuanImage-3.0-Instruct base model is built on the Hunyuan v1 architecture, specifically the tencent/Hunyuan-A13B-Instruct model. It utilizes two tokenizer delimiter templates:
+
+1) Pretrained template (default for gen_text mode), which concatenates system, image
+   tokens, and user question WITHOUT role delimiters:
+"<|startoftext|>{system_prompt}{image_tokens}{user_question}"
+
+   Example (before image token expansion):
+"<|startoftext|>You are an assistant that understands images and outputs text.<img>Describe the content of the picture."
+
+2) Instruct template, which uses explicit role prefixes and separators.
+"""
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate text from image using HunyuanImage-3.0-Instruct.")
+    parser.add_argument(
+        "--model",
+        default="tencent/HunyuanImage-3.0-Instruct",
+        help="Model name or local path.",
+    )
+    parser.add_argument(
+        "--image",
+        type=str,
+        required=True,
+        help="Path to input image file (PNG, JPG, etc.).",
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        required=True,
+        help="Pretrain template prompt: <|startoftext|>{system}<img>{question}",
+    )
+    return parser.parse_args()
+
+
+def load_image(image_path: str) -> Image.Image:
+    """Load an image from file path."""
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+    return Image.open(image_path).convert("RGB")
+
+
+def main(args: argparse.Namespace) -> None:
+    omni = Omni(model=args.model)
+
+    prompt_dict = {
+        "prompt": args.prompt,
+        "modalities": ["text"],
+    }
+
+    # Add image input if provided
+    if args.image:
+        if not os.path.exists(args.image):
+            raise FileNotFoundError(f"Input image not found: {args.image}")
+
+        input_image = load_image(args.image)
+        prompt_dict["multi_modal_data"] = {"image": input_image}
+
+    prompts = [prompt_dict]
+    omni_outputs = omni.generate(prompts=prompts)
+
+    prompt_text = omni_outputs[0].request_output[0].prompt
+    generated_text = omni_outputs[0].request_output[0].outputs[0].text
+    print(f"Prompt: {prompt_text}")
+    print(f"Text: {generated_text}")
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

examples/offline_inference/image_to_video/README.md

@@ -65,6 +65,12 @@ Key arguments:
 - `--vae-use-slicing`: Enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: Enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+- `--use-hsdp`: Enable Hybrid Sharded Data Parallel to shard model weights across GPUs.
+- `--hsdp-shard-size`: Number of GPUs to shard model weights across within each replica group. -1 (default) auto-calculates as world_size / replicate_size.
+- `--hsdp-replicate-size`: Number of replica groups for HSDP. Each replica holds a full sharded copy. Default 1 means pure sharding (no replication).
+
+
 
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.

examples/offline_inference/image_to_video/image_to_video.py

@@ -2,11 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-Image-to-Video generation example using Wan2.2 I2V/TI2V models.
+Image-to-Video generation example using Wan2.2 I2V/TI2V models or LTX2.
 
 Supports:
 - Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
 - Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B)
+- LTX2 image-to-video pipeline
 
 Usage:
     # I2V-A14B (MoE)
@@ -16,10 +17,18 @@ Usage:
     # TI2V-5B (unified)
     python image_to_video.py --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
         --image input.jpg --prompt "A cat playing with yarn"
+
+    # LTX2 image-to-video
+    python image_to_video.py --model /path/to/LTX-2 \
+        --model-class-name LTX2ImageToVideoPipeline \
+        --image input.jpg --prompt "A cinematic dolly shot of a boat" \
+        --num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
+        --frame-rate 24 --fps 24 --output ltx2_i2v.mp4
 """
 
 import argparse
 import os
+import time
 from pathlib import Path
 
 import numpy as np
@@ -34,12 +43,17 @@ from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Generate a video from an image with Wan2.2 I2V/TI2V.")
+    parser = argparse.ArgumentParser(description="Generate a video from an image with Wan2.2 or LTX2.")
     parser.add_argument(
         "--model",
         default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
         help="Diffusers Wan2.2 I2V model ID or local path.",
     )
+    parser.add_argument(
+        "--model-class-name",
+        default=None,
+        help="Override model class name (e.g., LTX2ImageToVideoPipeline).",
+    )
     parser.add_argument("--image", required=True, help="Path to input image.")
     parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.")
     parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
@@ -55,11 +69,17 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
     parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps.")
     parser.add_argument("--boundary-ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.")
+    parser.add_argument(
+        "--frame-rate",
+        type=float,
+        default=None,
+        help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
+    )
     parser.add_argument(
         "--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
     )
     parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
-    parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
+    parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.")
     parser.add_argument(
         "--vae-use-slicing",
         action="store_true",
@@ -80,6 +100,24 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--ulysses-degree",
+        type=int,
+        default=1,
+        help="Number of GPUs used for ulysses sequence parallelism.",
+    )
+    parser.add_argument(
+        "--ring-degree",
+        type=int,
+        default=1,
+        help="Number of GPUs used for ring sequence parallelism.",
+    )
+    parser.add_argument(
+        "--tensor-parallel-size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
+    )
     parser.add_argument(
         "--cfg-parallel-size",
         type=int,
@@ -87,18 +125,67 @@ def parse_args() -> argparse.Namespace:
         choices=[1, 2],
         help="Number of GPUs used for classifier free guidance parallel size.",
     )
+    parser.add_argument(
+        "--vae-patch-parallel-size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for VAE patch/tile parallelism (decode).",
+    )
     parser.add_argument(
         "--enforce-eager",
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
+    parser.add_argument(
+        "--audio-sample-rate",
+        type=int,
+        default=24000,
+        help="Sample rate for audio output when saved (default: 24000 for LTX2).",
+    )
+    parser.add_argument(
+        "--cache-backend",
+        type=str,
+        default=None,
+        choices=["cache_dit", "tea_cache"],
+        help=(
+            "Cache backend to use for acceleration. "
+            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
+            "Default: None (no cache acceleration)."
+        ),
+    )
+    parser.add_argument(
+        "--use-hsdp",
+        action="store_true",
+        help=("Enable Hybrid Sharded Data Parallel to shard model weights across GPUs. "),
+    )
+    parser.add_argument(
+        "--hsdp-shard-size",
+        type=int,
+        default=-1,
+        help=(
+            "Number of GPUs to shard model weights across within each replica group. "
+            "-1 (default) auto-calculates as world_size / replicate_size. "
+        ),
+    )
+    parser.add_argument(
+        "--hsdp-replicate-size",
+        type=int,
+        default=1,
+        help=(
+            "Number of replica groups for HSDP. Each replica holds a full sharded copy. "
+            "Default 1 means pure sharding (no replication). "
+        ),
+    )
     return parser.parse_args()
 
 
-def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> tuple[int, int]:
+def calculate_dimensions(
+    image: PIL.Image.Image,
+    max_area: int = 480 * 832,
+    mod_value: int = 16,
+) -> tuple[int, int]:
     """Calculate output dimensions maintaining aspect ratio."""
     aspect_ratio = image.height / image.width
-    mod_value = 16  # Must be divisible by 16
 
     height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
     width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
@@ -109,26 +196,77 @@ def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> t
 def main():
     args = parse_args()
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    model_name = str(args.model).lower() if args.model is not None else ""
+    model_class_name = args.model_class_name
+    is_ltx2 = "ltx2" in model_name or (model_class_name and "ltx2" in model_class_name.lower())
+    if model_class_name is None and is_ltx2:
+        model_class_name = "LTX2ImageToVideoPipeline"
 
     # Load input image
     image = PIL.Image.open(args.image).convert("RGB")
 
+    fps = args.fps if args.fps is not None else (24 if is_ltx2 else 16)
+    frame_rate = args.frame_rate if args.frame_rate is not None else float(fps)
+    guidance_scale = args.guidance_scale if args.guidance_scale is not None else (4.0 if is_ltx2 else 5.0)
+    num_frames = args.num_frames if args.num_frames is not None else (121 if is_ltx2 else 81)
+    num_inference_steps = args.num_inference_steps if args.num_inference_steps is not None else (40 if is_ltx2 else 50)
+
     # Calculate dimensions if not provided
     height = args.height
     width = args.width
     if height is None or width is None:
-        # Default to 480P area for I2V
-        calc_height, calc_width = calculate_dimensions(image, max_area=480 * 832)
+        # Default to 480P area for Wan2.2 I2V, 512x768 area for LTX2
+        max_area = 512 * 768 if is_ltx2 else 480 * 832
+        mod_value = 32 if is_ltx2 else 16
+        calc_height, calc_width = calculate_dimensions(image, max_area=max_area, mod_value=mod_value)
         height = height or calc_height
         width = width or calc_width
 
     # Resize image to target dimensions
     image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
 
+    # Configure cache based on backend type
+    cache_config = None
+    if args.cache_backend == "cache_dit":
+        if is_ltx2:
+            cache_config = {
+                "Fn_compute_blocks": 2,
+                "Bn_compute_blocks": 0,
+                "max_warmup_steps": 8,
+                "residual_diff_threshold": 0.12,
+                "max_continuous_cached_steps": 1,
+                "max_cached_steps": 20,
+                "enable_taylorseer": False,
+                "scm_steps_mask_policy": None,
+            }
+        else:
+            cache_config = {
+                "Fn_compute_blocks": 1,
+                "Bn_compute_blocks": 0,
+                "max_warmup_steps": 4,
+                "residual_diff_threshold": 0.24,
+                "max_continuous_cached_steps": 3,
+                "enable_taylorseer": False,
+                "taylorseer_order": 1,
+                "scm_steps_mask_policy": None,
+                "scm_steps_policy": "dynamic",
+            }
+    elif args.cache_backend == "tea_cache":
+        cache_config = {
+            "rel_l1_thresh": 0.2,
+        }
+
     # Check if profiling is requested via environment variable
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
     parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
         cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+        vae_patch_parallel_size=args.vae_patch_parallel_size,
+        use_hsdp=args.use_hsdp,
+        hsdp_shard_size=args.hsdp_shard_size,
+        hsdp_replicate_size=args.hsdp_replicate_size,
     )
     omni = Omni(
         model=args.model,
@@ -140,6 +278,9 @@ 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,
     )
 
     if profiler_enabled:
@@ -152,10 +293,14 @@ def main():
     print(f"  Model: {args.model}")
     print(f"  Inference steps: {args.num_inference_steps}")
     print(f"  Frames: {args.num_frames}")
-    print(f"  Parallel configuration: cfg_parallel_size={args.cfg_parallel_size}")
+    print(
+        f"  Parallel configuration: cfg_parallel_size={args.cfg_parallel_size},"
+        f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}"
+    )
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")
 
+    generation_start = time.perf_counter()
     # omni.generate() returns Generator[OmniRequestOutput, None, None]
     frames = omni.generate(
         {
@@ -167,38 +312,69 @@ def main():
             height=height,
             width=width,
             generator=generator,
-            guidance_scale=args.guidance_scale,
+            guidance_scale=guidance_scale,
             guidance_scale_2=args.guidance_scale_high,
-            num_inference_steps=args.num_inference_steps,
-            num_frames=args.num_frames,
+            num_inference_steps=num_inference_steps,
+            num_frames=num_frames,
+            frame_rate=frame_rate,
         ),
     )
+    generation_end = time.perf_counter()
+    generation_time = generation_end - generation_start
 
-    # Extract video frames from OmniRequestOutput
-    if isinstance(frames, list) and len(frames) > 0:
-        first_item = frames[0]
+    # Print profiling results
+    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
-        # Check if it's an OmniRequestOutput
-        if hasattr(first_item, "final_output_type"):
-            if first_item.final_output_type != "image":
-                raise ValueError(
-                    f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
-                )
+    audio = None
+    if isinstance(frames, list):
+        frames = frames[0] if frames else None
 
-            # Pipeline mode: extract from nested request_output
-            if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
-                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
-                    inner_output = first_item.request_output[0]
-                    if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
-                        frames = inner_output.images[0] if inner_output.images else None
-                        if frames is None:
-                            raise ValueError("No video frames found in output.")
-            # Diffusion mode: use direct images field
-            elif hasattr(first_item, "images") and first_item.images:
-                frames = first_item.images
+    if isinstance(frames, OmniRequestOutput):
+        if frames.final_output_type != "image":
+            raise ValueError(
+                f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
+            )
+        if frames.multimodal_output and "audio" in frames.multimodal_output:
+            audio = frames.multimodal_output["audio"]
+        if frames.is_pipeline_output and frames.request_output is not None:
+            inner_output = frames.request_output
+            if isinstance(inner_output, list):
+                inner_output = inner_output[0] if inner_output else None
+            if isinstance(inner_output, OmniRequestOutput):
+                if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
+                    audio = inner_output.multimodal_output["audio"]
+                frames = inner_output
+        if isinstance(frames, OmniRequestOutput):
+            if frames.images:
+                if len(frames.images) == 1 and isinstance(frames.images[0], tuple) and len(frames.images[0]) == 2:
+                    frames, audio = frames.images[0]
+                elif len(frames.images) == 1 and isinstance(frames.images[0], dict):
+                    audio = frames.images[0].get("audio")
+                    frames = frames.images[0].get("frames") or frames.images[0].get("video")
+                else:
+                    frames = frames.images
             else:
                 raise ValueError("No video frames found in OmniRequestOutput.")
 
+    if isinstance(frames, list) and frames:
+        first_item = frames[0]
+        if isinstance(first_item, tuple) and len(first_item) == 2:
+            frames, audio = first_item
+        elif isinstance(first_item, dict):
+            audio = first_item.get("audio")
+            frames = first_item.get("frames") or first_item.get("video")
+        elif isinstance(first_item, list):
+            frames = first_item
+
+    if isinstance(frames, tuple) and len(frames) == 2:
+        frames, audio = frames
+    elif isinstance(frames, dict):
+        audio = frames.get("audio")
+        frames = frames.get("frames") or frames.get("video")
+
+    if frames is None:
+        raise ValueError("No video frames found in output.")
+
     output_path = Path(args.output)
     output_path.parent.mkdir(parents=True, exist_ok=True)
 
@@ -207,32 +383,128 @@ def main():
     except ImportError:
         raise ImportError("diffusers is required for export_to_video.")
 
-    # frames may be np.ndarray (preferred) or torch.Tensor
+    def _normalize_frame(frame):
+        if isinstance(frame, torch.Tensor):
+            frame_tensor = frame.detach().cpu()
+            if frame_tensor.dim() == 4 and frame_tensor.shape[0] == 1:
+                frame_tensor = frame_tensor[0]
+            if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4):
+                frame_tensor = frame_tensor.permute(1, 2, 0)
+            if frame_tensor.is_floating_point():
+                frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5
+            return frame_tensor.float().numpy()
+        if isinstance(frame, np.ndarray):
+            frame_array = frame
+            if frame_array.ndim == 4 and frame_array.shape[0] == 1:
+                frame_array = frame_array[0]
+            if np.issubdtype(frame_array.dtype, np.integer):
+                frame_array = frame_array.astype(np.float32) / 255.0
+            return frame_array
+        try:
+            from PIL import Image
+        except ImportError:
+            Image = None
+        if Image is not None and isinstance(frame, Image.Image):
+            return np.asarray(frame).astype(np.float32) / 255.0
+        return frame
+
+    def _ensure_frame_list(video_array):
+        if isinstance(video_array, list):
+            if len(video_array) == 0:
+                return video_array
+            first_item = video_array[0]
+            if isinstance(first_item, np.ndarray):
+                if first_item.ndim == 5:
+                    return list(first_item[0])
+                if first_item.ndim == 4:
+                    if len(video_array) == 1:
+                        return list(first_item)
+                    return list(first_item)
+                if first_item.ndim == 3:
+                    return video_array
+            return video_array
+        if isinstance(video_array, np.ndarray):
+            if video_array.ndim == 5:
+                return list(video_array[0])
+            if video_array.ndim == 4:
+                return list(video_array)
+            if video_array.ndim == 3:
+                return [video_array]
+        return video_array
+
+    # frames may be np.ndarray, torch.Tensor, or list of tensors/arrays/images
     # export_to_video expects a list of frames with values in [0, 1]
     if isinstance(frames, torch.Tensor):
         video_tensor = frames.detach().cpu()
         if video_tensor.dim() == 5:
-            # [B, C, F, H, W] or [B, F, H, W, C]
             if video_tensor.shape[1] in (3, 4):
                 video_tensor = video_tensor[0].permute(1, 2, 3, 0)
             else:
                 video_tensor = video_tensor[0]
         elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
             video_tensor = video_tensor.permute(1, 2, 3, 0)
-        # If float, assume [-1,1] and normalize to [0,1]
         if video_tensor.is_floating_point():
             video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
         video_array = video_tensor.float().numpy()
+    elif isinstance(frames, np.ndarray):
+        video_array = frames
+        if video_array.ndim == 5:
+            video_array = video_array[0]
+        if np.issubdtype(video_array.dtype, np.integer):
+            video_array = video_array.astype(np.float32) / 255.0
+    elif isinstance(frames, list):
+        if len(frames) == 0:
+            raise ValueError("No video frames found in output.")
+        video_array = [_normalize_frame(frame) for frame in frames]
     else:
         video_array = frames
-        if hasattr(video_array, "shape") and video_array.ndim == 5:
-            video_array = video_array[0]
 
-    # Convert 4D array (frames, H, W, C) to list of frames for export_to_video
-    if isinstance(video_array, np.ndarray) and video_array.ndim == 4:
-        video_array = list(video_array)
+    video_array = _ensure_frame_list(video_array)
 
-    export_to_video(video_array, str(output_path), fps=args.fps)
+    use_ltx2_export = is_ltx2
+    encode_video = None
+    if use_ltx2_export:
+        try:
+            from diffusers.pipelines.ltx2.export_utils import encode_video
+        except ImportError:
+            encode_video = None
+
+    if use_ltx2_export and encode_video is not None:
+        if isinstance(video_array, list):
+            frames_np = np.stack(video_array, axis=0)
+        elif isinstance(video_array, np.ndarray):
+            frames_np = video_array
+        else:
+            frames_np = np.asarray(video_array)
+
+        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):
+                audio = audio[0] if audio else None
+            if isinstance(audio, np.ndarray):
+                audio = torch.from_numpy(audio)
+            if isinstance(audio, torch.Tensor):
+                audio_out = audio
+                if audio_out.dim() > 1:
+                    audio_out = audio_out[0]
+                audio_out = audio_out.float().cpu()
+
+        encode_video(
+            video_tensor,
+            fps=fps,
+            audio=audio_out,
+            audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
+            output_path=str(output_path),
+        )
+    else:
+        export_to_video(video_array, str(output_path), fps=fps)
     print(f"Saved generated video to {output_path}")
 
     if profiler_enabled:

examples/offline_inference/mammothmodal2_preview/README.md

@@ -0,0 +1,32 @@
+# MammothModa2-Preview
+
+## Run examples (MammothModa2-Preview)
+
+Download model
+```bash
+hf bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
+```
+
+### Text-to-Image (T2I)
+
+```bash
+python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py \
+  --model ./MammothModa2-Preview \
+  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
+  --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+  --height 1024 \
+  --width 1024 \
+  --num-inference-steps 50 \
+  --text-guidance-scale 4.0 \
+  --out output.png
+```
+
+### Image Summary
+
+```bash
+python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
+  --model ./MammothModa2-Preview \
+  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
+  --question "Summarize this image." \
+  --image ./image.png
+```

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -0,0 +1,115 @@
+"""
+Offline inference example: MammothModa2 image summarization (single AR stage).
+
+Example:
+  uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
+    --model path/to/MammothModa2-Preview \
+    --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
+    --image /path/to/input.jpg \
+    --question "Please summarize the content of this image."
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+from PIL import Image
+from vllm import SamplingParams
+from vllm.multimodal.image import convert_image_mode
+
+from vllm_omni import Omni
+
+DEFAULT_SYSTEM = "You are a helpful assistant."
+DEFAULT_QUESTION = "Please summarize the content of this image."
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
+    parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
+    parser.add_argument(
+        "--stage-config", type=str, required=True, help="Path to stage config yaml (single-stage AR->text)."
+    )
+    parser.add_argument("--image", type=str, required=True, help="Path to input image.")
+    parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
+    parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=512,
+        help="Max new tokens to generate.",
+    )
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--top-p", type=float, default=0.9)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--trust-remote-code", action="store_true")
+    return parser.parse_args()
+
+
+def build_prompt(system: str, question: str) -> str:
+    return (
+        f"<|im_start|>system\n{system}<|im_end|>\n"
+        "<|im_start|>user\n"
+        "<|vision_start|><|image_pad|><|vision_end|>"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+
+
+def main() -> None:
+    args = parse_args()
+
+    if not os.path.exists(args.image):
+        raise FileNotFoundError(f"Image file not found: {args.image}")
+
+    pil_image = Image.open(args.image)
+    image_data = convert_image_mode(pil_image, "RGB")
+    prompt = build_prompt(args.system, args.question)
+
+    omni = Omni(
+        model=args.model,
+        stage_configs_path=args.stage_config,
+        trust_remote_code=args.trust_remote_code,
+    )
+    try:
+        sp = SamplingParams(
+            temperature=float(args.temperature),
+            top_p=float(args.top_p),
+            top_k=-1,
+            max_tokens=int(args.max_tokens),
+            seed=int(args.seed),
+            detokenize=True,
+        )
+        # NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None].
+        # Consume it inside the try block so the worker isn't closed early.
+        outputs = list(
+            omni.generate(
+                [
+                    {
+                        "prompt": prompt,
+                        "multi_modal_data": {"image": image_data},
+                        "additional_information": {"omni_task": ["chat"]},
+                    }
+                ],
+                [sp],
+            )
+        )
+    finally:
+        omni.close()
+
+    lines: list[str] = []
+    for stage_outputs in outputs:
+        req_outputs = getattr(stage_outputs, "request_output", stage_outputs)
+        req_outputs = req_outputs if isinstance(req_outputs, list) else [req_outputs]
+        for ro in req_outputs:
+            text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
+            lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
+            lines.append("answer:\n")
+            lines.append(text.strip() + "\n")
+            lines.append("\n")
+
+    print("\n".join(lines))
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py

@@ -0,0 +1,251 @@
+"""
+Offline inference example for MammothModa2 Text-to-Image (T2I) generation.
+This script uses the vllm_omni.Omni pipeline with a multi-stage configuration.
+
+Workflow:
+1. Stage 0 (AR): Generates visual tokens and their corresponding hidden states.
+2. Stage 1 (DiT): Consumes the hidden states as conditions to perform diffusion
+   and VAE decoding to produce the final image.
+
+Example Usage:
+    uv run python examples/offline_inference/run_mammothmoda2_t2i.py \
+        --model path/to/MammothModa2-Preview \
+        --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
+        --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+        --out output.png
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+from pathlib import Path
+from typing import NamedTuple
+
+import torch
+from PIL import Image
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni import Omni
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+_PATCH_SIZE = 16  # AR image grid patch size (pixels per token)
+
+
+class T2IGenConfig(NamedTuple):
+    eol_token_id: int
+    visual_token_start_id: int
+    visual_token_end_id: int
+    top_k: int  # AR sampling top-k (covers the full visual generation vocabulary)
+    # Qwen2.5-VL special vision tokens: <|image_pad|>, <|video_pad|>, <|vision_start|>, <|vision_end|>
+    visual_ids: list[int]
+
+
+def load_t2i_generation_config(model_dir: str) -> T2IGenConfig:
+    """Load T2I token IDs from t2i_generation_config.json and config.json.
+
+    Supports both local directory paths and HuggingFace Hub model IDs.
+    """
+    model_path = Path(model_dir)
+
+    def _read_json(filename: str) -> dict:
+        local = model_path / filename
+        if local.exists():
+            with local.open(encoding="utf-8") as f:
+                return json.load(f)
+        # Fall back to HuggingFace Hub when model_dir is a repo ID.
+        try:
+            from huggingface_hub import hf_hub_download
+        except ImportError as exc:
+            raise ImportError(
+                "huggingface_hub is required to load configs from HF Hub. Install it with: pip install huggingface_hub"
+            ) from exc
+        cached = hf_hub_download(repo_id=model_dir, filename=filename)
+        with open(cached, encoding="utf-8") as f:
+            return json.load(f)
+
+    gen_cfg = _read_json("t2i_generation_config.json")
+    llm_cfg = _read_json("config.json").get("llm_config", {})
+
+    return T2IGenConfig(
+        eol_token_id=int(gen_cfg["eol_token_id"]),
+        visual_token_start_id=int(gen_cfg["visual_token_start_id"]),
+        visual_token_end_id=int(gen_cfg["visual_token_end_id"]),
+        top_k=int(gen_cfg["top_k"]),
+        visual_ids=[
+            int(llm_cfg["image_token_id"]),
+            int(llm_cfg["video_token_id"]),
+            int(llm_cfg["vision_start_token_id"]),
+            int(llm_cfg["vision_end_token_id"]),
+        ],
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Run MammothModa2 T2I (AR -> DiT) with vLLM-Omni.")
+    p.add_argument("--model", type=str, required=True, help="Path to the model directory.")
+    p.add_argument("--stage-config", type=str, required=True, help="Path to the multi-stage YAML configuration.")
+    p.add_argument(
+        "--prompt",
+        type=str,
+        action="append",
+        default=None,
+        help=(
+            "Text prompt for image generation. Can be provided multiple times "
+            "to generate multiple images with shared height/width/CFG settings."
+        ),
+    )
+    p.add_argument("--height", type=int, default=1024, help="Output image height (must be a multiple of 16).")
+    p.add_argument("--width", type=int, default=1024, help="Output image width (must be a multiple of 16).")
+    p.add_argument("--num-inference-steps", type=int, default=50, help="Number of diffusion steps for the DiT stage.")
+    p.add_argument(
+        "--text-guidance-scale", type=float, default=9.0, help="Classifier-Free Guidance (CFG) scale for DiT."
+    )
+    p.add_argument(
+        "--cfg-range",
+        type=float,
+        nargs=2,
+        default=(0.0, 1.0),
+        help="Relative step range [start, end] where CFG is active.",
+    )
+    p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
+    p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
+    args = p.parse_args()
+    if not args.prompt:
+        args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]
+    return args
+
+
+def tensor_to_pil(image: torch.Tensor) -> Image.Image:
+    """Convert a normalized torch tensor [-1, 1] to a PIL Image."""
+    if image.ndim == 4:
+        image = image[0]
+    image = image.detach().to("cpu")
+    image = (image / 2 + 0.5).clamp(0, 1)
+    image = (image * 255).to(torch.uint8)
+    image = image.permute(1, 2, 0).contiguous().numpy()
+    return Image.fromarray(image)
+
+
+def _format_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str:
+    """Build the AR-stage prompt string including the image grid header."""
+    return (
+        "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n"
+        f"<|im_start|>user\n{user_prompt}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+        f"<|image start|>{ar_width}*{ar_height}<|image token|>"
+    )
+
+
+def _collect_images(outputs: list) -> list[torch.Tensor]:
+    """Extract all image tensors produced by the final (DiT) stage."""
+    images: list[torch.Tensor] = []
+    for out in outputs:
+        ro_list = getattr(out, "request_output", out)
+        if not isinstance(ro_list, list):
+            ro_list = [ro_list]
+        for ro_item in ro_list:
+            for completion in getattr(ro_item, "outputs", None) or []:
+                mm = getattr(completion, "multimodal_output", None)
+                if not isinstance(mm, dict) or "image" not in mm:
+                    raise RuntimeError(f"Missing image in multimodal output: {mm}")
+                payload = mm["image"]
+                for tensor in payload if isinstance(payload, list) else [payload]:
+                    if not isinstance(tensor, torch.Tensor):
+                        raise TypeError(f"Expected image tensor, got {type(tensor)}")
+                    images.append(tensor)
+    return images
+
+
+def _save_images(images: list[torch.Tensor], out_path: str) -> list[str]:
+    """Save image tensors to disk.
+
+    Single image: written to *out_path* exactly.
+    Multiple images: suffixed as ``<base>_0<ext>``, ``<base>_1<ext>``, …
+    """
+    if not images:
+        raise RuntimeError("No images to save.")
+    base, ext = os.path.splitext(out_path)
+    ext = ext or ".png"
+    paths = []
+    for i, tensor in enumerate(images):
+        path = out_path if len(images) == 1 else f"{base}_{i}{ext}"
+        tensor_to_pil(tensor).save(path)
+        paths.append(path)
+    return paths
+
+
+def main() -> None:
+    args = parse_args()
+    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+
+    if args.height <= 0 or args.width <= 0:
+        raise ValueError(f"Height and width must be positive, got {args.height}x{args.width}")
+    if args.height % _PATCH_SIZE != 0 or args.width % _PATCH_SIZE != 0:
+        raise ValueError(f"Height and width must be multiples of {_PATCH_SIZE}, got {args.height}x{args.width}")
+
+    ar_height = args.height // _PATCH_SIZE
+    ar_width = args.width // _PATCH_SIZE
+    gen_cfg = load_t2i_generation_config(args.model)
+    expected_grid_tokens = ar_height * (ar_width + 1)
+
+    logger.info("Initializing Omni pipeline...")
+    omni = Omni(model=args.model, stage_configs_path=args.stage_config, trust_remote_code=args.trust_remote_code)
+    try:
+        ar_sampling = SamplingParams(
+            temperature=1.0,
+            top_p=1.0,
+            top_k=gen_cfg.top_k,
+            max_tokens=max(1, expected_grid_tokens + 1),  # +1 for hidden state of eoi
+            detokenize=False,
+        )
+        dit_sampling = SamplingParams(
+            temperature=0.0,
+            top_p=1.0,
+            top_k=-1,
+            max_tokens=1,
+            detokenize=False,
+        )
+
+        additional_information = {
+            "omni_task": ["t2i"],
+            "ar_width": [ar_width],
+            "ar_height": [ar_height],
+            "eol_token_id": [gen_cfg.eol_token_id],
+            "visual_token_start_id": [gen_cfg.visual_token_start_id],
+            "visual_token_end_id": [gen_cfg.visual_token_end_id],
+            "image_height": [args.height],
+            "image_width": [args.width],
+            "num_inference_steps": [args.num_inference_steps],
+            "text_guidance_scale": [args.text_guidance_scale],
+            "cfg_range": [args.cfg_range[0], args.cfg_range[1]],
+            "visual_ids": gen_cfg.visual_ids,
+        }
+        inputs = [
+            {
+                "prompt": _format_prompt(p, ar_width, ar_height),
+                "additional_information": dict(additional_information),
+            }
+            for p in args.prompt
+        ]
+
+        logger.info("Starting generation...")
+        # omni.generate() returns a Generator; consume it to run the full pipeline.
+        outputs = list(omni.generate(inputs, [ar_sampling, dit_sampling]))
+
+        logger.info("Post-processing and saving image(s)...")
+        for path in _save_images(_collect_images(outputs), args.out):
+            logger.info(f"Saved: {path}")
+    finally:
+        omni.close()
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/mimo_audio/README.md

@@ -0,0 +1,232 @@
+# MiMo-Audio Offline Inference
+
+This directory contains an offline demo for running MiMo-Audio models with vLLM Omni. It builds task-specific inputs and generates WAV files or text outputs locally.
+
+## Model Overview
+
+MiMo-Audio provides multiple task variants for audio understanding and generation:
+
+- **tts_sft**: Basic text-to-speech generation from text input.
+- **tts_sft_with_instruct**: TTS generation with explicit voice style instructions.
+- **tts_sft_with_audio**: TTS generation with audio reference for voice cloning.
+- **tts_sft_with_natural_instruction**: TTS generation from natural language descriptions embedded in text.
+- **audio_trancribing_sft**: Transcribe audio to text (speech-to-text).
+- **audio_understanding_sft**: Understand and analyze audio content with text queries.
+- **audio_understanding_sft_with_thinking**: Audio understanding with reasoning chain.
+- **spoken_dialogue_sft_multiturn**: Multi-turn spoken dialogue with audio input/output.
+- **speech2text_dialogue_sft_multiturn**: Multi-turn dialogue converting speech to text.
+- **text_dialogue_sft_multiturn**: Multi-turn text-only dialogue.
+
+## Setup
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+
+### Environment Variables
+
+The `MIMO_AUDIO_TOKENIZER_PATH` environment variable is mandatory due to the specialized architecture:
+
+```bash
+export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
+```
+
+## Quick Start
+
+Run a single sample for basic TTS:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type tts_sft
+```
+
+Run batch samples for basic TTS:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type tts_sft \
+  --num-prompts {batch_size}
+```
+
+When enabling multi-batch processing, if the total number of tokens passed to the next stage exceeds the `max_model_len` value in the `mimo_audio.yaml` configuration file, you must also synchronously update the `max_position_embeddings` value in `MiMo-Audio-7B-Instruct/config.json` to match the modified value.
+
+Generated audio files are saved to `output_audio/` by default. `--num-prompts` also can be used to all tasks below.
+
+## Task Usage
+
+### tts_sft (Basic Text-to-Speech)
+
+Generate speech from text input:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type tts_sft \
+  --text "The weather is so nice today."
+```
+
+### tts_sft_with_instruct (TTS with Voice Instructions)
+
+Generate speech with explicit voice style instructions:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type tts_sft_with_instruct \
+  --text "The weather is so nice today." \
+  --instruct "Speak happily in a child's voice"
+```
+
+### tts_sft_with_audio (TTS with Audio Reference)
+
+Generate speech using an audio reference for voice cloning:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type tts_sft_with_audio \
+  --text "The weather is so nice today." \
+  --audio-path "./spoken_dialogue_assistant_turn_1.wav"
+```
+
+### tts_sft_with_natural_instruction (Natural Language TTS)
+
+Generate speech from text containing natural voice descriptions:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type tts_sft_with_natural_instruction \
+  --text "In a panting young male voice, he said: I can't run anymore, wait for me!"
+```
+
+### audio_trancribing_sft (Speech-to-Text)
+
+Transcribe audio to text:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type audio_trancribing_sft \
+  --audio-path "./spoken_dialogue_assistant_turn_1.wav"
+```
+
+### audio_understanding_sft (Audio Understanding)
+
+Understand and analyze audio content with text queries:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type audio_understanding_sft \
+  --text "Summarize the audio." \
+  --audio-path "./spoken_dialogue_assistant_turn_1.wav"
+```
+
+### audio_understanding_sft_with_thinking (Audio Understanding with Reasoning)
+
+Audio understanding with reasoning chain:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type audio_understanding_sft_with_thinking \
+  --text "Summarize the audio." \
+  --audio-path "./spoken_dialogue_assistant_turn_1.wav"
+```
+
+### spoken_dialogue_sft_multiturn (Multi-turn Spoken Dialogue)
+
+Multi-turn dialogue with audio input and output:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type spoken_dialogue_sft_multiturn \
+  --audio-path "./prompt_speech_zh_m.wav"
+```
+
+Note: This task uses hardcoded audio files in the script. The audio files used in examples are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
+
+### speech2text_dialogue_sft_multiturn (Speech-to-Text Dialogue)
+
+Multi-turn dialogue converting speech to text:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type speech2text_dialogue_sft_multiturn
+```
+
+Note: This task uses hardcoded audio files and message lists in the script.
+
+### text_dialogue_sft_multiturn (Text Dialogue)
+
+Multi-turn text-only dialogue:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type text_dialogue_sft_multiturn
+```
+
+Note: This task uses hardcoded message lists in the script.
+
+## Troubleshooting
+
+### Audio dependencies (soundfile, librosa)
+
+This example depends on **soundfile** (read/write WAV) and **librosa** (load audio including MP3). Install the project requirements first:
+
+```bash
+pip install -r requirements/common.txt
+# or at least: pip install soundfile>=0.13.1 librosa>=0.11.0
+```
+
+- **`soundfile` / libsndfile not found**  
+  `soundfile` uses the C library **libsndfile**. On Linux, install the system package before pip:
+  - Debian/Ubuntu: `sudo apt-get install libsndfile1`
+  - For development builds: `sudo apt-get install libsndfile1-dev`
+  - Then: `pip install soundfile`
+
+- **`librosa` fails to load MP3 or reports "No backend available"**  
+  Loading MP3 (e.g. in `spoken_dialogue_sft_multiturn` with `.mp3` files) uses **ffmpeg** as the backend. Install ffmpeg:
+  - Debian/Ubuntu: `sudo apt-get install ffmpeg`
+  - macOS: `brew install ffmpeg`
+
+- **`ImportError: No module named 'soundfile'` or `ModuleNotFoundError: ... librosa`**  
+  Ensure you are in the same Python environment where vLLM Omni and the example dependencies are installed, and that `requirements/common.txt` (or the packages above) are installed.
+
+### Tokenizer path
+
+- **`MIMO_AUDIO_TOKENIZER_PATH` not set or model fails to find tokenizer**  
+  Export the tokenizer path before running:
+  ```bash
+  export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
+  ```
+  See [Environment Variables](#environment-variables) in Setup.
+
+### Other
+
+- If the model or stage config fails to load, check [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for memory and GPU settings.
+- For errors when reading/writing WAV (e.g. unsupported format), ensure input files are standard WAV/MP3 and that `soundfile` is linked to a working libsndfile (see above).
+
+## Notes
+
+- The script uses default model paths and audio files embedded in `end2end.py`. Update them if your local cache path differs.
+- Use `--output-dir` to change the output folder (default: `./output_audio`).
+- Use `--num-prompts` to generate multiple prompts in one run (default: 1).
+- Audio files used in multi-turn dialogue examples are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
+- The script supports various configuration options for initialization timeouts, batch timeouts, and shared memory thresholds. See `--help` for details.

examples/offline_inference/mimo_audio/end2end.py

@@ -0,0 +1,443 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+This example shows how to use vLLM for running offline inference
+with the correct prompt format on MiMo-Audio-Omni.
+"""
+
+import copy
+import json
+import os
+from typing import NamedTuple
+
+import soundfile as sf
+from message_convert import (
+    get_audio_data,
+    get_audio_understanding_sft_prompt,
+    get_s2t_dialogue_sft_multiturn_prompt,
+    get_spoken_dialogue_sft_multiturn_prompt,
+    get_text_dialogue_sft_multiturn_prompt,
+    get_tts_sft_prompt,
+    to_prompt,
+)
+from vllm import SamplingParams
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniTokensPrompt
+
+SEED = 42
+
+MAX_CODE2WAV_TOKENS = 18192  # Maximum tokens supported by code2wav model
+
+
+class QueryResult(NamedTuple):
+    inputs: dict
+    limit_mm_per_prompt: dict[str, int]
+
+
+def get_codes_query_from_json(codes_path: str) -> QueryResult:
+    with open(codes_path, encoding="utf-8") as f:
+        data = json.load(f)
+
+    if isinstance(data, list):
+        code_final = data
+    elif isinstance(data, dict) and "code_final" in data:
+        code_final = data["code_final"]
+    else:
+        raise ValueError(
+            f"Unsupported codes json format in {codes_path}.\n"
+            "Expect a JSON list[int] or {{'code_final': list[int]}}."
+        )
+
+    if not isinstance(code_final, list) or not all(isinstance(x, int) for x in code_final):
+        raise ValueError("code_final must be a list[int].")
+
+    if len(code_final) > MAX_CODE2WAV_TOKENS:
+        print(f"[Warn] code_final len={len(code_final)} > {MAX_CODE2WAV_TOKENS}, truncating.")
+        code_final = code_final[:MAX_CODE2WAV_TOKENS]
+
+    return QueryResult(
+        inputs=OmniTokensPrompt(
+            prompt_token_ids=code_final,
+            multi_modal_data=None,
+            mm_processor_kwargs=None,
+        ),
+        limit_mm_per_prompt={},
+    )
+
+
+def get_tts_sft(
+    text="The weather is so nice today.",
+    instruct=None,
+    read_text_only=True,
+    prompt_speech=None,
+    audio_list=None,
+):
+    res = get_tts_sft_prompt(
+        text,
+        instruct=instruct,
+        read_text_only=read_text_only,
+        prompt_speech=prompt_speech,
+    )
+
+    prompt = to_prompt(res)
+    final_prompt = {
+        "prompt": prompt,
+    }
+    if audio_list is not None:
+        final_prompt.update(
+            {
+                "multi_modal_data": {
+                    "audio": audio_list,
+                },
+            }
+        )
+    return final_prompt
+
+
+def get_audio_understanding_sft(audio_path, text="", thinking=False, use_sostm=False):
+    audio_list = []
+    audio_list.append(get_audio_data(audio_path))
+    res = get_audio_understanding_sft_prompt(
+        input_speech=audio_path, input_text=text, thinking=thinking, use_sostm=use_sostm
+    )
+    prompt = to_prompt(res)
+    final_prompt = {
+        "prompt": prompt,
+        "multi_modal_data": {
+            "audio": audio_list,
+        },
+    }
+    return final_prompt
+
+
+def get_spoken_dialogue_sft_multiturn(message_list, system_prompt=None, ref_audio_path=None, audio_list=None):
+    res = get_spoken_dialogue_sft_multiturn_prompt(
+        message_list, system_prompt=system_prompt, prompt_speech=ref_audio_path
+    )
+    prompt = to_prompt(res)
+    final_prompt = {
+        "prompt": prompt,
+        "multi_modal_data": {
+            "audio": audio_list,
+        },
+    }
+    return final_prompt
+
+
+def get_speech2text_dialogue_sft_multiturn(message_list, thinking=False, audio_list=None):
+    res = get_s2t_dialogue_sft_multiturn_prompt(
+        message_list,
+        thinking=thinking,
+    )
+    prompt = to_prompt(res)
+    final_prompt = {
+        "prompt": prompt,
+        "multi_modal_data": {
+            "audio": audio_list,
+        },
+    }
+    return final_prompt
+
+
+def get_text_dialogue_sft_multiturn(
+    message_list,
+):
+    res = get_text_dialogue_sft_multiturn_prompt(
+        message_list,
+    )
+    prompt = to_prompt(res)
+    final_prompt = {
+        "prompt": prompt,
+    }
+    return final_prompt
+
+
+query_map = {
+    "tts_sft": get_tts_sft,
+    "tts_sft_with_instruct": get_tts_sft,
+    "tts_sft_with_audio": get_tts_sft,
+    "tts_sft_with_natural_instruction": get_tts_sft,
+    "audio_trancribing_sft": get_audio_understanding_sft,
+    "audio_understanding_sft": get_audio_understanding_sft,
+    "audio_understanding_sft_with_thinking": get_audio_understanding_sft,
+    "spoken_dialogue_sft_multiturn": get_spoken_dialogue_sft_multiturn,
+    "speech2text_dialogue_sft_multiturn": get_speech2text_dialogue_sft_multiturn,
+    "text_dialogue_sft_multiturn": get_text_dialogue_sft_multiturn,
+}
+
+
+def main(args):
+    model_name = args.model_name
+
+    # Get paths from args
+    text = getattr(args, "text", None)
+    audio_path = getattr(args, "audio_path", None)
+
+    instruct = getattr(args, "instruct", None)
+
+    # Get the query function and call it with appropriate parameters
+    query_func = query_map[args.query_type]
+
+    omni_llm = Omni(
+        model=model_name,
+        stage_configs_path=args.stage_configs_path,
+        log_stats=args.enable_stats,
+        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
+        init_sleep_seconds=args.init_sleep_seconds,
+        batch_timeout=args.batch_timeout,
+        init_timeout=args.init_timeout,
+        shm_threshold_bytes=args.shm_threshold_bytes,
+    )
+
+    thinker_sampling_params = SamplingParams(
+        temperature=0.0,
+        top_p=1.0,
+        top_k=-1,
+        max_tokens=2048,
+        seed=SEED,
+        logit_bias={},
+        repetition_penalty=1.1,
+    )
+
+    code2wav_sampling_params = SamplingParams(
+        temperature=0.0,
+        top_p=1.0,
+        top_k=-1,
+        max_tokens=4096 * 16,
+        seed=SEED,
+        detokenize=True,
+        repetition_penalty=1.1,
+    )
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        code2wav_sampling_params,
+    ]
+
+    # Build query result based on query type
+    # Notice: The audio files used in this example are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
+    if args.query_type == "tts_sft":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type tts_sft
+        query_result = query_func(text=text, read_text_only=True)
+    elif args.query_type == "tts_sft_with_instruct":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type tts_sft_with_instruct --instruct "Speak happily in a child's voice"
+        query_result = query_func(text=text, instruct=instruct, read_text_only=True)
+    elif args.query_type == "tts_sft_with_audio":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type tts_sft_with_audio --audio_path "./spoken_dialogue_assistant_turn_1.wav"
+        audio_list = [get_audio_data(audio_path)]
+        query_result = query_func(text=text, read_text_only=True, prompt_speech=audio_path, audio_list=audio_list)
+    elif args.query_type == "tts_sft_with_natural_instruction":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type tts_sft_with_natural_instruction --text "In a panting young male voice, he said: I can't run anymore, wait for me!"
+        query_result = query_func(text=text, read_text_only=False)
+    elif args.query_type == "audio_trancribing_sft":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type audio_trancribing_sft --audio_path "./spoken_dialogue_assistant_turn_1.wav"
+        audio_path = "spoken_dialogue_assistant_turn_1.wav"
+        text = "Please transcribe this audio and repeat it once."
+        query_result = query_func(text=text, audio_path=audio_path, use_sostm=True)
+    elif args.query_type == "audio_understanding_sft":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type audio_understanding_sft --text "Summarize the audio." --audio_path "./spoken_dialogue_assistant_turn_1.wav"
+        query_result = query_func(text=text, audio_path=audio_path)
+    elif args.query_type == "audio_understanding_sft_with_thinking":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type audio_understanding_sft_with_thinking --text "Summarize the audio." --audio_path "./spoken_dialogue_assistant_turn_1.wav"
+        query_result = query_func(text=text, audio_path=audio_path, thinking=True)
+    elif args.query_type == "spoken_dialogue_sft_multiturn":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type spoken_dialogue_sft_multiturn  --audio_path "./prompt_speech_zh_m.wav"
+        first_turn_text_response = "我没办法获取实时的天气信息。不过呢，你可以试试几个方法来查看今天的天气。首先，你可以用手机自带的天气功能，比如苹果手机的天气应用，或者直接在系统设置里查看。其次，你也可以用一些专业的天气服务，像是国外的AccuWeather、Weather.com，或者国内的中国天气网、墨迹天气等等。再有就是，你还可以在谷歌或者百度里直接搜索你所在的城市加上天气这两个字。如果你能告诉我你所在的城市，我也可以帮你分析一下历史天气趋势，不过最新的数据还是需要你通过官方渠道去获取哦。"
+        audio_list = []
+        s1_audio_path = "weather_of_today.mp3"
+        s2_audio_path = "spoken_dialogue_assistant_turn_1.wav"
+        s3_audio_path = "beijing.mp3"
+        audio_list.append(get_audio_data(audio_path))
+        audio_list.append(get_audio_data(s1_audio_path))
+        audio_list.append(get_audio_data(s2_audio_path))
+        audio_list.append(get_audio_data(s3_audio_path))
+
+        message_list = [
+            {"role": "user", "content": s1_audio_path},
+            {"role": "assistant", "content": {"text": first_turn_text_response, "audio": s2_audio_path}},
+            {"role": "user", "content": s3_audio_path},
+        ]
+        query_result = query_func(message_list, system_prompt=None, ref_audio_path=audio_path, audio_list=audio_list)
+    elif args.query_type == "speech2text_dialogue_sft_multiturn":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type speech2text_dialogue_sft_multiturn
+        s1_audio_path = "weather_of_today.mp3"
+        s2_audio_path = "beijing.mp3"
+        audio_list = []
+        audio_list.append(get_audio_data(s1_audio_path))
+        audio_list.append(get_audio_data(s2_audio_path))
+        message_list = [
+            {"role": "user", "content": s1_audio_path},
+            {
+                "role": "assistant",
+                "content": "Hello, I can't get real-time weather information. If you can tell me your city, I can help you analyze historical weather trends, but for the latest data, you'll need to get it through official channels.",
+            },
+            {"role": "user", "content": s2_audio_path},
+        ]
+        query_result = query_func(message_list, thinking=True, audio_list=audio_list)
+    elif args.query_type == "text_dialogue_sft_multiturn":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type text_dialogue_sft_multiturn
+        message_list = [
+            {"role": "user", "content": "Could you recommend some tourist attractions in China?"},
+            {"role": "assistant", "content": "Hello, which city would you like to travel to?"},
+            {"role": "user", "content": "Beijing"},
+        ]
+        query_result = query_func(message_list=message_list)
+    else:
+        raise ValueError(f"Invalid query type: {args.query_type}")
+
+    prompts = [copy.deepcopy(query_result) for _ in range(args.num_prompts)]
+
+    print("prompts", prompts)
+    omni_outputs = omni_llm.generate(prompts, sampling_params_list)
+
+    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    if args.query_type is not None:
+        output_dir = os.path.join(output_dir, args.query_type)
+    os.makedirs(output_dir, exist_ok=True)
+
+    for stage_outputs in omni_outputs:
+        if stage_outputs.final_output_type == "text":
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
+                text_output = output.outputs[0].text
+                # Save aligned text file per request
+                prompt_text = output.prompt
+                out_txt = os.path.join(output_dir, f"{request_id}.txt")
+                lines = []
+                lines.append("Prompt:\n")
+                lines.append(str(prompt_text) + "\n")
+                lines.append("vllm_text_output:\n")
+                lines.append(str(text_output).strip() + "\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        print("lines", lines)
+                        f.writelines(lines)
+                except Exception as e:
+                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
+                print(f"Request ID: {request_id}, Text saved to {out_txt}\n")
+        elif stage_outputs.final_output_type == "audio":
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
+                audio_tensor = output.outputs[0].multimodal_output.get("audio")
+
+                if audio_tensor is None:
+                    continue
+
+                output_wav = os.path.join(output_dir, f"{request_id}.wav")
+
+                # Convert to numpy array and ensure correct format
+                audio_numpy = audio_tensor.float().detach().cpu().numpy()
+
+                # Ensure audio is 1D (flatten if needed)
+                if audio_numpy.ndim > 1:
+                    audio_numpy = audio_numpy.flatten()
+
+                # Save audio file with explicit WAV format
+                sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
+                print(f"Request ID: {request_id}, Audio saved to {output_wav}")
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser.add_argument(
+        "--model-name",
+        "-m",
+        type=str,
+        default="XiaomiMiMo/MiMo-Audio-7B-Instruct",
+        help="Backbone LLM path.",
+    )
+    parser.add_argument(
+        "--text",
+        "-t",
+        type=str,
+        default="The weather is so nice today.",
+        help="input text",
+    )
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="tts_sft",
+        choices=query_map.keys(),
+        help="Query type.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file. If not provided, uses default audio asset.",
+    )
+    parser.add_argument(
+        "--instruct",
+        type=str,
+        default=None,
+        help="instruct",
+    )
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        help="Enable writing detailed statistics (default: disabled)",
+    )
+    parser.add_argument(
+        "--init-sleep-seconds",
+        type=int,
+        default=20,
+        help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
+    )
+    parser.add_argument(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Timeout for batching in seconds (default: 5)",
+    )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=5000,
+        help="Timeout for initializing stages in seconds (default: 5000)",
+    )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Threshold for using shared memory in bytes (default: 65536)",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="./output_audio",
+        help="Output audio wav directory.",
+    )
+    parser.add_argument(
+        "--output-wav",
+        default="output_audio",
+        help="[Deprecated] Output wav directory (use --output-dir).",
+    )
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Number of prompts to generate.",
+    )
+
+    parser.add_argument(
+        "--sampling-rate",
+        type=int,
+        default=24000,
+        help="Sampling rate for audio.",
+    )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default="../../../model_executor/stage_configs/mimo_audio.yaml",
+        help="Path to a stage configs file.",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)