[BugFix] Restore vllm_c IR op priority and torch.nn.RMSNorm for Qwen-Image (Same As #5009) (#5017)

Signed-off-by: NumberWan <wantszkin2003@gmail.com>

.buildkite/test-merge.yml

@@ -627,10 +627,11 @@ steps:
       - label: "Diffusion · FLUX.2 Klein Test"
         source_file_dependencies:
           - tests/e2e/offline_inference/test_flux2_klein.py
+          - tests/e2e/online_serving/test_flux2_klein.py
           - vllm_omni/diffusion/models/flux2_klein/
         timeout_in_minutes: 30
         commands:
-          - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+          - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py tests/e2e/online_serving/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
         agents:
           queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
         plugins:

.buildkite/test-nightly.yml

@@ -598,38 +598,6 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-  - group: ":card_index_dividers: Tiny Model Tests [Multi-GPU]"
-    key: nightly-tiny-model-test-group
-    depends_on: upload-nightly-pipeline
-    if: >-
-      build.env("NIGHTLY") == "1" ||
-      build.pull_request.labels includes "nightly-test"
-    steps:
-      - label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations)"
-        timeout_in_minutes: 30
-        # NOTE: Single-GPU (core_model) tiny model tests run on the ready label,
-        # so nightly only adds the multi-GPU parallelism configurations.
-        #
-        # This is a bit confusing, but `full_model` is used to mark the multigpu tests, and
-        # --run-level `core_model` is used to indicate that we should use tiny model weights
-        # for the CI, since these tests should pass irrespective of whether or not we use
-        # the full weights or not.
-        commands:
-          - timeout 25m pytest -sv tests/model_tests/diffusion -m 'full_model and cuda' --run-level "core_model"
-        agents:
-          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              shm-size: "8gb"
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
-
   # Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below.
   - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
     key: nightly-diffusion-x2iat-group
@@ -1319,44 +1287,6 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: ":full_moon: Diffusion X2V · Doc Test"
-        timeout_in_minutes: 60
-        commands:
-          - pytest -s -v tests/examples/offline_inference/test_image_to_video.py -m "full_model and example and H100" --run-level "full_model"
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 1
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
-
       - label: ":full_moon: Diffusion X2V · Other Function Test"
         timeout_in_minutes: 90
         commands:

.buildkite/test-ready.yml

@@ -157,27 +157,6 @@ steps:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
 
-  - group: ":card_index_dividers: Tiny Model Tests [Single GPU]"
-    depends_on: upload-ready-pipeline
-    steps:
-      - label: "Tiny Model Tests · Diffusion (base)"
-        commands:
-          - timeout 15m pytest -sv tests/model_tests/diffusion -m 'core_model and cuda' --run-level "core_model"
-        agents:
-          # Ready label only runs the cases that don't need parallelism for tiny
-          # model tests, which is usually just a smoke test for TTI and I2I
-          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
-
   - label: "Engine&Model Executor Test"
     depends_on: upload-ready-pipeline
     commands:

benchmarks/diffusion/lingbot_video_parity.py

@@ -1,217 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import argparse
-import json
-import re
-import statistics
-import subprocess
-import sys
-import time
-from pathlib import Path
-from typing import Any
-
-import imageio.v3 as iio
-import numpy as np
-import torch
-
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.platforms import current_omni_platform
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="LingBot dense parity and timing harness.")
-    parser.add_argument("--model", default="/home/models/lingbot-video-dense-1.3b")
-    parser.add_argument("--official-repo", default="/tmp/lingbot-video")
-    parser.add_argument("--output-dir", default="/tmp/lingbot_dense_parity")
-    parser.add_argument("--prompt", default="a robotic arm picks up a red block")
-    parser.add_argument("--height", type=int, default=192)
-    parser.add_argument("--width", type=int, default=320)
-    parser.add_argument("--num-frames", type=int, default=9)
-    parser.add_argument("--steps", type=int, default=2)
-    parser.add_argument("--guidance-scale", type=float, default=3.0)
-    parser.add_argument("--shift", type=float, default=3.0)
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--fps", type=int, default=24)
-    parser.add_argument("--runs", type=int, default=0, help="Optional in-process steady-state native runs.")
-    return parser.parse_args()
-
-
-def _run(cmd: list[str], cwd: Path | None = None) -> tuple[float, str]:
-    start = time.perf_counter()
-    proc = subprocess.run(cmd, cwd=cwd, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
-    elapsed = time.perf_counter() - start
-    if proc.returncode != 0:
-        print(proc.stdout)
-        raise subprocess.CalledProcessError(proc.returncode, proc.args, output=proc.stdout)
-    return elapsed, proc.stdout
-
-
-def _read_video(path: Path) -> np.ndarray:
-    frames = iio.imread(path, plugin="pyav")
-    if frames.dtype != np.float32:
-        frames = frames.astype(np.float32) / 255.0
-    return frames
-
-
-def _compare(a: np.ndarray, b: np.ndarray) -> dict[str, Any]:
-    if a.shape != b.shape:
-        return {"shape_a": list(a.shape), "shape_b": list(b.shape), "shape_match": False}
-    diff = a - b
-    mse = float(np.mean(diff * diff))
-    mae = float(np.mean(np.abs(diff)))
-    psnr = float("inf") if mse == 0 else float(20.0 * np.log10(1.0 / np.sqrt(mse)))
-    return {
-        "shape": list(a.shape),
-        "shape_match": True,
-        "mae": mae,
-        "mse": mse,
-        "psnr": psnr,
-    }
-
-
-def _extract_total_generation_seconds(log: str) -> float | None:
-    match = re.search(r"Total generation time:\s*([0-9.]+)\s*seconds", log)
-    return float(match.group(1)) if match else None
-
-
-def _peak_memory_mb(output: Any) -> float:
-    result = output[0] if isinstance(output, list) and output else output
-    return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
-
-
-def _benchmark_native_requests(args: argparse.Namespace) -> dict[str, Any]:
-    load_start = time.perf_counter()
-    omni = Omni(
-        model=args.model,
-        model_class_name="LingBotVideoPipeline",
-        flow_shift=args.shift,
-        parallel_config=DiffusionParallelConfig(),
-    )
-    load_seconds = time.perf_counter() - load_start
-    prompt = {"prompt": args.prompt}
-    latencies: list[float] = []
-    peak_mb = 0.0
-
-    for _ in range(args.runs):
-        generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
-        sampling_params = OmniDiffusionSamplingParams(
-            height=args.height,
-            width=args.width,
-            num_frames=args.num_frames,
-            num_inference_steps=args.steps,
-            guidance_scale=args.guidance_scale,
-            generator=generator,
-            output_type="pt",
-            extra_args={"shift": args.shift},
-        )
-        start = time.perf_counter()
-        output = omni.generate(prompt, sampling_params)
-        latencies.append(time.perf_counter() - start)
-        peak_mb = max(peak_mb, _peak_memory_mb(output))
-        del output
-    omni.close()
-
-    return {
-        "load_seconds": load_seconds,
-        "run_seconds": latencies,
-        "mean_seconds": statistics.mean(latencies) if latencies else None,
-        "median_seconds": statistics.median(latencies) if latencies else None,
-        "min_seconds": min(latencies) if latencies else None,
-        "max_seconds": max(latencies) if latencies else None,
-        "peak_mb": peak_mb,
-    }
-
-
-def main() -> None:
-    args = parse_args()
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-    official_out = output_dir / "official_diffusers.mp4"
-    native_out = output_dir / "vllm_omni_native.mp4"
-
-    official_cmd = [
-        sys.executable,
-        str(Path(args.official_repo) / "scripts" / "inference.py"),
-        "--backend",
-        "diffusers",
-        "--model_dir",
-        args.model,
-        "--mode",
-        "t2v",
-        "--prompt",
-        args.prompt,
-        "--output",
-        str(official_out),
-        "--height",
-        str(args.height),
-        "--width",
-        str(args.width),
-        "--num_frames",
-        str(args.num_frames),
-        "--steps",
-        str(args.steps),
-        "--guidance_scale",
-        str(args.guidance_scale),
-        "--shift",
-        str(args.shift),
-        "--seed",
-        str(args.seed),
-        "--fps",
-        str(args.fps),
-        "--quiet_progress",
-    ]
-    native_cmd = [
-        sys.executable,
-        "examples/offline_inference/text_to_video/text_to_video_lingbot.py",
-        "--model",
-        args.model,
-        "--prompt",
-        args.prompt,
-        "--output",
-        str(native_out),
-        "--height",
-        str(args.height),
-        "--width",
-        str(args.width),
-        "--num-frames",
-        str(args.num_frames),
-        "--num-inference-steps",
-        str(args.steps),
-        "--guidance-scale",
-        str(args.guidance_scale),
-        "--flow-shift",
-        str(args.shift),
-        "--seed",
-        str(args.seed),
-        "--fps",
-        str(args.fps),
-    ]
-
-    official_time, official_log = _run(official_cmd, cwd=Path(args.official_repo))
-    native_time, native_log = _run(native_cmd, cwd=Path.cwd())
-
-    result = {
-        "official_seconds": official_time,
-        "native_process_seconds": native_time,
-        "native_request_seconds": _extract_total_generation_seconds(native_log),
-        "official_output": str(official_out),
-        "native_output": str(native_out),
-        "accuracy": _compare(_read_video(official_out), _read_video(native_out)),
-        "official_log_tail": official_log.splitlines()[-20:],
-        "native_log_tail": native_log.splitlines()[-20:],
-    }
-    if args.runs > 0:
-        result["native_steady_state"] = _benchmark_native_requests(args)
-
-    result_path = output_dir / "comparison.json"
-    result_path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
-    print(json.dumps(result, indent=2, sort_keys=True))
-
-
-if __name__ == "__main__":
-    main()

docs/models/supported_models.md

@@ -46,7 +46,6 @@ th {
 | `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
 | `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
-| `LingBotVideoPipeline` | LingBot-Video dense T2V | `robbyant/lingbot-video-dense-1.3b` | ✅︎ | | | |
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
 | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |

docs/user_guide/diffusion/attention_backends.md

@@ -20,8 +20,6 @@ The full set of backends and their platform defaults is in the **Backend Options
 | `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
 | `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
 | `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
-| `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. |
-| `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. |
 
 
 ## Configuration
@@ -212,36 +210,6 @@ Notes:
 - `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class.
 - SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness.
 
-## HuggingFace Kernels Hub Backends
-
-To achieve perfect numerical consistency between **training** (typically using HuggingFace Diffusers with Hub kernels) and **serving/rollout** (using vLLM-Omni), you can use the Hub-based attention backends. This eliminates numerical drift / sampling divergence caused by executing different local kernel versions during rollout.
-
-The following backend options are supported:
-- `FLASH_ATTN_HUB` (HuggingFace `kernels-community/flash-attn2`)
-- `FLASH_ATTN_3_HUB` (HuggingFace `kernels-community/flash-attn3`, Hopper sm_90+ only)
-
-### Installation
-
-To use these backends, you must install the `kernels` library:
-
-```bash
-pip install kernels==0.14.1
-```
-
-If the `kernels` library is not available in the environment, vLLM-Omni will log a warning and fall back gracefully to the corresponding local backend implementations (`FLASH_ATTN`). On CUDA GPUs below Hopper (compute capability < 9.0), `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
-
-### Usage
-
-Select a Hub backend using the global CLI flag or environment variables:
-
-```bash
-# Environment variable
-export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
-
-# CLI flag
-vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN_HUB
-```
-
 ## Usage Examples
 
 ### Default (auto-route)

docs/user_guide/diffusion_features.md

@@ -126,7 +126,7 @@ The following tables show which models support each feature:
 | **MagiHuman**            |     ❌     |     ❌      |           ❌           |       ❓        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **MammothModa2(T2I)**    |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **OmniGen2**             |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **OmniGen2**             |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Qwen-Image**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |
 | **Qwen-Image-2512**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |

docs/user_guide/examples/offline_inference/mammothmodal2_preview.md

@@ -22,7 +22,7 @@ Text-to-image now runs through the shared offline image example
 ```bash
 python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
   --model ./MammothModa2-Preview \
-  --deploy-config ./vllm_omni/deploy/mammoth_moda2_ar.yaml \
+  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
   --question "Summarize this image." \
   --image ./image.png
 ```

examples/offline_inference/image_to_image/image_edit.py

@@ -138,15 +138,6 @@ def parse_args() -> argparse.Namespace:
         required=True,
         help="Path(s) to input image file(s) (PNG, JPG, etc.). Can specify multiple images.",
     )
-    parser.add_argument(
-        "--deploy-config",
-        type=str,
-        default=None,
-        help=(
-            "Path to a deploy YAML. Required for multi-stage image-edit pipelines "
-            "whose deploy config is not auto-loaded."
-        ),
-    )
     parser.add_argument(
         "--prompt",
         type=str,
@@ -358,12 +349,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--enforce-eager",
         action="store_true",
-        default=None,
-        help=(
-            "Disable torch.compile and force eager execution. Left unset (None) "
-            "so it is only forwarded when explicitly given; "
-            "otherwise the per-stage deploy YAML value wins."
-        ),
+        help="Disable torch.compile and force eager execution.",
     )
     parser.add_argument(
         "--vae-use-slicing",
@@ -482,7 +468,7 @@ def main():
         }
 
     # Initialize Omni with appropriate pipeline
-    omni_kwargs: dict[str, Any] = dict(
+    omni = Omni(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
         vae_use_slicing=args.vae_use_slicing,
@@ -490,15 +476,11 @@ def main():
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         parallel_config=parallel_config,
+        enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         profiler_config=args.profiler_config,
     )
-    if args.enforce_eager is not None:
-        omni_kwargs["enforce_eager"] = args.enforce_eager
-    if args.deploy_config:
-        omni_kwargs["deploy_config"] = args.deploy_config
-    omni = Omni(**omni_kwargs)
     model_class_name = get_model_class_name(omni)
     declared_extra_body_params = get_extra_body_params(model_class_name)
     print("Pipeline loaded")

examples/offline_inference/image_to_video/README.md

@@ -3,140 +3,14 @@
 This shared example generates videos from images with VACE, Wan2.2,
 LTX-2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
 
-- `image_to_video.py`: command-line script for single video generation with advanced options.
+## Local CLI Usage
 
-## Table of Contents
-
-- [Overview](#overview)
-- [Prerequisites](#prerequisites)
-- [Quick Start](#quick-start)
-- [Key Arguments](#key-arguments)
-- [More CLI Examples](#more-cli-examples)
-- [Advanced Features](#advanced-features)
-- [FAQ](#faq)
-
-## Overview
-
-This folder provides a unified CLI script for image-to-video generation using vLLM-Omni diffusion/video pipelines. The script selects practical defaults for supported model families while still exposing common sampling, memory, and parallelism options.
-
-### Supported Models
-
-| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM Notes |
-| ----- | ------------------ | -------------- | ------------- | -------- | ---------- |
-| `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480 x 832 | 81 | 50 | 5.0 | Around 60 GiB BF16 for basic single-card usage |
-| `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 480 x 832 | 81 | 50 | 4.0 | Around 20–25 GiB BF16, smallest I2V model |
-| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v` | 480 x 832 | 121 | 50 | 6.0 | Around 100 GiB at default settings; the example enables `--enable-cpu-offload` + VAE tiling/slicing to fit an 80 GiB card |
-| LTX2 (local path + `--model-class-name LTX2ImageToVideoPipeline`) | 512 x 768 | 121 | 40 | 4.0 | Memory use depends on frame count and tensor parallelism |
-
-!!! info
-    Peak VRAM: based on basic single-card usage, batch size = 1, without any acceleration/optimization features. Some model weights cannot fit into one card with 80 GiB VRAM, which may need to use CPU offloading.
-
-Default model: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`.
-
-## Prerequisites
-
-Download the example image used in the snippets below:
+Download the example image:
 
 ```bash
-curl -L -o cherry_blossom.jpg https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
+wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
 ```
 
-## Quick Start
-
-### Python API
-
-Single-prompt generation using TI2V-5B (lightweight):
-
-```python
-import PIL.Image
-import torch
-
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-if __name__ == "__main__":
-    image = PIL.Image.open("cherry_blossom.jpg").convert("RGB")
-    image = image.resize((576, 320))
-
-    omni = Omni(
-        model="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
-        flow_shift=12.0,
-    )
-    outputs = omni.generate(
-        {
-            "prompt": "Cherry blossoms swaying gently in the breeze, petals falling",
-            "multi_modal_data": {"image": image},
-        },
-        OmniDiffusionSamplingParams(
-            height=320,
-            width=576,
-            num_frames=17,
-            num_inference_steps=20,
-            guidance_scale=4.0,
-            generator=torch.Generator(device="cuda").manual_seed(42),
-        ),
-    )
-    from diffusers.utils import export_to_video
-    frames = outputs[0].request_output.images
-    export_to_video(frames, "quick_test_i2v.mp4", fps=16)
-```
-
-### Local CLI Usage
-
-Quick test using TI2V-5B with a small resolution and few frames:
-
-```bash
-python image_to_video.py \
-  --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
-  --image cherry_blossom.jpg \
-  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --height 320 \
-  --width 576 \
-  --num-frames 17 \
-  --guidance-scale 4.0 \
-  --num-inference-steps 20 \
-  --flow-shift 12.0 \
-  --fps 16 \
-  --output quick_test_i2v.mp4
-```
-
-## Key Arguments
-
-| Argument | Type | Default | Description |
-| -------- | ---- | ------- | ----------- |
-| `--model` | str | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Diffusers I2V model ID or local path |
-| `--model-class-name` | str | `None` | Override model class name (e.g., `LTX2ImageToVideoPipeline`) |
-| `--image` | str | (required) | Path to input image |
-| `--prompt` | str | `""` | Text description of desired motion/animation |
-| `--negative-prompt` | str | `""` | Optional list of artifacts to suppress |
-| `--seed` | int | `42` | Random seed for deterministic sampling |
-| `--guidance-scale` | float | `5.0` | CFG scale |
-| `--guidance-scale-high` | float | `None` | Separate CFG for high-noise stage (MoE only) |
-| `--height` | int | auto | Video height (auto-calculated from image if not set). Multiples of 16 |
-| `--width` | int | auto | Video width (auto-calculated from image if not set). Multiples of 16 |
-| `--num-frames` | int | `81` | Number of frames |
-| `--num-inference-steps` | int | `50` | Number of denoising steps |
-| `--boundary-ratio` | float | `0.875` | Boundary split ratio for two-stage MoE models |
-| `--flow-shift` | float | `5.0` | Scheduler flow shift (5.0 for 720p, 12.0 for 480p) |
-| `--sample-solver` | str | `unipc` | Wan2.2 sampling solver (`unipc` or `euler` for Lightning/Distill) |
-| `--fps` | int | `None` | Frames per second for the saved MP4 |
-| `--frame-rate` | float | `None` | Generation frame rate for pipelines that require it (e.g., LTX2) |
-| `--output` | str | `i2v_output.mp4` | Path to save the generated video |
-| `--vae-use-slicing` | flag | off | Enable VAE slicing for memory optimization |
-| `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization |
-| `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
-| `--enable-layerwise-offload` | flag | off | Enable layerwise offloading on DiT modules |
-| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
-| `--tensor-parallel-size` | int | `1` | Tensor parallel size (effective for models that support TP, e.g. LTX2) |
-| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree |
-| `--ring-degree` | int | `1` | Ring sequence parallel degree |
-| `--cache-backend` | str | `None` | Cache backend: `cache_dit` or `tea_cache` |
-| `--use-hsdp` | flag | off | Enable Hybrid Sharded Data Parallel |
-| `--hsdp-shard-size` | int | `-1` | GPUs per shard group (-1 auto-calculates) |
-| `--hsdp-replicate-size` | int | `1` | Number of replica groups for HSDP |
-
-## More CLI Examples
-
 ### Wan2.2-I2V-A14B-Diffusers (MoE)
 
 ```bash
@@ -144,7 +18,7 @@ python image_to_video.py \
   --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
   --image cherry_blossom.jpg \
   --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --negative-prompt "low quality, blurry" \
+  --negative-prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
   --num-frames 48 \
@@ -154,7 +28,7 @@ python image_to_video.py \
   --boundary-ratio 0.875 \
   --flow-shift 12.0 \
   --fps 16 \
-  --output i2v_wan_moe.mp4
+  --output i2v_output.mp4
 ```
 
 ### Wan2.2-TI2V-5B-Diffusers (Unified)
@@ -164,7 +38,7 @@ python image_to_video.py \
   --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
   --image cherry_blossom.jpg \
   --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --negative-prompt "low quality, blurry" \
+  --negative-prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
   --num-frames 48 \
@@ -172,65 +46,7 @@ python image_to_video.py \
   --num-inference-steps 40 \
   --flow-shift 12.0 \
   --fps 16 \
-  --output i2v_wan_ti2v.mp4
-```
-
-### HunyuanVideo-1.5 I2V (480p)
-
-```bash
-python image_to_video.py \
-  --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \
-  --image cherry_blossom.jpg \
-  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --height 480 \
-  --width 832 \
-  --num-frames 121 \
-  --guidance-scale 6.0 \
-  --flow-shift 5.0 \
-  --num-inference-steps 50 \
-  --fps 24 \
-  --enable-cpu-offload \
-  --vae-use-tiling \
-  --vae-use-slicing \
-  --output hunyuan_i2v.mp4
-```
-
-### LTX2 Image-to-Video
-
-```bash
-python image_to_video.py \
-  --model /path/to/LTX-2 \
-  --model-class-name LTX2ImageToVideoPipeline \
-  --image cherry_blossom.jpg \
-  --prompt "A cinematic dolly shot of cherry blossoms" \
-  --height 512 \
-  --width 768 \
-  --num-frames 121 \
-  --num-inference-steps 40 \
-  --guidance-scale 4.0 \
-  --frame-rate 24 \
-  --fps 24 \
-  --output ltx2_i2v.mp4
-```
-
-## Advanced Features
-
-### CFG Parallel
-
-Set `--cfg-parallel-size 2` to enable CFG Parallel for faster inference on multi-GPU setups.
-See more examples in the [cfg_parallel user guide](../../../docs/user_guide/diffusion/parallelism/cfg_parallel.md).
-
-### Cache Acceleration
-
-Use `--cache-backend cache_dit` for Cache-DiT acceleration or `--cache-backend tea_cache` for Timestep Embedding Aware Cache:
-
-```bash
-python image_to_video.py \
-  --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
-  --image cherry_blossom.jpg \
-  --prompt "Cherry blossoms swaying gently in the breeze" \
-  --cache-backend cache_dit \
-  --output i2v_cached.mp4
+  --output i2v_output.mp4
 ```
 
 ### LTX-2.3
@@ -404,11 +220,3 @@ The VACE T2V command is documented in the shared
 
 For Wan2.2 LightX2V-converted local Diffusers directories and related LoRA
 assets, see the [LoRA guide](../../../docs/user_guide/diffusion/lora.md#wan22-lightx2v-offline-assembly).
-
-## FAQ
-
-**OOM errors**: Try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage. For very large models, add `--enable-cpu-offload` or `--enable-layerwise-offload`.
-
-**Auto-calculated resolution**: If `--height` and `--width` are not provided, the script calculates output dimensions from the input image while maintaining aspect ratio and targeting 480 x 832 area (or 512 x 768 for LTX2).
-
-**Wan2.2 MoE vs unified**: `I2V-A14B` is a larger Mixture-of-Experts model with separate low/high-noise DiT stages (use `--guidance-scale-high` and `--boundary-ratio`). `TI2V-5B` is a smaller unified T2V+I2V model that does not need these extra arguments.

examples/offline_inference/mammothmodal2_preview/README.md

@@ -19,7 +19,7 @@ hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-P
 ```bash
 python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
   --model ./MammothModa2-Preview \
-  --deploy-config ./vllm_omni/deploy/mammoth_moda2_ar.yaml \
+  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
   --question "Summarize this image." \
   --image ./image.png
 ```

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -4,7 +4,7 @@ Offline inference example: MammothModa2 image summarization (single AR stage).
 Example:
   uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
     --model path/to/MammothModa2-Preview \
-    --deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \
+    --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
     --image /path/to/input.jpg \
     --question "Please summarize the content of this image."
 """
@@ -27,7 +27,9 @@ DEFAULT_QUESTION = "Please summarize the content of this image."
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
     parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
-    parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.")
+    parser.add_argument(
+        "--stage-config", type=str, required=True, help="Path to stage config yaml (single-stage AR->text)."
+    )
     parser.add_argument("--image", type=str, required=True, help="Path to input image.")
     parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
     parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
@@ -71,7 +73,7 @@ def main() -> None:
 
     omni = Omni(
         model=args.model,
-        deploy_config=args.deploy_config,
+        stage_configs_path=args.stage_config,
         trust_remote_code=args.trust_remote_code,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )

examples/offline_inference/ming_flash_omni/README.md

@@ -90,40 +90,14 @@ Pass `--deploy-config /path/to/your_deploy.yaml` to any of the commands above to
 
 ### Image generation (text-to-image / img2img)
 
-Image generation is served through the standard task examples.
-The diffusion-stage knobs are declared centrally in `vllm_omni/model_extras/ming_flash_omni.py` and routed via `--extra-body`. With deploy yaml assigned properly, the model can run through the shared example scripts like any other diffusion model:
-
-Text-to-image (offline):
-
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-    --model Jonathan1909/Ming-flash-omni-2.0 \
-    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
-    --prompt "Please draw a cute cat." \
-    --height 1024 \
-    --width 1024 \
-    --extra-body '{"steps": 30, "cfg": 2.0, "seed": 42}' \
-    --output ming_flash_omni_t2i.png
-```
-
-Image-to-image (offline):
-```bash
-# Reference image: figures/cases/person_gen_05.png from the upstream Ming repo
-# https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/examples/vllm_demo.py
-wget https://raw.githubusercontent.com/inclusionAI/Ming/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/person_gen_05.png
-
-python examples/offline_inference/image_to_image/image_edit.py \
-    --model Jonathan1909/Ming-flash-omni-2.0 \
-    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
-    --image person_gen_05.png \
-    --prompt "Put a pair of sunglasses on the person." \
-    --extra-args '{"steps": 30, "cfg": 2.0, "seed": 42}' \
-    --output ming_flash_omni_i2i.png
-```
-
-For the online path and the full knob list
-(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`/`height`/`width`), see the
-[image-generation section in the recipe](../../../recipes/inclusionAI/Ming-flash-omni-2.0.md#image-generation-text-to-image--img2img).
+The diffusion-stage image-generation path is currently only wired through
+the online OpenAI-compatible chat endpoint (`/v1/chat/completions` with
+`"modalities": ["image"]`). For payloads, optional knobs
+(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`), and the img2img
+reference-image flow, see the
+[image-gen section in the online-serving README](../../online_serving/ming_flash_omni/README.md#image-generation-text-to-image).
+`end2end.py` does not yet exercise the imagegen stage; that is tracked as
+follow-up work.
 
 ## Online serving
 

examples/offline_inference/text_to_image/text_to_image.py

@@ -69,16 +69,7 @@ def parse_args() -> argparse.Namespace:
         "--stage-configs-path",
         type=str,
         default=None,
-        help="[Deprecated] Path to a legacy stage_args-format YAML. Prefer --deploy-config.",
-    )
-    parser.add_argument(
-        "--deploy-config",
-        type=str,
-        default=None,
-        help=(
-            "Path to a deploy YAML (new pipeline/stages format). Required for multi-stage "
-            "text-to-image pipelines whose deploy config is not auto-loaded."
-        ),
+        help="Path to a YAML file containing stage configurations for Omni.",
     )
     parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.")
     parser.add_argument(
@@ -164,12 +155,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--enforce-eager",
         action="store_true",
-        default=None,
-        help=(
-            "Disable torch.compile and force eager execution. Left unset (None) "
-            "so it is only forwarded when explicitly given; "
-            "otherwise the per-stage deploy YAML value wins."
-        ),
+        help="Disable torch.compile and force eager execution.",
     )
     parser.add_argument(
         "--enable-cpu-offload",
@@ -235,13 +221,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--tensor-parallel-size",
         type=int,
-        default=None,
-        help=(
-            "Number of GPUs used for tensor parallelism (TP) inside the DiT. "
-            "Left unset so it is only forwarded when explicitly given; "
-            "otherwise the per-stage deploy YAML (or engine default of 1) wins. "
-            "Passing it always overrides the deploy YAML."
-        ),
+        default=1,
+        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
     )
     parser.add_argument(
         "--enable-expert-parallel",
@@ -429,8 +410,10 @@ def main():
         "ring_degree": args.ring_degree,
         "ulysses_mode": args.ulysses_mode,
         "cfg_parallel_size": args.cfg_parallel_size,
+        "tensor_parallel_size": args.tensor_parallel_size,
         "vae_patch_parallel_size": args.vae_patch_parallel_size,
         "enable_expert_parallel": args.enable_expert_parallel,
+        "enforce_eager": args.enforce_eager,
         "enable_cpu_offload": args.enable_cpu_offload,
         "mode": "text-to-image",
         "log_stats": args.log_stats,
@@ -442,14 +425,8 @@ def main():
         **lora_args,
         **quant_kwargs,
     }
-    if args.tensor_parallel_size is not None:
-        omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
-    if args.enforce_eager is not None:
-        omni_kwargs["enforce_eager"] = args.enforce_eager
     if args.stage_configs_path:
         omni_kwargs["stage_configs_path"] = args.stage_configs_path
-    if args.deploy_config:
-        omni_kwargs["deploy_config"] = args.deploy_config
     if use_nextstep:
         # NextStep-1.1 requires explicit pipeline class
         omni_kwargs["model_class_name"] = "NextStep11Pipeline"
@@ -474,9 +451,8 @@ def main():
     print(f"  Quantization: {args.quantization if args.quantization else 'None (BF16)'}")
     if ignored_layers:
         print(f"  Ignored layers: {ignored_layers}")
-    tp_display = args.tensor_parallel_size if args.tensor_parallel_size is not None else "deploy/default"
     print(
-        f"  Parallel configuration: tensor_parallel_size={tp_display}, "
+        f"  Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, "
         f"ulysses_degree={args.ulysses_degree}, ulysses_mode={args.ulysses_mode}, "
         f"ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, "
         f"vae_patch_parallel_size={args.vae_patch_parallel_size}, "
@@ -546,8 +522,8 @@ def main():
         diffusion_params.extra_args.update({k: v for k, v in user_extra.items() if v is not None})
 
     if lora_request:
-        diffusion_params.lora_request = lora_request
-        diffusion_params.lora_scale = args.lora_scale
+        diffusion_params.extra_args["lora_request"] = lora_request
+        diffusion_params.extra_args["lora_scale"] = args.lora_scale
 
     # Build per-stage sampling params for multi-stage models (e.g. BAGEL),
     # or wrap single diffusion params for single-stage models.

examples/offline_inference/text_to_speech/omnivoice/end2end.py

@@ -34,7 +34,7 @@ def run_e2e():
         help="Model name or path (HuggingFace or local)",
     )
     parser.add_argument(
-        "--deploy-config",
+        "--stage-config",
         type=str,
         default="vllm_omni/deploy/omnivoice.yaml",
     )
@@ -87,14 +87,14 @@ def run_e2e():
     )
     args = parser.parse_args()
 
-    if not os.path.exists(args.deploy_config):
-        raise FileNotFoundError(f"Deploy config not found: {args.deploy_config}")
+    if not os.path.exists(args.stage_config):
+        raise FileNotFoundError(f"Stage config not found: {args.stage_config}")
 
     print(f"Initializing OmniVoice with model={args.model}")
 
     omni = Omni(
         model=args.model,
-        deploy_config=args.deploy_config,
+        stage_configs_path=args.stage_config,
         log_stats=True,
     )
 

examples/offline_inference/text_to_video/text_to_video_lingbot.py

@@ -1,124 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import argparse
-import time
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import torch
-
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.platforms import current_omni_platform
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Generate text-to-video with the LingBot dense preset.")
-    parser.add_argument("--model", default="/home/models/lingbot-video-dense-1.3b")
-    parser.add_argument("--prompt", default="a robotic arm picks up a red block")
-    parser.add_argument("--negative-prompt", default="")
-    parser.add_argument("--output", default="lingbot_video_output.mp4")
-    parser.add_argument("--height", type=int, default=192)
-    parser.add_argument("--width", type=int, default=320)
-    parser.add_argument("--num-frames", type=int, default=9)
-    parser.add_argument("--num-inference-steps", type=int, default=2)
-    parser.add_argument("--guidance-scale", type=float, default=3.0)
-    parser.add_argument("--flow-shift", type=float, default=3.0)
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--fps", type=int, default=24)
-    parser.add_argument("--output-type", choices=["pt", "np"], default="pt")
-    parser.add_argument("--batch-cfg", action="store_true")
-    return parser.parse_args()
-
-
-def _peak_memory_mb(output: Any) -> float:
-    result = output[0] if isinstance(output, list) and output else output
-    return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
-
-
-def _extract_video(output: Any) -> Any:
-    if isinstance(output, list):
-        output = output[0] if output else None
-    if isinstance(output, OmniRequestOutput):
-        if output.final_output_type != "image":
-            raise ValueError(f"Unexpected output type {output.final_output_type!r}.")
-        if not output.images:
-            raise ValueError("No video frames found in OmniRequestOutput.")
-        return output.images[0] if len(output.images) == 1 else output.images
-    return output
-
-
-def _to_frame_array(video: Any) -> np.ndarray:
-    if isinstance(video, torch.Tensor):
-        tensor = video.detach().cpu()
-        if tensor.dim() == 5:
-            tensor = tensor[0] if tensor.shape[1] not in (3, 4) else tensor[0].permute(1, 2, 3, 0)
-        elif tensor.dim() == 4 and tensor.shape[0] in (3, 4):
-            tensor = tensor.permute(1, 2, 3, 0)
-        if tensor.is_floating_point():
-            if float(tensor.min()) < 0.0:
-                tensor = tensor.clamp(-1, 1) * 0.5 + 0.5
-            else:
-                tensor = tensor.clamp(0, 1)
-        return tensor.float().numpy()
-    if isinstance(video, np.ndarray):
-        array = video[0] if video.ndim == 5 else video
-        if np.issubdtype(array.dtype, np.integer):
-            return array.astype(np.float32) / 255.0
-        return array.astype(np.float32, copy=False)
-    if isinstance(video, list):
-        return np.stack([_to_frame_array(frame) for frame in video], axis=0)
-    return np.asarray(video, dtype=np.float32)
-
-
-def main() -> None:
-    args = parse_args()
-    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
-    omni = Omni(
-        model=args.model,
-        model_class_name="LingBotVideoPipeline",
-        flow_shift=args.flow_shift,
-        parallel_config=DiffusionParallelConfig(),
-    )
-
-    prompt: dict[str, str] = {"prompt": args.prompt}
-    if args.negative_prompt:
-        prompt["negative_prompt"] = args.negative_prompt
-
-    sampling_params = OmniDiffusionSamplingParams(
-        height=args.height,
-        width=args.width,
-        num_frames=args.num_frames,
-        num_inference_steps=args.num_inference_steps,
-        guidance_scale=args.guidance_scale,
-        generator=generator,
-        output_type=args.output_type,
-        extra_args={"batch_cfg": args.batch_cfg, "shift": args.flow_shift},
-    )
-
-    start = time.perf_counter()
-    output = omni.generate(prompt, sampling_params)
-    elapsed = time.perf_counter() - start
-    print(f"Total generation time: {elapsed:.4f} seconds ({elapsed * 1000:.2f} ms)")
-    peak_mb = _peak_memory_mb(output)
-    if peak_mb:
-        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
-
-    video = _to_frame_array(_extract_video(output))
-    output_path = Path(args.output)
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-    from diffusers.utils import export_to_video
-
-    export_to_video(list(video), str(output_path), fps=args.fps)
-    print(f"Saved generated video to {output_path}")
-    omni.close()
-
-
-if __name__ == "__main__":
-    main()

examples/online_serving/ming_flash_omni/README.md

@@ -49,14 +49,22 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene
     --modalities text
 ```
 
+Parameterized curl wrapper in this directory:
 
-## Image generation (text-to-image / img2img)
+```bash
+bash run_curl_multimodal_generation.sh text
+bash run_curl_multimodal_generation.sh use_image
+bash run_curl_multimodal_generation.sh use_audio
+bash run_curl_multimodal_generation.sh use_video
+bash run_curl_multimodal_generation.sh use_mixed_modalities
+bash run_curl_multimodal_generation.sh use_image_gen
+```
 
-Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage. Launch with the image deploy YAML, which adds an image-gen stage behind the thinker.
+## Image generation (text-to-image)
 
-The image-generation stage is a standard vLLM-Omni diffusion pipeline (`MingImagePipeline`); its request knobs are declared in `vllm_omni/model_extras/ming_flash_omni.py` and routed through `extra_body`, so they no longer need a bespoke `sampling_params_list` recipe (that form is still available for per-stage thinker sampling — see below).
-
-### Launch
+Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage.
+Launch with the image deploy YAML, which adds an image-gen stage behind
+the thinker:
 
 ```bash
 vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
@@ -66,58 +74,32 @@ vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
     --port 8091
 ```
 
-
-### Text-to-image
-
-Request image output with `"modalities": ["image"]`:
+Then request image output by passing `"modalities": ["image"]`:
 
 ```bash
 curl -s http://127.0.0.1:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "Jonathan1909/Ming-flash-omni-2.0",
-    "messages": [
-      {
-        "role": "user",
-        "content": "Please draw a cute cat."
-      }
-    ],
-    "modalities": ["image"]
-  }' \
-  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
-  | base64 -d > ming_imagegen.png
+    -H "Content-Type: application/json" \
+    -d '{
+      "model": "Jonathan1909/Ming-flash-omni-2.0",
+      "messages": [{"role": "user", "content": "Please draw a cute cat."}],
+      "modalities": ["image"]
+    }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen.png
 ```
 
-Pass generation knobs under a literal `extra_body` object (the OpenAI Python client's `extra_body=` kwarg produces the same request). Keys are filtered against the declared set and routed into every stage's `extra_args`:
+### Optional knobs
 
-```bash
-curl http://127.0.0.1:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "Jonathan1909/Ming-flash-omni-2.0",
-    "modalities": ["image"],
-    "messages": [
-      {
-        "role": "user",
-        "content": "Draw a poster."
-      }
-    ],
-    "extra_body": {
-      "steps": 6,
-      "cfg": 1.5,
-      "height": 512,
-      "width": 512,
-      "seed": 123,
-      "byte5_text": ["理解与生成统一"],
-      "negative_prompt": "ugly, blurry, distorted"
-    }
-  }' \
-  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
-  | base64 -d > ming_imagegen_extra_body.png
-```
+Pass image-gen overrides as flat keys on the diffusion-stage `sampling_params_list[1].extra_args`:
 
-NOTE: `extra_body` does **not** set the thinker's own sampling params. To tune the thinker (stage-0) sampling (`temperature` / `top_p` / `top_k` / `max_tokens`) or place knobs explicitly per stage, use `sampling_params_list` (`[thinker, imagegen]`).
-`negative_prompt` must sit on the **stage-0** entry to trigger the real-CFG companion; the imagegen knobs go on **stage-1**:
+| Key | Default | Description |
+| --- | --- | --- |
+| `height` / `width` | from config (1024) | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). |
+| `steps` | 30 | Number of FlowMatchEuler denoise steps. |
+| `cfg` | 2.0 | Classifier-free guidance scale. |
+| `seed` | 42 | Per-request RNG seed (deterministic when ByT5 is also seed-stable). |
+| `byte5_text` | (auto) | Override the glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "...". ` format. |
+| `negative_prompt` | empty | Real CFG negative conditioning (set on **stage-0 thinker** `extra_args` so `expand_cfg_prompts` spawns the companion). |
+
+Example with all the knobs:
 
 ```bash
 curl http://127.0.0.1:8091/v1/chat/completions \
@@ -161,48 +143,27 @@ curl http://127.0.0.1:8091/v1/chat/completions \
 
 ### img2img (reference image)
 
-Add an `image_url` content part (a base64 data URL) to the user message; it is routed into the DiT stage as `extra[reference_image]`. base64-encode a local file and stream it through `jq` via stdin - piping `base64 -> jq -> curl` so that avoids the shell `ARG_MAX` limit that inlining a large base64 string in `-d '…'` hits.
+Add an `image_url` content part to the user message; the parser routes it
+into the diffusion stage as `extra[reference_image]`:
 
-```bash
-# Reference image: figures/cases/person_gen_05.png from the upstream Ming repo
-# Check https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/examples/vllm_demo.py
-wget https://raw.githubusercontent.com/inclusionAI/Ming/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/person_gen_05.png
-
-base64 -w0 person_gen_05.png \
-| jq -Rs --arg prompt "Put a pair of sunglasses on the person." '{
-    model: "Jonathan1909/Ming-flash-omni-2.0",
-    modalities: ["image"],
-    messages: [
-      {
-        role: "user",
-        content: [
-          { type: "text", text: $prompt },
-          { type: "image_url", image_url: { url: ("data:image/png;base64," + (. | rtrimstr("\n"))) } }
-        ]
-      }
-    ]
-  }' \
-| curl http://127.0.0.1:8091/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d @- \
-| jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
-| base64 -d > ming_img2img.png
+```jsonc
+"messages": [{
+  "role": "user",
+  "content": [
+    {"type": "text", "text": "Change the background to a sandy beach at sunset."},
+    {"type": "image_url", "image_url": {"url": "data:image/png;base64,<base64>"}}
+  ]
+}]
 ```
 
-The reference image can also be a public URL (`"url": "https://…/photo.jpg"`) or the simplified `{"image": "<base64>"}` content-part form — see the [image-to-image request formats](../image_to_image/README.md#request-format).
+### GPU layout
 
-### Knobs (declared `extra_body` params)
-
-| Key | Default | Description |
-| --- | --- | --- |
-| `height` / `width` | 1024 | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). |
-| `steps` | 30 | Number of FlowMatchEuler denoise steps. |
-| `cfg` | 2.0 | Classifier-free guidance scale. |
-| `seed` | 42 | Per-request RNG seed. |
-| `byte5_text` | (auto) | Glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "…". ` format. Auto-extracted from quoted spans in the prompt when omitted. |
-| `negative_prompt` | (empty) | Real CFG negative conditioning (text-to-image only). |
-
-For the offline `text_to_image.py` / `image_edit.py` scripts and the full knob reference, see the [image-generation section in the recipe](../../../recipes/inclusionAI/Ming-flash-omni-2.0.md#image-generation-text-to-image--img2img).
+The shipped `ming_flash_omni_image.yaml` allocates the thinker on GPUs 0–3
+(TP=4) and the diffusion stage on GPU 4 (TP=1). Copy the YAML and edit
+`devices` per stage to relocate; with fewer GPUs available, drop
+the thinker TP to 2 and run the diffusion stage on a free card. Image-gen
+warmup takes roughly an extra 30–60 s on top of the thinker — set
+`--stage-init-timeout 1800` if the default 300 s is too tight.
 
 ## Modality control
 

examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh

@@ -0,0 +1,169 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Server port
+PORT="${PORT:-8091}"
+# Default query type
+QUERY_TYPE="${1:-text}"
+
+# Validate query type
+if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen)$ ]]; then
+    echo "Error: Invalid query type '$QUERY_TYPE'"
+    echo "Usage: $0 [text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen]"
+    echo "  text: Text-only query"
+    echo "  use_audio: Audio + Text query"
+    echo "  use_image: Image + Text query"
+    echo "  use_video: Video + Text query"
+    echo "  use_mixed_modalities: Audio + Image + Video + Text query"
+    echo "  use_image_gen: Text-to-image (diffusion stage) — saves the generated PNG locally"
+    exit 1
+fi
+
+# Define URLs for assets
+MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
+CHERRY_BLOSSOM_IMAGE_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+
+# Build user content based on query type
+case "$QUERY_TYPE" in
+  text)
+    user_content='[
+      {
+        "type": "text",
+        "text": "请详细介绍鹦鹉的生活习性。"
+      }
+    ]'
+    ;;
+  use_image)
+    user_content='[
+        {
+          "type": "image_url",
+          "image_url": {
+            "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
+          }
+        },
+        {
+          "type": "text",
+          "text": "Describe this image in detail."
+        }
+      ]'
+    ;;
+  use_audio)
+    user_content='[
+        {
+          "type": "audio_url",
+          "audio_url": {
+            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
+          }
+        },
+        {
+          "type": "text",
+          "text": "Please recognize the language of this speech and transcribe it. Format: oral."
+        }
+      ]'
+    ;;
+  use_video)
+    user_content='[
+        {
+          "type": "video_url",
+          "video_url": {
+            "url": "'"$SAMPLE_VIDEO_URL"'"
+          }
+        },
+        {
+          "type": "text",
+          "text": "Describe what is happening in this video."
+        }
+      ]'
+    ;;
+  use_mixed_modalities)
+    user_content='[
+        {
+          "type": "image_url",
+          "image_url": {
+            "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
+          }
+        },
+        {
+          "type": "audio_url",
+          "audio_url": {
+            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
+          }
+        },
+        {
+          "type": "text",
+          "text": "Describe the image, and recognize the language of this speech and transcribe it. Format: oral"
+        }
+      ]'
+    ;;
+esac
+
+echo "Running query type: $QUERY_TYPE"
+echo ""
+
+if [[ "$QUERY_TYPE" == "use_image_gen" ]]; then
+    # Image-gen branch: dual-stage diffusion endpoint, modalities=["image"].
+    # The shipped ming_flash_omni_image.yaml runs the thinker on cards 0-3
+    # (TP=4) and the diffusion stage on card 4 (TP=1). See the image-gen
+    # section in README.md for knobs and img2img.
+    out_path="${OUT_PATH:-/tmp/ming_imagegen.png}"
+    request_body=$(cat <<'EOF'
+{
+  "model": "Jonathan1909/Ming-flash-omni-2.0",
+  "modalities": ["image"],
+  "messages": [
+    {"role": "user", "content": "Draw a beautiful girl with short black hair and red dress."}
+  ]
+}
+EOF
+)
+    output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
+        -X POST "http://localhost:${PORT}/v1/chat/completions" \
+        -H "Content-Type: application/json" \
+        -d "$request_body")
+    # Server returns base64-encoded PNG either as a data URL in
+    # ``choices[0].message.content[0].image_url.url`` or as a raw string.
+    b64=$(echo "$output" | jq -r '.choices[0].message.content
+        | if type == "array"
+            then (map(select(.image_url? != null)) | .[0].image_url.url // "")
+            elif type == "string" then . else "" end')
+    if [[ -z "$b64" ]]; then
+        echo "Error: no image returned. Raw response:"
+        echo "$output" | jq '.'
+        exit 1
+    fi
+    if [[ "$b64" == data:image* ]]; then b64="${b64#*,}"; fi
+    printf '%s' "$b64" | base64 --decode > "$out_path"
+    echo "Saved generated image: $out_path ($(stat -c%s "$out_path" 2>/dev/null || stat -f%z "$out_path") bytes)"
+    exit 0
+fi
+
+request_body=$(cat <<EOF
+{
+  "model": "Jonathan1909/Ming-flash-omni-2.0",
+  "modalities": ["text"],
+  "messages": [
+    {
+      "role": "system",
+      "content": [
+        {
+          "type": "text",
+          "text": "你是一个友好的AI助手。\n\ndetailed thinking off"
+        }
+      ]
+    },
+    {
+      "role": "user",
+      "content": $user_content
+    }
+  ]
+}
+EOF
+)
+
+output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
+    -X POST http://localhost:${PORT}/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d "$request_body")
+
+echo "Output of request: $(echo "$output" | jq '.choices[0].message.content')"

pyproject.toml

@@ -41,7 +41,6 @@ dev = [
     "pytest-cov>=4.0.0",
     "pytest-mock>=3.10.0",
     "websockets>=15.0.1",
-    "pytest-subtests>=0.11.0",
     "datasets>=2.14.0",
     "mypy==1.11.1",
     "pre-commit==4.0.1",

recipes/MammothModa2/MammothModa2-Preview.md

@@ -32,12 +32,12 @@ flags. Image size uses the standard `--height` / `--width` flags.
   [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
 - Declared parameters:
   [`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)
-- Deploy config:
-  [`vllm_omni/deploy/mammoth_moda2.yaml`](../../vllm_omni/deploy/mammoth_moda2.yaml)
+- Stage config:
+  [`vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml`](../../vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml)
 
 ## Hardware Support
 
-The default deploy config runs both the AR and DiT stages on a single GPU
+The default stage config runs both the AR and DiT stages on a single GPU
 (`devices: "0"`). The committed `gpu_memory_utilization` split (stage-0 AR `0.5`,
 stage-1 DiT `0.3`) is sized for an ~80 GB GPU. The model also fits on a 48 GB GPU
 after rebalancing the split so the AR weights (~23 GB) leave room for the KV
@@ -48,7 +48,7 @@ cache — see the note under *1x L40S 48GB*.
 ### 1x L40S 48GB
 
 > **48 GB config adjustment:** the committed
-> `vllm_omni/deploy/mammoth_moda2.yaml` uses
+> `vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml` uses
 > `gpu_memory_utilization` 0.5 / 0.3 (sized for ~80 GB). To fit on a 48 GB L40S,
 > set the stage-0 (AR) value to `0.8` and the stage-1 (DiT) value to `0.16`
 > before running. (On an ~80 GB GPU, leave the defaults unchanged.)
@@ -71,13 +71,13 @@ hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-P
 ```
 
 Run text-to-image with the shared offline example from the repository root. The
-deploy config sets `trust_remote_code`, so no extra flag is needed. Forward the
+stage config sets `trust_remote_code`, so no extra flag is needed. Forward the
 MammothModa2 generation parameters as a JSON object through `--extra-body`:
 
 ```bash
 python examples/offline_inference/text_to_image/text_to_image.py \
   --model ./MammothModa2-Preview \
-  --stage-configs-path vllm_omni/deploy/mammoth_moda2.yaml \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
   --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
   --height 1024 \
   --width 1024 \

recipes/OpenBMB/MiniCPM-o-4_5.md

@@ -102,9 +102,8 @@ curl http://localhost:8099/v1/chat/completions \
 ```
 
 **Text + speech in one response** (the headline 4.5 feature). The TTS
-path is gated by a Jinja flag on the chat template. Pass
-`use_tts_template=true` via the **top-level** `chat_template_kwargs`
-field (curl does not flatten nested `extra_body`):
+path is gated by a Jinja flag on the chat template, passed through
+`extra_body.chat_template_kwargs.use_tts_template=true`:
 
 ```bash
 curl http://localhost:8099/v1/chat/completions \
@@ -113,18 +112,12 @@ curl http://localhost:8099/v1/chat/completions \
         "model": "openbmb/MiniCPM-o-4_5",
         "messages": [{"role": "user", "content": "Say hello, then introduce vLLM in one sentence."}],
         "modalities": ["text", "audio"],
-        "chat_template_kwargs": {"use_tts_template": true}
+        "extra_body": {"chat_template_kwargs": {"use_tts_template": true}}
     }'
 ```
 
-When using the OpenAI Python SDK, the same flag can also be sent as
-`extra_body={"chat_template_kwargs": {"use_tts_template": True}}`
-because the client merges `extra_body` into the request root.
-
-Response carries text in one choice's `message.content` and base64 WAV
-in another choice's `message.audio.data` (24 kHz mono, see Notes). With
-`modalities: ["text", "audio"]` you typically get two `choices` entries
-(one text, one audio).
+Response carries text in `choices[0].message.content` and base64 WAV
+in `choices[0].message.audio.data` (24 kHz mono, see Notes).
 
 **Gradio demo (text + image + audio + video UI)**:
 
@@ -216,13 +209,9 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \
   missing dep raises `ImportError` at first request with the same
   install hint instead of silently emitting empty audio.
 
-- **TTS trigger**: speech output requires
-  `chat_template_kwargs.use_tts_template=true` so the chat template
-  appends `<|tts_bos|>` before generation. Without it, Stage-1 talker
-  receives no TTS token span and returns silent audio (not text-only).
-  For **curl**, put `chat_template_kwargs` at the request root; nested
-  `extra_body.chat_template_kwargs` is ignored. The OpenAI Python SDK
-  may use `extra_body` because it flattens those fields into the root.
+- **TTS trigger**: speech output is only emitted when the client passes
+  `extra_body.chat_template_kwargs.use_tts_template=true`. Without it,
+  the response is text-only (which is also faster).
 
 - **Output audio**: 24 kHz mono WAV inside the OpenAI-style
   `message.audio.data` (base64). The Gradio demo's WAV player decodes

recipes/README.md

@@ -48,7 +48,6 @@ recipes/
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
 | [`MammothModa2/MammothModa2-Preview.md`](./MammothModa2/MammothModa2-Preview.md) | Text-to-image with the shared offline image example (AR → DiT) | 1x L40S 48GB / 1x ≥40GB GPU |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
-| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense text-to-video serving | 1x L20X / A100 / H100 |
 | [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
 | [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
 | [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |

recipes/Robbyant/LingBot-Video.md

@@ -1,143 +0,0 @@
-# LingBot-Video
-
-> Native dense text-to-video serving for `robbyant/lingbot-video-dense-1.3b`
-
-## Summary
-
-- Vendor: Robbyant
-- Model: `robbyant/lingbot-video-dense-1.3b`
-- Task: Text-to-video generation
-- Mode: Offline generation and online serving with the OpenAI-compatible `/v1/videos` API
-- Maintainer: Community
-
-## When to use this recipe
-
-Use this recipe when you want to run the dense LingBot-Video checkpoint with
-vLLM-Omni's native pipeline. The runtime path does not import the upstream
-`lingbot_video` Python package; it loads the checkpoint components directly
-with the in-tree `LingBotVideoPipeline`, dense DiT transformer, shared FlowUniPC
-scheduler, Qwen3-VL text encoder, and Wan VAE.
-
-This recipe is intentionally text-to-video only. Text-to-image, image-conditioned
-video, TI2V, MoE checkpoints, and fused-expert kernels are not covered here.
-
-## References
-
-- Dense checkpoint: <https://huggingface.co/robbyant/lingbot-video-dense-1.3b>
-- Upstream project: <https://github.com/Robbyant/lingbot-video>
-- Related offline example: [`examples/offline_inference/text_to_video/text_to_video_lingbot.py`](../../examples/offline_inference/text_to_video/text_to_video_lingbot.py)
-- Related online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
-
-## Hardware Support
-
-This recipe documents the CUDA single-GPU dense checkpoint path. Multi-GPU
-parallelism, Cache-DiT, CPU offload, and the LingBot MoE path are not validated
-for this model in this PR.
-
-## GPU
-
-### 1 x NVIDIA L20X / A100 / H100
-
-The dense 1.3B checkpoint has been smoke-tested on a single NVIDIA L20X at a
-small validation shape (`192x320`, 9 frames, 2 steps). Use a larger GPU or lower
-resolution/frame count if your target workload runs out of memory.
-
-#### Offline T2V
-
-```bash
-CUDA_VISIBLE_DEVICES=0 \
-python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
-  --model robbyant/lingbot-video-dense-1.3b \
-  --prompt "a robotic arm picks up a red block" \
-  --output lingbot_t2v.mp4 \
-  --height 192 \
-  --width 320 \
-  --num-frames 9 \
-  --num-inference-steps 2 \
-  --guidance-scale 3.0 \
-  --flow-shift 3.0 \
-  --seed 42 \
-  --fps 24
-```
-
-#### Online Serving
-
-```bash
-CUDA_VISIBLE_DEVICES=0 \
-vllm serve robbyant/lingbot-video-dense-1.3b \
-  --omni \
-  --model-class-name LingBotVideoPipeline \
-  --port 8091
-```
-
-After the server is ready, submit a text-to-video job:
-
-```bash
-create_response=$(curl -s http://localhost:8091/v1/videos \
-  -F "model=robbyant/lingbot-video-dense-1.3b" \
-  -F "prompt=a robotic arm picks up a red block" \
-  -F "width=320" \
-  -F "height=192" \
-  -F "num_frames=9" \
-  -F "fps=24" \
-  -F "num_inference_steps=2" \
-  -F "guidance_scale=3.0" \
-  -F "flow_shift=3.0" \
-  -F "seed=42")
-
-video_id=$(echo "${create_response}" | jq -r '.id')
-while true; do
-  status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
-  if [ "${status}" = "completed" ]; then
-    break
-  fi
-  if [ "${status}" = "failed" ]; then
-    curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
-    exit 1
-  fi
-  sleep 2
-done
-
-curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o lingbot_t2v.mp4
-```
-
-## Key Parameters
-
-| Parameter | Suggested smoke value | Notes |
-|-----------|-----------------------|-------|
-| `height` | `192` | Must be a multiple of 16 |
-| `width` | `320` | Must be a multiple of 16 |
-| `num_frames` | `9` | Must be `1` or `4n + 1`; this PR validates T2V with video outputs |
-| `num_inference_steps` | `2` | Use more steps for quality sweeps |
-| `guidance_scale` | `3.0` | CFG is active when this is greater than `1.0` |
-| `flow_shift` | `3.0` | Scheduler flow-shift; aliases the pipeline's internal `shift` |
-| `negative_prompt` | model default | Optional text describing artifacts to avoid |
-| `fps` | `24` | Output MP4 frame rate |
-
-## Validation
-
-Local dense smoke run:
-
-- Shape: 9 frames at `192x320`
-- Steps: 2
-- Request generation time: `0.2923s`
-- Peak reserved GPU memory: `14548 MiB`
-
-Local parity harness against the upstream repository:
-
-- Shape: `[9, 192, 320, 3]`
-- MAE: `0.0065238`
-- MSE: `0.00006650`
-- PSNR: `41.77 dB`
-- Native request time: `0.2875s`
-
-Do not treat these as production benchmarks; they are draft-PR validation
-numbers for the small dense smoke configuration.
-
-## Known Limitations
-
-- T2V only. T2I, I2V, and TI2V are not claimed by this PR.
-- Dense checkpoint only. MoE/fused-expert support belongs in a separate PR.
-- No Cache-DiT, TeaCache, tensor parallelism, sequence parallelism, CFG
-  parallelism, HSDP, CPU offload, or VAE patch parallelism validation yet.
-- Only one request per LingBot pipeline batch is currently supported.

recipes/inclusionAI/Ming-flash-omni-2.0.md

@@ -18,7 +18,6 @@ Use this recipe when you want a known-good starting point for serving
 
 - **Thinker only** — multimodal understanding with text output.
 - **Thinker + Talker (omni-speech)** — multimodal understanding with text and spoken output.
-- **Thinker + Imagegen (DiT)** — text-to-image / img2img
 - **Talker only (TTS)** — standalone text-to-speech via the OpenAI `/v1/audio/speech` endpoint.
 
 ## References
@@ -161,129 +160,6 @@ deltas appear at `choices[0].delta.content`.
 - Memory usage: size depends on output modalities and multimodal input; leave
   headroom for video frames and audio caches.
 
-## Image generation (text-to-image / img2img)
-
-Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage. Launch with the image deploy YAML, which adds an image-gen stage behind the thinker.
-
-The image-generation stage is a standard vLLM-Omni diffusion pipeline (`MingImagePipeline`); its request knobs are declared in `vllm_omni/model_extras/ming_flash_omni.py` and routed through `extra_body`, so they no longer need a bespoke `sampling_params_list` recipe (that form is still available for per-stage thinker sampling — see below).
-
-```bash
-vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
-    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
-    --stage-init-timeout 1800 \
-    --init-timeout 1800 \
-    --port 8091
-```
-
-With fewer GPUs, copy the YAML and drop the thinker to TP=2 with the DiT on a
-free card.
-
-### Online (text-to-image)
-
-Request image output with `"modalities": ["image"]`:
-
-```bash
-curl -s http://127.0.0.1:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "Jonathan1909/Ming-flash-omni-2.0",
-    "messages": [
-      {
-        "role": "user",
-        "content": "Please draw a cute cat."
-      }
-    ],
-    "modalities": ["image"]
-  }' \
-  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
-  | base64 -d > ming_imagegen.png
-```
-
-**Quick form — `extra_body`** (keys are filtered against the declared set and routed into every stage's `extra_args`). Convenient, but it does **not** let you set the thinker (stage-0) sampling params. Wrap the knobs in a literal `"extra_body"` object for raw curl; the OpenAI Python client's `extra_body=` kwarg produces the same request:
-
-```bash
-curl http://127.0.0.1:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "Jonathan1909/Ming-flash-omni-2.0",
-    "modalities": ["image"],
-    "messages": [
-      {
-        "role": "user",
-        "content": "Draw a poster."
-      }
-    ],
-    "extra_body": {
-      "steps": 6,
-      "cfg": 1.5,
-      "height": 512,
-      "width": 512,
-      "seed": 123,
-      "byte5_text": ["理解与生成统一"],
-      "negative_prompt": "ugly, blurry, distorted"
-    }
-  }' \
-  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
-  | base64 -d > ming_imagegen_extra_body.png
-```
-
-**Full control — `sampling_params_list`** (one entry per stage: `[thinker, imagegen]`).
-Use this when you need to tune the thinker's own sampling (`temperature` / `top_p` / `top_k` / `max_tokens`), or to place knobs explicitly per stage.
-Note `negative_prompt` must sit on the **stage-0 thinker** `extra_args`; the imagegen-stage knobs (`steps` / `cfg` / `height` / `width` / `seed` / `byte5_text`) go on the **stage-1** entry:
-
-```bash
-curl http://127.0.0.1:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "Jonathan1909/Ming-flash-omni-2.0",
-    "modalities": ["image"],
-    "sampling_params_list": [
-      {
-        "temperature": 0.4,
-        "top_p": 0.9,
-        "top_k": 1,
-        "max_tokens": 1,
-        "seed": 42,
-        "extra_args": {
-          "negative_prompt": "ugly, blurry, distorted"
-        }
-      },
-      {
-        "seed": 42,
-        "extra_args": {
-          "steps": 6,
-          "cfg": 1.5,
-          "height": 512,
-          "width": 512,
-          "seed": 123,
-          "byte5_text": ["理解与生成统一"]
-        }
-      }
-    ],
-    "messages": [
-      {
-        "role": "user",
-        "content": "Draw a poster."
-      }
-    ]
-  }' \
-  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
-  | base64 -d > ming_imagegen_knobs.png
-```
-
-### Knobs (declared `extra_body` params)
-
-| Key | Default | Description |
-| --- | --- | --- |
-| `height` / `width` | 1024 | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). |
-| `steps` | 30 | Number of FlowMatchEuler denoise steps. |
-| `cfg` | 2.0 | Classifier-free guidance scale. |
-| `seed` | 42 | Per-request RNG seed. |
-| `byte5_text` | (auto) | Glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "…". ` format. Auto-extracted from quoted spans in the prompt when omitted. |
-| `negative_prompt` | (empty) | Real CFG negative conditioning. Spawns a CFG-text companion via `expand_cfg_prompts`; **online / text-to-image only** (offline uses Ming's default zero-negative). |
-
-For img2img, add an `image_url` content part to the user message (online) or pass `--image` (offline); the reference image is routed into the DiT stage as `extra[reference_image]`.
-
 ### 1x H100 80GB — standalone TTS (talker only)
 
 The bundled `ming_flash_omni_tts.yaml` runs the talker on a single GPU and exposes the OpenAI `/v1/audio/speech` endpoint.

requirements/common.txt

@@ -13,7 +13,7 @@ transformers >= 5.5.3
 av>=14.0.0
 omegaconf>=2.3.0
 diffusers==0.38.0
-safetensors>=0.8.0
+safetensors>=0.8.0rc0
 accelerate==1.12.0
 soundfile>=0.13.1
 cache-dit==1.3.0

tests/config/test_endpoint_policy.py

@@ -1,87 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for endpoint restrictions logic."""
-
-import pytest
-from fastapi import FastAPI
-from fastapi.testclient import TestClient
-
-from vllm_omni.config.endpoint_policy import (
-    EndpointRestriction,
-    OmniServingCapability,
-    shutdown_unsupported_routes,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-SERVER_ROUTE = "/v1/completions"
-REJECTION_REASON = "This route is banned!"
-
-
-def _make_app_with_server_route():
-    """Create a FastAPI app with a route to reject."""
-    app = FastAPI()
-
-    @app.post(SERVER_ROUTE)
-    async def existing_handler():
-        """Returns a response; we want to disable this."""
-        return {"ok": True}
-
-    return app
-
-
-def test_restricted_completions_returns_400():
-    """Ensure that simple route rejection works."""
-    app = _make_app_with_server_route()
-    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
-    shutdown_unsupported_routes(app, restrictions)
-
-    client = TestClient(app)
-    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
-    assert resp.status_code == 400
-    body = resp.json()
-    assert body["error"]["message"] == REJECTION_REASON
-    assert body["error"]["type"] == "BadRequestError"
-
-
-def test_unrestricted_completions_not_blocked():
-    """Ensure that if we don't shutdown any routes, the route stays."""
-    app = _make_app_with_server_route()
-    shutdown_unsupported_routes(app, ())
-
-    client = TestClient(app)
-    resp = client.post(SERVER_ROUTE)
-    assert resp.status_code == 200
-    assert resp.json() == {"ok": True}
-
-
-def test_route_rejection_is_idempotent():
-    """Ensure that calling rejection twice doesn't break things."""
-    app = _make_app_with_server_route()
-    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
-    shutdown_unsupported_routes(app, restrictions)
-    shutdown_unsupported_routes(app, restrictions)
-
-    client = TestClient(app)
-    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
-    assert resp.status_code == 400
-    body = resp.json()
-    assert body["error"]["message"] == REJECTION_REASON
-    assert body["error"]["type"] == "BadRequestError"
-
-
-def test_invalid_route_rejection():
-    """Ensure that a restriction won't crash if a route doesn't exist."""
-    # NOTE: This test is for safety since the route organization is still messy.
-    # Since the restriction needs a valid OmniServingCapability, its implied the route
-    # is a valid vLLM route, but was already removed from the app by earlier processing.
-    app = FastAPI()
-    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
-    shutdown_unsupported_routes(app, restrictions)
-
-    client = TestClient(app)
-    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
-    assert resp.status_code == 400
-    body = resp.json()
-    assert body["error"]["message"] == REJECTION_REASON
-    assert body["error"]["type"] == "BadRequestError"

tests/dfx/perf/tests/test_lingbot_video_vllm_omni.json

@@ -1,34 +0,0 @@
-[
-    {
-        "test_name": "test_lingbot_video_single_device",
-        "description": "Single-device LingBot dense T2V baseline at 320x192, 9 frames, 2 steps.",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "robbyant/lingbot-video-dense-1.3b",
-            "serve_args": {
-                "model-class-name": "LingBotVideoPipeline",
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "320x192_frames9_steps2",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 320,
-                "height": 192,
-                "num-frames": 9,
-                "fps": 24,
-                "num-inference-steps": 2,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "warmup-requests": 1,
-                "warmup-concurrency": 1,
-                "warmup-num-inference-steps": 2,
-                "enable-negative-prompt": true,
-                "skip-performance-assertion": true
-            }
-        ]
-    }
-]

tests/diffusion/attention/test_kernels_hub.py

@@ -1,104 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import sys
-import types
-
-import pytest
-import torch
-
-from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
-from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
-from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
-from vllm_omni.platforms import current_omni_platform
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_kernels_hub_platform_fallback(monkeypatch: pytest.MonkeyPatch):
-    """Test that when kernels package is not available, the platform fallback logic
-
-    routes to local FLASH_ATTN.
-    """
-    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
-
-    # Temporarily hide/remove kernels module if present
-    monkeypatch.setitem(sys.modules, "kernels", None)
-
-    # Use monkeypatch to mock CudaOmniPlatform capability and package checker to allow FLASH_ATTN
-    from vllm.platforms.interface import DeviceCapability
-
-    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
-
-    monkeypatch.setattr(
-        CudaOmniPlatform,
-        "get_device_capability",
-        classmethod(lambda cls, device_id=0: DeviceCapability(8, 0)),
-    )
-    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
-
-    # Test FLASH_ATTN_HUB falls back to FLASH_ATTN when kernels is unavailable
-    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
-    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
-
-    # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN when kernels is unavailable
-    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
-    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
-
-    # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN_HUB on pre-Hopper GPUs
-    kernels_module = types.ModuleType("kernels")
-    monkeypatch.setitem(sys.modules, "kernels", kernels_module)
-    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
-    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path()
-
-
-@pytest.mark.core_model
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_kernels_hub_execution():
-    """Verify basic forward of flash_attn_hub and flash_attn_3_hub, comparing with SDPA reference."""
-    device = torch.device(current_omni_platform.device_type)
-    dtype = torch.bfloat16
-
-    num_heads = 8
-    head_dim = 64
-    seq_len = 32
-    batch_size = 1
-
-    torch.manual_seed(42)
-    q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
-    k = q.clone()
-    v = q.clone()
-
-    # 1. Test PyTorch SDPA reference
-    sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
-    attn_metadata_sdpa = AttentionMetadata(attn_mask=None)
-    output_ref = sdpa_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
-
-    # 2. Test FlashAttentionHubBackend (FlashAttention 2)
-    from vllm_omni.diffusion.attention.backends.flash_attn_hub import FlashAttentionHubImpl
-
-    fa_hub_impl = FlashAttentionHubImpl(
-        num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
-    )
-    output_fa_hub = fa_hub_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
-    assert output_fa_hub.shape == q.shape
-    assert not torch.isnan(output_fa_hub).any()
-    max_diff = torch.max(torch.abs(output_ref - output_fa_hub)).item()
-    assert max_diff < 1e-2, f"FlashAttentionHub output differs too much from SDPA reference: {max_diff}"
-
-    # 3. Test FlashAttention3HubBackend (FlashAttention 3, Hopper+ only)
-    major, _minor = torch.cuda.get_device_capability()
-    if major < 9:
-        pytest.skip("FLASH_ATTN_3_HUB execution requires Hopper-class GPU (compute capability >= 9.0)")
-
-    from vllm_omni.diffusion.attention.backends.flash_attn_hub import FlashAttention3HubImpl
-
-    fa3_hub_impl = FlashAttention3HubImpl(
-        num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
-    )
-    output_fa3_hub = fa3_hub_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
-    assert output_fa3_hub.shape == q.shape
-    assert not torch.isnan(output_fa3_hub).any()
-    max_diff = torch.max(torch.abs(output_ref - output_fa3_hub)).item()
-    assert max_diff < 1e-2, f"FlashAttention3Hub output differs too much from SDPA reference: {max_diff}"

tests/diffusion/layers/test_rope_broadcast.py

@@ -1,236 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Shape broadcasting regression tests for RotaryEmbedding.forward_native.
-
-Covers the batch-dimension shape mismatch fix: forward_native now strips the
-batch dim from 3D cos/sin before applying RoPE, consistent with forward_cuda,
-forward_hip, and apply_rotary_emb_mindiesd paths.
-"""
-
-from __future__ import annotations
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.layers.rope import RotaryEmbedding
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _make_inputs(
-    batch: int,
-    seq_len: int,
-    num_heads: int,
-    head_dim: int,
-    cos_batch: int | None = None,
-    dtype: torch.dtype = torch.float32,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Generate x, cos, sin tensors for RoPE testing.
-
-    Args:
-        batch: Batch dimension of x.
-        seq_len: Sequence length.
-        num_heads: Number of attention heads.
-        head_dim: Head dimension (must be even).
-        cos_batch: Batch dimension of cos/sin. Defaults to batch if None.
-        dtype: Data type for the tensors.
-    """
-    if cos_batch is None:
-        cos_batch = batch
-
-    torch.manual_seed(42)
-    x = torch.randn(batch, seq_len, num_heads, head_dim, dtype=dtype)
-    cos = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype)
-    sin = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype)
-    return x, cos, sin
-
-
-class TestRotaryEmbeddingNativeShapeRegression:
-    """Regression tests for forward_native shape handling (#8e297d81)."""
-
-    @pytest.fixture
-    def rope_neox(self) -> RotaryEmbedding:
-        """NeoX-style (non-interleaved) RoPE."""
-        return RotaryEmbedding(is_neox_style=True)
-
-    @pytest.fixture
-    def rope_gptj(self) -> RotaryEmbedding:
-        """GPT-J style (interleaved) RoPE."""
-        return RotaryEmbedding(is_neox_style=False)
-
-    # ── 2D cos/sin (no batch dim) ──────────────────────────────────────
-
-    def test_2d_cos_sin_unchanged(self, rope_neox: RotaryEmbedding) -> None:
-        """2D cos/sin [S, D/2] pass through unchanged (no-op path)."""
-        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=2)
-        # Manually drop batch dim to create 2D cos/sin
-        cos_2d = cos[0]
-        sin_2d = sin[0]
-        assert cos_2d.dim() == 2
-        output = rope_neox.forward_native(x, cos_2d, sin_2d)
-        assert output.shape == x.shape
-        assert not torch.isnan(output).any()
-        assert not torch.isinf(output).any()
-
-    # ── 3D cos/sin — batch-dim stripping ──────────────────────────────
-
-    @pytest.mark.parametrize("cos_batch", [1, 2, 4])
-    def test_3d_cos_sin_no_error(self, rope_neox: RotaryEmbedding, cos_batch: int) -> None:
-        """3D cos/sin with batch dim should not raise errors."""
-        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch)
-        output = rope_neox.forward_native(x, cos, sin)
-        assert output.shape == x.shape
-
-    @pytest.mark.parametrize("interleaved", [False, True])
-    def test_3d_cos_identical_to_2d(self, interleaved: bool) -> None:
-        """3D cos/sin with identical batch slices == 2D cos/sin result."""
-        rope = RotaryEmbedding(is_neox_style=not interleaved)
-        B, S, H, D = 2, 16, 8, 64
-        x = torch.randn(B, S, H, D)
-        cos_2d = torch.randn(S, D // 2)
-        sin_2d = torch.randn(S, D // 2)
-
-        # 3D variant: expand the same 2D data across batch
-        cos_3d = cos_2d.unsqueeze(0).expand(B, -1, -1)
-        sin_3d = sin_2d.unsqueeze(0).expand(B, -1, -1)
-
-        out_2d = rope.forward_native(x, cos_2d, sin_2d)
-        out_3d = rope.forward_native(x, cos_3d, sin_3d)
-
-        torch.testing.assert_close(out_3d, out_2d, atol=0, rtol=0)
-
-    # ── Batch size mismatch — the regression scenario ─────────────────
-
-    @pytest.mark.parametrize(
-        "x_batch,cos_batch",
-        [
-            (2, 1),  # cos has fewer batch elements
-            (1, 2),  # cos has more batch elements (takes [0])
-            (4, 1),  # large x batch, single cos
-        ],
-    )
-    def test_batch_size_mismatch_regression(
-        self,
-        rope_neox: RotaryEmbedding,
-        x_batch: int,
-        cos_batch: int,
-    ) -> None:
-        """forward_native works when cos batch dim differs from x batch dim.
-
-        This is the exact regression scenario — before the fix, a 3D cos/sin
-        with a different batch size than x would cause a shape mismatch error
-        inside apply_rotary_emb_torch during broadcast.
-        """
-        x, cos, sin = _make_inputs(batch=x_batch, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch)
-        output = rope_neox.forward_native(x, cos, sin)
-        assert output.shape == x.shape
-        assert output.dtype == x.dtype
-
-    # ── Batch consistency — all elements get same rotation ────────────
-
-    def test_2d_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None:
-        """With 2D cos, identical x slices → identical output slices."""
-        B, S, H, D = 4, 16, 8, 64
-        x_template = torch.randn(1, S, H, D)
-        x = x_template.expand(B, -1, -1, -1)  # all batch elements identical
-        cos = torch.randn(S, D // 2)
-        sin = torch.randn(S, D // 2)
-
-        output = rope_neox.forward_native(x, cos, sin)
-
-        # All output batch slices must be identical
-        for i in range(1, B):
-            torch.testing.assert_close(output[i], output[0], atol=0, rtol=0)
-
-    def test_3d_identical_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None:
-        """With 3D cos (identical slices), identical x → identical output."""
-        B, S, H, D = 4, 16, 8, 64
-        x_template = torch.randn(1, S, H, D)
-        x = x_template.expand(B, -1, -1, -1)
-        cos = torch.randn(1, S, D // 2).expand(B, -1, -1)
-        sin = torch.randn(1, S, D // 2).expand(B, -1, -1)
-
-        output = rope_neox.forward_native(x, cos, sin)
-
-        for i in range(1, B):
-            torch.testing.assert_close(output[i], output[0], atol=0, rtol=0)
-
-    # ── Output shape / dtype invariants ───────────────────────────────
-
-    def test_output_shape_preserved(self, rope_neox: RotaryEmbedding) -> None:
-        """Output shape must match input shape for all cos/sin dimensionalities."""
-        B, S, H, D = 2, 32, 4, 128
-        x = torch.randn(B, S, H, D)
-
-        # 2D cos/sin
-        cos_2d = torch.randn(S, D // 2)
-        sin_2d = torch.randn(S, D // 2)
-        assert rope_neox.forward_native(x, cos_2d, sin_2d).shape == (B, S, H, D)
-
-        # 3D cos/sin, matching batch
-        cos_3d = torch.randn(B, S, D // 2)
-        sin_3d = torch.randn(B, S, D // 2)
-        assert rope_neox.forward_native(x, cos_3d, sin_3d).shape == (B, S, H, D)
-
-        # 3D cos/sin, different batch
-        cos_3d_1 = torch.randn(1, S, D // 2)
-        sin_3d_1 = torch.randn(1, S, D // 2)
-        assert rope_neox.forward_native(x, cos_3d_1, sin_3d_1).shape == (B, S, H, D)
-
-    def test_output_dtype_preserved(self, rope_neox: RotaryEmbedding) -> None:
-        """Output dtype must match input dtype."""
-        for dtype in [torch.float32, torch.float16, torch.bfloat16]:
-            x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, dtype=dtype)
-            output = rope_neox.forward_native(x, cos, sin)
-            assert output.dtype == dtype, f"dtype mismatch: {output.dtype} != {dtype}"
-
-    def test_output_no_nan_inf(
-        self,
-        rope_neox: RotaryEmbedding,
-        rope_gptj: RotaryEmbedding,
-    ) -> None:
-        """Output must be finite for both RoPE styles."""
-        for rope in (rope_neox, rope_gptj):
-            x, cos, sin = _make_inputs(batch=2, seq_len=32, num_heads=4, head_dim=64)
-            output = rope.forward_native(x, cos, sin)
-            assert not torch.isnan(output).any(), f"NaN in {rope}"
-            assert not torch.isinf(output).any(), f"Inf in {rope}"
-
-    def test_output_contiguous(self, rope_neox: RotaryEmbedding) -> None:
-        """Output should be contiguous."""
-        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64)
-        output = rope_neox.forward_native(x, cos, sin)
-        assert output.is_contiguous()
-
-    # ── CFG-style doubled batch ───────────────────────────────────────
-
-    def test_cfg_style_doubled_batch(self, rope_neox: RotaryEmbedding) -> None:
-        """CFG (classifier-free guidance) doubles batch — cos may have
-        original batch size while x has 2× batch."""
-        B, S, H, D = 2, 16, 8, 64
-        # Simulate CFG: x has 2× batch (cond + uncond concatenated)
-        x = torch.randn(B * 2, S, H, D)
-        # cos/sin only have the original batch
-        cos = torch.randn(B, S, D // 2)
-        sin = torch.randn(B, S, D // 2)
-
-        output = rope_neox.forward_native(x, cos, sin)
-        assert output.shape == (B * 2, S, H, D)
-
-    # ── RotaryEmbeddingWan forward_native is unchanged ────────────────
-
-    def test_wan_forward_native_unchanged(self) -> None:
-        """RotaryEmbeddingWan.forward_native still uses its own implementation
-        (unflatten + stack + flatten) and should not crash with 3D cos/sin."""
-        from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
-
-        rope = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True)
-        B, S, H, D = 2, 16, 8, 64
-        x = torch.randn(B, S, H, D)
-        cos = torch.randn(B, S, 1, D // 2)
-        sin = torch.randn(B, S, 1, D // 2)
-
-        output = rope.forward_native(x, cos, sin)
-        assert output.shape == x.shape
-        assert not torch.isnan(output).any()

tests/diffusion/models/lingbot_video/__init__.py

@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/diffusion/models/lingbot_video/test_lingbot_video_transformer.py

@@ -1,117 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-import torch
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
-
-
-def _tiny_transformer(**overrides):
-    from vllm_omni.diffusion.models.lingbot_video import LingBotVideoTransformer3DModel
-
-    config = {
-        "patch_size": (1, 1, 1),
-        "in_channels": 2,
-        "out_channels": 2,
-        "hidden_size": 16,
-        "num_attention_heads": 1,
-        "depth": 0,
-        "intermediate_size": 32,
-        "text_dim": 8,
-        "freq_dim": 8,
-        "axes_dims": (4, 4, 8),
-        "axes_lens": (32, 32, 32),
-    }
-    config.update(overrides)
-    return LingBotVideoTransformer3DModel(**config)
-
-
-def test_joint_position_ids_video_then_text_order():
-    from vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer import make_joint_position_ids
-
-    positions = make_joint_position_ids(text_len=3, grid_t=1, grid_h=2, grid_w=2, device=torch.device("cpu"))
-
-    assert positions.shape == (7, 3)
-    assert positions[:4, 0].tolist() == [4, 4, 4, 4]
-    assert positions[:4, 1:].tolist() == [[0, 0], [0, 1], [1, 0], [1, 1]]
-    assert positions[4:].tolist() == [[1, 0, 0], [2, 0, 0], [3, 0, 0]]
-
-
-def test_tiny_transformer_depth_zero_forward_shape():
-    model = _tiny_transformer()
-    hidden_states = torch.randn(1, 2, 1, 2, 2)
-    timestep = torch.tensor([300.0])
-    encoder_hidden_states = torch.randn(1, 3, 8)
-    encoder_attention_mask = torch.ones(1, 3, dtype=torch.long)
-
-    with torch.no_grad():
-        out = model(
-            hidden_states,
-            timestep,
-            encoder_hidden_states,
-            encoder_attention_mask=encoder_attention_mask,
-            return_dict=False,
-        )[0]
-
-    assert out.shape == hidden_states.shape
-    assert torch.isfinite(out).all()
-
-
-def test_packed_attention_uses_sdpa_fallback_without_flash_varlen(monkeypatch):
-    from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module
-
-    monkeypatch.setattr(module, "flash_attn_varlen_func_v3", None)
-    attn = module.LingBotVideoAttention(
-        hidden_size=8,
-        num_heads=2,
-        norm_eps=1e-6,
-        qkv_bias=False,
-        out_bias=False,
-    )
-    captured = {}
-
-    def fake_sdpa_forward(query, key, value, attn_metadata):
-        captured["mask"] = attn_metadata.attn_mask
-        return torch.zeros_like(query)
-
-    monkeypatch.setattr(attn.attn.sdpa_fallback, "forward", fake_sdpa_forward)
-    x = torch.randn(1, 5, 8)
-    rotary = torch.ones(1, 5, 2, dtype=torch.complex64)
-    packed_indices = {
-        "cu_seqlens_kv": torch.tensor([0, 2, 5], dtype=torch.int32),
-        "max_seqlen_in_batch_kv": 3,
-        "attention_mask": module._packed_block_attention_mask([2, 3], x.device),
-    }
-
-    out = attn(x, rotary, packed_indices=packed_indices)
-
-    assert out.shape == x.shape
-    mask = captured["mask"]
-    assert mask.shape == (1, 1, 5, 5)
-    assert mask[0, 0, :2, :2].all()
-    assert mask[0, 0, 2:, 2:].all()
-    assert not mask[0, 0, :2, 2:].any()
-    assert not mask[0, 0, 2:, :2].any()
-
-
-def test_tiny_transformer_rejects_invalid_rope_dims():
-    from vllm_omni.diffusion.models.lingbot_video import LingBotVideoTransformer3DModel
-
-    with pytest.raises(AssertionError, match="head_dim"):
-        LingBotVideoTransformer3DModel(
-            hidden_size=16,
-            num_attention_heads=1,
-            axes_dims=(4, 4, 4),
-            depth=0,
-        )
-
-
-def test_transformer_to_keeps_sensitive_modules_in_fp32():
-    model = _tiny_transformer()
-
-    model.to(dtype=torch.bfloat16)
-
-    assert model.patch_embedder.weight.dtype == torch.bfloat16
-    assert model.time_embedder.linear_1.weight.dtype == torch.float32
-    assert model.norm_out_modulation[1].weight.dtype == torch.float32

tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py

@@ -1,195 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from types import SimpleNamespace
-
-import numpy as np
-import pytest
-import torch
-from torch import nn
-
-from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
-
-
-def _make_pipeline():
-    from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import LingBotVideoPipeline
-
-    pipeline = object.__new__(LingBotVideoPipeline)
-    nn.Module.__init__(pipeline)
-    pipeline.device = torch.device("cpu")
-    pipeline.default_negative_prompt = "default negative"
-    pipeline.od_config = SimpleNamespace(flow_shift=None)
-    return pipeline
-
-
-def _make_request_batch(prompt, **sampling_overrides):
-    sampling = OmniDiffusionSamplingParams(**sampling_overrides)
-    request = OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id="lingbot-test")
-    return DiffusionRequestBatch([request])
-
-
-def test_lingbot_video_pipeline_import_and_registry():
-    from vllm_omni.diffusion.models.lingbot_video import (
-        LingBotVideoPipeline,
-        LingBotVideoTransformer3DModel,
-        get_lingbot_video_post_process_func,
-    )
-    from vllm_omni.diffusion.registry import _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS
-
-    assert LingBotVideoPipeline is not None
-    assert LingBotVideoTransformer3DModel is not None
-    assert get_lingbot_video_post_process_func is not None
-    assert _DIFFUSION_MODELS["LingBotVideoPipeline"] == (
-        "lingbot_video",
-        "pipeline_lingbot_video",
-        "LingBotVideoPipeline",
-    )
-    assert _DIFFUSION_POST_PROCESS_FUNCS["LingBotVideoPipeline"] == "get_lingbot_video_post_process_func"
-
-
-def test_component_discovery_declarations():
-    from vllm_omni.diffusion.models.lingbot_video import LingBotVideoPipeline
-
-    assert LingBotVideoPipeline._dit_modules == ["transformer"]
-    assert LingBotVideoPipeline._encoder_modules == ["text_encoder"]
-    assert LingBotVideoPipeline._vae_modules == ["vae"]
-    assert LingBotVideoPipeline.supports_step_execution is False
-
-
-def test_extra_body_params_include_video_flow_shift_alias():
-    from vllm_omni.model_extras import get_extra_body_params
-
-    params = get_extra_body_params("LingBotVideoPipeline")
-    assert "shift" in params
-    assert "flow_shift" in params
-    assert "batch_cfg" in params
-
-
-def test_extract_prompt_string_and_mapping():
-    from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import _extract_prompt
-
-    string_req = SimpleNamespace(prompt="a robot waves")
-    assert _extract_prompt(string_req) == ("a robot waves", None)
-
-    mapping_req = SimpleNamespace(prompt={"prompt": "a robot walks", "negative_prompt": "low quality"})
-    assert _extract_prompt(mapping_req) == ("a robot walks", "low quality")
-
-    with pytest.raises(TypeError, match="Unsupported LingBot prompt type"):
-        _extract_prompt(SimpleNamespace(prompt=["not", "supported"]))
-
-
-def test_check_inputs_accepts_t2v_shapes_and_rejects_invalid_values():
-    from vllm_omni.diffusion.models.lingbot_video import LingBotVideoPipeline
-
-    LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=1)
-    LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=9)
-
-    with pytest.raises(ValueError, match="num_frames"):
-        LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=8)
-    with pytest.raises(ValueError, match="height"):
-        LingBotVideoPipeline.check_inputs(height=190, width=320, num_frames=9)
-
-
-def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
-    from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
-
-    frames = torch.ones(2, 4, 4, 3)
-    post = get_lingbot_video_post_process_func(SimpleNamespace())
-
-    assert post(frames, SimpleNamespace(output_type="pt")) is frames
-    array = post(frames, SimpleNamespace(output_type="np"))
-    assert isinstance(array, np.ndarray)
-    assert array.shape == (2, 4, 4, 3)
-
-
-def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
-    pipeline = _make_pipeline()
-    calls = []
-
-    def fake_generate(**kwargs):
-        calls.append(kwargs)
-        return torch.zeros(9, 192, 320, 3)
-
-    pipeline._generate = fake_generate
-    req = _make_request_batch(
-        {"prompt": "a robot arm", "negative_prompt": "bad hands"},
-        height=192,
-        width=320,
-        num_frames=9,
-        num_inference_steps=2,
-        guidance_scale=3.0,
-        seed=42,
-        output_type="pt",
-        extra_args={"flow_shift": 4.0, "batch_cfg": True},
-    )
-
-    out = pipeline.forward(req)
-
-    assert torch.equal(out.output, torch.zeros(9, 192, 320, 3))
-    assert len(calls) == 1
-    call = calls[0]
-    assert call["prompt"] == "a robot arm"
-    assert call["negative_prompt"] == "bad hands"
-    assert call["height"] == 192
-    assert call["width"] == 320
-    assert call["num_frames"] == 9
-    assert call["num_inference_steps"] == 2
-    assert call["guidance_scale"] == 3.0
-    assert call["shift"] == 4.0
-    assert call["batch_cfg"] is True
-    assert call["output_type"] == "pt"
-
-
-def test_forward_prefers_shift_over_flow_shift_and_defaults_negative_prompt():
-    pipeline = _make_pipeline()
-    calls = []
-
-    def fake_generate(**kwargs):
-        calls.append(kwargs)
-        return torch.zeros(5, 192, 192, 3)
-
-    pipeline._generate = fake_generate
-    req = _make_request_batch(
-        "a robot arm",
-        height=192,
-        width=192,
-        num_frames=5,
-        num_inference_steps=2,
-        guidance_scale=1.0,
-        seed=42,
-        extra_args={"shift": 5.0, "flow_shift": 4.0},
-    )
-
-    pipeline.forward(req)
-
-    assert calls[0]["shift"] == 5.0
-    assert calls[0]["negative_prompt"] == "default negative"
-
-
-def test_forward_rejects_multi_request_batches():
-    pipeline = _make_pipeline()
-    first = OmniDiffusionRequest(
-        prompt="first",
-        sampling_params=OmniDiffusionSamplingParams(seed=1),
-        request_id="lingbot-first",
-    )
-    second = OmniDiffusionRequest(
-        prompt="second",
-        sampling_params=OmniDiffusionSamplingParams(seed=2),
-        request_id="lingbot-second",
-    )
-
-    with pytest.raises(ValueError, match="only supports one request"):
-        pipeline.forward(DiffusionRequestBatch([first, second]))
-
-
-def test_load_weights_rejects_external_weight_stream():
-    pipeline = _make_pipeline()
-
-    assert pipeline.load_weights([]) == set()
-    with pytest.raises(RuntimeError, match="components are loaded directly"):
-        pipeline.load_weights([("transformer.weight", torch.zeros(1))])

tests/e2e/offline_inference/test_flux2_klein.py

@@ -18,8 +18,6 @@ from tests.helpers.runtime import DiffusionResponse, OmniRunnerHandler
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
-# NOTE: Most of the flux2klein tests have migrated to the common diffusion tests.
-# This model will also use the tiny model builder when run at the core model level.
 MODEL = "black-forest-labs/FLUX.2-klein-4B"
 
 _OMNI_RUNNER_PARAM = (MODEL, None)
@@ -72,6 +70,41 @@ def _send_inpaint_with_generator(
     )
 
 
+# Regression test for https://github.com/vllm-project/vllm-omni/issues/3097
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+def test_flux2_klein_can_accept_text_inputs(omni_runner_handler: OmniRunnerHandler):
+    omni_runner_handler.send_diffusion_request(
+        {
+            "model": MODEL,
+            "prompt": "a cup of coffee on the table",
+            "sampling_params": OmniDiffusionSamplingParams(num_inference_steps=2, seed=42),
+        }
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_flux2_klein_inpaint_basic(omni_runner_handler: OmniRunnerHandler):
+    input_image, mask_image = _create_test_inputs()
+    omni_runner_handler.send_diffusion_request(
+        {
+            "model": MODEL,
+            "prompt": "Fill in the masked area with a beautiful garden",
+            "multi_modal_data": {"image": input_image, "mask_image": mask_image},
+            "sampling_params": OmniDiffusionSamplingParams(
+                height=_HEIGHT,
+                width=_WIDTH,
+                num_inference_steps=_NUM_INFERENCE_STEPS,
+                guidance_scale=0.0,
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
+                num_outputs_per_prompt=1,
+            ),
+        }
+    )
+
+
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "L4"}, num_cards=1)

tests/e2e/offline_inference/test_mammoth_moda2_expansion.py

@@ -5,7 +5,7 @@ Verifies that the AR->DiT pipeline produces an image tensor whose pixel values
 match a golden reference.
 
 Model Hub repo id: ``bytedance-research/MammothModa2-Preview``.
-Deploy config: ``get_deploy_config_path("mammoth_moda2.yaml")`` -> ``vllm_omni/deploy/mammoth_moda2.yaml``
+Stage config: ``get_deploy_config_path("mammoth_moda2.yaml")`` → ``vllm_omni/deploy/mammoth_moda2.yaml``
 
 Golden pixel file: ``tests/e2e/offline_inference/fixtures/mammoth_moda2_t2i_golden.json``
   Regenerate with: ``UPDATE_GOLDEN=1 pytest tests/e2e/offline_inference/test_mammoth_moda2_expansion.py``
@@ -36,9 +36,9 @@ _VISION_END_TOKEN_ID = 151653  # "<|vision_end|>"
 _AR_PATCH_SIZE = 16
 
 MODEL_PATH = "bytedance-research/MammothModa2-Preview"
-T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
+T2I_STAGE_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
 
-_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG})
+_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_STAGE_CONFIG)
 
 # Golden pixel reference file.  Set UPDATE_GOLDEN=1 to regenerate.
 _GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"

tests/e2e/offline_inference/test_omnivoice_expansion.py

@@ -18,15 +18,13 @@ from tests.helpers.runtime import OmniRunner
 from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL = "k2-fsa/OmniVoice"
-DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml")
+STAGE_CONFIG = get_deploy_config_path("omnivoice.yaml")
 
-# OmniRunner tuple: model, legacy stage config path, extra Omni kwargs.
-# The migrated test passes deploy_config through extra Omni kwargs.
+# (model, stage_config_path, extra_omni_kwargs) — see ``omni_runner`` in tests.helpers.fixtures.runtime
 _OMNI_RUNNER_PARAM = (
     MODEL,
-    None,
+    STAGE_CONFIG,
     {
-        "deploy_config": DEPLOY_CONFIG,
         "trust_remote_code": True,
         "log_stats": True,
     },

tests/e2e/online_serving/test_flux2_klein.py

@@ -0,0 +1,62 @@
+"""
+Tests for Flux2 Klein; currently Dev is implemented separately,
+but ideally these models will fold together in the future.
+"""
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import (
+    OmniServer,
+    OmniServerParams,
+    OpenAIClientHandler,
+    dummy_messages_from_mix_data,
+)
+
+FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4)
+POSITIVE_PROMPT = "A cat sitting on a windowsill"
+NEGATIVE_PROMPT = "blurry, low quality"
+
+
+# For now, we only test a good configuration for this model to keep the CI
+# light, but in the near future may explore using random models to check
+# cross-feature compatibility more generally.
+def _get_diffusion_feature_cases(model: str):
+    return [
+        # FP8 / Hybrid sequence parallelism
+        pytest.param(
+            OmniServerParams(
+                model=model,
+            ),
+            id="default",
+            marks=FOUR_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases(
+        model="black-forest-labs/FLUX.2-klein-4B",
+    ),
+    indirect=True,
+)
+def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT)
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "true_cfg_scale": 4.0,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)

tests/e2e/online_serving/test_flux2_klein_expansion.py

@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end tests for Flux2 Klein inpainting in online serving mode.
+
+Uses /v1/images/edits endpoint which is the correct API for image inpainting.
+"""
+
+import base64
+from io import BytesIO
+
+import httpx
+import pytest
+from PIL import Image, ImageDraw
+
+from tests.helpers.runtime import OmniServer, OmniServerParams
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+
+MODEL = "black-forest-labs/FLUX.2-klein-4B"
+
+_HEIGHT = 512
+_WIDTH = 512
+_NUM_INFERENCE_STEPS = 4
+
+
+def _get_diffusion_feature_cases(model: str):
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--tensor-parallel-size", "2", "--enable-cpu-offload"],
+            ),
+            id="tp2_cpu_offload",
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--tensor-parallel-size", "2"],
+            ),
+            id="tp2_basic",
+        ),
+    ]
+
+
+def _image_to_base64_jpeg(image: Image.Image) -> str:
+    buffer = BytesIO()
+    image.save(buffer, format="JPEG")
+    buffer.seek(0)
+    return base64.b64encode(buffer.read()).decode("utf-8")
+
+
+def _create_test_mask_base64(width: int = _WIDTH, height: int = _HEIGHT) -> str:
+    mask = Image.new("L", (width, height), 0)
+    draw = ImageDraw.Draw(mask)
+    draw.rectangle([width // 4, height // 4, width * 3 // 4, height * 3 // 4], fill=255)
+    return _image_to_base64_jpeg(mask)
+
+
+def _compare_images(img1: Image.Image, img2: Image.Image) -> bool:
+    return list(img1.getdata()) == list(img2.getdata())
+
+
+def _send_edit_request(host: str, port: int, model: str, image_b64: str, mask_b64: str, prompt: str, **kwargs):
+    url = f"http://{host}:{port}/v1/images/edits"
+    files = {
+        "image": ("image.jpg", base64.b64decode(image_b64), "image/jpeg"),
+        "mask_image": ("mask.jpg", base64.b64decode(mask_b64), "image/jpeg"),
+    }
+    data = {"prompt": prompt, "model": model, **kwargs}
+    with httpx.Client(timeout=60.0) as client:
+        response = client.post(url, files=files, data=data)
+        response.raise_for_status()
+        return response.json()
+
+
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_flux2_klein_inpaint_basic(omni_server: OmniServer):
+    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
+    mask_b64 = _create_test_mask_base64()
+
+    result = _send_edit_request(
+        host=omni_server.host,
+        port=omni_server.port,
+        model=MODEL,
+        image_b64=input_image_b64,
+        mask_b64=mask_b64,
+        prompt="Fill in the masked area with a beautiful garden",
+        guidance_scale=1.0,
+        num_inference_steps=_NUM_INFERENCE_STEPS,
+        n=1,
+        seed=42,
+    )
+
+    assert "data" in result and len(result["data"]) == 1
+    img_data = result["data"][0].get("b64_json") or result["data"][0].get("url", "").split(",")[-1]
+    img = Image.open(BytesIO(base64.b64decode(img_data)))
+    assert img.size == (_WIDTH, _HEIGHT)
+
+
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_flux2_klein_inpaint_deterministic(omni_server: OmniServer):
+    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
+    mask_b64 = _create_test_mask_base64()
+    prompt = "A red flower in a field"
+
+    result1 = _send_edit_request(
+        host=omni_server.host,
+        port=omni_server.port,
+        model=MODEL,
+        image_b64=input_image_b64,
+        mask_b64=mask_b64,
+        prompt=prompt,
+        guidance_scale=1.0,
+        num_inference_steps=_NUM_INFERENCE_STEPS,
+        n=1,
+        seed=12345,
+    )
+
+    result2 = _send_edit_request(
+        host=omni_server.host,
+        port=omni_server.port,
+        model=MODEL,
+        image_b64=input_image_b64,
+        mask_b64=mask_b64,
+        prompt=prompt,
+        guidance_scale=1.0,
+        num_inference_steps=_NUM_INFERENCE_STEPS,
+        n=1,
+        seed=12345,
+    )
+
+    img1_data = result1["data"][0].get("b64_json") or result1["data"][0].get("url", "").split(",")[-1]
+    img2_data = result2["data"][0].get("b64_json") or result2["data"][0].get("url", "").split(",")[-1]
+
+    img1 = Image.open(BytesIO(base64.b64decode(img1_data)))
+    img2 = Image.open(BytesIO(base64.b64decode(img2_data)))
+
+    assert _compare_images(img1, img2), (
+        "Same input with same seed should produce identical output. This is critical for offline/online consistency."
+    )
+
+
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_flux2_klein_inpaint_multiple_outputs(omni_server: OmniServer):
+    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
+    mask_b64 = _create_test_mask_base64()
+
+    result = _send_edit_request(
+        host=omni_server.host,
+        port=omni_server.port,
+        model=MODEL,
+        image_b64=input_image_b64,
+        mask_b64=mask_b64,
+        prompt="A beautiful landscape",
+        guidance_scale=1.0,
+        num_inference_steps=_NUM_INFERENCE_STEPS,
+        n=2,
+        seed=42,
+    )
+
+    assert "data" in result and len(result["data"]) == 2

tests/e2e/online_serving/test_lingbot_video.py

@@ -1,61 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Online serving smoke for ``robbyant/lingbot-video-dense-1.3b``.
-
-The LingBot dense PR supports text-to-video only, so this module exercises the
-native ``/v1/videos`` path with a small 9-frame request.
-"""
-
-import os
-
-import pytest
-
-from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
-MODEL = "robbyant/lingbot-video-dense-1.3b"
-PROMPT = "a robotic arm picks up a red block"
-NEGATIVE_PROMPT = "low quality, blurry, watermark, text"
-
-SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
-
-
-def _get_diffusion_feature_cases(model: str):
-    return [
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=["--model-class-name", "LingBotVideoPipeline"],
-            ),
-            id="default",
-            marks=SINGLE_CARD_FEATURE_MARKS,
-        ),
-    ]
-
-
-@pytest.mark.core_model
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_text_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
-    request_config = {
-        "model": omni_server.model,
-        "form_data": {
-            "model": omni_server.model,
-            "prompt": PROMPT,
-            "negative_prompt": NEGATIVE_PROMPT,
-            "height": 192,
-            "width": 320,
-            "num_frames": 9,
-            "fps": 24,
-            "num_inference_steps": 2,
-            "guidance_scale": 3.0,
-            "flow_shift": 3.0,
-            "seed": 42,
-        },
-    }
-    openai_client.send_video_diffusion_request(request_config)

tests/e2e/online_serving/test_lingbot_video_expansion.py

@@ -1,82 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-L4 expansion coverage for ``robbyant/lingbot-video-dense-1.3b``.
-
-This is intentionally T2V-only. T2I, I2V, TI2V, MoE, and multi-GPU feature
-rows should be added in follow-up PRs when those paths are implemented.
-"""
-
-import json
-import os
-
-import pytest
-
-from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
-
-MODEL = "robbyant/lingbot-video-dense-1.3b"
-PROMPT = "a robotic arm picks up a red block"
-NEGATIVE_PROMPT = "low quality, blurry, watermark, text"
-
-SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
-
-
-def _get_diffusion_feature_cases(model: str):
-    return [
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=["--model-class-name", "LingBotVideoPipeline"],
-            ),
-            id="default",
-            marks=SINGLE_CARD_FEATURE_MARKS,
-        ),
-    ]
-
-
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_cfg_off(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
-    request_config = {
-        "model": omni_server.model,
-        "form_data": {
-            "model": omni_server.model,
-            "prompt": PROMPT,
-            "height": 192,
-            "width": 192,
-            "num_frames": 9,
-            "fps": 24,
-            "num_inference_steps": 2,
-            "guidance_scale": 1.0,
-            "flow_shift": 3.0,
-            "seed": 42,
-        },
-    }
-    openai_client.send_video_diffusion_request(request_config)
-
-
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_batch_cfg_extra_params(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
-    request_config = {
-        "model": omni_server.model,
-        "form_data": {
-            "model": omni_server.model,
-            "prompt": PROMPT,
-            "negative_prompt": NEGATIVE_PROMPT,
-            "height": 192,
-            "width": 320,
-            "num_frames": 9,
-            "fps": 24,
-            "num_inference_steps": 2,
-            "guidance_scale": 3.0,
-            "flow_shift": 3.0,
-            "seed": 42,
-            "extra_params": json.dumps({"batch_cfg": True}, separators=(",", ":")),
-        },
-    }
-    openai_client.send_video_diffusion_request(request_config)

tests/e2e/online_serving/test_qwen3_omni.py

@@ -226,27 +226,3 @@ def test_thinker_prefix_caching_audio_output(omni_server, openai_client) -> None
     }
 
     _run_prefix_cache_check(openai_client, request_config)
-
-
-@pytest.mark.advanced_model
-@pytest.mark.core_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_completions_rejected_for_thinker_talker(omni_server, openai_client) -> None:
-    """Ensure Thinker-talker models reject /v1/completions; we do this because the
-    thinker-talker handoff implementations currently use ChatML <|im_start|> and
-    <|im_end|> markers to segment the input sequence; when we don't have them,
-    the talker does not get any embeddings, which breaks the server.
-    """
-    responses = openai_client.send_completions_http_request(
-        {
-            "json": {
-                "model": omni_server.model,
-                "prompt": "Hello, how are you?",
-                "max_tokens": 10,
-            },
-        },
-        err_code=400,
-    )
-    assert not responses[0].success

tests/entrypoints/test_omni_entrypoints.py

@@ -107,7 +107,6 @@ class FakeAsyncOmniEngine:
         self.stage_vllm_configs = [None for _ in range(self.num_stages)]
         self.output_processors = [SimpleNamespace(tokenizer=None) for _ in range(self.num_stages)]
         self.input_processor = None
-        self.endpoint_restrictions = ()
 
         self.output_q: queue.Queue[Any] = queue.Queue()
         self.submitted: list[dict[str, Any]] = []

tests/entrypoints/test_resolve_dreamzero_config.py

@@ -24,8 +24,8 @@ def test_dreamzero_vla_resolves_to_dreamzero_config(monkeypatch):
 
 def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
     monkeypatch.setattr(
-        "vllm_omni.config.config_factory.StageConfigFactory._try_infer_model_type",
-        classmethod(lambda _cls, model, trust_remote_code=True: "vla"),
+        "vllm_omni.config.config_factory.StageConfigFactory._auto_detect_model_type",
+        classmethod(lambda _cls, _model, trust_remote_code=True: ("vla", None)),
     )
     monkeypatch.setattr(
         "vllm_omni.diffusion.utils.hf_utils._looks_like_dreamzero",

tests/examples/offline_inference/test_image_to_video.py

@@ -1,81 +0,0 @@
-"""
-Offline inference tests: image-to-video.
-See examples/offline_inference/image_to_video/README.md
-"""
-
-import shutil
-from pathlib import Path
-
-import pytest
-import requests
-
-from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet
-from tests.helpers.assertions import assert_video_valid
-from tests.helpers.mark import hardware_marks
-
-pytestmark = [
-    pytest.mark.usefixtures("clean_gpu_memory_between_tests"),
-    pytest.mark.full_model,
-    pytest.mark.example,
-    *hardware_marks(res={"cuda": "H100"}),
-]
-
-I2V_SCRIPT = EXAMPLES / "offline_inference" / "image_to_video" / "image_to_video.py"
-README_PATH = I2V_SCRIPT.with_name("README.md")
-EXAMPLE_OUTPUT_SUBFOLDER = "example_offline_i2v"
-
-_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
-_IMAGE_NAME = "cherry_blossom.jpg"
-
-_SKIP_SECTIONS = {
-    "Prerequisites",
-    "Advanced Features",
-    "FAQ",
-    # VACE conditional tasks are covered by test_vace_video_generation.py, which
-    # runs the same README snippets with synthetic assets and smoke settings. The
-    # prep block here downloads inputs via wget (not portable, not collected), so
-    # exercising this section from the shared runner would only duplicate and fail.
-    "Wan2.1 VACE Conditional Tasks",
-}
-
-
-def _ensure_test_image(run_dir: Path) -> None:
-    """Download or copy the example image into run_dir so CLI snippets can find it."""
-    dest = run_dir / _IMAGE_NAME
-    if dest.exists():
-        return
-    src = I2V_SCRIPT.parent / _IMAGE_NAME
-    if src.exists():
-        shutil.copy2(src, dest)
-        return
-    response = requests.get(_IMAGE_URL, timeout=60)
-    response.raise_for_status()
-    dest.write_bytes(response.content)
-
-
-def _skip_readme_snippet(language: str, code: str, h2_title: str) -> tuple[bool, str]:
-    if h2_title in _SKIP_SECTIONS:
-        return True, f"README section '{h2_title}' is intentionally excluded for examples tests"
-    if language == "python":
-        return True, "Python API snippets produce video files that ExampleRunner does not auto-collect"
-    if "/path/to/" in code:
-        return True, "Snippet references a placeholder local model path"
-    return False, ""
-
-
-README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH, skipif=_skip_readme_snippet)
-
-
-@pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id)
-def test_image_to_video(snippet: ReadmeSnippet, example_runner: ExampleRunner):
-    should_skip, reason = snippet.skip
-    if should_skip:
-        pytest.skip(reason)
-
-    run_dir = example_runner.output_root / EXAMPLE_OUTPUT_SUBFOLDER / snippet.test_id
-    run_dir.mkdir(parents=True, exist_ok=True)
-    _ensure_test_image(run_dir)
-
-    result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER))
-    for asset in result.assets:
-        assert_video_valid(asset)

tests/helpers/assertions.py

@@ -19,20 +19,19 @@ from PIL import Image
 from tests.helpers.media import (
     convert_audio_bytes_to_text,
     cosine_similarity_text,
-    preprocess_text,
 )
 
 _GENDER_PIPELINE = None
 _GENDER_PIPELINE_LOCK = threading.Lock()
 # Transcript gates default to whisper ``small`` for speed. ``small`` mishears a
 # short TTS clip ~0.5% of the time (e.g. "Hello"->"fellow", or hallucinating a
-# leading SFX token), which flakes the deterministic similarity gate. Short
-# clips first get a conservative containment fallback for minor ASR repeats/noise.
-# A test can also opt in to ASR escalation by setting
-# ``transcript_escalation_model`` to a whisper model name (e.g. ``"large-v3"``)
-# in its request_config: on a failed fast pass the clip is re-transcribed with
-# that stronger ASR before the test fails, so a weak-ASR mishear is rescued while
-# a genuine model artifact still fails (the strong ASR mismatches too).
+# leading SFX token), which flakes the deterministic similarity gate. A test can
+# opt in to ASR escalation by setting ``transcript_escalation_model`` to a
+# whisper model name (e.g. ``"large-v3"``) in its request_config: on a failed
+# fast pass the clip is re-transcribed with that stronger ASR before the test
+# fails, so a weak-ASR mishear is rescued while a genuine model artifact still
+# fails (the strong ASR mismatches too). Tests that do not opt in keep the
+# strict behaviour.
 _PCM_SPEECH_SAMPLE_RATE_HZ = 24_000
 _MIN_PCM_SPEECH_HNR_DB = 1.0
 _PRESET_VOICE_GENDER_MAP: dict[str, str] = {
@@ -44,23 +43,6 @@ _PRESET_VOICE_GENDER_MAP: dict[str, str] = {
 }
 
 
-def _short_transcript_contains_expected(transcript: str, expected: str) -> bool:
-    """Allow minor ASR repeats/noise for very short speech clips."""
-    transcript_clean = preprocess_text(transcript)
-    expected_clean = preprocess_text(expected)
-    if not transcript_clean or not expected_clean:
-        return False
-
-    transcript_words = transcript_clean.split()
-    expected_words = expected_clean.split()
-    if not transcript_words or not expected_words:
-        return False
-
-    short_text = min(len(transcript_clean), len(expected_clean)) <= 15
-    small_word_delta = len(transcript_words) <= len(expected_words) + 2
-    return short_text and small_word_delta and expected_clean in transcript_clean
-
-
 def assert_image_diffusion_response(
     response: "DiffusionResponse",
     request_config: dict[str, Any],
@@ -638,16 +620,13 @@ def _assert_transcript_matches(
     ``transcript`` is the fast whisper-``small`` result. If it clears
     ``threshold`` the check passes immediately.
 
-    If the cosine check fails, very short clips get a conservative containment
-    fallback that accepts minor ASR repeats/noise only when the expected text is
-    still present and the transcript has few extra words.
-
     When ``escalation_model`` is set (opt-in via the ``transcript_escalation_model``
-    request_config key) and the fast check plus containment fallback fail, the
-    clip is re-transcribed with that stronger ASR and the assertion is decided on
-    its verdict -- so a weak whisper-``small`` mishear on a short clip does not
-    flake the gate, while a genuine model artifact still fails (the strong ASR
-    mismatches too).
+    request_config key) and the fast check fails, the clip is re-transcribed with
+    that stronger ASR and the assertion is decided on its verdict -- so a weak
+    whisper-``small`` mishear on a short clip does not flake the gate, while a
+    genuine model artifact still fails (the strong ASR mismatches too). When
+    ``escalation_model`` is ``None`` the original strict behaviour is preserved,
+    so other tests are unaffected.
     """
     expected = str(expected_text).strip().lower()
     similarity = cosine_similarity_text(transcript.strip().lower(), expected)
@@ -655,10 +634,6 @@ def _assert_transcript_matches(
     if similarity > threshold:
         return
 
-    if _short_transcript_contains_expected(transcript, expected):
-        print("short speech containment check passed")
-        return
-
     if escalation_model and audio_bytes:
         print(
             f"whisper-small below threshold ({similarity:.2f} <= {threshold}); "

tests/helpers/runtime.py

@@ -47,7 +47,6 @@ from tests.helpers.media import (
     _merge_base64_audio_to_segment,
     decode_b64_image,
 )
-from tests.model_tests.diffusion.utils import resolve_tiny_model_path
 from vllm_omni.config.stage_config import resolve_deploy_yaml
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.outputs import OmniRequestOutput
@@ -1065,29 +1064,6 @@ class OpenAIClientHandler:
         )
         return [resp]
 
-    def send_completions_http_request(
-        self,
-        request_config: dict[str, Any],
-        *,
-        err_code: int | tuple[int, ...] | list[int] | None = None,
-        err_message: str | tuple[str, ...] | list[str] | None = None,
-    ) -> list[HttpResponse]:
-        """POST ``/v1/completions`` with ``json`` or ``raw_body``."""
-        # TODO (Alex): A lot of these helpers should be consolidated as they differ only by endpoint
-        cfg = _merge_http_expectation_kwargs(
-            request_config,
-            err_code=err_code,
-            err_message=err_message,
-        )
-        r = self._post_json_endpoint("/v1/completions", cfg, default_timeout=120.0)
-        resp = self._http_response_from_requests(r)
-        assert_http_error(
-            resp,
-            err_code=cfg.get("err_code"),
-            err_message=cfg.get("err_message"),
-        )
-        return [resp]
-
     def send_omni_sleep_http_request(
         self,
         request_config: dict[str, Any],
@@ -2816,24 +2792,11 @@ def iter_omni_server(
 
     with omni_fixture_lock:
         params: OmniServerParams = request.param
-        # For now, when a tiny model is substituted, we preserve the original model
-        # name via --served-model-name (so that the server still accepts requests with
-        # the original name). We also do the same for server.model so that tests reading
-        # server.model send the correct name in requests.
-        #
-        # TODO: core models on this path currently do not clean up tiny models, although
-        # tiny model paths are deterministic, so it's not a huge footprint. Still, it would
-        # be ideal to cleanup consistently everywhere.
-        original_model = model_prefix + params.model
-        model = original_model
-        if run_level == "core_model":
-            model = resolve_tiny_model_path(model)
+        model = model_prefix + params.model
         port = params.port
         stage_config_path = stage_config_path_for_run_level(params.stage_config_path, run_level)
 
         server_args = params.server_args or []
-        if model != original_model:
-            server_args = [*server_args, "--served-model-name", original_model]
         if params.use_omni and params.stage_init_timeout is not None:
             server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)]
         else:
@@ -2859,8 +2822,6 @@ def iter_omni_server(
                 port=port,
                 env_dict=params.env_dict,
             ) as server:
-                if model != original_model:
-                    server.model = original_model
                 print("OmniServer started successfully")
                 yield server
                 print("OmniServer stopping...")
@@ -2884,8 +2845,6 @@ def iter_omni_server(
                     use_omni=params.use_omni,
                 )
             ) as server:
-                if model != original_model:
-                    server.model = original_model
                 print("OmniServer started successfully")
                 yield server
                 print("OmniServer stopping...")
@@ -2917,8 +2876,6 @@ def iter_omni_runner(
             extra_omni_kwargs = dict(extra) if extra is not None else {}
         stage_config_path = stage_config_path_for_run_level(stage_config_path, run_level)
         model = model_prefix + model
-        if run_level == "core_model":
-            model = resolve_tiny_model_path(model)
         with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner:
             print("OmniRunner started successfully")
             yield runner

tests/helpers/test_assertions.py

@@ -1,25 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-
-from tests.helpers.assertions import _assert_transcript_matches
-
-
-def test_short_transcript_repeat_passes_containment_fallback():
-    _assert_transcript_matches(
-        " How... how are you?",
-        audio_bytes=None,
-        expected_text="how are you",
-        threshold=0.9,
-    )
-
-
-def test_short_transcript_unrelated_text_still_fails():
-    with pytest.raises(AssertionError, match="Transcript doesn't match input"):
-        _assert_transcript_matches(
-            " I don't know, sorry.",
-            audio_bytes=None,
-            expected_text="how are you",
-            threshold=0.9,
-        )

tests/model_executor/models/voxtral_tts/test_tts_preprocess.py

@@ -1,78 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import functools
-
-import pytest
-import torch
-import torch.nn as nn
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-@functools.lru_cache(maxsize=1)
-def _voxtral_tts_model_cls():
-    from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules
-
-    bootstrap_vllm_layer_custom_op_modules()
-    import vllm.model_executor.models.utils  # noqa: F401
-
-    from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts import (
-        VoxtralTTSForConditionalGeneration,
-    )
-
-    return VoxtralTTSForConditionalGeneration
-
-
-class FakeAudioGeneration(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.embed_multimodal_calls = []
-
-    def embed_multimodal(self, **kwargs):
-        self.embed_multimodal_calls.append(kwargs)
-        return [torch.full((1, 4), 7.0)]
-
-
-def _make_voxtral_tts_model():
-    model_cls = _voxtral_tts_model_cls()
-    model = model_cls.__new__(model_cls)
-    nn.Module.__init__(model)
-    model.model_stage = "audio_generation"
-    model.model = FakeAudioGeneration()
-    model._audio_token_id = 42
-    return model
-
-
-def test_tts_preprocess_consumes_nested_codes_audio_feedback():
-    model = _make_voxtral_tts_model()
-    input_ids = torch.tensor([model._audio_token_id])
-    input_embeds = torch.zeros((1, 4))
-    audio_tokens = torch.tensor([[1, 2, 3, 4]])
-
-    _, output_embeds, _ = model.tts_preprocess(
-        input_ids=input_ids,
-        input_embeds=input_embeds,
-        codes={"audio": audio_tokens},
-    )
-
-    assert len(model.model.embed_multimodal_calls) == 1
-    torch.testing.assert_close(model.model.embed_multimodal_calls[0]["audio_tokens"], audio_tokens)
-    torch.testing.assert_close(output_embeds, torch.full((1, 4), 7.0))
-
-
-def test_tts_preprocess_keeps_legacy_top_level_audio_feedback():
-    model = _make_voxtral_tts_model()
-    input_ids = torch.tensor([model._audio_token_id])
-    input_embeds = torch.zeros((1, 4))
-    audio_tokens = torch.tensor([[5, 6, 7, 8]])
-
-    _, output_embeds, _ = model.tts_preprocess(
-        input_ids=input_ids,
-        input_embeds=input_embeds,
-        audio=audio_tokens,
-    )
-
-    assert len(model.model.embed_multimodal_calls) == 1
-    torch.testing.assert_close(model.model.embed_multimodal_calls[0]["audio_tokens"], audio_tokens)
-    torch.testing.assert_close(output_embeds, torch.full((1, 4), 7.0))

tests/model_extras/test_model_extras.py

@@ -126,102 +126,6 @@ def test_magi_human_extra_registry_declares_request_and_response_params() -> Non
     assert should_init_extra_args_for_non_diffusion_stages("MagiHumanPipeline") is False
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_ming_flash_omni_extra_registry_declares_request_and_response_params() -> None:
-    assert get_extra_body_params("MingImagePipeline") == frozenset(
-        {
-            "height",
-            "width",
-            "steps",
-            "cfg",
-            "seed",
-            "byte5_text",
-            "negative_prompt",
-        }
-    )
-    assert get_extra_output_params("MingImagePipeline") == frozenset()
-    assert should_init_extra_args_for_non_diffusion_stages("MingImagePipeline") is True
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_ming_flash_omni_declared_extra_args_route_into_sampling_params() -> None:
-    params = OmniDiffusionSamplingParams()
-    declared = get_extra_body_params("MingImagePipeline")
-    apply_declared_extra_args(
-        params,
-        declared,
-        {
-            "steps": 6,
-            "cfg": 1.5,
-            "byte5_text": ["理解与生成统一"],
-            "negative_prompt": "ugly, blurry",
-            "unknown": "ignored",
-        },
-    )
-    assert params.extra_args == {
-        "steps": 6,
-        "cfg": 1.5,
-        "byte5_text": ["理解与生成统一"],
-        "negative_prompt": "ugly, blurry",
-    }
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_ming_flash_omni_text_to_image_prompt_builder() -> None:
-    assert build_text_to_image_prompt(
-        "MingImagePipeline",
-        prompt="Please draw a cute cat.",
-        negative_prompt="ugly, blurry",
-        height=512,
-        width=768,
-    ) == {
-        "prompt": "Please draw a cute cat.",
-        "modalities": ["image"],
-        "mm_processor_kwargs": {
-            "modalities": ["image"],
-            "target_h": 512,
-            "target_w": 768,
-        },
-        "negative_prompt": "ugly, blurry",
-    }
-    # target_h/w are omitted when height/width are not supplied so the
-    # pipeline's 1024x1024 default still applies.
-    assert build_text_to_image_prompt(
-        "MingImagePipeline",
-        prompt="Draw a poster.",
-        negative_prompt=None,
-    ) == {
-        "prompt": "Draw a poster.",
-        "modalities": ["image"],
-        "mm_processor_kwargs": {"modalities": ["image"]},
-    }
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_ming_flash_omni_image_to_image_prompt_builder() -> None:
-    dummy_image = Image.new("RGB", (64, 64))
-    result = build_image_to_image_prompt(
-        "MingImagePipeline",
-        prompt="Change the background to a sandy beach at sunset.",
-        negative_prompt=None,
-        input_image=dummy_image,
-        height=256,
-        width=256,
-    )
-    assert result["modalities"] == ["img2img"]
-    assert result["multi_modal_data"]["img2img"] is dummy_image
-    assert result["mm_processor_kwargs"] == {
-        "modalities": ["img2img"],
-        "target_h": 256,
-        "target_w": 256,
-    }
-    assert "negative_prompt" not in result
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_cosmos3_text_to_image_prompt_builder_selects_image_modality() -> None:

tests/model_tests/conftest.py

@@ -1,41 +0,0 @@
-from shutil import rmtree
-
-import pytest
-
-from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
-
-
-@pytest.fixture(scope="session")
-def tiny_model_paths(request, run_level):
-    """Build or download the tiny models for the selected tests.
-
-    At core_model level, builds tiny models via the builder function.
-    At advanced_model / full_model level, uses the real HF model path.
-
-    NOTE: this is session scoped to avoid churn in tiny model creation,
-    but will ensure all the tiny models you need are created for the selected tests
-    before it starts to execute them."""
-    model_paths = {}
-    built_paths = []
-    print("Initializing tiny models...")
-    for item in request.session.items:
-        if not hasattr(item, "callspec"):
-            continue
-        model_name = item.callspec.params.get("model_name")
-        if model_name is None:
-            continue
-        if model_name not in model_paths:
-            settings = DIFFUSION_TEST_SETTINGS[model_name]
-            if run_level == "core_model":
-                print(f"Calling tiny model builder for: {model_name}")
-                path = settings.builder()
-                built_paths.append(path)
-            else:
-                print(f"Run level is {run_level}; {model_name} will use the full model")
-                path = settings.model
-            model_paths[model_name] = path
-
-    yield model_paths
-    for path in built_paths:
-        print(f"Removing tiny model: {path}")
-        rmtree(path, ignore_errors=True)

tests/model_tests/diffusion/case_filtering.py

@@ -1,104 +0,0 @@
-"""
-Analogous to: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/vlm_utils/case_filtering.py
-
-This follows a similar pattern so that we can use `get_parametrized_options` to expand our test
-configurations out, similar to the way in which vLLM expands supported input type tests. For now,
-this mostly just means parametrizing over the test groups, which may be unique per model, but
-doing so in this way lets us cleanly define test groups that may be heterogeneous + avoid having
-excessive skip marks for features that may not be supported in every pipeline.
-"""
-
-# Engine initialization is currently quite expensive. For efficiency,
-# when we run the tests, we should try to collapse compatible accelerations,
-# because even with tiny models, the overhead of spinning up a new instance
-# stacks up fast.
-import itertools
-
-import pytest
-
-from tests.helpers.mark import hardware_marks
-from tests.model_tests.diffusion.config_types import DiffusionAccs, DiffusionModelTestOpts, get_required_device_count
-from vllm_omni.platforms import current_omni_platform
-
-# These tests are intended to run on L4 GPUs since they test parallelism (i.e., need GPUs),
-# but are intended to be smoke tests that we can properly validate without real model
-# weights in the CI.
-#
-# As such, no DiffusionAccs may require more than 4 GPUs, since we would otherwise skip
-# them in the CI.
-MAX_CI_DEVICES = 4
-
-
-def get_test_group_marks(model_name: str, test_group: list[DiffusionAccs] | None, model_marks: list | None) -> list:
-    """Build the full set of pytest marks for a test group.
-
-    For now, single device groups default to core model, and multi device groups
-    default to advanced_model marks. We also validate and throw if a test group exceeds
-    the max number of GPUs for the CI to prevent accidentally adding configurations that
-    would be skipped in the CI."""
-    marks = list(model_marks) if model_marks is not None else []
-
-    required_devices = get_required_device_count(test_group)
-    if required_devices > MAX_CI_DEVICES:
-        raise ValueError(
-            f"Test group {test_group} for {model_name} requires {required_devices} devices, "
-            f"but the max CI device count is {MAX_CI_DEVICES}. "
-        )
-
-    marks.extend(hardware_marks(res={"cuda": "L4"}, num_cards=required_devices))
-    # hardware_marks only adds skipif for num_cards > 1, so we currently handle the single-device case
-    # directly here. This should probably be handled in a more common way later on.
-    assert current_omni_platform is not None and current_omni_platform.device_count is not None
-    device_count = current_omni_platform.device_count()
-    if current_omni_platform.is_cuda() and device_count < required_devices:
-        marks.append(pytest.mark.skip(reason=f"Need {required_devices} devices, got {device_count}"))
-    if required_devices > 1:
-        marks.append(pytest.mark.full_model)
-    else:
-        marks.append(pytest.mark.core_model)
-    return marks
-
-
-def get_model_parametrization(model_name: str, test_info: DiffusionModelTestOpts, online: bool):
-    """Given a model & its corresponding test options, build the list of pytest params
-    to be run for this model. The base case (no accelerations) is always included. Extra
-    test groups are always appended for offline tests, but are only added to the online
-    tests if online_base_only is overridden to False to avoid redundant testing.
-    """
-    test_groups: list[list[DiffusionAccs] | None] = [None]
-    if not (online and test_info.online_base_only) and test_info.extra_test_groups:
-        test_groups.extend(test_info.extra_test_groups)
-
-    return [
-        pytest.param(
-            model_name,
-            test_group,
-            test_info.supported_tasks,
-            test_info.check_multi_output and test_group is None,
-            test_info.check_determinism and test_group is None,
-            id=f"{model_name}[{'+'.join(test_group)}]" if test_group else model_name,
-            marks=get_test_group_marks(model_name, test_group, test_info.marks),
-        )
-        for test_group in test_groups
-    ]
-
-
-def get_parametrized_options(
-    test_settings: dict[str, DiffusionModelTestOpts],
-    online: bool = False,
-):
-    """Converts all the DiffusionModelTestOpts into an expanded list of parameters
-    based on which accelerations are available.
-
-    When online=True and a model has online_base_only=True, only the base case
-    (no accelerations) is included for that model's parametrization to minimize
-    redundancy between offline / online tests.
-    """
-    # Get a list per model type, where each entry contains a tuple of all of
-    # that model type's cases, then flatten them into the top level so that
-    # we can consume them in one mark.parametrize call.
-    parametrization = [
-        get_model_parametrization(model_name, test_info, online=online)
-        for model_name, test_info in test_settings.items()
-    ]
-    return list(itertools.chain(*parametrization))

tests/model_tests/diffusion/config_types.py

@@ -1,187 +0,0 @@
-"""
-Common definitions for controlling what tests run where.
-"""
-
-from collections.abc import Callable
-from enum import StrEnum, auto
-from typing import NamedTuple, TypeAlias
-
-from pytest import MarkDecorator
-
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.entrypoints.omni import Omni
-
-# All builder funcs take no params and return a path
-TinyDiffusionBuilder: TypeAlias = Callable[[], str]
-
-
-class DiffusionAccs(StrEnum):
-    """Supported acceleration types / test settings for Diffusion Models."""
-
-    HSDP = auto()
-    TEA_CACHE = auto()
-    CACHE_DIT = auto()
-    SEQUENCE_PARALLEL = auto()
-    CFG_PARALLEL = auto()
-    TENSOR_PARALLEL = auto()
-    CPU_OFFLOAD = auto()
-    LAYERWISE_OFFLOAD = auto()
-    VAE_PATCH_PARALLEL = auto()
-
-
-class DiffusionTasks(StrEnum):
-    """Supported tasks for Diffusion Models."""
-
-    TEXT_TO_IMAGE = auto()
-    IMAGE_TO_IMAGE = auto()
-    # Text to video, text to audio, etc should be added here as needed
-
-
-class DiffusionModelTestOpts(NamedTuple):
-    """Configuration for one Diffusion model's tests."""
-
-    # HF model name for real-weight tests (advanced_model / full_model level).
-    # For now, whether we use the real weights vs tiny weights in the common tests
-    # depends on the run level.
-    model: str
-
-    # Creates a tiny model for the given architecture. We should always use tiny
-    # model weights for tests that do not require us to check the model quality.
-    builder: TinyDiffusionBuilder
-
-    # Actual tasks which controls the tests actually run
-    supported_tasks: list[DiffusionTasks]
-
-    # Additional acceleration groups to run beyond the base case (no acceleration).
-    # The base case is always run for every model in the test settings. None means
-    # we only run the base case.
-    extra_test_groups: list[list[DiffusionAccs]] | None = None
-
-    # Pytest Marks for this model. This may be useful for selecting which models
-    # we want to run where, similar to the way vLLM's multimodal tests mark some
-    # as core models to always run in the CI.
-    # Example: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/test_common.py#L131
-    marks: list[MarkDecorator] | None = None
-
-    # When True (default), online tests only run the base case (no accelerations).
-    # When False, online tests run all test_groups, same as offline. This should be
-    # True unless there is a good reason for it not to be, because the execution of
-    # the acceleration should be the same on both codepaths, and CLI parsing etc should
-    # be tested by tests adding the acceleration, and not per model.
-    online_base_only: bool = True
-
-    # Additional checks to run for the base case.
-    check_multi_output: bool = True  # Runs multiple generations in one request
-    check_determinism: bool = True  # Runs 2 generations with the same seed and check determinism
-
-
-### Mappings & utils for building offline Omni() instances given a list of enabled accelerations
-ACC_OMNI_KWARGS = {
-    DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_use_tiling": True},
-    DiffusionAccs.CPU_OFFLOAD: {"enable_cpu_offload": True},
-    DiffusionAccs.LAYERWISE_OFFLOAD: {"enable_layerwise_offload": True},
-    DiffusionAccs.CACHE_DIT: {"cache_backend": "cache_dit"},
-    DiffusionAccs.TEA_CACHE: {"cache_backend": "tea_cache"},
-}
-
-ACC_PARALLEL_KWARGS = {
-    DiffusionAccs.HSDP: {"use_hsdp": True, "hsdp_shard_size": 2},
-    DiffusionAccs.TENSOR_PARALLEL: {"tensor_parallel_size": 2},
-    DiffusionAccs.CFG_PARALLEL: {"cfg_parallel_size": 2},
-    DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_patch_parallel_size": 2},
-    # For SP, we don't run ring here to conserve devices, since it is easy
-    # to blow up the number of needed devices fast. Compatibility for ring
-    # and ulysses together should be tested generically outside of these tests.
-    DiffusionAccs.SEQUENCE_PARALLEL: {"ulysses_degree": 2},
-}
-
-ACC_DEVICE_COUNT_KEYS = {
-    DiffusionAccs.TENSOR_PARALLEL: "tensor_parallel_size",
-    DiffusionAccs.CFG_PARALLEL: "cfg_parallel_size",
-    DiffusionAccs.SEQUENCE_PARALLEL: "ulysses_degree",
-    DiffusionAccs.VAE_PATCH_PARALLEL: "vae_patch_parallel_size",
-    DiffusionAccs.HSDP: "hsdp_shard_size",
-}
-
-### CLI args for launching an OmniServer subprocess per acceleration.
-# Follows the same pattern as e2e online serving tests, which hand-code
-# their server_args lists in OmniServerParams.
-ACC_SERVER_ARGS: dict[DiffusionAccs, list[str]] = {
-    DiffusionAccs.HSDP: ["--use-hsdp", "--hsdp-shard-size", "2"],
-    DiffusionAccs.TEA_CACHE: ["--cache-backend", "tea_cache"],
-    DiffusionAccs.CACHE_DIT: ["--cache-backend", "cache_dit"],
-    DiffusionAccs.SEQUENCE_PARALLEL: ["--usp", "2"],
-    DiffusionAccs.CFG_PARALLEL: ["--cfg-parallel-size", "2"],
-    DiffusionAccs.TENSOR_PARALLEL: ["--tensor-parallel-size", "2"],
-    DiffusionAccs.CPU_OFFLOAD: ["--enable-cpu-offload"],
-    DiffusionAccs.LAYERWISE_OFFLOAD: ["--enable-layerwise-offload"],
-    DiffusionAccs.VAE_PATCH_PARALLEL: ["--vae-use-tiling", "--vae-patch-parallel-size", "2"],
-}
-
-## TODO ^ These should be in the same object, it's getting messy.
-
-
-def get_required_device_count(accelerations: list[DiffusionAccs] | None) -> int:
-    """Compute the minimum number of devices needed for a set of accelerations.
-    The total is the product of all parallel dimensions (defaulting to 1).
-
-    If not enough devices are available for a test group's accelerations,
-    that test will be skipped."""
-    count = 1
-    if accelerations is None:
-        return count
-
-    for acc in accelerations:
-        key = ACC_DEVICE_COUNT_KEYS.get(acc)
-        if key is not None:
-            count *= ACC_PARALLEL_KWARGS[acc][key]
-    return count
-
-
-def build_parallel_config_from_diff_accelerations(accelerations: list[DiffusionAccs]) -> DiffusionParallelConfig | None:
-    """Given a list of accelerations pertaining to the current test group,
-    build the parallel config needed for the Omni() object (if any)."""
-    config_kwargs = {}
-    for acc in accelerations:
-        update_dict = ACC_PARALLEL_KWARGS.get(acc, {})
-        config_kwargs.update(update_dict)
-    if config_kwargs:
-        return DiffusionParallelConfig(**config_kwargs)
-    return None
-
-
-### Offline Omni() object builder
-def build_omni_from_diff_accelerations(accelerations: list[DiffusionAccs] | None, **kwargs) -> Omni:
-    """Given one or more acceleration types, build the corresponding Omni() object."""
-    # Coerce to a list and build the parallel config, since that depends on the accelerations
-    if accelerations is None:
-        accelerations = []
-    parallel_config = build_parallel_config_from_diff_accelerations(accelerations)
-
-    # Then add anything else that's a top-level kwarg
-    acc_kwargs = {}
-    if parallel_config is not None:
-        acc_kwargs["parallel_config"] = parallel_config
-    for acc in accelerations:
-        update_dict = ACC_OMNI_KWARGS.get(acc, {})
-        acc_kwargs.update(update_dict)
-
-    # Keys passed through should mostly be things like enforce_eager;
-    # if there's overlap, it's probably due to a misconfiguration
-    shared_keys = acc_kwargs.keys() & kwargs.keys()
-    if shared_keys:
-        raise ValueError(f"Explicit Omni kwargs and inferred Omni kwargs for accelerations overlap: {shared_keys}")
-    omni_kwargs = {**acc_kwargs, **kwargs}
-    return Omni(**omni_kwargs)
-
-
-### Online server flag builder
-def build_server_args_from_diff_accelerations(accelerations: list[DiffusionAccs] | None) -> list[str]:
-    """Given one or more acceleration types, build the corresponding CLI args
-    for launching an OmniServer subprocess."""
-    if accelerations is None:
-        return []
-    args = []
-    for acc in accelerations:
-        args.extend(ACC_SERVER_ARGS.get(acc, []))
-    return args

tests/model_tests/diffusion/diff_model_builders.py

@@ -1,63 +0,0 @@
-import os
-import tempfile
-
-import torch
-from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
-from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel
-from diffusers.pipelines.flux2.pipeline_flux2_klein import Flux2KleinPipeline
-from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
-from transformers import AutoTokenizer, Qwen3Config, Qwen3ForCausalLM
-
-TINY_MODEL_DIR = os.path.join(tempfile.gettempdir(), "vllm-omni-tiny-models")
-
-
-def _get_tiny_model_path(name: str) -> str:
-    path = os.path.join(TINY_MODEL_DIR, name)
-    os.makedirs(path, exist_ok=True)
-    return path
-
-
-def tiny_flux2_klein_builder() -> str:
-    """Build a tiny Flux2Klein model."""
-    model_dir = _get_tiny_model_path("Flux2KleinPipeline")
-
-    pipe = Flux2KleinPipeline(
-        scheduler=FlowMatchEulerDiscreteScheduler(),
-        vae=AutoencoderKLFlux2(
-            down_block_types=("DownEncoderBlock2D",),
-            up_block_types=("UpDecoderBlock2D",),
-            block_out_channels=(32,),
-            layers_per_block=1,
-            latent_channels=16,
-            norm_num_groups=16,
-        ),
-        # NOTE: For now we need 28 layers because of hardcoded stuff in the model :(
-        text_encoder=Qwen3ForCausalLM(
-            Qwen3Config(
-                hidden_size=32,
-                intermediate_size=64,
-                num_hidden_layers=28,
-                num_attention_heads=2,
-                num_key_value_heads=2,
-                head_dim=16,
-                vocab_size=151936,
-                max_position_embeddings=512,
-            )
-        ),
-        tokenizer=AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B"),
-        # NOTE: For now we need at least 2 layers for the transformer
-        # due to hardcoded hacks in CacheDiT for Flux2Klein specifically.
-        transformer=Flux2Transformer2DModel(
-            in_channels=64,
-            num_layers=2,
-            num_single_layers=2,
-            attention_head_dim=16,
-            num_attention_heads=2,
-            joint_attention_dim=96,
-            timestep_guidance_channels=32,
-            axes_dims_rope=(4, 4, 4, 4),
-        ),
-    )
-    # Need dtypes to be consistent; for now we just put it on bfloat16
-    pipe.to(torch.bfloat16).save_pretrained(model_dir)
-    return model_dir

tests/model_tests/diffusion/model_settings.py

@@ -1,42 +0,0 @@
-from tests.model_tests.diffusion import diff_model_builders
-from tests.model_tests.diffusion.config_types import (
-    DiffusionAccs,
-    DiffusionModelTestOpts,
-    DiffusionTasks,
-)
-
-# This object defines the (tiny) model configurations for common tests.
-#
-# TO ADD A NEW MODEL:
-# Map the pipeline name (i.e., the name of the pipeline in _DIFFUSION_MODELS)
-# to a DiffusionModelTestOpts and define the supported tasks. A base case with no
-# accelerations is always run. Use extra_test_groups to define additional
-# acceleration configurations to run a basic smoke test for (i.e., simple inference).
-#
-# The model will only be loaded once per test group, and will execute each of its
-# supported tasks as a pytest subtest. The base case also runs additional checks
-# that are independent of accelerations, e.g., for determinism and compatibility with
-# multi-output, unless explicitly disabled in the test settings.
-#
-# TO RUN A SUBSET OF THE TESTS:
-# These tests should be fast, but if you'd like to run a subset of the tests, you can do so
-# with `-k`; the IDs of the tests are the name of the pipeline joined by +.
-#
-# Example:
-# $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline]
-#   ^ Runs only the base (no accelerations) case for Flux2KleinPipeline
-#
-# $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline
-#   ^ Runs all test groups for Flux2KleinPipeline only
-DIFFUSION_TEST_SETTINGS = {
-    "Flux2KleinPipeline": DiffusionModelTestOpts(
-        model="black-forest-labs/FLUX.2-klein-4B",
-        builder=diff_model_builders.tiny_flux2_klein_builder,
-        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE, DiffusionTasks.IMAGE_TO_IMAGE],
-        extra_test_groups=[
-            [DiffusionAccs.HSDP, DiffusionAccs.TEA_CACHE],
-            [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
-            [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
-        ],
-    )
-}

tests/model_tests/diffusion/task_runners.py

@@ -1,172 +0,0 @@
-"""
-Definitions running and validating individual tasks, e.g., text to image,
-image to image, and so on. These are called by the core test runner.
-"""
-
-import base64
-import io
-from dataclasses import replace
-
-import numpy as np
-from PIL import Image
-
-from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler, dummy_messages_from_mix_data
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.outputs import OmniRequestOutput
-
-PROMPT = "Dummy prompt"
-IMAGE_DIMS = (512, 512)
-HEIGHT, WIDTH = IMAGE_DIMS
-INPUT_IMAGE = Image.new("RGB", IMAGE_DIMS)
-
-# Offline sampling params
-IMAGE_GEN_SAMPLING_PARAMS = OmniDiffusionSamplingParams(
-    num_inference_steps=4,
-    height=HEIGHT,
-    width=WIDTH,
-    seed=42,
-)
-
-# Online extra_body for diffusion requests
-IMAGE_GEN_EXTRA_BODY = {
-    "height": HEIGHT,
-    "width": WIDTH,
-    "num_inference_steps": 4,
-    "seed": 42,
-}
-
-
-### Shared validation
-def _validate_images(images: list[Image.Image], expected_n: int = 1):
-    """Given a set of images, ensure we got the expected count of images,
-    and that all provided images match the expected dimensions."""
-    assert len(images) == expected_n
-    for img in images:
-        assert isinstance(img, Image.Image)
-        assert img.size == IMAGE_DIMS
-    return images
-
-
-def _validate_image_gen_determinism(images_a: list[Image.Image], images_b: list[Image.Image]):
-    """Ensure that two image sets are valid and that the results match."""
-    _validate_images(images_a)
-    _validate_images(images_b)
-    assert np.array_equal(np.array(images_a[0]), np.array(images_b[0]))
-
-
-### Output extractor utils for offline / online paths respectively
-def _get_offline_images(outputs: list[OmniRequestOutput]) -> list[Image.Image]:
-    """Extract the images from an Omni .generate() call."""
-    assert len(outputs) == 1
-    return outputs[0].images
-
-
-def _get_online_images(responses: list[DiffusionResponse]) -> list[Image.Image]:
-    """Extract the images from a server response."""
-    assert len(responses) == 1
-    images = responses[0].images
-    assert images is not None
-    return images
-
-
-### Offline helpers
-def _run_offline_t2i(omni: Omni, params: OmniDiffusionSamplingParams = IMAGE_GEN_SAMPLING_PARAMS):
-    return omni.generate({"prompt": PROMPT}, params)
-
-
-def _run_offline_i2i(omni: Omni):
-    return omni.generate(
-        {"prompt": PROMPT, "multi_modal_data": {"image": INPUT_IMAGE}},
-        IMAGE_GEN_SAMPLING_PARAMS,
-    )
-
-
-### Online helpers
-def _build_online_image_data_url() -> str:
-    """Get a valid base 64 encoded data URL corresponding to an image."""
-    buf = io.BytesIO()
-    INPUT_IMAGE.save(buf, format="PNG")
-    b64 = base64.b64encode(buf.getvalue()).decode()
-    return f"data:image/png;base64,{b64}"
-
-
-def _run_online_t2i(
-    server: OmniServer, client: OpenAIClientHandler, extra_body: dict | None = None
-) -> list[DiffusionResponse]:
-    """Run a text to image request through the server."""
-    messages = dummy_messages_from_mix_data(content_text=PROMPT)
-    request_config = {
-        "model": server.model,
-        "messages": messages,
-        "extra_body": extra_body or IMAGE_GEN_EXTRA_BODY,
-    }
-    return client.send_diffusion_request(request_config)
-
-
-def _run_online_i2i(server: OmniServer, client: OpenAIClientHandler) -> list[DiffusionResponse]:
-    """Run an image to image request through the server."""
-    image_data_url = _build_online_image_data_url()
-    messages = dummy_messages_from_mix_data(
-        content_text=PROMPT,
-        image_data_url=image_data_url,
-    )
-    request_config = {
-        "model": server.model,
-        "messages": messages,
-        "extra_body": IMAGE_GEN_EXTRA_BODY,
-    }
-    return client.send_diffusion_request(request_config)
-
-
-### Offline task runners
-def run_and_validate_text_to_image_request(omni: Omni):
-    """Run and validate a text to image request."""
-    _validate_images(_get_offline_images(_run_offline_t2i(omni)))
-
-
-def run_and_validate_image_to_image_request(omni: Omni):
-    """Run and validate an image to image request."""
-    _validate_images(_get_offline_images(_run_offline_i2i(omni)))
-
-
-def run_and_validate_text_to_image_determinism(omni: Omni):
-    """Checks for determinism; for now we just keep this for TTI."""
-    _validate_image_gen_determinism(
-        _get_offline_images(_run_offline_t2i(omni)),
-        _get_offline_images(_run_offline_t2i(omni)),
-    )
-
-
-def run_and_validate_text_to_image_multi_output(omni: Omni):
-    """Checks for multi-output; for now we just keep this for TTI."""
-    params = replace(IMAGE_GEN_SAMPLING_PARAMS, num_outputs_per_prompt=2)
-    _validate_images(_get_offline_images(_run_offline_t2i(omni, params)), expected_n=2)
-
-
-### Online task runners
-def run_and_validate_online_text_to_image_request(server: OmniServer, client: OpenAIClientHandler):
-    """Run and validate a text to image request through the server."""
-    _validate_images(_get_online_images(_run_online_t2i(server, client)))
-
-
-def run_and_validate_online_image_to_image_request(server: OmniServer, client: OpenAIClientHandler):
-    """Run and validate an image to image request through the server."""
-    _validate_images(_get_online_images(_run_online_i2i(server, client)))
-
-
-def run_and_validate_online_text_to_image_determinism(server: OmniServer, client: OpenAIClientHandler):
-    """Checks for determinism through the server; for now we just keep this for TTI."""
-    _validate_image_gen_determinism(
-        _get_online_images(_run_online_t2i(server, client)),
-        _get_online_images(_run_online_t2i(server, client)),
-    )
-
-
-def run_and_validate_online_text_to_image_multi_output(server: OmniServer, client: OpenAIClientHandler):
-    """Checks for multi-output through the server; for now we just keep this for TTI."""
-    extra_body = {**IMAGE_GEN_EXTRA_BODY, "num_outputs_per_prompt": 2}
-    _validate_images(
-        _get_online_images(_run_online_t2i(server, client, extra_body=extra_body)),
-        expected_n=2,
-    )

tests/model_tests/diffusion/test_alignment.py

@@ -1,99 +0,0 @@
-"""
-Tests for alignment between _DIFFUSION_MODELS and DIFFUSION_TEST_SETTINGS; if
-tests in this file are failing, you are probably adding a new model, and need
-to account for it in DIFFUSION_TEST_SETTINGS
-"""
-
-import pytest
-
-from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
-from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
-
-# These diffusion models currently do not have any added (tiny) model tests.
-# In general, pipelines should only be in the list below if they have not been
-# migrated yet or there is a compelling reason that they should not be (which for
-# now may just be CI pressure etc due to the volume of models and the cost of
-# initializing the server).
-#
-# The tests here validate that all keys in the diffusion registry are accounted for;
-# if you are adding a new model and see the tests below fail, please follow the pattern
-# for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS.
-EXCLUDED_MODELS = [
-    "QwenImagePipeline",
-    "QwenImageEditPipeline",
-    "QwenImageEditPlusPipeline",
-    "QwenImageLayeredPipeline",
-    "GlmImagePipeline",
-    "ZImagePipeline",
-    "OvisImagePipeline",
-    "WanPipeline",
-    "WanVACEPipeline",
-    "LTX2Pipeline",
-    "LTX2ImageToVideoPipeline",
-    "LTX2TwoStagesPipeline",
-    "LTX2ImageToVideoTwoStagesPipeline",
-    "LTX2T2VDMD2Pipeline",
-    "LTX2I2VDMD2Pipeline",
-    "LTX23Pipeline",
-    "LTX23ImageToVideoPipeline",
-    "StableAudioPipeline",
-    "WanImageToVideoPipeline",
-    "WanS2VPipeline",
-    "WanT2VDMD2Pipeline",
-    "WanI2VDMD2Pipeline",
-    "LongCatImagePipeline",
-    "BagelPipeline",
-    "LancePipeline",
-    "MingImagePipeline",
-    "InternVLAA1Pipeline",
-    "LongCatImageEditPipeline",
-    "StableDiffusion3Pipeline",
-    "FluxKontextPipeline",
-    "HunyuanImage3ForCausalMM",
-    "ErnieImagePipeline",
-    "NextStep11Pipeline",
-    "FluxPipeline",
-    "FluxDMD2Pipeline",
-    "QwenImageDMD2Pipeline",
-    "OmniGen2Pipeline",
-    "HeliosPipeline",
-    "HeliosPyramidPipeline",
-    "Flux2Pipeline",
-    "DreamIDOmniPipeline",
-    "SenseNovaU1Pipeline",
-    "AudioXPipeline",
-    "HunyuanVideo15Pipeline",
-    "HunyuanVideo15ImageToVideoPipeline",
-    "LingBotVideoPipeline",
-    "MagiHumanPipeline",
-    "OmniVoicePipeline",
-    "OmniVoice",
-    "Cosmos3OmniDiffusersPipeline",
-    "DiffusersAdapterPipeline",
-    "HiDreamImagePipeline",
-    "DreamZeroPipeline",
-    "StableDiffusionXLPipeline",
-    "Gr00tN1d7Pipeline",
-    "SoulXSingerPipeline",
-    "SoulXSingerSVCPipeline",
-]
-
-
-def test_all_non_excluded_pipelines_have_tests():
-    """Ensure that every pipeline in the diffusion registry either has a test
-    configuration in DIFFUSION_TEST_SETTINGS, or is explicitly excluded for a
-    known reason."""
-    non_excluded_diff_pipes = _DIFFUSION_MODELS.keys() - set(EXCLUDED_MODELS)
-    missing_tests = non_excluded_diff_pipes - DIFFUSION_TEST_SETTINGS.keys()
-    assert len(missing_tests) == 0, f"Pipelines missing test settings: {missing_tests}"
-
-
-def test_no_excluded_models_have_test_settings():
-    """Ensure that no models in the exclude list have a test configuration in
-    DIFFUSION_TEST_SETTINGS (i.e., if they have test settings, we should not
-    exclude them)."""
-    non_excluded_diff_pipes = _DIFFUSION_MODELS.keys() - set(EXCLUDED_MODELS)
-    missing_pipes = DIFFUSION_TEST_SETTINGS.keys() - set(non_excluded_diff_pipes)
-    assert len(missing_pipes) == 0

tests/model_tests/diffusion/test_common_offline.py

@@ -1,69 +0,0 @@
-import pytest
-
-from tests.model_tests.diffusion.case_filtering import get_parametrized_options
-from tests.model_tests.diffusion.config_types import (
-    DiffusionAccs,
-    DiffusionTasks,
-    build_omni_from_diff_accelerations,
-)
-from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
-from tests.model_tests.diffusion.task_runners import (
-    run_and_validate_image_to_image_request,
-    run_and_validate_text_to_image_determinism,
-    run_and_validate_text_to_image_multi_output,
-    run_and_validate_text_to_image_request,
-)
-
-# NOTE : Hardware marks are added dynamically based on test requirements
-pytestmark = [pytest.mark.diffusion]
-
-
-@pytest.mark.parametrize(
-    "model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
-    get_parametrized_options(DIFFUSION_TEST_SETTINGS),
-)
-def test_pipeline_on_supported_tasks(
-    model_name,
-    accelerations: list[DiffusionAccs] | None,
-    supported_tasks: list[DiffusionTasks],
-    check_multioutput: bool,
-    check_determinism: bool,
-    tiny_model_paths: dict[str, str],
-    subtests,
-):
-    """Run a smoke test on all of the pipelines supported tasks using a set of enabled accelerations."""
-    assert len(supported_tasks) > 0
-    # We initialize the Omni object before running the tasks, then run each task as a pytest subtask.
-    # This lets us init the model once, but display separate failures in pytest, and avoid halting the
-    # checks on other tasks if one fails.
-    #
-    # This allows us to have some degree of test isolation without the cost of redundant initialization,
-    # since starting the server can take 10+ seconds, even for tiny models.
-    #
-    # NOTE: Be sure to install pytest-subtests if you're running on pytest < 9
-    omni = build_omni_from_diff_accelerations(
-        accelerations=accelerations,
-        model=tiny_model_paths[model_name],
-        enforce_eager=True,
-    )
-    try:
-        for task_type in supported_tasks:
-            with subtests.test(msg=task_type):
-                if task_type == DiffusionTasks.TEXT_TO_IMAGE:
-                    run_and_validate_text_to_image_request(omni)
-                elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
-                    run_and_validate_image_to_image_request(omni)
-                else:
-                    raise ValueError(f"Task type {task_type} is not yet supported")
-
-        # NOTE: For now, we only check determinism + multi output for the base case,
-        # since checking it on every extra acceleration configuration is redundant
-        # (see case_filtering).
-        if check_determinism:
-            with subtests.test(msg="determinism"):
-                run_and_validate_text_to_image_determinism(omni)
-        if check_multioutput:
-            with subtests.test(msg="multi_output"):
-                run_and_validate_text_to_image_multi_output(omni)
-    finally:
-        omni.close()

tests/model_tests/diffusion/test_common_online.py

@@ -1,75 +0,0 @@
-"""
-Analogous to test_common_offline, but for server tests. Validates the full
-online serving stack (CLI arg parsing, subprocess, API routing, response
-encoding) using tiny models.
-"""
-
-import pytest
-
-from tests.helpers.runtime import OmniServer, OpenAIClientHandler
-from tests.model_tests.diffusion.case_filtering import get_parametrized_options
-from tests.model_tests.diffusion.config_types import (
-    DiffusionAccs,
-    DiffusionTasks,
-    build_server_args_from_diff_accelerations,
-)
-from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
-from tests.model_tests.diffusion.task_runners import (
-    run_and_validate_online_image_to_image_request,
-    run_and_validate_online_text_to_image_determinism,
-    run_and_validate_online_text_to_image_multi_output,
-    run_and_validate_online_text_to_image_request,
-)
-
-# NOTE : Hardware marks are added dynamically based on test requirements
-pytestmark = [pytest.mark.diffusion]
-
-
-@pytest.mark.parametrize(
-    "model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
-    get_parametrized_options(DIFFUSION_TEST_SETTINGS, online=True),
-)
-def test_online_on_supported_tasks(
-    model_name: str,
-    accelerations: list[DiffusionAccs] | None,
-    supported_tasks: list[DiffusionTasks],
-    check_multioutput: bool,
-    check_determinism: bool,
-    tiny_model_paths: dict[str, str],
-    run_level: str,
-    subtests,
-):
-    """Smoke test: start a tiny model server and run each supported task via the API."""
-    model_path = tiny_model_paths[model_name]
-    server_args = build_server_args_from_diff_accelerations(accelerations)
-    server_args.append("--enforce-eager")
-
-    with OmniServer(model_path, server_args) as server:
-        # TODO: We may want to revisit run_level validation here,
-        # because checks for things like image size etc should not
-        # depend on whether or not the weights are real or random
-        client = OpenAIClientHandler(
-            host=server.host,
-            port=server.port,
-            api_key="EMPTY",
-            run_level=run_level,
-            log_stats=server.log_stats,
-        )
-        for task_type in supported_tasks:
-            with subtests.test(msg=task_type):
-                if task_type == DiffusionTasks.TEXT_TO_IMAGE:
-                    run_and_validate_online_text_to_image_request(server, client)
-                elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
-                    run_and_validate_online_image_to_image_request(server, client)
-                else:
-                    raise ValueError(f"Task type {task_type} is not yet supported")
-
-        # NOTE: For now, we only check determinism + multi output for the base case,
-        # since checking it on every extra acceleration configuration is redundant
-        # (see case_filtering).
-        if check_determinism:
-            with subtests.test(msg="determinism"):
-                run_and_validate_online_text_to_image_determinism(server, client)
-        if check_multioutput:
-            with subtests.test(msg="multi_output"):
-                run_and_validate_online_text_to_image_multi_output(server, client)

tests/model_tests/diffusion/utils.py

@@ -1,33 +0,0 @@
-"""
-Utilities for resolving real models to their tiny model equivalents.
-"""
-
-import logging
-
-from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
-from vllm_omni.diffusion.data import resolve_model_class_name
-
-logger = logging.getLogger(__name__)
-
-
-def resolve_tiny_model_path(model: str) -> str:
-    """Given a real model name/path, resolve it to a tiny model path.
-
-    Raises ValueError if the pipeline class cannot be determined (invalid
-    model). Returns the original model path if no tiny builder exists yet."""
-    pipeline_class = resolve_model_class_name(model)
-    if pipeline_class is None:
-        raise ValueError(
-            f"Cannot resolve pipeline class for model: {model}. The model path may be invalid or its config unreadable."
-        )
-
-    test_opts = DIFFUSION_TEST_SETTINGS.get(pipeline_class)
-    if test_opts is None:
-        logger.warning(
-            "No tiny model builder for pipeline %s (model: %s). Using original model.",
-            pipeline_class,
-            model,
-        )
-        return model
-
-    return test_opts.builder()

tests/platforms/test_cuda_ir_op_priority.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+from vllm.config import CompilationConfig, CompilationMode, DeviceConfig, VllmConfig
+
+from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _vllm_config(*, backend: str = "inductor", mode: CompilationMode) -> VllmConfig:
+    return VllmConfig(
+        device_config=DeviceConfig(device="cpu"),
+        compilation_config=CompilationConfig(backend=backend, mode=mode),
+    )
+
+
+@pytest.mark.parametrize(
+    "mode",
+    [CompilationMode.NONE, CompilationMode.VLLM_COMPILE, CompilationMode.STOCK_TORCH_COMPILE],
+)
+def test_cuda_default_ir_op_priority_prefers_vllm_c_when_inductor_backend(mode: CompilationMode) -> None:
+    """Regression for #4964: inductor-active configs must not switch to native-only."""
+    priority = CudaOmniPlatform.get_default_ir_op_priority(_vllm_config(mode=mode))
+    assert priority.rms_norm == ["vllm_c", "native"]
+    assert priority.fused_add_rms_norm == ["vllm_c", "native"]
+
+
+def test_cuda_default_ir_op_priority_with_oink(monkeypatch: pytest.MonkeyPatch) -> None:
+    import vllm.envs as envs
+
+    monkeypatch.setattr(envs, "VLLM_USE_OINK_OPS", True)
+    priority = CudaOmniPlatform.get_default_ir_op_priority(
+        _vllm_config(mode=CompilationMode.VLLM_COMPILE),
+    )
+    assert priority.rms_norm == ["oink", "vllm_c", "native"]
+    assert priority.fused_add_rms_norm == ["oink", "vllm_c", "native"]

tests/test_config_factory.py

@@ -16,7 +16,6 @@ from transformers import PretrainedConfig, Qwen3OmniMoeConfig
 
 from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage
 from vllm_omni.config.config_factory import StageConfigFactory
-from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline
 from vllm_omni.config.stage_config import (
     DeployConfig,
@@ -40,15 +39,6 @@ from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_black
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-@pytest.fixture(autouse=True)
-def clear_config_factory_caches():
-    """Clear cached classmethods from the StageConfigFactory to prevent test pollution."""
-    yield
-    StageConfigFactory.get_hf_config.cache_clear()
-    StageConfigFactory.try_infer_model_type.cache_clear()
-    StageConfigFactory.get_pipeline_config.cache_clear()
-
-
 Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True)
 Q3_OMNI_THINKER_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=False)
 
@@ -676,43 +666,6 @@ class TestPipelineRegistration:
         assert resolved_config is not None
         assert len(resolved_config) > 0
 
-    def test_deploy_override_uses_correct_endpoint_restrictions(self, clean_pipeline_registry, tmp_path):
-        """Ensure endpoint restrictions must come from the final pipeline
-        after deploy config overrides, not the auto-detected pipeline.
-        """
-        # Register two pipeline configs, where one has an endpoint restriction, and one doesn't
-        restriction = EndpointRestriction(
-            OmniServingCapability.COMPLETIONS,
-            "pipeline_a blocks completions",
-        )
-        pipe_a = PipelineConfig(model_type="detect_type", endpoint_restrictions=(restriction,))
-        pipe_b = PipelineConfig(model_type="override_type", endpoint_restrictions=())
-        register_pipeline(pipe_a)
-        register_pipeline(pipe_b)
-
-        # Create a config with the autodetected type, and write the
-        # deploy config specifying the override type to a temp path
-        class FakeConfig(PretrainedConfig):
-            model_type = "detect_type"
-
-        deploy_yaml = tmp_path / "override.yaml"
-        deploy_yaml.write_text("pipeline: override_type\n")
-
-        # Get the endpoint restrictions, passing the deploy config with the override
-        # type + patching the config for the detected type. Ensure that the endpoint
-        # restrictions correspond to the type in the deploy config.
-        with patch(
-            "vllm_omni.config.config_factory.get_config",
-            return_value=FakeConfig(),
-        ):
-            restrictions = StageConfigFactory.get_pipeline_endpoint_restrictions(
-                model="fake/model",
-                trust_remote_code=False,
-                deploy_config_path=str(deploy_yaml),
-            )
-
-        assert restrictions == ()
-
 
 class TestResolveScheduler:
     def test_all_execution_types_handled(self):
@@ -796,32 +749,6 @@ stages:
         runtime_config = deploy.stages[0].engine_extras["hf_overrides"]["voxcpm2_runtime_config"]
         assert runtime_config == expected_runtime_config
 
-    @pytest.mark.parametrize(
-        ("deploy_name", "pipeline_name", "stage_count", "final_output_type"),
-        [
-            ("mammoth_moda2.yaml", "mammoth_moda2", 2, "image"),
-            ("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text"),
-            ("omnivoice.yaml", "omnivoice", 1, "audio"),
-        ],
-    )
-    def test_load_new_registry_backed_deploy_configs(
-        self,
-        deploy_name: str,
-        pipeline_name: str,
-        stage_count: int,
-        final_output_type: str,
-    ):
-        deploy_path = Path(get_deploy_config_path(deploy_name))
-        deploy = load_deploy_config(deploy_path)
-        assert deploy.pipeline == pipeline_name
-
-        with patch("vllm_omni.platforms.current_omni_platform") as platform:
-            platform.device_name = "cuda"
-            stages = merge_pipeline_deploy(OMNI_PIPELINES[pipeline_name], deploy)
-        assert len(stages) == stage_count
-        assert stages[-1].final_output is True
-        assert stages[-1].final_output_type == final_output_type
-
     def test_merge_pipeline_deploy(self):
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():

tests/worker/test_min_tokens_stop_ids_guard.py

@@ -1,115 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Guard against out-of-vocabulary stop ids in min-tokens masking (#4962).
-
-vLLM's input processor folds the stage tokenizer's EOS id into
-``SamplingParams.all_stop_token_ids``. For AR stages whose lm_head is
-narrower than the tokenizer vocabulary (codec talkers such as Qwen3-TTS:
-3072 logits vs text EOS 151645), any ``min_tokens >= 1`` makes
-``MinTokensLogitsProcessor`` write ``-inf`` at an out-of-range index and
-the engine dies with a CUDA device-side assert.
-"""
-
-import pytest
-import torch
-from vllm import SamplingParams
-from vllm.v1.sample.logits_processor import (
-    BatchUpdate,
-    LogitsProcessors,
-    MinTokensLogitsProcessor,
-)
-
-from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-TALKER_VOCAB = 3072
-CODEC_EOS = 2150
-TEXT_EOS = 151645
-
-
-def _make_min_tokens_proc(stop_token_ids: list[int], eos_token_id: int | None) -> MinTokensLogitsProcessor:
-    """Build a MinTokensLogitsProcessor with one active request, mirroring
-    the production path (engine folds the tokenizer EOS via
-    ``update_from_generation_config``)."""
-    params = SamplingParams(min_tokens=2, stop_token_ids=stop_token_ids)
-    if eos_token_id is not None:
-        params.update_from_generation_config({}, eos_token_id)
-    proc = MinTokensLogitsProcessor(None, device=torch.device("cpu"), is_pin_memory=False)
-    proc.update_state(
-        BatchUpdate(
-            batch_size=1,
-            removed=[],
-            added=[(0, params, None, [])],
-            moved=[],
-        )
-    )
-    return proc
-
-
-def test_oob_stop_id_crashes_without_guard():
-    """Document the failure mode: the folded text EOS is out of range for a
-    narrow codec head and index_put_ raises (device-side assert on CUDA)."""
-    proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS)
-    logits = torch.zeros(1, TALKER_VOCAB)
-    with pytest.raises((IndexError, RuntimeError)):
-        proc.apply(logits)
-
-
-def test_guard_filters_oob_and_keeps_in_range_mask():
-    proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS)
-    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
-
-    logits = torch.zeros(1, TALKER_VOCAB)
-    out = proc.apply(logits)
-
-    assert out[0, CODEC_EOS] == -float("inf")
-    masked = (out == -float("inf")).nonzero()
-    assert masked.tolist() == [[0, CODEC_EOS]]
-
-
-def test_guard_noop_when_all_stop_ids_in_range():
-    """Full-width heads (e.g. CosyVoice3 talker) must be untouched: no
-    tensor rebuild, mask unchanged."""
-    proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=None)
-    slice_before = proc.logits_slice
-
-    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
-
-    assert proc.logits_slice is slice_before
-    out = proc.apply(torch.zeros(1, TALKER_VOCAB))
-    assert out[0, CODEC_EOS] == -float("inf")
-
-
-def test_guard_sanitizes_params_only_once():
-    """The stop-id set is shared with the request's SamplingParams, so the
-    second call finds nothing to drop (no repeated rebuilds/warnings)."""
-    proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS)
-    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
-    slice_after_first = proc.logits_slice
-
-    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
-
-    assert proc.logits_slice is slice_after_first
-
-
-def test_guard_handles_multiple_requests():
-    params_oob = SamplingParams(min_tokens=4, stop_token_ids=[CODEC_EOS])
-    params_oob.update_from_generation_config({}, TEXT_EOS)
-    params_ok = SamplingParams(min_tokens=4, stop_token_ids=[7])
-    proc = MinTokensLogitsProcessor(None, device=torch.device("cpu"), is_pin_memory=False)
-    proc.update_state(
-        BatchUpdate(
-            batch_size=2,
-            removed=[],
-            added=[(0, params_oob, None, []), (1, params_ok, None, [])],
-            moved=[],
-        )
-    )
-
-    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
-    out = proc.apply(torch.zeros(2, TALKER_VOCAB))
-
-    assert out[0, CODEC_EOS] == -float("inf")
-    assert out[1, 7] == -float("inf")
-    assert int((out == -float("inf")).sum()) == 2

vllm_omni/config/config_factory.py

@@ -5,7 +5,6 @@
 from __future__ import annotations
 
 import dataclasses
-import functools
 from dataclasses import asdict
 from pathlib import Path
 from typing import Any
@@ -15,7 +14,6 @@ from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_config
 from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
 
-from vllm_omni.config.endpoint_policy import EndpointRestriction
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 from vllm_omni.config.stage_config import (
     _DEPLOY_DIR,
@@ -48,174 +46,83 @@ class StageConfigFactory:
     """
 
     @classmethod
-    def get_pipeline_endpoint_restrictions(
+    def create_from_model(
         cls,
         model: str,
-        trust_remote_code: bool,
-        deploy_config_path: str | None,
-    ) -> tuple[EndpointRestriction, ...]:
-        """Given a model string, determine the corresponding endpoint restrictions.
+        cli_overrides: dict[str, Any] | None = None,
+        deploy_config_path: str | None = None,
+        **deprecated_kwargs: Any,
+    ) -> list[StageConfig] | None:
+        """Load pipeline + deploy config, merge with CLI overrides.
 
-        Args:
-            model: Model name or path.
-            trust_remote_code: Whether to trust remote code for HF config loading.
-            deploy_config_path: Optional path to the deploy config for the pipeline.
-
-        Returns:
-            A tuple of model specific endpoint restrictions.
+        Checks OMNI_PIPELINES first, since supported models should be explicitly
+        registered. If a model is not registered in OMNI_PIPELINES, tries to fall
+        back to using the Transformers config & finding pipelines that have overlapping
+        supported architectures.
         """
-        pipeline_cfg = StageConfigFactory.get_pipeline_config(
-            model=model,
-            trust_remote_code=trust_remote_code,
-            deploy_config_path=deploy_config_path,
-        )
-        return pipeline_cfg.endpoint_restrictions if pipeline_cfg else ()
+        if cli_overrides is None:
+            cli_overrides = {}
 
-    @classmethod
-    @functools.cache
-    def get_hf_config(cls, model: str, trust_remote_code: bool) -> PretrainedConfig | None:
-        """Fetch the HF config (if it exists) from the model directory.
+        trust_remote_code = cli_overrides.get("trust_remote_code", True)
+        if trust_remote_code is None:
+            trust_remote_code = False
 
-        Args:
-            model: Model name or path.
-            trust_remote_code: Whether to trust remote code for HF config loading.
-
-        Returns:
-            the model's config or None.
-        """
-        hf_config = None
-        try:
-            return get_config(model, trust_remote_code=trust_remote_code)
-        except Exception as e:
-            logger.debug(f"`get_config` failed with exception {e}; inferred HF config is None")
-        return hf_config
-
-    @classmethod
-    @functools.cache
-    def try_infer_model_type(cls, model: str, trust_remote_code: bool) -> str | None:
-        """Auto-detect model_type from model directory and apply any model
-        specific patches to get the correct model_type str. If we are unable
-        to infer it from the model directory, we fall back to the PipelineConfig.
-
-        Args:
-            model: Model name or path.
-            trust_remote_code: Whether to trust remote code for HF config loading.
-
-        Returns:
-            model_type as a string; may be None on failure.
-        """
-        model_type = cls._try_infer_model_type(
-            model=model,
-            trust_remote_code=trust_remote_code,
-        )
+        model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code)
         if model_type == "vla":
             if _looks_like_dreamzero(model):
                 model_type = "dreamzero"
-        return model_type
 
-    @classmethod
-    def _try_infer_model_type(cls, model: str, trust_remote_code: bool) -> str | None:
-        """Auto-detect model_type from model directory.
+        # --- 1. Explicit deploy-config pipeline override (highest precedence) ---
+        # A deploy YAML may set ``pipeline: <model_type>`` to force routing for
+        # models whose HF config has a generic/colliding ``model_type`` or no
+        # matching architectures. Ming-omni-tts, for example, reports
+        # model_type="dense" with arch BailingMMNative…, which matches no
+        # pipeline — without this it falls through to the diffusion default and
+        # dies with "Model class BailingMMNativeForConditionalGeneration not
+        # found in diffusion model registry". Honor the key before auto-detection.
+        explicit_pipeline = None
+        if deploy_config_path:
+            try:
+                explicit_pipeline = load_deploy_config(deploy_config_path).pipeline
+            except Exception:
+                logger.exception("Failed to read 'pipeline' key from deploy config %s", deploy_config_path)
+        if explicit_pipeline:
+            pipeline_cfg = cls.resolve_pipeline_config(explicit_pipeline, hf_config)
+            if pipeline_cfg is not None:
+                return cls._create_from_registry(
+                    explicit_pipeline,
+                    pipeline_cfg,
+                    cli_overrides,
+                    deploy_config_path,
+                )
+            logger.warning(
+                "Deploy config %s requested pipeline %r which is not in OMNI_PIPELINES; "
+                "falling back to auto-detection.",
+                deploy_config_path,
+                explicit_pipeline,
+            )
 
-        Args:
-            model: Model name or path.
-            trust_remote_code: Whether to trust remote code for HF config loading.
-
-        Returns:
-            model_type as a string; may be None on failure.
-        """
-        hf_config = cls.get_hf_config(
-            model=model,
-            trust_remote_code=trust_remote_code,
-        )
-        if hf_config is not None:
-            return hf_config.model_type
-
-        # Fallback: read config.json directly for custom model types that
-        # are not registered with transformers (e.g. qwen3_tts).
-        try:
-            config_dict = get_hf_file_to_dict("config.json", model, revision=None)
-            if config_dict:
-                if "model_type" in config_dict:
-                    return config_dict["model_type"]
-                # VoxCPM2-style configs use singular ``architecture`` rather
-                # than HF's standard ``model_type`` / ``architectures``. Accept
-                # it as a fallback so the pipeline registry can still match.
-                if "architecture" in config_dict and isinstance(config_dict["architecture"], str):
-                    return config_dict["architecture"]
-        except Exception as e:
-            logger.debug(f"Failed to auto-detect model type for {model}: {e}")
-
-        # Fallback for diffusers-style models: check model_index.json.
-        # Some models (e.g. GLM-Image) have no root config.json but ship a
-        # model_index.json with _class_name that maps to a pipeline key via
-        # PipelineConfig.diffusers_class_name.
-        try:
-            model_index = get_hf_file_to_dict("model_index.json", model, revision=None)
-            if model_index and "_class_name" in model_index:
-                class_name = model_index["_class_name"]
-                for obj in OMNI_PIPELINES.values():
-                    # If we have a resolver, call it with the optional hf_config
-                    # to get the default pipeline config for this key
-                    pipeline_cfg = obj(hf_config) if callable(obj) else obj
-                    if pipeline_cfg is not None and pipeline_cfg.diffusers_class_name == class_name:
-                        logger.info(
-                            "Detected pipeline %r from model_index.json (_class_name=%r)",
-                            pipeline_cfg.model_type,
-                            class_name,
-                        )
-                        return pipeline_cfg.model_type
-        except Exception as e:
-            logger.debug(f"Failed to detect model type for diffusers-style models: {e}")
-
-        # Final fallback: some models (e.g. CosyVoice3) ship an empty
-        # config.json and rely on naming conventions. Match the model path
-        # basename against registered pipeline keys — longest match wins
-        # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9).
-        model_lower = model.lower().replace("-", "").replace("_", "")
-        best: str | None = None
-        best_len = 0
-        for registered_key in OMNI_PIPELINES.keys():
-            candidate = registered_key.lower().replace("-", "").replace("_", "")
-            if candidate and candidate in model_lower and len(candidate) > best_len:
-                best = registered_key
-                best_len = len(candidate)
-        if best is not None:
-            return best
-
-        return None
-
-    @classmethod
-    @functools.cache
-    def get_pipeline_config(
-        cls,
-        model: str,
-        trust_remote_code: bool,
-        deploy_config_path: str | None = None,
-    ) -> PipelineConfig | None:
-        """Resolve the PipelineConfig for a model path/name."""
-        model_type = cls.try_infer_model_type(model=model, trust_remote_code=trust_remote_code)
-        hf_config = cls.get_hf_config(model=model, trust_remote_code=trust_remote_code)
-
-        # Resolve the deploy config & check if the user set the pipeline;
-        # If the pipeline is explicitly set, it takes highest priority
-        deploy_config_pipe = cls._get_deploy_override_pipe_config(hf_config, deploy_config_path)
-        if deploy_config_pipe is not None:
-            return deploy_config_pipe
-
-        # Pipeline isn't set in the yaml spec, so we need infer it ourselves.
+        # --- 2. Auto-detected model_type registered in OMNI_PIPELINES ---
         if model_type and model_type in OMNI_PIPELINES:
             pipeline_cfg = cls.resolve_pipeline_config(model_type, hf_config)
             if pipeline_cfg is not None:
-                return pipeline_cfg
+                return cls._create_from_registry(
+                    model_type,
+                    pipeline_cfg,
+                    cli_overrides,
+                    deploy_config_path,
+                )
 
+        # --- HF architecture fallback: some models report a generic
+        # model_type that collides with another model. Match by the
+        # hf_architectures declared on each registered PipelineConfig.
         if hf_config is not None:
-            if model_type is not None:
-                logger.warning("Inferred model type %s is not registered to an Omni pipeline", model_type)
+            logger.warning("Inferred model type %s is not registered to an Omni pipeline", model_type)
             hf_archs = set(getattr(hf_config, "architectures", []) or [])
             if hf_archs:
                 for registered in OMNI_PIPELINES.values():
                     pipeline_cfg = registered if isinstance(registered, PipelineConfig) else registered(hf_config)
+                    # Resolvers that get configs of the incorrect type should return None
                     if pipeline_cfg is None:
                         continue
                     predicate = pipeline_cfg.hf_config_predicate
@@ -236,61 +143,36 @@ class StageConfigFactory:
                                 pipeline_cfg.model_type,
                             )
                             continue
+
                     if isinstance(pipeline_cfg, PipelineConfig) and hf_archs.intersection(
                         pipeline_cfg.hf_architectures
                     ):
-                        return pipeline_cfg
-        return None
+                        return cls._create_from_registry(
+                            pipeline_cfg.model_type,
+                            pipeline_cfg,
+                            cli_overrides,
+                            deploy_config_path,
+                        )
 
-    @classmethod
-    def _get_deploy_override_pipe_config(
-        cls,
-        hf_config: PretrainedConfig | None,
-        deploy_config_path: str | None,
-    ) -> PipelineConfig | None:
-        """Load the deploy config and extract + resolve its pipeline field."""
-        if deploy_config_path is None:
-            return None
+        # --- Explicit deploy-config pipeline ---
+        # When auto-detection above resolves nothing (generic/missing HF model_type
+        # and no matching architecture), honor an explicit pipeline key in the deploy config
+        if deploy_config_path is not None:
+            deploy_path = Path(deploy_config_path)
+            if deploy_path.exists():
+                deploy_cfg = load_deploy_config(deploy_path)
+                if deploy_cfg.pipeline:
+                    pipeline_cfg = cls.resolve_pipeline_config(deploy_cfg.pipeline, hf_config)
+                    if pipeline_cfg is not None:
+                        return cls._create_from_registry(
+                            pipeline_cfg.model_type,
+                            pipeline_cfg,
+                            cli_overrides,
+                            deploy_config_path,
+                        )
 
-        deploy_path = Path(deploy_config_path)
-        if deploy_path.exists():
-            deploy_cfg = load_deploy_config(deploy_path)
-            if deploy_cfg.pipeline is not None:
-                return cls.resolve_pipeline_config(deploy_cfg.pipeline, hf_config)
-        return None
-
-    @classmethod
-    def create_from_model(
-        cls,
-        model: str,
-        *,
-        trust_remote_code: bool = False,
-        cli_overrides: dict[str, Any] | None = None,
-        deploy_config_path: str | None = None,
-        **deprecated_kwargs: Any,
-    ) -> list[StageConfig] | None:
-        """Load pipeline + deploy config, merge with CLI overrides.
-
-        Checks OMNI_PIPELINES first, since supported models should be explicitly
-        registered. If a model is not registered in OMNI_PIPELINES, tries to fall
-        back to using the Transformers config & finding pipelines that have overlapping
-        supported architectures.
-        """
-        if cli_overrides is None:
-            cli_overrides = {}
-
-        pipeline_cfg = cls.get_pipeline_config(
-            model=model,
-            trust_remote_code=trust_remote_code,
-            deploy_config_path=deploy_config_path,
-        )
-        if pipeline_cfg is not None:
-            return cls._create_from_registry(
-                pipeline_cfg.model_type,
-                pipeline_cfg,
-                cli_overrides,
-                deploy_config_path,
-            )
+        # Not in the pipeline registry — let the caller fall back to the
+        # legacy ``stage_configs/*.yaml`` path (resolve_model_config_path).
         return None
 
     @classmethod
@@ -406,6 +288,79 @@ class StageConfigFactory:
 
         return [config_dict]
 
+    @classmethod
+    def _auto_detect_model_type(cls, model: str, trust_remote_code: bool = True) -> tuple[str | None, Any]:
+        """Auto-detect model_type from model directory.
+
+        Args:
+            model: Model name or path.
+            trust_remote_code: Whether to trust remote code for HF config loading.
+
+        Returns:
+            Tuple of (model_type, hf_config). Both may be None on failure.
+        """
+        hf_config = None
+
+        try:
+            hf_config = get_config(model, trust_remote_code=trust_remote_code)
+            return hf_config.model_type, hf_config
+        except Exception as e:
+            logger.debug(f"`get_config` failed for {e}; Falling back to raw config.json path")
+
+        # Fallback: read config.json directly for custom model types that
+        # are not registered with transformers (e.g. qwen3_tts).
+        try:
+            config_dict = get_hf_file_to_dict("config.json", model, revision=None)
+            if config_dict:
+                if "model_type" in config_dict:
+                    return config_dict["model_type"], None
+                # VoxCPM2-style configs use singular ``architecture`` rather
+                # than HF's standard ``model_type`` / ``architectures``. Accept
+                # it as a fallback so the pipeline registry can still match.
+                if "architecture" in config_dict and isinstance(config_dict["architecture"], str):
+                    return config_dict["architecture"], None
+        except Exception as e:
+            logger.debug(f"Failed to auto-detect model type for {model}: {e}")
+
+        # Fallback for diffusers-style models: check model_index.json.
+        # Some models (e.g. GLM-Image) have no root config.json but ship a
+        # model_index.json with _class_name that maps to a pipeline key via
+        # PipelineConfig.diffusers_class_name.
+        try:
+            model_index = get_hf_file_to_dict("model_index.json", model, revision=None)
+            if model_index and "_class_name" in model_index:
+                class_name = model_index["_class_name"]
+                for obj in OMNI_PIPELINES.values():
+                    # If we have a resolver, call it with the optional hf_config
+                    # to get the default pipeline config for this key
+                    pipeline_cfg = obj(hf_config) if callable(obj) else obj
+                    if pipeline_cfg is not None and pipeline_cfg.diffusers_class_name == class_name:
+                        logger.info(
+                            "Detected pipeline %r from model_index.json (_class_name=%r)",
+                            pipeline_cfg.model_type,
+                            class_name,
+                        )
+                        return pipeline_cfg.model_type, None
+        except Exception as e:
+            logger.debug(f"Failed to detect model type for diffusers-style models: {e}")
+
+        # Final fallback: some models (e.g. CosyVoice3) ship an empty
+        # config.json and rely on naming conventions. Match the model path
+        # basename against registered pipeline keys — longest match wins
+        # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9).
+        model_lower = model.lower().replace("-", "").replace("_", "")
+        best: str | None = None
+        best_len = 0
+        for registered_key in OMNI_PIPELINES.keys():
+            candidate = registered_key.lower().replace("-", "").replace("_", "")
+            if candidate and candidate in model_lower and len(candidate) > best_len:
+                best = registered_key
+                best_len = len(candidate)
+        if best is not None:
+            return best, None
+
+        return None, None
+
     @classmethod
     def _merge_cli_overrides(
         cls,
@@ -426,7 +381,6 @@ class StageConfigFactory:
         """Given a model type, resolve to the pipeline to be used. If the pipeline
         maps to a callable we resolve based on the HF config."""
         if model_type not in OMNI_PIPELINES:
-            logger.warning("Model type %s is not registered to OMNI_PIPELINES", model_type)
             return None
         obj = OMNI_PIPELINES[model_type]
         return obj(hf_config) if callable(obj) else obj

vllm_omni/config/endpoint_policy.py

@@ -1,91 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Endpoint restriction policy for omni pipelines."""
-
-from dataclasses import dataclass
-from enum import Enum
-from typing import NamedTuple
-
-from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse
-from vllm.entrypoints.serve.utils.error_response import create_error_response
-
-
-class RouteTarget(NamedTuple):
-    """A server path & supported methods."""
-
-    path: str
-    methods: frozenset[str]
-
-
-class OmniServingCapability(Enum):
-    """Serving capabilities that pipelines can shut down."""
-
-    CHAT_COMPLETIONS_BATCH = RouteTarget("/v1/chat/completions/batch", frozenset({"POST"}))
-    COMPLETIONS = RouteTarget("/v1/completions", frozenset({"POST"}))
-
-    @property
-    def path(self) -> str:
-        return self.value.path
-
-    @property
-    def methods(self) -> frozenset[str]:
-        return self.value.methods
-
-
-@dataclass(frozen=True)
-class EndpointRestriction:
-    capability: OmniServingCapability
-    reason: str
-
-
-# Routes that are not supported for any model, but are supported in vLLM.
-# This is only temporary to avoid 500s for batched chat completions.
-UNSUPPORTED_ROUTES: tuple[EndpointRestriction, ...] = (
-    EndpointRestriction(
-        OmniServingCapability.CHAT_COMPLETIONS_BATCH,
-        "Batched chat completions are not yet supported by vLLM Omni.",
-    ),
-)
-
-
-def build_rejection_handler(reason: str):
-    """Build a rejection handler for a given endpoint for the provided reason."""
-
-    async def rejection_handler(raw_request: Request):
-        error = create_error_response(message=reason)
-        return JSONResponse(
-            content=error.model_dump(),
-            status_code=error.error.code,
-        )
-
-    return rejection_handler
-
-
-def shutdown_unsupported_routes(
-    app: FastAPI,
-    endpoint_restrictions: tuple[EndpointRestriction, ...],
-):
-    """Given an initialized FastAPI server instance and a set of model specific endpoint
-    restrictions, remove the restricted routes and patch a handler that returns 400.
-    """
-    from vllm_omni.entrypoints.openai.api_server import _remove_route_from_app
-
-    # Generally these should not overlap since there is no point. If they do,
-    # we use the reason message in UNSUPPORTED_ROUTES, for consistent error messages.
-    restricted_endpoints = (*endpoint_restrictions, *UNSUPPORTED_ROUTES)
-
-    for end_restrict in restricted_endpoints:
-        capability = end_restrict.capability
-        # Remove the route from the app
-        _remove_route_from_app(app, capability.path, capability.methods)
-
-        # Patch the bad request error with the model specific
-        # reason for shutting down this endpoint
-        rejection_handler = build_rejection_handler(end_restrict.reason)
-
-        app.add_api_route(
-            capability.path,
-            rejection_handler,
-            methods=list(capability.methods),
-        )

vllm_omni/config/pipeline_registry.py

@@ -56,10 +56,6 @@ from vllm_omni.model_executor.models.hunyuan_image3.pipeline import (
 )
 from vllm_omni.model_executor.models.indextts2.pipeline import INDEXTTS2_PIPELINE
 from vllm_omni.model_executor.models.lance.pipeline import LANCE_PIPELINE
-from vllm_omni.model_executor.models.mammoth_moda2.pipeline import (
-    MAMMOTH_MODA2_AR_PIPELINE,
-    MAMMOTH_MODA2_PIPELINE,
-)
 from vllm_omni.model_executor.models.mimo_audio.pipeline import MIMO_AUDIO_PIPELINE
 from vllm_omni.model_executor.models.ming_flash_omni.pipeline import (
     MING_FLASH_OMNI_IMAGE_PIPELINE,
@@ -78,7 +74,6 @@ from vllm_omni.model_executor.models.moss_tts.pipeline import (
     MOSS_TTS_REALTIME_PIPELINE,
 )
 from vllm_omni.model_executor.models.moss_tts_nano.pipeline import MOSS_TTS_NANO_PIPELINE
-from vllm_omni.model_executor.models.omnivoice.pipeline import OMNIVOICE_PIPELINE
 from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import (
     QWEN2_5_OMNI_PIPELINE,
     QWEN2_5_OMNI_THINKER_ONLY_PIPELINE,
@@ -123,9 +118,6 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "ming_flash_omni_thinker_only": MING_FLASH_OMNI_THINKER_ONLY_PIPELINE,
     "ming_flash_omni_image": MING_FLASH_OMNI_IMAGE_PIPELINE,
     "moss_tts_nano": MOSS_TTS_NANO_PIPELINE,
-    "omnivoice": OMNIVOICE_PIPELINE,
-    "mammoth_moda2": MAMMOTH_MODA2_PIPELINE,
-    "mammoth_moda2_ar": MAMMOTH_MODA2_AR_PIPELINE,
     "moss_tts_delay": MOSS_TTS_PIPELINE,
     "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE,
     "moss_tts_local": MOSS_TTS_LOCAL_PIPELINE,

vllm_omni/config/stage_config.py

@@ -18,7 +18,6 @@ from transformers import PretrainedConfig
 from vllm.logger import init_logger
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 
-from vllm_omni.config.endpoint_policy import EndpointRestriction
 from vllm_omni.config.yaml_util import create_config, load_yaml_config, to_dict
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
@@ -269,7 +268,6 @@ class PipelineConfig:
     # this value to auto-detect the pipeline.  Only needed for diffusers-style
     # multi-component repos (e.g. GLM-Image).  ``None`` = not a diffusers model.
     diffusers_class_name: str | None = None
-    endpoint_restrictions: tuple[EndpointRestriction, ...] = ()
 
     def get_stage(self, stage_id: int) -> StagePipelineConfig | None:
         """Look up a stage by its ID."""

vllm_omni/deploy/mammoth_moda2.yaml

@@ -1,27 +1,42 @@
-# MammothModa2 AR→DiT (text-to-image).
+# MammothModa2 AR→DiT (text-to-image). Canonical copy lives at
+# ``vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml``; keep in sync.
+#
+# Must remain legacy ``stage_args`` layout (no top-level ``stages``): Omni resolves
+# explicit ``--stage-configs-path`` via ``load_stage_configs_from_yaml``.
 
-async_chunk: false
-pipeline: mammoth_moda2
-
-stages:
+stage_args:
   - stage_id: 0
-    devices: "0"
-    max_num_seqs: 100
-    max_model_len: 8192
-    gpu_memory_utilization: 0.5
-    enforce_eager: true
-    trust_remote_code: true
-    enable_prefix_caching: false
+    runtime:
+      devices: "0"
+    engine_args:
+      model_stage: ar
+      max_num_seqs: 100
+      model_arch: MammothModa2ForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      max_model_len: 8192
+      gpu_memory_utilization: 0.5
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      enable_prefix_caching: false
+    final_output: false
 
   - stage_id: 1
-    devices: "0"
-    max_num_seqs: 1
-    gpu_memory_utilization: 0.3
-    enforce_eager: true
-    trust_remote_code: true
-    enable_prefix_caching: false
-    default_sampling_params:
-      extra_args:
-        text_guidance_scale: 9.0
-        cfg_range: [0.0, 1.0]
-        num_inference_steps: 50
+    runtime:
+      devices: "0"
+    engine_args:
+      model_stage: dit
+      max_num_seqs: 1
+      model_arch: MammothModa2ForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      gpu_memory_utilization: 0.3
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: image
+      enable_prefix_caching: false
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit
+    final_output: true
+    final_output_type: image

vllm_omni/deploy/mammoth_moda2_ar.yaml

@@ -1,14 +0,0 @@
-# MammothModa2 AR-only (image/video understanding and summarization).
-
-async_chunk: false
-pipeline: mammoth_moda2_ar
-
-stages:
-  - stage_id: 0
-    devices: "0"
-    max_num_seqs: 16
-    max_model_len: 8192
-    gpu_memory_utilization: 0.5
-    enforce_eager: true
-    trust_remote_code: true
-    enable_prefix_caching: false

vllm_omni/deploy/omnivoice.yaml

@@ -1,13 +1,19 @@
 # OmniVoice TTS deploy: diffusion engine single stage (dit).
 # Text → iterative unmasking → codebook tokens → DAC decode → 24kHz audio.
+# (Migrated from model_executor/stage_configs/omnivoice.yaml for deploy/ layout.)
 
-async_chunk: false
-pipeline: omnivoice
-
-stages:
+stage_args:
   - stage_id: 0
-    devices: "0"
-    enforce_eager: true
-    trust_remote_code: true
-    distributed_executor_backend: "mp"
-    dtype: "float32"
+    stage_type: diffusion
+    is_comprehension: true
+    runtime:
+      devices: 0
+    engine_args:
+      model_stage: dit
+      model_class_name: "OmniVoicePipeline"
+      enforce_eager: true
+      trust_remote_code: true
+      distributed_executor_backend: "mp"
+      dtype: "float32"
+    final_output: true
+    final_output_type: audio

vllm_omni/deploy/qwen3_tts.yaml

@@ -54,10 +54,7 @@ stages:
       temperature: 0.9
       top_k: 50
       max_tokens: 4096
-      # Official qwen-tts generates the talker with min_new_tokens=2; without
-      # it, checkpoints that sample codec EOS on the first frame return empty
-      # audio (HTTP 400). See #4962.
-      min_tokens: 2
+      seed: 42
       repetition_penalty: 1.05
     subtalker_sampling_params:
       do_sample: true
@@ -88,6 +85,7 @@ stages:
       top_p: 1.0
       top_k: -1
       max_tokens: 65536
+      seed: 42
       repetition_penalty: 1.0
 
 platforms:

vllm_omni/deploy/qwen3_tts_high_concurrency.yaml

@@ -60,10 +60,6 @@ stages:
       temperature: 0.9
       top_k: 50
       max_tokens: 4096
-      # Official qwen-tts generates the talker with min_new_tokens=2; without
-      # it, checkpoints that sample codec EOS on the first frame return empty
-      # audio (HTTP 400). See #4962.
-      min_tokens: 2
       repetition_penalty: 1.05
     subtalker_sampling_params:
       do_sample: true

vllm_omni/diffusion/attention/backends/flash_attn_hub.py

@@ -1,406 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import threading
-from functools import partial
-
-import torch
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.attention.backends.abstract import (
-    AttentionBackend,
-    AttentionImpl,
-    AttentionMetadata,
-)
-from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
-    piecewise_attn,
-)
-
-logger = init_logger(__name__)
-
-
-_hub_modules: dict[str, object] = {}
-_hub_lock = threading.Lock()
-
-
-def _load_hub_module(repo_id: str):
-    from kernels import get_kernel
-
-    logger.info("Loading %s kernel from HuggingFace Hub...", repo_id)
-    last_error = None
-    for version in (1, 2, None):
-        try:
-            if version is not None:
-                return get_kernel(repo_id, version=version)
-            return get_kernel(repo_id)
-        except Exception as exc:
-            logger.info("Failed to load %s version %s: %s", repo_id, version, exc)
-            last_error = exc
-    raise RuntimeError(f"Failed to load HuggingFace Hub kernel {repo_id!r}") from last_error
-
-
-def _get_hub_module(repo_id: str):
-    with _hub_lock:
-        if repo_id not in _hub_modules:
-            _hub_modules[repo_id] = _load_hub_module(repo_id)
-        return _hub_modules[repo_id]
-
-
-def _run_varlen_dense(
-    varlen_func,
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    *,
-    causal: bool,
-    softmax_scale: float,
-) -> torch.Tensor:
-    batch_size, q_len = query.size()[:2]
-    k_len = key.size(1)
-    cu_seqlens_q = torch.arange(0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=query.device)
-    cu_seqlens_k = torch.arange(0, (batch_size + 1) * k_len, step=k_len, dtype=torch.int32, device=query.device)
-
-    out = varlen_func(
-        q=query.flatten(0, 1),
-        k=key.flatten(0, 1),
-        v=value.flatten(0, 1),
-        cu_seqlens_q=cu_seqlens_q,
-        cu_seqlens_k=cu_seqlens_k,
-        max_seqlen_q=q_len,
-        max_seqlen_k=k_len,
-        causal=causal,
-        softmax_scale=softmax_scale,
-    )
-    if isinstance(out, tuple):
-        out = out[0]
-    return out.reshape(batch_size, q_len, *out.shape[1:])
-
-
-class FlashAttentionHubBackend(AttentionBackend):
-    accept_output_buffer: bool = True
-
-    @classmethod
-    def supports_attention_mask(cls) -> bool:
-        return True
-
-    @staticmethod
-    def get_supported_head_sizes() -> list[int]:
-        return [64, 96, 128, 192, 256]
-
-    @staticmethod
-    def get_name() -> str:
-        return "FLASH_ATTN_HUB"
-
-    @staticmethod
-    def get_impl_cls() -> type["FlashAttentionHubImpl"]:
-        return FlashAttentionHubImpl
-
-
-class FlashAttentionHubImpl(AttentionImpl):
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        softmax_scale: float,
-        causal: bool = False,
-        num_kv_heads: int | None = None,
-        prefix: str = "",
-        qkv_layout: str | None = None,
-        backend_kwargs: dict | None = None,
-        **extra_impl_args,
-    ) -> None:
-        self.num_heads = num_heads
-        self.causal = causal
-        self.softmax_scale = softmax_scale
-        self.qkv_layout = qkv_layout
-        if backend_kwargs:
-            logger.warning("FlashAttentionHubImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
-
-        hub_module = _get_hub_module("kernels-community/flash-attn2")
-        self.flash_attn_func = getattr(hub_module, "flash_attn_func", None)
-        self.flash_attn_varlen_func = getattr(hub_module, "flash_attn_varlen_func", None)
-        if self.flash_attn_func is None and self.flash_attn_varlen_func is None:
-            raise RuntimeError("Failed to load flash-attn2 kernel from HuggingFace Hub: no functions found")
-
-    @staticmethod
-    def _unwrap_flash_output(out: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor:
-        # FA3 may return (out, lse), FA2 returns out
-        return out[0] if isinstance(out, tuple) else out
-
-    def _flash_wrapper(self, q, k, v, *, attn_func, **kwargs):
-        if attn_func is not None:
-            return self._unwrap_flash_output(attn_func(q, k, v, **kwargs))
-        return _run_varlen_dense(
-            self.flash_attn_varlen_func,
-            q,
-            k,
-            v,
-            causal=kwargs.get("causal", self.causal),
-            softmax_scale=kwargs.get("softmax_scale", self.softmax_scale),
-        )
-
-    def _forward_varlen_masked(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attention_mask: torch.Tensor,
-    ) -> torch.Tensor:
-        from vllm_omni.diffusion.attention.backends.utils.fa import (
-            _pad_input,
-            _unpad_input,
-            _upad_input,
-        )
-
-        assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)"
-        query_length = query.size(1)
-        q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input(
-            query, key, value, attention_mask, query_length, _unpad_input
-        )
-
-        out_unpad = self.flash_attn_varlen_func(
-            q,
-            k,
-            v,
-            cu_seqlens_q=cu_seq_lens_q,
-            cu_seqlens_k=cu_seq_lens_k,
-            max_seqlen_q=max_length_q,
-            max_seqlen_k=max_length_k,
-            **{
-                "causal": self.causal,
-                "softmax_scale": self.softmax_scale,
-            },
-        )
-        out_unpad = self._unwrap_flash_output(out_unpad)
-        return _pad_input(out_unpad, indices_q, query.size(0), query_length)
-
-    def _forward_varlen_dense(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-    ) -> torch.Tensor:
-        return _run_varlen_dense(
-            self.flash_attn_varlen_func,
-            query,
-            key,
-            value,
-            causal=self.causal,
-            softmax_scale=self.softmax_scale,
-        )
-
-    def forward_cuda(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
-    ) -> torch.Tensor:
-        attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
-        full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
-
-        # Try piecewise attention
-        if full_attn_spans is not None:
-            logger.debug("Using piecewise Flash Attention for mixed causal/full mask")
-            attn_func = partial(
-                self._flash_wrapper,
-                attn_func=self.flash_attn_func,
-            )
-
-            return piecewise_attn(
-                query,
-                key,
-                value,
-                full_attn_spans,
-                self.softmax_scale,
-                attn_func,
-            )
-
-        if attention_mask is not None and torch.any(~attention_mask):
-            return self._forward_varlen_masked(
-                query,
-                key,
-                value,
-                attention_mask,
-            )
-
-        if self.flash_attn_func is not None:
-            out = self.flash_attn_func(
-                query,
-                key,
-                value,
-                causal=self.causal,
-                softmax_scale=self.softmax_scale,
-            )
-            return self._unwrap_flash_output(out)
-
-        return self._forward_varlen_dense(
-            query,
-            key,
-            value,
-        )
-
-
-class FlashAttention3HubBackend(AttentionBackend):
-    accept_output_buffer: bool = True
-
-    @classmethod
-    def supports_attention_mask(cls) -> bool:
-        return True
-
-    @staticmethod
-    def get_supported_head_sizes() -> list[int]:
-        return [64, 96, 128, 192, 256]
-
-    @staticmethod
-    def get_name() -> str:
-        return "FLASH_ATTN_3_HUB"
-
-    @staticmethod
-    def get_impl_cls() -> type["FlashAttention3HubImpl"]:
-        return FlashAttention3HubImpl
-
-
-class FlashAttention3HubImpl(AttentionImpl):
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        softmax_scale: float,
-        causal: bool = False,
-        num_kv_heads: int | None = None,
-        prefix: str = "",
-        qkv_layout: str | None = None,
-        backend_kwargs: dict | None = None,
-        **extra_impl_args,
-    ) -> None:
-        self.num_heads = num_heads
-        self.causal = causal
-        self.softmax_scale = softmax_scale
-        self.qkv_layout = qkv_layout
-        if backend_kwargs:
-            logger.warning("FlashAttention3HubImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
-
-        hub_module = _get_hub_module("kernels-community/flash-attn3")
-        self.flash_attn_func = getattr(hub_module, "flash_attn_func", None)
-        self.flash_attn_varlen_func = getattr(hub_module, "flash_attn_varlen_func", None)
-        if self.flash_attn_func is None and self.flash_attn_varlen_func is None:
-            raise RuntimeError("Failed to load flash-attn3 kernel from HuggingFace Hub: no functions found")
-
-    @staticmethod
-    def _unwrap_flash_output(out: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor:
-        # FA3 returns (out, lse)
-        return out[0] if isinstance(out, tuple) else out
-
-    def _flash_wrapper(self, q, k, v, *, attn_func, **kwargs):
-        if attn_func is not None:
-            return self._unwrap_flash_output(attn_func(q, k, v, **kwargs))
-        return _run_varlen_dense(
-            self.flash_attn_varlen_func,
-            q,
-            k,
-            v,
-            causal=kwargs.get("causal", self.causal),
-            softmax_scale=kwargs.get("softmax_scale", self.softmax_scale),
-        )
-
-    def _forward_varlen_masked(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attention_mask: torch.Tensor,
-    ) -> torch.Tensor:
-        from vllm_omni.diffusion.attention.backends.utils.fa import (
-            _pad_input,
-            _unpad_input,
-            _upad_input,
-        )
-
-        assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)"
-        query_length = query.size(1)
-        q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input(
-            query, key, value, attention_mask, query_length, _unpad_input
-        )
-
-        out_unpad = self.flash_attn_varlen_func(
-            q,
-            k,
-            v,
-            cu_seqlens_q=cu_seq_lens_q,
-            cu_seqlens_k=cu_seq_lens_k,
-            max_seqlen_q=max_length_q,
-            max_seqlen_k=max_length_k,
-            **{
-                "causal": self.causal,
-                "softmax_scale": self.softmax_scale,
-            },
-        )
-        out_unpad = self._unwrap_flash_output(out_unpad)
-        return _pad_input(out_unpad, indices_q, query.size(0), query_length)
-
-    def _forward_varlen_dense(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-    ) -> torch.Tensor:
-        return _run_varlen_dense(
-            self.flash_attn_varlen_func,
-            query,
-            key,
-            value,
-            causal=self.causal,
-            softmax_scale=self.softmax_scale,
-        )
-
-    def forward_cuda(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
-    ) -> torch.Tensor:
-        attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
-        full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
-
-        # Try piecewise attention
-        if full_attn_spans is not None:
-            logger.debug("Using piecewise Flash Attention for mixed causal/full mask")
-            attn_func = partial(
-                self._flash_wrapper,
-                attn_func=self.flash_attn_func,
-            )
-
-            return piecewise_attn(
-                query,
-                key,
-                value,
-                full_attn_spans,
-                self.softmax_scale,
-                attn_func,
-            )
-
-        if attention_mask is not None and torch.any(~attention_mask):
-            return self._forward_varlen_masked(
-                query,
-                key,
-                value,
-                attention_mask,
-            )
-
-        if self.flash_attn_func is not None:
-            out = self.flash_attn_func(
-                query,
-                key,
-                value,
-                causal=self.causal,
-                softmax_scale=self.softmax_scale,
-            )
-            return self._unwrap_flash_output(out)
-
-        return self._forward_varlen_dense(
-            query,
-            key,
-            value,
-        )

vllm_omni/diffusion/attention/backends/registry.py

@@ -62,8 +62,6 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
     SAGE_ATTN_3 = "vllm_omni.diffusion.attention.backends.sage_attn3.SageAttention3Backend"
     CUDNN_ATTN = "vllm_omni.diffusion.attention.backends.cudnn_attn.CuDNNAttentionBackend"
     FLASHINFER_ATTN = "vllm_omni.diffusion.attention.backends.flashinfer_attn.FlashInferAttentionBackend"
-    FLASH_ATTN_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttentionHubBackend"
-    FLASH_ATTN_3_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttention3HubBackend"
 
     def get_path(self, include_classname: bool = True) -> str:
         """Get the class path for this backend (respects overrides).

vllm_omni/diffusion/layers/rope.py

@@ -188,12 +188,6 @@ class RotaryEmbedding(CustomOp):
         cos: torch.Tensor,
         sin: torch.Tensor,
     ) -> torch.Tensor:
-        # All batch elements share the same rotary position encoding.
-        # Strip the batch dim so the underlying op broadcasts over the batch,
-        # consistent with forward_cuda / forward_hip / apply_rotary_emb_mindiesd.
-        if cos.dim() == 3:
-            cos = cos[0]
-            sin = sin[0]
         return apply_rotary_emb_torch(
             x,
             cos,

vllm_omni/diffusion/models/lingbot_video/__init__.py

@@ -1,16 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer import (
-    LingBotVideoTransformer3DModel,
-)
-from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import (
-    LingBotVideoPipeline,
-    get_lingbot_video_post_process_func,
-)
-
-__all__ = [
-    "LingBotVideoPipeline",
-    "LingBotVideoTransformer3DModel",
-    "get_lingbot_video_post_process_func",
-]

vllm_omni/diffusion/models/lingbot_video/lingbot_video_transformer.py

@@ -1,702 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# Adapted from LingBot-Video (https://github.com/Robbyant/lingbot-video).
-
-import math
-import os
-
-import torch
-import torch.distributed as dist
-import torch.nn as nn
-import torch.nn.functional as F
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.models.embeddings import TimestepEmbedding, Timesteps
-from diffusers.models.modeling_outputs import Transformer2DModelOutput
-from diffusers.models.modeling_utils import ModelMixin
-
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
-from vllm_omni.diffusion.attention.backends.utils.fa import (
-    flash_attn_varlen_func as flash_attn_varlen_func_v3,
-)
-from vllm_omni.diffusion.attention.layer import Attention
-
-LINGBOT_VIDEO_FP32_MODULES = (
-    "time_embedder",
-    "time_modulation",
-    "scale_shift_table",
-    "norm",
-    "norm1",
-    "norm2",
-    "norm_q",
-    "norm_k",
-    "norm_post_attn",
-    "norm_post_ffn",
-    "norm_out",
-    "norm_out_modulation",
-    "router",
-)
-
-
-def should_keep_in_fp32(name: str) -> bool:
-    return any(module_name in name.split(".") for module_name in LINGBOT_VIDEO_FP32_MODULES)
-
-
-def _all_to_all_split_cat(
-    local_input: torch.Tensor,
-    scatter_dim: int,
-    gather_dim: int,
-    group: dist.ProcessGroup,
-) -> torch.Tensor:
-    world_size = dist.get_world_size(group)
-    input_list = [tensor.contiguous() for tensor in torch.tensor_split(local_input, world_size, scatter_dim)]
-    output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)]
-    dist.all_to_all(output_list, input_list, group=group)
-    return torch.cat(output_list, dim=gather_dim).contiguous()
-
-
-class LingBotVideoRMSNorm(nn.Module):
-    """RMSNorm with fp32 accumulation."""
-
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.variance_epsilon = eps
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return (self.weight * hidden_states).to(input_dtype)
-
-
-def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
-    """Apply complex RoPE to `(B, S, H, D)` attention tensors."""
-    with torch.amp.autocast("cuda", enabled=False):
-        x_c = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
-        out = torch.view_as_real(x_c * freqs_cis.unsqueeze(2)).flatten(3)
-        return out.type_as(x)
-
-
-class LingBotVideoRotaryEmbedding(nn.Module):
-    """Complex64 RoPE table indexed by position ids."""
-
-    def __init__(self, axes_dims: tuple[int, ...], axes_lens: tuple[int, ...], theta: float):
-        super().__init__()
-        self.axes_dims = tuple(axes_dims)
-        self.axes_lens = list(axes_lens)
-        self.theta = theta
-        self.freqs_cis = None
-
-    @staticmethod
-    def precompute_freqs_cis(dim: tuple[int, ...], end: tuple[int, ...], theta: float):
-        freqs_cis = []
-        for d, e in zip(dim, end):
-            freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d))
-            timestep = torch.arange(e, device=freqs.device, dtype=torch.float64)
-            freqs = torch.outer(timestep, freqs).float()
-            freqs_cis.append(torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64))
-        return freqs_cis
-
-    def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
-        # position_ids: (S, 3) int -> (S, head_dim/2) complex64
-        device = position_ids.device
-        max_vals = position_ids.max(dim=0).values.tolist()
-        needs_rebuild = self.freqs_cis is None or any(
-            max_val >= axis_len for max_val, axis_len in zip(max_vals, self.axes_lens)
-        )
-        if needs_rebuild:
-            for i in range(len(self.axes_lens)):
-                if max_vals[i] >= self.axes_lens[i]:
-                    self.axes_lens[i] = int(max_vals[i] * 1.5) + 1
-            self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, tuple(self.axes_lens), theta=self.theta)
-            self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
-        elif self.freqs_cis[0].device != device:
-            self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
-
-        return torch.cat([self.freqs_cis[i][position_ids[:, i]] for i in range(len(self.axes_dims))], dim=-1)
-
-
-def make_joint_position_ids(text_len: int, grid_t: int, grid_h: int, grid_w: int, device: torch.device) -> torch.Tensor:
-    """3D positions in [video; text] order. Text t-axis is 1..text_len; video t-axis starts at text_len+1.
-
-    Matches patchify_and_embed: cap start (1,0,0); vision start (cap_len+1,0,0);
-    freqs ordered with x first and cap second (same order as cat_interleave).
-    """
-    tt = torch.arange(grid_t, device=device, dtype=torch.int32) + (text_len + 1)
-    hh = torch.arange(grid_h, device=device, dtype=torch.int32)
-    ww = torch.arange(grid_w, device=device, dtype=torch.int32)
-    grid = torch.stack(torch.meshgrid(tt, hh, ww, indexing="ij"), dim=-1).flatten(0, 2)
-    text_t = torch.arange(text_len, device=device, dtype=torch.int32) + 1
-    text_pos = torch.stack([text_t, torch.zeros_like(text_t), torch.zeros_like(text_t)], dim=-1)
-    return torch.cat([grid, text_pos], dim=0)  # (Nx + L, 3)
-
-
-def _cat_interleave(
-    a: torch.Tensor,
-    len_a: list[int],
-    b: torch.Tensor,
-    len_b: list[int],
-) -> torch.Tensor:
-    a_split = torch.split(a, len_a, dim=1)
-    b_split = torch.split(b, len_b, dim=1)
-    blocks: list[torch.Tensor] = []
-    for x_part, text_part in zip(a_split, b_split):
-        blocks.extend([x_part, text_part])
-    return torch.cat(blocks, dim=1)
-
-
-def _packed_block_attention_mask(sample_seq_lens: list[int], device: torch.device) -> torch.Tensor:
-    total_seq_len = sum(sample_seq_lens)
-    mask = torch.zeros(
-        1,
-        1,
-        total_seq_len,
-        total_seq_len,
-        dtype=torch.bool,
-        device=device,
-    )
-    start = 0
-    for seq_len in sample_seq_lens:
-        end = start + seq_len
-        mask[:, :, start:end, start:end] = True
-        start = end
-    return mask
-
-
-class LingBotVideoTextEmbedder(nn.Module):
-    """Matches CondProjection: RMSNorm(text_dim, eps=1e-6 fixed) -> Linear-SiLU-Linear."""
-
-    def __init__(self, text_dim: int, hidden_size: int):
-        super().__init__()
-        self.norm = LingBotVideoRMSNorm(text_dim, eps=1e-6)
-        self.linear_1 = nn.Linear(text_dim, hidden_size, bias=True)
-        self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.norm(x)
-        return self.linear_2(F.silu(self.linear_1(x)))
-
-
-class LingBotVideoAttention(nn.Module):
-    def __init__(self, hidden_size, num_heads, norm_eps, qkv_bias, out_bias):
-        super().__init__()
-        self.num_heads = num_heads
-        self.head_dim = hidden_size // num_heads
-        self.to_q = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
-        self.to_k = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
-        self.to_v = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
-        self.norm_q = LingBotVideoRMSNorm(self.head_dim, norm_eps)
-        self.norm_k = LingBotVideoRMSNorm(self.head_dim, norm_eps)
-        self.to_out = nn.Linear(hidden_size, hidden_size, bias=out_bias)
-        self.attn = Attention(
-            num_heads=num_heads,
-            head_size=self.head_dim,
-            softmax_scale=1.0 / math.sqrt(self.head_dim),
-            causal=False,
-            num_kv_heads=num_heads,
-        )
-
-    def forward(
-        self,
-        x,
-        rotary_emb,
-        attention_mask=None,
-        packed_indices: dict[str, torch.Tensor] | None = None,
-        parallel_config=None,
-    ):
-        B, S, _ = x.shape
-        if os.environ.get("LINGBOT_FUSED_QKV_LINEAR") == "1":
-            weight = torch.cat(
-                (self.to_q.weight, self.to_k.weight, self.to_v.weight),
-                dim=0,
-            )
-            bias = None
-            if self.to_q.bias is not None:
-                bias = torch.cat(
-                    (self.to_q.bias, self.to_k.bias, self.to_v.bias),
-                    dim=0,
-                )
-            qkv = F.linear(x, weight, bias)
-            q, k, v = qkv.view(B, S, 3, self.num_heads, self.head_dim).unbind(2)
-        else:
-            q = self.to_q(x).unflatten(2, (self.num_heads, self.head_dim))
-            k = self.to_k(x).unflatten(2, (self.num_heads, self.head_dim))
-            v = self.to_v(x).unflatten(2, (self.num_heads, self.head_dim))
-        q = apply_rotary_emb(self.norm_q(q), rotary_emb)
-        k = apply_rotary_emb(self.norm_k(k), rotary_emb)
-        if packed_indices is None:
-            out = self.attn(q, k, v, attn_metadata=AttentionMetadata(attn_mask=attention_mask))
-        else:
-            packed_attention_mask = packed_indices.get("attention_mask")
-            if packed_attention_mask is not None and parallel_config is None:
-                out = self.attn.sdpa_fallback.forward(
-                    q,
-                    k,
-                    v,
-                    AttentionMetadata(attn_mask=packed_attention_mask),
-                )
-                return self.to_out(out.flatten(2, 3).type_as(x))
-            if flash_attn_varlen_func_v3 is None:
-                raise RuntimeError(
-                    "A flash attention varlen function is required for packed context parallel attention."
-                )
-            if parallel_config is None:
-                result = flash_attn_varlen_func_v3(
-                    q=q.reshape(-1, self.num_heads, self.head_dim),
-                    k=k.reshape(-1, self.num_heads, self.head_dim),
-                    v=v.reshape(-1, self.num_heads, self.head_dim),
-                    cu_seqlens_q=packed_indices["cu_seqlens_kv"],
-                    cu_seqlens_k=packed_indices["cu_seqlens_kv"],
-                    max_seqlen_q=packed_indices["max_seqlen_in_batch_kv"],
-                    max_seqlen_k=packed_indices["max_seqlen_in_batch_kv"],
-                    causal=False,
-                )
-                out = result[0] if isinstance(result, tuple) else result
-                out = out.reshape(B, S, self.num_heads, self.head_dim)
-            else:
-                group = parallel_config.context_parallel_config._ulysses_mesh.get_group()
-                world_size = dist.get_world_size(group)
-                local_heads = self.num_heads // world_size
-                q_global = _all_to_all_split_cat(
-                    q.reshape(B, S, self.num_heads * self.head_dim),
-                    scatter_dim=2,
-                    gather_dim=1,
-                    group=group,
-                ).view(B, S * world_size, local_heads, self.head_dim)
-                k_global = _all_to_all_split_cat(
-                    k.reshape(B, S, self.num_heads * self.head_dim),
-                    scatter_dim=2,
-                    gather_dim=1,
-                    group=group,
-                ).view(B, S * world_size, local_heads, self.head_dim)
-                v_global = _all_to_all_split_cat(
-                    v.reshape(B, S, self.num_heads * self.head_dim),
-                    scatter_dim=2,
-                    gather_dim=1,
-                    group=group,
-                ).view(B, S * world_size, local_heads, self.head_dim)
-                q_flat = q_global.reshape(-1, local_heads, self.head_dim)
-                k_flat = k_global.reshape(-1, local_heads, self.head_dim)
-                v_flat = v_global.reshape(-1, local_heads, self.head_dim)
-                result = flash_attn_varlen_func_v3(
-                    q=q_flat,
-                    k=k_flat,
-                    v=v_flat,
-                    cu_seqlens_q=packed_indices["cu_seqlens_kv"],
-                    cu_seqlens_k=packed_indices["cu_seqlens_kv"],
-                    max_seqlen_q=packed_indices["max_seqlen_in_batch_kv"],
-                    max_seqlen_k=packed_indices["max_seqlen_in_batch_kv"],
-                    causal=False,
-                )
-                out_global = result[0] if isinstance(result, tuple) else result
-                out_global = out_global.reshape(B, S * world_size, local_heads * self.head_dim)
-                out = _all_to_all_split_cat(
-                    out_global,
-                    scatter_dim=1,
-                    gather_dim=2,
-                    group=group,
-                ).view(B, S, self.num_heads, self.head_dim)
-        return self.to_out(out.flatten(2, 3).type_as(x))
-
-
-class LingBotVideoMLP(nn.Module):
-    def __init__(self, hidden_size, intermediate_size):
-        super().__init__()
-        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
-
-    def forward(self, x):
-        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
-
-
-class LingBotVideoBlock(nn.Module):
-    def __init__(
-        self,
-        hidden_size,
-        num_attention_heads,
-        intermediate_size,
-        norm_eps,
-        qkv_bias,
-        out_bias,
-        num_experts,
-        num_experts_per_tok,
-        moe_intermediate_size,
-        decoder_sparse_step,
-        mlp_only_layers,
-        n_shared_experts,
-        score_func,
-        norm_topk_prob,
-        n_group,
-        topk_group,
-        routed_scaling_factor,
-        layer_idx: int,
-    ):
-        super().__init__()
-        self.layer_idx = layer_idx
-        h = hidden_size
-        self.scale_shift_table = nn.Parameter(torch.zeros(1, 6 * h))
-        self.norm1 = LingBotVideoRMSNorm(h, norm_eps)
-        self.attn = LingBotVideoAttention(h, num_attention_heads, norm_eps, qkv_bias, out_bias)
-        self.norm_post_attn = LingBotVideoRMSNorm(h, norm_eps)
-        self.norm2 = LingBotVideoRMSNorm(h, norm_eps)
-        if num_experts > 0:
-            raise NotImplementedError(
-                "LingBotVideoTransformer3DModel in vLLM-Omni currently supports "
-                "the dense checkpoint only. MoE/fused-expert support belongs in "
-                "the follow-up MoE PR."
-            )
-        self.ffn = LingBotVideoMLP(h, intermediate_size)
-        self.norm_post_ffn = LingBotVideoRMSNorm(h, norm_eps)
-
-    def forward(
-        self,
-        x,
-        temb6,
-        rotary_emb,
-        attention_mask=None,
-        moe_padding_mask=None,
-        packed_indices: dict[str, torch.Tensor] | None = None,
-        parallel_config=None,
-    ):
-        expected_tokens = x.shape[0] * x.shape[1]
-        if temb6.ndim != 2 or temb6.shape[0] != expected_tokens:
-            raise ValueError(
-                "LingBotVideoBlock expects token-level temb6 with shape "
-                f"(B*S, 6D); got {tuple(temb6.shape)} for hidden states {tuple(x.shape)}."
-            )
-        # AdaLN modulation and normalization stay in fp32 for the sensitive path.
-        mod = temb6.view(x.shape[0], x.shape[1], -1) + self.scale_shift_table.unsqueeze(0)
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=-1)
-        gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
-        scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
-
-        # AdaLN modulation / norms run in fp32 (sensitive path); cast to the bulk
-        # compute dtype only at the bf16 Linear boundary. This replaces the old
-        # ambient autocast, which rounded Linear inputs to bf16 at the same point.
-        bulk_dtype = self.attn.to_q.weight.dtype
-        attn_in = (self.norm1(x) * scale_msa + shift_msa).to(bulk_dtype)
-        attn_out = self.attn(
-            attn_in,
-            rotary_emb,
-            attention_mask,
-            packed_indices=packed_indices,
-            parallel_config=parallel_config,
-        )
-        x = x + (gate_msa * self.norm_post_attn(attn_out)).to(x.dtype)
-
-        ffn_in = (self.norm2(x) * scale_mlp + shift_mlp).to(bulk_dtype)
-        ffn_out = self.ffn(ffn_in)
-        ffn_normed = self.norm_post_ffn(ffn_out)
-        x = x + (gate_mlp * ffn_normed).to(x.dtype)
-        return x
-
-
-class LingBotVideoTransformer3DModel(ModelMixin, ConfigMixin):
-    _supports_gradient_checkpointing = False
-    _repeated_blocks = ["LingBotVideoBlock"]
-    _layerwise_offload_blocks_attr = "blocks"
-    _no_split_modules = ["LingBotVideoBlock"]
-    _keep_in_fp32_modules = list(LINGBOT_VIDEO_FP32_MODULES)
-
-    def to(self, *args, **kwargs):
-        device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
-        if dtype is None or dtype == torch.float32:
-            return super().to(*args, **kwargs)
-
-        dtype_is_floating = torch.is_floating_point(torch.empty((), dtype=dtype))
-        if not dtype_is_floating:
-            return super().to(*args, **kwargs)
-
-        if device is not None:
-            super().to(device=device, non_blocking=non_blocking)
-
-        for name, param in self.named_parameters():
-            if not torch.is_floating_point(param):
-                continue
-            target_dtype = torch.float32 if should_keep_in_fp32(name) else dtype
-            param.data = param.data.to(dtype=target_dtype, non_blocking=non_blocking)
-            if param.grad is not None:
-                param.grad.data = param.grad.data.to(dtype=target_dtype, non_blocking=non_blocking)
-
-        for name, buffer in self.named_buffers():
-            if not torch.is_floating_point(buffer):
-                continue
-            target_dtype = torch.float32 if should_keep_in_fp32(name) else dtype
-            buffer.data = buffer.data.to(dtype=target_dtype, non_blocking=non_blocking)
-
-        return self
-
-    @register_to_config
-    def __init__(
-        self,
-        patch_size: tuple[int, int, int] = (1, 2, 2),
-        in_channels: int = 16,
-        out_channels: int = 16,
-        hidden_size: int = 2048,
-        num_attention_heads: int = 16,
-        depth: int = 24,
-        intermediate_size: int = 6144,
-        text_dim: int = 2560,
-        freq_dim: int = 256,
-        norm_eps: float = 1e-6,
-        rope_theta: float = 256.0,
-        axes_dims: tuple[int, int, int] = (32, 48, 48),
-        axes_lens: tuple[int, int, int] = (8192, 1024, 1024),
-        qkv_bias: bool = False,
-        out_bias: bool = True,
-        patch_embed_bias: bool = True,
-        timestep_mlp_bias: bool = True,
-        num_experts: int = 0,
-        num_experts_per_tok: int = 8,
-        moe_intermediate_size: int = 512,
-        decoder_sparse_step: int = 1,
-        mlp_only_layers: tuple[int, ...] = (),
-        n_shared_experts: int | None = None,
-        score_func: str = "sigmoid",
-        norm_topk_prob: bool = True,
-        n_group: int | None = None,
-        topk_group: int | None = None,
-        routed_scaling_factor: float = 1.0,
-    ):
-        super().__init__()
-        head_dim = hidden_size // num_attention_heads
-        assert head_dim == sum(axes_dims), f"head_dim {head_dim} != sum(axes_dims) {sum(axes_dims)}"
-        mlp_only_layers = tuple(mlp_only_layers)
-
-        self.patch_embedder = nn.Linear(in_channels * math.prod(patch_size), hidden_size, bias=patch_embed_bias)
-        self.time_proj = Timesteps(freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
-        self.time_embedder = TimestepEmbedding(freq_dim, hidden_size, act_fn="silu", sample_proj_bias=timestep_mlp_bias)
-        self.time_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size))
-        self.text_embedder = LingBotVideoTextEmbedder(text_dim, hidden_size)
-        self.rope = LingBotVideoRotaryEmbedding(axes_dims, axes_lens, rope_theta)
-        self.blocks = nn.ModuleList(
-            [
-                LingBotVideoBlock(
-                    hidden_size=hidden_size,
-                    num_attention_heads=num_attention_heads,
-                    intermediate_size=intermediate_size,
-                    norm_eps=norm_eps,
-                    qkv_bias=qkv_bias,
-                    out_bias=out_bias,
-                    num_experts=num_experts,
-                    num_experts_per_tok=num_experts_per_tok,
-                    moe_intermediate_size=moe_intermediate_size,
-                    decoder_sparse_step=decoder_sparse_step,
-                    mlp_only_layers=mlp_only_layers,
-                    n_shared_experts=n_shared_experts,
-                    score_func=score_func,
-                    norm_topk_prob=norm_topk_prob,
-                    n_group=n_group,
-                    topk_group=topk_group,
-                    routed_scaling_factor=routed_scaling_factor,
-                    layer_idx=i,
-                )
-                for i in range(depth)
-            ]
-        )
-        self.norm_out = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps)
-        self.norm_out_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size))
-        self.proj_out = nn.Linear(hidden_size, math.prod(patch_size) * out_channels)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,  # (B, C, T, H, W)
-        timestep: torch.Tensor,  # (B,) in [0, 1000](= sigma*1000)
-        encoder_hidden_states: torch.Tensor,  # (B, L, text_dim)
-        encoder_attention_mask: torch.Tensor | None = None,  # (B, L) 1=valid
-        return_dict: bool = True,
-    ):
-        B, C, T, H, W = hidden_states.shape
-        pF, pH, pW = self.config.patch_size
-        gt, gh, gw = T // pF, H // pH, W // pW
-        n_video = gt * gh * gw
-        L = encoder_hidden_states.shape[1]
-        device = hidden_states.device
-        if encoder_attention_mask is not None:
-            text_lens = encoder_attention_mask.sum(dim=-1).long()
-        else:
-            text_lens = torch.full((B,), L, dtype=torch.long, device=device)
-        text_lens_list = [int(v) for v in text_lens.detach().cpu().tolist()]
-        packed_batch = B > 1
-
-        # patchify: token order (f h w), feature order (pf ph pw c) -- matches patchify_and_embed
-        patch_tokens = hidden_states.reshape(B, C, gt, pF, gh, pH, gw, pW)
-        patch_tokens = patch_tokens.permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(
-            B,
-            n_video,
-            pF * pH * pW * C,
-        )
-        if packed_batch:
-            x = torch.cat(
-                [self.patch_embedder(patch_tokens[i : i + 1]) for i in range(B)],
-                dim=1,
-            )
-        else:
-            x = self.patch_embedder(patch_tokens)
-
-        if packed_batch:
-            text_parts = [
-                self.text_embedder(encoder_hidden_states[i : i + 1, : text_lens_list[i], :]) for i in range(B)
-            ]
-            text = torch.cat(text_parts, dim=1)
-            joint = _cat_interleave(
-                x,
-                [n_video] * B,
-                text,
-                text_lens_list,
-            )
-        else:
-            text = self.text_embedder(encoder_hidden_states)
-            joint = torch.cat([x, text], dim=1)  # [video; text]
-        joint_seq_len = joint.shape[1]
-
-        # Per-sample RoPE: video t-axis start = real text length of this sample + 1
-        rotary_parts = [self.rope(make_joint_position_ids(text_lens_list[i], gt, gh, gw, device)) for i in range(B)]
-        if packed_batch:
-            rotary = torch.cat(rotary_parts, dim=0).unsqueeze(0)
-        else:
-            rotary = torch.stack(rotary_parts, dim=0)  # (B, S, head_dim/2) complex64
-
-        parallel_config = getattr(self, "_parallel_config", None)
-        use_packed_attention = parallel_config is not None
-
-        attention_mask = None
-        moe_padding_mask = None
-        packed_indices = None
-        has_padding = encoder_attention_mask is not None and bool((text_lens < L).any())
-        if packed_batch or use_packed_attention:
-            sample_seq_lens = [n_video + text_len for text_len in text_lens_list]
-            cu_seqlens = torch.zeros(B + 1, device=device, dtype=torch.int32)
-            cu_seqlens[1:] = torch.cumsum(
-                torch.tensor(sample_seq_lens, device=device, dtype=torch.int32),
-                dim=0,
-            )
-            packed_indices = {
-                "cu_seqlens_kv": cu_seqlens,
-                "max_seqlen_in_batch_kv": max(sample_seq_lens),
-            }
-            if packed_batch and not use_packed_attention:
-                packed_indices["attention_mask"] = _packed_block_attention_mask(sample_seq_lens, device)
-            has_padding = False
-        if has_padding:
-            key_mask = torch.cat(
-                [torch.ones(B, n_video, dtype=torch.bool, device=device), encoder_attention_mask.bool()],
-                dim=1,
-            )
-            attention_mask = key_mask[:, None, None, :]  # (B,1,1,S) -> SDPA broadcast
-            moe_padding_mask = key_mask.reshape(-1).float()  # (B*S,)
-        packed_cp = packed_indices is not None and parallel_config is not None
-        padding_size = 0
-        if packed_cp:
-            cp_config = parallel_config.context_parallel_config
-            cp_world_size = int(getattr(cp_config, "ulysses_degree", getattr(cp_config, "_world_size", 1)))
-            padding_size = (cp_world_size - (joint_seq_len % cp_world_size)) % cp_world_size
-            if padding_size:
-                joint = torch.cat(
-                    [
-                        joint,
-                        torch.zeros(
-                            joint.shape[0],
-                            padding_size,
-                            joint.shape[2],
-                            device=joint.device,
-                            dtype=joint.dtype,
-                        ),
-                    ],
-                    dim=1,
-                )
-                rotary = torch.cat(
-                    [
-                        rotary,
-                        torch.zeros(
-                            rotary.shape[0],
-                            padding_size,
-                            rotary.shape[2],
-                            device=rotary.device,
-                            dtype=rotary.dtype,
-                        ),
-                    ],
-                    dim=1,
-                )
-                if packed_indices is None:
-                    raise RuntimeError("packed_indices must be initialized for packed context parallel.")
-                packed_indices["cu_seqlens_kv"] = torch.cat(
-                    [
-                        packed_indices["cu_seqlens_kv"],
-                        packed_indices["cu_seqlens_kv"][-1:] + padding_size,
-                    ],
-                    dim=0,
-                )
-                packed_indices["max_seqlen_in_batch_kv"] = max(
-                    int(packed_indices["max_seqlen_in_batch_kv"]),
-                    int(padding_size),
-                )
-                joint_seq_len = joint.shape[1]
-
-        timestep_for_embed = timestep.float()
-        timestep_proj = self.time_proj(timestep_for_embed)
-        t_emb = self.time_embedder(timestep_proj)  # (B, D)
-        if packed_batch:
-            temb_input = torch.cat(
-                [t_emb[i : i + 1].unsqueeze(1).expand(1, n_video + text_lens_list[i], -1) for i in range(B)],
-                dim=1,
-            )
-            if padding_size:
-                temb_input = torch.cat(
-                    [
-                        temb_input,
-                        torch.zeros(
-                            temb_input.shape[0],
-                            padding_size,
-                            temb_input.shape[2],
-                            device=temb_input.device,
-                            dtype=temb_input.dtype,
-                        ),
-                    ],
-                    dim=1,
-                )
-            temb6 = self.time_modulation(temb_input.reshape(joint_seq_len, -1))
-            temb6 = temb6.reshape(1, joint_seq_len, -1)
-        else:
-            temb_input = t_emb.unsqueeze(1).expand(B, joint_seq_len, -1)  # (B, S, D)
-            temb6 = self.time_modulation(temb_input.reshape(B * joint_seq_len, -1))
-            temb6 = temb6.reshape(B, joint_seq_len, -1)  # (B, S, 6D)
-
-        temb6 = temb6.reshape(temb6.shape[0] * temb6.shape[1], -1)
-
-        for block in self.blocks:
-            joint = block(
-                joint,
-                temb6,
-                rotary,
-                attention_mask,
-                moe_padding_mask,
-                packed_indices=packed_indices,
-                parallel_config=parallel_config,
-            )
-        final_mod = self.norm_out_modulation(temb_input.reshape(joint.shape[0] * joint.shape[1], -1))
-        shift, scale = final_mod.reshape(joint.shape[0], joint.shape[1], -1).chunk(2, dim=-1)
-        final_hidden = self.norm_out(joint) * (1.0 + scale) + shift
-        projected = self.proj_out(final_hidden.to(self.proj_out.weight.dtype))
-        if packed_cp:
-            if padding_size:
-                projected = projected[:, :-padding_size, :]
-        if packed_batch:
-            split_lengths: list[int] = []
-            for text_len in text_lens_list:
-                split_lengths.extend([n_video, text_len])
-            parts = torch.split(projected, split_lengths, dim=1)
-            x = torch.cat(parts[::2], dim=1).reshape(B, n_video, -1)
-        else:
-            x = projected[:, :n_video]
-
-        # unpatchify (matches the rearrange in postprocess)
-        Cout = self.config.out_channels
-        x = x.reshape(B, gt, gh, gw, pF, pH, pW, Cout)
-        x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(B, Cout, T, H, W)
-
-        if not return_dict:
-            return (x,)
-        return Transformer2DModelOutput(sample=x)

vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py

@@ -1,695 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# Adapted from LingBot-Video (https://github.com/Robbyant/lingbot-video).
-
-from __future__ import annotations
-
-import os
-from collections.abc import Iterable, Mapping
-from contextlib import nullcontext
-from typing import Any, ClassVar
-
-import numpy as np
-import torch
-import torch.distributed as dist
-from diffusers import AutoencoderKLWan
-from diffusers.utils.torch_utils import randn_tensor
-from torch import nn
-from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
-
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
-from vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer import LingBotVideoTransformer3DModel
-from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
-from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
-from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-
-TOKEN_LENGTH = 37698
-HIDDEN_STATE_SKIP_LAYER = 0
-LOW_NOISE_TAIL_V1_DEFAULT_STEPS = 2
-
-PROMPT_TEMPLATE = (
-    "<|im_start|>system\nGiven a user input that may include a text prompt alone, "
-    "a text prompt with an image reference, or a text prompt with a video reference "
-    'or a video reference alone, generate an "Enhanced prompt" that provides detailed '
-    "visual descriptions suitable for video generation. Evaluate the level of detail "
-    "in the user's input: if it is simple, enrich it by adding specifics about colors, "
-    "shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal "
-    "progression, and spatial relationships to create vivid, concrete, and temporally "
-    "coherent scenes to create vivid and concrete scenes. Please generate only the "
-    "enhanced description for the prompt below and avoid including any additional "
-    "commentary or evaluations:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n"
-    "<|im_start|>assistant\n"
-)
-DEFAULT_NEGATIVE_PROMPT = (
-    '{"universal_negative": {"visual_quality": ["low quality", "worst quality", "blurry", '
-    '"pixelated", "jpeg artifacts", "low resolution", "unstable color", "color flicker", '
-    '"underexposed", "overexposed", "invisible subject", "subject hidden in darkness"], '
-    '"artistic_style": ["painting", "illustration", "drawing", "cartoon", "3d render", '
-    '"cgi", "sketch", "digital art"], "composition_and_content": ["text", "watermark", '
-    '"signature", "logo", "subtitles", "pillarboxed", "side bars", '
-    '"portrait image in landscape frame"], "temporal_and_motion_stability": ["flickering", '
-    '"jittery", "motion blur", "temporal inconsistency", "warping", "morphing", '
-    '"incoherent motion", "unnatural movement", "static object with sudden jump", '
-    '"frame-to-frame inconsistency"], "material_and_structure": ["plastic-like glass", '
-    '"unrealistic texture", "deformed bottle", "liquid freezing improperly", '
-    '"distorted reflections"]}}'
-)
-
-
-def _dtype_from_name(value: Any, default: torch.dtype) -> torch.dtype:
-    if isinstance(value, torch.dtype):
-        return value
-    if value is None:
-        return default
-    normalized = str(value).lower()
-    if normalized in {"bf16", "bfloat16", "torch.bfloat16"}:
-        return torch.bfloat16
-    if normalized in {"fp16", "float16", "torch.float16"}:
-        return torch.float16
-    if normalized in {"fp32", "float32", "torch.float32"}:
-        return torch.float32
-    raise ValueError(f"Unsupported LingBot dtype: {value!r}.")
-
-
-def _extract_prompt(req: OmniDiffusionRequest) -> tuple[str, str | None]:
-    prompt_obj = req.prompt
-    if isinstance(prompt_obj, str):
-        return prompt_obj, None
-    if isinstance(prompt_obj, Mapping):
-        prompt = prompt_obj.get("prompt", "")
-        negative_prompt = prompt_obj.get("negative_prompt")
-        return str(prompt), None if negative_prompt is None else str(negative_prompt)
-    raise TypeError(f"Unsupported LingBot prompt type: {type(prompt_obj)!r}.")
-
-
-def _module_dtype(module: torch.nn.Module) -> torch.dtype:
-    try:
-        return next(module.parameters()).dtype
-    except StopIteration:
-        return torch.float32
-
-
-def _module_device(module: torch.nn.Module) -> torch.device:
-    try:
-        return next(module.parameters()).device
-    except StopIteration:
-        return torch.device("cpu")
-
-
-def _transformer_timestep(timestep: torch.Tensor, transformer_dtype: torch.dtype) -> torch.Tensor:
-    sigma = timestep.float() / 1000.0
-    if transformer_dtype in {torch.bfloat16, torch.float16}:
-        sigma = sigma.to(transformer_dtype)
-    return (sigma * 1000.0).float()
-
-
-def _transformer_autocast(device: torch.device, transformer_dtype: torch.dtype):
-    if device.type != "cuda" or transformer_dtype not in {torch.bfloat16, torch.float16}:
-        return nullcontext()
-    return torch.autocast(device_type="cuda", dtype=transformer_dtype)
-
-
-def _group_global_rank(group: Any | None, group_rank: int) -> int:
-    if group is None:
-        return group_rank
-    get_global_rank = getattr(dist, "get_global_rank", None)
-    if get_global_rank is None:
-        return group_rank
-    return int(get_global_rank(group, group_rank))
-
-
-def _validate_refiner_sigmas(sigmas: np.ndarray, t_thresh: float | None = None) -> np.ndarray:
-    arr = np.asarray(list(sigmas), dtype=np.float64)
-    if arr.ndim != 1 or arr.size == 0:
-        raise ValueError("refiner sigma schedule must be a non-empty 1D list")
-    if not np.all(np.isfinite(arr)):
-        raise ValueError("refiner sigma schedule contains non-finite values")
-    if np.any(arr < 0.0) or np.any(arr > 1.0):
-        raise ValueError(f"refiner sigma schedule values must be in [0, 1], got {arr.tolist()}")
-    if arr.size > 1 and not np.all(np.diff(arr) < 0.0):
-        raise ValueError(f"refiner sigma schedule must be strictly descending, got {arr.tolist()}")
-    if t_thresh is not None and abs(float(arr[0]) - float(t_thresh)) > 1e-6:
-        raise ValueError(f"refiner sigma schedule must start at t_thresh={float(t_thresh)}, got {float(arr[0])}")
-    return arr
-
-
-def _compute_refiner_sigmas(
-    *,
-    sigma_max: float,
-    sigma_min: float,
-    num_inference_steps: int,
-    shift: float,
-    t_thresh: float | None,
-    tail_steps: int = 0,
-) -> np.ndarray | None:
-    if t_thresh is None:
-        return None
-    t_value = float(t_thresh)
-    if not (0.0 < t_value <= 1.0):
-        raise ValueError(f"refiner t_thresh must lie in (0, 1], got {t_value}")
-    steps = int(num_inference_steps)
-    if steps < 1:
-        raise ValueError(f"num_inference_steps must be >= 1, got {steps}")
-    tail = int(tail_steps or 0)
-    if tail < 0:
-        raise ValueError(f"refiner_sigma_tail_steps must be >= 0, got {tail}")
-
-    base = np.linspace(float(sigma_max), float(sigma_min), steps + 1).copy()[:-1]
-    shift_value = float(shift)
-    shifted = shift_value * base / (1.0 + (shift_value - 1.0) * base)
-    eps = 1e-6
-    sigmas = shifted[shifted <= t_value + eps]
-    if sigmas.size == 0 or abs(float(sigmas[0]) - t_value) > eps:
-        sigmas = np.concatenate([[t_value], sigmas])
-    if tail > 0:
-        start = float(sigmas[-1])
-        stop = min(float(sigma_min), start)
-        extra = np.linspace(start, stop, tail + 2, dtype=np.float64)[1:-1]
-        sigmas = np.concatenate([sigmas, extra])
-    return _validate_refiner_sigmas(sigmas, t_value).astype(np.float32)
-
-
-def _pad_prompt_embeds(
-    embeds: torch.Tensor,
-    mask: torch.Tensor,
-    target_length: int,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    if embeds.shape[0] != 1:
-        raise ValueError(f"batched CFG helper expects batch=1 inputs, got {embeds.shape[0]}")
-    if embeds.shape[1] > target_length:
-        raise ValueError(f"cannot pad length {embeds.shape[1]} down to {target_length}")
-    pad_len = target_length - embeds.shape[1]
-    if pad_len == 0:
-        return embeds, mask
-    embed_pad = torch.zeros(embeds.shape[0], pad_len, embeds.shape[2], dtype=embeds.dtype, device=embeds.device)
-    mask_pad = torch.zeros(mask.shape[0], pad_len, dtype=mask.dtype, device=mask.device)
-    return torch.cat([embeds, embed_pad], dim=1), torch.cat([mask, mask_pad], dim=1)
-
-
-def _batch_cfg_prompt_inputs(
-    prompt_embeds: torch.Tensor,
-    prompt_mask: torch.Tensor,
-    negative_embeds: torch.Tensor,
-    negative_mask: torch.Tensor,
-    *,
-    null_cond_clone_zero: bool = False,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    if null_cond_clone_zero:
-        zero_negative = torch.zeros_like(prompt_embeds)
-        return (
-            torch.cat([prompt_embeds, zero_negative], dim=0),
-            torch.cat([prompt_mask, prompt_mask.clone()], dim=0),
-        )
-
-    target_length = max(int(prompt_embeds.shape[1]), int(negative_embeds.shape[1]))
-    prompt_padded, prompt_mask_padded = _pad_prompt_embeds(prompt_embeds, prompt_mask, target_length)
-    negative_padded, negative_mask_padded = _pad_prompt_embeds(negative_embeds, negative_mask, target_length)
-    return (
-        torch.cat([prompt_padded, negative_padded], dim=0),
-        torch.cat([prompt_mask_padded, negative_mask_padded], dim=0),
-    )
-
-
-def get_lingbot_video_post_process_func(od_config: OmniDiffusionConfig):
-    del od_config
-
-    def post_process_func(frames: torch.Tensor, sampling_params=None):
-        output_type = getattr(sampling_params, "output_type", None) or "pt"
-        if output_type == "np" and isinstance(frames, torch.Tensor):
-            return frames.float().cpu().numpy()
-        return frames
-
-    return post_process_func
-
-
-class LingBotVideoPipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery):
-    """Native vLLM-Omni entry for the dense LingBot-Video checkpoint.
-
-    This PR supports the dense checkpoint only. MoE and fused-expert kernels are
-    intentionally left for the follow-up MoE PR.
-    """
-
-    supports_step_execution: ClassVar[bool] = False
-    _dit_modules: ClassVar[list[str]] = ["transformer"]
-    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
-    _vae_modules: ClassVar[list[str]] = ["vae"]
-
-    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
-        super().__init__()
-        del prefix
-        self.od_config = od_config
-        self.device = get_local_device()
-        self.vae_scale_factor_temporal = 4
-        self.vae_scale_factor_spatial = 8
-        self.token_length = TOKEN_LENGTH
-        self.hidden_state_skip_layer = HIDDEN_STATE_SKIP_LAYER
-        self.prompt_template = PROMPT_TEMPLATE
-        self._crop_start: int | None = None
-
-        model = od_config.model
-        local_files_only = os.path.exists(model)
-        dtype = getattr(od_config, "dtype", torch.bfloat16)
-        model_config = getattr(od_config, "model_config", None) or {}
-        transformer_dtype = _dtype_from_name(model_config.get("transformer_dtype"), dtype)
-        text_encoder_dtype = _dtype_from_name(model_config.get("text_encoder_dtype"), dtype)
-        vae_dtype = _dtype_from_name(model_config.get("vae_dtype"), torch.float32)
-
-        transformer_subfolder = str(model_config.get("transformer_subfolder", "transformer"))
-        text_encoder_subfolder = str(model_config.get("text_encoder_subfolder", "text_encoder"))
-        processor_subfolder = str(model_config.get("processor_subfolder", "processor"))
-        vae_subfolder = str(model_config.get("vae_subfolder", "vae"))
-        scheduler_subfolder = str(model_config.get("scheduler_subfolder", "scheduler"))
-
-        self.transformer = LingBotVideoTransformer3DModel.from_pretrained(
-            model,
-            subfolder=transformer_subfolder,
-            torch_dtype=transformer_dtype,
-            local_files_only=local_files_only,
-        ).to(self.device)
-        text_encoder_kwargs: dict[str, Any] = {
-            "dtype": text_encoder_dtype,
-            "local_files_only": local_files_only,
-        }
-        self.text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
-            model,
-            subfolder=text_encoder_subfolder,
-            **text_encoder_kwargs,
-        ).to(self.device)
-        self.processor = Qwen3VLProcessor.from_pretrained(
-            model,
-            subfolder=processor_subfolder,
-            local_files_only=local_files_only,
-        )
-        self.vae = AutoencoderKLWan.from_pretrained(
-            model,
-            subfolder=vae_subfolder,
-            torch_dtype=vae_dtype,
-            local_files_only=local_files_only,
-        ).to(self.device)
-        self.scheduler = FlowUniPCMultistepScheduler.from_pretrained(
-            model,
-            subfolder=scheduler_subfolder,
-            local_files_only=local_files_only,
-        )
-        self.set_progress_bar_config(disable=bool(model_config.get("quiet_progress", True)))
-        self.default_negative_prompt = DEFAULT_NEGATIVE_PROMPT
-
-    def to(self, *args, **kwargs):
-        device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
-        super().to(*args, **kwargs)
-        if device is not None:
-            self.device = torch.device(device)
-            self.transformer.to(device=self.device, non_blocking=non_blocking)
-            self.text_encoder.to(device=self.device, non_blocking=non_blocking)
-            self.vae.to(device=self.device, non_blocking=non_blocking)
-        if dtype is not None:
-            self.transformer.to(dtype=dtype, non_blocking=non_blocking)
-        return self
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        consumed = list(weights)
-        if consumed:
-            raise RuntimeError(
-                f"{self.__class__.__name__}.load_weights received {len(consumed)} weight tensors; "
-                "LingBot components are loaded directly from subfolders during __init__."
-            )
-        return {name for name, _ in self.named_parameters()}
-
-    @staticmethod
-    def check_inputs(height: int, width: int, num_frames: int) -> None:
-        if num_frames != 1 and (num_frames - 1) % 4 != 0:
-            raise ValueError(f"`num_frames` must be 1 or 4n+1, got {num_frames}.")
-        if height % 16 != 0 or width % 16 != 0:
-            raise ValueError(f"`height` and `width` must be multiples of 16, got {height}x{width}.")
-
-    @staticmethod
-    def apply_text_to_template(text: str, template: str = PROMPT_TEMPLATE) -> str:
-        return template.format(text)
-
-    def _compute_crop_start(self) -> int:
-        if self._crop_start is None:
-            marker = "<|USER_INPUT_MARKER|>"
-            marked = self.prompt_template.format(marker)
-            marker_pos = marked.find(marker)
-            if marker_pos < 0:
-                self._crop_start = 0
-            else:
-                prefix = self.processor(
-                    text=marked[:marker_pos],
-                    images=None,
-                    videos=None,
-                    return_tensors="pt",
-                )
-                self._crop_start = int(prefix["input_ids"].shape[1])
-        return self._crop_start
-
-    def _build_prompt_inputs(self, prompt: str | list[str]):
-        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
-        texts = [self.apply_text_to_template(text, self.prompt_template) for text in prompts]
-        return self.processor(
-            text=texts,
-            images=None,
-            videos=None,
-            do_resize=False,
-            truncation=True,
-            max_length=self.token_length,
-            padding="longest",
-            return_tensors="pt",
-        )
-
-    @torch.no_grad()
-    def encode_prompt(
-        self,
-        prompt: str | list[str],
-        *,
-        device: str | torch.device | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        device = torch.device(device) if device is not None else self.device
-        inputs = self._build_prompt_inputs(prompt).to(device)
-        outputs = self.text_encoder(
-            **inputs,
-            output_hidden_states=self.hidden_state_skip_layer is not None,
-        )
-        if self.hidden_state_skip_layer is not None:
-            prompt_embeds = outputs.hidden_states[-(self.hidden_state_skip_layer + 1)]
-        else:
-            prompt_embeds = outputs.last_hidden_state
-
-        prompt_mask = inputs["attention_mask"]
-        crop_start = self._compute_crop_start()
-        if crop_start > 0:
-            prompt_embeds = prompt_embeds[:, crop_start:]
-            prompt_mask = prompt_mask[:, crop_start:]
-
-        if prompt_embeds.shape[0] == 1:
-            true_len = int(prompt_mask[0].sum().item())
-            prompt_embeds = prompt_embeds[:, :true_len]
-            prompt_mask = prompt_mask[:, :true_len]
-        return prompt_embeds, prompt_mask
-
-    def prepare_latents(
-        self,
-        num_frames: int,
-        height: int,
-        width: int,
-        generator: torch.Generator | None,
-        latents: torch.Tensor | None,
-        device: torch.device,
-    ) -> torch.Tensor:
-        latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
-        latent_height = height // self.vae_scale_factor_spatial
-        latent_width = width // self.vae_scale_factor_spatial
-        shape = (1, self.transformer.config.in_channels, latent_frames, latent_height, latent_width)
-        if latents is None:
-            return randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
-        if tuple(latents.shape) != shape:
-            raise ValueError(f"`latents` shape must be {shape}, got {tuple(latents.shape)}.")
-        return latents.to(device=device, dtype=torch.float32)
-
-    def _dit_latent_to_vae(self, latents: torch.Tensor) -> torch.Tensor:
-        mean = torch.tensor(self.vae.config.latents_mean, device=latents.device, dtype=torch.float32)
-        std_inv = 1.0 / torch.tensor(self.vae.config.latents_std, device=latents.device, dtype=torch.float32)
-        mean = mean.view(1, -1, 1, 1, 1)
-        std_inv = std_inv.view(1, -1, 1, 1, 1)
-        return latents.float() / std_inv + mean
-
-    @torch.no_grad()
-    def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
-        vae_device = _module_device(self.vae)
-        vae_dtype = _module_dtype(self.vae)
-        vae_latents = self._dit_latent_to_vae(latents).to(device=vae_device, dtype=torch.float32)
-        if vae_latents.ndim == 5:
-            vae_latents = vae_latents.contiguous(memory_format=torch.channels_last_3d)
-        autocast_dtype = (
-            vae_dtype if vae_device.type == "cuda" and vae_dtype in {torch.bfloat16, torch.float16} else None
-        )
-        with torch.autocast("cuda", dtype=autocast_dtype or torch.bfloat16, enabled=autocast_dtype is not None):
-            decoded = self.vae.decode(vae_latents)
-        frames = decoded[0] if isinstance(decoded, tuple) else decoded.sample
-        frames = frames.float().clamp_(-1, 1)
-        frames = (frames + 1.0) / 2.0
-        frames = frames.permute(0, 2, 3, 4, 1).cpu()
-        return frames[0]
-
-    @torch.no_grad()
-    def _generate(
-        self,
-        *,
-        prompt: str,
-        negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
-        height: int = 480,
-        width: int = 480,
-        num_frames: int = 81,
-        num_inference_steps: int = 40,
-        guidance_scale: float = 6.0,
-        shift: float = 3.0,
-        generator: torch.Generator | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        prompt_mask: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_mask: torch.Tensor | None = None,
-        output_type: str = "pt",
-        cfg_parallel_group: Any | None = None,
-        batch_cfg: bool = False,
-        null_cond_clone_zero: bool = False,
-        t_thresh: float | None = None,
-        refiner_sigma_tail_steps: int = LOW_NOISE_TAIL_V1_DEFAULT_STEPS,
-        offload_vae_during_denoise: bool = False,
-        **extra_args,
-    ) -> torch.Tensor:
-        del extra_args
-        self.check_inputs(height, width, num_frames)
-        device = self.device
-        do_cfg = guidance_scale > 1.0
-        effective_batch_cfg = bool(batch_cfg)
-        cfg_parallel = cfg_parallel_group is not None
-        cfg_parallel_rank = 0
-        if cfg_parallel:
-            if not dist.is_available() or not dist.is_initialized():
-                raise ValueError("`cfg_parallel_group` requires an initialized process group.")
-            if effective_batch_cfg:
-                raise ValueError("`cfg_parallel_group` and `batch_cfg` are mutually exclusive.")
-            if not do_cfg:
-                raise ValueError("CFG parallel requires `guidance_scale > 1.0`.")
-            cfg_parallel_rank = dist.get_rank(cfg_parallel_group)
-            cfg_parallel_world_size = dist.get_world_size(cfg_parallel_group)
-            if cfg_parallel_world_size != 2:
-                raise ValueError(f"CFG parallel currently requires exactly 2 ranks, got {cfg_parallel_world_size}.")
-
-        if prompt_embeds is not None:
-            if prompt_mask is None:
-                raise ValueError("`prompt_mask` is required when `prompt_embeds` is provided.")
-            prompt_embeds = prompt_embeds.to(device=device)
-            prompt_mask = prompt_mask.to(device=device)
-        if negative_prompt_embeds is not None:
-            if negative_prompt_mask is None:
-                raise ValueError("`negative_prompt_mask` is required when `negative_prompt_embeds` is provided.")
-            negative_prompt_embeds = negative_prompt_embeds.to(device=device)
-            negative_prompt_mask = negative_prompt_mask.to(device=device)
-
-        negative_embeds = None
-        negative_mask = None
-        if cfg_parallel and cfg_parallel_rank == 1:
-            if negative_prompt_embeds is not None:
-                negative_embeds, negative_mask = negative_prompt_embeds, negative_prompt_mask
-            else:
-                negative_embeds, negative_mask = self.encode_prompt(negative_prompt, device=device)
-            prompt_embeds = prompt_mask = None
-        else:
-            if prompt_embeds is None:
-                prompt_embeds, prompt_mask = self.encode_prompt(prompt, device=device)
-            if do_cfg and not cfg_parallel:
-                if null_cond_clone_zero:
-                    negative_embeds = torch.zeros_like(prompt_embeds)
-                    negative_mask = prompt_mask.clone()
-                elif negative_prompt_embeds is not None:
-                    negative_embeds, negative_mask = negative_prompt_embeds, negative_prompt_mask
-                else:
-                    negative_embeds, negative_mask = self.encode_prompt(negative_prompt, device=device)
-
-        latents = self.prepare_latents(num_frames, height, width, generator, latents, device)
-        sigmas = _compute_refiner_sigmas(
-            sigma_max=float(self.scheduler.sigma_max),
-            sigma_min=float(self.scheduler.sigma_min),
-            num_inference_steps=num_inference_steps,
-            shift=shift,
-            t_thresh=t_thresh,
-            tail_steps=refiner_sigma_tail_steps,
-        )
-        if sigmas is None:
-            self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
-        else:
-            self.scheduler.set_timesteps(int(sigmas.shape[0]), device=device, sigmas=sigmas, shift=1.0)
-
-        transformer_dtype = _module_dtype(self.transformer)
-        vae_restore_device: torch.device | None = None
-        vae_offloaded = False
-        if offload_vae_during_denoise and output_type != "latent":
-            vae_device = _module_device(self.vae)
-            if vae_device.type == "cuda":
-                self.vae.to("cpu")
-                torch.accelerator.empty_cache()
-                vae_restore_device = vae_device
-                vae_offloaded = True
-
-        cfg_latent_src = _group_global_rank(cfg_parallel_group, 0)
-        cfg_uncond_src = _group_global_rank(cfg_parallel_group, 1)
-        for timestep in self.progress_bar(self.scheduler.timesteps):
-            if cfg_parallel:
-                dist.broadcast(latents, src=cfg_latent_src, group=cfg_parallel_group)
-            timestep_batch = _transformer_timestep(timestep, transformer_dtype).expand(1).to(device)
-            latent_model_input = latents
-            if cfg_parallel:
-                if cfg_parallel_rank == 0:
-                    branch_embeds = prompt_embeds
-                    branch_mask = prompt_mask
-                else:
-                    branch_embeds = negative_embeds
-                    branch_mask = negative_mask
-                if branch_embeds is None:
-                    raise RuntimeError("CFG branch embeddings were not initialized.")
-                with _transformer_autocast(device, transformer_dtype):
-                    branch_noise_pred = self.transformer(
-                        latent_model_input,
-                        timestep_batch,
-                        branch_embeds.to(transformer_dtype),
-                        encoder_attention_mask=branch_mask,
-                        return_dict=False,
-                    )[0].float()
-                if cfg_parallel_rank == 0:
-                    noise_pred = branch_noise_pred
-                    noise_pred_uncond = torch.empty_like(noise_pred)
-                else:
-                    noise_pred_uncond = branch_noise_pred
-                dist.broadcast(noise_pred_uncond, src=cfg_uncond_src, group=cfg_parallel_group)
-                if cfg_parallel_rank != 0:
-                    continue
-                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
-            else:
-                if prompt_embeds is None:
-                    raise RuntimeError("Prompt embeddings were not initialized.")
-                prompt_model_input = prompt_embeds.to(transformer_dtype)
-                if do_cfg and effective_batch_cfg:
-                    if negative_embeds is None or negative_mask is None:
-                        raise RuntimeError("Negative embeddings were not initialized for CFG.")
-                    cfg_embeds, cfg_mask = _batch_cfg_prompt_inputs(
-                        prompt_model_input,
-                        prompt_mask,
-                        negative_embeds.to(transformer_dtype),
-                        negative_mask,
-                        null_cond_clone_zero=False,
-                    )
-                    cfg_latents = torch.cat([latent_model_input, latent_model_input], dim=0)
-                    cfg_timesteps = torch.cat([timestep_batch, timestep_batch], dim=0)
-                    with _transformer_autocast(device, transformer_dtype):
-                        noise_batched = self.transformer(
-                            cfg_latents,
-                            cfg_timesteps,
-                            cfg_embeds,
-                            encoder_attention_mask=cfg_mask,
-                            return_dict=False,
-                        )[0].float()
-                    noise_pred, noise_pred_uncond = noise_batched.chunk(2, dim=0)
-                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
-                else:
-                    with _transformer_autocast(device, transformer_dtype):
-                        noise_pred = self.transformer(
-                            latent_model_input,
-                            timestep_batch,
-                            prompt_model_input,
-                            encoder_attention_mask=prompt_mask,
-                            return_dict=False,
-                        )[0].float()
-
-                if do_cfg and not effective_batch_cfg:
-                    if negative_embeds is None or negative_mask is None:
-                        raise RuntimeError("Negative embeddings were not initialized for CFG.")
-                    with _transformer_autocast(device, transformer_dtype):
-                        noise_pred_uncond = self.transformer(
-                            latent_model_input,
-                            timestep_batch,
-                            negative_embeds.to(transformer_dtype),
-                            encoder_attention_mask=negative_mask,
-                            return_dict=False,
-                        )[0].float()
-                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
-
-            latents = self.scheduler.step(noise_pred, timestep, latents, return_dict=False, generator=generator)[0]
-
-        if cfg_parallel:
-            dist.barrier(group=cfg_parallel_group)
-            if cfg_parallel_rank != 0:
-                return latents if output_type == "latent" else []
-
-        if output_type == "latent":
-            return latents
-        if output_type in {"pt", "np"}:
-            if vae_offloaded and vae_restore_device is not None:
-                self.vae.to(device=vae_restore_device)
-                torch.accelerator.empty_cache()
-            return self._decode_latents(latents)
-        raise ValueError(f"Unsupported output_type: {output_type}")
-
-    @torch.inference_mode()
-    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
-        if req.num_reqs != 1:
-            raise ValueError(f"LingBotVideoPipeline only supports one request per batch, got {req.num_reqs}.")
-        request = req.requests[0]
-        prompt, prompt_negative = _extract_prompt(request)
-        sampling = request.sampling_params
-        extra_args = dict(sampling.extra_args or {})
-
-        generator = sampling.generator
-        if isinstance(generator, list):
-            generator = generator[0] if generator else None
-        if generator is None and sampling.seed is not None:
-            generator = torch.Generator(device=self.device).manual_seed(int(sampling.seed))
-
-        height = sampling.height if sampling.height is not None else extra_args.pop("height", 480)
-        width = sampling.width if sampling.width is not None else extra_args.pop("width", 480)
-        num_frames = sampling.num_frames or extra_args.pop("num_frames", 81)
-        num_inference_steps = (
-            sampling.num_inference_steps
-            if sampling.num_inference_steps is not None
-            else extra_args.pop("num_inference_steps", 40)
-        )
-        guidance_scale = (
-            sampling.guidance_scale
-            if sampling.guidance_scale_provided or sampling.guidance_scale > 0
-            else extra_args.pop("guidance_scale", 6.0)
-        )
-        shift = extra_args.pop(
-            "shift",
-            extra_args.pop("flow_shift", getattr(self.od_config, "flow_shift", None) or 3.0),
-        )
-        negative_prompt = extra_args.pop("negative_prompt", prompt_negative or self.default_negative_prompt)
-        output_type = (
-            sampling.output_type or getattr(self.od_config, "output_type", None) or extra_args.pop("output_type", "pt")
-        )
-        if output_type not in {"pt", "np", "latent"}:
-            output_type = "pt"
-        sampling.output_type = output_type
-
-        frames = self._generate(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            height=height,
-            width=width,
-            num_frames=num_frames,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            shift=shift,
-            generator=generator,
-            latents=sampling.latents,
-            output_type=output_type,
-            batch_cfg=bool(extra_args.pop("batch_cfg", False)),
-            null_cond_clone_zero=bool(extra_args.pop("null_cond_clone_zero", False)),
-            offload_vae_during_denoise=bool(extra_args.pop("offload_vae_during_denoise", False)),
-            t_thresh=extra_args.pop("t_thresh", None),
-            refiner_sigma_tail_steps=int(extra_args.pop("refiner_sigma_tail_steps", LOW_NOISE_TAIL_V1_DEFAULT_STEPS)),
-        )
-        return DiffusionOutput(output=frames)

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -9,7 +9,6 @@ from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
 from vllm.config import VllmConfig
-from vllm.logger import init_logger
 from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
 
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
@@ -20,8 +19,6 @@ from .mammothmoda2_dit_model import SimpleQFormerImageRefiner, Transformer2DMode
 from .rope_real import RotaryPosEmbedReal
 from .schedulers import FlowMatchEulerDiscreteScheduler
 
-logger = init_logger(__name__)
-
 
 class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
     """
@@ -63,14 +60,14 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
         self.gen_transformer = Transformer2DModel.from_config(self.config.gen_dit_config)
 
         # llm_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config
-        llm_hidden_size = 0
-        text_config = self.config.get_text_config()
-        if text_config is None:
-            logger.warning("No text config; failed to infer llm_hidden_size.")
-        elif not hasattr(text_config, "hidden_size"):
-            logger.warning("Text config exists, but has no hidden_size attribute; failed to infer llm_hidden_size.")
-        else:
-            llm_hidden_size = int(text_config.hidden_size or 0)
+        llm_hidden_size = int(
+            getattr(
+                getattr(self.config.llm_config, "text_config", None),
+                "hidden_size",
+                0,
+            )
+            or 0
+        )
         if llm_hidden_size <= 0:
             raise ValueError(
                 "Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.text_config.hidden_size"

vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py

@@ -22,10 +22,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.distributed.sp_plan import (
-    SequenceParallelInput,
-    SequenceParallelOutput,
-)
 from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
@@ -812,36 +808,6 @@ class OmniGen2TransformerBlock(nn.Module):
         return hidden_states
 
 
-class _OmniGen2SPInputBoundary(nn.Module):
-    """SP input-sharding boundary marker for ``OmniGen2Transformer2DModel._sp_plan``.
-
-    Resets ``sp_original_seq_len`` on every call: OmniGen2 runs the transformer
-    once per CFG branch and the unconditional / conditional branches have
-    different joint sequence lengths (text prefix differs). The SP framework's
-    auto-pad records the length only on the first call, so without this reset
-    a stale length from an earlier branch truncates the gather of later ones.
-    """
-
-    def forward(self, hidden_states: torch.Tensor, rotary_emb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        from vllm_omni.diffusion.forward_context import (
-            get_forward_context,
-            is_forward_context_available,
-        )
-
-        if is_forward_context_available():
-            ctx = get_forward_context()
-            ctx.sp_original_seq_len = None
-            ctx.sp_padding_size = 0
-        return hidden_states, rotary_emb
-
-
-class _OmniGen2SPOutputBoundary(nn.Module):
-    """SP output-gathering boundary marker for ``OmniGen2Transformer2DModel._sp_plan``."""
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        return hidden_states
-
-
 class OmniGen2Transformer2DModel(nn.Module):
     """
     OmniGen2 Transformer 2D Model.
@@ -870,17 +836,6 @@ class OmniGen2Transformer2DModel(nn.Module):
         timestep_scale: Scale factor for timestep embeddings
     """
 
-    # ``attention_mask`` is intentionally omitted from the plan: OmniGen2Attention
-    # accepts but never reads it. ``auto_pad=True`` requires Ulysses SP — Ring SP
-    # cannot use attention masks.
-    _sp_plan = {
-        "sp_input_boundary": {
-            0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
-            1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
-        },
-        "sp_output_boundary": SequenceParallelOutput(gather_dim=1, expected_dims=3),
-    }
-
     def __init__(
         self,
         patch_size: int = 2,
@@ -1028,9 +983,6 @@ class OmniGen2Transformer2DModel(nn.Module):
         # Add learnable embeddings to distinguish different images
         self.image_index_embedding = nn.Parameter(torch.randn(5, hidden_size))  # support max 5 ref images
 
-        self.sp_input_boundary = _OmniGen2SPInputBoundary()
-        self.sp_output_boundary = _OmniGen2SPOutputBoundary()
-
     def img_patch_embed_and_refine(
         self,
         hidden_states,
@@ -1299,13 +1251,9 @@ class OmniGen2Transformer2DModel(nn.Module):
 
         hidden_states = joint_hidden_states
 
-        hidden_states, rotary_emb = self.sp_input_boundary(hidden_states, rotary_emb)
-
         for layer in self.layers:
             hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb)
 
-        hidden_states = self.sp_output_boundary(hidden_states)
-
         # 4. Output norm & projection
         hidden_states = self.norm_out(hidden_states, temb)
 

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -17,7 +17,6 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous
 from vllm.logger import init_logger
-from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
@@ -555,8 +554,8 @@ class QwenImageCrossAttention(nn.Module):
         self.query_num_heads = self.to_qkv.num_heads
         self.kv_num_heads = self.to_qkv.num_kv_heads
 
-        self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_q = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
 
         self.inner_dim = out_dim if out_dim is not None else head_dim * self.total_num_heads
 
@@ -593,8 +592,8 @@ class QwenImageCrossAttention(nn.Module):
             prefix=f"{prefix}.to_out.0",
         )
 
-        self.norm_added_q = RMSNorm(head_dim, eps=eps)
-        self.norm_added_k = RMSNorm(head_dim, eps=eps)
+        self.norm_added_q = nn.RMSNorm(head_dim, eps=eps)
+        self.norm_added_k = nn.RMSNorm(head_dim, eps=eps)
 
         self.attn = Attention(
             num_heads=self.query_num_heads,
@@ -1007,7 +1006,7 @@ class QwenImageTransformer2DModel(CachedTransformer):
             quant_config=quant_config,
         )
 
-        self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
+        self.txt_norm = nn.RMSNorm(joint_attention_dim, eps=1e-6)
 
         # Entry projections (image/text) are kept full precision —
         # small sensitive layers at the network boundary (see #2728).

vllm_omni/diffusion/registry.py

@@ -256,11 +256,6 @@ _DIFFUSION_MODELS = {
         "pipeline_hunyuan_video_1_5_i2v",
         "HunyuanVideo15I2VPipeline",
     ),
-    "LingBotVideoPipeline": (
-        "lingbot_video",
-        "pipeline_lingbot_video",
-        "LingBotVideoPipeline",
-    ),
     "MagiHumanPipeline": (
         "magi_human",
         "pipeline_magi_human",
@@ -531,7 +526,6 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "Flux2Pipeline": "get_flux2_post_process_func",
     "HunyuanVideo15Pipeline": "get_hunyuan_video_15_post_process_func",
     "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_post_process_func",
-    "LingBotVideoPipeline": "get_lingbot_video_post_process_func",
     "MagiHumanPipeline": "get_magi_human_post_process_func",
     "OmniVoicePipeline": "get_omnivoice_post_process_func",
     "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func",

vllm_omni/engine/async_omni_engine.py

@@ -30,7 +30,6 @@ from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.input_processor import InputProcessor
 
-from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.stage_config import strip_parent_engine_args
 from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config
 from vllm_omni.diffusion.diffusion_engine import supports_audio_output
@@ -217,7 +216,6 @@ class AsyncOmniEngine:
         transfer_emitter: Any = None,
         log_stats: bool = False,
         tokenizer: str | None = None,
-        trust_remote_code: bool = False,
         **kwargs: Any,
     ) -> None:
         self.model = model
@@ -275,17 +273,6 @@ class AsyncOmniEngine:
                 self._omni_master_port,
             )
 
-        # Stage resolution pops deploy_config, so get the pipeline endpoint
-        # restriction beforehand. TODO (Alex) make this cleaner and refactor
-        # stage config resolution to remove kwargs hacks.
-        deploy_config_path = kwargs.get("deploy_config")
-        self.endpoint_restrictions = StageConfigFactory.get_pipeline_endpoint_restrictions(
-            model=model,
-            trust_remote_code=trust_remote_code,
-            deploy_config_path=deploy_config_path,
-        )
-
-        kwargs["trust_remote_code"] = trust_remote_code
         self.config_path, self.stage_configs = self._resolve_stage_configs(model, kwargs)
 
         self.num_stages = len(self.stage_configs)

vllm_omni/entrypoints/async_omni.py

@@ -148,7 +148,6 @@ class AsyncOmni(EngineClient, OmniBase):
         self.config_path = self.engine.config_path
         self.tts_max_instructions_length = kwargs.get("tts_max_instructions_length", None)
         self.input_processor = self.engine.input_processor
-        self.endpoint_restrictions = self.engine.endpoint_restrictions
 
         stage_index = self._get_comprehension_stage_index()
         if stage_index is None:

vllm_omni/entrypoints/openai/api_server.py

@@ -91,7 +91,6 @@ from vllm.utils import random_uuid
 from vllm.utils.system_utils import decorate_logs
 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
-from vllm_omni.config.endpoint_policy import shutdown_unsupported_routes
 from vllm_omni.diffusion.models.interface import ReferenceVideoDecodeSpec
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
@@ -260,7 +259,7 @@ async def _get_vllm_config(engine_client: EngineClient) -> Any:
     return getattr(engine_client, "vllm_config", None)
 
 
-def _remove_route_from_app(app, path: str, methods: frozenset[str] | None = None):
+def _remove_route_from_app(app, path: str, methods: set[str] | None = None):
     """Remove a route from the app by path and optionally by methods.
 
     OMNI: used to override upstream /v1/chat/completions with omni behavior.
@@ -510,12 +509,6 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
 
         await omni_init_app_state(engine_client, app.state, args)
 
-        # After initializing the app state, shut down any endpoints that are model specific
-        if hasattr(engine_client, "endpoint_restrictions"):
-            shutdown_unsupported_routes(app, engine_client.endpoint_restrictions)
-        else:
-            logger.warning("engine client has no endpoint restrictions attribute")
-
         # Start background processes
         await STORAGE_MANAGER.start()
 

vllm_omni/entrypoints/utils.py

@@ -345,7 +345,6 @@ def load_stage_configs_from_model(
 
     stages = StageConfigFactory.create_from_model(
         model,
-        trust_remote_code=cli_overrides.get("trust_remote_code", False),
         cli_overrides=cli_overrides,
         deploy_config_path=deploy_config_path,
     )

vllm_omni/model_executor/models/mammoth_moda2/pipeline.py

@@ -1,66 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""MammothModa2 pipeline topology (frozen).
-
-Stage 0: AR  — multimodal understanding + latent generation
-Stage 1: DiT — latent → image
-
-For text/image understanding tasks (text output only), use MAMMOTH_MODA2_AR_PIPELINE.
-"""
-
-from vllm_omni.config.stage_config import (
-    PipelineConfig,
-    StageExecutionType,
-    StagePipelineConfig,
-)
-
-_PROC = "vllm_omni.model_executor.stage_input_processors.mammoth_moda2"
-
-MAMMOTH_MODA2_PIPELINE = PipelineConfig(
-    model_type="mammoth_moda2",
-    model_arch="MammothModa2ForConditionalGeneration",
-    hf_architectures=("Mammothmoda2Model", "MammothModa2ForConditionalGeneration"),
-    stages=(
-        StagePipelineConfig(
-            stage_id=0,
-            model_stage="ar",
-            execution_type=StageExecutionType.LLM_AR,
-            input_sources=(),
-            final_output=False,
-            owns_tokenizer=True,
-            requires_multimodal_data=True,
-            engine_output_type="latent",
-        ),
-        StagePipelineConfig(
-            stage_id=1,
-            model_stage="dit",
-            execution_type=StageExecutionType.LLM_GENERATION,
-            input_sources=(0,),
-            final_output=True,
-            final_output_type="image",
-            owns_tokenizer=False,
-            requires_multimodal_data=False,
-            engine_output_type="image",
-            custom_process_input_func=f"{_PROC}.ar2dit",
-        ),
-    ),
-)
-
-# Single-stage AR variant for understanding / summarization tasks
-MAMMOTH_MODA2_AR_PIPELINE = PipelineConfig(
-    model_type="mammoth_moda2_ar",
-    model_arch="MammothModa2ForConditionalGeneration",
-    stages=(
-        StagePipelineConfig(
-            stage_id=0,
-            model_stage="ar",
-            execution_type=StageExecutionType.LLM_AR,
-            input_sources=(),
-            final_output=True,
-            final_output_type="text",
-            owns_tokenizer=True,
-            requires_multimodal_data=True,
-            engine_output_type="text",
-        ),
-    ),
-)

vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py

@@ -1,6 +1,6 @@
 """CUDA Graph acceleration for the MOSS Audio Tokenizer codec decoder.
 
-Captures the codec decode method (_decode_frame, or _decode on the vendored codec) for a set of fixed frame-count
+Captures MossAudioTokenizerModel._decode for a set of fixed frame-count
 bucket sizes, then replays the captured graph at inference time to eliminate
 kernel-launch overhead.  Inputs that exceed all captured sizes fall back to
 eager execution transparently.
@@ -24,7 +24,7 @@ logger = init_logger(__name__)
 
 
 class MossTTSCUDAGraphCodecWrapper:
-    """CUDA Graph wrapper for the codec decode method (_decode_frame, or _decode on the vendored codec).
+    """CUDA Graph wrapper for MossAudioTokenizerModel._decode.
 
     Graphs are keyed by padded_T (int).  On each call the actual T is
     bucket-matched to the smallest pre-captured size >= T.  The static code
@@ -51,15 +51,6 @@ class MossTTSCUDAGraphCodecWrapper:
         enabled: bool = True,
     ) -> None:
         self.model = model
-        # Upstream remote-code tokenizer exposes _decode_frame; the vendored
-        # MossAudioTokenizerModel still exposes _decode. Resolve once here so
-        # both codec implementations work.
-        decode_fn = getattr(model, "_decode_frame", None) or getattr(model, "_decode", None)
-        if decode_fn is None:
-            raise AttributeError(
-                f"{type(model).__module__}.{type(model).__name__} exposes neither _decode_frame nor _decode"
-            )
-        self._decode_fn = decode_fn
         self.capture_sizes: list[int] = sorted(capture_sizes)
         self.num_quantizers = num_quantizers
         self.enabled = enabled
@@ -108,7 +99,7 @@ class MossTTSCUDAGraphCodecWrapper:
             dummy_codes = torch.zeros(nq, 1, size, dtype=torch.long, device=device)
             dummy_lengths = torch.tensor([size], dtype=torch.long, device=device)
             with torch.no_grad():
-                _ = self._decode_fn(dummy_codes, dummy_lengths)
+                _ = self.model._decode(dummy_codes, dummy_lengths)
 
         torch.accelerator.synchronize(device)
 
@@ -141,13 +132,13 @@ class MossTTSCUDAGraphCodecWrapper:
 
         # Extra eager warmup inside capture to ensure all kernels are compiled.
         with torch.no_grad():
-            _ = self._decode_fn(static_codes, static_lengths)
+            _ = self.model._decode(static_codes, static_lengths)
         torch.accelerator.synchronize(device)
 
         graph = CUDAGraph()
         with torch.no_grad():
             with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
-                static_out = self._decode_fn(static_codes, static_lengths)
+                static_out = self.model._decode(static_codes, static_lengths)
 
         self.graphs[size] = graph
         self.static_codes[size] = static_codes

vllm_omni/model_executor/models/omnivoice/pipeline.py

@@ -1,31 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""OmniVoice pipeline topology (frozen).
-
-Single-stage diffusion model for text-to-speech.
-Stage 0: DiT — text → iterative unmasking → codebook tokens → DAC decode → 24kHz audio
-"""
-
-from vllm_omni.config.stage_config import (
-    PipelineConfig,
-    StageExecutionType,
-    StagePipelineConfig,
-)
-
-OMNIVOICE_PIPELINE = PipelineConfig(
-    model_type="omnivoice",
-    model_arch="OmniVoicePipeline",
-    stages=(
-        StagePipelineConfig(
-            stage_id=0,
-            model_stage="dit",
-            execution_type=StageExecutionType.DIFFUSION,
-            input_sources=(),
-            final_output=True,
-            final_output_type="audio",
-            owns_tokenizer=False,
-            requires_multimodal_data=False,
-            engine_output_type="audio",
-        ),
-    ),
-)

vllm_omni/model_executor/models/qwen3_omni/pipeline.py

@@ -9,7 +9,6 @@ Stage 2: Code2Wav — RVQ codes → audio waveform
 
 from transformers import Qwen3OmniMoeConfig
 
-from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
 from vllm_omni.config.stage_config import (
     PipelineConfig,
     StageExecutionType,
@@ -22,12 +21,6 @@ _PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_omni"
 QWEN3_OMNI_PIPELINE = PipelineConfig(
     model_type="qwen3_omni_moe",
     model_arch="Qwen3OmniMoeForConditionalGeneration",
-    endpoint_restrictions=(
-        EndpointRestriction(
-            OmniServingCapability.COMPLETIONS,
-            "Qwen3-Omni requires chat template structure for thinker-talker handoff. Use /v1/chat/completions instead.",
-        ),
-    ),
     stages=(
         StagePipelineConfig(
             stage_id=0,

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -309,12 +309,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # tensor read; postprocess receives the tail-only slice instead, which
         # avoids ~18 ms merge + ~6 ms write per step (Sy0307 profile, #3665).
         self.requires_full_prefix_cached_hidden_states = False
-        # Stage 1 (code2wav) consumes runtime audio codes, not hidden states,
-        # so the inter-stage pooler payload never needs a CPU hidden-states
-        # view. Opting out lets the runner skip the per-step blocking
-        # ``hidden_states[:n].to("cpu")`` (measured at 24% of stage-0 on-CPU
-        # samples at c=64 on H20).
-        self.omni_pooler_payload_include_hidden = False
         # ``codes.audio`` is only needed for future prefix-hit reconstruction
         # after a request has produced codec rows. Keep per-step rows on GPU and
         # materialize the CPU OmniTensorPrefixCache entry once at completion.

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py

@@ -209,12 +209,7 @@ class VoxtralTTSForConditionalGeneration(
 
     def tts_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict | None):
         self.post_process_idx = 0
-        codes = info_dict.get("codes")
-        audio_tokens = codes.get("audio") if isinstance(codes, Mapping) else None
-        if audio_tokens is None:
-            # Backward compatible with request state written before #4527 moved
-            # Voxtral TTS feedback frames under ``codes.audio``.
-            audio_tokens = info_dict.pop("audio", None)
+        audio_tokens = info_dict.pop("audio", None)
         if audio_tokens is not None:
             kwargs = {"audio_tokens": audio_tokens.to(input_ids.device)}
             multimodal_embeddings = self.model.embed_multimodal(**kwargs)

vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml

@@ -0,0 +1,51 @@
+# Stage config for running MammothModa2 with full AR→DiT pipeline (text-to-image).
+# Stage 0: AR  (multimodal understanding + latent generation)
+# Stage 1: DiT (latent → image)
+#
+# For text/image understanding tasks (text output only), use mammoth_moda2_ar.yaml.
+
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0"
+    engine_args:
+      model_stage: ar
+      max_num_seqs: 100
+      model_arch: MammothModa2ForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      max_model_len: 8192
+      gpu_memory_utilization: 0.5
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      enable_prefix_caching: false
+    final_output: false
+
+  - stage_id: 1
+    runtime:
+      devices: "0"
+    engine_args:
+      model_stage: dit
+      max_num_seqs: 1
+      model_arch: MammothModa2ForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      gpu_memory_utilization: 0.3
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: image
+      enable_prefix_caching: false
+    # Declares extra_args so has_sampling_extra_args is enabled, routing the t2i knobs
+    # (text_guidance_scale, cfg_range, num_inference_steps) declared in
+    # vllm_omni/model_extras/mammothmodal2_preview.py from extra_body into the DiT
+    # forward. Values here are the defaults applied when a request omits a knob.
+    default_sampling_params:
+      extra_args:
+        text_guidance_scale: 9.0
+        cfg_range: [0.0, 1.0]
+        num_inference_steps: 50
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit
+    final_output: true
+    final_output_type: image

vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml

@@ -0,0 +1,24 @@
+# Stage config for running MammothModa2 with single AR stage (understanding / summarization).
+# Stage 0: AR (multimodal understanding + text generation)
+#
+# Use this config for image/video understanding tasks where only text output is needed.
+# For text-to-image generation (AR → DiT), use mammoth_moda2.yaml instead.
+
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0"
+    engine_args:
+      model_stage: ar
+      max_num_seqs: 16
+      model_arch: MammothModa2ForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      max_model_len: 8192
+      gpu_memory_utilization: 0.5
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: text
+      enable_prefix_caching: false
+    final_output: true
+    final_output_type: text

vllm_omni/model_extras/lingbot_video.py

@@ -1,16 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-LINGBOT_VIDEO_EXTRA_BODY_PARAMS = frozenset(
-    {
-        "batch_cfg",
-        "flow_shift",
-        "negative_prompt",
-        "null_cond_clone_zero",
-        "offload_vae_during_denoise",
-        "output_type",
-        "refiner_sigma_tail_steps",
-        "shift",
-        "t_thresh",
-    }
-)

vllm_omni/model_extras/ming_flash_omni.py

@@ -1,66 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-from typing import Any
-
-from PIL import Image
-
-
-def build_text_to_image_prompt(
-    prompt: str,
-    negative_prompt: str | None,
-    height: int | None = None,
-    width: int | None = None,
-) -> dict[str, Any]:
-    text_prompt: dict[str, Any] = {
-        "prompt": prompt,
-        "modalities": ["image"],
-        "mm_processor_kwargs": {"modalities": ["image"]},
-    }
-    if height is not None:
-        text_prompt["mm_processor_kwargs"]["target_h"] = height
-    if width is not None:
-        text_prompt["mm_processor_kwargs"]["target_w"] = width
-    if negative_prompt is not None:
-        text_prompt["negative_prompt"] = negative_prompt
-    return text_prompt
-
-
-def build_image_to_image_prompt(
-    prompt: str,
-    negative_prompt: str | None,
-    input_image: Image.Image | list[Image.Image],
-    height: int | None = None,
-    width: int | None = None,
-) -> dict[str, Any]:
-    img_prompt: dict[str, Any] = {
-        "prompt": prompt,
-        "modalities": ["img2img"],
-        "multi_modal_data": {"img2img": input_image},
-        "mm_processor_kwargs": {"modalities": ["img2img"]},
-    }
-    if height is not None:
-        img_prompt["mm_processor_kwargs"]["target_h"] = height
-    if width is not None:
-        img_prompt["mm_processor_kwargs"]["target_w"] = width
-    if negative_prompt is not None:
-        img_prompt["negative_prompt"] = negative_prompt
-    return img_prompt
-
-
-# For Image-generation (diffusion-stage)
-MING_FLASH_OMNI_EXTRA_BODY_PARAMS = frozenset(
-    {
-        "height",
-        "width",
-        "steps",
-        "cfg",
-        "seed",
-        "byte5_text",
-        "negative_prompt",
-    }
-)
-MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS = frozenset()
-MING_FLASH_OMNI_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True

vllm_omni/model_extras/registry.py

@@ -34,7 +34,6 @@ from vllm_omni.model_extras.helios import (
     HELIOS_EXTRA_BODY_PARAMS,
     HELIOS_EXTRA_OUTPUT_PARAMS,
 )
-from vllm_omni.model_extras.lingbot_video import LINGBOT_VIDEO_EXTRA_BODY_PARAMS
 from vllm_omni.model_extras.magi_human import (
     MAGI_HUMAN_EXTRA_BODY_PARAMS,
     MAGI_HUMAN_EXTRA_OUTPUT_PARAMS,
@@ -47,17 +46,6 @@ from vllm_omni.model_extras.mammothmodal2_preview import (
 from vllm_omni.model_extras.mammothmodal2_preview import (
     build_text_to_image_prompt as build_mammothmoda2_text_to_image_prompt,
 )
-from vllm_omni.model_extras.ming_flash_omni import (
-    MING_FLASH_OMNI_EXTRA_BODY_PARAMS,
-    MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS,
-    MING_FLASH_OMNI_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES,
-)
-from vllm_omni.model_extras.ming_flash_omni import (
-    build_image_to_image_prompt as build_ming_flash_omni_image_to_image_prompt,
-)
-from vllm_omni.model_extras.ming_flash_omni import (
-    build_text_to_image_prompt as build_ming_flash_omni_text_to_image_prompt,
-)
 from vllm_omni.model_extras.sensenova_u1 import (
     SENSENOVA_U1_EXTRA_BODY_PARAMS,
     SENSENOVA_U1_EXTRA_OUTPUT_PARAMS,
@@ -166,9 +154,6 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
         "extra_body_params": HELIOS_EXTRA_BODY_PARAMS,
         "extra_output_params": HELIOS_EXTRA_OUTPUT_PARAMS,
     },
-    "LingBotVideoPipeline": {
-        "extra_body_params": LINGBOT_VIDEO_EXTRA_BODY_PARAMS,
-    },
     "WanVACEPipeline": {
         "extra_body_params": VACE_EXTRA_BODY_PARAMS,
         "extra_output_params": VACE_EXTRA_OUTPUT_PARAMS,
@@ -180,13 +165,6 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
         "init_extra_args_for_non_diffusion_stages": MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES,
         "text_to_image_prompt_builder": build_mammothmoda2_text_to_image_prompt,
     },
-    "MingImagePipeline": {
-        "extra_body_params": MING_FLASH_OMNI_EXTRA_BODY_PARAMS,
-        "extra_output_params": MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS,
-        "init_extra_args_for_non_diffusion_stages": MING_FLASH_OMNI_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES,
-        "text_to_image_prompt_builder": build_ming_flash_omni_text_to_image_prompt,
-        "image_to_image_prompt_builder": build_ming_flash_omni_image_to_image_prompt,
-    },
 }
 
 

vllm_omni/platforms/cuda/platform.py

@@ -121,31 +121,6 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
-            if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
-                try:
-                    importlib.import_module("kernels")
-                    logger.info("Using HuggingFace kernels-backed attention backend '%s'", backend_upper)
-                except ImportError:
-                    if backend_upper == "FLASH_ATTN_HUB":
-                        logger.warning(
-                            "HuggingFace `kernels` library is not available. Falling back to local FLASH_ATTN."
-                        )
-                        backend_upper = "FLASH_ATTN"
-                    elif backend_upper == "FLASH_ATTN_3_HUB":
-                        logger.warning(
-                            "HuggingFace `kernels` library is not available. Falling back to local FLASH_ATTN."
-                        )
-                        backend_upper = "FLASH_ATTN"
-
-            if backend_upper == "FLASH_ATTN_3_HUB":
-                fa3_hub_supported = compute_capability is not None and compute_capability.major >= 9
-                if not fa3_hub_supported:
-                    logger.warning(
-                        "FLASH_ATTN_3_HUB requires a Hopper-class GPU with compute capability >= 9.0. "
-                        "Falling back to FLASH_ATTN_HUB."
-                    )
-                    backend_upper = "FLASH_ATTN_HUB"
-
             if backend_upper == "FLASH_ATTN" and not flash_attn_supported:
                 if not compute_supported:
                     logger.warning(
@@ -249,16 +224,8 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
 
     @classmethod
     def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
-        """Copied from upstream CudaPlatformBase with inductor-aware logic.
-
-        When inductor is active (compiling) use native as the default;
-        otherwise prefer vllm_c kernels where available.
-        """
-        from vllm.config.compilation import CompilationMode
-
-        cc = vllm_config.compilation_config
-        using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
-        default = ["native"] if using_inductor else ["vllm_c", "native"]
+        """Prefer ``vllm_c`` CUDA kernels over ``native`` for diffusion IR ops."""
+        default = ["vllm_c", "native"]
 
         # Use oink if enabled for rms_norm
         # TODO(Laurawly/luka): remove this env var,

vllm_omni/platforms/musa/platform.py

@@ -78,14 +78,6 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
-            if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
-                logger.warning(
-                    "HuggingFace kernels-backed FlashAttention is "
-                    "not supported on MUSA. Falling back to local "
-                    "FLASH_ATTN."
-                )
-                backend_upper = "FLASH_ATTN"
-
             if backend_upper == "FLASH_ATTN" and not flash_attn_supported:
                 if not compute_supported:
                     logger.warning(