cherry-pick pr3067 (#3087 )

Signed-off-by: lyj-jjj <liuyingjun5@huawei.com>
[release v0.18.0.post1] Fix default sampling params for /v1/videos (#3021 )
2026-04-24 09:56:49 +08:00 · 2026-04-23 11:29:54 +08:00 · 2026-04-22 11:09:27 +08:00 · 2026-04-21 20:36:08 +08:00 · 2026-04-20 19:22:11 +08:00 · 2026-04-20 14:51:38 +08:00
116 changed files with 4778 additions and 1367 deletions
@@ -1,92 +0,0 @@
-steps:
-  - label: ":buildkit: Build and Push NPU Test Image"
-    key: image-build
-    agents:
-      queue: "ascend-a2b3"
-      resource_class: "npu-2"
-    plugins:
-      - kubernetes:
-          metadata:
-            annotations:
-              vault.hashicorp.com/agent-init-first: "true"
-              vault.hashicorp.com/agent-inject: "true"
-              vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
-              vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
-              vault.hashicorp.com/agent-inject-perms-config.json: "0400"
-              vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
-              vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
-              vault.hashicorp.com/agent-pre-populate-only: "true"
-              vault.hashicorp.com/agent-run-as-group: "1000"
-              vault.hashicorp.com/agent-run-as-user: "1000"
-              vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
-              vault.hashicorp.com/role: ascend-gha-runners
-              vault.hashicorp.com/secret-volume-path: /home/user/.docker/
-              vault.hashicorp.com/tls-skip-verify: "true"
-          podSpecPatch:
-            volumes:
-            - name: token-vol
-              projected:
-                defaultMode: 420
-                sources:
-                - serviceAccountToken:
-                    audience: api
-                    expirationSeconds: 600
-                    path: token
-    env:
-      VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-      IMAGE_NAME: "vllm-omni-ci-npu"
-      IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
-      BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
-    command: |
-      set -ex
-
-      echo "--- Building and pushing NPU Test Image"
-      echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
-      echo "buildkitd address: $${BUILDKITD_ADDR}"
-
-      git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf "https://github.com/"
-
-      if ! command -v buildctl &> /dev/null; then
-        echo "Installing buildctl..."
-        mkdir -p /tmp/buildkit
-        BUILDKIT_VERSION="v0.29.0"
-        wget -q "https://gh-proxy.test.osinfra.cn/https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
-        tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
-        cp /tmp/buildkit/bin/buildctl /usr/local/bin/
-      fi
-
-      export DOCKER_CONFIG=/home/user/.docker
-      buildctl \
-        --addr="$${BUILDKITD_ADDR}" \
-        --tlscacert=/home/user/.docker/ca.pem \
-        --tlscert=/home/user/.docker/cert.pem \
-        --tlskey=/home/user/.docker/key.pem \
-        build \
-        --frontend dockerfile.v0 \
-        --local context=. \
-        --local dockerfile=./docker \
-        --opt filename=Dockerfile.npu \
-        --secret id=dockerconfig,src=/home/user/.docker/config.json \
-        --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
-        --progress=plain
-
-      echo "--- Image pushed successfully"
-      echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
-
-  - label: "Upload Ready Pipeline"
-    depends_on: image-build
-    key: upload-ready-pipeline
-    if: build.branch != "main" && build.pull_request.labels includes "ready"
-    agents:
-      queue: "ascend-a2b3"
-      resource_class: "npu-2"
-    command: |
-      buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml
-      
@@ -64,22 +64,15 @@ while true; do
 done

 echo "--- Pulling container"
-## Temporary change to use AMD Docker Hub to store the vllm-ci image
-# to bypass the rate limit issue with ECR Public Gallery.
-# TODO: @tjtanaa point back to ECR Public Gallery
-# once the amd agents are configured to use ECR Public Gallery.
-# image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
-image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}-rocm-omni"
+image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
 container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"

-# TODO: @tjtanaa uncomment this once the amd agents are configured to use ECR Public Gallery.
-# # Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
-# sudo apt-get update && sudo apt-get install -y awscli
-##  Use safe docker login helper to prevent race conditions
-# source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
-# safe_docker_login_ecr_public
-
-## Pull the container from AMD Docker Hub
+# Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
+sudo apt-get update && sudo apt-get install -y awscli
+# Use safe docker login helper to prevent race conditions
+source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
+safe_docker_login_ecr_public
+# Pull the container from ECR Public Gallery

 docker pull "${image_name}"

@@ -133,6 +126,8 @@ if [[ $commands == *"--shard-id="* ]]; then
        -e VLLM_ROCM_USE_AITER=0 \
        -e HIP_VISIBLE_DEVICES="${GPU}" \
        -e HF_TOKEN \
+        -e AWS_ACCESS_KEY_ID \
+        -e AWS_SECRET_ACCESS_KEY \
        -v "${HF_CACHE}:${HF_MOUNT}" \
        -e "HF_HOME=${HF_MOUNT}" \
        -e "PYTHONPATH=${MYPYTHONPATH}" \
@@ -165,6 +160,8 @@ else
          -e MIOPEN_DEBUG_CONV_GEMM=0 \
          -e VLLM_ROCM_USE_AITER=0 \
          -e HF_TOKEN \
+          -e AWS_ACCESS_KEY_ID \
+          -e AWS_SECRET_ACCESS_KEY \
          -v "${HF_CACHE}:${HF_MOUNT}" \
          -e "HF_HOME=${HF_MOUNT}" \
          -e "PYTHONPATH=${MYPYTHONPATH}" \
@@ -1,14 +0,0 @@
-steps:
-  - label: "Simple Unit Test"
-    depends_on: upload-ready-pipeline
-    agents:
-      queue: "ascend-a2b3"
-      resource_class: "npu-2"
-    image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}"
-    plugins:
-      - kubernetes:
-          podSpecPatch:
-            imagePullSecrets:
-            - name: swr-secret
-    commands:
-      - "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"
@@ -3,7 +3,7 @@
   Last synced: 2025-12-15
   Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests
 #}
-{% set docker_image_amd = "rocm/vllm-ci:$BUILDKITE_COMMIT-rocm-omni" %}
+{% set docker_image_amd = "public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT-rocm-omni" %}
 {% set default_working_dir = "/app/vllm-omni" %}

  - group: "AMD Tests"
@@ -13,6 +13,7 @@
        depends_on: ~
        soft_fail: false
        commands:
+          - "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target test --build-arg ARG_PYTORCH_ROCM_ARCH=gfx942 --progress plain ."
          - "docker push {{ docker_image_amd }}"
        key: "amd-build"
@@ -21,15 +22,13 @@
        retry:
          automatic:
            - exit_status: -1  # Agent was lost
-              limit: 2
+              limit: 1
            - exit_status: -10  # Agent was lost
-              limit: 2
-            - exit_status: 128  # Git connectivity issues
-              limit: 2
+              limit: 1
            - exit_status: 1  # Machine occasionally fail
              limit: 1
        agents:
-          queue: amd-cpu
+          queue: cpu_queue_premerge

    {% for step in steps %}
    {% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
@@ -5,34 +5,26 @@ FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 WORKDIR /vllm-workspace/vllm
 RUN git fetch origin --tags && git checkout v0.18.0

-RUN git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf "https://github.com/"
-
 WORKDIR /vllm-workspace/vllm-ascend
 RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
 # Install vllm-ascend
-RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
    source /usr/local/Ascend/nnal/atb/set_env.sh && \
    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -U pip && \
-    python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ \
-      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
-      --extra-index https://download.pytorch.org/whl/cpu/ \
-      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
-      --default-timeout=60
+    python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
    python3 -m pip cache purge

 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 COPY . .

-RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
    source /usr/local/Ascend/nnal/atb/set_env.sh && \
    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -U pip && \
-    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \
-      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
-      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
-      --default-timeout=60
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation

 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn

@@ -5,34 +5,26 @@ FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 WORKDIR /vllm-workspace/vllm
 RUN git fetch origin --tags && git checkout v0.18.0

-RUN git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf "https://github.com/"
-
 WORKDIR /vllm-workspace/vllm-ascend
 RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
 # Install vllm-ascend
-RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
    source /usr/local/Ascend/nnal/atb/set_env.sh && \
    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -U pip && \
-    python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ \
-      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
-      --extra-index https://download.pytorch.org/whl/cpu/ \
-      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
-      --default-timeout=60
+    python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
    python3 -m pip cache purge

 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 COPY . .

-RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
    source /usr/local/Ascend/nnal/atb/set_env.sh && \
    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -U pip && \
-    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \
-      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
-      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
-      --default-timeout=60
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation

 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn

@@ -62,6 +62,7 @@ nav:
        - FP8: user_guide/diffusion/quantization/fp8.md
        - Int8: user_guide/diffusion/quantization/int8.md
        - GGUF: user_guide/diffusion/quantization/gguf.md
+      - Frame Interpolation: user_guide/diffusion/frame_interpolation.md
      - Step Execution: user_guide/diffusion/step_execution.md
      - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
      - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
@@ -567,15 +567,15 @@ L4 level testing is a comprehensive quality audit before a version release. It e
            {
                "dataset_name": "random",
                "num_prompts": [10, 20],
-                "max_concurrency": [1, 4],
+                "request_rate": [0.5, 1],
                "random_input_len": 2500,
                "random_output_len": 900,
                "ignore_eos": true,
                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                "baseline": {
-                    "mean_ttft_ms": [500, 800],
-                    "mean_audio_ttfp_ms": [2000, 3500],
-                    "mean_audio_rtf": [0.25, 0.35]
+                    "mean_ttft_ms": 100000,
+                    "mean_audio_ttfp_ms": 100000,
+                    "mean_audio_rtf": 100000
                }
            }
        ]
@@ -634,15 +634,15 @@ L4 level testing is a comprehensive quality audit before a version release. It e

    1.  Change the ---xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
    2.  For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
-    3.  Optionally add a `baseline` object (see **Baseline thresholds** below). If you omit `baseline` or leave it empty, the performance test still runs but does not assert metric thresholds from this field.
+    3.  Add the baseline parameter to specify the required validation values, ensuring the validation metric names match those in the result.json generated by the benchmark.
    4.  The qps and concurrency modes are mutually exclusive. For detailed explanations, see the table below:

    | Parameter       | Type        | Required | Example/Values  | Description                                                                                                                                                                                                                                                          |
    | --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
    | num_prompts     | int / array | Yes      | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
    | request_rate    | int / array | No       | 1, [1, 2, 3]    | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
-    | max_concurrency | int / array | No       | 1, [1, 2, 3]    | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode).                                                                                                                                                                                                             |
-    | baseline        | object      | No       | see above       | Optional per-metric thresholds; keys must match benchmark output fields. Scalar, list (per sweep step), or object (keyed by concurrency or QPS string).                                                                                                               |
+    | max_concurrency | int / array | No       | 1, [1, 2, 3]    | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
+    </details>

    -   -   ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).

@@ -1,106 +1,10 @@
 # CI Failures

-## Overview of CI Checks
+What should I do when a CI job fails on my PR, but I don't think my PR caused the failure?

-When you open a PR against vLLM-Omni, several CI checks run automatically:

-| Check | Platform | What it does |
-| ----- | -------- | ------------ |
-| **pre-commit** | GitHub Actions | Runs linting (Ruff), formatting, spell-checking (typos), and YAML validation. |
-| **Build Wheel** | GitHub Actions | Builds Python wheels for Python 3.11 and 3.12 on Ubuntu. Skipped for docs-only or Markdown-only changes (controlled by `paths-ignore` in the workflow). |
-| **DCO** | GitHub | Verifies every commit has a `Signed-off-by` line. |
-| **docs/readthedocs.org:vllm-omni** | Read the Docs | Builds the MkDocs documentation site. |
-| **buildkite/vllm-omni** | Buildkite | Runs GPU-based tests on NVIDIA CUDA hardware (L4, H100). |
-| **buildkite/vllm-omni-amd** | Buildkite | Runs GPU-based tests on AMD ROCm hardware (MI325). |
-| **buildkite/vllm-omni-intel** | Buildkite | Runs GPU-based tests on Intel XPU hardware (Intel Arc BMG). |
+## Common Case 1:

-## Step 1: Identify the Failing Check
+🚨 Error: The command was interrupted by a signal: signal: terminated

-Click the **Details** link next to the failing check on your PR to open the build log. The most common failures fall into these categories:
-
-### pre-commit failures
-
-These are typically formatting or linting issues introduced by your PR. Fix them locally:
-
-```bash
-uv pip install pre-commit
-pre-commit run --all-files
-```
-
-Then commit the fixes and push.
-
-### DCO failures
-
-Every commit must include a `Signed-off-by` line. If you forgot, amend your commits:
-
-```bash
-git commit --amend -s
-git push --force-with-lease
-```
-
-For multiple commits, use an interactive rebase to add the sign-off to each one.
-
-### Read the Docs failures
-
-The documentation build uses MkDocs with `fail_on_warning: true`, so even a minor warning (not just errors) will cause the build to fail. To reproduce locally:
-
-```bash
-uv pip install -e ".[docs]"
-mkdocs build --strict
-```
-
-Common causes include broken cross-references, invalid admonition syntax, or missing files referenced by `--8<--` includes.
-
-### Buildkite failures
-
-Buildkite runs GPU tests in Docker containers. These are the most complex checks and can fail for reasons unrelated to your PR (infrastructure issues, flaky tests, etc.). See the sections below for how to investigate.
-
-## Step 2: Check if the Failure Is a Known Issue
-
-Before spending time debugging, check whether the failure already exists on the `main` branch:
-
-1. **Look at the Buildkite build log** — the test name and error message are usually enough to identify the issue.
-2. **Check recent CI runs on `main`** — if the same test is failing there, the failure is not caused by your PR.
-3. **Search existing issues** — look for open issues in the [vllm-omni issue tracker](https://github.com/vllm-project/vllm-omni/issues) with the test name or error message.
-
-If the failure is already tracked, leave a comment on your PR noting that the failure is pre-existing and link the issue.
-
-## Step 3: Investigate the Failure
-
-If the failure appears to be new, investigate whether your changes caused it.
-
-### Reading Buildkite Logs
-
-1. Click **Details** next to the Buildkite check on your PR.
-2. Find the failing step in the pipeline (e.g., "Diffusion Model Test", "Simple Unit Test").
-3. Expand the step to see the full test output with the traceback.
-
-### Running Tests Locally
-
-For instructions on running tests locally (including specific test files, functions, and markers), see the [Running Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_guide/#running-tests) section in the Test Guide.
-
-## Step 4: Raise an Issue or Fix It
-
-### If the failure is pre-existing (not caused by your PR)
-
-1. **Raise a new issue** if one doesn't already exist, using the title format:
-   `[CI Failure]: [job-name] - [test-path]`
-2. Include the error message, relevant log excerpts, and the commit hash where the failure occurs (e.g., "Still failing on main as of commit `abc1234`").
-3. Leave a comment on your PR linking to the issue and noting that the failure is unrelated to your changes.
-
-### If the failure is caused by your PR
-
-1. Fix the issue in your branch and push the update.
-2. If the fix is non-trivial, consider adding a test to prevent regression.
-
-## Common Failure Patterns
-
-| Symptom | Likely Cause | Fix |
-| ------- | ------------ | --- |
-| `ruff` or formatting errors | Code style violation | Run `pre-commit run --all-files` |
-| `Signed-off-by` missing | DCO check | Amend commits with `git commit --amend -s` |
-| MkDocs build warning | Broken docs reference | Run `mkdocs build --strict` locally |
-| `OOM` or `CUDA out of memory` | Test exceeds GPU memory | Check if your changes increased memory usage; use `--vae_use_slicing` / `--vae_use_tiling` for diffusion tests |
-| Import errors | Missing or changed dependency | Check `pyproject.toml` and make sure dependencies are correct |
-| Timeout (step exceeded N minutes) | Test is too slow or hangs | Profile the test; check for infinite loops or deadlocks |
-| `Agent lost` in Buildkite | Infrastructure issue (not your fault) | Re-trigger the build; comment on your PR |
+Reason: The test is terminated due to exceed the time limits. Sometimes rebuild the test will pass.
@@ -1,14 +1,15 @@
 # VAE Patch Parallelism

 This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
-We use **Qwen-Image** as the reference implementation.
+We use **Qwen-Image** as the reference implementation for decode parallel, and **Wan2.2** for encode parallel.

 ---

 ## Table of Contents

 - [Overview](#overview)
- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Step-by-Step Implementation (Decode)](#step-by-step-implementation-decode)
+- [Encode Parallel Implementation](#encode-parallel-implementation)
 - [Testing](#testing)
 - [Reference Implementations](#reference-implementations)
 - [Summary](#summary)
@@ -19,13 +20,13 @@ We use **Qwen-Image** as the reference implementation.

 ### What is Vae Patch parallel?

-**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is:
+**VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is:

 + Split into multiple spatial tiles

 + Distributed across multiple ranks

-+ Decoded in parallel
+ Encoded/Decoded in parallel

 + Merged to reconstruct the final output

@@ -35,10 +36,17 @@ This approach:

 + Reduces peak memory usage per device

-+ Accelerates decoding latency
+ Accelerates encoding/decoding latency
+
+### When to Use Encode vs Decode Parallel
+
+| Operation | Use Case | Example |
+|-----------|----------|---------|
+| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
+| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |

 ### Architecture
-We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding.
+We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding.

 The executor is model-agnostic and accepts three function parameters:

@@ -84,7 +92,7 @@ Therefore:

 + Merge must perform blending to avoid seams

-## Step-by-Step Implementation
+## Step-by-Step Implementation (Decode)

 ### Step 1: Implement DistributedAutoencoderKLQwenImage
 `QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
@@ -205,14 +213,14 @@ def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid
 We need to override tiled_decode, the main logic is:
 + check distributed is enabled
 + select split/exec/merge
-+ Invoke self.distributed_decoder.execute to decode
+ Invoke self.distributed_executor.execute to decode
 ```
 def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
    if not self.is_distributed_enabled():
        return super().tiled_decode(z, return_dict=return_dict)

    logger.info("Decode run with distributed executor")
-    result = self.distributed_decoder.execute(
+    result = self.distributed_executor.execute(
        z,
        DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
        broadcast_result=True,
@@ -243,6 +251,166 @@ class YourModelPipeline(nn.Module):
 +       ).to(self.device)
 ```

+## Encode Parallel Implementation
+
+For models that require VAE encoding (e.g., Image-to-Video), you can also parallelize the encode operation. We use **Wan2.2** as the reference implementation.
+
+### Step 1: Implement encode_tile_split
+
+Similar to decode, split the input tensor into tiles. Key considerations:
+
+ **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly
+ **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x)
+
+```python
+def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+    _, _, num_frames, height, width = x.shape
+    encode_spatial_compression_ratio = self.spatial_compression_ratio
+
+    # Scale tile parameters for patchified coordinate system
+    tile_sample_min_height = self.tile_sample_min_height
+    tile_sample_min_width = self.tile_sample_min_width
+    tile_sample_stride_height = self.tile_sample_stride_height
+    tile_sample_stride_width = self.tile_sample_stride_width
+
+    if self.config.patch_size is not None:
+        # When input is patchified, scale tile parameters accordingly
+        encode_spatial_compression_ratio = self.spatial_compression_ratio // self.config.patch_size
+        tile_sample_min_height = tile_sample_min_height // self.config.patch_size
+        tile_sample_min_width = tile_sample_min_width // self.config.patch_size
+        tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
+        tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
+
+    latent_height = height // encode_spatial_compression_ratio
+    latent_width = width // encode_spatial_compression_ratio
+
+    tile_latent_min_height = tile_sample_min_height // encode_spatial_compression_ratio
+    tile_latent_min_width = tile_sample_min_width // encode_spatial_compression_ratio
+    tile_latent_stride_height = tile_sample_stride_height // encode_spatial_compression_ratio
+    tile_latent_stride_width = tile_sample_stride_width // encode_spatial_compression_ratio
+
+    blend_height = tile_latent_min_height - tile_latent_stride_height
+    blend_width = tile_latent_min_width - tile_latent_stride_width
+
+    tiletask_list = []
+    # Use temporal compression ratio from config instead of hardcoding
+    temporal_compression = self.config.scale_factor_temporal
+
+    for i in range(0, height, tile_sample_stride_height):
+        for j in range(0, width, tile_sample_stride_width):
+            time_list = []
+            frame_range = 1 + (num_frames - 1) // temporal_compression
+            for k in range(frame_range):
+                if k == 0:
+                    tile = x[:, :, :1, i : i + tile_sample_min_height, j : j + tile_sample_min_width]
+                else:
+                    tile = x[
+                        :, :,
+                        1 + temporal_compression * (k - 1) : 1 + temporal_compression * k,
+                        i : i + tile_sample_min_height,
+                        j : j + tile_sample_min_width,
+                    ]
+                time_list.append(tile)
+            tiletask_list.append(
+                TileTask(len(tiletask_list), (i // tile_sample_stride_height, j // tile_sample_stride_width),
+                         time_list, workload=time_list[0].shape[3] * time_list[0].shape[4])
+            )
+
+    grid_spec = GridSpec(
+        split_dims=(3, 4),
+        grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+        tile_spec={
+            "latent_height": latent_height, "latent_width": latent_width,
+            "blend_height": blend_height, "blend_width": blend_width,
+            "tile_latent_stride_height": tile_latent_stride_height,
+            "tile_latent_stride_width": tile_latent_stride_width,
+        },
+        output_dtype=self.dtype,
+    )
+    return tiletask_list, grid_spec
+```
+
+### Step 2: Implement encode_tile_exec
+
+```python
+def encode_tile_exec(self, task: TileTask) -> torch.Tensor:
+    """Encode a single sample tile into latent space."""
+    self.clear_cache()
+    time = []
+    for k, tile in enumerate(task.tensor):
+        self._enc_conv_idx = [0]
+        encoded = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+        encoded = self.quant_conv(encoded)
+        time.append(encoded)
+    result = torch.cat(time, dim=2)
+    self.clear_cache()
+    return result
+```
+
+### Step 3: Implement encode_tile_merge
+
+```python
+def encode_tile_merge(
+    self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec
+) -> torch.Tensor:
+    """Merge encoded tiles into a full latent tensor."""
+    grid_h, grid_w = grid_spec.grid_shape
+    result_rows = []
+    for i in range(grid_h):
+        result_row = []
+        for j in range(grid_w):
+            tile = coord_tensor_map[(i, j)]
+            if i > 0:
+                tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+            if j > 0:
+                tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+            result_row.append(tile[:, :, :,
+                : grid_spec.tile_spec["tile_latent_stride_height"],
+                : grid_spec.tile_spec["tile_latent_stride_width"]])
+        result_rows.append(torch.cat(result_row, dim=-1))
+
+    enc = torch.cat(result_rows, dim=3)[
+        :, :, :, : grid_spec.tile_spec["latent_height"], : grid_spec.tile_spec["latent_width"]
+    ]
+    return enc
+```
+
+### Step 4: Override tiled_encode method
+
+Override `tiled_encode` instead of `encode`. The parent's `_encode()` handles patchify before calling `tiled_encode()`, so input `x` is already patchified.
+
+```python
+def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+    """
+    Encode using distributed VAE executor.
+
+    Note: x is already patchified by parent's _encode() before calling this method.
+    """
+    if not self.is_distributed_enabled():
+        return super().tiled_encode(x)
+
+    self.clear_cache()
+    result = self.distributed_executor.execute(
+        x,
+        DistributedOperator(
+            split=self.encode_tile_split,
+            exec=self.encode_tile_exec,
+            merge=self.encode_tile_merge,
+        ),
+        broadcast_result=True,  # Latents needed by all ranks for diffusion
+    )
+    self.clear_cache()
+    return result
+```
+
+**Key differences from decode parallel:**
+
+| Aspect | Decode Parallel | Encode Parallel |
+|--------|-----------------|-----------------|
+| `broadcast_result` | Often `False` (only rank 0 needs output) | `True` (all ranks need latents for diffusion) |
+| Patchify | Applied in merge (unpatchify) | Handled by parent `_encode()` before `tiled_encode()` |
+| Temporal chunking | Frame-by-frame | Chunk-based (e.g., 1 + 4n frames) |
+
 ## Testing
 Verify numerical consistency between:
 + vae_patch_parallel_size = 1
@@ -272,18 +440,20 @@ When vae_patch_parallel_size is larger than the DiT world size, it will automati

 Complete examples in the codebase:

-| Model | Path | Notes |
-|-------|------|-------|
-| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
-| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
-| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage |
+| Model | Path | Decode Parallel | Encode Parallel |
+|-------|------|-----------------|-----------------|
+| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | ✅ | ❌ |
+| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | ✅ | ✅ |
+| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | ✅ | ❌ |

 ---

 ## Summary

-Adding Vae Patch Parallel support to diffusion model:
+Adding VAE Patch Parallel support to diffusion model:

-1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
-2. **Change vae model in pipeline to Distributed Vae**
-3. **Test** - Verify with `tensor_parallel_size=N` quality
+1. **Implement Distributed VAE** - Inherit from base VAE class and `DistributedVaeMixin`
+2. **Decode Parallel** - Refactor `tiled_decode` into `tile_split`/`tile_exec`/`tile_merge`
+3. **Encode Parallel** (optional) - Implement `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models
+4. **Change VAE model in pipeline** - Use the distributed version
+5. **Test** - Verify numerical consistency with `vae_patch_parallel_size=1` vs `N`
@@ -29,11 +29,6 @@ uv pip install vllm --torch-backend=auto
 uv pip install vllm-omni
 ```

-To run Gradio demos, also install the optional extras:
-```bash
-uv pip install 'vllm-omni[demo]'
-```
-
 # --8<-- [end:pre-built-wheels]

 # --8<-- [start:build-wheel-from-source]
@@ -55,11 +50,6 @@ cd vllm-omni
 uv pip install -e .
 ```

-To run Gradio demos, install with optional extras:
-```bash
-uv pip install -e '.[demo]'
-```
-
 <details><summary>(Optional) Installation of vLLM from source</summary>
 If you want to check, modify or debug with source code of vLLM, install the library from source with the following instructions:

@@ -0,0 +1,92 @@
+# Frame Interpolation
+
+## Overview
+
+vLLM-Omni supports post-generation frame interpolation for supported video
+diffusion pipelines. This feature inserts synthesized intermediate frames
+between adjacent generated frames to improve temporal smoothness without
+rerunning the diffusion denoising loop.
+
+Frame interpolation runs in the diffusion worker post-processing path instead
+of the API server encoding path. This allows the interpolation step to reuse
+the worker's current accelerator device and keeps the FastAPI event loop free
+from heavy synchronous PyTorch work.
+
+For an input video with `N` generated frames and interpolation exponent `exp`,
+the output frame count is:
+
+```text
+(N - 1) * 2**exp + 1
+```
+
+The output FPS is multiplied by `2**exp` so the clip duration remains close to
+the original generated video.
+
+## Supported Pipelines
+
+Frame interpolation is currently supported for:
+
+- `WanPipeline` (Wan2.2 text-to-video)
+- `WanImageToVideoPipeline`
+- `Wan22TI2VPipeline`
+
+## Request Parameters
+
+The video APIs `/v1/videos` and `/v1/videos/sync` accept:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `enable_frame_interpolation` | bool | `false` | Enable post-generation frame interpolation |
+| `frame_interpolation_exp` | int | `1` | Interpolation exponent. `1=2x`, `2=4x`, etc. |
+| `frame_interpolation_scale` | float | `1.0` | RIFE inference scale |
+| `frame_interpolation_model_path` | str | `None` | Local directory or Hugging Face repo ID containing `flownet.pkl` |
+
+## Execution Flow
+
+For supported Wan2.2 pipelines, the execution order is:
+
+1. Diffusion worker finishes denoising and decodes the raw video tensor.
+2. Worker-side model-specific post-processing runs.
+3. If frame interpolation is enabled, RIFE interpolates the decoded video
+   tensor on the worker side and records a FPS multiplier in `custom_output`.
+4. The API server receives the already-interpolated video and only performs
+   MP4 export.
+
+This design keeps interpolation close to the generated tensor and avoids
+introducing another heavyweight GPU context in the API server process.
+
+## Example
+
+Start the server:
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091
+```
+
+Run a sync request with interpolation enabled:
+
+```bash
+curl -X POST http://localhost:8091/v1/videos/sync \
+  -F "prompt=A dog running through a park" \
+  -F "num_frames=81" \
+  -F "width=832" \
+  -F "height=480" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=1.0" \
+  -F "guidance_scale_2=1.0" \
+  -F "enable_frame_interpolation=true" \
+  -F "frame_interpolation_exp=1" \
+  -F "frame_interpolation_scale=1.0" \
+  -F "seed=42" \
+  -o sync_t2v_interpolated.mp4
+```
+
+## Notes
+
+- This is a post-processing feature. It does not modify the diffusion denoising
+  schedule.
+- Higher interpolation exponents increase post-processing time and memory usage.
+- If the interpolation model weights are not available locally,
+  `frame_interpolation_model_path` may point to a Hugging Face repo containing
+  `flownet.pkl`.
@@ -65,7 +65,6 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
 | Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
 | HunyuanImage-3 | `tencent/HunyuanImage3` | All layers | None |
 | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None |
-| Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | All layers | None |

 ## Combining with Other Features

@@ -125,7 +125,6 @@ th {
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |
-| `HunyuanImage3Pipeline` | HunyuanImage3 | `tencent/HunyuanImage-3.0-Instruct` |

 ### VideoGen

@@ -256,15 +256,6 @@ lora_adapter/

 ## Web UI Demo

-!!! note "Gradio is an optional dependency"
-    The Gradio demo requires the `[demo]` extras. Install them first:
-
-    ```bash
-    pip install 'vllm-omni[demo]'
-    ```
-
-    Or, if installing from source: `pip install -e '.[demo]'`
-
 Launch the Gradio demo:

 ```bash
@@ -11,15 +11,6 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serv

 ## Gradio Demo

-!!! note "Gradio is an optional dependency"
-    The Gradio demo requires the `[demo]` extras. Install them first:
-
-    ```bash
-    pip install 'vllm-omni[demo]'
-    ```
-
-    Or, if installing from source: `pip install -e '.[demo]'`
-
 An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support.

 ```bash
@@ -122,15 +122,6 @@ python openai_chat_client.py --input input1.png input2.png --prompt "Combine the

 ### Method 4: Using Gradio Demo

-!!! note "Gradio is an optional dependency"
-    The Gradio demo requires the `[demo]` extras. Install them first:
-
-    ```bash
-    pip install 'vllm-omni[demo]'
-    ```
-
-    Or, if installing from source: `pip install -e '.[demo]'`
-
 ```bash
 python gradio_demo.py
 # Visit http://localhost:7861
@@ -149,15 +149,6 @@ python openai_chat_completion_client_for_multimodal_generation.py \

 ## Run Local Web UI Demo

-!!! note "Gradio is an optional dependency"
-    The Gradio demo requires the `[demo]` extras. Install them first:
-
-    ```bash
-    pip install 'vllm-omni[demo]'
-    ```
-
-    Or, if installing from source: `pip install -e '.[demo]'`
-
 This Web UI demo allows users to interact with the model through a web browser.

 ### Running Gradio Demo
@@ -154,7 +154,7 @@ print(response.choices[1].message.audio)    # Audio response
 ```

 ## Streaming Output
-If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. We support both text streaming output and audio streaming output. Other modalities can output normally.
+If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py \
    --query-type use_image \
@@ -163,15 +163,6 @@ python openai_chat_completion_client_for_multimodal_generation.py \

 ## Run Local Web UI Demo

-!!! note "Gradio is an optional dependency"
-    The Gradio demo requires the `[demo]` extras. Install them first:
-
-    ```bash
-    pip install 'vllm-omni[demo]'
-    ```
-
-    Or, if installing from source: `pip install -e '.[demo]'`
-
 This Web UI demo allows users to interact with the model through a web browser.

 ### Running Gradio Demo
@@ -19,15 +19,6 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/

 ## Gradio Demo

-!!! note "Gradio is an optional dependency"
-    The Gradio demo requires the `[demo]` extras. Install them first:
-
-    ```bash
-    pip install 'vllm-omni[demo]'
-    ```
-
-    Or, if installing from source: `pip install -e '.[demo]'`
-
 Two interactive Gradio demos are available, both supporting all 3 task types:

 | Demo     | File                     | Transport    | Streaming Quality                                  |
@@ -105,15 +105,6 @@ python openai_chat_client.py --prompt "A beautiful landscape painting" --output

 ### Method 4: Using Gradio Demo

-!!! note "Gradio is an optional dependency"
-    The Gradio demo requires the `[demo]` extras. Install them first:
-
-    ```bash
-    pip install 'vllm-omni[demo]'
-    ```
-
-    Or, if installing from source: `pip install -e '.[demo]'`
-
 ```bash
 python gradio_demo.py
 # Visit http://localhost:7860
@@ -123,6 +123,9 @@ curl -X POST http://localhost:8091/v1/videos \
  -F "guidance_scale_2=4.0" \
  -F "boundary_ratio=0.875" \
  -F "flow_shift=5.0" \
+  -F "enable_frame_interpolation=true" \
+  -F "frame_interpolation_exp=1" \
+  -F "frame_interpolation_scale=1.0" \
  -F "seed=42"
 ```

@@ -145,6 +148,35 @@ curl -X POST http://localhost:8091/v1/videos \
 | `flow_shift`          | float  | None    | Scheduler flow shift (Wan2.2)                    |
 | `seed`                | int    | None    | Random seed (reproducible)                       |
 | `lora`                | object | None    | LoRA configuration                               |
+| `enable_frame_interpolation` | bool | false | Enable RIFE frame interpolation before MP4 encoding |
+| `frame_interpolation_exp` | int | 1 | Interpolation exponent; 1=2x temporal resolution, 2=4x |
+| `frame_interpolation_scale` | float | 1.0 | RIFE inference scale; use 0.5 for high-resolution inputs |
+| `frame_interpolation_model_path` | str | None | Local directory or Hugging Face repo ID with `flownet.pkl`; defaults to `elfgum/RIFE-4.22.lite` |
+
+## Frame Interpolation
+
+Frame interpolation is an optional post-processing step for `/v1/videos` and
+`/v1/videos/sync`. It synthesizes intermediate frames between generated frames
+without rerunning the diffusion model. If the generated video has `N` frames,
+the interpolated output frame count is `(N - 1) * 2**exp + 1`. The encoder FPS
+is multiplied by `2**exp` so the output duration remains close to the original.
+
+Frame interpolation runs in the diffusion worker post-processing path instead of
+the API server encoding path, so it can reuse the worker's current accelerator
+device without blocking the FastAPI event loop.
+
+Example: generate 5 frames and interpolate to 9 frames:
+
+```bash
+curl -X POST http://localhost:8091/v1/videos/sync \
+  -F "prompt=A dog running through a park" \
+  -F "num_frames=5" \
+  -F "fps=8" \
+  -F "enable_frame_interpolation=true" \
+  -F "frame_interpolation_exp=1" \
+  -F "frame_interpolation_scale=1.0" \
+  -o sync_t2v_interpolated.mp4
+```

 ## Create Response Format

@@ -129,17 +129,8 @@ The first run loads model weights. Subsequent runs are faster:
 --stage-init-timeout 900  # Increase timeout for slow storage
 ```

-### `Transformers does not recognize this architecture` Error
-
-Your have to upgrade `transformers` package to `5.3.0` or above:
-
-```
-pip install --upgrade transformers
-```
-
 ## Requirements

 - vLLM-Omni with GLM-Image support
 - CUDA-capable GPU (recommended: H100/A100 with 80GB)
 - GLM-Image model weights
- `transformers` v5.3.0 or above
@@ -178,12 +178,6 @@ def parse_args() -> argparse.Namespace:
        action="store_true",
        help="Enable diffusion pipeline profiler to display stage durations.",
    )
-    parser.add_argument(
-        "--quantization",
-        type=str,
-        default=None,
-        help="Quantization method (e.g. 'fp8').",
-    )

    # Memory & parallelism
    parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
@@ -219,7 +213,6 @@ def main():
        parallel_config=parallel_config,
        enforce_eager=args.enforce_eager,
        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
-        quantization=args.quantization,
    )

    # Validate I2V / V2V arguments
@@ -1,10 +1,7 @@
 import argparse
 from functools import lru_cache

-try:
-    import gradio as gr
-except ImportError:
-    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+import gradio as gr
 import torch

 from vllm_omni.entrypoints.omni import Omni
@@ -64,8 +64,8 @@ def parse_args() -> argparse.Namespace:
    parser.add_argument(
        "--guidance-scale",
        type=float,
-        default=4.0,
-        help="Classifier-free guidance scale. HunyuanImage3 recommends 4.0-5.0.",
+        default=1.0,
+        help="Classifier-free guidance scale.",
    )
    parser.add_argument("--height", type=int, default=1024, help="Height of generated image.")
    parser.add_argument("--width", type=int, default=1024, help="Width of generated image.")
@@ -16,10 +16,7 @@ import argparse
 import base64
 import io

-try:
-    import gradio as gr
-except ImportError:
-    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+import gradio as gr
 import httpx
 import numpy as np
 import soundfile as sf
@@ -196,9 +196,3 @@ gpu_memory_utilization: 0.5
 ```

 - The first request may be slow due to model warmup. Subsequent requests will be faster.
-
- If you encounter `Transformers does not recognize this architecture` error, your have to upgrade `transformers` package to `5.3.0` or above:
-
-```
-pip install --upgrade transformers
-```
@@ -10,10 +10,7 @@ import argparse
 import base64
 from io import BytesIO

-try:
-    import gradio as gr
-except ImportError:
-    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+import gradio as gr
 import requests
 from PIL import Image

@@ -6,10 +6,7 @@ import random
 from pathlib import Path
 from typing import Any

-try:
-    import gradio as gr
-except ImportError:
-    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+import gradio as gr
 import numpy as np
 import soundfile as sf
 import torch
@@ -211,7 +211,7 @@ print(response.choices[1].message.audio)
 Supported speaker names depend on the model (e.g. `Ethan`, `Chelsie`, `Aiden`). Omit `speaker` to use the default.

 ## Streaming Output
-If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. We support both text streaming output and audio streaming output. Other modalities can output normally.
+If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.
 ```bash
 python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
    --query-type use_image \
@@ -6,10 +6,7 @@ import random
 from pathlib import Path
 from typing import Any

-try:
-    import gradio as gr
-except ImportError:
-    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+import gradio as gr
 import numpy as np
 import soundfile as sf
 import torch
@@ -24,10 +24,7 @@ import io
 import json
 import logging

-try:
-    import gradio as gr
-except ImportError:
-    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+import gradio as gr
 import httpx
 import numpy as np
 import soundfile as sf
@@ -3,10 +3,7 @@
 import base64
 import io

-try:
-    import gradio as gr
-except ImportError:
-    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+import gradio as gr
 import httpx
 import numpy as np
 import soundfile as sf
@@ -10,10 +10,7 @@ import argparse
 import base64
 from io import BytesIO

-try:
-    import gradio as gr
-except ImportError:
-    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+import gradio as gr
 import requests
 from PIL import Image

@@ -20,10 +20,7 @@ from datetime import datetime
 from pathlib import Path
 from typing import Any

-try:
-    import gradio as gr
-except ImportError:
-    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+import gradio as gr
 import httpx
 import numpy as np
 import soundfile as sf
@@ -56,10 +56,6 @@ dev = [
    "mistune>=3.2.0", # for example tests
 ]

-demo = [
-    "gradio>=6.7.0",
-]
-
 docs = [
    "mkdocs>=1.5.0",
    "mkdocs-api-autonav",
@@ -4,6 +4,7 @@ librosa>=0.11.0
 resampy>=0.4.3
 diffusers>=0.36.0
 accelerate==1.12.0
+gradio==5.50
 soundfile>=0.13.1
 cache-dit==1.3.0
 tqdm>=4.66.0
@@ -17,4 +18,3 @@ prettytable>=3.8.0
 aenum==3.1.16
 pyzmq>=25.0.0
 janus>=1.0.0
-pydub
@@ -118,6 +118,10 @@ def benchmark_params(request, omni_server):
    if param_index >= len(all_params):
        raise ValueError(f"No benchmark parameters found for index {param_index} in test: {test_name}")

+    if all_params[param_index]["dataset_name"] == "random-mm":
+        # TODO: Due to known issues, skip the random-mm dataset.
+        pytest.skip("Skipping parameter for random-mm dataset.")
+
    current = param_index + 1
    total = len(all_params)
    print(f"\n  Running benchmark {current}/{total} for {test_name}")
@@ -128,70 +132,15 @@ def benchmark_params(request, omni_server):
    }


-def _resolve_baseline_value(
-    baseline_raw: Any,
-    *,
-    sweep_index: int | None,
-    max_concurrency: Any = None,
-    request_rate: Any = None,
-) -> Any:
-    """Pick the baseline threshold for this sweep step.
-
-    Supported shapes per metric:
-    - **Scalar** — same threshold for every concurrency / QPS.
-    - **List** — aligned with ``max_concurrency`` / ``request_rate`` sweep order; use ``sweep_index``.
-    - **Dict** — keyed by concurrency or rate, e.g. ``{"1": 500, "4": 800}`` (keys are strings in JSON).
-
-    For dict lookup, ``max_concurrency`` is preferred when both are set (concurrency sweep).
-    """
-    if baseline_raw is None:
-        # If no baseline is set, the maximum value will be used.
-        return 100000
-    if isinstance(baseline_raw, dict):
-        if max_concurrency is not None:
-            for key in (max_concurrency, str(max_concurrency)):
-                if key in baseline_raw:
-                    return baseline_raw[key]
-        if request_rate is not None:
-            for key in (request_rate, str(request_rate)):
-                if key in baseline_raw:
-                    return baseline_raw[key]
-        raise KeyError(
-            f"baseline dict has no key for max_concurrency={max_concurrency!r} "
-            f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}"
-        )
-    if isinstance(baseline_raw, (list, tuple)):
-        return baseline_raw[sweep_index]
-    return baseline_raw
-
-
-def assert_result(
-    result,
-    params,
-    num_prompt,
-    *,
-    sweep_index: int | None = None,
-    max_concurrency: Any = None,
-    request_rate: Any = None,
-) -> None:
+def assert_result(result, params, num_prompt):
    assert result["completed"] == num_prompt, "Request failures exist"
    baseline_data = params.get("baseline", {})
-    for metric_name, baseline_raw in baseline_data.items():
+    for metric_name, baseline_value in baseline_data.items():
        current_value = result[metric_name]
-        baseline_value = _resolve_baseline_value(
-            baseline_raw,
-            sweep_index=sweep_index,
-            max_concurrency=max_concurrency,
-            request_rate=request_rate,
-        )
        if "throughput" in metric_name:
-            if current_value <= baseline_value:
-                print(
-                    f"ERROR: Throughput test results were below baseline: {metric_name}: {current_value} > {baseline_value}"
-                )
+            assert current_value >= baseline_value, f"{metric_name}: {current_value} < {baseline_value}"
        else:
-            if current_value >= baseline_value:
-                print(f"ERROR: Test results exceeded baseline: {metric_name}: {current_value} < {baseline_value}")
+            assert current_value <= baseline_value, f"{metric_name}: {current_value} > {baseline_value}"


@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@@ -246,8 +195,8 @@ def test_performance_benchmark(omni_server, benchmark_params):
        elif not isinstance(value, bool):
            args.extend([arg_name, str(value)])

-    # QPS test (sweep_index aligns with qps_list / num_prompt_list for this loop)
-    for i, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
+    # QPS test
+    for qps, num_prompt in zip(qps_list, num_prompt_list):
        args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
        result = run_benchmark(
            args=args,
@@ -256,16 +205,10 @@ def test_performance_benchmark(omni_server, benchmark_params):
            dataset_name=dataset_name,
            num_prompt=num_prompt,
        )
-        assert_result(
-            result,
-            params,
-            num_prompt=num_prompt,
-            sweep_index=i,
-            request_rate=qps,
-        )
+        assert_result(result, params, num_prompt=num_prompt)

-    # concurrency test (sweep_index aligns with max_concurrency_list for separate thresholds per concurrency)
-    for i, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
+    # concurrency test
+    for concurrency, num_prompt in zip(max_concurrency_list, num_prompt_list):
        args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
        result = run_benchmark(
            args=args,
@@ -274,10 +217,4 @@ def test_performance_benchmark(omni_server, benchmark_params):
            dataset_name=dataset_name,
            num_prompt=num_prompt,
        )
-        assert_result(
-            result,
-            params,
-            num_prompt=num_prompt,
-            sweep_index=i,
-            max_concurrency=concurrency,
-        )
+        assert_result(result, params, num_prompt=num_prompt)
@@ -25,9 +25,9 @@
                "ignore_eos": true,
                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                "baseline": {
-                    "mean_ttft_ms": [1000, 3000, 5000],
-                    "mean_audio_ttfp_ms": [8000, 10000, 13000],
-                    "mean_audio_rtf": [0.2, 0.25, 0.45]
+                    "mean_ttft_ms": 100000,
+                    "mean_audio_ttfp_ms": 100000,
+                    "mean_audio_rtf": 100000
                }
            },
            {
@@ -39,10 +39,10 @@
                    40,
                    100
                ],
-                "request_rate": [
-                    0.1,
-                    0.3,
-                    0.5
+                "max_concurrency": [
+                    1,
+                    4,
+                    10
                ],
                "random_input_len": 100,
                "random_output_len": 100,
@@ -62,31 +62,9 @@
                },
                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                "baseline": {
-                    "mean_ttft_ms": [2000, 4000, 6000],
-                    "mean_audio_ttfp_ms": [10000, 13000, 15000],
-                    "mean_audio_rtf": [0.25, 0.35, 0.45]
-                }
-            },
-            {
-                "dataset_name": "random",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    4,
-                    16
-                ],
-                "max_concurrency": [
-                    1,
-                    4
-                ],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [1000, 3000],
-                    "mean_audio_ttfp_ms": [30000, 60000],
-                    "mean_audio_rtf": [0.35, 0.45]
+                    "mean_ttft_ms": 100000,
+                    "mean_audio_ttfp_ms": 100000,
+                    "mean_audio_rtf": 100000
                }
            }
        ]
@@ -135,9 +113,9 @@
                "ignore_eos": true,
                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                "baseline": {
-                    "mean_ttft_ms": [1000, 3000, 5000],
-                    "mean_audio_ttfp_ms": [1000, 3000, 5000],
-                    "mean_audio_rtf": [0.2, 0.35, 0.6]
+                    "mean_ttft_ms": 100000,
+                    "mean_audio_ttfp_ms": 100000,
+                    "mean_audio_rtf": 100000
                }
            },
            {
@@ -149,10 +127,10 @@
                    40,
                    100
                ],
-                "request_rate": [
-                    0.1,
-                    0.3,
-                    0.5
+                "max_concurrency": [
+                    1,
+                    4,
+                    10
                ],
                "random_input_len": 100,
                "random_output_len": 100,
@@ -172,31 +150,9 @@
                },
                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                "baseline": {
-                    "mean_ttft_ms": [2000, 4000, 6000],
-                    "mean_audio_ttfp_ms": [2000, 4000, 6000],
-                    "mean_audio_rtf": [0.25, 0.4, 0.7]
-                }
-            },
-            {
-                "dataset_name": "random",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    4,
-                    16
-                ],
-                "max_concurrency": [
-                    1,
-                    4
-                ],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [1000, 3000],
-                    "mean_audio_ttfp_ms": [1000, 3000],
-                    "mean_audio_rtf": [0.35, 0.45]
+                    "mean_ttft_ms": 100000,
+                    "mean_audio_ttfp_ms": 100000,
+                    "mean_audio_rtf": 100000
                }
            }
        ]
@@ -227,8 +183,8 @@
                },
                "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
                "baseline": {
-                    "mean_audio_ttfp_ms": [6000, 6000],
-                    "mean_audio_rtf": [0.3, 0.3]
+                    "mean_audio_ttfp_ms": 100000,
+                    "mean_audio_rtf": 100000
                }
            }
        ]
@@ -0,0 +1,43 @@
+import pytest
+import torch
+
+from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_wan as wan_vae_module
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummyOmniAutoencoderKLWan(OmniAutoencoderKLWan):
+    def __init__(self, *, dtype: torch.dtype):
+        torch.nn.Module.__init__(self)
+        self.register_parameter("dummy_weight", torch.nn.Parameter(torch.ones(1, dtype=dtype)))
+
+
+def test_wan_vae_execution_context_handles_fp32():
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.float32)
+    with model._execution_context():
+        output = model.dummy_weight + 1
+    assert output.dtype == torch.float32
+
+
+def test_wan_vae_execution_context_handles_bf16():
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
+    with model._execution_context():
+        output = model.dummy_weight + 1
+    assert output.dtype == torch.bfloat16
+
+
+def test_wan_vae_execution_context_uses_platform_autocast(mocker):
+    sentinel = object()
+    platform = mocker.Mock()
+    platform.create_autocast_context.return_value = sentinel
+    mocker.patch.object(wan_vae_module, "current_omni_platform", platform)
+
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
+
+    assert model._execution_context() is sentinel
+    platform.create_autocast_context.assert_called_once_with(
+        device_type=model.dummy_weight.device.type,
+        dtype=torch.bfloat16,
+        enabled=True,
+    )
@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for DistributedAutoencoderKLWan encode parallel (CPU-only)."""
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _DummyConfig:
+    def __init__(self, patch_size=None, scale_factor_temporal=4):
+        self.patch_size = patch_size
+        self.scale_factor_temporal = scale_factor_temporal
+
+
+class _DummyWanVae:
+    """Minimal mock of DistributedAutoencoderKLWan for testing encode_tile_split."""
+
+    def __init__(
+        self,
+        config=None,
+        spatial_compression_ratio=8,
+        tile_sample_min_height=256,
+        tile_sample_min_width=256,
+        tile_sample_stride_height=192,
+        tile_sample_stride_width=192,
+    ):
+        self.config = config or _DummyConfig()
+        self.spatial_compression_ratio = spatial_compression_ratio
+        self.tile_sample_min_height = tile_sample_min_height
+        self.tile_sample_min_width = tile_sample_min_width
+        self.tile_sample_stride_height = tile_sample_stride_height
+        self.tile_sample_stride_width = tile_sample_stride_width
+        self.dtype = torch.float32
+
+        # Mock caches
+        self._enc_feat_map = None
+        self._enc_conv_idx = [0]
+
+    def clear_cache(self):
+        self._enc_feat_map = None
+        self._enc_conv_idx = [0]
+
+    def encoder(self, x, feat_cache=None, feat_idx=None):  # noqa: ARG002
+        # Simple mock: just return the input
+        return x
+
+    def quant_conv(self, x):
+        return x
+
+    def blend_v(self, _a, b, _blend_extent):
+        return b
+
+    def blend_h(self, _a, b, _blend_extent):
+        return b
+
+
+def _import_encode_tile_split():
+    """Import the encode_tile_split method from the module."""
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+        DistributedAutoencoderKLWan,
+    )
+
+    return DistributedAutoencoderKLWan.encode_tile_split
+
+
+def _import_encode_tile_exec():
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+        DistributedAutoencoderKLWan,
+    )
+
+    return DistributedAutoencoderKLWan.encode_tile_exec
+
+
+def _import_encode_tile_merge():
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+        DistributedAutoencoderKLWan,
+    )
+
+    return DistributedAutoencoderKLWan.encode_tile_merge
+
+
+class TestEncodeTileSplit:
+    """Tests for encode_tile_split method."""
+
+    def test_basic_split_without_patch_size(self):
+        """Test basic tile splitting without patch_size."""
+        encode_tile_split = _import_encode_tile_split()
+
+        vae = _DummyWanVae(
+            config=_DummyConfig(patch_size=None, scale_factor_temporal=4),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=256,
+            tile_sample_min_width=256,
+            tile_sample_stride_height=192,
+            tile_sample_stride_width=192,
+        )
+
+        # Input: (B, C, T, H, W) = (1, 3, 5, 256, 256)
+        x = torch.randn(1, 3, 5, 256, 256)
+
+        tiletask_list, grid_spec = encode_tile_split(vae, x)
+
+        # With stride 192 and input size 256, we should get:
+        # Height: ceil(256/192) = 2 positions (0, 192) but 192+256 > 256, so only 1
+        # Actually for i in range(0, 256, 192): i = 0, 192 but 192 is out of bounds
+        # So we get 1x1 grid
+        assert len(tiletask_list) >= 1
+        assert grid_spec.grid_shape[0] >= 1
+        assert grid_spec.grid_shape[1] >= 1
+
+        # Check temporal chunking: 5 frames -> 1 + (5-1)//4 = 2 chunks
+        first_task = tiletask_list[0]
+        assert len(first_task.tensor) == 2  # 2 temporal chunks
+
+    def test_split_with_patch_size_scales_coordinates(self):
+        """Test that patch_size properly scales tile coordinates."""
+        encode_tile_split = _import_encode_tile_split()
+
+        # Without patch_size
+        vae_no_patch = _DummyWanVae(
+            config=_DummyConfig(patch_size=None, scale_factor_temporal=4),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=256,
+            tile_sample_min_width=256,
+            tile_sample_stride_height=128,
+            tile_sample_stride_width=128,
+        )
+
+        # With patch_size=2 (simulating patchified input)
+        vae_with_patch = _DummyWanVae(
+            config=_DummyConfig(patch_size=2, scale_factor_temporal=4),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=256,
+            tile_sample_min_width=256,
+            tile_sample_stride_height=128,
+            tile_sample_stride_width=128,
+        )
+
+        # Same patchified input size
+        x = torch.randn(1, 3, 5, 256, 256)
+
+        tasks_no_patch, _ = encode_tile_split(vae_no_patch, x)
+        tasks_with_patch, _ = encode_tile_split(vae_with_patch, x)
+
+        # With patch_size=2, stride becomes 128//2=64, so more tiles
+        assert len(tasks_with_patch) >= len(tasks_no_patch)
+
+    def test_temporal_compression_from_config(self):
+        """Test that temporal compression ratio is read from config."""
+        encode_tile_split = _import_encode_tile_split()
+
+        # temporal_compression=4 (default)
+        vae_4x = _DummyWanVae(
+            config=_DummyConfig(scale_factor_temporal=4),
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=512,
+            tile_sample_stride_width=512,
+        )
+
+        # temporal_compression=2
+        vae_2x = _DummyWanVae(
+            config=_DummyConfig(scale_factor_temporal=2),
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=512,
+            tile_sample_stride_width=512,
+        )
+
+        # 9 frames input
+        x = torch.randn(1, 3, 9, 512, 512)
+
+        tasks_4x, _ = encode_tile_split(vae_4x, x)
+        tasks_2x, _ = encode_tile_split(vae_2x, x)
+
+        # With 4x compression: 1 + (9-1)//4 = 3 chunks
+        assert len(tasks_4x[0].tensor) == 3
+
+        # With 2x compression: 1 + (9-1)//2 = 5 chunks
+        assert len(tasks_2x[0].tensor) == 5
+
+    def test_grid_spec_latent_dimensions(self):
+        """Test that grid_spec contains correct latent dimensions."""
+        encode_tile_split = _import_encode_tile_split()
+
+        vae = _DummyWanVae(
+            config=_DummyConfig(patch_size=None),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=512,
+            tile_sample_stride_width=512,
+        )
+
+        # Input: 512x512 with compression 8 -> 64x64 latent
+        x = torch.randn(1, 3, 5, 512, 512)
+
+        _, grid_spec = encode_tile_split(vae, x)
+
+        assert grid_spec.tile_spec["latent_height"] == 64
+        assert grid_spec.tile_spec["latent_width"] == 64
+
+
+class TestEncodeTileExec:
+    """Tests for encode_tile_exec method."""
+
+    def test_basic_exec(self):
+        """Test basic tile execution."""
+        encode_tile_exec = _import_encode_tile_exec()
+
+        vae = _DummyWanVae()
+
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+            TileTask,
+        )
+
+        # Create a simple task with 2 temporal chunks
+        tile1 = torch.randn(1, 3, 1, 32, 32)
+        tile2 = torch.randn(1, 3, 4, 32, 32)
+        task = TileTask(tile_id=0, grid_coord=(0, 0), tensor=[tile1, tile2])
+
+        result = encode_tile_exec(vae, task)
+
+        # Result should concatenate temporal dimension
+        assert result.shape[2] == 5  # 1 + 4 frames
+
+
+class TestEncodeTileMerge:
+    """Tests for encode_tile_merge method."""
+
+    def test_basic_merge(self):
+        """Test basic tile merging."""
+        encode_tile_merge = _import_encode_tile_merge()
+
+        vae = _DummyWanVae()
+
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+            GridSpec,
+        )
+
+        # Create 2x2 grid of tiles
+        tile_00 = torch.ones(1, 16, 2, 32, 32) * 0
+        tile_01 = torch.ones(1, 16, 2, 32, 32) * 1
+        tile_10 = torch.ones(1, 16, 2, 32, 32) * 2
+        tile_11 = torch.ones(1, 16, 2, 32, 32) * 3
+
+        coord_tensor_map = {
+            (0, 0): tile_00,
+            (0, 1): tile_01,
+            (1, 0): tile_10,
+            (1, 1): tile_11,
+        }
+
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(2, 2),
+            tile_spec={
+                "latent_height": 48,
+                "latent_width": 48,
+                "blend_height": 8,
+                "blend_width": 8,
+                "tile_latent_stride_height": 24,
+                "tile_latent_stride_width": 24,
+            },
+        )
+
+        result = encode_tile_merge(vae, coord_tensor_map, grid_spec)
+
+        # Output should be (1, 16, 2, 48, 48)
+        assert result.shape == (1, 16, 2, 48, 48)
@@ -59,9 +59,9 @@ class E2EOperator:
 class DummyMixin(DistributedVaeMixin):
    def __init__(self):
        self.use_tiling = True
-        self.distributed_decoder = MagicMock()
-        self.distributed_decoder.parallel_size = 2
-        self.distributed_decoder.group = None
+        self.distributed_executor = MagicMock()
+        self.distributed_executor.parallel_size = 2
+        self.distributed_executor.group = None


@pytest.fixture(autouse=True)
@@ -7,49 +7,6 @@ import pytest
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


-class TestSetForwardContextNumTokens:
-    """Test _set_forward_context_num_tokens defensive fix for vLLM 0.18.0."""
-
-    def test_sets_num_tokens_when_context_available(self, mocker):
-        """num_tokens should be set on ForwardContext when available."""
-        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
-
-        mock_ctx = mocker.MagicMock()
-        del mock_ctx.in_profile_run  # simulate missing attr
-        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
-        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
-
-        hunyuan_moe._set_forward_context_num_tokens(1024)
-
-        assert mock_ctx.num_tokens == 1024
-        assert mock_ctx.in_profile_run is False
-
-    def test_sets_in_profile_run_only_if_missing(self, mocker):
-        """in_profile_run should not be overwritten if already set."""
-        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
-
-        mock_ctx = mocker.MagicMock()
-        mock_ctx.in_profile_run = True  # already set
-        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
-        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
-
-        hunyuan_moe._set_forward_context_num_tokens(512)
-
-        assert mock_ctx.num_tokens == 512
-        assert mock_ctx.in_profile_run is True  # not overwritten
-
-    def test_noop_when_context_unavailable(self, mocker):
-        """Should do nothing when ForwardContext is not available."""
-        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
-
-        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=False)
-        mock_get = mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context")
-
-        hunyuan_moe._set_forward_context_num_tokens(256)
-
-        mock_get.assert_not_called()
-
-
 class TestHunyuanFusedMoEPlatformDispatch:
    """Test platform dispatch via platform qualname hooks."""

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit_plus import (
+    get_qwen_image_edit_plus_pre_process_func,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_qwen_image_edit_plus_rejects_too_many_input_images(tmp_path: Path):
+    vae_dir = tmp_path / "vae"
+    vae_dir.mkdir()
+    # Keep the mock config intentionally minimal: this test only needs the
+    # fields touched during pre-process initialization.
+    (vae_dir / "config.json").write_text(json.dumps({"z_dim": 16}))
+
+    pre_process = get_qwen_image_edit_plus_pre_process_func(SimpleNamespace(model=str(tmp_path)))
+    image = Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8))
+    request = SimpleNamespace(
+        prompts=[
+            {
+                "prompt": "combine",
+                "multi_modal_data": {"image": [image, image, image, image, image]},
+            }
+        ],
+        sampling_params=SimpleNamespace(height=None, width=None),
+    )
+
+    with pytest.raises(ValueError, match=r"At most 4 images are supported by this model"):
+        pre_process(request)
@@ -0,0 +1,260 @@
+import inspect
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import (
+    QwenImagePipeline,
+)
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import (
+    QwenImageEditPipeline,
+)
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit_plus import (
+    QwenImageEditPlusPipeline,
+)
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_layered import (
+    QwenImageLayeredPipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _RejectingTextEncoder:
+    dtype = torch.float32
+
+    def __call__(self, *args, **kwargs):
+        raise AssertionError("text encoder should not run for prompts that exceed max_sequence_length")
+
+
+class _FakeModelInputs:
+    def __init__(self, total_sequence_length: int):
+        attention_mask = torch.ones((1, total_sequence_length), dtype=torch.long)
+        self.input_ids = attention_mask.clone()
+        self.attention_mask = attention_mask
+        self.pixel_values = None
+        self.image_grid_thw = None
+
+    def to(self, device):
+        return self
+
+
+class _FakeTokenizer:
+    def __init__(self, total_sequence_length: int | list[int]):
+        if isinstance(total_sequence_length, list):
+            self.total_sequence_lengths = list(total_sequence_length)
+        else:
+            self.total_sequence_lengths = [total_sequence_length]
+
+    def __call__(self, *args, **kwargs):
+        if len(self.total_sequence_lengths) > 1:
+            total_sequence_length = self.total_sequence_lengths.pop(0)
+        else:
+            total_sequence_length = self.total_sequence_lengths[0]
+        return _FakeModelInputs(total_sequence_length)
+
+
+class _FakeProcessor(_FakeTokenizer):
+    pass
+
+
+class _FakeScheduler:
+    def __init__(self):
+        self.begin_index = None
+
+    def set_begin_index(self, begin_index: int):
+        self.begin_index = begin_index
+
+
+PIPELINE_CASES = [
+    pytest.param(QwenImagePipeline, 34, "tokenizer", id="qwen-image"),
+    pytest.param(QwenImageLayeredPipeline, 34, "tokenizer", id="qwen-image-layered"),
+    pytest.param(QwenImageEditPipeline, 64, "processor", id="qwen-image-edit"),
+    pytest.param(QwenImageEditPlusPipeline, 64, "processor", id="qwen-image-edit-plus"),
+]
+
+
+def _make_pipeline(
+    pipeline_class: type,
+    *,
+    total_sequence_length: int,
+    drop_idx: int,
+    input_kind: str,
+):
+    pipeline = object.__new__(pipeline_class)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.text_encoder = _RejectingTextEncoder()
+    pipeline.tokenizer_max_length = 1024
+    pipeline.prompt_template_encode = "{}"
+    pipeline.prompt_template_encode_start_idx = drop_idx
+    pipeline.tokenizer = _FakeTokenizer([total_sequence_length, 0])
+    if input_kind == "processor":
+        pipeline.processor = _FakeProcessor(total_sequence_length)
+    return pipeline
+
+
+@pytest.mark.parametrize(("pipeline_class", "drop_idx", "input_kind"), PIPELINE_CASES)
+def test_encode_prompt_rejects_prompt_longer_than_default_max_sequence_length(
+    pipeline_class: type,
+    drop_idx: int,
+    input_kind: str,
+):
+    pipeline = _make_pipeline(
+        pipeline_class,
+        total_sequence_length=1025,
+        drop_idx=drop_idx,
+        input_kind=input_kind,
+    )
+
+    with pytest.raises(ValueError, match=r"got 1025 tokens, but `max_sequence_length` is 1024"):
+        pipeline.encode_prompt(prompt="prompt")
+
+
+@pytest.mark.parametrize(("pipeline_class", "drop_idx", "input_kind"), PIPELINE_CASES)
+def test_encode_prompt_rejects_prompt_longer_than_explicit_max_sequence_length(
+    pipeline_class: type,
+    drop_idx: int,
+    input_kind: str,
+):
+    pipeline = _make_pipeline(
+        pipeline_class,
+        total_sequence_length=17,
+        drop_idx=drop_idx,
+        input_kind=input_kind,
+    )
+
+    with pytest.raises(ValueError, match=r"got 17 tokens, but `max_sequence_length` is 16"):
+        pipeline.encode_prompt(prompt="prompt", max_sequence_length=16)
+
+
+def test_prepare_encode_defaults_to_tokenizer_max_length():
+    pipeline = object.__new__(QwenImagePipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.tokenizer_max_length = 1024
+    pipeline.vae_scale_factor = 8
+    pipeline.default_sample_size = 128
+    pipeline.scheduler = _FakeScheduler()
+    pipeline._extract_prompts = lambda prompts: (["prompt"], None)
+
+    captured = {}
+
+    def _fake_prepare_generation_context(**kwargs):
+        captured["max_sequence_length"] = kwargs["max_sequence_length"]
+        embeds = torch.ones((1, 1, 1))
+        mask = torch.ones((1, 1), dtype=torch.long)
+        return {
+            "prompt_embeds": embeds,
+            "prompt_embeds_mask": mask,
+            "negative_prompt_embeds": None,
+            "negative_prompt_embeds_mask": None,
+            "latents": embeds,
+            "timesteps": torch.tensor([1]),
+            "do_true_cfg": False,
+            "guidance": None,
+            "img_shapes": [[(1, 1, 1)]],
+            "txt_seq_lens": [1],
+            "negative_txt_seq_lens": None,
+        }
+
+    pipeline._prepare_generation_context = _fake_prepare_generation_context
+    state = SimpleNamespace(
+        prompts=["prompt"],
+        sampling=SimpleNamespace(
+            height=None,
+            width=None,
+            num_inference_steps=None,
+            sigmas=None,
+            guidance_scale_provided=False,
+            num_outputs_per_prompt=0,
+            generator=None,
+            true_cfg_scale=None,
+            max_sequence_length=None,
+        ),
+    )
+
+    pipeline.prepare_encode(state)
+
+    assert captured["max_sequence_length"] == 1024
+
+
+@pytest.mark.parametrize(
+    ("pipeline_class", "drop_idx"),
+    [
+        pytest.param(QwenImageEditPipeline, 64, id="qwen-image-edit"),
+        pytest.param(QwenImageEditPlusPipeline, 64, id="qwen-image-edit-plus"),
+    ],
+)
+def test_edit_pipelines_validate_text_prompt_length_before_image_token_expansion(
+    pipeline_class: type,
+    drop_idx: int,
+):
+    pipeline = object.__new__(pipeline_class)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.text_encoder = _RejectingTextEncoder()
+    pipeline.tokenizer_max_length = 1024
+    pipeline.prompt_template_encode = "{}"
+    pipeline.prompt_template_encode_start_idx = drop_idx
+    pipeline.tokenizer = _FakeTokenizer([8, 0])
+    pipeline.processor = _FakeProcessor(drop_idx + 1500)
+
+    with pytest.raises(AssertionError, match="text encoder should not run"):
+        pipeline.encode_prompt(prompt="short prompt")
+
+
+@pytest.mark.parametrize(
+    "pipeline_class",
+    [
+        pytest.param(QwenImagePipeline, id="qwen-image"),
+        pytest.param(QwenImageLayeredPipeline, id="qwen-image-layered"),
+    ],
+)
+def test_qwen_generation_validator_excludes_template_suffix_from_budget(pipeline_class: type):
+    pipeline = object.__new__(pipeline_class)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.text_encoder = _RejectingTextEncoder()
+    pipeline.tokenizer_max_length = 1024
+    pipeline.prompt_template_encode = "{}"
+    pipeline.prompt_template_encode_start_idx = 34
+    pipeline.tokenizer = _FakeTokenizer([1029, 5])
+
+    with pytest.raises(AssertionError, match="text encoder should not run"):
+        pipeline.encode_prompt(prompt="boundary prompt")
+
+
+@pytest.mark.parametrize(
+    "pipeline_class",
+    [
+        pytest.param(QwenImageEditPipeline, id="qwen-image-edit"),
+        pytest.param(QwenImageEditPlusPipeline, id="qwen-image-edit-plus"),
+    ],
+)
+def test_qwen_edit_validator_excludes_image_placeholders_from_budget(pipeline_class: type):
+    pipeline = object.__new__(pipeline_class)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.text_encoder = _RejectingTextEncoder()
+    pipeline.tokenizer_max_length = 1024
+    pipeline.prompt_template_encode = "{}"
+    pipeline.prompt_template_encode_start_idx = 64
+    pipeline.tokenizer = _FakeTokenizer([30, 20])
+    pipeline.processor = _FakeProcessor(1500)
+
+    with pytest.raises(AssertionError, match="text encoder should not run"):
+        pipeline.encode_prompt(prompt="short prompt")
+
+
+@pytest.mark.parametrize(
+    "pipeline_class",
+    [
+        QwenImagePipeline,
+        QwenImageLayeredPipeline,
+        QwenImageEditPipeline,
+        QwenImageEditPlusPipeline,
+    ],
+)
+def test_forward_max_sequence_length_default_is_1024(pipeline_class: type):
+    assert inspect.signature(pipeline_class.forward).parameters["max_sequence_length"].default == 1024
@@ -0,0 +1,26 @@
+import pytest
+
+from vllm_omni.diffusion.utils.size_utils import (
+    normalize_min_aligned_size,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    ("height", "width", "expected"),
+    [
+        (1, 1, (16, 16)),
+        (15, 15, (16, 16)),
+        (17, 17, (16, 16)),
+        (31, 33, (16, 32)),
+        (64, 80, (64, 80)),
+    ],
+)
+def test_normalize_min_aligned_size_clamps_to_minimum_aligned_shape(height, width, expected):
+    assert normalize_min_aligned_size(height, width, alignment=16) == expected
+
+
+def test_normalize_min_aligned_size_rejects_invalid_alignment():
+    with pytest.raises(ValueError, match="positive alignment"):
+        normalize_min_aligned_size(16, 16, alignment=0)
@@ -6,7 +6,7 @@ from unittest.mock import Mock, patch

 import pytest

-from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.data import DiffusionOutput, OmniRequestError, normalize_omni_error
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
@@ -245,6 +245,39 @@ class TestDiffusionEngine:
        assert output is expected
        engine.executor.add_req.assert_called_once_with(request)

+    def test_add_req_and_wait_for_response_wraps_executor_errors(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = RequestScheduler()
+        engine.scheduler.initialize(Mock())
+        engine.executor = Mock()
+        engine.executor.add_req.side_effect = RuntimeError("worker oom")
+        engine._rpc_lock = threading.Lock()
+
+        request = _make_request("engine-error")
+
+        with pytest.raises(OmniRequestError, match="worker oom") as exc_info:
+            engine.add_req_and_wait_for_response(request)
+
+        assert exc_info.value.error_type == "RuntimeError"
+
+    def test_add_req_and_wait_for_response_raises_scheduler_error_when_empty_without_requests(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = Mock()
+        engine.scheduler.add_request.return_value = "sched-1"
+        engine.scheduler.schedule.return_value = Mock(is_empty=True)
+        engine.scheduler.has_requests.return_value = False
+        engine.scheduler.pop_request_state.return_value = None
+        engine.executor = Mock()
+        engine._rpc_lock = threading.Lock()
+
+        request = _make_request("engine-empty")
+
+        with pytest.raises(OmniRequestError, match="no runnable requests") as exc_info:
+            engine.add_req_and_wait_for_response(request)
+
+        assert exc_info.value.error_type == "SchedulerError"
+        engine.scheduler.pop_request_state.assert_called_once_with("sched-1")
+
    def test_supports_scheduler_interface_injection(self) -> None:
        request = _make_request("engine_iface")
        expected = DiffusionOutput(output=None)
@@ -297,3 +330,15 @@ class TestDiffusionEngine:

        with pytest.raises(RuntimeError, match="Dummy run failed: boom"):
            engine._dummy_run()
+
+
+def test_normalize_omni_error_preserves_existing_metadata() -> None:
+    err = OmniRequestError("boom", status_code=409, error_type="Conflict")
+
+    normalized = normalize_omni_error(err, request_id="req-1", stage_id=3)
+
+    assert normalized is err
+    assert normalized.request_id == "req-1"
+    assert normalized.stage_id == 3
+    assert normalized.status_code == 409
+    assert normalized.error_type == "Conflict"
@@ -443,14 +443,42 @@ class TestWorker:
 class TestIPC:
    def test_pack_unpack_runner_output_shm(self):
        tensor = torch.zeros(300_000, dtype=torch.float32)
-        output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
+        log_probs = torch.ones(300_000, dtype=torch.float32)
+        output = RunnerOutput(
+            req_id="req-1",
+            finished=True,
+            result=DiffusionOutput(output=tensor, trajectory_log_probs=log_probs),
+        )

        packed = pack_diffusion_output_shm(output)
        assert isinstance(packed.result.output, dict)
        assert packed.result.output["__tensor_shm__"] is True
+        assert isinstance(packed.result.trajectory_log_probs, dict)
+        assert packed.result.trajectory_log_probs["__tensor_shm__"] is True

        unpacked = unpack_diffusion_output_shm(packed)
        assert isinstance(unpacked.result.output, torch.Tensor)
+        assert isinstance(unpacked.result.trajectory_log_probs, torch.Tensor)
+        torch.testing.assert_close(unpacked.result.output, tensor)
+        torch.testing.assert_close(unpacked.result.trajectory_log_probs, log_probs)
+
+    def test_pack_unpack_runner_output_shm_bfloat16(self):
+        tensor = torch.arange(600_000, dtype=torch.bfloat16)
+        output = RunnerOutput(
+            req_id="req-1",
+            finished=True,
+            result=DiffusionOutput(output=tensor),
+        )
+
+        packed = pack_diffusion_output_shm(output)
+        assert isinstance(packed.result.output, dict)
+        assert packed.result.output["__tensor_shm__"] is True
+        assert packed.result.output["torch_dtype"] == str(torch.bfloat16)
+        assert packed.result.output["numpy_dtype"] == "uint16"
+
+        unpacked = unpack_diffusion_output_shm(packed)
+        assert isinstance(unpacked.result.output, torch.Tensor)
+        assert unpacked.result.output.dtype == torch.bfloat16
        torch.testing.assert_close(unpacked.result.output, tensor)


@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
+from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def mock_engine():
+    with patch("vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine") as mock:
+        engine_instance = MagicMock()
+        mock.make_engine.return_value = engine_instance
+        yield engine_instance
+
+
+@pytest.fixture
+def client(mock_engine):
+    metadata = StageMetadata(
+        stage_id=0,
+        stage_type="diffusion",
+        engine_output_type="image",
+        is_comprehension=False,
+        requires_multimodal_data=False,
+        engine_input_source="prompt",
+        final_output=True,
+        final_output_type="image",
+        default_sampling_params={},
+        custom_process_input_func=None,
+        model_stage=None,
+        runtime_cfg=None,
+    )
+    with patch.object(InlineStageDiffusionClient, "_enrich_config"):
+        od_config = MagicMock(spec=OmniDiffusionConfig)
+        c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1)
+        yield c
+        c.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_inline_dispatch_request_success(client, mock_engine):
+    mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()])
+    mock_engine.step.return_value = [mock_result]
+
+    sampling_params = OmniDiffusionSamplingParams()
+    await client.add_request_async("req-1", "A test prompt", sampling_params)
+
+    # Wait for the task to be processed
+    for _ in range(10):
+        output = client.get_diffusion_output_nowait()
+        if output is not None:
+            break
+        await asyncio.sleep(0.01)
+
+    assert output is not None
+    assert output.request_id == "req-1"
+    mock_engine.step.assert_called_once()
+    request = mock_engine.step.call_args.args[0]
+    assert request.prompts == ["A test prompt"]
+    assert request.request_ids == ["req-1"]
+
+
+@pytest.mark.asyncio
+async def test_inline_dispatch_batch_success(client, mock_engine):
+    mock_engine.step.return_value = [
+        OmniRequestOutput.from_diffusion(request_id="req-batch-0", images=[MagicMock()]),
+        OmniRequestOutput.from_diffusion(request_id="req-batch-1", images=[MagicMock()]),
+    ]
+
+    sampling_params = OmniDiffusionSamplingParams()
+    prompts = ["prompt-0", "prompt-1"]
+    await client.add_batch_request_async("req-batch", prompts, sampling_params)
+
+    for _ in range(10):
+        output = client.get_diffusion_output_nowait()
+        if output is not None:
+            break
+        await asyncio.sleep(0.01)
+
+    assert output is not None
+    assert output.request_id == "req-batch"
+    assert len(output.images) == 2
+    request = mock_engine.step.call_args.args[0]
+    assert request.prompts == prompts
+    assert request.request_ids == ["req-batch-0", "req-batch-1"]
+
+
+@pytest.mark.asyncio
+async def test_inline_dispatch_request_error(client, mock_engine):
+    mock_engine.step.side_effect = RuntimeError("Engine failure")
+
+    sampling_params = OmniDiffusionSamplingParams()
+    await client.add_request_async("req-err", "A test prompt", sampling_params)
+
+    for _ in range(10):
+        output = client.get_diffusion_output_nowait()
+        if output is not None:
+            break
+        await asyncio.sleep(0.01)
+
+    assert output is not None
+    assert output["request_id"] == "req-err"
+    assert output["error"] == "Engine failure"
+    assert output["error_type"] == "RuntimeError"
+
+
+def test_inline_shutdown(client, mock_engine):
+    assert not client._shutting_down
+
+    # Shutting down should cleanly cancel anything queued and close engine
+    client.shutdown()
+
+    assert client._shutting_down
+    mock_engine.close.assert_called_once()
@@ -71,6 +71,67 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke
            os.environ[env_var] = old_env


+def test_initialize_stages_uses_inline_diffusion_client_for_single_stage(monkeypatch):
+    """Single-stage diffusion init should request the inline client path."""
+    import vllm_omni.engine.async_omni_engine as engine_mod
+    from vllm_omni.platforms import current_omni_platform
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "dummy-model"
+    engine.config_path = "dummy-config"
+    engine.num_stages = 1
+    engine.async_chunk = False
+    engine.diffusion_batch_size = 1
+    engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion")]
+
+    env_var = current_omni_platform.device_control_env_var
+    old_env = os.environ.get(env_var)
+    os.environ[env_var] = "0"
+
+    metadata = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="diffusion",
+        runtime_cfg={"devices": "0"},
+        prompt_expand_func=None,
+    )
+    captured_use_inline = None
+    diffusion_client = types.SimpleNamespace(is_comprehension=False)
+
+    monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None)
+    monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "extract_stage_metadata", lambda _cfg: metadata)
+    monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {})
+    monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None))
+    monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "_inject_kv_stage_info", lambda *_: None)
+
+    def _fake_initialize_diffusion_stage(*args, **kwargs):
+        nonlocal captured_use_inline
+        captured_use_inline = kwargs.get("use_inline")
+        return diffusion_client
+
+    monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", _fake_initialize_diffusion_stage)
+    monkeypatch.setattr(
+        engine_mod,
+        "finalize_initialized_stages",
+        lambda stage_clients, _input_processor: (
+            stage_clients,
+            [types.SimpleNamespace()],
+            [{"final_output_type": "image"}],
+        ),
+    )
+
+    try:
+        engine._initialize_stages(stage_init_timeout=1)
+    finally:
+        if old_env is None:
+            os.environ.pop(env_var, None)
+        else:
+            os.environ[env_var] = old_env
+
+    assert captured_use_inline is True
+
+
 def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch):
    """Regression test for GLM-Image t2i preprocessing path.

@@ -169,16 +169,10 @@ def test_client(mock_async_diffusion):
    app.state.engine_client = mock_async_diffusion
    app.state.diffusion_engine = mock_async_diffusion  # Also set for health endpoint
    app.state.stage_configs = [SimpleNamespace(stage_type="diffusion")]
-    from vllm.entrypoints.openai.models.protocol import BaseModelPath
-
-    from vllm_omni.entrypoints.openai.api_server import _DiffusionServingModels
-
-    app.state.openai_serving_models = _DiffusionServingModels(
-        [BaseModelPath(name="Qwen/Qwen-Image", model_path="Qwen/Qwen-Image")]
-    )
+    app.state.diffusion_model_name = "Qwen/Qwen-Image"  # For models endpoint
    app.state.args = Namespace(
-        default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}',
-        max_generated_image_size=4096,  # 64*64
+        default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}',
+        max_generated_image_size=1024 * 1792,
    )

    return TestClient(app)
@@ -200,7 +194,7 @@ def async_omni_test_client():
        SimpleNamespace(stage_type="diffusion"),
    ]
    app.state.args = Namespace(
-        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
+        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}',
        max_generated_image_size=1048576,  # 1024*1024 to support resolution tests
    )
    return TestClient(app)
@@ -222,7 +216,7 @@ def async_omni_rgba_test_client():
        SimpleNamespace(stage_type="diffusion"),
    ]
    app.state.args = Namespace(
-        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
+        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}',
        max_generated_image_size=1048576,
    )
    return TestClient(app)
@@ -244,8 +238,8 @@ def async_omni_stage_configs_only_client():
    # Intentionally do not populate app.state.stage_configs. Refactored
    # AsyncOmni exposes stage_configs on the engine instance.
    app.state.args = Namespace(
-        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
-        max_generated_image_size=4096,  # 64*64
+        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}',
+        max_generated_image_size=1024 * 1792,
    )
    return TestClient(app)

@@ -392,6 +386,18 @@ def test_image_edits_async_omni_stage_configs_only(async_omni_stage_configs_only
    assert len(captured) == 2


+def test_generate_images_max_size_rejected(async_omni_test_client):
+    """Test that a size exceeding max_generated_image_size returns 400."""
+    response = async_omni_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a cat",
+            "size": "2048x2048",  # 4,194,304 pixels > max_generated_image_size (1,048,576)
+        },
+    )
+    assert response.status_code == 400
+
+
 def test_generate_multiple_images(test_client):
    """Test generating multiple images"""
    response = test_client.post(
@@ -621,6 +627,13 @@ def test_parameter_validation():
    with pytest.raises(ValueError):
        ImageGenerationRequest(prompt="test", guidance_scale=21.0)

+    # Invalid layers for layered models (must stay within the backend-supported range)
+    with pytest.raises(ValueError):
+        ImageGenerationRequest(prompt="test", layers=2)
+
+    with pytest.raises(ValueError):
+        ImageGenerationRequest(prompt="test", layers=11)
+

 # Pass-Through Tests

@@ -740,6 +753,86 @@ def test_image_edit_images_processing(async_omni_test_client):
    assert processed_images[1].size == (24, 24)


+def test_image_edit_rejects_multiple_images_when_model_does_not_support_them(async_omni_test_client):
+    img_bytes_1 = make_test_image_bytes((16, 16))
+    img_bytes_2 = make_test_image_bytes((32, 32))
+
+    engine = async_omni_test_client.app.state.engine_client
+    engine.get_diffusion_od_config = lambda: SimpleNamespace(supports_multimodal_inputs=False)
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[
+            ("image", img_bytes_1),
+            ("image", img_bytes_2),
+        ],
+        data={"prompt": "hello world."},
+    )
+
+    assert response.status_code == 400
+    assert (
+        response.json()["detail"] == "Received multiple input images. Only a single image is supported by this model."
+    )
+    assert engine.captured_prompt is None
+
+
+def test_image_edit_rejects_too_many_images_for_qwen_image_edit_2511(async_omni_test_client):
+    engine = async_omni_test_client.app.state.engine_client
+    engine.get_diffusion_od_config = lambda: SimpleNamespace(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=4,
+    )
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+        ],
+        data={"prompt": "hello world."},
+    )
+
+    assert response.status_code == 400
+    assert response.json()["detail"] == "Received 5 input images. At most 4 images are supported by this model."
+    assert engine.captured_prompt is None
+
+
+def test_image_edit_rejects_too_many_images_for_qwen_image_edit_2511_before_loading(
+    async_omni_test_client, monkeypatch: pytest.MonkeyPatch
+):
+    import vllm_omni.entrypoints.openai.api_server as api_server_module
+
+    engine = async_omni_test_client.app.state.engine_client
+    engine.get_diffusion_od_config = lambda: SimpleNamespace(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=4,
+    )
+
+    def _fail_load(*args, **kwargs):
+        raise AssertionError("_load_input_images should not run for over-limit requests")
+
+    monkeypatch.setattr(api_server_module, "_load_input_images", _fail_load)
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+            ("image", make_test_image_bytes((16, 16))),
+        ],
+        data={"prompt": "hello world."},
+    )
+
+    assert response.status_code == 400
+    assert response.json()["detail"] == "Received 5 input images. At most 4 images are supported by this model."
+    assert engine.captured_prompt is None
+
+
 def test_image_edit_parameter_pass(async_omni_test_client):
    img_bytes_1 = make_test_image_bytes((16, 16))

@@ -845,35 +938,36 @@ def test_image_edit_invalid_resolution(async_omni_test_client):


 def test_image_edit_invalid_layers(async_omni_test_client):
-    """Test that invalid layers values (< 1) are rejected with 400."""
+    """Test that layered image edits reject out-of-range layers with 400."""
    img_bytes = make_test_image_bytes((16, 16))

-    # Test layers = 0
+    # Test layers below the supported range
    response = async_omni_test_client.post(
        "/v1/images/edits",
        files=[("image", img_bytes)],
        data={
            "prompt": "test",
-            "layers": 0,
+            "layers": 2,
        },
    )
    assert response.status_code == 400
    detail = response.json()["detail"]
    assert "Invalid layers" in detail
-    assert "layers must be >= 1" in detail
+    assert "layers must be between 3 and 10 inclusive" in detail

-    # Test layers = -1
+    # Test layers above the supported range
    response = async_omni_test_client.post(
        "/v1/images/edits",
        files=[("image", img_bytes)],
        data={
            "prompt": "test",
-            "layers": -1,
+            "layers": 11,
        },
    )
    assert response.status_code == 400
    detail = response.json()["detail"]
    assert "Invalid layers" in detail
+    assert "layers must be between 3 and 10 inclusive" in detail


 def test_image_edit_resolution_and_size_conflict(async_omni_test_client):
@@ -917,6 +1011,7 @@ def test_image_edit_parameter_default(async_omni_test_client):
    assert captured_sampling_params.num_outputs_per_prompt == 1
    assert captured_sampling_params.num_inference_steps == 4
    assert captured_sampling_params.guidance_scale == 7.5
+    assert captured_sampling_params.generator_device == "cpu"

    # Test that a size exceeding max_generated_image_size returns 400
    response = async_omni_test_client.post(
@@ -950,13 +1045,15 @@ def test_image_edit_parameter_default_single_stage(test_client):
    assert captured_sampling_params.num_outputs_per_prompt == 1
    assert captured_sampling_params.num_inference_steps == 4
    assert captured_sampling_params.guidance_scale == 7.5
+    assert captured_sampling_params.generator_device == "cpu"

+    # Size exceeding max_generated_image_size (1024*1792) returns 400
    response = test_client.post(
        "/v1/images/edits",
        files=[("image", img_bytes_1)],
        data={
            "prompt": "hello world.",
-            "size": "96x96",
+            "size": "2048x2048",
        },
    )
    assert response.status_code == 400
@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for OpenAI-compatible video API encoding helpers."""
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.diffusion.postprocess import rife_interpolator
+from vllm_omni.entrypoints.openai import video_api_utils
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _install_fake_video_mux(monkeypatch, mux_calls):
+    def _fake_mux_video_audio_bytes(frames, audio, fps, audio_sample_rate, video_codec_options=None):
+        mux_calls.append(
+            {
+                "frames": frames,
+                "audio": audio,
+                "fps": fps,
+                "audio_sample_rate": audio_sample_rate,
+                "video_codec_options": video_codec_options,
+            }
+        )
+        return b"fake-video"
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.utils.media_utils.mux_video_audio_bytes",
+        _fake_mux_video_audio_bytes,
+    )
+
+
+def _install_fake_export_to_video(monkeypatch, export_calls):
+    def _fake_export_to_video(frames, output_path, fps):
+        export_calls.append(
+            {
+                "frames": frames,
+                "output_path": output_path,
+                "fps": fps,
+            }
+        )
+        with open(output_path, "wb") as f:
+            f.write(b"fake-video")
+
+    monkeypatch.setattr("diffusers.utils.export_to_video", _fake_export_to_video)
+
+
+def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch):
+    export_calls = []
+    _install_fake_export_to_video(monkeypatch, export_calls)
+
+    frames = [np.full((2, 2, 3), fill_value=i / 5, dtype=np.float32) for i in range(5)]
+    video_bytes = video_api_utils._encode_video_bytes(
+        frames,
+        fps=8,
+    )
+
+    assert video_bytes == b"fake-video"
+    assert len(export_calls) == 1
+    assert len(export_calls[0]["frames"]) == 5
+    assert export_calls[0]["frames"][0].shape == (2, 2, 3)
+    assert export_calls[0]["fps"] == 8
+
+
+def test_encode_video_bytes_uses_array_fast_path(monkeypatch):
+    export_calls = []
+    _install_fake_export_to_video(monkeypatch, export_calls)
+    monkeypatch.setattr(
+        video_api_utils,
+        "_normalize_frames",
+        lambda frames: (_ for _ in ()).throw(AssertionError("ndarray input should not go through frame list path")),
+    )
+
+    video = np.linspace(0.0, 1.0, num=5 * 2 * 2 * 3, dtype=np.float32).reshape(5, 2, 2, 3)
+    video_bytes = video_api_utils._encode_video_bytes(video, fps=12)
+
+    assert video_bytes == b"fake-video"
+    assert len(export_calls) == 1
+    assert len(export_calls[0]["frames"]) == 5
+    assert export_calls[0]["frames"][0].flags["C_CONTIGUOUS"]
+    assert export_calls[0]["fps"] == 12.0
+
+
+def test_encode_video_bytes_without_audio_uses_diffusers_export(monkeypatch):
+    export_calls = []
+    _install_fake_export_to_video(monkeypatch, export_calls)
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.utils.media_utils.mux_video_audio_bytes",
+        lambda *args, **kwargs: (_ for _ in ()).throw(
+            AssertionError("no-audio path should not use mux_video_audio_bytes")
+        ),
+    )
+
+    video = np.linspace(0.0, 1.0, num=4 * 2 * 2 * 3, dtype=np.float32).reshape(4, 2, 2, 3)
+    video_bytes = video_api_utils._encode_video_bytes(video, fps=10)
+
+    assert video_bytes == b"fake-video"
+    assert len(export_calls) == 1
+    assert export_calls[0]["fps"] == 10
+    assert len(export_calls[0]["frames"]) == 4
+
+
+def test_rife_model_inference_runs_on_dummy_tensors():
+    model = rife_interpolator.Model().eval()
+    img0 = torch.rand(1, 3, 32, 32)
+    img1 = torch.rand(1, 3, 32, 32)
+
+    output = model.inference(img0, img1, scale=1.0)
+
+    assert output.shape == (1, 3, 32, 32)
+    assert torch.isfinite(output).all()
+
+
+def test_frame_interpolator_runs_actual_torch_tensor_path(monkeypatch):
+    model = rife_interpolator.Model().eval()
+    interpolator = rife_interpolator.FrameInterpolator()
+    monkeypatch.setattr(interpolator, "_ensure_model_loaded", lambda preferred_device=None: model)
+
+    video = torch.zeros(1, 3, 2, 32, 32)
+    output_video, multiplier = interpolator.interpolate_tensor(video, exp=1, scale=1.0)
+
+    assert multiplier == 2
+    assert output_video.shape == (1, 3, 3, 32, 32)
+    assert torch.isfinite(output_video).all()
+
+
+def test_frame_interpolator_uses_platform_device_when_tensor_is_cpu(monkeypatch):
+    chosen_devices = []
+    model = rife_interpolator.Model().eval()
+
+    def _fake_ensure_model_loaded(*, preferred_device=None):
+        chosen_devices.append(preferred_device)
+        return model
+
+    interpolator = rife_interpolator.FrameInterpolator()
+    monkeypatch.setattr(interpolator, "_ensure_model_loaded", _fake_ensure_model_loaded)
+    monkeypatch.setattr(model.flownet, "to", lambda device: model.flownet)
+    monkeypatch.setattr(rife_interpolator, "_select_torch_device", lambda: torch.device("cuda"))
+
+    video = torch.zeros(1, 3, 2, 32, 32)
+    output_video, multiplier = interpolator.interpolate_tensor(video, exp=1, scale=1.0)
+
+    assert chosen_devices == [torch.device("cuda")]
+    assert multiplier == 2
+    assert output_video.shape == (1, 3, 3, 32, 32)
@@ -14,11 +14,12 @@ import time
 from types import SimpleNamespace

 import pytest
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from fastapi.testclient import TestClient
 from PIL import Image
 from pytest_mock import MockerFixture

+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.entrypoints.openai import api_server
 from vllm_omni.entrypoints.openai.api_server import router
 from vllm_omni.entrypoints.openai.protocol.videos import (
@@ -29,22 +30,39 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
 from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
 from vllm_omni.entrypoints.openai.storage import LocalStorageManager
 from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


 class MockVideoResult:
-    def __init__(self, videos, audios=None, sample_rate=None):
+    def __init__(
+        self,
+        videos,
+        audios=None,
+        sample_rate=None,
+        custom_output=None,
+        stage_durations=None,
+        peak_memory_mb=0.0,
+    ):
        self.multimodal_output = {"video": videos}
        if audios is not None:
            self.multimodal_output["audio"] = audios
        if sample_rate is not None:
            self.multimodal_output["audio_sample_rate"] = sample_rate
+        self._custom_output = custom_output or {}
+        self.stage_durations = stage_durations or {}
+        self.peak_memory_mb = peak_memory_mb
+
+    @property
+    def custom_output(self):
+        return self._custom_output


 class FakeAsyncOmni:
    def __init__(self):
        self.stage_configs = [SimpleNamespace(stage_type="diffusion")]
+        self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
        self.captured_prompt = None
        self.captured_sampling_params_list = None

@@ -56,6 +74,18 @@ class FakeAsyncOmni:
        yield MockVideoResult(videos)


+class ErrorAsyncOmni(FakeAsyncOmni):
+    def __init__(self, exc: Exception):
+        super().__init__()
+        self._exc = exc
+
+    async def generate(self, prompt, request_id, sampling_params_list):
+        self.captured_prompt = prompt
+        self.captured_sampling_params_list = sampling_params_list
+        raise self._exc
+        yield  # pragma: no cover
+
+
 class BlockingVideoHandler:
    def __init__(self):
        self.model_name = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
@@ -67,7 +97,7 @@ class BlockingVideoHandler:
        if self.stage_configs is None:
            self.stage_configs = stage_configs

-    async def generate_videos(self, request, reference_id, *, reference_image=None):
+    async def generate_video_bytes(self, request, reference_id, *, reference_image=None):
        self.started.set()
        try:
            await asyncio.Future()
@@ -135,15 +165,81 @@ def _wait_until(predicate, timeout_s: float = 2.0, interval_s: float = 0.02):
    raise AssertionError("Timed out waiting for condition")


+def test_async_video_generation_bypasses_base64(test_client, mocker: MockerFixture):
+    """Regression test: Ensure async video generation saves raw bytes directly
+    without bouncing through base64 encoding."""
+    # We mock _encode_video_bytes (the correct path)
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"raw-mp4-bytes",
+    )
+
+    # We assert that encode_video_base64 is never called
+    mock_base64 = mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        side_effect=RuntimeError("Regression: async video path should not base64 encode"),
+    )
+
+    response = test_client.post(
+        "/v1/videos",
+        data={"prompt": "A base64 test."},
+    )
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+
+    # Wait for completion. If it used base64, the RuntimeError would fail the task
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    mock_base64.assert_not_called()
+
+
+def test_async_video_generation_with_audio_bypasses_base64(test_client, mocker: MockerFixture):
+    """Regression test: Ensure async video generation passes audio through
+    generate_video_bytes without bouncing through base64 encoding."""
+    mock_encode = mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"raw-mp4-bytes",
+    )
+
+    mock_base64 = mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        side_effect=RuntimeError("Regression: async video path should not base64 encode"),
+    )
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+
+    async def _generate(prompt, request_id, sampling_params_list):
+        engine.captured_prompt = prompt
+        engine.captured_sampling_params_list = sampling_params_list
+        yield MockVideoResult([object()], audios=[object()], sample_rate=48000)
+
+    engine.generate = _generate
+
+    response = test_client.post(
+        "/v1/videos",
+        data={"prompt": "A base64 test with audio."},
+    )
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    mock_base64.assert_not_called()
+
+    mock_encode.assert_called_once()
+    kwargs = mock_encode.call_args.kwargs
+    assert "audio" in kwargs
+    assert kwargs["audio"] is not None
+    assert kwargs["audio_sample_rate"] == 48000
+
+
 def test_t2v_video_generation_form(test_client, mocker: MockerFixture):
    fps_values = []

-    def _fake_encode(video, fps):
+    def _fake_encode(video, fps, audio=None, audio_sample_rate=None, **kwargs):
        fps_values.append(fps)
-        return "Zg=="
+        return b"fake-video"

    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
        side_effect=_fake_encode,
    )
    response = test_client.post(
@@ -175,8 +271,8 @@ def test_i2v_video_generation_form(test_client, mocker: MockerFixture):
    image_bytes = _make_test_image_bytes((48, 32))

    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
    )
    response = test_client.post(
        "/v1/videos",
@@ -201,8 +297,8 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client,
    image_bytes = _make_test_image_bytes((48, 32))

    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
    )
    response = test_client.post(
        "/v1/videos",
@@ -227,8 +323,8 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client,

 def test_i2v_video_generation_with_image_reference_form(test_client, mocker: MockerFixture):
    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
    )
    response = test_client.post(
        "/v1/videos",
@@ -252,12 +348,12 @@ def test_i2v_video_generation_with_image_reference_form(test_client, mocker: Moc
 def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
    fps_values = []

-    def _fake_encode(video, fps):
+    def _fake_encode(video, fps, audio=None, audio_sample_rate=None, **kwargs):
        fps_values.append(fps)
-        return "Zg=="
+        return b"fake-video"

    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
        side_effect=_fake_encode,
    )
    response = test_client.post(
@@ -281,8 +377,8 @@ def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):

 def test_size_param_sets_width_height(test_client, mocker: MockerFixture):
    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
    )
    response = test_client.post(
        "/v1/videos",
@@ -303,8 +399,8 @@ def test_size_param_sets_width_height(test_client, mocker: MockerFixture):

 def test_sampling_params_pass_through(test_client, mocker: MockerFixture):
    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
    )
    response = test_client.post(
        "/v1/videos",
@@ -332,13 +428,155 @@ def test_sampling_params_pass_through(test_client, mocker: MockerFixture):
    assert captured.extra_args["flow_shift"] == 0.25


+def test_frame_interpolation_params_pass_to_diffusion_sampling_params(test_client, mocker: MockerFixture):
+    """Frame interpolation parameters should be forwarded to diffusion worker sampling params."""
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "smooth motion",
+            "fps": "8",
+            "enable_frame_interpolation": "true",
+            "frame_interpolation_exp": "2",
+            "frame_interpolation_scale": "0.5",
+            "frame_interpolation_model_path": "local-rife",
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+    captured = engine.captured_sampling_params_list[0]
+    assert captured.enable_frame_interpolation is True
+    assert captured.frame_interpolation_exp == 2
+    assert captured.frame_interpolation_scale == 0.5
+    assert captured.frame_interpolation_model_path == "local-rife"
+
+
+def test_default_sampling_params_apply_to_video_requests(test_client, mocker: MockerFixture):
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    engine = test_client.app.state.openai_serving_video._engine_client
+    engine.default_sampling_params_list = [
+        OmniDiffusionSamplingParams(
+            num_inference_steps=4,
+            guidance_scale=7.5,
+            generator_device="cpu",
+            enable_frame_interpolation=True,
+            frame_interpolation_exp=2,
+            frame_interpolation_scale=0.5,
+            frame_interpolation_model_path="default-rife",
+        )
+    ]
+
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "default param pass-through",
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    captured = engine.captured_sampling_params_list[0]
+    assert captured.num_inference_steps == 4
+    assert captured.guidance_scale == 7.5
+    assert captured.generator_device == "cpu"
+    assert captured.enable_frame_interpolation is True
+    assert captured.frame_interpolation_exp == 2
+    assert captured.frame_interpolation_scale == 0.5
+    assert captured.frame_interpolation_model_path == "default-rife"
+
+
+def test_request_params_override_default_video_sampling_params(test_client, mocker: MockerFixture):
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    engine = test_client.app.state.openai_serving_video._engine_client
+    engine.default_sampling_params_list = [
+        OmniDiffusionSamplingParams(
+            num_inference_steps=4,
+            guidance_scale=7.5,
+            enable_frame_interpolation=True,
+            frame_interpolation_exp=2,
+            frame_interpolation_scale=0.5,
+            frame_interpolation_model_path="default-rife",
+        )
+    ]
+
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "explicit override",
+            "num_inference_steps": "8",
+            "enable_frame_interpolation": "false",
+            "frame_interpolation_exp": "1",
+            "frame_interpolation_scale": "1.0",
+            "frame_interpolation_model_path": "custom-rife",
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    captured = engine.captured_sampling_params_list[0]
+    assert captured.num_inference_steps == 8
+    assert captured.guidance_scale == 7.5
+    assert captured.enable_frame_interpolation is False
+    assert captured.frame_interpolation_exp == 1
+    assert captured.frame_interpolation_scale == 1.0
+    assert captured.frame_interpolation_model_path == "custom-rife"
+
+
+def test_worker_fps_multiplier_is_applied_to_async_encoding(test_client, mocker: MockerFixture):
+    fps_values = []
+    engine = test_client.app.state.openai_serving_video._engine_client
+
+    async def _generate(prompt, request_id, sampling_params_list):
+        engine.captured_prompt = prompt
+        engine.captured_sampling_params_list = sampling_params_list
+        import numpy as np
+
+        yield MockVideoResult([np.zeros((1, 64, 64, 3), dtype=np.uint8)], custom_output={"video_fps_multiplier": 2})
+
+    engine.generate = _generate
+
+    def _fake_encode(video, fps, **kwargs):
+        del video, kwargs
+        fps_values.append(fps)
+        return b"fake-video"
+
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        side_effect=_fake_encode,
+    )
+
+    response = test_client.post("/v1/videos", data={"prompt": "fps multiplier", "fps": "8"})
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    assert fps_values == [16]
+
+
 def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFixture):
    audio_sample_rates = []

-    def _fake_encode(video, fps, audio=None, audio_sample_rate=None):
-        del video, fps, audio
+    def _fake_encode(video, fps, audio=None, audio_sample_rate=None, video_codec_options=None):
+        del video, fps, audio, video_codec_options
        audio_sample_rates.append(audio_sample_rate)
-        return "Zg=="
+        return b"fake-video"

    engine = test_client.app.state.openai_serving_video._engine_client
    engine.model_config = SimpleNamespace(
@@ -352,12 +590,14 @@ def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFi
    async def _generate(prompt, request_id, sampling_params_list):
        engine.captured_prompt = prompt
        engine.captured_sampling_params_list = sampling_params_list
-        yield MockVideoResult([object()], audios=[object()])
+        import numpy as np
+
+        yield MockVideoResult([np.zeros((1, 64, 64, 3), dtype=np.uint8)], audios=[object()])

    engine.generate = _generate

    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
        side_effect=_fake_encode,
    )
    response = test_client.post(
@@ -371,6 +611,33 @@ def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFi
    assert audio_sample_rates == [16000]


+def test_video_job_persists_profiler_metadata(test_client, mocker: MockerFixture):
+    engine = test_client.app.state.openai_serving_video._engine_client
+
+    async def _generate(prompt, request_id, sampling_params_list):
+        engine.captured_prompt = prompt
+        engine.captured_sampling_params_list = sampling_params_list
+        yield MockVideoResult(
+            [object()],
+            stage_durations={"diffuse": 2.5, "vae.decode": 0.3},
+            peak_memory_mb=4096.5,
+        )
+
+    engine.generate = _generate
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+
+    response = test_client.post("/v1/videos", data={"prompt": "profile me"})
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    completed = _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    assert completed["stage_durations"] == {"diffuse": 2.5, "vae.decode": 0.3}
+    assert completed["peak_memory_mb"] == 4096.5
+
+
 def test_missing_handler_returns_503():
    app = FastAPI()
    app.include_router(router)
@@ -428,8 +695,8 @@ def test_invalid_seconds_returns_422(test_client):

 def test_negative_prompt_and_seed_pass_through(test_client, mocker: MockerFixture):
    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
    )
    response = test_client.post(
        "/v1/videos",
@@ -498,12 +765,16 @@ def test_video_request_validation():

    with pytest.raises(ValueError):
        VideoGenerationRequest(prompt="test", image_reference={"file_id": "file-1", "image_url": "https://example.com"})
+    with pytest.raises(ValueError):
+        VideoGenerationRequest(prompt="test", frame_interpolation_exp=0)
+    with pytest.raises(ValueError):
+        VideoGenerationRequest(prompt="test", frame_interpolation_scale=0)


 def test_list_videos_supports_order_after_and_limit(test_client, mocker: MockerFixture):
    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
    )
    ids = []
    for i in range(3):
@@ -571,8 +842,8 @@ def test_list_videos_supports_order_after_and_limit(test_client, mocker: MockerF

 def test_delete_completed_job_removes_file_and_metadata(test_client, mocker: MockerFixture):
    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
    )
    create_resp = test_client.post("/v1/videos", data={"prompt": "Delete this video"})
    assert create_resp.status_code == 200
@@ -643,8 +914,8 @@ def test_video_response_file_extension_is_robust():
 def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture):
    """extra_params JSON object is merged into sampling_params.extra_args."""
    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
    )
    extra_params = {
        "is_enable_stage2": True,
@@ -674,8 +945,8 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture)
 def test_extra_params_none_by_default(test_client, mocker: MockerFixture):
    """When extra_params is omitted, extra_args stays empty."""
    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
    )
    response = test_client.post(
        "/v1/videos",
@@ -715,8 +986,8 @@ def test_extra_params_invalid_json(test_client):
 def test_extra_params_merged_with_existing_extra_args(test_client, mocker: MockerFixture):
    """extra_params is merged on top of existing extra_args (e.g. flow_shift)."""
    mocker.patch(
-        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
-        return_value="Zg==",
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
    )
    response = test_client.post(
        "/v1/videos",
@@ -735,3 +1006,59 @@ def test_extra_params_merged_with_existing_extra_args(test_client, mocker: Mocke
    assert captured.extra_args["flow_shift"] == 0.5
    assert captured.extra_args["use_zero_init"] is True
    assert captured.extra_args["zero_steps"] == 2
+
+
+def test_sample_solver_forwarded_via_extra_params(test_client, mocker: MockerFixture):
+    """sample_solver can be passed through existing extra_params for Wan2.2 online serving."""
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "A fox running through snow.",
+            "extra_params": json.dumps({"sample_solver": "euler"}),
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    engine = test_client.app.state.openai_serving_video._engine_client
+    captured = engine.captured_sampling_params_list[0]
+    assert captured.extra_args["sample_solver"] == "euler"
+
+
+@pytest.mark.asyncio
+async def test_run_generation_maps_omni_request_error_to_http_exception():
+    serving = OmniOpenAIServingVideo.for_diffusion(
+        diffusion_engine=ErrorAsyncOmni(
+            OmniRequestError(
+                "oom",
+                status_code=507,
+                request_id="req-oom",
+                stage_id=2,
+                error_type="OutOfMemoryError",
+                detail={"retryable": False},
+            )
+        ),
+        model_name="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    )
+
+    with pytest.raises(HTTPException) as exc_info:
+        await serving._run_generation(
+            {"prompt": "test"},
+            OmniDiffusionSamplingParams(),
+            "req-oom",
+        )
+
+    exc = exc_info.value
+    assert exc.status_code == 507
+    assert exc.detail == {
+        "message": "oom",
+        "request_id": "req-oom",
+        "stage_id": 2,
+        "error_type": "OutOfMemoryError",
+        "detail": {"retryable": False},
+    }
@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_get_diffusion_od_config_returns_direct_config():
+    diffusion = object.__new__(AsyncOmniDiffusion)
+    diffusion.od_config = object()
+
+    assert diffusion.get_diffusion_od_config() is diffusion.od_config
@@ -69,6 +69,20 @@ def test_default_stage_config_propagates_ulysses_mode():
    assert parallel_config.ulysses_mode == "advanced_uaa"


+def test_default_stage_config_includes_default_sampling_params():
+    """Ensure default sampling params survive the default diffusion-stage builder."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "default_sampling_params": '{"0": {"generator_device":"cpu", "guidance_scale":7.5}}',
+        }
+    )[0]
+
+    assert stage_cfg["default_sampling_params"] == {
+        "generator_device": "cpu",
+        "guidance_scale": 7.5,
+    }
+
+
 def test_serve_cli_accepts_ulysses_mode():
    """Ensure diffusion serve CLI exposes ulysses_mode and wires it to parallel_config."""
    parser = FlexibleArgumentParser()
@@ -4,6 +4,8 @@ from unittest.mock import MagicMock, patch

 import pytest

+from vllm_omni.diffusion.data import OmniRequestError
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


@@ -204,6 +206,45 @@ class TestOmniBaseProfilerSignatureConsistency:
        # stages should default to None
        assert sig.parameters["stages"].default is None

+
+class TestOmniBaseOutputHandling:
+    def test_handle_output_message_raises_structured_omni_error(self):
+        mock_engine = MagicMock()
+        mock_engine.num_stages = 1
+        mock_engine.default_sampling_params_list = [MagicMock()]
+        mock_engine.get_stage_metadata.return_value = {
+            "final_output_type": "image",
+            "final_output": True,
+        }
+
+        with (
+            patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine),
+            patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x),
+            patch("vllm_omni.entrypoints.omni_base.weakref.finalize"),
+        ):
+            from vllm_omni.entrypoints.omni_base import OmniBase
+
+            omni_base = OmniBase(model="test-model")
+
+        with pytest.raises(OmniRequestError, match="oom") as exc_info:
+            omni_base._handle_output_message(
+                {
+                    "type": "error",
+                    "error": "oom",
+                    "status_code": 507,
+                    "request_id": "req-1",
+                    "stage_id": 2,
+                    "error_type": "OutOfMemoryError",
+                    "detail": {"retryable": False},
+                }
+            )
+
+        assert exc_info.value.status_code == 507
+        assert exc_info.value.request_id == "req-1"
+        assert exc_info.value.stage_id == 2
+        assert exc_info.value.error_type == "OutOfMemoryError"
+        assert exc_info.value.detail == {"retryable": False}
+
    def test_stop_profile_default_values(self):
        """Verify stop_profile has correct default parameter values."""
        import inspect
@@ -10,6 +10,7 @@ from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.sampling_params import RequestOutputKind, SamplingParams

+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.omni import Omni

@@ -379,6 +380,19 @@ def test_openai_serving_models_can_consume_async_omni_compat_attrs():
    assert serving_models.input_processor is input_processor


+def test_get_diffusion_od_config_returns_diffusion_stage_config():
+    diffusion_od_config = object()
+    omni = object.__new__(AsyncOmni)
+    omni.engine = SimpleNamespace(
+        stage_clients=[
+            SimpleNamespace(stage_type="llm"),
+            SimpleNamespace(stage_type="diffusion", _engine=SimpleNamespace(od_config=diffusion_od_config)),
+        ]
+    )
+
+    assert omni.get_diffusion_od_config() is diffusion_od_config
+
+
@pytest.mark.asyncio
 async def test_async_omni_yields_only_final_stage_outputs(monkeypatch: pytest.MonkeyPatch):
    engine = FakeAsyncOmniEngine(
@@ -527,12 +541,16 @@ async def test_async_omni_propagates_engine_error(monkeypatch: pytest.MonkeyPatc

    app = AsyncOmni("dummy-model")
    try:
-        with pytest.raises(RuntimeError, match="engine boom"):
+        with pytest.raises(OmniRequestError, match="engine boom") as exc_info:
            async for _ in app.generate(prompt="hello", request_id="req-1"):
                pass
    finally:
        app.shutdown()

+    assert exc_info.value.request_id == "req-1"
+    assert exc_info.value.stage_id == 0
+    assert exc_info.value.status_code == 500
+

 def test_omni_generate_py_generator_yields_final_outputs_for_each_request(monkeypatch: pytest.MonkeyPatch):
    sampling_params = [SamplingParams(max_tokens=8) for _ in range(3)]
@@ -14,6 +14,7 @@ from vllm_omni.diffusion.data import (
    DiffusionParallelConfig,
    OmniDiffusionConfig,
 )
+from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES


 def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
@@ -106,3 +107,12 @@ class TestCreateDefaultDiffusion:
        ea = stages[0]["engine_args"]
        assert ea["enforce_eager"] is True
        assert ea["lora_path"] == "/tmp/lora"
+
+
+def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
+    od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")
+
+    od_config.update_multimodal_support()
+
+    assert od_config.supports_multimodal_inputs is True
+    assert od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
@@ -36,6 +36,7 @@ class Attention(nn.Module):
        scatter_idx: int = 2,
        gather_idx: int = 1,
        use_sync: bool = False,
+        skip_sequence_parallel: bool = False,
    ):
        super().__init__()
        self.attn_backend = get_attn_backend(-1)
@@ -62,6 +63,7 @@ class Attention(nn.Module):
        self.gather_idx = gather_idx
        self.use_sync = use_sync
        self.causal = causal
+        self.skip_sequence_parallel = skip_sequence_parallel

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


-def enable_hunyuan_image3_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
-    """
-    Enable TeaCache for HunyuanImage3 model.
-
-    HunyuanImage3 uses a GPT-based architecture with KV cache, which is incompatible
-    with the standard hook-based TeaCache approach. Instead, we store the TeaCacheConfig
-    on the pipeline so the denoising loop can implement caching directly.
-    """
-    teacache_config = TeaCacheConfig(
-        transformer_type="HunyuanImage3Pipeline",
-        rel_l1_thresh=config.rel_l1_thresh,
-        coefficients=config.coefficients,
-    )
-    pipeline._tea_cache_config = teacache_config
-
-    logger.info(f"TeaCache enabled for HunyuanImage3 with rel_l1_thresh={teacache_config.rel_l1_thresh}")
-
-
 def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
    """
    Enable TeaCache for Bagel model.
@@ -88,7 +70,6 @@ def enable_flux2_klein_teacache(pipeline: Any, config: DiffusionCacheConfig) ->
 CUSTOM_TEACACHE_ENABLERS = {
    "BagelPipeline": enable_bagel_teacache,
    "Flux2KleinPipeline": enable_flux2_klein_teacache,
-    "HunyuanImage3Pipeline": enable_hunyuan_image3_teacache,
 }


@@ -176,17 +157,6 @@ class TeaCacheBackend(CacheBackend):
                                Currently not used by TeaCache but accepted for interface consistency.
            verbose: Whether to log refresh operations (default: True)
        """
-        # HunyuanImage3: tea cache state is managed inside the denoising loop,
-        # so refresh is a no-op (state is re-initialized every __call__).
-        if (
-            hasattr(pipeline, "_tea_cache_config")
-            and isinstance(pipeline._tea_cache_config, TeaCacheConfig)
-            and pipeline.__class__.__name__ == "HunyuanImage3Pipeline"
-        ):
-            if verbose:
-                logger.debug(f"TeaCache state refreshed for HunyuanImage3 (num_inference_steps={num_inference_steps})")
-            return
-
        # Extract transformer from pipeline
        transformer = pipeline.transformer

@@ -54,16 +54,6 @@ _MODEL_COEFFICIENTS = {
        -12.281286412689623,
        1.0733905006198015,
    ],
-    # HunyuanImage3 pipeline coefficients
-    # Calibrated via polyfit on 3920 data points (80 prompts × 49 steps).
-    # Maps time_embed rel_l1 (range ~0.12-0.54) to output rel_l1 (range ~0.01-0.53).
-    "HunyuanImage3Pipeline": [
-        1.04117826e02,
-        -1.26848482e02,
-        5.68168652e01,
-        -1.04182570e01,
-        6.78098549e-01,
-    ],
 }


@@ -221,7 +221,8 @@ def extract_qwen_context(
    block = module.transformer_blocks[0]
    img_mod_params = block.img_mod(temb)
    img_mod1, _ = img_mod_params.chunk(2, dim=-1)
-    img_modulated, _ = block.img_norm1(hidden_states, img_mod1)
+    img_scale1, img_shift1, _ = block._modulate(img_mod1)
+    img_modulated = block.img_norm1(hidden_states, img_scale1, img_shift1)

    # ============================================================================
    # DEFINE TRANSFORMER EXECUTION (Qwen-specific)
@@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig,
 )

+from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 from vllm_omni.diffusion.utils.network_utils import is_port_available
 from vllm_omni.quantization import build_quant_config

@@ -468,8 +469,10 @@ class OmniDiffusionConfig:
    # Scheduler flow_shift for Wan2.2 (12.0 for 480p, 5.0 for 720p)
    flow_shift: float | None = None

-    # support multi images input
+    # Support multi-image inputs and expose any model-specific request limit
+    # through a generic config field so serving code stays model-agnostic.
    supports_multimodal_inputs: bool = False
+    max_multimodal_image_inputs: int | None = None

    log_level: str = "info"

@@ -615,7 +618,54 @@ class OmniDiffusionConfig:
            raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA")

    def update_multimodal_support(self) -> None:
-        self.supports_multimodal_inputs = self.model_class_name in {"QwenImageEditPlusPipeline"}
+        # Resolve serving-visible multimodal behavior from shared metadata
+        # instead of importing concrete pipeline modules into the config layer.
+        metadata = get_diffusion_model_metadata(self.model_class_name)
+        self.supports_multimodal_inputs = metadata.supports_multimodal_inputs
+        self.max_multimodal_image_inputs = metadata.max_multimodal_image_inputs
+
+    def enrich_config(self) -> None:
+        """Load model metadata from HuggingFace and populate config fields.
+
+        Diffusers-style models expose ``model_index.json`` with ``_class_name``.
+        Non-diffusers models (e.g. Bagel, NextStep) only have ``config.json``,
+        so we fall back to reading that and mapping model_type manually.
+        """
+        from vllm.transformers_utils.config import get_hf_file_to_dict
+
+        try:
+            config_dict = get_hf_file_to_dict("model_index.json", self.model)
+            if config_dict is not None:
+                if self.model_class_name is None:
+                    self.model_class_name = config_dict.get("_class_name", None)
+                self.update_multimodal_support()
+
+                tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model)
+                self.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
+            else:
+                raise FileNotFoundError("model_index.json not found")
+        except (AttributeError, OSError, ValueError, FileNotFoundError):
+            cfg = get_hf_file_to_dict("config.json", self.model)
+            if cfg is None:
+                raise ValueError(f"Could not find config.json or model_index.json for model {self.model}")
+
+            self.tf_model_config = TransformerConfig.from_dict(cfg)
+            model_type = cfg.get("model_type")
+            architectures = cfg.get("architectures") or []
+
+            if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
+                self.model_class_name = "BagelPipeline"
+                self.tf_model_config = TransformerConfig()
+                self.update_multimodal_support()
+            elif model_type == "nextstep":
+                if self.model_class_name is None:
+                    self.model_class_name = "NextStep11Pipeline"
+                self.tf_model_config = TransformerConfig()
+                self.update_multimodal_support()
+            elif architectures and len(architectures) == 1:
+                self.model_class_name = architectures[0]
+            else:
+                raise

    @classmethod
    def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
@@ -656,6 +706,7 @@ class DiffusionOutput:
    output: torch.Tensor | None = None
    trajectory_timesteps: list[torch.Tensor] | None = None
    trajectory_latents: torch.Tensor | None = None
+    trajectory_log_probs: torch.Tensor | None = None
    trajectory_decoded: list[torch.Tensor] | None = None
    error: str | None = None

@@ -675,6 +726,49 @@ class DiffusionOutput:
    peak_memory_mb: float = 0.0


+@dataclass
+class OmniRequestError(RuntimeError):
+    def __init__(
+        self,
+        message: str,
+        *,
+        status_code: int = 500,
+        request_id: str | None = None,
+        stage_id: int | None = None,
+        error_type: str | None = None,
+        detail: dict[str, Any] | None = None,
+    ):
+        super().__init__(message)
+        self.status_code = status_code
+        self.request_id = request_id
+        self.stage_id = stage_id
+        self.error_type = error_type or self.__class__.__name__
+        self.detail = detail or {}
+
+
+def normalize_omni_error(
+    exc: Exception,
+    *,
+    status_code: int = 500,
+    request_id: str | None = None,
+    stage_id: int | None = None,
+) -> OmniRequestError:
+    if isinstance(exc, OmniRequestError):
+        if exc.request_id is None:
+            exc.request_id = request_id
+        if exc.stage_id is None:
+            exc.stage_id = stage_id
+        return exc
+
+    return OmniRequestError(
+        str(exc),
+        status_code=status_code,
+        request_id=request_id,
+        stage_id=stage_id,
+        error_type=type(exc).__name__,
+    )
+
+
 class AttentionBackendEnum(enum.Enum):
    FA = enum.auto()
    SLIDING_TILE_ATTN = enum.auto()
@@ -1,6 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+from __future__ import annotations
+
+import inspect
 import threading
 import time
 from collections.abc import Iterable
@@ -11,7 +14,7 @@ import PIL.Image
 import torch
 from vllm.logger import init_logger

-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, OmniRequestError, normalize_omni_error
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.registry import (
    DiffusionModelRegistry,
@@ -69,6 +72,12 @@ class DiffusionEngine:

        self.post_process_func = get_diffusion_post_process_func(od_config)
        self.pre_process_func = get_diffusion_pre_process_func(od_config)
+        # Cache whether the model-specific postprocess accepts request-level
+        # sampling params so step() can support both legacy and extended hooks.
+        self._post_process_accepts_sampling_params = bool(
+            self.post_process_func is not None
+            and "sampling_params" in inspect.signature(self.post_process_func).parameters
+        )

        executor_class = DiffusionExecutor.get_class(od_config)
        self.executor = executor_class(od_config)
@@ -99,7 +108,12 @@ class DiffusionEngine:
        exec_total_time = time.perf_counter() - exec_start_time

        if output.error:
-            raise Exception(f"{output.error}")
+            # raise Exception(f"{output.error}")
+            raise OmniRequestError(
+                output.error,
+                status_code=500,
+                error_type="DiffusionExecutionError",
+            )
        logger.info("Generation completed successfully.")

        if output.output is None:
@@ -127,10 +141,24 @@ class DiffusionEngine:
            output_data = output_data.cpu()

        postprocess_start_time = time.perf_counter()
-        outputs = self.post_process_func(output_data) if self.post_process_func is not None else output_data
+        if self.post_process_func is not None:
+            # Some video pipelines need request-level controls during
+            # postprocess (for example worker-side frame interpolation).
+            if self._post_process_accepts_sampling_params:
+                outputs = self.post_process_func(output_data, sampling_params=request.sampling_params)
+            else:
+                outputs = self.post_process_func(output_data)
+        else:
+            outputs = output_data
        audio_payload = None
+        custom_output = output.custom_output or {}
+        model_audio_sample_rate = None
+        model_fps = None
        if isinstance(outputs, dict):
            audio_payload = outputs.get("audio")
+            custom_output.update(outputs.get("custom_output") or {})
+            model_audio_sample_rate = outputs.get("audio_sample_rate")
+            model_fps = outputs.get("fps")
            outputs = outputs.get("video", outputs)
        postprocess_time = time.perf_counter() - postprocess_start_time
        logger.info(f"Post-processing completed in {postprocess_time:.4f} seconds")
@@ -186,6 +214,10 @@ class DiffusionEngine:
                mm_output = {}
                if audio_payload is not None:
                    mm_output["audio"] = audio_payload
+                if model_audio_sample_rate is not None:
+                    mm_output["audio_sample_rate"] = model_audio_sample_rate
+                if model_fps is not None:
+                    mm_output["fps"] = model_fps
                return [
                    OmniRequestOutput.from_diffusion(
                        request_id=request_id,
@@ -193,7 +225,11 @@ class DiffusionEngine:
                        prompt=prompt,
                        metrics=metrics,
                        latents=output.trajectory_latents,
-                        custom_output=output.custom_output or {},
+                        trajectory_latents=output.trajectory_latents,
+                        trajectory_timesteps=output.trajectory_timesteps,
+                        trajectory_log_probs=output.trajectory_log_probs,
+                        trajectory_decoded=output.trajectory_decoded,
+                        custom_output=custom_output,
                        multimodal_output=mm_output,
                        stage_durations=output.stage_durations,
                        peak_memory_mb=output.peak_memory_mb,
@@ -244,6 +280,10 @@ class DiffusionEngine:
                                if num_outputs == 1:
                                    sliced_audio = sliced_audio[0]
                        mm_output["audio"] = sliced_audio
+                    if model_audio_sample_rate is not None:
+                        mm_output["audio_sample_rate"] = model_audio_sample_rate
+                    if model_fps is not None:
+                        mm_output["fps"] = model_fps
                    results.append(
                        OmniRequestOutput.from_diffusion(
                            request_id=request_id,
@@ -251,7 +291,11 @@ class DiffusionEngine:
                            prompt=prompt,
                            metrics=metrics,
                            latents=output.trajectory_latents,
-                            custom_output=output.custom_output or {},
+                            trajectory_latents=output.trajectory_latents,
+                            trajectory_timesteps=output.trajectory_timesteps,
+                            trajectory_log_probs=output.trajectory_log_probs,
+                            trajectory_decoded=output.trajectory_decoded,
+                            custom_output=custom_output,
                            multimodal_output=mm_output,
                            stage_durations=output.stage_durations,
                            peak_memory_mb=output.peak_memory_mb,
@@ -264,7 +308,7 @@ class DiffusionEngine:
    def make_engine(
        config: OmniDiffusionConfig,
        scheduler: SchedulerInterface | None = None,
-    ) -> "DiffusionEngine":
+    ) -> DiffusionEngine:
        """Factory method to create a DiffusionEngine instance.

        Args:
@@ -280,33 +324,53 @@ class DiffusionEngine:
            target_sched_req_id = self.scheduler.add_request(request)

            # keep scheduling and executing until the target request is finished
-            while True:
-                sched_output = self.scheduler.schedule()
-                if sched_output.is_empty:
-                    if not self.scheduler.has_requests():
-                        raise RuntimeError("Diffusion scheduler has no runnable requests.")
-                    continue
+            try:
+                while True:
+                    sched_output = self.scheduler.schedule()
+                    if sched_output.is_empty:
+                        if not self.scheduler.has_requests():
+                            # raise RuntimeError("Diffusion scheduler has no runnable requests.")
+                            raise OmniRequestError(
+                                "Diffusion scheduler has no runnable requests.",
+                                status_code=500,
+                                error_type="SchedulerError",
+                            )
+                        continue

-                # NOTE: add_req_and_wait_for_response() is synchronous, and
-                # the scheduler currently enforces _max_batch_size = 1 (see
-                # vllm_omni/diffusion/sched/base_scheduler.py), so we directly
-                # take the single scheduled request here.
-                sched_req_id = sched_output.scheduled_req_ids[0]
-                req = sched_output.scheduled_new_reqs[0].req
+                    # NOTE: add_req_and_wait_for_response() is synchronous, and
+                    # the scheduler currently enforces _max_batch_size = 1 (see
+                    # vllm_omni/diffusion/sched/base_scheduler.py), so we directly
+                    # take the single scheduled request here.
+                    sched_req_id = sched_output.scheduled_req_ids[0]
+                    req = sched_output.scheduled_new_reqs[0].req
+                    try:
+                        output = self.executor.add_req(req)
+                    except Exception as exc:
+                        logger.error(
+                            "Execution failed for diffusion request %s",
+                            sched_req_id,
+                            exc_info=True,
+                        )
+                        # output = DiffusionOutput(error=str(exc))
+                        raise normalize_omni_error(exc) from exc
+
+                    finished_req_ids = self.scheduler.update_from_output(sched_output, output)
+
+                    if output.error:
+                        status_code = 400 if "`max_sequence_length`" in output.error else 500
+                        raise OmniRequestError(
+                            output.error,
+                            status_code=status_code,
+                            error_type="DiffusionExecutionError",
+                        )
+                    if target_sched_req_id in finished_req_ids:
+                        # self.scheduler.pop_request_state(target_sched_req_id)
+                        return output
+            finally:
                try:
-                    output = self.executor.add_req(req)
-                except Exception as exc:
-                    logger.error(
-                        "Execution failed for diffusion request %s",
-                        sched_req_id,
-                        exc_info=True,
-                    )
-                    output = DiffusionOutput(error=str(exc))
-
-                finished_req_ids = self.scheduler.update_from_output(sched_output, output)
-                if target_sched_req_id in finished_req_ids:
                    self.scheduler.pop_request_state(target_sched_req_id)
-                    return output
+                except Exception:
+                    logger.debug("Request state already removed: %s", target_sched_req_id, exc_info=True)

    def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
        """Start or stop torch profiling on all diffusion workers.
@@ -93,7 +93,7 @@ class DistributedAutoencoderKL_base(DistributedVaeMixin):

        _, _, latent_h, latent_w = z.shape
        scale = int(2 ** (len(self.config.block_out_channels) - 1))
-        max_parallel_size = self.distributed_decoder.parallel_size
+        max_parallel_size = self.distributed_executor.parallel_size

        root = int(math.sqrt(max_parallel_size))
        for rows in range(root, 0, -1):
@@ -187,7 +187,7 @@ class DistributedAutoencoderKL_base(DistributedVaeMixin):
        if split is not None:
            strategy = "tile" if split == self.tile_split else "patch"
            logger.info(f"Decode run with distributed executor, split strategy is {strategy}")
-            result = self.distributed_decoder.execute(
+            result = self.distributed_executor.execute(
                z, DistributedOperator(split=split, exec=exec, merge=merge), broadcast_result=False
            )
            if not return_dict:
@@ -108,8 +108,8 @@ class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMi
        if not self.is_distributed_enabled():
            return super().tiled_decode(z, return_dict=return_dict)

-        logger.info("Decode run with distributed executor")
-        result = self.distributed_decoder.execute(
+        logger.debug("Decode running with distributed executor")
+        result = self.distributed_executor.execute(
            z,
            DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
            broadcast_result=True,
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+from contextlib import nullcontext
 from typing import Any

 import torch
@@ -15,11 +16,38 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
    GridSpec,
    TileTask,
 )
+from vllm_omni.platforms import current_omni_platform

 logger = init_logger(__name__)


-class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
+class OmniAutoencoderKLWan(AutoencoderKLWan):
+    def _execution_context(self):
+        try:
+            first_param = next(self.parameters())
+        except StopIteration:
+            return nullcontext()
+
+        dtype = first_param.dtype
+        if dtype not in (torch.float16, torch.bfloat16):
+            return nullcontext()
+
+        return current_omni_platform.create_autocast_context(
+            device_type=first_param.device.type,
+            dtype=dtype,
+            enabled=True,
+        )
+
+    def encode(self, x: torch.Tensor, return_dict: bool = True):
+        with self._execution_context():
+            return super().encode(x, return_dict=return_dict)
+
+    def decode(self, z: torch.Tensor, return_dict: bool = True):
+        with self._execution_context():
+            return super().decode(z, return_dict=return_dict)
+
+
+class DistributedAutoencoderKLWan(OmniAutoencoderKLWan, DistributedVaeMixin):
    @classmethod
    def from_pretrained(cls, *args: Any, **kwargs: Any):
        model = super().from_pretrained(*args, **kwargs)
@@ -84,14 +112,128 @@ class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
        """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, first_chunk=(k == 0))
-            time.append(decoded)
+        with self._execution_context():
+            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, first_chunk=(k == 0))
+                time.append(decoded)
        result = torch.cat(time, dim=2)
        return result

+    def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, num_frames, height, width = x.shape
+        encode_spatial_compression_ratio = self.spatial_compression_ratio
+        # Scale tile parameters for patchified coordinate system
+        tile_sample_min_height = self.tile_sample_min_height
+        tile_sample_min_width = self.tile_sample_min_width
+        tile_sample_stride_height = self.tile_sample_stride_height
+        tile_sample_stride_width = self.tile_sample_stride_width
+        if self.config.patch_size is not None:
+            assert encode_spatial_compression_ratio % self.config.patch_size == 0
+            encode_spatial_compression_ratio = self.spatial_compression_ratio // self.config.patch_size
+            # When input is patchified, scale tile parameters accordingly
+            tile_sample_min_height = tile_sample_min_height // self.config.patch_size
+            tile_sample_min_width = tile_sample_min_width // self.config.patch_size
+            tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
+            tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
+
+        latent_height = height // encode_spatial_compression_ratio
+        latent_width = width // encode_spatial_compression_ratio
+
+        tile_latent_min_height = tile_sample_min_height // encode_spatial_compression_ratio
+        tile_latent_min_width = tile_sample_min_width // encode_spatial_compression_ratio
+        tile_latent_stride_height = tile_sample_stride_height // encode_spatial_compression_ratio
+        tile_latent_stride_width = tile_sample_stride_width // encode_spatial_compression_ratio
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        tiletask_list = []
+        temporal_compression = self.config.scale_factor_temporal
+        for i in range(0, height, tile_sample_stride_height):
+            for j in range(0, width, tile_sample_stride_width):
+                time_list = []
+                frame_range = 1 + (num_frames - 1) // temporal_compression
+                for k in range(frame_range):
+                    if k == 0:
+                        tile = x[:, :, :1, i : i + tile_sample_min_height, j : j + tile_sample_min_width]
+                    else:
+                        tile = x[
+                            :,
+                            :,
+                            1 + temporal_compression * (k - 1) : 1 + temporal_compression * k,
+                            i : i + tile_sample_min_height,
+                            j : j + tile_sample_min_width,
+                        ]
+                    time_list.append(tile)
+                tiletask_list.append(
+                    TileTask(
+                        len(tiletask_list),
+                        (i // tile_sample_stride_height, j // tile_sample_stride_width),
+                        time_list,
+                        workload=time_list[0].shape[3] * time_list[0].shape[4],
+                    )
+                )
+
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec={
+                "latent_height": latent_height,
+                "latent_width": latent_width,
+                "blend_height": blend_height,
+                "blend_width": blend_width,
+                "tile_latent_stride_height": tile_latent_stride_height,
+                "tile_latent_stride_width": tile_latent_stride_width,
+            },
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def encode_tile_exec(self, task: TileTask) -> torch.Tensor:
+        """Encode a single sample tile into latent space."""
+        self.clear_cache()
+        time = []
+        for k, tile in enumerate(task.tensor):
+            self._enc_conv_idx = [0]
+            encoded = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+            encoded = self.quant_conv(encoded)
+            time.append(encoded)
+        result = torch.cat(time, dim=2)
+        self.clear_cache()
+        return result
+
+    def encode_tile_merge(
+        self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec
+    ) -> torch.Tensor:
+        """Merge encoded tiles into a full latent tensor."""
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+                result_row.append(
+                    tile[
+                        :,
+                        :,
+                        :,
+                        : grid_spec.tile_spec["tile_latent_stride_height"],
+                        : grid_spec.tile_spec["tile_latent_stride_width"],
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        enc = torch.cat(result_rows, dim=3)[
+            :, :, :, : grid_spec.tile_spec["latent_height"], : grid_spec.tile_spec["latent_width"]
+        ]
+        return enc
+
    def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
        """Merge decoded tiles into a full image."""
        grid_h, grid_w = grid_spec.grid_shape
@@ -130,8 +272,8 @@ class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
        if not self.is_distributed_enabled():
            return super().tiled_decode(z, return_dict=return_dict)

-        logger.info("Decode run with distributed executor")
-        result = self.distributed_decoder.execute(
+        logger.debug("Decode running with distributed executor")
+        result = self.distributed_executor.execute(
            z,
            DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
            broadcast_result=False,
@@ -140,3 +282,26 @@ class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin):
            return (result,)

        return DecoderOutput(sample=result)
+
+    def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode using distributed VAE executor.
+
+        Note: x is already patchified by parent's _encode() before calling this method.
+        """
+        if not self.is_distributed_enabled():
+            return super().tiled_encode(x)
+
+        logger.debug("Encode running with distributed executor")
+        self.clear_cache()
+        result = self.distributed_executor.execute(
+            x,
+            DistributedOperator(
+                split=self.encode_tile_split,
+                exec=self.encode_tile_exec,
+                merge=self.encode_tile_merge,
+            ),
+            broadcast_result=True,
+        )
+        self.clear_cache()
+        return result
@@ -54,9 +54,9 @@ class DistributedVaeExecutor:
        self.parallel_size = parallel_size

    def gather_tensors(self, tensor: torch.Tensor):
-        gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)] if self.rank == 0 else None
-        dist.gather(tensor, gather_list=gather_list, dst=0, group=self.group)
-        return gather_list
+        gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)]
+        dist.all_gather(gather_list, tensor, group=self.group)
+        return gather_list if self.rank == 0 else None

    def broadcast_tensor(self, tensor: torch.Tensor):
        dist.broadcast(tensor, src=0, group=self.group)
@@ -168,25 +168,25 @@ class DistributedVaeExecutor:

 class DistributedVaeMixin:
    def init_distributed(self):
-        self.distributed_decoder = DistributedVaeExecutor()
+        self.distributed_executor = DistributedVaeExecutor()

-    def set_parallel_size(self, parallel_size: int) -> bool:
-        return self.distributed_decoder.set_parallel_size(parallel_size)
+    def set_parallel_size(self, parallel_size: int) -> None:
+        self.distributed_executor.set_parallel_size(parallel_size)

    def is_distributed_enabled(self) -> bool:
        if (
-            self.distributed_decoder.parallel_size <= 1
+            self.distributed_executor.parallel_size <= 1
            or not dist.is_initialized()
            or not getattr(self, "use_tiling", False)
        ):
            return False
-        world_size = dist.get_world_size(group=self.distributed_decoder.group)
-        pp_size = min(int(self.distributed_decoder.parallel_size), int(world_size))
+        world_size = dist.get_world_size(group=self.distributed_executor.group)
+        pp_size = min(int(self.distributed_executor.parallel_size), int(world_size))
        if pp_size <= 1:
            return False
-        if self.distributed_decoder.parallel_size > pp_size:
+        if self.distributed_executor.parallel_size > pp_size:
            logger.warning(
-                f"vae_patch_parallel_size={self.distributed_decoder.parallel_size} "
+                f"vae_patch_parallel_size={self.distributed_executor.parallel_size} "
                f"is greater than dit_group={world_size};"
                f" using dit_group size={world_size}"
            )
@@ -0,0 +1,362 @@
+"""Inline Stage Diffusion Client for vLLM-Omni multi-stage runtime.
+
+Runs DiffusionEngine in a ThreadPoolExecutor inside the Orchestrator process
+instead of spawning a separate StageDiffusionProc subprocess, eliminating ZMQ
+IPC overhead. Used when there is only a single diffusion stage.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING, Any
+
+import torch
+from PIL import Image
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import normalize_omni_error
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.inputs.data import OmniPromptType
+
+logger = init_logger(__name__)
+
+
+class InlineStageDiffusionClient:
+    """Runs DiffusionEngine in a thread executor inside the Orchestrator."""
+
+    stage_type: str = "diffusion"
+
+    def __init__(
+        self,
+        model: str,
+        od_config: OmniDiffusionConfig,
+        metadata: StageMetadata,
+        batch_size: int = 1,
+    ) -> None:
+        self.model = model
+        self.od_config = od_config
+        self.stage_id = metadata.stage_id
+        self.final_output = metadata.final_output
+        self.final_output_type = metadata.final_output_type
+        self.default_sampling_params = metadata.default_sampling_params
+        self.custom_process_input_func = metadata.custom_process_input_func
+        self.engine_input_source = metadata.engine_input_source
+        self.batch_size = batch_size
+
+        self._enrich_config()
+        self._engine = DiffusionEngine.make_engine(self.od_config)
+        self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="inline-diffusion")
+
+        self._output_queue: asyncio.Queue[OmniRequestOutput | dict[str, Any]] = asyncio.Queue()
+        self._tasks: dict[str, asyncio.Task[Any]] = {}
+        self._shutting_down = False
+
+        logger.info(
+            "[InlineStageDiffusionClient] Stage-%s initialized inline (batch_size=%d)",
+            self.stage_id,
+            self.batch_size,
+        )
+
+    def _enrich_config(self) -> None:
+        """Load model metadata from HuggingFace and populate od_config fields."""
+        self.od_config.enrich_config()
+
+    def _queue_error(self, request_id: str, exc: Exception, log_prefix: str) -> None:
+        err = normalize_omni_error(exc, request_id=request_id, stage_id=self.stage_id)
+        logger.exception(
+            "[InlineStageDiffusionClient] Stage-%s %s req=%s failed: %s",
+            self.stage_id,
+            log_prefix,
+            request_id,
+            err,
+        )
+        self._output_queue.put_nowait(
+            {
+                "type": "error",
+                "request_id": request_id,
+                "stage_id": self.stage_id,
+                "status_code": err.status_code,
+                "error": str(err),
+                "error_type": err.error_type,
+                "detail": err.detail,
+                "finished": True,
+            }
+        )
+
+    # ------------------------------------------------------------------
+    # Request processing
+    # ------------------------------------------------------------------
+
+    async def add_request_async(
+        self,
+        request_id: str,
+        prompt: OmniPromptType,
+        sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[int, dict[str, Any]] | None = None,
+    ) -> None:
+        task = asyncio.create_task(
+            self._dispatch_request(
+                request_id,
+                prompt,
+                sampling_params,
+                kv_sender_info,
+            )
+        )
+        self._tasks[request_id] = task
+
+    async def _dispatch_request(
+        self,
+        request_id: str,
+        prompt: Any,
+        sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[str, Any] | None = None,
+    ) -> None:
+        try:
+            # Inline diffusion uses the same request envelope as AsyncOmniDiffusion.
+            # kv_sender_info is accepted at the interface layer for compatibility
+            # with other stage clients but is not consumed by DiffusionEngine here.
+            _ = kv_sender_info
+            request = OmniDiffusionRequest(
+                prompts=[prompt],
+                sampling_params=sampling_params,
+                request_ids=[request_id],
+            )
+
+            loop = asyncio.get_running_loop()
+            results = await loop.run_in_executor(self._executor, self._engine.step, request)
+            result = results[0]
+            if not result.request_id:
+                result.request_id = request_id
+
+            self._output_queue.put_nowait(result)
+        except asyncio.CancelledError:
+            logger.info("request_id: %s cancelled", request_id)
+        except Exception as e:
+            self._queue_error(request_id, e, "request")
+        finally:
+            self._tasks.pop(request_id, None)
+
+    async def add_batch_request_async(
+        self,
+        request_id: str,
+        prompts: list[OmniPromptType],
+        sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[int, dict[str, Any]] | None = None,
+    ) -> None:
+        task = asyncio.create_task(
+            self._dispatch_batch(
+                request_id,
+                prompts,
+                sampling_params,
+                kv_sender_info,
+            )
+        )
+        self._tasks[request_id] = task
+
+    async def _dispatch_batch(
+        self,
+        request_id: str,
+        prompts: list[Any],
+        sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[str, Any] | None = None,
+    ) -> None:
+        try:
+            _ = kv_sender_info
+            request = OmniDiffusionRequest(
+                prompts=prompts,
+                sampling_params=sampling_params,
+                request_ids=[f"{request_id}-{i}" for i in range(len(prompts))],
+            )
+
+            loop = asyncio.get_running_loop()
+            results = await loop.run_in_executor(self._executor, self._engine.step, request)
+
+            all_images: list = []
+            merged_mm: dict[str, Any] = {}
+            merged_metrics: dict[str, Any] = {}
+            merged_durations: dict[str, float] = {}
+            merged_custom: dict[str, Any] = {}
+            peak_mem = 0.0
+            latents = None
+            trajectory_latents: torch.Tensor | None = None
+            trajectory_timesteps: list[torch.Tensor] | None = None
+            trajectory_log_probs: torch.Tensor | None = None
+            trajectory_decoded: list[torch.Tensor] | list[Image.Image] | None = None
+            final_output_type = "image"
+
+            for r in results:
+                all_images.extend(r.images)
+                merged_mm.update(r._multimodal_output)
+                merged_metrics.update(r.metrics)
+                merged_durations.update(r.stage_durations)
+                merged_custom.update(r._custom_output)
+                peak_mem = max(peak_mem, r.peak_memory_mb)
+                if latents is None and r.latents is not None:
+                    latents = r.latents
+                if trajectory_latents is None:
+                    trajectory_latents = r.trajectory_latents
+                if trajectory_timesteps is None:
+                    trajectory_timesteps = r.trajectory_timesteps
+                if trajectory_log_probs is None:
+                    trajectory_log_probs = r.trajectory_log_probs
+                if trajectory_decoded is None:
+                    trajectory_decoded = r.trajectory_decoded
+                if r.final_output_type != "image":
+                    final_output_type = r.final_output_type
+
+            result = OmniRequestOutput.from_diffusion(
+                request_id=request_id,
+                images=all_images,
+                prompt=prompts[0] if len(prompts) == 1 else None,
+                metrics=merged_metrics,
+                latents=latents,
+                trajectory_latents=trajectory_latents,
+                trajectory_timesteps=trajectory_timesteps,
+                trajectory_log_probs=trajectory_log_probs,
+                trajectory_decoded=trajectory_decoded,
+                custom_output=merged_custom or None,
+                multimodal_output=merged_mm or None,
+                final_output_type=final_output_type,
+                stage_durations=merged_durations,
+                peak_memory_mb=peak_mem,
+            )
+
+            self._output_queue.put_nowait(result)
+        except asyncio.CancelledError:
+            logger.info("request_id: %s cancelled", request_id)
+        except Exception as e:
+            self._queue_error(request_id, e, "batch request")
+        finally:
+            self._tasks.pop(request_id, None)
+
+    def get_diffusion_output_async(self) -> OmniRequestOutput | dict[str, Any] | None:
+        try:
+            return self._output_queue.get_nowait()
+        except asyncio.QueueEmpty:
+            return None
+
+    def get_diffusion_output_nowait(self) -> OmniRequestOutput | dict[str, Any] | None:
+        return self.get_diffusion_output_async()
+
+    async def abort_requests_async(self, request_ids: list[str]) -> None:
+        for rid in request_ids:
+            task = self._tasks.pop(rid, None)
+            if task:
+                task.cancel()
+            self._engine.abort(rid)
+
+    async def collective_rpc_async(
+        self,
+        method: str,
+        timeout: float | None = None,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> Any:
+        loop = asyncio.get_running_loop()
+
+        if method == "profile":
+            is_start = args[0] if args else True
+            profile_prefix = args[1] if len(args) > 1 else None
+            if is_start and profile_prefix is None:
+                profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}"
+            return await loop.run_in_executor(
+                self._executor,
+                self._engine.profile,
+                is_start,
+                profile_prefix,
+            )
+
+        kwargs = kwargs or {}
+
+        # LoRA methods
+        if method == "add_lora":
+            lora_request = args[0] if args else kwargs.get("lora_request")
+            results = await loop.run_in_executor(
+                self._executor,
+                self._engine.collective_rpc,
+                "add_lora",
+                timeout,
+                (),
+                {"lora_request": lora_request},
+                None,
+            )
+            return all(results) if isinstance(results, list) else results
+
+        if method == "remove_lora":
+            results = await loop.run_in_executor(
+                self._executor,
+                self._engine.collective_rpc,
+                "remove_lora",
+                timeout,
+                args,
+                kwargs,
+                None,
+            )
+            return all(results) if isinstance(results, list) else results
+
+        if method == "list_loras":
+            results = await loop.run_in_executor(
+                self._executor,
+                self._engine.collective_rpc,
+                "list_loras",
+                timeout,
+                (),
+                {},
+                None,
+            )
+            if not isinstance(results, list):
+                return results or []
+            merged: set[int] = set()
+            for part in results:
+                merged.update(part or [])
+            return sorted(merged)
+
+        if method == "pin_lora":
+            lora_id = args[0] if args else kwargs.get("adapter_id")
+            results = await loop.run_in_executor(
+                self._executor,
+                self._engine.collective_rpc,
+                "pin_lora",
+                timeout,
+                (),
+                {"adapter_id": lora_id},
+                None,
+            )
+            return all(results) if isinstance(results, list) else results
+
+        return await loop.run_in_executor(
+            self._executor,
+            self._engine.collective_rpc,
+            method,
+            timeout,
+            args,
+            kwargs,
+            None,
+        )
+
+    def shutdown(self) -> None:
+        self._shutting_down = True
+
+        # Cancel all pending tasks
+        for task in self._tasks.values():
+            task.cancel()
+
+        try:
+            # Cancel queued futures and wait for the running one to complete deterministically
+            self._executor.shutdown(wait=True, cancel_futures=True)
+        except Exception:
+            pass
+
+        try:
+            self._engine.close()
+        except Exception:
+            pass
@@ -31,7 +31,13 @@ def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]:
    import numpy as np

    tensor = tensor.detach().cpu().contiguous()
-    arr = tensor.numpy()
+    if tensor.dtype == torch.bfloat16:
+        # NumPy cannot materialize BF16 tensors directly, so store the raw
+        # 16-bit payload and reconstruct the original dtype on unpack.
+        tensor_for_storage = tensor.view(torch.uint16)
+        arr = tensor_for_storage.numpy()
+    else:
+        arr = tensor.numpy()
    nbytes = arr.nbytes
    shm = shared_memory.SharedMemory(create=True, size=nbytes)
    shm_arr = np.ndarray(arr.shape, dtype=arr.dtype, buffer=shm.buf[:nbytes])
@@ -59,6 +65,8 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
        np_dtype = np.dtype(handle["numpy_dtype"])
        arr = np.ndarray(handle["shape"], dtype=np_dtype, buffer=shm.buf[: handle["nbytes"]])
        tensor = torch.from_numpy(arr.copy())
+        if handle["torch_dtype"] == str(torch.bfloat16):
+            tensor = tensor.view(torch.bfloat16)
    finally:
        shm.close()
        shm.unlink()
@@ -72,6 +80,9 @@ def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
    if output.trajectory_latents is not None and isinstance(output.trajectory_latents, torch.Tensor):
        if output.trajectory_latents.nelement() * output.trajectory_latents.element_size() > _SHM_TENSOR_THRESHOLD:
            output.trajectory_latents = _tensor_to_shm(output.trajectory_latents)
+    if output.trajectory_log_probs is not None and isinstance(output.trajectory_log_probs, torch.Tensor):
+        if output.trajectory_log_probs.nelement() * output.trajectory_log_probs.element_size() > _SHM_TENSOR_THRESHOLD:
+            output.trajectory_log_probs = _tensor_to_shm(output.trajectory_log_probs)
    return output


@@ -95,6 +106,8 @@ def _unpack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
        output.output = _tensor_from_shm(output.output)
    if isinstance(output.trajectory_latents, dict) and output.trajectory_latents.get("__tensor_shm__"):
        output.trajectory_latents = _tensor_from_shm(output.trajectory_latents)
+    if isinstance(output.trajectory_log_probs, dict) and output.trajectory_log_probs.get("__tensor_shm__"):
+        output.trajectory_log_probs = _tensor_from_shm(output.trajectory_log_probs)
    return output


@@ -1,10 +1,10 @@
 from importlib.util import find_spec

 import torch
-import torch.nn as nn
 from vllm.logger import init_logger

 from vllm_omni.diffusion.layers.custom_op import CustomOp
+from vllm_omni.diffusion.layers.norm import LayerNorm

 logger = init_logger(__name__)

@@ -22,104 +22,60 @@ class AdaLayerNorm(CustomOp):
        self.eps = eps
        self.elementwise_affine = elementwise_affine
        self.hidden_size = hidden_size
-        self.layernorm = nn.LayerNorm(self.hidden_size, elementwise_affine=self.elementwise_affine, eps=self.eps)
-
-    def preprocess(
-        self,
-        mod_params: torch.Tensor,
-        index: torch.Tensor = None,
-    ) -> torch.Tensor:
-        # shift: b d, scale: b d, gate: b d
-        shift, scale, gate = mod_params.chunk(3, dim=-1)
-
-        if index is not None:
-            # Assuming mod_params batch dim is 2*actual_batch (chunked into 2 parts)
-            # So shift, scale, gate have shape [2*actual_batch, d]
-            actual_batch = shift.size(0) // 2
-            shift_0, shift_1 = shift[:actual_batch], shift[actual_batch:]  # each: [actual_batch, d]
-            scale_0, scale_1 = scale[:actual_batch], scale[actual_batch:]
-            gate_0, gate_1 = gate[:actual_batch], gate[actual_batch:]
-
-            # index: [b, l] where b is actual batch size
-            # Expand to [b, l, 1] to match feature dimension
-            index_expanded = index.unsqueeze(-1)  # [b, l, 1]
-
-            # Expand chunks to [b, 1, d] then broadcast to [b, l, d]
-            shift_0_exp = shift_0.unsqueeze(1)  # [b, 1, d]
-            shift_1_exp = shift_1.unsqueeze(1)  # [b, 1, d]
-            scale_0_exp = scale_0.unsqueeze(1)
-            scale_1_exp = scale_1.unsqueeze(1)
-            gate_0_exp = gate_0.unsqueeze(1)
-            gate_1_exp = gate_1.unsqueeze(1)
-
-            # Use torch.where to select based on index
-            shift_result = torch.where(index_expanded == 0, shift_0_exp, shift_1_exp)
-            scale_result = torch.where(index_expanded == 0, scale_0_exp, scale_1_exp)
-            gate_result = torch.where(index_expanded == 0, gate_0_exp, gate_1_exp)
-        else:
-            shift_result = shift.unsqueeze(1)
-            scale_result = scale.unsqueeze(1)
-            gate_result = gate.unsqueeze(1)
-
-        return shift_result, scale_result, gate_result
+        self.layernorm = LayerNorm(self.hidden_size, elementwise_affine=self.elementwise_affine, eps=self.eps)

    def forward_cuda(
        self,
        x: torch.Tensor,
-        mod_params: torch.Tensor,
-        index: torch.Tensor = None,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
    ) -> torch.Tensor:
-        return self.forward_native(x, mod_params, index)
+        return self.forward_native(x, scale, shift)

    def forward_hip(
        self,
        x: torch.Tensor,
-        mod_params: torch.Tensor,
-        index: torch.Tensor = None,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
    ) -> torch.Tensor:
-        return self.forward_native(x, mod_params, index)
+        return self.forward_native(x, scale, shift)

    def forward_npu(
        self,
        x: torch.Tensor,
-        mod_params: torch.Tensor,
-        index: torch.Tensor = None,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
    ) -> torch.Tensor:
-        shift_result, scale_result, gate_result = self.preprocess(mod_params, index)
-
        if _HAS_MINDIESD:
            try:
                from mindiesd import layernorm_scale_shift

-                output = layernorm_scale_shift(self.layernorm, x, scale_result, shift_result, fused=True)
+                output = layernorm_scale_shift(self.layernorm, x, scale, shift, fused=True)

-                return output, gate_result
+                return output
            except ImportError as e:
                logger.warning_once(f"mindiesd import failed, falling back to torch_npu: {e}")

        import torch_npu

        output = (
-            torch_npu.npu_layer_norm_eval(x, normalized_shape=[self.hidden_size], eps=self.eps) * (1 + scale_result)
-            + shift_result
+            torch_npu.npu_layer_norm_eval(x, normalized_shape=[self.hidden_size], eps=self.eps) * (1 + scale) + shift
        )

-        return output, gate_result
+        return output

    def forward_xpu(
        self,
        x: torch.Tensor,
-        mod_params: torch.Tensor,
-        index: torch.Tensor = None,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
    ) -> torch.Tensor:
-        return self.forward_native(x, mod_params, index)
+        return self.forward_native(x, scale, shift)

    def forward_native(
        self,
        x: torch.Tensor,
-        mod_params: torch.Tensor,
-        index: torch.Tensor = None,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
    ) -> torch.Tensor:
-        shift_result, scale_result, gate_result = self.preprocess(mod_params, index)
-
-        return self.layernorm(x) * (1 + scale_result) + shift_result, gate_result
+        return self.layernorm(x) * (1 + scale) + shift
@@ -0,0 +1,185 @@
+from importlib.util import find_spec
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.layers.custom_op import CustomOp
+
+logger = init_logger(__name__)
+
+_HAS_MINDIESD = find_spec("mindiesd") is not None
+
+
+class LayerNorm(nn.LayerNorm, CustomOp):
+    """
+    LayerNorm implementation that inherits from both ``nn.LayerNorm`` and ``CustomOp``.
+    NPU:
+        Uses ``mindiesd.fast_layernorm(self, x)`` when MindIE-SD is installed.
+    CUDA / HIP / XPU / native:
+        Falls back to FP32 nn.LayerNorm implementation.
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True):
+        super().__init__(normalized_shape=dim, eps=eps, elementwise_affine=elementwise_affine)
+        self._forward_method = CustomOp.dispatch_forward(self)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self._forward_method(x)
+
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_hip(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_npu(self, x: torch.Tensor) -> torch.Tensor:
+        if _HAS_MINDIESD:
+            try:
+                from mindiesd import fast_layernorm
+
+                return fast_layernorm(self, x)
+            except ImportError as e:
+                logger.warning_once(
+                    "mindiesd.fast_layernorm import failed, falling back to FP32 layer_norm: %s",
+                    e,
+                )
+
+        return self.forward_native(x)
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        origin_dtype = x.dtype
+        return F.layer_norm(
+            x.float(),
+            self.normalized_shape,
+            self.weight.float() if self.weight is not None else None,
+            self.bias.float() if self.bias is not None else None,
+            self.eps,
+        ).to(origin_dtype)
+
+
+class RMSNorm(CustomOp):
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward_cuda(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_hip(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_npu(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        import torch_npu
+
+        output = torch_npu.npu_rms_norm(x, gamma=self.weight, epsilon=self.variance_epsilon)[0]
+
+        return output
+
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        input_dtype = x.dtype
+        x = x.to(torch.float32)
+        variance = x.pow(2).mean(-1, keepdim=True)
+        out = x * torch.rsqrt(variance + self.variance_epsilon)
+        out = self.weight.to(torch.float32) * out
+        return out.to(input_dtype)
+
+class RMSNormVAE(CustomOp):
+    """Root Mean Square Layer Normalization for Channel-First or Last"""
+
+    def __init__(
+        self,
+        dim: int,
+        channel_first: bool = True,
+        images: bool = True,
+        bias: bool = False,
+        epsilon: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else None
+        self.epsilon = epsilon
+
+        self.gamma_rmsnorm = None
+
+    def forward_cuda(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_hip(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_npu(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        import torch_npu
+
+        if self.gamma_rmsnorm is None:
+            self.gamma_rmsnorm = self.gamma.reshape(-1)
+
+        if self.channel_first:
+            x = x.transpose(1, -1)
+            out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0].transpose(1, -1)
+        else:
+            out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0]
+
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        out = (
+            F.normalize(
+                x,
+                dim=(1 if self.channel_first else -1),
+                eps=self.epsilon,
+            )
+            * self.scale
+            * self.gamma
+        )
+        if self.bias is not None:
+            out = out + self.bias
+        return out
@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class DiffusionModelMetadata:
+    # Keep serving-facing capability metadata in a lightweight shared module so
+    # config/model plumbing can read it without importing concrete pipelines.
+    supports_multimodal_inputs: bool = False
+    max_multimodal_image_inputs: int | None = None
+
+
+QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4
+
+
+_DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
+    "QwenImageEditPlusPipeline": DiffusionModelMetadata(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
+    ),
+}
+
+
+def get_diffusion_model_metadata(model_class_name: str | None) -> DiffusionModelMetadata:
+    # Unknown models fall back to "no special multimodal capabilities" so new
+    # pipelines do not accidentally inherit limits meant for other models.
+    if model_class_name is None:
+        return DiffusionModelMetadata()
+    return _DIFFUSION_MODEL_METADATA.get(model_class_name, DiffusionModelMetadata())
@@ -196,8 +196,6 @@ class FluxAttention(torch.nn.Module):
        attention_mask: torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        # Ensure contiguous for FP8 quantized linear layers
-        hidden_states = hidden_states.contiguous()
        qkv, _ = self.to_qkv(hidden_states)
        q_size = self.to_qkv.num_heads * self.head_dim
        kv_size = self.to_qkv.num_kv_heads * self.head_dim
@@ -211,7 +209,6 @@ class FluxAttention(torch.nn.Module):
        key = self.norm_k(key)

        if self.added_kv_proj_dim is not None:
-            encoder_hidden_states = encoder_hidden_states.contiguous()
            encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states)
            add_q_size = self.add_kv_proj.num_heads * self.head_dim
            add_kv_size = self.add_kv_proj.num_kv_heads * self.head_dim
@@ -723,7 +720,6 @@ class FluxKontextTransformer2DModel(FluxTransformer2DModel):
        guidance_embeds: bool = True,
        axes_dims_rope: tuple[int, int, int] = (16, 56, 56),
        theta: float = 10000.0,
-        quant_config: "QuantizationConfig | None" = None,
    ):
        super().__init__(
            od_config=od_config,
@@ -739,7 +735,6 @@ class FluxKontextTransformer2DModel(FluxTransformer2DModel):
            guidance_embeds=guidance_embeds,
            axes_dims_rope=axes_dims_rope,
            theta=theta,
-            quant_config=quant_config,
        )

    def forward(
@@ -130,8 +130,6 @@ class FluxKontextPipeline(nn.Module, FluxPipelineMixin, SupportImageInput):
        ).to(self._execution_device)

        transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, FluxKontextTransformer2DModel)
-        transformer_kwargs["od_config"] = od_config
-        transformer_kwargs["quant_config"] = od_config.quantization_config
        self.transformer = FluxKontextTransformer2DModel(**transformer_kwargs)

        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
@@ -5,7 +5,7 @@
 import math
 from collections.abc import Iterable
 from functools import lru_cache
-from typing import TYPE_CHECKING, Any
+from typing import Any

 import torch
 import torch.nn as nn
@@ -29,11 +29,6 @@ from vllm_omni.diffusion.distributed.sp_plan import (
    SequenceParallelOutput,
 )

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


@@ -101,15 +96,7 @@ class DistributedRMSNorm(nn.Module):
 class ColumnParallelGELU(nn.Module):
    """Column parallel linear with GELU activation."""

-    def __init__(
-        self,
-        dim_in: int,
-        dim_out: int,
-        *,
-        approximate: str = "tanh",
-        bias: bool = True,
-        quant_config: "QuantizationConfig | None" = None,
-    ):
+    def __init__(self, dim_in: int, dim_out: int, *, approximate: str = "tanh", bias: bool = True):
        super().__init__()
        self.proj = ColumnParallelLinear(
            dim_in,
@@ -117,7 +104,6 @@ class ColumnParallelGELU(nn.Module):
            bias=bias,
            gather_output=False,
            return_bias=False,
-            quant_config=quant_config,
        )
        self.approximate = approximate

@@ -135,15 +121,18 @@ class HeliosFeedForward(nn.Module):
        inner_dim: int,
        dim_out: int | None = None,
        bias: bool = True,
-        quant_config: "QuantizationConfig | None" = None,
    ) -> None:
        super().__init__()
        dim_out = dim_out or dim

-        self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias, quant_config=quant_config)
+        self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias)
        self.net_1 = nn.Identity()
        self.net_2 = RowParallelLinear(
-            inner_dim, dim_out, bias=bias, input_is_parallel=True, return_bias=False, quant_config=quant_config
+            inner_dim,
+            dim_out,
+            bias=bias,
+            input_is_parallel=True,
+            return_bias=False,
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -273,7 +262,6 @@ class HeliosSelfAttention(nn.Module):
        dropout: float = 0.0,
        is_amplify_history: bool = False,
        history_scale_mode: str = "per_head",
-        quant_config: "QuantizationConfig | None" = None,
    ):
        super().__init__()

@@ -287,7 +275,6 @@ class HeliosSelfAttention(nn.Module):
            head_size=head_dim,
            total_num_heads=num_heads,
            bias=True,
-            quant_config=quant_config,
        )

        self.num_heads = self.to_qkv.num_heads
@@ -303,7 +290,6 @@ class HeliosSelfAttention(nn.Module):
            bias=True,
            input_is_parallel=True,
            return_bias=False,
-            quant_config=quant_config,
        )
        self.dropout = nn.Dropout(dropout)

@@ -382,7 +368,6 @@ class HeliosCrossAttention(nn.Module):
        head_dim: int,
        eps: float = 1e-5,
        dropout: float = 0.0,
-        quant_config: "QuantizationConfig | None" = None,
    ):
        super().__init__()

@@ -397,7 +382,6 @@ class HeliosCrossAttention(nn.Module):
            bias=True,
            gather_output=False,
            return_bias=False,
-            quant_config=quant_config,
        )
        self.to_k = ColumnParallelLinear(
            dim,
@@ -405,7 +389,6 @@ class HeliosCrossAttention(nn.Module):
            bias=True,
            gather_output=False,
            return_bias=False,
-            quant_config=quant_config,
        )
        self.to_v = ColumnParallelLinear(
            dim,
@@ -413,7 +396,6 @@ class HeliosCrossAttention(nn.Module):
            bias=True,
            gather_output=False,
            return_bias=False,
-            quant_config=quant_config,
        )

        tp_size = get_tensor_model_parallel_world_size()
@@ -429,7 +411,6 @@ class HeliosCrossAttention(nn.Module):
            bias=True,
            input_is_parallel=True,
            return_bias=False,
-            quant_config=quant_config,
        )
        self.dropout = nn.Dropout(dropout)

@@ -480,7 +461,6 @@ class HeliosTransformerBlock(nn.Module):
        guidance_cross_attn: bool = False,
        is_amplify_history: bool = False,
        history_scale_mode: str = "per_head",
-        quant_config: "QuantizationConfig | None" = None,
    ):
        super().__init__()

@@ -495,17 +475,19 @@ class HeliosTransformerBlock(nn.Module):
            eps=eps,
            is_amplify_history=is_amplify_history,
            history_scale_mode=history_scale_mode,
-            quant_config=quant_config,
        )

        # 2. Cross-attention
        self.attn2 = HeliosCrossAttention(
-            dim=dim, num_heads=num_heads, head_dim=head_dim, eps=eps, quant_config=quant_config
+            dim=dim,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            eps=eps,
        )
        self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()

        # 3. Feed-forward
-        self.ffn = HeliosFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim, quant_config=quant_config)
+        self.ffn = HeliosFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim)
        self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)

        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -620,7 +602,6 @@ class HeliosTransformer3DModel(nn.Module):
        has_multi_term_memory_patch: bool = True,
        is_amplify_history: bool = False,
        history_scale_mode: str = "per_head",
-        quant_config: "QuantizationConfig | None" = None,
    ) -> None:
        super().__init__()

@@ -708,7 +689,6 @@ class HeliosTransformer3DModel(nn.Module):
                    guidance_cross_attn=guidance_cross_attn,
                    is_amplify_history=is_amplify_history,
                    history_scale_mode=history_scale_mode,
-                    quant_config=quant_config,
                )
                for _ in range(num_layers)
            ]
@@ -9,7 +9,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, Any
+from typing import Any

 import numpy as np
 import torch
@@ -31,11 +31,6 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.platforms import current_omni_platform

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


@@ -85,9 +80,7 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc
    return load_json_config(model_path, subfolder, "config.json", local_files_only)


-def create_transformer_from_config(
-    config: dict, quant_config: QuantizationConfig | None = None
-) -> HeliosTransformer3DModel:
+def create_transformer_from_config(config: dict) -> HeliosTransformer3DModel:
    kwargs = {}

    key_map = [
@@ -119,7 +112,7 @@ def create_transformer_from_config(
                val = tuple(val)
            kwargs[key] = val

-    return HeliosTransformer3DModel(quant_config=quant_config, **kwargs)
+    return HeliosTransformer3DModel(**kwargs)


 def get_helios_post_process_func(
@@ -199,9 +192,7 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip
        ).to(self.device)

        transformer_config = load_transformer_config(model, "transformer", local_files_only)
-        self.transformer = create_transformer_from_config(
-            transformer_config, quant_config=od_config.quantization_config
-        )
+        self.transformer = create_transformer_from_config(transformer_config)

        # Read scheduler config to determine scheduler type
        sched_cfg = load_json_config(model, "scheduler", "scheduler_config.json", local_files_only)
@@ -3,28 +3,12 @@

 from typing import Any

-import vllm.forward_context as _vllm_fc
 from vllm.model_executor.layers.fused_moe import SharedFusedMoE
 from vllm.utils.import_utils import resolve_obj_by_qualname

 from vllm_omni.platforms import current_omni_platform


-def _set_forward_context_num_tokens(num_tokens: int) -> None:
-    """Set num_tokens on the vLLM ForwardContext for MoE routing.
-
-    After the rebase to vLLM 0.18.0, SharedFusedMoE expects
-    ForwardContext.num_tokens to be set. Without it, MoE expert
-    routing may produce incorrect results (silent correctness bug).
-    """
-    if not _vllm_fc.is_forward_context_available():
-        return
-    forward_context = _vllm_fc.get_forward_context()
-    forward_context.num_tokens = num_tokens
-    if not hasattr(forward_context, "in_profile_run"):
-        forward_context.in_profile_run = False
-
-
 class HunyuanFusedMoEDefault(SharedFusedMoE):
    def __init__(self, *, prefix: str = "", **kwargs: Any) -> None:
        super().__init__(prefix=prefix, **kwargs)
@@ -37,7 +21,6 @@ class HunyuanFusedMoEDefault(SharedFusedMoE):
        self._init_hook_handle.remove()

    def forward(self, hidden_states: Any, router_logits: Any) -> Any:
-        _set_forward_context_num_tokens(hidden_states.shape[0])
        return super().forward(hidden_states, router_logits)


@@ -2667,15 +2667,6 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
        self._num_timesteps = len(timesteps)

-        # ---- TeaCache initialization ----
-        tea_cache_config = getattr(self.model, "_tea_cache_config", None)
-        if tea_cache_config is not None:
-            tc_rescale = np.poly1d(tea_cache_config.coefficients)
-            tc_acc_dist = 0.0
-            tc_prev_mod = None
-            tc_prev_pred = None
-            tc_cnt = 0
-
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                # expand the latents if we are doing classifier free guidance
@@ -2684,42 +2675,17 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):

                t_expand = t.repeat(latent_model_input.shape[0])

-                # ---- TeaCache: decide whether to compute or reuse ----
-                should_compute = True
-                if tea_cache_config is not None:
-                    with torch.no_grad():
-                        # Use timestep embedding as the modulated input for cache decision
-                        cur_mod = self.model.time_embed(t.unsqueeze(0))
-                    if tc_cnt > 0 and tc_prev_mod is not None:
-                        rel_dist = (
-                            ((cur_mod - tc_prev_mod).abs().mean() / (tc_prev_mod.abs().mean() + 1e-8)).cpu().item()
-                        )
-                        rescaled = float(tc_rescale(rel_dist))
-                        tc_acc_dist += abs(rescaled)
-                        if tc_acc_dist < tea_cache_config.rel_l1_thresh:
-                            should_compute = False
-                        else:
-                            tc_acc_dist = 0.0
-                    tc_prev_mod = cur_mod.detach()
+                model_inputs = self.model.prepare_inputs_for_generation(
+                    input_ids,
+                    images=latent_model_input,
+                    timestep=t_expand,
+                    **model_kwargs,
+                )

-                if should_compute:
-                    model_inputs = self.model.prepare_inputs_for_generation(
-                        input_ids,
-                        images=latent_model_input,
-                        timestep=t_expand,
-                        **model_kwargs,
-                    )
-
-                    with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
-                        model_output = self.model.forward_call(**model_inputs, first_step=(i == 0))
-                        pred = model_output["diffusion_prediction"]
-                    pred = pred.to(dtype=torch.float32)
-
-                    if tea_cache_config is not None:
-                        tc_prev_pred = pred.clone()
-                else:
-                    # TeaCache fast path: reuse previous prediction
-                    pred = tc_prev_pred
+                with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
+                    model_output = self.model.forward_call(**model_inputs, first_step=(i == 0))
+                    pred = model_output["diffusion_prediction"]
+                pred = pred.to(dtype=torch.float32)

                # perform guidance
                if self.do_classifier_free_guidance:
@@ -2729,7 +2695,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                # compute the previous noisy sample x_t -> x_t-1
                latents = self.scheduler.step(pred, t, latents, **_scheduler_step_extra_kwargs, return_dict=False)[0]

-                if i != len(timesteps) - 1 and should_compute:
+                if i != len(timesteps) - 1:
                    model_kwargs = self.model._update_model_kwargs_for_generation(  # noqa
                        model_output,
                        model_kwargs,
@@ -2743,9 +2709,6 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                    model_kwargs["query_lens"] = query_lens
                    model_kwargs["seq_lens"] = seq_lens

-                if tea_cache_config is not None:
-                    tc_cnt += 1
-
                if callback_on_step_end is not None:
                    callback_kwargs = {}
                    for k in callback_on_step_end_tensor_inputs:
@@ -34,6 +34,12 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
 )
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.prompt_utils import (
+    validate_prompt_sequence_lengths,
+)
+from vllm_omni.diffusion.utils.size_utils import (
+    normalize_min_aligned_size,
+)
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs

 if TYPE_CHECKING:
@@ -360,8 +366,10 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
                "that was used to generate `negative_prompt_embeds`."
            )

-        if max_sequence_length is not None and max_sequence_length > 1024:
-            raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
+        if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length:
+            raise ValueError(
+                f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}"
+            )

    def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
        bool_mask = mask.bool()
@@ -375,6 +383,8 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
        self,
        prompt: str | list[str] = None,
        dtype: torch.dtype | None = None,
+        max_sequence_length: int | None = None,
+        prompt_name: str = "prompt",
    ):
        dtype = dtype or self.text_encoder.dtype

@@ -385,12 +395,27 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
        txt = [template.format(e) for e in prompt]
        txt_tokens = self.tokenizer(
            txt,
-            max_length=self.tokenizer_max_length + drop_idx,
            padding=True,
-            truncation=True,
+            truncation=False,
            return_tensors="pt",
        ).to(self.device)
-        # print(f"attention mask: {txt_tokens.attention_mask}")
+        # Validate only the user prompt contribution. The Qwen template also
+        # adds a fixed suffix after the user text, so subtracting only
+        # prompt_template_encode_start_idx would overcount near-limit prompts.
+        template_tokens = self.tokenizer(
+            [template.format("")],
+            padding=True,
+            truncation=False,
+            return_tensors="pt",
+        ).to(self.device)
+        validate_prompt_sequence_lengths(
+            txt_tokens.attention_mask,
+            max_sequence_length=max_sequence_length or self.tokenizer_max_length,
+            supported_max_sequence_length=self.tokenizer_max_length,
+            prompt_name=prompt_name,
+            baseline_attention_mask=template_tokens.attention_mask,
+            error_context="after applying the Qwen prompt template",
+        )
        encoder_hidden_states = self.text_encoder(
            input_ids=txt_tokens.input_ids,
            attention_mask=txt_tokens.attention_mask,
@@ -419,6 +444,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
        prompt_embeds: torch.Tensor | None = None,
        prompt_embeds_mask: torch.Tensor | None = None,
        max_sequence_length: int = 1024,
+        prompt_name: str = "prompt",
    ):
        r"""

@@ -436,7 +462,11 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]

        if prompt_embeds is None:
-            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt)
+            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(
+                prompt,
+                max_sequence_length=max_sequence_length,
+                prompt_name=prompt_name,
+            )

        prompt_embeds = prompt_embeds[:, :max_sequence_length]
        prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
@@ -629,6 +659,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
                prompt_embeds_mask=negative_prompt_embeds_mask,
                num_images_per_prompt=num_images_per_prompt,
                max_sequence_length=max_sequence_length,
+                prompt_name="negative_prompt",
            )
        else:
            negative_prompt_embeds = None
@@ -700,7 +731,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
            num_images_per_prompt=sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else 1,
            generator=sampling.generator,
            true_cfg_scale=sampling.true_cfg_scale or 4.0,
-            max_sequence_length=sampling.max_sequence_length or 512,
+            max_sequence_length=sampling.max_sequence_length or self.tokenizer_max_length,
            attention_kwargs=kwargs.get("attention_kwargs"),
        )

@@ -931,13 +962,14 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
        output_type: str | None = "pil",
        attention_kwargs: dict[str, Any] | None = None,
        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 512,
+        max_sequence_length: int = 1024,
    ) -> DiffusionOutput:
        extracted_prompt, negative_prompt = self._extract_prompts(req.prompts)
        prompt = extracted_prompt or prompt

        height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
        width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
+        height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
        sigmas = req.sampling_params.sigmas or sigmas
        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
@@ -37,6 +37,12 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
 )
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.prompt_utils import (
+    validate_prompt_sequence_lengths,
+)
+from vllm_omni.diffusion.utils.size_utils import (
+    normalize_min_aligned_size,
+)
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.model_loader.weight_utils import (
@@ -97,9 +103,7 @@ def get_qwen_image_edit_pre_process_func(
            width = request.sampling_params.width or calculated_width

            # Ensure dimensions are multiples of vae_scale_factor * 2
-            multiple_of = vae_scale_factor * 2
-            height = height // multiple_of * multiple_of
-            width = width // multiple_of * multiple_of
+            height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2)

            # Store calculated dimensions in request
            prompt["additional_information"]["calculated_height"] = calculated_height
@@ -322,8 +326,10 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
                "that was used to generate `negative_prompt_embeds`."
            )

-        if max_sequence_length is not None and max_sequence_length > 1024:
-            raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
+        if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length:
+            raise ValueError(
+                f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}"
+            )

    def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
        bool_mask = mask.bool()
@@ -383,6 +389,8 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
        prompt: str | list[str] = None,
        image: PIL.Image.Image | torch.Tensor | None = None,
        dtype: torch.dtype | None = None,
+        max_sequence_length: int | None = None,
+        prompt_name: str = "prompt",
    ):
        """Get prompt embeddings with image support for editing."""
        dtype = dtype or self.text_encoder.dtype
@@ -392,6 +400,33 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
        template = self.prompt_template_encode
        drop_idx = self.prompt_template_encode_start_idx
        txt = [template.format(e) for e in prompt]
+        txt_tokens = self.tokenizer(
+            txt,
+            padding=True,
+            truncation=False,
+            return_tensors="pt",
+        ).to(self.device)
+        # The edit template contains fixed multimodal scaffolding around the
+        # instruction. Validate against the empty-template baseline so image
+        # placeholder text does not consume the user's text budget.
+        template_tokens = self.tokenizer(
+            [template.format("")],
+            padding=True,
+            truncation=False,
+            return_tensors="pt",
+        ).to(self.device)
+        # Qwen-Image-Edit expands image placeholders into many vision tokens
+        # inside the processor. `max_sequence_length` is meant to constrain the
+        # prompt text length, so validate on the text template before image
+        # token expansion.
+        validate_prompt_sequence_lengths(
+            txt_tokens.attention_mask,
+            max_sequence_length=max_sequence_length or self.tokenizer_max_length,
+            supported_max_sequence_length=self.tokenizer_max_length,
+            prompt_name=prompt_name,
+            baseline_attention_mask=template_tokens.attention_mask,
+            error_context="after applying the Qwen prompt template",
+        )

        # Use processor to handle both text and image inputs
        model_inputs = self.processor(
@@ -433,6 +468,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
        prompt_embeds: torch.Tensor | None = None,
        prompt_embeds_mask: torch.Tensor | None = None,
        max_sequence_length: int = 1024,
+        prompt_name: str = "prompt",
    ):
        r"""

@@ -452,7 +488,12 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]

        if prompt_embeds is None:
-            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image)
+            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(
+                prompt,
+                image,
+                max_sequence_length=max_sequence_length,
+                prompt_name=prompt_name,
+            )

        _, seq_len, _ = prompt_embeds.shape
        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
@@ -623,7 +664,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
        output_type: str | None = "pil",
        attention_kwargs: dict[str, Any] | None = None,
        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 512,
+        max_sequence_length: int = 1024,
    ) -> DiffusionOutput:
        """Forward pass for image editing."""
        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
@@ -661,9 +702,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
            height = height or calculated_height
            width = width or calculated_width

-            multiple_of = self.vae_scale_factor * 2
-            width = width // multiple_of * multiple_of
-            height = height // multiple_of * multiple_of
+            height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)

            if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
                image = self.image_processor.resize(image, calculated_height, calculated_width)
@@ -740,6 +779,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
                prompt_embeds_mask=negative_prompt_embeds_mask,
                num_images_per_prompt=num_images_per_prompt,
                max_sequence_length=max_sequence_length,
+                prompt_name="negative_prompt",
            )

        num_channels_latents = self.transformer.in_channels // 4
@@ -25,6 +25,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
    QwenImageCFGParallelMixin,
@@ -40,6 +41,12 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
 )
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.prompt_utils import (
+    validate_prompt_sequence_lengths,
+)
+from vllm_omni.diffusion.utils.size_utils import (
+    normalize_min_aligned_size,
+)
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.model_loader.weight_utils import (
@@ -50,6 +57,12 @@ logger = logging.getLogger(__name__)

 CONDITION_IMAGE_SIZE = 384 * 384
 VAE_IMAGE_SIZE = 1024 * 1024
+# Keep this in sync with the practical conditioning-token budget for
+# Qwen-Image-Edit-2511. Empirically, 4 images stays within the supported range
+# while 5 images overflows the prompt/conditioning path and fails downstream.
+# Re-export the shared metadata value locally so this pipeline keeps a nearby,
+# descriptive constant for validation and tests without becoming the source of truth.
+MAX_QWEN_IMAGE_EDIT_PLUS_INPUT_IMAGES = QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES


 def get_qwen_image_edit_plus_pre_process_func(
@@ -87,6 +100,11 @@ def get_qwen_image_edit_plus_pre_process_func(

            if not isinstance(raw_image, list):
                raw_image = [raw_image]
+            if len(raw_image) > MAX_QWEN_IMAGE_EDIT_PLUS_INPUT_IMAGES:
+                raise ValueError(
+                    f"Received {len(raw_image)} input images. "
+                    f"At most {MAX_QWEN_IMAGE_EDIT_PLUS_INPUT_IMAGES} images are supported by this model."
+                )
            image = [
                PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image | np.ndarray | torch.Tensor, im)
                for im in raw_image
@@ -99,9 +117,7 @@ def get_qwen_image_edit_plus_pre_process_func(
            width = request.sampling_params.width or calculated_width

            # Ensure dimensions are multiples of vae_scale_factor * 2
-            multiple_of = vae_scale_factor * 2
-            height = height // multiple_of * multiple_of
-            width = width // multiple_of * multiple_of
+            height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2)

            # Store calculated dimensions in request
            prompt["additional_information"]["calculated_height"] = calculated_height
@@ -282,8 +298,10 @@ class QwenImageEditPlusPipeline(
                "that was used to generate `negative_prompt_embeds`."
            )

-        if max_sequence_length is not None and max_sequence_length > 1024:
-            raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
+        if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length:
+            raise ValueError(
+                f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}"
+            )

    def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
        bool_mask = mask.bool()
@@ -298,6 +316,8 @@ class QwenImageEditPlusPipeline(
        prompt: str | list[str],
        image: list[torch.Tensor] | torch.Tensor | None = None,
        dtype: torch.dtype | None = None,
+        max_sequence_length: int | None = None,
+        prompt_name: str = "prompt",
    ):
        """Get prompt embeddings with support for multiple images."""
        dtype = dtype or self.text_encoder.dtype
@@ -318,6 +338,32 @@ class QwenImageEditPlusPipeline(
        template = self.prompt_template_encode
        drop_idx = self.prompt_template_encode_start_idx
        txt = [template.format(base_img_prompt + e) for e in prompt]
+        txt_tokens = self.tokenizer(
+            txt,
+            padding=True,
+            truncation=False,
+            return_tensors="pt",
+        ).to(self.device)
+        # Multi-image edit prepends "Picture N" placeholders before the user
+        # instruction. Subtract the placeholder-aware baseline so attached
+        # images do not reduce the remaining prompt budget.
+        template_tokens = self.tokenizer(
+            [template.format(base_img_prompt)],
+            padding=True,
+            truncation=False,
+            return_tensors="pt",
+        ).to(self.device)
+        # The processor expands image placeholders into many vision tokens.
+        # `max_sequence_length` should guard the prompt text length before that
+        # multimodal expansion happens.
+        validate_prompt_sequence_lengths(
+            txt_tokens.attention_mask,
+            max_sequence_length=max_sequence_length or self.tokenizer_max_length,
+            supported_max_sequence_length=self.tokenizer_max_length,
+            prompt_name=prompt_name,
+            baseline_attention_mask=template_tokens.attention_mask,
+            error_context="after applying the Qwen prompt template",
+        )

        # Use processor to handle both text and image inputs
        model_inputs = self.processor(
@@ -359,6 +405,7 @@ class QwenImageEditPlusPipeline(
        prompt_embeds: torch.Tensor | None = None,
        prompt_embeds_mask: torch.Tensor | None = None,
        max_sequence_length: int = 1024,
+        prompt_name: str = "prompt",
    ):
        r"""

@@ -378,7 +425,12 @@ class QwenImageEditPlusPipeline(
        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]

        if prompt_embeds is None:
-            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image)
+            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(
+                prompt,
+                image,
+                max_sequence_length=max_sequence_length,
+                prompt_name=prompt_name,
+            )

        _, seq_len, _ = prompt_embeds.shape
        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
@@ -556,7 +608,7 @@ class QwenImageEditPlusPipeline(
        output_type: str | None = "pil",
        attention_kwargs: dict[str, Any] | None = None,
        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 512,
+        max_sequence_length: int = 1024,
    ) -> DiffusionOutput:
        """Forward pass for image editing with support for multiple images."""
        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
@@ -604,9 +656,7 @@ class QwenImageEditPlusPipeline(
            height = height or calculated_height
            width = width or calculated_width

-            multiple_of = self.vae_scale_factor * 2
-            width = width // multiple_of * multiple_of
-            height = height // multiple_of * multiple_of
+            height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)

            condition_images = []
            vae_images = []
@@ -693,6 +743,7 @@ class QwenImageEditPlusPipeline(
                prompt_embeds_mask=negative_prompt_embeds_mask,
                num_images_per_prompt=num_images_per_prompt,
                max_sequence_length=max_sequence_length,
+                prompt_name="negative_prompt",
            )

        num_channels_latents = self.transformer.in_channels // 4
@@ -36,6 +36,12 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
 )
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.prompt_utils import (
+    validate_prompt_sequence_lengths,
+)
+from vllm_omni.diffusion.utils.size_utils import (
+    normalize_min_aligned_size,
+)
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.model_loader.weight_utils import (
@@ -109,9 +115,7 @@ def get_qwen_image_layered_pre_process_func(
            height = calculated_height
            width = calculated_width

-            multiple_of = vae_scale_factor * 2
-            width = width // multiple_of * multiple_of
-            height = height // multiple_of * multiple_of
+            height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2)

            # Store calculated dimensions in request
            prompt["additional_information"]["calculated_height"] = calculated_height
@@ -339,8 +343,10 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
                "generate `negative_prompt_embeds`."
            )

-        if max_sequence_length is not None and max_sequence_length > 1024:
-            raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
+        if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length:
+            raise ValueError(
+                f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}"
+            )

    def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
        bool_mask = mask.bool()
@@ -355,6 +361,8 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
        prompt: str | list[str] | None = None,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
+        max_sequence_length: int | None = None,
+        prompt_name: str = "prompt",
    ):
        device = device or self.device
        dtype = dtype or self.text_encoder.dtype
@@ -367,8 +375,26 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
        txt_tokens = self.tokenizer(
            txt,
            padding=True,
+            truncation=False,
            return_tensors="pt",
        ).to(device)
+        # The layered template also appends fixed non-user tokens after the
+        # editable text, so use the empty-template tokenized baseline instead of
+        # counting everything after prompt_template_encode_start_idx.
+        template_tokens = self.tokenizer(
+            [template.format("")],
+            padding=True,
+            truncation=False,
+            return_tensors="pt",
+        ).to(device)
+        validate_prompt_sequence_lengths(
+            txt_tokens.attention_mask,
+            max_sequence_length=max_sequence_length or self.tokenizer_max_length,
+            supported_max_sequence_length=self.tokenizer_max_length,
+            prompt_name=prompt_name,
+            baseline_attention_mask=template_tokens.attention_mask,
+            error_context="after applying the Qwen prompt template",
+        )
        encoder_hidden_states = self.text_encoder(
            input_ids=txt_tokens.input_ids,
            attention_mask=txt_tokens.attention_mask,
@@ -398,6 +424,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
        prompt_embeds: torch.Tensor | None = None,
        prompt_embeds_mask: torch.Tensor | None = None,
        max_sequence_length: int = 1024,
+        prompt_name: str = "prompt",
    ):
        r"""

@@ -418,7 +445,11 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]

        if prompt_embeds is None:
-            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt)
+            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(
+                prompt,
+                max_sequence_length=max_sequence_length,
+                prompt_name=prompt_name,
+            )

        prompt_embeds = prompt_embeds[:, :max_sequence_length]
        prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
@@ -602,7 +633,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
        negative_prompt_embeds_mask: torch.Tensor | None = None,
        output_type: str | None = "pil",
        attention_kwargs: dict[str, Any] | None = None,
-        max_sequence_length: int = 512,
+        max_sequence_length: int = 1024,
        resolution: int = 640,
        cfg_normalize: bool = False,
        use_en_prompt: bool = False,
@@ -665,9 +696,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
            height = calculated_height
            width = calculated_width

-            multiple_of = self.vae_scale_factor * 2
-            width = width // multiple_of * multiple_of
-            height = height // multiple_of * multiple_of
+            height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)

            if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
                image = self.image_processor.resize(image, calculated_height, calculated_width)
@@ -737,6 +766,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
                device=self.device,
                num_images_per_prompt=num_images_per_prompt,
                max_sequence_length=max_sequence_length,
+                prompt_name="negative_prompt",
            )

        # 4. Prepare latent variables
@@ -743,9 +743,9 @@ class QwenImageTransformerBlock(nn.Module):

        self.zero_cond_t = zero_cond_t

-    def _modulate(self, x, mod_params, index=None):
+    def _modulate(self, mod_params, index=None):
        """Apply modulation to input tensor"""
-        # x: b l d, shift: b d, scale: b d, gate: b d
+        # shift: b d, scale: b d, gate: b d
        shift, scale, gate = mod_params.chunk(3, dim=-1)

        if index is not None:
@@ -777,7 +777,7 @@ class QwenImageTransformerBlock(nn.Module):
            scale_result = scale.unsqueeze(1)
            gate_result = gate.unsqueeze(1)

-        return x * (1 + scale_result) + shift_result, gate_result
+        return scale_result, shift_result, gate_result

    def forward(
        self,
@@ -803,10 +803,12 @@ class QwenImageTransformerBlock(nn.Module):
        txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]

        # Process image stream - norm1 + modulation
-        img_modulated, img_gate1 = self.img_norm1(hidden_states, img_mod1, modulate_index)
+        img_scale1, img_shift1, img_gate1 = self._modulate(img_mod1, modulate_index)
+        img_modulated = self.img_norm1(hidden_states, img_scale1, img_shift1)

        # Process text stream - norm1 + modulation
-        txt_modulated, txt_gate1 = self.txt_norm1(encoder_hidden_states, txt_mod1)
+        txt_scale1, txt_shift1, txt_gate1 = self._modulate(txt_mod1)
+        txt_modulated = self.txt_norm1(encoder_hidden_states, txt_scale1, txt_shift1)

        # Use QwenAttnProcessor2_0 for joint attention computation
        # This directly implements the DoubleStreamLayerMegatron logic:
@@ -831,13 +833,16 @@ class QwenImageTransformerBlock(nn.Module):
        encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output

        # Process image stream - norm2 + MLP
-        img_modulated2, img_gate2 = self.img_norm2(hidden_states, img_mod2, modulate_index)
+        img_scale2, img_shift2, img_gate2 = self._modulate(img_mod2, modulate_index)
+        img_modulated2 = self.img_norm2(hidden_states, img_scale2, img_shift2)

        img_mlp_output = self.img_mlp(img_modulated2)
        hidden_states = hidden_states + img_gate2 * img_mlp_output

        # Process text stream - norm2 + MLP
-        txt_modulated2, txt_gate2 = self.txt_norm2(encoder_hidden_states, txt_mod2)
+        txt_scale2, txt_shift2, txt_gate2 = self._modulate(txt_mod2)
+        txt_modulated2 = self.txt_norm2(encoder_hidden_states, txt_scale2, txt_shift2)
+
        txt_mlp_output = self.txt_mlp(txt_modulated2)
        encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output

@@ -1,3 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .patch_diffusers import patch_wan_rms_norm
 from .pipeline_wan2_2 import (
    Wan22Pipeline,
    create_transformer_from_config,
@@ -33,3 +37,5 @@ __all__ = [
    "get_wan22_ti2v_pre_process_func",
    "WanTransformer3DModel",
 ]
+
+patch_wan_rms_norm()
@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import sys
+
+from vllm_omni.diffusion.layers.norm import RMSNormVAE
+
+
+def patch_wan_rms_norm():
+    """Patch diffusers Wan RMSNorm implementation with RMSNormVAE."""
+
+    for module_name, module in sys.modules.items():
+        if hasattr(module, "WanRMS_norm"):
+            setattr(module, "WanRMS_norm", RMSNormVAE)
@@ -25,6 +25,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
+from vllm_omni.diffusion.postprocess import interpolate_video_tensor
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniTextPrompt
@@ -121,10 +122,23 @@ def get_wan22_post_process_func(
    def post_process_func(
        video: torch.Tensor,
        output_type: str = "np",
+        sampling_params=None,
    ):
        if output_type == "latent":
            return video
-        return video_processor.postprocess_video(video, output_type=output_type)
+        custom_output = {}
+        if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
+            video, multiplier = interpolate_video_tensor(
+                video,
+                exp=sampling_params.frame_interpolation_exp,
+                scale=sampling_params.frame_interpolation_scale,
+                model_path=sampling_params.frame_interpolation_model_path,
+            )
+            custom_output["video_fps_multiplier"] = multiplier
+        return {
+            "video": video_processor.postprocess_video(video, output_type=output_type),
+            "custom_output": custom_output,
+        }

    return post_process_func

@@ -134,9 +148,6 @@ def get_wan22_pre_process_func(
 ):
    """Pre-process function for Wan2.2: optionally load and resize input image for I2V mode."""
    import numpy as np
-    from diffusers.video_processor import VideoProcessor
-
-    video_processor = VideoProcessor(vae_scale_factor=8)

    def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
        for i, prompt in enumerate(request.prompts):
@@ -180,10 +191,6 @@ def get_wan22_pre_process_func(
            )
            prompt["multi_modal_data"]["image"] = image  # type: ignore # key existence already checked above

-            # Preprocess for VAE
-            prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess(
-                image, height=request.sampling_params.height, width=request.sampling_params.width
-            )
            request.prompts[i] = prompt
        return request

@@ -272,7 +279,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
            model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
        ).to(self.device)
        self.vae = DistributedAutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
        ).to(self.device)

        # Initialize transformers with correct config (weights loaded via load_weights)
@@ -12,6 +12,7 @@ from typing import Any, cast
 import numpy as np
 import PIL.Image
 import torch
+import torchvision.transforms.functional as TF
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel
@@ -30,6 +31,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
    load_transformer_config,
    retrieve_latents,
 )
+from vllm_omni.diffusion.postprocess import interpolate_video_tensor
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniTextPrompt
@@ -72,10 +74,23 @@ def get_wan22_i2v_post_process_func(
    def post_process_func(
        video: torch.Tensor,
        output_type: str = "np",
+        sampling_params=None,
    ):
        if output_type == "latent":
            return video
-        return video_processor.postprocess_video(video, output_type=output_type)
+        custom_output = {}
+        if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
+            video, multiplier = interpolate_video_tensor(
+                video,
+                exp=sampling_params.frame_interpolation_exp,
+                scale=sampling_params.frame_interpolation_scale,
+                model_path=sampling_params.frame_interpolation_model_path,
+            )
+            custom_output["video_fps_multiplier"] = multiplier
+        return {
+            "video": video_processor.postprocess_video(video, output_type=output_type),
+            "custom_output": custom_output,
+        }

    return post_process_func

@@ -84,9 +99,6 @@ def get_wan22_i2v_pre_process_func(
    od_config: OmniDiffusionConfig,
 ):
    """Pre-process function for I2V: load and resize input image."""
-    from diffusers.video_processor import VideoProcessor
-
-    video_processor = VideoProcessor(vae_scale_factor=8)

    def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
        for i, prompt in enumerate(request.prompts):
@@ -132,10 +144,6 @@ def get_wan22_i2v_pre_process_func(
            )
            prompt["multi_modal_data"]["image"] = image  # type: ignore # key existence already checked above

-            # Preprocess for VAE
-            prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess(
-                image, height=request.sampling_params.height, width=request.sampling_params.width
-            )
            request.prompts[i] = prompt
        return request

@@ -217,7 +225,7 @@ class Wan22I2VPipeline(

        # VAE
        self.vae = DistributedAutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
        ).to(self.device)

        # Transformers (weights loaded via load_weights)
@@ -459,6 +467,7 @@ class Wan22I2VPipeline(
        video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)

        if isinstance(image, PIL.Image.Image):
+            image = TF.to_tensor(image).to(device)
            image_tensor = video_processor.preprocess(image, height=height, width=width)
        else:
            image_tensor = image
@@ -467,6 +476,7 @@ class Wan22I2VPipeline(
        # Handle last_image if provided
        if last_image is not None:
            if isinstance(last_image, PIL.Image.Image):
+                last_image = TF.to_tensor(last_image).to(device)
                last_image_tensor = video_processor.preprocess(last_image, height=height, width=width)
            else:
                last_image_tensor = last_image
@@ -789,12 +799,14 @@ class Wan22I2VPipeline(
            return latents, latent_condition, first_frame_mask

        # Wan2.1 style: create mask and concatenate with condition
-        mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, latent_width)
+        mask_lat_size = torch.ones(
+            batch_size, 1, num_frames, latent_height, latent_width, device=latent_condition.device
+        )

        if last_image is None:
-            mask_lat_size[:, :, list(range(1, num_frames))] = 0
+            mask_lat_size[:, :, 1:] = 0
        else:
-            mask_lat_size[:, :, list(range(1, num_frames - 1))] = 0
+            mask_lat_size[:, :, 1 : num_frames - 1] = 0

        first_frame_mask = mask_lat_size[:, :, 0:1]
        first_frame_mask = torch.repeat_interleave(first_frame_mask, dim=2, repeats=self.vae_scale_factor_temporal)
@@ -24,13 +24,13 @@ from typing import Any, cast
 import numpy as np
 import PIL.Image
 import torch
-from diffusers import AutoencoderKLWan
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import AutoTokenizer, UMT5EncoderModel
 from vllm.model_executor.models.utils import AutoWeightsLoader

 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -42,6 +42,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
    load_transformer_config,
    retrieve_latents,
 )
+from vllm_omni.diffusion.postprocess import interpolate_video_tensor
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.platforms import current_omni_platform
@@ -59,10 +60,23 @@ def get_wan22_ti2v_post_process_func(
    def post_process_func(
        video: torch.Tensor,
        output_type: str = "np",
+        sampling_params=None,
    ):
        if output_type == "latent":
            return video
-        return video_processor.postprocess_video(video, output_type=output_type)
+        custom_output = {}
+        if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
+            video, multiplier = interpolate_video_tensor(
+                video,
+                exp=sampling_params.frame_interpolation_exp,
+                scale=sampling_params.frame_interpolation_scale,
+                model_path=sampling_params.frame_interpolation_model_path,
+            )
+            custom_output["video_fps_multiplier"] = multiplier
+        return {
+            "video": video_processor.postprocess_video(video, output_type=output_type),
+            "custom_output": custom_output,
+        }

    return post_process_func

@@ -71,9 +85,6 @@ def get_wan22_ti2v_pre_process_func(
    od_config: OmniDiffusionConfig,
 ):
    """Pre-process function for TI2V: optionally load and resize input image."""
-    from diffusers.video_processor import VideoProcessor
-
-    video_processor = VideoProcessor(vae_scale_factor=8)

    def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
        for i, prompt in enumerate(request.prompts):
@@ -119,10 +130,6 @@ def get_wan22_ti2v_pre_process_func(
            )
            prompt["multi_modal_data"]["image"] = image  # type: ignore # key existence already checked above

-            # Preprocess for VAE
-            prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess(
-                image, height=request.sampling_params.height, width=request.sampling_params.width
-            )
            request.prompts[i] = prompt
        return request

@@ -174,8 +181,8 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress
        ).to(self.device)

        # VAE
-        self.vae = AutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+        self.vae = OmniAutoencoderKLWan.from_pretrained(
+            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
        ).to(self.device)

        # Single transformer (TI2V uses dense 5B model, not MoE)
@@ -11,7 +11,6 @@ import torch.nn.functional as F
 from diffusers.models.attention import FeedForward
 from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
-from diffusers.models.normalization import FP32LayerNorm
 from vllm.distributed import (
    get_tensor_model_parallel_rank,
    get_tensor_model_parallel_world_size,
@@ -29,6 +28,8 @@ from vllm_omni.diffusion.distributed.sp_plan import (
    SequenceParallelOutput,
 )
 from vllm_omni.diffusion.forward_context import get_forward_context
+from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
+from vllm_omni.diffusion.layers.norm import LayerNorm, RMSNorm

 logger = init_logger(__name__)

@@ -52,10 +53,14 @@ def apply_rotary_emb_wan(
    x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
    cos = freqs_cos[..., 0::2]
    sin = freqs_sin[..., 1::2]
-    out = torch.empty_like(hidden_states)
-    out[..., 0::2] = x1 * cos - x2 * sin
-    out[..., 1::2] = x1 * sin + x2 * cos
-    return out.type_as(hidden_states)
+    rotated = torch.stack(
+        (
+            x1 * cos - x2 * sin,
+            x1 * sin + x2 * cos,
+        ),
+        dim=-1,
+    )
+    return rotated.flatten(-2, -1).to(hidden_states.dtype)


 class DistributedRMSNorm(nn.Module):
@@ -230,9 +235,9 @@ class WanImageEmbedding(nn.Module):
    def __init__(self, in_features: int, out_features: int, pos_embed_seq_len: int | None = None):
        super().__init__()

-        self.norm1 = FP32LayerNorm(in_features)
+        self.norm1 = LayerNorm(in_features)
        self.ff = FeedForward(in_features, out_features, mult=1, activation_fn="gelu")
-        self.norm2 = FP32LayerNorm(out_features)
+        self.norm2 = LayerNorm(out_features)
        if pos_embed_seq_len is not None:
            self.pos_embed = nn.Parameter(torch.zeros(1, pos_embed_seq_len, in_features))
        else:
@@ -372,8 +377,12 @@ class WanSelfAttention(nn.Module):
        self.tp_inner_dim = self.num_heads * head_dim

        # QK normalization using vLLM's RMSNorm
-        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
-        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+        if get_tensor_model_parallel_world_size() > 1:
+            self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+            self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+        else:
+            self.norm_q = RMSNorm(self.tp_inner_dim, eps=eps)
+            self.norm_k = RMSNorm(self.tp_inner_dim, eps=eps)

        self.to_out = RowParallelLinear(
            self.inner_dim,
@@ -492,8 +501,12 @@ class WanCrossAttention(nn.Module):
        self.tp_inner_dim = self.num_heads * head_dim

        # QK normalization
-        self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
-        self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+        if get_tensor_model_parallel_world_size() > 1:
+            self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+            self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+        else:
+            self.norm_q = RMSNorm(self.tp_inner_dim, eps=eps)
+            self.norm_k = RMSNorm(self.tp_inner_dim, eps=eps)

        # Optional added KV projections for I2V (image embeddings)
        self.added_kv_proj_dim = added_kv_proj_dim
@@ -512,7 +525,10 @@ class WanCrossAttention(nn.Module):
                gather_output=False,
                return_bias=False,
            )
-            self.norm_added_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+            if get_tensor_model_parallel_world_size() > 1:
+                self.norm_added_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps)
+            else:
+                self.norm_added_k = RMSNorm(self.tp_inner_dim, eps=eps)
        else:
            self.add_k_proj = None
            self.add_v_proj = None
@@ -535,6 +551,7 @@ class WanCrossAttention(nn.Module):
            num_kv_heads=self.num_heads,
            softmax_scale=1.0 / (head_dim**0.5),
            causal=False,
+            skip_sequence_parallel=True,
        )

    def forward(
@@ -614,7 +631,8 @@ class WanTransformerBlock(nn.Module):
        head_dim = dim // num_heads

        # 1. Self-attention
-        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
+
        self.attn1 = WanSelfAttention(
            dim=dim,
            num_heads=num_heads,
@@ -630,11 +648,11 @@ class WanTransformerBlock(nn.Module):
            eps=eps,
            added_kv_proj_dim=added_kv_proj_dim,
        )
-        self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
+        self.norm2 = LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()

        # 3. Feed-forward
        self.ffn = WanFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim)
-        self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm3 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)

        # Scale-shift table for modulation
        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -650,7 +668,7 @@ class WanTransformerBlock(nn.Module):
        if temb.ndim == 4:
            # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
-                self.scale_shift_table.unsqueeze(0) + temb.float()
+                self.scale_shift_table.unsqueeze(0) + temb
            ).chunk(6, dim=2)
            shift_msa = shift_msa.squeeze(2)
            scale_msa = scale_msa.squeeze(2)
@@ -661,25 +679,23 @@ class WanTransformerBlock(nn.Module):
        else:
            # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B)
            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
-                self.scale_shift_table + temb.float()
+                self.scale_shift_table + temb
            ).chunk(6, dim=1)

        # 1. Self-attention
-        norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states)
+        norm_hidden_states = self.norm1(hidden_states, scale_msa, shift_msa).type_as(hidden_states)
        attn_output = self.attn1(norm_hidden_states, rotary_emb, hidden_states_mask)
-        hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states)
+        hidden_states = (hidden_states + attn_output * gate_msa).type_as(hidden_states)

        # 2. Cross-attention
-        norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states)
+        norm_hidden_states = self.norm2(hidden_states).type_as(hidden_states)
        attn_output = self.attn2(norm_hidden_states, encoder_hidden_states)
        hidden_states = hidden_states + attn_output

        # 3. Feed-forward
-        norm_hidden_states = (self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as(
-            hidden_states
-        )
+        norm_hidden_states = self.norm3(hidden_states, c_scale_msa, c_shift_msa).type_as(hidden_states)
        ff_output = self.ffn(norm_hidden_states)
-        hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states)
+        hidden_states = (hidden_states + ff_output * c_gate_msa).type_as(hidden_states)

        return hidden_states

@@ -848,7 +864,7 @@ class WanTransformer3DModel(nn.Module):
        )

        # 4. Output norm & projection
-        self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False)
+        self.norm_out = AdaLayerNorm(inner_dim, elementwise_affine=False, eps=eps)
        self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))

        # SP helper modules
@@ -936,7 +952,7 @@ class WanTransformer3DModel(nn.Module):
            shift = shift.unsqueeze(1)
            scale = scale.unsqueeze(1)

-        hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states)
+        hidden_states = self.norm_out(hidden_states, scale, shift).type_as(hidden_states)
        hidden_states = self.proj_out(hidden_states)

        hidden_states = hidden_states.reshape(
@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Diffusion post-processing helpers."""
+
+from vllm_omni.diffusion.postprocess.rife_interpolator import (
+    FrameInterpolator,
+    interpolate_video_tensor,
+)
+
+__all__ = ["FrameInterpolator", "interpolate_video_tensor"]
@@ -0,0 +1,443 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+RIFE 4.22.lite frame interpolation for vLLM-Omni video generation.
+
+RIFE model code is vendored and adapted from:
+  - https://github.com/hzwer/ECCV2022-RIFE  (MIT License)
+  - https://github.com/hzwer/Practical-RIFE  (MIT License)
+  Copyright (c) 2021 Zhewei Huang
+
+The FrameInterpolator wrapper and vLLM-Omni integration code are original work.
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_DEFAULT_RIFE_HF_REPO = "elfgum/RIFE-4.22.lite"
+_MODEL_CACHE: dict[tuple[str, str], Model] = {}
+_MODEL_CACHE_LOCK = threading.Lock()
+
+
+def warp(ten_input: torch.Tensor, ten_flow: torch.Tensor) -> torch.Tensor:
+    """Warp input tensor by optical flow using grid_sample."""
+    ten_horizontal = (
+        torch.linspace(-1.0, 1.0, ten_flow.shape[3], device=ten_flow.device)
+        .view(1, 1, 1, ten_flow.shape[3])
+        .expand(ten_flow.shape[0], -1, ten_flow.shape[2], -1)
+    )
+    ten_vertical = (
+        torch.linspace(-1.0, 1.0, ten_flow.shape[2], device=ten_flow.device)
+        .view(1, 1, ten_flow.shape[2], 1)
+        .expand(ten_flow.shape[0], -1, -1, ten_flow.shape[3])
+    )
+    ten_grid = torch.cat([ten_horizontal, ten_vertical], dim=1)
+
+    ten_flow = torch.cat(
+        [
+            ten_flow[:, 0:1, :, :] / ((ten_input.shape[3] - 1.0) / 2.0),
+            ten_flow[:, 1:2, :, :] / ((ten_input.shape[2] - 1.0) / 2.0),
+        ],
+        dim=1,
+    )
+    grid = (ten_grid + ten_flow).permute(0, 2, 3, 1)
+    return F.grid_sample(
+        input=ten_input,
+        grid=grid,
+        mode="bilinear",
+        padding_mode="border",
+        align_corners=True,
+    )
+
+
+def _conv(
+    in_planes: int,
+    out_planes: int,
+    kernel_size: int = 3,
+    stride: int = 1,
+    padding: int = 1,
+    dilation: int = 1,
+) -> nn.Sequential:
+    return nn.Sequential(
+        nn.Conv2d(
+            in_planes,
+            out_planes,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=True,
+        ),
+        nn.LeakyReLU(0.2, True),
+    )
+
+
+class ResConv(nn.Module):
+    """Residual convolution block with learnable beta scaling."""
+
+    def __init__(self, c: int, dilation: int = 1):
+        super().__init__()
+        self.conv = nn.Conv2d(c, c, 3, 1, dilation, dilation=dilation, groups=1)
+        self.beta = nn.Parameter(torch.ones((1, c, 1, 1)), requires_grad=True)
+        self.relu = nn.LeakyReLU(0.2, True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.relu(self.conv(x) * self.beta + x)
+
+
+class IFBlock(nn.Module):
+    """Single-scale optical flow, mask, and feature block."""
+
+    def __init__(self, in_planes: int, c: int = 64):
+        super().__init__()
+        self.conv0 = nn.Sequential(
+            _conv(in_planes, c // 2, 3, 2, 1),
+            _conv(c // 2, c, 3, 2, 1),
+        )
+        self.convblock = nn.Sequential(
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+            ResConv(c),
+        )
+        self.lastconv = nn.Sequential(
+            nn.ConvTranspose2d(c, 4 * 13, 4, 2, 1),
+            nn.PixelShuffle(2),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        flow: torch.Tensor | None = None,
+        scale: float = 1.0,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False)
+        if flow is not None:
+            flow = (
+                F.interpolate(
+                    flow,
+                    scale_factor=1.0 / scale,
+                    mode="bilinear",
+                    align_corners=False,
+                )
+                * 1.0
+                / scale
+            )
+            x = torch.cat((x, flow), 1)
+        feat = self.conv0(x)
+        feat = self.convblock(feat)
+        tmp = self.lastconv(feat)
+        tmp = F.interpolate(tmp, scale_factor=scale, mode="bilinear", align_corners=False)
+        flow = tmp[:, :4] * scale
+        mask = tmp[:, 4:5]
+        feat = tmp[:, 5:]
+        return flow, mask, feat
+
+
+class Head(nn.Module):
+    """Feature encoder producing four-channel features at full resolution."""
+
+    def __init__(self):
+        super().__init__()
+        self.cnn0 = nn.Conv2d(3, 16, 3, 2, 1)
+        self.cnn1 = nn.Conv2d(16, 16, 3, 1, 1)
+        self.cnn2 = nn.Conv2d(16, 16, 3, 1, 1)
+        self.cnn3 = nn.ConvTranspose2d(16, 4, 4, 2, 1)
+        self.relu = nn.LeakyReLU(0.2, True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x0 = self.cnn0(x)
+        x = self.relu(x0)
+        x1 = self.cnn1(x)
+        x = self.relu(x1)
+        x2 = self.cnn2(x)
+        x = self.relu(x2)
+        x3 = self.cnn3(x)
+        return x3
+
+
+class IFNet(nn.Module):
+    """Four-scale IFNet optical flow network."""
+
+    def __init__(self):
+        super().__init__()
+        self.block0 = IFBlock(7 + 8, c=192)
+        self.block1 = IFBlock(8 + 4 + 8 + 8, c=128)
+        self.block2 = IFBlock(8 + 4 + 8 + 8, c=64)
+        self.block3 = IFBlock(8 + 4 + 8 + 8, c=32)
+        self.encode = Head()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        timestep: float = 0.5,
+        scale_list: list[float] | None = None,
+    ) -> tuple[list[torch.Tensor], torch.Tensor, list[tuple[torch.Tensor, torch.Tensor] | torch.Tensor]]:
+        if scale_list is None:
+            scale_list = [8, 4, 2, 1]
+
+        channel = x.shape[1] // 2
+        img0 = x[:, :channel]
+        img1 = x[:, channel:]
+
+        if not torch.is_tensor(timestep):
+            timestep = (x[:, :1].clone() * 0 + 1) * timestep
+        else:
+            timestep = timestep.repeat(1, 1, img0.shape[2], img0.shape[3])
+
+        f0 = self.encode(img0[:, :3])
+        f1 = self.encode(img1[:, :3])
+
+        flow_list: list[torch.Tensor] = []
+        merged: list[tuple[torch.Tensor, torch.Tensor] | torch.Tensor] = []
+        mask_list: list[torch.Tensor] = []
+        warped_img0 = img0
+        warped_img1 = img1
+        flow = None
+        mask = None
+
+        for i, block in enumerate([self.block0, self.block1, self.block2, self.block3]):
+            if flow is None:
+                flow, mask, feat = block(
+                    torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1),
+                    None,
+                    scale=scale_list[i],
+                )
+            else:
+                wf0 = warp(f0, flow[:, :2])
+                wf1 = warp(f1, flow[:, 2:4])
+                fd, m0, feat = block(
+                    torch.cat(
+                        (
+                            warped_img0[:, :3],
+                            warped_img1[:, :3],
+                            wf0,
+                            wf1,
+                            timestep,
+                            mask,
+                            feat,
+                        ),
+                        1,
+                    ),
+                    flow,
+                    scale=scale_list[i],
+                )
+                mask = m0
+                flow = flow + fd
+
+            mask_list.append(mask)
+            flow_list.append(flow)
+            warped_img0 = warp(img0, flow[:, :2])
+            warped_img1 = warp(img1, flow[:, 2:4])
+            merged.append((warped_img0, warped_img1))
+
+        mask = torch.sigmoid(mask)
+        merged[3] = warped_img0 * mask + warped_img1 * (1 - mask)
+        return flow_list, mask_list[3], merged
+
+
+class Model:
+    """Wraps IFNet and exposes RIFE-compatible load/inference helpers."""
+
+    def __init__(self):
+        self.flownet = IFNet()
+
+    def eval(self) -> Model:
+        self.flownet.eval()
+        return self
+
+    def device(self) -> torch.device:
+        return next(self.flownet.parameters()).device
+
+    def load_model(self, path: str) -> None:
+        flownet_path = os.path.join(path, "flownet.pkl")
+        if not os.path.isfile(flownet_path):
+            raise FileNotFoundError(
+                f"RIFE weight file not found: {flownet_path}. Expected layout: <model_path>/flownet.pkl"
+            )
+
+        state = torch.load(flownet_path, map_location="cpu", weights_only=False)
+        state = {k.removeprefix("module."): v for k, v in state.items()}
+        self.flownet.load_state_dict(state, strict=False)
+        logger.info("Loaded RIFE weights from %s", flownet_path)
+
+    def inference(
+        self,
+        img0: torch.Tensor,
+        img1: torch.Tensor,
+        scale: float = 1.0,
+        timestep: float = 0.5,
+    ) -> torch.Tensor:
+        _n, _c, h, w = img0.shape
+        ph = ((h - 1) // 32 + 1) * 32
+        pw = ((w - 1) // 32 + 1) * 32
+        pad = (0, pw - w, 0, ph - h)
+        img0 = F.pad(img0, pad)
+        img1 = F.pad(img1, pad)
+
+        imgs = torch.cat((img0, img1), 1)
+        scale_list = [8 / scale, 4 / scale, 2 / scale, 1 / scale]
+        with torch.no_grad():
+            _flow_list, _mask, merged = self.flownet(
+                imgs,
+                timestep=timestep,
+                scale_list=scale_list,
+            )
+        return merged[3][:, :, :h, :w]
+
+
+def _resolve_rife_model_path(model_path: str | None) -> str:
+    model_path = model_path or _DEFAULT_RIFE_HF_REPO
+    if os.path.isdir(model_path):
+        return model_path
+    from vllm_omni.model_executor.model_loader.weight_utils import (
+        download_weights_from_hf_specific,
+    )
+
+    return download_weights_from_hf_specific(
+        model_path,
+        cache_dir=None,
+        allow_patterns=["flownet.pkl"],
+        require_all=True,
+    )
+
+
+def _select_torch_device() -> torch.device:
+    try:
+        from vllm_omni.platforms import current_omni_platform
+
+        return current_omni_platform.get_torch_device()
+    except Exception as exc:
+        logger.warning("Failed to resolve current vLLM-Omni torch device: %s", exc)
+
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+def _normalize_video_tensor_layout(video: torch.Tensor) -> tuple[torch.Tensor, Any]:
+    if video.ndim == 5:
+        if video.shape[1] in (3, 4):
+            return video, lambda out: out
+        if video.shape[2] in (3, 4):
+            return video.permute(0, 2, 1, 3, 4), lambda out: out.permute(0, 2, 1, 3, 4)
+    elif video.ndim == 4:
+        if video.shape[0] in (3, 4):
+            return video.unsqueeze(0), lambda out: out.squeeze(0)
+        if video.shape[1] in (3, 4):
+            return video.permute(1, 0, 2, 3).unsqueeze(0), lambda out: out.squeeze(0).permute(1, 0, 2, 3)
+    raise ValueError(f"Unsupported video tensor shape for interpolation: {tuple(video.shape)}")
+
+
+def _normalize_video_tensor_range(video: torch.Tensor) -> tuple[torch.Tensor, Any]:
+    original_dtype = video.dtype
+    video = video.detach()
+    if video.is_floating_point():
+        video = video.to(torch.float32)
+        if torch.amin(video) < 0.0 or torch.amax(video) > 1.0:
+            return video.clamp(-1.0, 1.0) * 0.5 + 0.5, lambda out: (out * 2.0 - 1.0).to(original_dtype)
+        return video.clamp(0.0, 1.0), lambda out: out.to(original_dtype)
+    return video.to(torch.float32) / 255.0, lambda out: (out * 255.0).round().clamp(0, 255).to(original_dtype)
+
+
+class FrameInterpolator:
+    """Lazy-loaded RIFE 4.22.lite frame interpolator."""
+
+    def __init__(self, model_path: str | None = None):
+        self._model_path = model_path
+        self._resolved_path: str | None = None
+
+    def _ensure_model_loaded(self, preferred_device: torch.device | None = None) -> Model:
+        resolved_path = _resolve_rife_model_path(self._model_path)
+        self._resolved_path = resolved_path
+        device = preferred_device or _select_torch_device()
+        cache_key = (resolved_path, str(device))
+
+        with _MODEL_CACHE_LOCK:
+            if cache_key in _MODEL_CACHE:
+                return _MODEL_CACHE[cache_key]
+
+            model = Model()
+            model.load_model(resolved_path)
+            model.eval()
+            model.flownet = model.flownet.to(device)
+            _MODEL_CACHE[cache_key] = model
+            logger.info("RIFE model loaded on device: %s", device)
+            return model
+
+    def _make_inference(
+        self,
+        model: Model,
+        img0: torch.Tensor,
+        img1: torch.Tensor,
+        n: int,
+        scale: float,
+    ) -> list[torch.Tensor]:
+        if n == 1:
+            return [model.inference(img0, img1, scale=scale)]
+        mid = model.inference(img0, img1, scale=scale)
+        return (
+            self._make_inference(model, img0, mid, n // 2, scale)
+            + [mid]
+            + self._make_inference(model, mid, img1, n // 2, scale)
+        )
+
+    def interpolate_tensor(
+        self,
+        video: torch.Tensor,
+        exp: int = 1,
+        scale: float = 1.0,
+    ) -> tuple[torch.Tensor, int]:
+        if exp < 1:
+            raise ValueError(f"frame interpolation exp must be >= 1, got {exp}")
+        if scale <= 0:
+            raise ValueError(f"frame interpolation scale must be > 0, got {scale}")
+
+        video, restore_layout = _normalize_video_tensor_layout(video)
+        if video.shape[2] < 2:
+            return restore_layout(video), 1
+
+        video, restore_range = _normalize_video_tensor_range(video)
+        # A CPU tensor may be transport/offload state rather than an execution
+        # choice, so only trust it when it is already on an accelerator.
+        preferred_device = video.device
+        if preferred_device.type == "cpu":
+            preferred_device = _select_torch_device()
+        model = self._ensure_model_loaded(preferred_device=preferred_device)
+        video = video.to(model.device())
+        intermediates_per_pair = 2**exp // 2
+
+        result_frames: list[torch.Tensor] = []
+        for idx in range(video.shape[2] - 1):
+            img0 = video[:, :, idx, :, :]
+            img1 = video[:, :, idx + 1, :, :]
+            result_frames.append(img0)
+            result_frames.extend(self._make_inference(model, img0, img1, intermediates_per_pair, scale))
+        result_frames.append(video[:, :, -1, :, :])
+        result = torch.stack(result_frames, dim=2)
+        return restore_layout(restore_range(result)), 2**exp
+
+
+def interpolate_video_tensor(
+    video: torch.Tensor,
+    exp: int = 1,
+    scale: float = 1.0,
+    model_path: str | None = None,
+) -> tuple[torch.Tensor, int]:
+    """Interpolate a video tensor and return the FPS multiplier."""
+    interpolator = FrameInterpolator(model_path=model_path)
+    return interpolator.interpolate_tensor(video, exp=exp, scale=scale)
@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any

 from vllm.logger import init_logger

+from vllm_omni.diffusion.data import normalize_omni_error
 from vllm_omni.engine.stage_init_utils import StageMetadata
 from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
 from vllm_omni.outputs import OmniRequestOutput
@@ -23,6 +24,21 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)


+def create_diffusion_client(
+    model: str,
+    od_config: OmniDiffusionConfig,
+    metadata: StageMetadata,
+    batch_size: int = 1,
+    use_inline: bool = False,
+) -> Any:
+    """Factory to create either an inline or the default diffusion client."""
+    if use_inline:
+        from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
+
+        return InlineStageDiffusionClient(model, od_config, metadata, batch_size=batch_size)
+    return StageDiffusionClient(model, od_config, metadata, batch_size=batch_size)
+
+
 class StageDiffusionClient:
    """Wraps AsyncOmniDiffusion for use inside the Orchestrator.

@@ -75,11 +91,24 @@ class StageDiffusionClient:
            result = await self._engine.generate(prompt, sampling_params, request_id)
            await self._output_queue.put(result)
        except Exception as e:
+            err = normalize_omni_error(e, request_id=request_id, stage_id=self.stage_id)
            logger.exception(
                "[StageDiffusionClient] Stage-%s req=%s failed: %s",
                self.stage_id,
                request_id,
-                e,
+                err,
+            )
+            await self._output_queue.put(
+                {
+                    "type": "error",
+                    "request_id": request_id,
+                    "stage_id": self.stage_id,
+                    "status_code": err.status_code,
+                    "error": str(err),
+                    "error_type": err.error_type,
+                    "detail": err.detail,
+                    "finished": True,
+                }
            )
        finally:
            self._tasks.pop(request_id, None)
@@ -162,10 +191,8 @@ class StageDiffusionClient:
                    "todo": True,
                    "reason": f"AsyncOmniDiffusion.{method} is not implemented",
                }
-            # Extract is_start and profile_prefix from args
            is_start = args[0] if args else True
            profile_prefix = args[1] if len(args) > 1 else None
-            # Generate profile_prefix with stage_id if starting and no prefix provided
            if is_start and profile_prefix is None:
                profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}"
            result = target(is_start, profile_prefix)
@@ -186,7 +213,6 @@ class StageDiffusionClient:
                return await asyncio.wait_for(result, timeout=timeout)
            return await result

-        # Fall back to collective RPC for other methods
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(
            self._engine._executor,
@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Video/audio muxing utilities using PyAV (no ffmpeg binary dependency)."""
+
+from __future__ import annotations
+
+import io
+from fractions import Fraction
+
+import av
+import numpy as np
+
+
+def mux_video_audio_bytes(
+    video_frames: np.ndarray,
+    audio_waveform: np.ndarray | None = None,
+    *,
+    fps: float = 25.0,
+    audio_sample_rate: int = 44100,
+    video_codec: str = "h264",
+    audio_codec: str = "aac",
+    crf: str = "18",
+    video_codec_options: dict[str, str] | None = None,
+) -> bytes:
+    """Mux video frames and optional audio waveform into MP4 bytes.
+
+    Args:
+        video_frames: uint8 array of shape ``(T, H, W, 3)`` (RGB).
+        audio_waveform: float32 array – mono ``(N,)`` or ``(N, C)`` / ``(C, N)``.
+        fps: Video frame rate.
+        audio_sample_rate: Audio sample rate in Hz.
+        video_codec: Video codec name.
+        audio_codec: Audio codec name.
+        crf: Constant rate factor for the video encoder.
+
+    Returns:
+        Raw MP4 bytes ready to be written to disk or streamed.
+    """
+    buf = io.BytesIO()
+    container = av.open(buf, mode="w", format="mp4")
+
+    v_stream = container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000))
+    v_stream.width = video_frames.shape[2]
+    v_stream.height = video_frames.shape[1]
+    v_stream.pix_fmt = "yuv420p"
+
+    options = {"crf": str(crf)}
+    if video_codec_options:
+        options.update(video_codec_options)
+    v_stream.options = options
+
+    a_stream = None
+    if audio_waveform is not None:
+        samples = audio_waveform.astype(np.float32)
+        if samples.ndim == 1:
+            samples = samples.reshape(1, -1)
+        elif samples.ndim == 2 and samples.shape[0] > samples.shape[1]:
+            samples = np.ascontiguousarray(samples.T)
+        num_channels = samples.shape[0]
+        layout = "stereo" if num_channels >= 2 else "mono"
+        a_stream = container.add_stream(audio_codec, rate=audio_sample_rate)
+        a_stream.layout = layout
+
+    for frame_data in video_frames:
+        frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24")
+        for packet in v_stream.encode(frame):
+            container.mux(packet)
+    for packet in v_stream.encode():
+        container.mux(packet)
+
+    if a_stream is not None and audio_waveform is not None:
+        audio_frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout=layout)
+        audio_frame.sample_rate = audio_sample_rate
+        for packet in a_stream.encode(audio_frame):
+            container.mux(packet)
+        for packet in a_stream.encode():
+            container.mux(packet)
+
+    container.close()
+    return buf.getvalue()
@@ -0,0 +1,38 @@
+import torch
+
+
+def validate_prompt_sequence_lengths(
+    attention_mask: torch.Tensor,
+    *,
+    max_sequence_length: int,
+    supported_max_sequence_length: int,
+    prompt_name: str = "prompt",
+    length_offset: int = 0,
+    baseline_attention_mask: torch.Tensor | None = None,
+    error_context: str,
+) -> None:
+    sequence_lengths = attention_mask.sum(dim=1)
+    if baseline_attention_mask is not None:
+        # Some callers need to validate only the user-controlled portion of a
+        # templated prompt. In those cases we subtract the fully-tokenized
+        # template baseline instead of only removing a fixed prefix length,
+        # because the template may also contribute a suffix or image markers.
+        baseline_lengths = baseline_attention_mask.sum(dim=1)
+        if baseline_lengths.shape[0] == 1 and sequence_lengths.shape[0] > 1:
+            baseline_lengths = baseline_lengths.expand(sequence_lengths.shape[0])
+        sequence_lengths = sequence_lengths - baseline_lengths
+    if length_offset:
+        sequence_lengths = sequence_lengths - length_offset
+    sequence_lengths = torch.clamp(sequence_lengths, min=0)
+    too_long = torch.nonzero(sequence_lengths > max_sequence_length, as_tuple=False)
+    if too_long.numel() == 0:
+        return
+
+    batch_idx = int(too_long[0].item())
+    actual_length = int(sequence_lengths[batch_idx].item())
+    prompt_ref = f"`{prompt_name}` at batch index {batch_idx}" if attention_mask.shape[0] > 1 else f"`{prompt_name}`"
+    raise ValueError(
+        f"{prompt_ref} is too long {error_context}: got {actual_length} tokens, but "
+        f"`max_sequence_length` is {max_sequence_length}. Shorten the prompt or increase "
+        f"`max_sequence_length` up to {supported_max_sequence_length}."
+    )
@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Shared size normalization helpers for diffusion pipelines."""
+
+
+def normalize_min_aligned_size(height: int, width: int, alignment: int) -> tuple[int, int]:
+    """Clamp dimensions to the minimum valid aligned size.
+
+    This preserves floor-to-alignment behavior for normal requests while
+    preventing very small dimensions from collapsing to zero after alignment.
+    """
+
+    alignment = int(alignment)
+    if alignment <= 0:
+        raise ValueError(f"Expected positive alignment, got {alignment}")
+
+    normalized_height = max(alignment, (int(height) // alignment) * alignment)
+    normalized_width = max(alignment, (int(width) // alignment) * alignment)
+    return normalized_height, normalized_width
@@ -487,11 +487,13 @@ class AsyncOmniEngine:

                                    inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to)
                                _inject_kv_stage_info(stage_cfg, stage_id)
+                                use_inline = self.num_stages == 1
                                stage_clients[stage_id] = initialize_diffusion_stage(
                                    self.model,
                                    stage_cfg,
                                    metadata,
                                    batch_size=self.diffusion_batch_size,
+                                    use_inline=use_inline,
                                )
                                logger.info(
                                    "[AsyncOmniEngine] Stage %s initialized (diffusion, batch_size=%d)",
@@ -798,6 +800,16 @@ class AsyncOmniEngine:
        # We temporally create a default config for diffusion stage.
        # In the future, we should merge the default config with the user-provided config.
        normalized_kwargs = dict(kwargs)
+        default_sampling_params = normalized_kwargs.get("default_sampling_params")
+        if isinstance(default_sampling_params, str):
+            try:
+                default_sampling_params = json.loads(default_sampling_params)
+            except json.JSONDecodeError:
+                logger.warning("Invalid default_sampling_params JSON, ignoring stage defaults.")
+                default_sampling_params = None
+        if not isinstance(default_sampling_params, dict):
+            default_sampling_params = None
+        stage_default_sampling_params = default_sampling_params.get("0", {}) if default_sampling_params else {}

        # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error.
        if "dtype" in normalized_kwargs and not isinstance(normalized_kwargs["dtype"], str):
@@ -886,6 +898,7 @@ class AsyncOmniEngine:
                        else {}
                    ),
                },
+                "default_sampling_params": stage_default_sampling_params,
                "final_output": True,
                "final_output_type": "image",
            }
@@ -242,8 +242,48 @@ class Orchestrator:
                    output = stage_client.get_diffusion_output_async()
                    if output is not None:
                        idle = False
+
+                        if isinstance(output, dict) and output.get("type") == "error":
+                            req_id = output.get("request_id")
+                            req_state = self.request_states.get(req_id)
+
+                            await self.output_async_queue.put(
+                                {
+                                    "type": "error",
+                                    "request_id": req_id,
+                                    "stage_id": stage_id,
+                                    "status_code": output.get("status_code", 500),
+                                    "error": output.get("error", "Unknown diffusion stage error"),
+                                    "error_type": output.get("error_type"),
+                                    "detail": output.get("detail") or {},
+                                    "finished": True,
+                                    "stage_submit_ts": req_state.stage_submit_ts.get(stage_id) if req_state else None,
+                                }
+                            )
+
+                            self._cleanup_companion_state(req_id)
+                            self.request_states.pop(req_id, None)
+                            continue
+
                        req_state = self.request_states.get(output.request_id)
                        if req_state is not None:
+                            if getattr(output, "error", None) is not None:
+                                parent_id = self._companion_to_parent.get(output.request_id, output.request_id)
+                                await self.output_async_queue.put(
+                                    {
+                                        "type": "error",
+                                        "request_id": parent_id,
+                                        "stage_id": stage_id,
+                                        "error": output.error,
+                                    }
+                                )
+                                role_map = self._companion_map.get(parent_id, {})
+                                for cid in role_map.values():
+                                    self.request_states.pop(cid, None)
+                                self._cleanup_companion_state(parent_id)
+                                self.request_states.pop(parent_id, None)
+                                continue
+
                            stage_metrics = self._build_stage_metrics(stage_id, output.request_id, [output], req_state)
                            await self._route_output(stage_id, output, req_state, stage_metrics)
                    continue
@@ -437,6 +437,7 @@ def initialize_diffusion_stage(
    stage_cfg: Any,
    metadata: StageMetadata,
    batch_size: int = 1,
+    use_inline: bool = False,
 ) -> Any:
    """Build a diffusion stage client.

@@ -447,9 +448,10 @@ def initialize_diffusion_stage(
        batch_size: Maximum number of requests to batch together in the
            diffusion engine.  Passed through to ``StageDiffusionClient``
            and ultimately to ``AsyncOmniDiffusion``.
+        use_inline: If True, uses the inline diffusion client instead of subprocess.
    """
    from vllm_omni.diffusion.data import OmniDiffusionConfig
-    from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+    from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client

    od_config = OmniDiffusionConfig.from_kwargs(
        model=model,
@@ -457,7 +459,7 @@ def initialize_diffusion_stage(
    )
    if metadata.cfg_kv_collect_func is not None:
        od_config.cfg_kv_collect_func = metadata.cfg_kv_collect_func
-    return StageDiffusionClient(model, od_config, metadata, batch_size=batch_size)
+    return create_diffusion_client(model, od_config, metadata, batch_size=batch_size, use_inline=use_inline)


 def close_started_llm_stage(started: StartedLlmStage) -> None:
Author	SHA1	Message	Date
lyj-jjj	5fd676f086	cherry-pick pr3067 (#3087 ) Signed-off-by: lyj-jjj <liuyingjun5@huawei.com>	2026-04-24 09:56:49 +08:00
WeiQing Chen	2980618f70	[release v0.18.0.post1] Fix default sampling params for /v1/videos (#3021 ) Signed-off-by: david6666666 <530634352@qq.com>	2026-04-23 11:29:54 +08:00
bjf-frz	6d76e7d343	[cherry-pick][release/v.0.18.0.post1]remove duplicate video preprocess in Wan2.2 pipeline (#3005 ) Signed-off-by: bjf-frz <frz123db@gmail.com>	2026-04-22 11:09:27 +08:00
fan2956	caa26e1d04	[Cherry-pick][v0.18.0.post1][Bugfix] Fix VAE parallelism dist.gather performance bottleneck on NPU (#2970 ) Signed-off-by: fan2956 <zhoufan53@huawei.com>	2026-04-21 20:36:08 +08:00
WeiQing Chen	89f733d941	[codex][release/v0.18.0.post1] revert Wan2.2 pipeline changes from #2878 (#2937 ) Signed-off-by: david6666666 <530634352@qq.com>	2026-04-20 19:22:11 +08:00
fan2956	8a1ff4e994	[Cherry-pick][v0.18.0.post1][Perf] Optimize Wan2.2 device free on image preprocess (#2854 ) Signed-off-by: fan2956 <zhoufan53@huawei.com>	2026-04-20 14:51:38 +08:00
WeiQing Chen	2116e88638	[cherry-pick][release/v0.18.0.post1] cherry-pick #2847 #2780 #2840 #2876 #2877 (#2878 ) Signed-off-by: david6666666 <530634352@qq.com> Signed-off-by: David Chen <530634352@qq.com> Signed-off-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>	2026-04-20 11:14:12 +08:00
fan2956	783086bcbf	[PERF] Wan2.2 support RMSNorm NPU (#2576 ) Signed-off-by: fan2956 <zhoufan53@huawei.com>	2026-04-18 15:46:37 +08:00
WeiQing ChenandSamit	fa5d358c62	[cherry-pick releasev0.18.0post1] Codex/release v0.18.0.post1 cherry-pick #2735/#2736/#2555 (#2813 ) Signed-off-by: samithuang <285365963@qq.com> Signed-off-by: David Chen <530634352@qq.com> Signed-off-by: Samit <285365963@qq.com> Signed-off-by: david6666666 <530634352@qq.com> Co-authored-by: Samit <285365963@qq.com>	2026-04-17 15:44:45 +08:00
fan2956andgcanlin	f70fd1f58b	[Cherry-pick][v0.18.0.post1][PERF] Wan2.2 support adalayernorm on NPU (#2575 ) Signed-off-by: fan2956 <zhoufan53@huawei.com> Signed-off-by: gcanlin <canlinguosdu@gmail.com> Co-authored-by: gcanlin <canlinguosdu@gmail.com>	2026-04-16 08:47:38 +08:00
fan2956	91fa5c51d7	[Bugfix] fix mindiesd laserattention unsupported error (#2674 ) Signed-off-by: fan2956 <zhoufan53@huawei.com>	2026-04-10 16:19:17 +08:00
WeiQing Chen	a109f777fd	[cherry-pick release/v0.18.0post1][Bugfix] Fix Qwen-Image min-size normalization for tiny requests (#2637 ) (#2648 ) Signed-off-by: David Chen <530634352@qq.com>	2026-04-09 21:04:39 +08:00
bjf-frz	ca7ee8d6e7	[release/v0.18.0.post1][Bugfix] fix oom no response to user (#2638 ) Signed-off-by: bjf-frz <frz123db@gmail.com>	2026-04-09 20:31:58 +08:00
WeiQing Chen	589e664947	[cherry-pick release/v0.18.0post1][Bugfix] Enforce --max-generated-image-size on /v1/images/generations… (#2625 ) Signed-off-by: David Chen <530634352@qq.com>	2026-04-09 14:42:23 +08:00
WeiQing Chen Y. Fisher Canlin Guo	73d31edf63	[cherry-pick][release/v0.18.0post1][Perf] Wan2.2 I2V optimization: convert datatype from FP32 to BF16 in… (#2618 ) Signed-off-by: KexiongYu <yukexiong1@huawei.com> Signed-off-by: Canlin Guo <961750412@qq.com> Signed-off-by: David Chen <530634352@qq.com> Co-authored-by: Y. Fisher <yukexiong1@huawei.com> Co-authored-by: Canlin Guo <961750412@qq.com>	2026-04-09 09:57:14 +08:00
WeiQing Chen Canlin Guo Hongsheng Liu	2f289c78db	[cherry-pick][Feature] Support vae tiling parallel encode (#2368 ) (#2538 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com> Signed-off-by: David Chen <530634352@qq.com> Co-authored-by: Canlin Guo <canlinguosdu@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-04-07 11:16:44 +08:00
WeiQing ChenandCanlin Guo	9b23589c80	[cherry-pick][Perf] Skip Wan2.2 cross attn Ulysses SP (#2459 ) (#2537 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com> Signed-off-by: David Chen <530634352@qq.com> Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>	2026-04-07 10:57:02 +08:00
WeiQing ChenandCanlin Guo	5845eb743c	[cherry-pick][Perf] Optimize Wan2.2 rotary embedding (#2393 ) (#2536 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com> Signed-off-by: David Chen <530634352@qq.com> Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>	2026-04-07 10:56:51 +08:00
WeiQing Chen	300411c9bd	[cherry-pick][release/v0.18.0post1] fix: return 400 for unsupported multi-image edits such as Qw… (#2352 ) Signed-off-by: David Chen <530634352@qq.com>	2026-03-31 11:19:23 +08:00
WeiQing Chen	3485569742	[cherry-pick] [release/v0.18.0post1] fix: validate layered image layers range (#2334 ) (#2349 ) Signed-off-by: David Chen <530634352@qq.com>	2026-03-31 11:18:42 +08:00
WeiQing ChenandCanlin Guo	3edc2a9379	[cherry-pick] [release/v0.18.0.post1]fix: handle Qwen-Image-Layered layered RGBA output for jpeg … (#2329 ) Signed-off-by: David Chen <530634352@qq.com> Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>	2026-03-30 20:41:03 +08:00