OmniRequestOutput inherits RequestOutput

Signed-off-by: Boao Shi <aoibosh@connect.hku.hk>

.buildkite/pipeline-intel.yaml

@@ -10,7 +10,7 @@ steps:
           DOCKER_BUILDKIT: "1"
           # Buildkite will automatically replace this with the actual commit hash
           VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.24.0"
+          VLLM_VERSION: "v0.25.0"
         priority: 100
         timeout_in_minutes: 60
         soft_fail: false

.buildkite/scripts/hardware_ci/run-xpu-test.sh

@@ -13,12 +13,33 @@ container_name="xpu_${BUILDKITE_COMMIT}_$(
 )"
 
 cd "${omni_source_dir}"
+
+# The XPU base image is ~37GB; the default gzip layer exporter is single-threaded
+# and takes ~16min to compress it. zstd is multi-threaded (uses all cores) at a
+# similar ratio; set EXPORT_COMPRESSION=uncompressed to skip compression entirely
+# for local-only images. This requires the containerd image store (buildx docker
+# driver), which is the default here.
+EXPORT_COMPRESSION="${EXPORT_COMPRESSION:-zstd}"
+if [ "${EXPORT_COMPRESSION}" = "uncompressed" ]; then
+    export_args=(--output "type=image,name={{IMAGE}},compression=uncompressed")
+else
+    export_args=(--output "type=image,name={{IMAGE}},compression=${EXPORT_COMPRESSION},compression-level=3,force-compression=true")
+fi
+
+docker_build() {
+    # $1 = image name; remaining args passed through to docker build.
+    local image="$1"
+    shift
+    local out=("${export_args[@]/'{{IMAGE}}'/${image}}")
+    docker build "${out[@]}" "$@" -f docker/Dockerfile.xpu .
+}
+
 if [ -z "$(docker images -q "${base_image_name}")" ]; then
-    docker build --target vllm-base -t "${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -f docker/Dockerfile.xpu .
+    docker_build "${base_image_name}" --target vllm-base --build-arg "VLLM_VERSION=${VLLM_VERSION}"
 fi
 
 # Try building the docker image
-docker build --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -t "${image_name}" -f docker/Dockerfile.xpu .
+docker_build "${image_name}" --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}"
 
 # Setup cleanup
 remove_docker_container() {

.buildkite/test-nightly.yml

@@ -650,6 +650,8 @@ steps:
             tests/e2e/online_serving/test_qwen_image_edit_expansion.py
             tests/e2e/online_serving/test_qwen_image_layered_expansion.py
             tests/e2e/online_serving/test_bagel_expansion.py
+            tests/e2e/online_serving/test_krea2_expansion.py
+            tests/e2e/offline_inference/test_krea2_expansion.py
             -m "full_model and diffusion and H100 and not distributed_cuda"
             --run-level "full_model"
         agents:
@@ -693,6 +695,7 @@ steps:
             pytest -sv
             tests/e2e/online_serving/test_qwen_image_expansion.py
             tests/e2e/online_serving/test_qwen_image_edit_expansion.py
+            tests/e2e/offline_inference/test_krea2_expansion.py
             -m "full_model and diffusion and H100 and distributed_cuda"
             --run-level "full_model"
         agents:
@@ -1486,9 +1489,12 @@ steps:
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json
             EXIT1=$$?
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json
+            EXIT2=$$?
             buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
             buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
-            exit $$EXIT1
+            if [ $$EXIT1 -ne 0 ]; then exit $$EXIT1; fi
+            exit $$EXIT2
         agents:
           queue: "mithril-h100-pool"
         plugins:

.gitignore

@@ -248,6 +248,7 @@ datasets/
 
 # Output files
 outputs/
+!vllm_omni/outputs/
 results/
 generated/
 output_*/

README.md

@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---
 
 *Latest News* 🔥
+- [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
 - [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
 - [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
 - [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
@@ -31,9 +32,9 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 [vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni is a framework that extends its support for omni-modality model inference and serving:
 
-- **Omni-modality**: Text, image, video, and audio data processing
+- **Omni-modality**: Text, image, audio, video, and action data processing
 - **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models
-- **Heterogeneous outputs**: from traditional text generation to multimodal outputs
+- **Heterogeneous outputs**: from traditional text generation to multimodal and action outputs
 
 <p align="center">
   <picture>
@@ -57,9 +58,10 @@ vLLM-Omni is flexible and easy to use with:
 
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
-- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos, HunyuanImage, BAGEL)
+- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos3, HunyuanImage, BAGEL)
 - **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
 - **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
+- **Robot-policy and action models** (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
 
 ## Getting Started
 

benchmarks/diffusion/backends.py

@@ -28,6 +28,7 @@ class RequestFuncInput:
     slo_ms: float | None = None
     extra_body: dict[str, Any] = field(default_factory=dict)
     image_paths: list[str] | None = None
+    video_paths: list[str] | None = None
     request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
     default_bot_task: str | None = DEFAULT_EDITS_BOT_TASK
 
@@ -338,6 +339,11 @@ async def async_request_v1_videos(
         form.add_field(k, str(v))
 
     image_file = None
+    if input.image_paths and input.video_paths:
+        output.error = "Only one of image_paths or video_paths can be provided"
+        output.success = False
+        return output
+
     if input.image_paths and len(input.image_paths) > 0:
         image_path = input.image_paths[0]
         image_file = open(image_path, "rb")
@@ -345,7 +351,16 @@ async def async_request_v1_videos(
             "input_reference",
             image_file,
             filename=os.path.basename(image_path),
-            content_type="application/octet-stream",
+            content_type=_guess_mime_type(image_path),
+        )
+    elif input.video_paths and len(input.video_paths) > 0:
+        video_path = input.video_paths[0]
+        image_file = open(video_path, "rb")
+        form.add_field(
+            "input_reference",
+            image_file,
+            filename=os.path.basename(video_path),
+            content_type=_guess_mime_type(video_path),
         )
 
     job_id = None

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -73,6 +73,10 @@ Usage:
             {"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1}
         ]'
 
+    v2v:
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+        --endpoint /v1/videos --dataset random --task v2v --num-prompts 1 \
+        --height 720 --width 1280 --fps 24 --num-frames 189 --num-inference-steps 35
 
 """
 
@@ -542,11 +546,17 @@ class TraceDataset(BaseDataset):
         if not image_paths:
             single = row.get("image_path")
             image_paths = [single] if single else None
+        video_paths = row.get("video_paths")
+        if not video_paths:
+            single_video = row.get("video_path")
+            video_paths = [single_video] if single_video else None
 
         if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i", "it2i"]:
             raise ValueError(
                 f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row."
             )
+        if not video_paths and self.args.task == "v2v":
+            raise ValueError("Task v2v requires video input, but no video_path or video_paths found in trace row.")
 
         override_w = self.args.width
         override_h = self.args.height
@@ -570,6 +580,7 @@ class TraceDataset(BaseDataset):
             timestamp=timestamp,
             slo_ms=slo_ms,
             image_paths=image_paths,
+            video_paths=video_paths,
             request_id=str(row.get("request_id")) if row.get("request_id") is not None else str(uuid.uuid4()),
         )
 
@@ -589,6 +600,8 @@ class CustomDataset(BaseDataset):
     - seed (optional): Random seed
     - image_paths (optional): List of input image paths for i2i/i2v/ti2i tasks
     - image_urls (optional): List of input image URLs (alternative to image_paths)
+    - video_paths (optional): List of input video paths for v2v tasks
+    - video_urls (optional): List of input video URLs (alternative to video_paths)
 
     Example JSONL for ti2i:
     {"prompt": "Add sunset lighting", "width": 1024, "height": 1024, "image_urls": ["https://example.com/image.jpg"]}
@@ -660,6 +673,42 @@ class CustomDataset(BaseDataset):
 
         return None
 
+    def _resolve_video_paths(self, item: dict) -> list[str] | None:
+        if "video_paths" in item and item["video_paths"]:
+            paths = item["video_paths"]
+            if isinstance(paths, str):
+                paths = [paths]
+
+            valid_paths = []
+            for path in paths:
+                if os.path.exists(path):
+                    valid_paths.append(path)
+                else:
+                    raise ValueError(f"Video file not found: {path}")
+
+            return valid_paths if valid_paths else None
+
+        if "video_urls" in item and item["video_urls"]:
+            urls = item["video_urls"]
+            if isinstance(urls, str):
+                urls = [urls]
+
+            downloaded_paths = []
+            for url in urls:
+                try:
+                    response = requests.get(url, timeout=30)
+                    response.raise_for_status()
+                    suffix = os.path.splitext(url)[1] or ".mp4"
+                    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+                        tmp.write(response.content)
+                        downloaded_paths.append(tmp.name)
+                except Exception as e:
+                    raise ValueError(f"Failed to download video from {url}: {e}")
+
+            return downloaded_paths if downloaded_paths else None
+
+        return None
+
     def __len__(self) -> int:
         return len(self.data)
 
@@ -672,11 +721,22 @@ class CustomDataset(BaseDataset):
         height = item.get("height", self.args.height)
         num_inference_steps = item.get("num_inference_steps", self.args.num_inference_steps)
         seed = item.get("seed", self.args.seed)
-        reserved_keys = {"prompt", "width", "height", "num_inference_steps", "seed", "image_paths", "image_urls"}
+        reserved_keys = {
+            "prompt",
+            "width",
+            "height",
+            "num_inference_steps",
+            "seed",
+            "image_paths",
+            "image_urls",
+            "video_paths",
+            "video_urls",
+        }
         extra_body = {k: v for k, v in item.items() if k not in reserved_keys}
 
         # Handle image paths/URLs
         image_paths = self._resolve_image_paths(item)
+        video_paths = self._resolve_video_paths(item)
 
         return RequestFuncInput(
             prompt=prompt,
@@ -684,6 +744,7 @@ class CustomDataset(BaseDataset):
             model=self.model,
             seed=seed,
             image_paths=image_paths,
+            video_paths=video_paths,
             extra_body=extra_body,
             width=width,
             height=height,
@@ -737,6 +798,10 @@ class RandomDataset(BaseDataset):
             self._random_image_path = self._generate_random_image_paths()
         else:
             self._random_image_path = None
+        if self.args.task == "v2v":
+            self._random_video_path = self._generate_random_video_paths()
+        else:
+            self._random_video_path = None
 
     def __len__(self) -> int:
         return self.num_prompts
@@ -756,13 +821,15 @@ class RandomDataset(BaseDataset):
         if self._sampled_requests:
             profile = self._sampled_requests[idx]
             params.update(profile)
+        prompt = str(params.pop("prompt", f"Random prompt {idx} for benchmarking diffusion models"))
         return RequestFuncInput(
-            prompt=f"Random prompt {idx} for benchmarking diffusion models",
+            prompt=prompt,
             api_url=self.api_url,
             model=self.model,
             seed=self.args.seed,
             extra_body=extra_body,
             image_paths=self._random_image_path,
+            video_paths=self._random_video_path,
             **params,
         )
 
@@ -781,6 +848,35 @@ class RandomDataset(BaseDataset):
             image_paths.append(image_path)
         return image_paths
 
+    def _generate_random_video_paths(self) -> list[str]:
+        video_path = os.path.join(tempfile.gettempdir(), "diffusion_benchmark_random_video.mp4")
+        try:
+            import cv2
+
+            profile = self.random_request_config[0] if self.random_request_config else {}
+            width = int(self.args.width or profile.get("width") or 1280)
+            height = int(self.args.height or profile.get("height") or 720)
+            num_frames = int(self.args.num_frames or profile.get("num_frames") or 16)
+            fps = float(self.args.fps or profile.get("fps") or 8)
+            writer = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
+            if not writer.isOpened():
+                raise RuntimeError("cv2.VideoWriter failed to open")
+            for frame_idx in range(num_frames):
+                frame = np.zeros((height, width, 3), dtype=np.uint8)
+                x = int((frame_idx / max(num_frames - 1, 1)) * max(width - 96, 1))
+                y = height // 2
+                frame[:, :, 0] = 24
+                frame[:, :, 1] = 32
+                frame[:, :, 2] = 48
+                frame[max(y - 36, 0) : min(y + 36, height), x : min(x + 96, width), :] = (64, 128, 220)
+                writer.write(frame)
+            writer.release()
+            return [video_path]
+        except Exception as e:
+            raise RuntimeError(
+                "Failed to generate synthetic v2v input video. Install opencv-python or provide a custom dataset."
+            ) from e
+
 
 def _compute_expected_latency_ms_from_base(req: RequestFuncInput, args, base_time_ms: float | None) -> float | None:
     """Compute expected execution time (ms) based on a base per-step-per-frame unit time.
@@ -1112,7 +1208,7 @@ def wait_for_service(base_url: str, timeout: int = 120) -> None:
 
 
 def _default_endpoint_for_task(task: str) -> str:
-    if task in {"t2v", "i2v", "ti2v"}:
+    if task in {"t2v", "i2v", "ti2v", "v2v"}:
         return "/v1/videos"
     if task in {"i2i", "ti2i", "it2i"}:
         return "/v1/images/edits"
@@ -1126,7 +1222,7 @@ async def benchmark(args):
     if args.base_url is None:
         args.base_url = f"http://{args.host}:{args.port}"
 
-    VIDEO_TASKS = {"t2v", "i2v", "ti2v"}
+    VIDEO_TASKS = {"t2v", "i2v", "ti2v", "v2v"}
     IMAGE_TASKS = {"t2i", "i2i", "ti2i", "it2i"}
 
     if args.task in VIDEO_TASKS:
@@ -1177,6 +1273,10 @@ async def benchmark(args):
     if args.return_stage_metrics and args.endpoint in _STAGE_METRICS_ENDPOINTS:
         for req in requests_list:
             req.extra_body.setdefault(_RETURN_STAGE_METRICS_FIELD, True)
+    if args.extra_body:
+        extra_body = json.loads(args.extra_body)
+        for req in requests_list:
+            req.extra_body.update(extra_body)
 
     if args.endpoint == "/v1/images/edits":
         for req in requests_list:
@@ -1328,7 +1428,7 @@ if __name__ == "__main__":
         "--task",
         type=str,
         default="t2v",
-        choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "it2i", "t2i"],
+        choices=["t2v", "i2v", "ti2v", "v2v", "ti2i", "i2i", "it2i", "t2i"],
         help="Task type.",
     )
     parser.add_argument(
@@ -1443,6 +1543,12 @@ if __name__ == "__main__":
             '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
         ),
     )
+    parser.add_argument(
+        "--extra-body",
+        type=str,
+        default=None,
+        help="JSON object merged into every generated request body/form.",
+    )
     parser.add_argument(
         "--num-input-images",
         type=int,

docker/Dockerfile.ci

@@ -2,7 +2,7 @@
 # The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must
 # be updated on every rebase to match the new upstream vllm commit/tag.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.24.0
+ARG VLLM_BASE_TAG=v0.25.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.cuda

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.24.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.25.0
 FROM ${BASE_IMAGE}
 
 ARG COMMON_WORKDIR=/app

docker/Dockerfile.npu

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.24.0
+ARG VLLM_ASCEND_TAG=v0.25.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm

docker/Dockerfile.npu.a3

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.24.0-a3
+ARG VLLM_ASCEND_TAG=v0.25.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm

docker/Dockerfile.npu.ci

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.24.0
+ARG VLLM_ASCEND_TAG=v0.25.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni

docker/Dockerfile.npu.ci.a3

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.24.0-a3
+ARG VLLM_ASCEND_TAG=v0.25.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.24.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.25.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.

docker/Dockerfile.xpu

@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"
 
-ARG VLLM_VERSION=v0.24.0
+ARG VLLM_VERSION=v0.25.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm
 

docs/.nav.yml

@@ -23,7 +23,7 @@ nav:
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
       - InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
       - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
-      - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
+      - MiMo-Audio: user_guide/examples/offline_inference/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
       - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
@@ -38,7 +38,7 @@ nav:
       - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
       - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
-      - Online serving Example of vLLM-Omni for MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
+      - MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
       - Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
@@ -80,6 +80,7 @@ nav:
       - Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
       - FP8 W8A8: user_guide/quantization/fp8.md
       - Int8 W8A8: user_guide/quantization/int8.md
+      - BitsAndBytes W4: user_guide/quantization/bitsandbytes.md
       - MXFP8 W8A8: user_guide/quantization/mxfp8.md
       - MXFP4 W4A4: user_guide/quantization/mxfp4.md
       - ModelOpt: user_guide/quantization/modelopt.md

docs/README.md

@@ -28,9 +28,9 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 [vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni is a framework that extends its support for omni-modality model inference and serving:
 
-- **Omni-modality**: Text, image, video, and audio data processing
+- **Omni-modality**: Text, image, audio, video, and action data processing
 - **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models
-- **Heterogeneous outputs**: from traditional text generation to multimodal outputs
+- **Heterogeneous outputs**: from traditional text generation to multimodal and action outputs
 
 <p align="center">
   <picture>
@@ -55,8 +55,10 @@ vLLM-Omni is flexible and easy to use with:
 
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
-- Omni-modality models (e.g. Qwen2.5-Omni, Qwen3-Omni)
-- Multi-modality generation models (e.g. Qwen-Image)
+- Omni-modality models (e.g. Qwen3-Omni, Cosmos3, HunyuanImage, BAGEL)
+- TTS models (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
+- Diffusion models — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
+- Robot-policy and action models (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
 
 For more information, checkout the following:
 

docs/api/README.md

@@ -46,17 +46,17 @@ Engine classes for offline and online inference.
 - [vllm_omni.engine.PromptEmbedsPayload][]
 - [vllm_omni.engine.arg_utils.OmniEngineArgs][]
 - [vllm_omni.engine.async_omni_engine.AsyncOmniEngine][]
-- [vllm_omni.engine.mm_outputs.MultimodalCompletionOutput][]
-- [vllm_omni.engine.mm_outputs.MultimodalPayload][]
 - [vllm_omni.engine.orchestrator.Orchestrator][]
 - [vllm_omni.engine.orchestrator.OrchestratorRequestState][]
-- [vllm_omni.engine.output_modality.OutputModality][]
-- [vllm_omni.engine.output_modality.TensorAccumulationStrategy][]
-- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
-- [vllm_omni.engine.output_processor.OmniRequestState][]
 - [vllm_omni.engine.stage_engine_core_client.StageEngineCoreClient][]
 - [vllm_omni.engine.stage_init_utils.StageMetadata][]
 - [vllm_omni.engine.stage_init_utils.StartedLlmStage][]
+- [vllm_omni.outputs.mm_outputs.MultimodalCompletionOutput][]
+- [vllm_omni.outputs.mm_outputs.MultimodalPayload][]
+- [vllm_omni.outputs.output_modality.OutputModality][]
+- [vllm_omni.outputs.output_modality.TensorAccumulationStrategy][]
+- [vllm_omni.outputs.output_processor.MultimodalOutputProcessor][]
+- [vllm_omni.outputs.output_processor.OmniRequestState][]
 
 ## Core
 

docs/cli/README.md

@@ -18,11 +18,17 @@ Specify the port:
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
 
-If you have custom stage configs file, launch the server with command below
+For a migrated model, load a custom deploy configuration with `--deploy-config`:
+
 ```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni --stage-configs-path /path/to/stage_configs_file
+vllm serve Qwen/Qwen2.5-Omni-7B --omni --deploy-config /path/to/deploy_config.yaml
 ```
 
+The deprecated `--stage-configs-path` flag is retained for models that still use the legacy `stage_args` schema:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --stage-configs-path /path/to/legacy_stage_config.yaml
+```
 
 ## bench
 

docs/cli/bench/serve.md

@@ -2,7 +2,7 @@
 The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
 
 ## Notes
-We currently only support using the "openai-chat-omni" and "openai-image-edits-omni" backend.
+vLLM-Omni registers the `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` serving benchmark backends.
 
 ## Basic Parameter Description
 You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
@@ -10,7 +10,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
   Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio.
 
 - `--backend`
-  Specify the backend adapter as openai-chat-omni, using OpenAI Chat compatible API behavior as the protocol. Currently only openai-chat-omni is supported.
+  Specify the backend adapter. vLLM-Omni adds `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` to the upstream vLLM backend choices.
 
 - `--model`
   The model identifier to load, filled according to the models supported by vLLM-Omni.
@@ -55,7 +55,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
                     'Allowed metric names are "ttft", "tpot", "itl", "ttfc", "tpoc", "icl", '
                     '"tpop", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
 
-- `-print-stage`
+- `--print-stage`
 Print per-stage benchmark metrics for --omni serving when stage metrics are returned by the server. Disabled by default.
 
 - `--save-result`
@@ -312,7 +312,7 @@ vllm bench serve \
   --ignore-eos \
   --percentile-metrics ttft,tpot,itl \
   --random-output-len 2 \
-  --extra_body '{"modalities": ["text"]}'
+  --extra-body '{"modalities": ["text"]}'
 ```
 
 If successful, you will see the following output:
@@ -390,7 +390,7 @@ vllm bench serve --omni \
   --ignore-eos \
   --percentile-metrics ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,ttfc,tpoc,icl \
   --random-output-len 900 \
-  --extra_body '{"modalities": ["text","audio"]}' \
+  --extra-body '{"modalities": ["text","audio"]}' \
   --print-stage
 ```
 

docs/configuration/composable_parallel.md

@@ -92,12 +92,12 @@ Full pattern hierarchy: `TakeRank`, `Union`, `GatherDim`, `AllGather`, `StitchSp
 
 From highest to lowest, when the same axis is touched by multiple layers:
 
-1. **CLI overrides** (`--stage_N_tensor_parallel_size`, `--stage-overrides`, `--omni-lb-policy`, etc.). A CLI value always wins; a warning is emitted whenever it overrides a strategy-declared axis (see `_reconcile_strategy_with_cli` in `vllm_omni/config/config_factory.py`).
+1. **CLI overrides** (`--tensor-parallel-size`, `--stage-overrides`, `--omni-lb-policy`, etc.). A CLI value always wins; a warning is emitted whenever it overrides a strategy-declared axis (see `_reconcile_strategy_with_cli` in `vllm_omni/config/config_factory.py`).
 2. **Strategy YAML** (`--strategy-config`). Derived sizing is written onto the merged stages with *conflict-on-explicit* semantics — if the deploy YAML already set a value to something different, application raises `StrategyApplyError`.
 3. **Deploy YAML** (`--deploy-config` or the bundled `vllm_omni/deploy/<model_type>.yaml`).
 4. **Parser defaults.**
 
-After CLI overrides land, the device-layout guard re-runs (`check_device_layout`, called from `_reconcile_strategy_with_cli`) against the *effective* `tp * dp * pp * num_replicas` and the resolved `devices` string, so a `--devices` or `--tensor-parallel-size` that conflicts with the strategy's world size cannot slip past the pre-spawn check.
+After CLI overrides land, the device-layout guard re-runs (`check_device_layout`, called from `_reconcile_strategy_with_cli`) against the *effective* `tp * dp * pp * num_replicas` and the resolved `devices` string, so a `devices` value passed through `--stage-overrides` or a `--tensor-parallel-size` value that conflicts with the strategy's world size cannot slip past the pre-spawn check.
 
 ## Worked example
 
@@ -214,7 +214,7 @@ Strategy keys MUST be `model_stage` *names* (strings), the same labels stage con
     Use the `model_stage` name (e.g. `thinker`) instead.
 
 !!! warning "Mixing CLI overrides with a strategy on the same axis"
-    A strategy is meant to be the single writer for the axes it declares. If a CLI override (`--stage_0_tensor_parallel_size`, `--stage-overrides '{"0": {...}}'`) sets the same axis, the **CLI value wins** and a warning is emitted from `_reconcile_strategy_with_cli`. The device-layout guard then re-runs against the effective layout, so a CLI override that breaks `tp * dp * pp * num_replicas` will still fail fast at config time.
+    A strategy is meant to be the single writer for the axes it declares. If a CLI override (`--tensor-parallel-size 2`, `--stage-overrides '{"0": {"tensor_parallel_size": 2}}'`) sets the same axis, the **CLI value wins** and a warning is emitted from `_reconcile_strategy_with_cli`. The device-layout guard then re-runs against the effective layout, so a CLI override that breaks `tp * dp * pp * num_replicas` will still fail fast at config time.
 
 !!! warning "Passing `--strategy-config` against the legacy `--stage-configs-path`"
     Composable parallel only applies on the registry-based deploy path. If a model still resolves through the legacy `stage_args` YAML (loaded via `--stage-configs-path`), `--strategy-config` is silently inapplicable and emits:

docs/configuration/stage_configs.md

@@ -1,11 +1,11 @@
-# Stage configs for vLLM-Omni
+# Deploy and legacy stage configurations
 
-In vLLM-Omni, the target model is separated into multiple stages, which are processed by different LLMEngines, DiffusionEngines or other types of engines. Depending on different types of stages, such as Autoregressive (AR) stage or Diffusion transformer (DiT) stage, each can choose corresponding schedulers, model workers to load with the Engines in a plug-in fashion.
+In vLLM-Omni, a model's `PipelineConfig` defines its fixed stage topology, while a deploy configuration controls how those stages run. Models that have not migrated to this split still use legacy stage configurations with a `stage_args` schema.
 
 !!! note
     Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when neither `--stage-configs-path` nor `--deploy-config` is provided — the model registry resolves the right pipeline + deploy YAML by `model_type`. The bundled defaults have been verified on 1xH100 for Qwen2.5-Omni and 2xH100 for Qwen3-Omni. Models that have not yet migrated to the new schema continue to use the legacy `vllm_omni/model_executor/stage_configs/<model>.yaml` files via `--stage-configs-path`.
 
-## New deploy schema reference
+## Deploy configuration schema
 
 The new deploy schema lives under `vllm_omni/deploy/` and is paired with a frozen `PipelineConfig` registered by the model's `pipeline.py`. Each deploy YAML has these top-level fields:
 
@@ -82,14 +82,14 @@ stages:
     input_connectors:  {from_stage_0: shm}
 ```
 
-### CLI flags introduced in this refactor
+### CLI flags
 
 | Flag | Description |
 |------|-------------|
-| `--deploy-config PATH` | Load a new-schema deploy YAML. Takes precedence over `--stage-configs-path`. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
+| `--deploy-config PATH` | Load a new-schema deploy YAML. It is mutually exclusive with `--stage-configs-path`. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
 | `--stage-overrides JSON` | Per-stage JSON overrides, e.g. `'{"0":{"gpu_memory_utilization":0.5}}'`. Per-stage values always win over global flags. |
 | `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. |
-| `--stage-configs-path` | **Deprecated.** Accepts legacy `stage_args` yamls and (auto-detected) new deploy yamls; emits a deprecation warning. Migrate to `--deploy-config`. To be removed in a follow-up PR. |
+| `--stage-configs-path` | **Deprecated.** Load a legacy `stage_args` YAML. New-schema YAMLs passed to this flag are auto-detected for compatibility, but should use `--deploy-config`. It is mutually exclusive with `--deploy-config` and will be removed in a future release. |
 
 ### Stage-Based CLI Paradigm
 
@@ -146,7 +146,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 
 From highest to lowest:
 
-1. Per-stage flags (`--stage-overrides` JSON, `--stage-<id>-<key>` if registered)
+1. Per-stage overrides (`--stage-overrides` JSON)
 2. Explicit global CLI flags (`--gpu-memory-utilization 0.85`, etc.)
 3. Platform section (`platforms.npu.stages`, etc.) on top of the base `stages:`
 4. Overlay YAML (via `base_config:`) on top of the base YAML
@@ -210,7 +210,7 @@ Effective config per stage after the merge:
 | 1 | `max_model_len` | `16384` | global CLI |
 | 2 | (all defaults) | — | base YAML (no overrides apply) |
 
-Therefore, as a core part of vLLM-Omni, the stage configs for a model have several main functions:
+Therefore, as a core part of vLLM-Omni, a model's pipeline and deployment configurations have several main functions:
 
 - Claim partition of stages and their corresponding class implementation in `model_executor/models`.
 - The disaggregated configuration for each stage and the communication topology among them.
@@ -219,16 +219,16 @@ Therefore, as a core part of vLLM-Omni, the stage configs for a model have sever
 - Default input parameters.
 
 To override specific parameters, explicitly inject the customized configuration schema
-in both online and offline instantiation flows. Prioritize the `--deploy-config` flag
-when loading the new-schema deploy YAML schemas, reserving the `--stage-configs-path` parameter
+in both online and offline instantiation flows. Use the `--deploy-config` flag
+when loading a deploy configuration, reserving the `--stage-configs-path` parameter
 exclusively to maintain compatibility with legacy `stage_args` YAML constructs.
 
 Examples:
 
-For offline (Assume necessary dependencies have been imported):
+For offline inference (assuming the necessary dependencies have been imported):
 ```python
 model_name = "Qwen/Qwen2.5-Omni-7B"
-omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
+omni = Omni(model=model_name, deploy_config="/path/to/deploy_config.yaml")
 ```
 
 For online serving:
@@ -242,11 +242,14 @@ Legacy online serving:
 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
 !!! important
-    We are actively iterating on the definition of stage configs, and we welcome all feedbacks from both community users and developers to help us shape the development!
+    We are actively iterating on the definition of deployment configurations, and we welcome feedback from users and developers.
+
+## Legacy stage configuration reference
+
+The remainder of this page documents the deprecated `stage_args` schema for models that have not migrated to deploy configurations. The following historical Qwen2.5-Omni configuration illustrates that schema; current Qwen2.5-Omni deployments should use `vllm_omni/deploy/qwen2_5_omni.yaml` instead.
 
-Below is a specific example of stage_configs.yaml in Qwen2.5-omni.
 ```python
-# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
+# Historical stage configuration for Qwen2.5-Omni.
 stage_args:
   - stage_id: 0 # mark the unique id for each stage
     runtime: # The disaggregated configuration
@@ -338,7 +341,7 @@ runtime:
 
 ```
 
-## Stage Configuration Arguments
+## Legacy stage configuration arguments
 
 Each stage in the `stage_args` list contains the following configuration options:
 

docs/contributing/DOCS_GUIDE.md

@@ -67,6 +67,21 @@ docs/
 └── stylesheets/         # Custom CSS
 ```
 
+## Naming Model Examples
+
+Offline inference and online serving examples for the same model use the same
+directory name and the shared display name in
+`examples/model_display_names.yml`. Use these title forms:
+
+- `# <Model>: Offline inference`
+- `# <Model>: Online serving`
+
+The documentation generator uses the full title for the page H1 and the shared
+display name alone for navigation, and fails the build if a mapped README uses
+a different H1. Keep checkpoint identifiers in commands and prose rather than
+in the display name, and do not rename an existing example directory solely to
+adjust its title because the directory defines its public documentation URL.
+
 ## Publishing Documentation
 
 ### GitHub Pages (Recommended)

docs/contributing/ci/test_guide.md

@@ -12,12 +12,37 @@ When running test cases, you may need to install the following dependencies:
 
 ```bash
 uv pip install ".[dev]"
-apt-get install -y ffmpeg
+apt-get install -y espeak-ng jq
 ```
 
 ## Running Tests
 Our test scripts use the pytest framework. First, please use `git clone https://github.com/vllm-project/vllm-omni.git` to download the vllm-omni source code. Then, in the root directory of vllm-omni, you can run the following commands in your local test environment to execute the corresponding test cases.
 
+### CI job runners (L2–L4): logs, timing, and timeouts
+
+[`tools/run_ready_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_ready_jobs.sh), [`tools/run_merge_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_merge_jobs.sh), and [`tools/nightly/run_nightly_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/nightly/run_nightly_jobs.sh) share the same run-time behavior (via [`tools/run_jobs_common.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_jobs_common.sh)):
+
+**Log layout** (default under `logs/ready_jobs`, `logs/merge_jobs`, or `logs/nightly_jobs`; override with `--log-dir`):
+
+| Path | Description |
+|------|-------------|
+| `jobs/*.sh` | Generated per-step bash wrappers |
+| `jobs/.job_timeouts` | Job key → `timeout_in_minutes` from Buildkite YAML (when set) |
+| `<job_key>.log` | Tee output for each job |
+| `timing_summary.log` | Per-job duration and total wall time after the run completes |
+
+**Timing:** while each job runs, stderr prints `finished in …` or `failed after …`. When all jobs finish, a summary is printed and written to `timing_summary.log`, for example:
+
+```
+=== Job timing summary ===
+  TTS_VoxCPM2_Test  25m 40s  OK
+  Omni_Qwen3-Omni_Test  8m 05s  TIMED OUT
+Total wall time: 33m 45s (2 jobs)
+Failed jobs: 1/2
+```
+
+**Timeouts:** if a Buildkite step defines `timeout_in_minutes`, the runner wraps that job with `timeout ${N}m` (aligned with CI). Steps without `timeout_in_minutes` rely on any inline `timeout …` already present in the extracted pytest command. On timeout the process exits with code `124` and the summary marks the job as `TIMED OUT`.
+
 === "L1 level"
 
     ```bash
@@ -28,61 +53,109 @@ Our test scripts use the pytest framework. First, please use `git clone https://
 
 === "L2 level"
 
+    **Recommended:** run CI-aligned jobs from the repo root via [`tools/run_ready_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_ready_jobs.sh). The script reads [`.buildkite/test-ready.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml), generates per-step bash wrappers, runs pytest, and tees logs under `logs/ready_jobs/` (see [CI job runners](#ci-job-runners-l2l4-logs-timing-and-timeouts) for timing and timeout behavior). Requires `bash`, `python3`, and PyYAML (`pip install pyyaml`).
+
+    ```bash
+    # All L2 and L1 jobs (default: every pytest step in test-ready.yml)
+    bash tools/run_ready_jobs.sh
+
+    # Skip L1-style "Simple Test" steps in test-ready.yml (group or "Simple ·" labels)
+    bash tools/run_ready_jobs.sh --skip-simple
+    bash tools/run_ready_jobs.sh --skip-simple --model-type diffusion
+
+    # Preview extracted commands without running
+    bash tools/run_ready_jobs.sh --dry-run
+
+    # Filter by model area (OR semantics): omni | tts | diffusion | all
+    bash tools/run_ready_jobs.sh --model-type omni
+    bash tools/run_ready_jobs.sh --model-type tts,diffusion
+
+    # Match a Buildkite label substring (e.g. one model suite)
+    bash tools/run_ready_jobs.sh --model-type tts --label-substr "VoxCPM2"
+
+    ```
+
+    **Ad hoc:** run a single test file or marker expression manually:
+
     ```bash
     cd tests
-    pytest -s -v -m "core_model and not cpu" --run-level=core_model
-    ```
-    If you only want to run a specific test case, you can use:
-    ```bash
     pytest -s -v test_xxxx.py --run-level=core_model
+    pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
     ```
-    If you only want to run specific test cases on a particular platform, you can use:
-    ```bash
-    pytest -s -v -m "core_model and distributed_cuda and L4"  --run-level=core_model
-    ```
-    The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
 
 === "L3 level"
 
+    **Recommended:** run CI-aligned jobs from the repo root via [`tools/run_merge_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_merge_jobs.sh). The script reads [`.buildkite/test-merge.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml), generates per-step bash wrappers, runs pytest, and tees logs under `logs/merge_jobs/` (see [CI job runners](#ci-job-runners-l2l4-logs-timing-and-timeouts)). Requires `bash`, `python3`, and PyYAML (`pip install pyyaml`).
+
     ```bash
-    pytest -s -v -m "advanced_model" --run-level=advanced_model
+    # All L3 and L1 jobs (default: every pytest step in test-merge.yml)
+    bash tools/run_merge_jobs.sh
+
+    # Skip L1-style "Simple Test" steps in test-merge.yml (group or "Simple ·" labels)
+    bash tools/run_merge_jobs.sh --skip-simple
+    bash tools/run_merge_jobs.sh --skip-simple --model-type omni
+
+    # Preview extracted commands without running
+    bash tools/run_merge_jobs.sh --dry-run
+
+    # Filter by model area (OR semantics): omni | tts | diffusion | all
+    bash tools/run_merge_jobs.sh --model-type diffusion
+    bash tools/run_merge_jobs.sh --model-type omni,tts
+
+    # Match a Buildkite label substring
+    bash tools/run_merge_jobs.sh --model-type diffusion --label-substr "Wan22"
+
     ```
-    If you only want to run a specific test case, you can use:
+
+    **Ad hoc:** run a single test file or marker expression manually:
+
     ```bash
     pytest -s -v test_xxxx.py --run-level=advanced_model
+    pytest -s -v -m "advanced_model and distributed_cuda and L4" --run-level=advanced_model
     ```
-    If you only want to run specific test cases on a particular platform, you can use:
-    ```bash
-    pytest -s -v -m "advanced_model and distributed_cuda and L4"  --run-level=advanced_model
-    ```
-    The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
 
 
 === "L4 level"
 
+    **Recommended:** run CI-aligned nightly jobs from the repo root via [`tools/nightly/run_nightly_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/nightly/run_nightly_jobs.sh). The script reads [`.buildkite/test-nightly.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml), generates per-step bash wrappers, runs pytest (perf jobs first, then aggregates perf JSON into Excel), and tees logs under `logs/nightly_jobs/` (see [CI job runners](#ci-job-runners-l2l4-logs-timing-and-timeouts); `generate_nightly_perf_excel` is included in the timing summary). Requires `bash`, `python3`, and PyYAML (`pip install pyyaml`).
+
     ```bash
-    cd tests
-    pytest -s -v -m "full_model" --run-level=full_model
+    # All nightly jobs (default: test-type all, model-type all)
+    bash tools/nightly/run_nightly_jobs.sh
+
+    # Preview extracted commands without running
+    bash tools/nightly/run_nightly_jobs.sh --dry-run
+
+    # Test kind (--test-type, OR semantics): perf | acc | function | stability | local | all
+    bash tools/nightly/run_nightly_jobs.sh --test-type function
+    bash tools/nightly/run_nightly_jobs.sh --test-type perf,acc
+
+    # Model area (--model-type, OR semantics): omni | tts | diffusion | all
+    bash tools/nightly/run_nightly_jobs.sh --test-type function --model-type omni
+    bash tools/nightly/run_nightly_jobs.sh --test-type perf --model-type diffusion
+
+    # Stability scripts (tests/dfx/stability/scripts/) and local marker tests
+    bash tools/nightly/run_nightly_jobs.sh --test-type stability --model-type omni
+    bash tools/nightly/run_nightly_jobs.sh --test-type local --model-type tts
     ```
-    If you only want to run a specific test case, you can use:
-    ```bash
-    pytest -s -v test_xxxx.py --run-level=full_model
-    ```
-    If you only want to run specific test cases on a particular platform, you can use:
-    ```bash
-    pytest -s -v -m "full_model and distributed_cuda and L4"  --run-level=full_model
-    ```
-    Note: To run performance tests (defaults to ``test_qwen_omni.json``; use ``--test-config-file tests/dfx/perf/tests/test_tts.json`` for TTS):
-    ```bash
-    pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
-    ```
-    The latest L4 (nightly) test commands use the `full_model` marker and `--run-level full_model` (see [test-nightly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml) and [test-nightly-diffusion.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly-diffusion.yml)). Example:
+
+    Perf steps (label contains ``Perf Test``) run before other jobs; Excel output is written under ``logs/``. Diffusion nightly steps may also appear in [test-nightly-diffusion.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly-diffusion.yml) on CI.
+
+    **Ad hoc:** run a single test file or marker expression manually:
 
     ```bash
     cd tests
+    pytest -s -v test_xxxx.py --run-level=full_model
+    pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model
     pytest -s -v -m "full_model and (omni or tts) and H100" --run-level=full_model
     ```
 
+    To run a standalone perf benchmark (defaults to ``test_qwen_omni.json``; use ``--test-config-file tests/dfx/perf/tests/test_tts.json`` for TTS):
+
+    ```bash
+    pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
+    ```
+
 === "L5 level"
 
     L5 includes stability and reliability testing. Typical commands:

docs/contributing/model/adding_diffusion_model.md

@@ -638,7 +638,7 @@ mkdir -p examples/online_serving/your_model_name
 
 - Script: `examples/offline_inference/your_model_name/end2end.py`
   - Parse args like BAGEL (`--model`, `--modality`, optional `--image-path`, `--steps`, etc.)
-  - Use `from vllm_omni.entrypoints.omni import Omni` (or `OmniDiffusion` if your model is diffusion-only)
+  - Use `from vllm_omni.entrypoints.omni import Omni` for both multi-stage and diffusion-only models
   - Save outputs (images/audio/video/text) with deterministic filenames (e.g., `output_0_0.png`)
 - Doc: `examples/offline_inference/your_model_name/README.md`
   - Include at least one runnable command, e.g.:

docs/design/architecture_overview.md

@@ -13,7 +13,7 @@ This document outlines the architectural design for vLLM-Omni.
 
 The primary goal of the vLLM-Omni project is to build the fastest and easiest-to-use open-source Omni-Modality model inference & serving engine. vLLM-Omni extends the original vLLM, which was created to support large language models for text-based autoregressive (AR) generation tasks. vLLM-Omni is designed to support:
 
-* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, and video, alongside text.
+* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, video, and action trajectories, alongside text.
 * **Non-Autoregressive Structure:** Support model structure beyond autoregressive, especially Diffusion Transformer (DiT), which is widely used in visual and audio generation.
 * **Integration with vLLM Core:** Maintain compatibility and leverage existing vLLM key modules and optimizations where applicable.
 * **Extensibility:** Design a modular and flexible architecture that can easily accommodate new modalities, model architectures, and output formats.

docs/design/module/dit_module.md

@@ -470,7 +470,7 @@ class Attention(nn.Module):
                  role="self", role_category=None, ...):
         # Resolve backend for this role from the active OmniDiffusionConfig
         config = get_current_diffusion_config_or_none()
-        attention_config = config.attention_config if config is not None else None
+        attention_config = config.diffusion_attention_config if config is not None else None
         attn_backend_cls, spec = get_attn_backend_for_role(
             role=role,
             head_size=head_size,
@@ -493,7 +493,7 @@ These backends provide the **kernel implementations** for attention computation.
 
 #### Backend Selection Mechanism
 
-Selection is driven by `AttentionConfig` on `OmniDiffusionConfig`, which carries a global `default` plus a `per_role` map of `AttentionSpec` entries. A role string is resolved in this order:
+Selection is driven by `OmniDiffusionConfig.diffusion_attention_config`, which carries a global `default` plus a `per_role` map of `AttentionSpec` entries. A role string is resolved in this order:
 
 ```python
 def get_attn_backend_for_role(role, head_size, attention_config=None, role_category=None):
@@ -909,7 +909,7 @@ def initialize_model_parallel(
 
 ```
 1. User Request
-   └─> OmniDiffusion.generate(prompt)
+   └─> Omni.generate(prompt)
        └─> Prepare OmniDiffusionRequest
            └─> DiffusionEngine.step(requests)
 

docs/features/custom_pipeline.md

@@ -10,7 +10,7 @@ This guide demonstrates how to use the newly added features for extending vLLM-O
 Three main features enable custom pipeline extension:
 
 1. **`WorkerWrapperBase`**: A wrapper class that enables dynamic worker extension with custom functionality
-2. **`load_format`**: A parameter that controls how diffusion models are loaded, including support for custom pipelines
+2. **`diffusion_load_format`**: A parameter that controls how diffusion models are loaded, including support for custom pipelines
 3. **`CustomPipelineWorkerExtension`**: An extension class that enables pipeline re-initialization with custom implementations
 
 ## Features
@@ -27,15 +27,15 @@ Three main features enable custom pipeline extension:
 
 **Location:** `vllm_omni/diffusion/worker/diffusion_worker.py`
 
-### load_format Parameter
+### diffusion_load_format parameter
 
-The `load_format` parameter controls how diffusion models are loaded. It supports the following values:
+The `diffusion_load_format` parameter controls the initial diffusion model load. The relevant values are:
 
 - **`"default"`**: Standard model loading using the model registry (default behavior)
-- **`"custom_pipeline"`**: Load a custom pipeline class specified by `custom_pipeline_name`
-- **`"dummy"`**: Skip model loading (useful for testing or when pipeline will be initialized separately)
+- **`"dummy"`**: Skip the initial model load; use this with `custom_pipeline_args["pipeline_class"]` for the custom pipeline workflow below
+- **`"diffusers"`**: Load through the Hugging Face Diffusers adapter
 
-**Location:** `vllm_omni/diffusion/model_loader/diffusers_loader.py`
+**Location:** `vllm_omni/diffusion/data.py` (`OmniDiffusionConfig`)
 
 ### CustomPipelineWorkerExtension
 

docs/getting_started/installation/gpu.md

@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, and 0.24.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, and 0.25.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

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

@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.24.0 --torch-backend=auto
+uv pip install vllm==0.25.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.24.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library
 
 ```bash
-uv pip install vllm==0.24.0 --torch-backend=auto
+uv pip install vllm==0.25.0 --torch-backend=auto
 ```
 
-The 0.24.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.25.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.24.0
+git checkout v0.25.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0 (the default for v0.24.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.24.0/vllm-0.24.0-cp38-abi3-manylinux_2_28_x86_64.whl
+#For CUDA 13.0 (the default for v0.25.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.25.0/vllm-0.25.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.24.0 \
+    vllm/vllm-omni:v0.25.0 \
     vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 

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

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722
+uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.24.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library
 
 ```bash
-uv pip install vllm==0.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722
+uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
 ```
 
-The pre-built 0.24.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.25.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.24.0
+git checkout v0.25.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.24.0 \
+  vllm/vllm-omni-rocm:v0.25.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -149,7 +149,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.24.0
+  vllm/vllm-omni-rocm:v0.25.0
 ```
 
 # --8<-- [end:pre-built-images]

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

@@ -9,8 +9,8 @@ The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built
 
 ```bash
 # vLLM-Ascend image
-# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.24.0
-# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.24.0-a3
+# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.25.0
+# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.25.0-a3
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \
@@ -28,11 +28,11 @@ docker run --rm \
     -v /etc/ascend_install.info:/etc/ascend_install.info \
     -v /root/.cache:/root/.cache \
     -p 8000:8000 \
-    -it quay.io/atlas-ci/vllm-ascend:0.24.0-a3 bash
+    -it quay.io/atlas-ci/vllm-ascend:0.25.0-a3 bash
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.24.0 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.25.0 https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 pip install -v -e . --no-build-isolation
 # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
@@ -84,8 +84,8 @@ Supported images as following.
 Here's an example deployment command that has been verified on 4 x NPUs:
 
 ```bash
-# Atlas A2: quay.io/ascend/vllm-omni:v0.24.0
-# Atlas A3: quay.io/ascend/vllm-omni:v0.24.0-a3
+# Atlas A2: quay.io/ascend/vllm-omni:v0.25.0
+# Atlas A3: quay.io/ascend/vllm-omni:v0.25.0-a3
 docker run --rm \
     --name vllm-omni-a3 \
     --shm-size=4g \
@@ -110,7 +110,7 @@ docker run --rm \
     -v /etc/ascend_install.info:/etc/ascend_install.info \
     -v ~/.cache:/root/.cache \
     -p 8091:8091 \
-    -it quay.io/ascend/vllm-omni:v0.24.0-a3 bash
+    -it quay.io/ascend/vllm-omni:v0.25.0-a3 bash
 ```
 
 !!! tip

docs/getting_started/quickstart.md

@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.24.0 --torch-backend=auto
+uv pip install vllm==0.25.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722
+uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -34,7 +34,7 @@ For additional installation methods — please see the [installation guide](inst
 !!! note
     It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
 
-    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.24.0` and vLLM Omni `0.24.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
+    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.25.0` and vLLM Omni `0.25.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
 
 ## Offline Inference
 

docs/mkdocs/hooks/generate_examples.py

@@ -16,8 +16,35 @@ ROOT_DIR_RELATIVE = "../../../../.."
 EXAMPLE_DIR = ROOT_DIR / "examples"
 EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
 NAV_FILE = ROOT_DIR / "docs/.nav.yml"
+MODEL_DISPLAY_NAMES_FILE = EXAMPLE_DIR / "model_display_names.yml"
 MAX_INLINE_MATERIAL_SIZE = 8 * 1024
 
+SERVING_MODE_TITLES = {
+    "offline_inference": "Offline inference",
+    "online_serving": "Online serving",
+}
+
+
+def load_model_display_names() -> dict[str, str]:
+    try:
+        with open(MODEL_DISPLAY_NAMES_FILE, encoding="utf-8") as f:
+            model_display_names = yaml.safe_load(f)
+    except (OSError, yaml.YAMLError) as exc:
+        raise ValueError(f"Failed to load MODEL_DISPLAY_NAMES_FILE at {MODEL_DISPLAY_NAMES_FILE}: {exc}") from exc
+
+    if not isinstance(model_display_names, dict) or not all(
+        isinstance(key, str) and isinstance(value, str) for key, value in model_display_names.items()
+    ):
+        raise ValueError(
+            "MODEL_DISPLAY_NAMES_FILE "
+            f"at {MODEL_DISPLAY_NAMES_FILE} must contain a YAML mapping "
+            "of string model IDs to string display names"
+        )
+    return model_display_names
+
+
+MODEL_DISPLAY_NAMES = load_model_display_names()
+
 
 def fix_case(text: str) -> str:
     subs = {
@@ -54,6 +81,7 @@ class Example:
         main_file (Path): The main file in the directory.
         other_files (list[Path]): list of other files in the directory.
         title (str): The title of the document.
+        nav_title (str): The concise title used in navigation.
 
     Methods:
         __post_init__(): Initializes the main_file, other_files, and title attributes.
@@ -78,6 +106,16 @@ class Example:
     def is_code(self) -> bool:
         return self.main_file.suffix != ".md"
 
+    @property
+    def model_display_name(self) -> str | None:
+        if self.category not in SERVING_MODE_TITLES:
+            return None
+        return MODEL_DISPLAY_NAMES.get(self.path.stem)
+
+    @property
+    def nav_title(self) -> str:
+        return self.model_display_name or self.title
+
     def determine_main_file(self) -> Path:
         """
         Determines the main file in the given path.
@@ -151,13 +189,25 @@ class Example:
         return [file for file in self.path.rglob("*") if is_other_file(file)]
 
     def determine_title(self) -> str:
+        model_display_name = self.model_display_name
         if not self.is_code:
             # Specify encoding for building on Windows
             with open(self.main_file, encoding="utf-8") as f:
                 first_line = f.readline().strip()
             match = re.match(r"^#\s+(?P<title>.+)$", first_line)
+            if model_display_name:
+                expected_title = f"{model_display_name}: {SERVING_MODE_TITLES[self.category]}"
+                actual_title = match.group("title") if match else first_line
+                if actual_title != expected_title:
+                    raise ValueError(
+                        f"Model example title mismatch in {self.main_file}: "
+                        f"expected '# {expected_title}', got {first_line!r}"
+                    )
+                return expected_title
             if match:
                 return match.group("title")
+        elif model_display_name:
+            raise ValueError(f"Mapped model example must use a Markdown README: {self.path}")
         return fix_case(self.path.stem.replace("_", " ").title())
 
     def fix_relative_links(self, content: str) -> str:
@@ -300,7 +350,7 @@ def update_nav_file(examples: list[Example]):
         for example in category_examples:
             doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
             rel_path = doc_path.relative_to(ROOT_DIR / "docs")
-            category_items.append({example.title: rel_path.as_posix()})
+            category_items.append({example.nav_title: rel_path.as_posix()})
 
         if category_items:
             # Format category name (e.g., "offline_inference" -> "Offline Inference")

docs/models/supported_models.md

@@ -34,6 +34,7 @@ th {
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | |
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | |
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |

docs/serving/audio_generate_api.md

@@ -66,7 +66,7 @@ Content-Type: application/json
 |-----------|------|---------|-------------|
 | `input` | string | **required** | Text prompt describing the audio to generate |
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
 | `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
 
 #### Diffusion Parameters
@@ -90,7 +90,6 @@ Returns binary audio data with the appropriate `Content-Type` header:
 | `mp3` | `audio/mpeg` |
 | `flac` | `audio/flac` |
 | `pcm` | `audio/pcm` |
-| `aac` | `audio/aac` |
 | `opus` | `audio/opus` |
 
 ## Examples
@@ -285,7 +284,7 @@ Pydantic validation failure (e.g. invalid `response_format`, `speed` out of rang
     "detail": [
         {
             "type": "literal_error",
-            "msg": "Input should be 'wav', 'pcm', 'flac', 'mp3', 'aac' or 'opus'",
+            "msg": "Input should be 'wav', 'pcm', 'flac', 'mp3' or 'opus'",
             ...
         }
     ]

docs/serving/image_edit_api.md

@@ -252,7 +252,7 @@ Validation errors (missing required fields):
 
 ```bash
 # Check if server is responding
-curl -X http://localhost:8000/v1/images/edit \
+curl http://localhost:8000/v1/images/edits \
   -F "prompt='test'"
 ```
 

docs/serving/image_generation_api.md

@@ -138,7 +138,7 @@ Content-Type: application/json
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
 | `n` | integer | 1 | Number of images to generate (1-10) |
 | `size` | string | model defaults | Image dimensions in WxH format (e.g., "1024x1024", "512x512") |
-| `response_format` | string | "b64_json" | Response format (only "b64_json" supported) |
+| `response_format` | string | "b64_json" | Response format (`"b64_json"` or `"file"`) |
 | `user` | string | null | User identifier for tracking |
 
 #### vllm-omni Extension Parameters

docs/serving/speech_api.md

@@ -109,7 +109,7 @@ Content-Type: application/json
 | `input` | string | **required** | The text to synthesize into speech |
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
 | `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
 | `speed` | float | 1.0 | Playback speed (0.25-4.0) |
 
 #### vLLM-Omni Extension Parameters

docs/user_guide/diffusion/attention_backends.md

@@ -83,7 +83,7 @@ When constructing `OmniDiffusionConfig` directly:
 from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
 
 config = OmniDiffusionConfig(
-    attention_config=AttentionConfig(
+    diffusion_attention_config=AttentionConfig(
         default=AttentionSpec(backend="FLASH_ATTN"),
         per_role={
             "cross": AttentionSpec(backend="TORCH_SDPA"),

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -33,7 +33,7 @@ m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)
 
 **CLI:**
 ```bash
-vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-cpu-offload
 ```
 
 ### To Support a Model
@@ -82,6 +82,33 @@ mutual exclusion: when one group runs, the other moves to CPU.
 - Support single GPU only for now
 
 
+### Component offloading for split models (e.g. Cosmos3)
+
+Some models split their transformer into mutually-exclusive *components* that run
+in different phases of a single forward pass rather than as separate pipeline
+components -- e.g. Cosmos3's understanding (reasoner) component runs once per
+generation while the generation (generator) component runs every denoising step.
+Such models have no separate text encoder to swap against, so the transformer
+owns a small model-local offload path and wraps each phase with
+`with self._offload_context(name):`
+
+```python
+class Cosmos3VFMTransformer(nn.Module):
+    def forward(self, ...):
+        with self._offload_context("reasoner"):
+            ...  # understanding pass, runs once
+        with self._offload_context("generator"):
+            ...  # denoising pass, runs every step
+```
+
+Model-level offloading then keeps exactly one component GPU-resident at a time
+(the other on CPU), reusing the same `SequentialOffloadHook` `.to()` movers. The
+pipeline opts in by exposing `enable_omni_model_cpu_offload` (which drives the
+transformer's `enable_model_cpu_offload` and pins the VAE). Layerwise offloading
+works for these models too -- each component declares its own block container via
+`_layerwise_offload_blocks_attrs`.
+
+
 ## Layerwise (Blockwise) Offloading
 
 ### How It Works
@@ -119,10 +146,10 @@ m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True
 **CLI:**
 ```bash
 # Text-to-video
-vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-layerwise-offload
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-layerwise-offload
 
 # Or image-to-video
-vllm-omni serve diffusion Wan-AI/Wan2.2-I2V-A14B-Diffusers --enable-layerwise-offload
+vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --enable-layerwise-offload
 ```
 
 ### To Support a Model
@@ -145,9 +172,9 @@ class Flux2Transformer2DModel(nn.Module):
 ```
 
 ### Limitations
-- Cold start latency increases because of
-    1) components are loaded to CPU first at the very first during initialization,
-    2) weight consolidation and pinning
+- Cold start latency increases because offloaded components must be moved to CPU
+  during setup; layerwise offload may add extra weight consolidation and pinning
+  work.
 - Performance depends on compute cost and H2D bandwidth as well
 - Support single GPU only for now
 
@@ -178,11 +205,18 @@ Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelH
 
 ```
 OffloadBackend (base class)
-├── ModelLevelOffloadBackend → uses SequentialOffloadHook
+├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap)
+│                              (delegates to a pipeline's enable_omni_model_cpu_offload
+│                               for split models like Cosmos3)
 └── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
 ```
 
-Factory function `get_offload_backend()` selects the appropriate backend based on configuration.
+Factory function `get_offload_backend()` selects the appropriate backend based on
+configuration.
+
+For split models, `ModelLevelOffloadBackend.enable()` detects a pipeline's
+`enable_omni_model_cpu_offload` hook and delegates to it; Cosmos3 then swaps its
+reasoner/generator components inside the model forward pass.
 
 
 ## Supported Models

docs/user_guide/diffusion/frame_interpolation.md

@@ -47,7 +47,8 @@ For supported Wan2.2 pipelines, the execution order is:
 1. Diffusion worker finishes denoising and decodes the raw video tensor.
 2. Worker-side model-specific post-processing runs.
 3. If frame interpolation is enabled, RIFE interpolates the decoded video
-   tensor on the worker side and records a FPS multiplier in `custom_output`.
+   tensor on the worker side and records a FPS multiplier in
+   `metadata.video.video_fps_multiplier`.
 4. The API server receives the already-interpolated video and only performs
    MP4 export.
 

docs/user_guide/diffusion_features.md

@@ -105,7 +105,7 @@ The following tables show which models support each feature:
 
 > Notes:
 
-> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only.
+> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload_diffusion.md).
 > 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](quantization/overview.md) for per-method and per-model support details.
 
 ### ImageGen
@@ -121,6 +121,7 @@ The following tables show which models support each feature:
 | **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Hidream-I1-Full**        |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ❌         |
+| **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **MagiHuman**            |     ❌     |     ❌      |           ❌           |       ❓        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
@@ -143,7 +144,8 @@ The following tables show which models support each feature:
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
 > 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
-> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload is not supported; use layerwise offload.
+> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload.
+> 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
 
 ### VideoGen
 

docs/user_guide/examples/offline_inference/bagel.md

@@ -257,7 +257,7 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
 | :------- | :--- | :------ | :---------- |
 | `--deploy-config` | string | `None` | Path to deploy YAML (auto-detected if omitted) |
 | `--stage-configs-path` | string | `None` | [Deprecated] Legacy path to `stage_args` YAML; prefer `--deploy-config` |
-| `--worker-backend` | choice | `process` | `process` or `ray` |
+| `--worker-backend` | choice | `multi_process` | `multi_process` or `ray` |
 | `--ray-address` | string | `None` | Ray cluster address |
 | `--quantization` | string | `None` | Quantization method (e.g. `fp8`) |
 | `--log-stats` | flag | `False` | Enable statistics logging |

docs/user_guide/examples/offline_inference/glm_image.md

@@ -26,18 +26,22 @@ Stage 0 (AR Model)                Stage 1 (Diffusion)
 
 ```python
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 if __name__ == "__main__":
     omni = Omni(model="zai-org/GLM-Image")
     outputs = omni.generate(
         "A photorealistic mountain landscape at sunset",
-        sampling_params={
-            "height": 1024,
-            "width": 1024,
-            "num_inference_steps": 50,
-            "guidance_scale": 1.5,
-            "seed": 42,
-        },
+        sampling_params_list=[
+            omni.default_sampling_params_list[0],
+            OmniDiffusionSamplingParams(
+                height=1024,
+                width=1024,
+                num_inference_steps=50,
+                guidance_scale=1.5,
+                seed=42,
+            ),
+        ],
     )
     outputs[0].request_output.images[0].save("output.png")
 ```
@@ -46,6 +50,7 @@ if __name__ == "__main__":
 
 ```python
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 if __name__ == "__main__":
     omni = Omni(model="zai-org/GLM-Image")
@@ -56,13 +61,16 @@ if __name__ == "__main__":
                 "image": "input.png",
             },
         },
-        sampling_params={
-            "height": 1024,
-            "width": 1024,
-            "num_inference_steps": 50,
-            "guidance_scale": 1.5,
-            "seed": 42,
-        },
+        sampling_params_list=[
+            omni.default_sampling_params_list[0],
+            OmniDiffusionSamplingParams(
+                height=1024,
+                width=1024,
+                num_inference_steps=50,
+                guidance_scale=1.5,
+                seed=42,
+            ),
+        ],
     )
     outputs[0].request_output.images[0].save("output.png")
 ```

docs/user_guide/examples/offline_inference/mimo_audio.md

@@ -1,4 +1,4 @@
-# MiMo-Audio Offline Inference
+# MiMo-Audio: Offline inference
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/mimo_audio>.
 

docs/user_guide/examples/offline_inference/qwen2_5_omni.md

@@ -1,4 +1,4 @@
-# Qwen2.5-Omni
+# Qwen2.5-Omni: Offline inference
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen2_5_omni>.
 

docs/user_guide/examples/offline_inference/qwen3_omni.md

@@ -1,4 +1,4 @@
-# Qwen3-Omni
+# Qwen3-Omni: Offline inference
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen3_omni>.
 

docs/user_guide/examples/online_serving/aura_omni.md

@@ -1,4 +1,4 @@
-# AURA Omni Pipeline
+# AURA Omni: Online serving
 
 `aura_omni` wires ASR, AURA, and Qwen3-TTS into one vLLM-Omni pipeline:
 

docs/user_guide/examples/online_serving/mimo_audio.md

@@ -1,4 +1,4 @@
-# Online serving Example of vLLM-Omni for MiMo-Audio
+# MiMo-Audio: Online serving
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/mimo_audio>.
 

docs/user_guide/examples/online_serving/qwen2_5_omni.md

@@ -1,4 +1,4 @@
-# Qwen2.5-Omni
+# Qwen2.5-Omni: Online serving
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen2_5_omni>.
 

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -1,4 +1,4 @@
-# Qwen3-Omni
+# Qwen3-Omni: Online serving
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_omni>.
 

docs/user_guide/examples/online_serving/text_to_audio.md

@@ -139,7 +139,7 @@ Content-Type: application/json
 |-----------|------|---------|-------------|
 | `input` | string | **required** | Text prompt describing the audio to generate |
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
 | `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
 | `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) |
 | `audio_start` | float | 0.0 | Audio start time in seconds |

docs/user_guide/quantization/autoround.md

@@ -3,7 +3,7 @@
 ## Overview
 
 [AutoRound](https://github.com/intel/auto-round) produces pre-quantized
-checkpoints for LLMs, VLMs, and diffusion models. vLLM-Omni reads the
+checkpoints for LLMs, VLMs, diffusion models, and world models. vLLM-Omni reads the
 checkpoint's `config.json` and auto-detects
 `quantization_config.quant_method = "auto-round"`.
 
@@ -36,6 +36,13 @@ guide. AutoRound is Intel-supported.
 | Wan2.2-T2V | `Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
 | Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
 
+### World Model (Cosmos3)
+
+| Model | Checkpoint | Scope | Scheme | Backend |
+|-------|------------|-------|--------|---------|
+| Cosmos3-Nano | `Intel/Cosmos3-Nano-int4-AutoRound` | World-model transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
+| Cosmos3-Super | `Intel/Cosmos3-Super-int4-AutoRound` | World-model transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
+
 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
 
 | Model | Checkpoint | Scope | Scheme | Backend |

docs/user_guide/quantization/bitsandbytes.md

@@ -0,0 +1,120 @@
+# BitsAndBytes W4 Quantization
+
+## Overview
+
+BitsAndBytes 4-bit quantization supports weight-only NF4/FP4 diffusion transformer
+inference on CUDA GPUs. It quantizes BF16/FP16 weights at load time; activations
+stay in BF16/FP16. No pre-quantized checkpoint is required.
+
+This is an online (dynamic) path: load the normal HuggingFace checkpoint and
+quantize during model loading via the `bitsandbytes` CUDA kernels.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA CUDA GPU (SM 75+) | ✅ |
+| NVIDIA Blackwell GPU (SM 100+) | ✅ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ |
+| AMD ROCm | ❌ |
+| Intel XPU | ❌ |
+| Ascend NPU | ❌ |
+
+Legend: `✅` supported, `❌` unsupported. Non-CUDA platforms raise in
+`get_quant_method()`; CUDA requires compute capability 7.5+ (`SM 75+`).
+
+Requires the optional `bitsandbytes` package (`pip install bitsandbytes`).
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Model | HF models | CUDA | Mode | Recommendation |
+|-------|-----------|:----:|------|----------------|
+| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Online W4 weight-only | All heavy linear layers; sensitive embedders stay BF16 |
+| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Not validated | Online W4 weight-only | Compare vs BF16 before enabling |
+| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Online W4 weight-only | Validate before enabling in docs |
+
+Other diffusion models may work if their transformer uses supported linear
+layers, but they are not validated in this guide.
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| Qwen3-Omni | Thinker language-model stage | Not validated | Prefer checkpoint-supported ModelOpt FP8 or AutoRound paths |
+| Qwen3-TTS | TTS language-model stage | Not validated | No BitsAndBytes TTS stage support is documented |
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| BAGEL | Stage-specific transformer or DiT module | Not validated | Requires explicit stage routing |
+| GLM-Image | Stage-specific transformer or DiT module | Not validated | Requires quality comparison with BF16 |
+
+## Configuration
+
+Python API:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(model="Tongyi-MAI/Z-Image-Turbo", quantization="bitsandbytes")
+
+omni_with_skips = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    quantization_config={
+        "method": "bitsandbytes",
+        "ignored_layers": ["to_out"],
+    },
+)
+
+outputs = omni.generate(
+    "A cup of coffee on the table",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+CLI:
+
+```bash
+python text_to_image.py --model Tongyi-MAI/Z-Image-Turbo --quantization bitsandbytes
+python text_to_image.py --model Tongyi-MAI/Z-Image-Turbo --quantization bitsandbytes --ignored-layers "to_out"
+vllm serve Tongyi-MAI/Z-Image-Turbo --omni --quantization bitsandbytes
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `method` | str | - | Quantization method (`"bitsandbytes"`) |
+| `quant_type` | str | `"nf4"` | 4-bit data type: `"nf4"` (recommended) or `"fp4"` |
+| `compress_statistics` | bool | `True` | Double-quantize block scaling statistics for better accuracy |
+| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16/FP16 |
+
+## Validation and Notes
+
+On Z-Image-Turbo (single GPU, 1024×1024, 50 steps), BitsAndBytes W4 typically
+reduces peak VRAM from roughly 24.5 GiB (BF16) to roughly 17 GiB. Compare output
+quality against a BF16 baseline before enabling on new models.
+
+Multi-GPU tensor parallelism (`tensor_parallel_size` > 1) is not validated for
+BitsAndBytes in diffusion models; each rank quantizes its own weight shard
+independently.
+
+If quality regresses, use `ignored_layers` to keep sensitive projections in BF16
+(for example `to_out` or `w2`).
+
+```python
+omni = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    quantization="bitsandbytes",
+    cache_backend="tea_cache",
+    cache_config={"rel_l1_thresh": 0.2},
+)
+```
+
+Only add a new model to the supported table after comparing the BitsAndBytes
+output against a BF16 baseline and documenting any required `ignored_layers`.

examples/model_display_names.yml

@@ -0,0 +1,12 @@
+# Canonical display names for model examples that have both offline and online pages.
+# The docs generator combines these names with the serving mode for each page H1.
+aura_omni: AURA Omni
+covo_audio: Covo-Audio-Chat
+dynin_omni: Dynin-Omni
+lance: Lance
+mimo_audio: MiMo-Audio
+ming_flash_omni: Ming-flash-omni 2.0
+qwen2_5_omni: Qwen2.5-Omni
+qwen3_omni: Qwen3-Omni
+sensenova_u1: SenseNova-U1
+step_audio2: Step-Audio2

examples/offline_inference/aura_omni/README.md

@@ -1,4 +1,4 @@
-# AURA Omni
+# AURA Omni: Offline inference
 
 This example runs the native AURA Omni pipeline offline:
 

examples/offline_inference/covo_audio/README.md

@@ -1,4 +1,4 @@
-# Covo-Audio-Chat (Offline Inference)
+# Covo-Audio-Chat: Offline inference
 
 ## Setup
 

examples/offline_inference/dynin_omni/README.md

@@ -1,4 +1,4 @@
-# Dynin-Omni Offline End2End Example
+# Dynin-Omni: Offline inference
 
 This folder contains a unified offline inference entrypoint:
 

examples/offline_inference/lance/README.md

@@ -1,4 +1,4 @@
-# Lance (ByteDance) — offline inference
+# Lance: Offline inference
 
 [Lance](https://huggingface.co/bytedance-research/Lance) is a 3B unified
 autoregressive + diffusion multimodal model on a Qwen2.5-VL backbone. It is

examples/offline_inference/lance/end2end.py

@@ -293,16 +293,19 @@ def main():
 
     if args.modality in ("img2text", "video2text"):
         for i, req_output in enumerate(outputs):
-            text = getattr(req_output, "output", None) or getattr(req_output, "text", None)
-            if text is None and hasattr(req_output, "custom_output"):
-                text = (req_output.custom_output or {}).get("text_output")
+            multimodal_output = getattr(req_output, "multimodal_output", {}) or {}
+            metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {}
+            text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {}
+            text = (multimodal_output.get("text") if isinstance(multimodal_output, dict) else None) or (
+                text_metadata.get("text_output") if isinstance(text_metadata, dict) else None
+            )
+            text = text or getattr(req_output, "output", None) or getattr(req_output, "text", None)
             print(f"[Output {i}] {text}")
         return
 
     if args.modality in ("text2video", "video2video", "image2video"):
         for i, req_output in enumerate(outputs):
-            custom = getattr(req_output, "custom_output", None) or {}
-            frames = custom.get("video_frames")
+            frames = getattr(req_output, "images", None)
             if not frames:
                 print(f"[Output {i}] no video frames returned")
                 continue

examples/offline_inference/lance/gradio_demo.py

@@ -835,12 +835,15 @@ def run_task(
     print(f"[unified] task={task} {width}x{height} frames={num_frames} cfg={cfg} steps={steps} seed={seed}", flush=True)
 
     outputs = list(omni.generate(prompts=formatted, sampling_params_list=params))
-    custom = getattr(outputs[0], "custom_output", None) or {}
-    # Lance/Bagel return text at outputs[0].output (string) or
-    # custom_output["text_output"].
+    multimodal_output = getattr(outputs[0], "multimodal_output", {}) or {}
+    metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {}
+    text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {}
     text_out = None
     if task in {"x2t_image", "x2t_video"}:
-        text_out = custom.get("text_output") or getattr(outputs[0], "output", None)
+        text_out = (multimodal_output.get("text") if isinstance(multimodal_output, dict) else None) or (
+            text_metadata.get("text_output") if isinstance(text_metadata, dict) else None
+        )
+        text_out = text_out or getattr(outputs[0], "output", None)
         if not isinstance(text_out, str):
             text_out = None
 
@@ -856,13 +859,11 @@ def run_task(
             gr.update(value=str(text_out or "(empty)"), visible=True),
         )
     if task in {"t2v", "i2v", "video_edit"}:
-        frames = custom.get("video_frames")
+        frames = getattr(outputs[0], "images", None)
         if frames is None:
             raise gr.Error("No video frames in output.")
         return (gr.update(visible=False), gr.update(value=_save_video(frames), visible=True), gr.update(visible=False))
-    # image tasks — orchestrator strips ``outputs[0].output`` across the
-    # IPC boundary but preserves ``custom_output``, so prefer the latter.
-    img_out = custom.get("image") or getattr(outputs[0], "output", None)
+    img_out = (getattr(outputs[0], "images", None) or [None])[0] or getattr(outputs[0], "output", None)
     if img_out is None:
         raise gr.Error("No image in output.")
     return (gr.update(value=_save_image(img_out), visible=True), gr.update(visible=False), gr.update(visible=False))

examples/offline_inference/mimo_audio/README.md

@@ -1,4 +1,4 @@
-# MiMo-Audio Offline Inference
+# MiMo-Audio: Offline inference
 
 This directory contains an offline demo for running MiMo-Audio models with vLLM Omni. It builds task-specific inputs and generates WAV files or text outputs locally.
 

examples/offline_inference/ming_flash_omni/README.md

@@ -1,4 +1,4 @@
-# Ming-flash-omni 2.0
+# Ming-flash-omni 2.0: Offline inference
 
 [Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with text and speech outputs.
 

examples/offline_inference/qwen2_5_omni/README.md

@@ -1,4 +1,4 @@
-# Qwen2.5-Omni
+# Qwen2.5-Omni: Offline inference
 
 ## Setup
 Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.

examples/offline_inference/qwen3_omni/README.md

@@ -1,4 +1,4 @@
-# Qwen3-Omni
+# Qwen3-Omni: Offline inference
 
 ## Setup
 Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.

examples/offline_inference/sensenova_u1/README.md

@@ -1,4 +1,4 @@
-# SenseNova-U1-8B-MoT
+# SenseNova-U1: Offline inference
 
 ## Architecture
 

examples/offline_inference/sensenova_u1/end2end.py

@@ -278,12 +278,17 @@ def main():
     )
 
     for req_output in outputs:
-        custom = getattr(req_output, "_custom_output", {}) or {}
-        if args.print_think and custom.get("think_text"):
-            print(f"[Think]\n{custom['think_text']}\n")
+        multimodal_output = getattr(req_output, "multimodal_output", {}) or {}
+        metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {}
+        text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {}
+        think_text = text_metadata.get("think_text") if isinstance(text_metadata, dict) else None
+        if args.print_think and think_text:
+            print(f"[Think]\n{think_text}\n")
 
         if is_text_output:
-            text = custom.get("text_output", "")
+            text = multimodal_output.get("text", "") if isinstance(multimodal_output, dict) else ""
+            if not text:
+                text = text_metadata.get("text_output", "") if isinstance(text_metadata, dict) else ""
             if not text:
                 text = getattr(req_output, "text", "") or ""
             print(f"[Response]\n{text}")

examples/offline_inference/step_audio2/README.md

@@ -1,4 +1,4 @@
-# Step-Audio2 Offline Inference Examples
+# Step-Audio2: Offline inference
 
 This directory contains examples for running offline inference with Step-Audio2 using vLLM-Omni.
 

examples/offline_inference/text_to_image/README.md

@@ -36,6 +36,7 @@ This folder provides several entrypoints for experimenting with text-to-image di
 | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
 | `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3)  | 160 |
 | `HiDream-I1-Full` | 1024 x 1024 | 63.7  | 57.7 |
+| `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | 1024 (Raw) / 2048 (Turbo) | — | ~30 |
 
 !!! info
 *Peak VRAM:  based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU.
@@ -196,6 +197,37 @@ python examples/offline_inference/text_to_image/text_to_image.py \
   --output /output.png
 ```
 
+### Krea 2
+
+Krea 2 is published as diffusers-format checkpoints: `krea/Krea-2-Turbo` (few-step distilled) and `krea/Krea-2-Raw`.
+The pipeline class (`Krea2Pipeline`) is auto-detected from `model_index.json`. The distilled Turbo checkpoint sets
+`is_distilled=true` and is best run at 2048x2048 with few steps and guidance disabled:
+
+```bash
+# Few-step distilled (Turbo) checkpoint
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model krea/Krea-2-Turbo \
+  --prompt "a fox in the snow" \
+  --num-inference-steps 8 \
+  --guidance-scale 0.0 \
+  --height 2048 --width 2048 \
+  --output krea2.png
+```
+
+For the Raw checkpoint, use more steps with CFG enabled:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model krea/Krea-2-Raw \
+  --prompt "a fox in the snow" \
+  --num-inference-steps 28 \
+  --guidance-scale 4.5 \
+  --output krea2.png
+```
+
+`guidance-scale` follows the Krea 2 convention (`velocity = cond + guidance_scale * (cond - uncond)`); CFG is enabled
+whenever `guidance-scale > 0`.
+
 ### Batch Requests (Multiple Prompts)
 
 You can pass multiple prompts in a single `Omni.generate` call. `Omni`

examples/offline_inference/text_to_image/text_to_image.py

@@ -202,9 +202,10 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8", "int8"],
+        choices=["fp8", "int8", "bitsandbytes"],
         help="Quantization method for the transformer. "
-        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8). "
+        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), "
+        "'int8' (Int8 W8A8), 'bitsandbytes' (NF4 4-bit weight-only, CUDA SM 75+). "
         "Default: None (no quantization, uses BF16).",
     )
     parser.add_argument(

examples/online_serving/aura_omni/README.md

@@ -1,4 +1,4 @@
-# AURA Omni Native Pipeline
+# AURA Omni: Online serving
 
 `aura_omni` serves AURA as a native multi-stage vLLM-Omni pipeline:
 

examples/online_serving/covo_audio/README.md

@@ -1,4 +1,4 @@
-# Covo-Audio-Chat
+# Covo-Audio-Chat: Online serving
 
 ## Installation
 

examples/online_serving/dynin_omni/README.md

@@ -1,4 +1,4 @@
-# Dynin-Omni Online Serving Example
+# Dynin-Omni: Online serving
 
 ## Installation
 

examples/online_serving/lance/README.md

@@ -1,4 +1,4 @@
-# Lance — online serving
+# Lance: Online serving
 
 OpenAI-compatible chat completions API for Lance, served by `vllm-omni`.
 

examples/online_serving/mimo_audio/README.md

@@ -1,4 +1,4 @@
-# Online serving Example of vLLM-Omni for MiMo-Audio
+# MiMo-Audio: Online serving
 
 ## 🛠️ Installation
 

examples/online_serving/ming_flash_omni/README.md

@@ -1,4 +1,4 @@
-# Ming-flash-omni 2.0
+# Ming-flash-omni 2.0: Online serving
 
 ## Installation
 

examples/online_serving/qwen2_5_omni/README.md

@@ -1,4 +1,4 @@
-# Qwen2.5-Omni
+# Qwen2.5-Omni: Online serving
 
 ## 🛠️ Installation
 

examples/online_serving/qwen3_omni/README.md

@@ -1,4 +1,4 @@
-# Qwen3-Omni
+# Qwen3-Omni: Online serving
 
 ## 🛠️ Installation
 

examples/online_serving/sensenova_u1/README.md

@@ -1,4 +1,4 @@
-# SenseNova-U1 Online Serving
+# SenseNova-U1: Online serving
 
 ## Launch the Server
 

examples/online_serving/step_audio2/README.md

@@ -1,4 +1,4 @@
-# Step-Audio2 Online Serving
+# Step-Audio2: Online serving
 
 This directory contains examples for running Step-Audio2 with vLLM-Omni's online serving API.
 

recipes/README.md

@@ -44,6 +44,7 @@ recipes/
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
 | [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
 | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
+| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`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 |

recipes/cosmos3/Cosmos3-Nano.md

@@ -422,3 +422,185 @@ ffprobe -v error -show_entries stream=codec_type,nb_frames,width,height cosmos3_
   `generate_sound`/`sound_duration`, `guardrails`, `action_*`, ...) are declared
   once in `vllm_omni/model_extras/cosmos3.py` and forwarded through `--extra-body`;
   unknown keys for the model are dropped.
+
+## NPU
+
+### 1x Ascend 910B / 910C (Atlas A2 / A3) — Online serving
+
+#### Environment
+
+- OS: Linux (aarch64)
+- Python: 3.12+
+- Driver / runtime: CANN 8.5.1 + NNAL + Ascend 910B / 910C
+- vLLM version: match the repository requirements from your current checkout
+- vLLM-Ascend version: match the repository requirements from your current checkout
+- vLLM-Omni version or commit: use the commit you are deploying from
+
+#### Command
+
+Requires the `vllm-omni` package (or the `quay.io/atlas-ci/vllm-ascend` A2 / A3 container),
+which provides the `vllm serve … --omni` entrypoint used below.
+
+Safety guardrails are **on by default** (NVIDIA Open Model License). They load
+the **gated** `nvidia/Cosmos-1.0-Guardrail` model, so to keep them on you must:
+
+1. `pip install cosmos-guardrail`
+2. Accept the license at <https://huggingface.co/nvidia/Cosmos-1.0-Guardrail>
+3. Export a token with access: `export HF_TOKEN=hf_...`
+
+Then launch the recommended server:
+
+```bash
+vllm serve nvidia/Cosmos3-Nano \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --init-timeout 1800
+```
+
+To run **without** guardrails (you are responsible for license compliance),
+add `--no-guardrails` (no token/`cosmos-guardrail` needed). For tensor parallel
+add `--tensor-parallel-size 8`. `--quantization fp8` and
+`--enable-layerwise-offload` are not supported on NPU.
+The pipeline auto-resolves from `model_index.json`; pass
+`--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly.
+
+#### Verification
+
+Best quality uses the JSON-upsampled prompts from `assets/` (download with
+`hf download nvidia/Cosmos3-Nano assets/ --local-dir Cosmos3-Nano`). Minimal
+self-contained examples:
+
+```bash
+curl http://localhost:8000/v1/models
+
+# Text-to-image -> /v1/images/generations  (1024x1024, 10 steps; base64 PNG)
+curl -sS -X POST http://localhost:8000/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "nvidia/Cosmos3-Nano",
+    "prompt": "A photorealistic red sports car on a city street at golden hour, cinematic lighting.",
+    "negative_prompt": "blurry, distorted, low quality",
+    "size": "1024x1024", "n": 1, "response_format": "b64_json",
+    "num_inference_steps": 10, "guidance_scale": 7.0, "seed": 42
+  }' | python -c "import sys,json,base64; open('cosmos3_t2i.png','wb').write(base64.b64decode(json.load(sys.stdin)['data'][0]['b64_json']))"
+
+# Text-to-video -> /v1/videos/sync  (720p, 49 frames @ 24fps)
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=A robot arm is cleaning a plate in the kitchen" \
+  -F "negative_prompt=blurry, distorted, low quality, jittery, deformed" \
+  -F "size=1280x720" -F "num_frames=49" -F "fps=24" \
+  -F "num_inference_steps=20" -F "guidance_scale=6.0" \
+  -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F "seed=123" \
+  -o cosmos3_t2v.mp4
+
+# Image-to-video -> /v1/videos/sync with an uploaded reference image
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=The scene comes to life with smooth, natural motion." \
+  -F "size=1280x720" -F "num_frames=25" -F "fps=8" \
+  -F "num_inference_steps=10" -F "guidance_scale=6.0" \
+  -F "seed=42" \
+  -F "input_reference=@reference.jpg;type=image/jpeg" \
+  -o cosmos3_i2v.mp4
+
+# Video-to-video -> /v1/videos/sync with an uploaded reference video
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
+  -F "size=1280x720" -F "num_frames=17" -F "fps=5" \
+  -F "num_inference_steps=10" -F "guidance_scale=6.0" \
+  -F "seed=42" \
+  -F "input_reference=@reference.mp4;type=video/mp4" \
+  -o cosmos3_v2v.mp4
+```
+
+#### Notes
+
+- **Measured latency (1x Ascend 910B / 910C, bf16, guardrails off):**
+  - T2I 1024² — 10 steps → ~8 s
+  - T2V 1280×720 @ 20 steps — 49 frames → ~55 s
+  - I2V 1280×720 @ 10 steps — 25 frames → ~25 s
+  - V2V 480×320 @ 10 steps — 17 frames → ~12 s
+- **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 NPU;
+  full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer)
+  ~33 GB on disk.
+- **Determinism:** identical seed reproduces identical output on the same
+  hardware; outputs are not bit-identical across different GPU/NPU types.
+- **Supported sizes (per model card):** 256p / 480p / 720p at 16:9, 4:3, 1:1,
+  3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V/V2V
+  1280×720, 35 steps, guidance 6.0, `flow_shift=10.0`.
+- **Key flags / params:** `--no-guardrails` (optional, to disable guardrails), `--init-timeout 1800`
+  (for model loading), `--tensor-parallel-size 8` for multi-NPU, and
+  `--model-class-name Cosmos3OmniDiffusersPipeline` to force the pipeline class.
+- **Known limitations:**
+  - Transfer V2V with `extra_params` (`edge`/`blur`/`depth`/`seg`/`wsm`) hits a
+    resolution-parsing bug; basic V2V without transfer hints works.
+  - FP8 online quantization and layerwise offload are not supported on NPU.
+
+### 1x Ascend 910B / 910C (Atlas A2 / A3) — Offline generation
+
+#### Environment
+
+- OS: Linux (aarch64)
+- Python: 3.12+
+- Driver / runtime: CANN 8.5.1 + NNAL + Ascend 910B / 910C
+- vLLM-Omni version or commit: use the commit you are deploying from
+
+#### Command
+
+The same offline task examples run on NPU; pass model-specific knobs via
+`--extra-body`. Guardrails are on by default — pass `"guardrails": false` for a
+quick local run (install `cosmos-guardrail` + accept the gated repo to enable them).
+
+```bash
+# Text-to-image -> examples/offline_inference/text_to_image
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model nvidia/Cosmos3-Nano \
+  --prompt "A photorealistic red sports car at golden hour, cinematic lighting." \
+  --negative-prompt "blurry, distorted, low quality" \
+  --height 1024 --width 1024 --num-inference-steps 50 --guidance-scale 7.0 \
+  --extra-body '{"flow_shift": 3.0, "guardrails": false}' \
+  --output cosmos3_t2i.png
+
+# Text-to-video -> examples/offline_inference/text_to_video
+python examples/offline_inference/text_to_video/text_to_video.py \
+  --model nvidia/Cosmos3-Nano \
+  --prompt "A robot arm is cleaning a plate in the kitchen." \
+  --negative-prompt "blurry, distorted, low quality, jittery, deformed" \
+  --height 720 --width 1280 --num-frames 189 --fps 24 \
+  --num-inference-steps 35 --guidance-scale 6.0 \
+  --extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false,
+                 "use_resolution_template": false, "use_duration_template": false}' \
+  --output cosmos3_t2v.mp4
+
+# Image-to-video -> examples/offline_inference/image_to_video
+# (Cosmos3 bundles example frames under assets/; any RGB image works too.)
+python examples/offline_inference/image_to_video/image_to_video.py \
+  --model nvidia/Cosmos3-Nano \
+  --image /path/to/Cosmos3-Nano/assets/example_i2v_input.jpg \
+  --prompt "The scene comes to life with smooth, natural motion." \
+  --height 720 --width 1280 --num-frames 189 --fps 24 \
+  --num-inference-steps 35 --guidance-scale 6.0 \
+  --extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}' \
+  --output cosmos3_i2v.mp4
+```
+
+#### Verification
+
+```bash
+python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image', im.size, im.mode)"
+ffprobe -v error -show_entries stream=codec_type,nb_frames,width,height cosmos3_t2v.mp4
+```
+
+#### Notes
+
+- Guardrails are on by default on NPU with `cosmos-guardrail` installed. Pass
+  `"guardrails": false` in `--extra-body` to disable them (there is no
+  `--no-guardrails` flag for the offline scripts).
+- Video at the 189-frame default takes ~15 min/clip on 1 NPU; reduce
+  `--num-frames` for faster iteration.

recipes/krea/Krea-2.md

@@ -0,0 +1,143 @@
+# Krea 2 for text-to-image (base + LoRA)
+
+## Summary
+
+- Vendor: Krea
+- Model: `krea/Krea-2-Turbo` (few-step distilled) and `krea/Krea-2-Raw`
+- Task: Text-to-image (T2I)
+- Mode: Offline engine and online OpenAI-compatible serving
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe for a known-good starting point to run the Krea 2 single-stream
+MMDiT with vLLM-Omni on a **single 80 GB NVIDIA H100** — both the offline
+`text_to_image.py` entrypoint and the online `/v1/images/generations` route,
+including PEFT LoRA adapters. `Krea2Pipeline` is auto-detected from
+`model_index.json`; guidance is checkpoint-aware (distilled Turbo runs no-CFG,
+Raw runs CFG).
+
+## References
+
+- Related example under `examples/`:
+  [`examples/offline_inference/text_to_image/README.md`](../../examples/offline_inference/text_to_image/README.md)
+  (see the "Krea 2" section)
+- Feature support:
+  [`docs/user_guide/diffusion_features.md`](../../docs/user_guide/diffusion_features.md)
+- Related PR: [#4730](https://github.com/vllm-project/vllm-omni/pull/4730)
+
+## Hardware Support
+
+This recipe documents a **single-GPU** CUDA layout on H100 80 GB. Add more
+platforms (ROCm / NPU) or multi-GPU layouts as community validation lands.
+
+## GPU
+
+### 1× H100 80GB
+
+#### Environment
+
+Versions from a working editable install (activate `vllm-omni/.venv` or your equivalent):
+
+- OS: Linux
+- Python: 3.12
+- Driver / runtime: NVIDIA driver **595.71.05** (CUDA 13.x)
+- torch: **2.11.0+cu130**
+- vLLM: **0.24.0**
+- vLLM-Omni: editable install from this repo (Git **`960ce8f6`**)
+- Transformers: **5.12.1**
+
+Krea 2 weights are ~30 GB and fit comfortably on a single 80 GB H100.
+
+#### Command
+
+Offline — few-step distilled (Turbo) checkpoint (2048×2048, no CFG):
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model krea/Krea-2-Turbo \
+  --prompt "a fox in the snow" \
+  --num-inference-steps 8 \
+  --guidance-scale 0.0 \
+  --height 2048 --width 2048 \
+  --output krea2_turbo.png
+```
+
+Offline — Raw checkpoint (1024×1024, CFG enabled):
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model krea/Krea-2-Raw \
+  --prompt "a fox in the snow" \
+  --num-inference-steps 28 \
+  --guidance-scale 4.5 \
+  --height 1024 --width 1024 \
+  --output krea2_raw.png
+```
+
+LoRA (offline) — pass a PEFT adapter with `--lora-path` / `--lora-scale`.
+`NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA` is a vLLM-Omni-compatible repackaging
+of `krea/Krea-2-LoRA-darkbrush`; its trigger word is `monochrome ink wash style`:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model krea/Krea-2-Turbo \
+  --prompt "a fox in the snow, monochrome ink wash style" \
+  --num-inference-steps 8 --guidance-scale 0.0 \
+  --lora-path NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA --lora-scale 1.0 \
+  --output krea2_darkbrush.png
+```
+
+Online — serve and query the OpenAI-compatible image route:
+
+```bash
+vllm serve krea/Krea-2-Turbo --omni --enforce-eager --port 8091
+```
+
+```bash
+curl -s http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "prompt": "a fox in the snow, monochrome ink wash style",
+    "size": "1024x1024",
+    "num_inference_steps": 8,
+    "guidance_scale": 0.0,
+    "seed": 42,
+    "lora": {"name": "darkbrush", "path": "NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA", "scale": 1.0}
+  }' | jq -r '.data[0].b64_json' | base64 -d > krea2_online.png
+```
+
+Cache-DiT acceleration (offline) — add `--cache-backend cache_dit`:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model krea/Krea-2-Turbo --prompt "a fox in the snow" \
+  --num-inference-steps 8 --guidance-scale 0.0 \
+  --cache-backend cache_dit --output krea2_cachedit.png
+```
+
+#### Verification
+
+```bash
+# Offline: the script prints "Saved generated image to <path>"; confirm the file exists.
+ls -lh krea2_turbo.png
+```
+
+For the online route, decode the returned `b64_json` (see the `curl` above) and
+confirm a valid 1024×1024 PNG is written.
+
+#### Notes
+
+- Key flags: distilled Turbo → few steps (~8) with `guidance_scale=0.0`; Raw →
+  more steps (~28) with `guidance_scale>0` (CFG). `guidance-scale` follows the
+  Krea 2 convention `velocity = cond + guidance_scale * (cond - uncond)`.
+- LoRA: adapters must be PEFT format (`adapter_config.json` +
+  `adapter_model.safetensors`); the projections are `ReplicatedLinear`
+  (264 modules). Offline uses `--lora-path` (which sets
+  `OmniDiffusionSamplingParams.lora_request` / `lora_scale`); online uses the
+  top-level `lora` object. Style LoRAs like darkbrush need their trigger word in
+  the prompt.
+- Cache-DiT is supported; `has_separate_cfg` is checkpoint-aware (False for Turbo
+  no-CFG, True for Raw CFG).
+- Known limitations: TP / SP / CFG-Parallel are not yet wired for Krea 2. HSDP,
+  layerwise CPU offload, and VAE-patch-parallel (decode) are supported.

tests/ar_diffusion/test_engine_routing.py

@@ -134,3 +134,17 @@ def test_runner_rejects_multi_sequence_config():
     fake = SimpleNamespace(od_config=SimpleNamespace(max_num_seqs=4))
     with pytest.raises(ValueError, match="max_num_seqs=1"):
         ARDiffusionModelRunner._preallocate_kv_cache(fake)
+
+
+def test_execute_model_accepts_kv_prefetch_jobs_kwarg():
+    """Regression: the worker passes ``kv_prefetch_jobs`` to every runner
+    (added with the KV prefetch path, #4448); the AR-Diffusion override must
+    accept and forward it or DreamZero crashes on the warm-up dummy run."""
+    import inspect
+
+    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    for cls in (DiffusionModelRunner, ARDiffusionModelRunner):
+        params = inspect.signature(cls.execute_model).parameters
+        assert "kv_prefetch_jobs" in params, cls

tests/ar_diffusion/test_kv_cache.py

@@ -9,14 +9,9 @@ KV stack on CPU (block bookkeeping only, no GPU tensors).
 import pytest
 import torch
 from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind
+from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 from vllm.v1.request import RequestStatus
 
-try:
-    from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
-except ModuleNotFoundError:
-    KVCacheSpecRegistry = None
-    from vllm.v1.core.single_type_kv_cache_manager import spec_manager_map
-
 from vllm_omni.experimental.ar_diffusion.kv_cache import (
     ARDiffusionKVCache,
     ARDiffusionKVConfig,
@@ -53,10 +48,7 @@ def test_spec_registration_resolves_to_chunk_window_manager():
     # Without explicit registration the MRO walk would fall back to the parent
     # SlidingWindowManager; assert the subclass manager wins.
     spec = make_spec()
-    if KVCacheSpecRegistry is not None:
-        assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
-    else:
-        assert spec_manager_map[type(spec)] is ChunkWindowManager
+    assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
 
 
 def test_spec_kind_is_sliding_window():

tests/ar_diffusion/test_pipeline_bridge.py

@@ -203,7 +203,7 @@ def test_failed_forward_tears_down_session(monkeypatch):
     runner.device = None
     runner._perf_e2e_times = []
 
-    def boom(self, req):
+    def boom(self, req, kv_prefetch_jobs=None):
         raise RuntimeError("layer 17 exploded")
 
     monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -36,7 +36,7 @@ import pytest
 
 from benchmarks.diffusion.backends import endpoint_filename_token, normalize_endpoint
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model, pytest.mark.local_model]
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
@@ -807,7 +807,7 @@ def assert_result(
 
 def _default_benchmark_endpoint_for_task(task: str) -> str:
     """Return the default client-side benchmark endpoint for a diffusion task."""
-    if task in {"t2v", "i2v", "ti2v"}:
+    if task in {"t2v", "i2v", "ti2v", "v2v"}:
         return "/v1/videos"
     if task in {"t2i", "i2i", "ti2i"}:
         return "/v1/chat/completions"

tests/dfx/perf/tests/test_cosmos3_vllm_omni.json

@@ -0,0 +1,213 @@
+[
+    {
+        "test_name": "test_cosmos3_t2i_official_demo_2gpu",
+        "description": "Cosmos3-Nano text-to-image official demo workload on 2 GPUs",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/images/generations",
+        "server_params": {
+            "model": "nvidia/Cosmos3-Nano",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "use-hsdp": true,
+                "hsdp-shard-size": 2,
+                "vae-patch-parallel-size": 2,
+                "vae-use-tiling": true,
+                "model-class-name": "Cosmos3OmniDiffusersPipeline",
+                "no-guardrails": true,
+                "stage-init-timeout": 900,
+                "init-timeout": 1200,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "1024x1024_steps4",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 1024,
+                "height": 1024,
+                "num-inference-steps": 4,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "seed": 42,
+                "extra-body": {
+                    "guidance_scale": 7.0,
+                    "flow_shift": 3.0,
+                    "guardrails": false
+                },
+                "baseline": {
+                    "throughput_qps": 1.107,
+                    "latency_mean": 0.8929,
+                    "peak_memory_mb_mean": 80000.0
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_cosmos3_t2v_official_demo_2gpu",
+        "description": "Cosmos3-Nano text-to-video official demo workload on 2 GPUs",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/videos",
+        "server_params": {
+            "model": "nvidia/Cosmos3-Nano",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "use-hsdp": true,
+                "hsdp-shard-size": 2,
+                "vae-patch-parallel-size": 2,
+                "vae-use-tiling": true,
+                "model-class-name": "Cosmos3OmniDiffusersPipeline",
+                "no-guardrails": true,
+                "stage-init-timeout": 900,
+                "init-timeout": 1200,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "1280x720_frames189_steps4",
+                "dataset": "random",
+                "task": "t2v",
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "seed": 42,
+                "random-request-config": [
+                    {
+                        "width": 1280,
+                        "height": 720,
+                        "prompt": "A robot arm is cleaning a plate in the kitchen.",
+                        "num_inference_steps": 4,
+                        "num_frames": 189,
+                        "fps": 24,
+                        "weight": 1
+                    }
+                ],
+                "extra-body": {
+                    "guidance_scale": 6.0,
+                    "flow_shift": 10.0,
+                    "max_sequence_length": 4096,
+                    "guardrails": false,
+                    "use_resolution_template": false,
+                    "use_duration_template": false
+                },
+                "baseline": {
+                    "throughput_qps": 0.027,
+                    "latency_mean": 36.6824,
+                    "peak_memory_mb_mean": 30155.4
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_cosmos3_i2v_official_demo_2gpu",
+        "description": "Cosmos3-Nano image-to-video official demo workload on 2 GPUs",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/videos",
+        "server_params": {
+            "model": "nvidia/Cosmos3-Nano",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "use-hsdp": true,
+                "hsdp-shard-size": 2,
+                "vae-patch-parallel-size": 2,
+                "vae-use-tiling": true,
+                "model-class-name": "Cosmos3OmniDiffusersPipeline",
+                "no-guardrails": true,
+                "stage-init-timeout": 900,
+                "init-timeout": 1200,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "1280x720_frames189_steps4",
+                "dataset": "random",
+                "task": "i2v",
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "num-input-images": 1,
+                "seed": 42,
+                "random-request-config": [
+                    {
+                        "width": 1280,
+                        "height": 720,
+                        "prompt": "The scene comes to life with smooth, natural motion.",
+                        "num_inference_steps": 4,
+                        "num_frames": 189,
+                        "fps": 24,
+                        "weight": 1
+                    }
+                ],
+                "extra-body": {
+                    "guidance_scale": 6.0,
+                    "flow_shift": 10.0,
+                    "max_sequence_length": 4096,
+                    "guardrails": false
+                },
+                "baseline": {
+                    "throughput_qps": 0.027,
+                    "latency_mean": 36.968,
+                    "peak_memory_mb_mean": 30443.6
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_cosmos3_v2v_official_demo_2gpu",
+        "description": "Cosmos3-Nano video-to-video official demo workload on 2 GPUs",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/videos",
+        "server_params": {
+            "model": "nvidia/Cosmos3-Nano",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "use-hsdp": true,
+                "hsdp-shard-size": 2,
+                "vae-patch-parallel-size": 2,
+                "vae-use-tiling": true,
+                "model-class-name": "Cosmos3OmniDiffusersPipeline",
+                "no-guardrails": true,
+                "stage-init-timeout": 900,
+                "init-timeout": 1200,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "1280x720_frames189_steps4",
+                "dataset": "random",
+                "task": "v2v",
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "seed": 42,
+                "random-request-config": [
+                    {
+                        "width": 1280,
+                        "height": 720,
+                        "prompt": "Continue the same scene with smooth natural motion.",
+                        "num_inference_steps": 4,
+                        "num_frames": 189,
+                        "fps": 24,
+                        "weight": 1
+                    }
+                ],
+                "extra-body": {
+                    "guidance_scale": 6.0,
+                    "flow_shift": 10.0,
+                    "max_sequence_length": 4096,
+                    "guardrails": false,
+                    "condition_frame_indexes_vision": [
+                        0,
+                        1
+                    ],
+                    "condition_video_keep": "first"
+                },
+                "baseline": {
+                    "throughput_qps": 0.027,
+                    "latency_mean": 37.2117,
+                    "peak_memory_mb_mean": 30135.6
+                }
+            }
+        ]
+    }
+]

tests/dfx/perf/tests/test_hunyuan_image3_it2i.json

@@ -7,6 +7,7 @@
         "server_params": {
             "model": "tencent/HunyuanImage-3.0-Instruct",
             "serve_args": {
+                "trust-remote-code": true,
                 "deploy-config-inline": {
                     "type": "yaml",
                     "content": {

tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py

@@ -62,8 +62,8 @@ def _finalize_edits_form_data(template: dict[str, Any], server_model: str) -> di
         ),
         pytest.param(
             True,
-            {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "layers": "2"},
-            ("Invalid layers", "3", "1"),
+            {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "layers": "1"},
+            ("Invalid layers", "2", "10"),
             400,
             id="invalid_layers",
         ),

tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py

@@ -87,8 +87,8 @@ def _minimal_images_gen_json(omni_server: OmniServer) -> dict[str, object]:
         pytest.param(
             {"true_cfg_scale": 21.0}, ("true_cfg_scale", "less_than_equal", "20"), id="true_cfg_scale_above_max"
         ),
-        pytest.param({"layers": 2}, ("layers", "value_error", "3", "10"), id="layers_below_min"),
-        pytest.param({"layers": 11}, ("layers", "value_error", "3", "10"), id="layers_above_max"),
+        pytest.param({"layers": 1}, ("layers", "value_error", "2", "10"), id="layers_below_min"),
+        pytest.param({"layers": 11}, ("layers", "value_error", "2", "10"), id="layers_above_max"),
         pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative", marks=_SKIP_ISSUE_3649),
         pytest.param(
             {"seed": 2**32},

tests/diffusion/distributed/test_hsdp.py

@@ -1,19 +1,97 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities.
+"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities."""
 
-These tests verify HSDP configuration logic without requiring a distributed environment.
-"""
+import gc
+import os
+import socket
 
 import pytest
+import torch
+import torch.distributed as dist
 import torch.nn as nn
+from torch.distributed.tensor import DeviceMesh, DTensor
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
+from vllm_omni.diffusion.distributed.hsdp import (
+    HSDPInferenceConfig,
+    _unshardable_parameters,
+    shard_model,
+)
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model]
 
 
+class _PackedBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(2, 2))
+        self.packed_weight = nn.Parameter(torch.ones(2, 2, dtype=torch.uint8), requires_grad=False)
+        self.input_global_scale = nn.Parameter(torch.tensor(1.0), requires_grad=False)
+
+
+class _PackedModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.block = _PackedBlock()
+        self.root_weight = nn.Parameter(torch.ones(2))
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return int(sock.getsockname()[1])
+
+
+@pytest.fixture(scope="module")
+def cpu_process_group():
+    if dist.is_initialized():
+        yield
+        return
+
+    master_port = _find_free_port()
+    os.environ.update(
+        {
+            "RANK": "0",
+            "LOCAL_RANK": "0",
+            "WORLD_SIZE": "1",
+            "MASTER_ADDR": "127.0.0.1",
+            "MASTER_PORT": str(master_port),
+        }
+    )
+    dist.init_process_group("gloo", rank=0, world_size=1)
+    try:
+        yield
+    finally:
+        dist.destroy_process_group()
+        for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
+            os.environ.pop(key, None)
+        gc.collect()
+
+
+def test_hsdp_keeps_packed_and_scalar_parameters_local(cpu_process_group):
+    model = _PackedModel()
+    ignored_params = _unshardable_parameters(model)
+    expected_ignored = {model.block.packed_weight, model.block.input_global_scale}
+    assert ignored_params == expected_ignored
+
+    packed_weight = model.block.packed_weight
+    input_global_scale = model.block.input_global_scale
+    shard_model(
+        model,
+        mesh=DeviceMesh("cpu", [0]),
+        hsdp_shard_conditions=[lambda name, _module: name == "block"],
+        ignored_params=ignored_params,
+    )
+
+    assert isinstance(model.block.weight, DTensor)
+    assert isinstance(model.root_weight, DTensor)
+    assert model.block.packed_weight is packed_weight
+    assert model.block.input_global_scale is input_global_scale
+    assert not isinstance(model.block.packed_weight, DTensor)
+    assert not isinstance(model.block.input_global_scale, DTensor)
+
+
 class TestHSDPInferenceConfig:
     """Tests for HSDPInferenceConfig dataclass."""
 

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -179,6 +179,59 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
     assert get_current_diffusion_config_or_none() is None
 
 
+def test_hsdp_processes_quantized_weights_before_sharding(mocker):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    loader.quant_config = object()
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    events: list[str] = []
+
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load"))  # type: ignore[method-assign]
+    loader._process_weights_after_loading = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model, _device: events.append("process")
+    )
+    mocker.patch.object(
+        loader_mod.ModuleDiscovery,
+        "discover",
+        return_value=PipelineModules(
+            dits=[model.transformer],
+            dit_names=["transformer"],
+            vaes=[],
+            encoders=[],
+            encoder_names=[],
+            resident_modules=[],
+            resident_names=[],
+        ),
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.quantization.hsdp_fp8.prepare_fp8_layers_for_fsdp",
+        side_effect=lambda _model: events.append("prepare"),
+    )
+    mocker.patch.object(
+        loader_mod,
+        "apply_hsdp_to_model",
+        side_effect=lambda *_args, **_kwargs: events.append("shard"),
+    )
+
+    loader._load_model_with_hsdp(torch.device("cpu"))
+
+    assert events == ["load", "process", "prepare", "shard"]
+
+
 def test_get_all_weights(prefetch_helios_model, mock_tp_group):
     """Ensure that get all weights on a tiny model resolves to nonempty weights."""
     od_config = OmniDiffusionConfig(

tests/diffusion/model_loader/test_modelopt_fp8_adapter.py

@@ -42,6 +42,53 @@ class _QuantizedPackedModelOptModel(nn.Module):
         )
 
 
+class _RemappedModelOptModel(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.runtime = nn.Module()
+        self.runtime.proj = nn.Linear(2, 2, bias=False)
+
+    @staticmethod
+    def remap_checkpoint_key(name: str) -> str:
+        return {
+            "transformer.orig.proj.weight": "runtime.proj.weight",
+        }.get(name, name)
+
+
+class _QuantizedRemappedModelOptModel(nn.Module):
+    """FP8 target whose params live under a namespace the raw checkpoint keys
+    never use, reachable only via ``remap_checkpoint_key``.
+
+    Mirrors the Cosmos3 case that regressed: the generic weights-mapper cannot
+    resolve ``transformer.orig.*`` to ``runtime.*``, so without the remap hook
+    the adapter treats every tensor as unresolved and silently drops the scales.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.runtime = nn.Module()
+        self.runtime.proj = nn.Module()
+        self.runtime.proj.register_parameter(
+            "weight",
+            nn.Parameter(torch.empty(2, 2, dtype=torch.float8_e4m3fn), requires_grad=False),
+        )
+        self.runtime.proj.register_parameter(
+            "weight_scale",
+            nn.Parameter(torch.empty(1), requires_grad=False),
+        )
+        self.runtime.proj.register_parameter(
+            "input_scale",
+            nn.Parameter(torch.empty(1), requires_grad=False),
+        )
+
+    @staticmethod
+    def remap_checkpoint_key(name: str) -> str:
+        prefix = "transformer.orig.proj."
+        if name.startswith(prefix):
+            return "runtime.proj." + name[len(prefix) :]
+        return name
+
+
 def _make_source() -> SimpleNamespace:
     return SimpleNamespace(
         subfolder="transformer",
@@ -92,3 +139,62 @@ def test_modelopt_adapter_keeps_scale_tensors_for_quantized_target():
         "transformer.block.to_q.weight_scale",
         "transformer.block.to_q.input_scale",
     ]
+
+
+def test_modelopt_adapter_uses_checkpoint_key_remap_for_target_dtype():
+    model = _RemappedModelOptModel()
+    adapter = ModelOptFp8CheckpointAdapter(model, _make_source())
+    fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn)
+    scale = torch.tensor([0.5], dtype=torch.float32)
+
+    adapted = list(
+        adapter.adapt(
+            iter(
+                [
+                    ("transformer.orig.proj.weight", fp8_weight),
+                    ("transformer.orig.proj.weight_scale", scale),
+                ]
+            )
+        )
+    )
+
+    assert [name for name, _ in adapted] == ["runtime.proj.weight"]
+    assert adapted[0][1].dtype == model.runtime.proj.weight.dtype
+    assert torch.allclose(adapted[0][1], fp8_weight.to(torch.float32) * scale)
+
+
+def test_modelopt_adapter_remaps_and_keeps_scales_for_quantized_target():
+    """Regression: scales were dropped for a quantized (FP8) target that only the
+    model's ``remap_checkpoint_key`` hook can resolve. Assert both scales are
+    emitted and the FP8 weight passes through unchanged (no dequantization)."""
+    model = _QuantizedRemappedModelOptModel()
+    adapter = ModelOptFp8CheckpointAdapter(model, _make_source())
+    fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn)
+    weight_scale = torch.tensor([0.5], dtype=torch.float32)
+    input_scale = torch.tensor([1.0], dtype=torch.float32)
+
+    adapted = list(
+        adapter.adapt(
+            iter(
+                [
+                    ("transformer.orig.proj.weight_scale", weight_scale),
+                    ("transformer.orig.proj.input_scale", input_scale),
+                    ("transformer.orig.proj.weight", fp8_weight),
+                ]
+            )
+        )
+    )
+
+    emitted = dict(adapted)
+    assert len(adapted) == len(emitted) == 3
+
+    # Both scales are retained (the dropped-scale regression). They keep their
+    # source-checkpoint names; the model's load_weights remaps them downstream.
+    assert torch.equal(emitted["transformer.orig.proj.weight_scale"], weight_scale)
+    assert torch.equal(emitted["transformer.orig.proj.input_scale"], input_scale)
+
+    # The FP8 weight is passed through unchanged under the remapped target name
+    # (the quantized param is FP8, so it must NOT be dequantized). Compare via
+    # float32 since torch.equal has no FP8 CPU kernel.
+    assert emitted["runtime.proj.weight"].dtype == torch.float8_e4m3fn
+    assert torch.equal(emitted["runtime.proj.weight"].to(torch.float32), fp8_weight.to(torch.float32))

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -20,6 +20,16 @@ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def test_pipeline_declares_layerwise_offload_components() -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+
+    assert Cosmos3OmniDiffusersPipeline._dit_modules == ["transformer.language_model", "transformer"]
+    assert Cosmos3OmniDiffusersPipeline._encoder_modules == []
+    assert Cosmos3OmniDiffusersPipeline._vae_modules == ["vae"]
+    assert Cosmos3OmniDiffusersPipeline._resident_modules == []
+    assert hasattr(Cosmos3OmniDiffusersPipeline, "enable_omni_model_cpu_offload")
+
+
 class StubScheduler:
     def __init__(
         self,
@@ -284,7 +294,6 @@ def test_pipeline_registered_and_exported() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
     from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
     from vllm_omni.diffusion.registry import (
-        _DIFFUSION_ACTION_POST_PROCESS_FUNCS,
         _DIFFUSION_IR_OP_PRIORITY_FUNCS,
         _DIFFUSION_MODELS,
         _DIFFUSION_POST_PROCESS_FUNCS,
@@ -301,9 +310,6 @@ def test_pipeline_registered_and_exported() -> None:
     )
     assert _DIFFUSION_PRE_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_pre_process_func"
     assert _DIFFUSION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_post_process_func"
-    assert (
-        _DIFFUSION_ACTION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_action_post_process_func"
-    )
     assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_ir_op_priority_func"
     assert "Cosmos3OmniDiffusersPipeline" in CUSTOM_DIT_ENABLERS
     assert "Cosmos3OmniDiffusersPipeline" in cosmos3.__all__
@@ -708,11 +714,11 @@ def test_postprocess_handles_image_video_audio_and_validation() -> None:
 def test_action_postprocess_handles_robolab_policy_outputs() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
         RoboLabPolicyInputs,
-        get_cosmos3_action_post_process_func,
+        get_cosmos3_post_process_func,
         make_robolab_action_postprocess_inputs,
     )
 
-    func = get_cosmos3_action_post_process_func(SimpleNamespace())
+    func = get_cosmos3_post_process_func(SimpleNamespace())
     inputs = RoboLabPolicyInputs(
         prompt="Pick the cube.",
         video_tensor=torch.zeros(1, 3, 3, 16, 16),
@@ -736,13 +742,41 @@ def test_action_postprocess_handles_robolab_policy_outputs() -> None:
     )
 
     action = torch.tensor([[[0.0, 0.25], [1.0, 0.75]]])
-    custom_output = {"robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs)}
-    processed = func(action, custom_output=custom_output)
+    processed = func(
+        {
+            "payload": {
+                "actions": action,
+            },
+            "metadata": {
+                "actions": {
+                    "raw_action_dim": 2,
+                    "action_mode": "policy",
+                    "domain_id": 7,
+                },
+                "common": {
+                    "action_only_output": True,
+                },
+                "internal": {
+                    "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs),
+                },
+            },
+        }
+    )
 
-    assert processed.shape == (1, 2)
-    assert processed.dtype == torch.zeros((), dtype=torch.float32).numpy().dtype
-    torch.testing.assert_close(torch.from_numpy(processed), torch.tensor([[1.0, 0.25]]))
-    assert "robolab_action_postprocess" not in custom_output
+    processed_action = processed["payload"]["actions"]
+    assert processed_action.shape == (1, 2)
+    assert processed_action.dtype == torch.zeros((), dtype=torch.float32).numpy().dtype
+    torch.testing.assert_close(torch.from_numpy(processed_action), torch.tensor([[1.0, 0.25]]))
+    assert processed["metadata"] == {
+        "actions": {
+            "raw_action_dim": 2,
+            "action_mode": "policy",
+            "domain_id": 7,
+        },
+        "common": {
+            "action_only_output": True,
+        },
+    }
 
 
 def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -1198,8 +1232,8 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint
     assert captured["flow_shifts"] == [10.0]
     # Transfer uses the V2V flow-sigma schedule (karras off) so flow_shift applies.
     assert captured["scheduler_use_karras_sigmas"] == [False]
-    assert output.custom_output["fps"] == expected_fps
-    assert output.output["video"].device.type == "meta"
+    assert output.output["metadata"]["video"]["fps"] == expected_fps
+    assert output.output["payload"]["video"].device.type == "meta"
 
 
 def test_forward_transfer_runs_multichunk_overlap_path(
@@ -1268,12 +1302,12 @@ def test_forward_transfer_runs_multichunk_overlap_path(
 
     assert captured["conditional_frames"] == [2, 1]
     assert len(captured["decode_calls"]) == 2
-    assert output.output["video"].shape == (1, 3, 8, 16, 16)
+    assert output.output["payload"]["video"].shape == (1, 3, 8, 16, 16)
     torch.testing.assert_close(
-        output.output["video"][0, 0, :, 0, 0],
+        output.output["payload"]["video"][0, 0, :, 0, 0],
         torch.tensor([-0.6, -0.5, -0.4, -0.3, -0.2, 0.2, 0.3, 0.4]),
     )
-    assert output.custom_output["transfer_controls"]["edge"].shape == (1, 3, 8, 16, 16)
+    assert output.output["metadata"]["transfer"]["controls"]["edge"].shape == (1, 3, 8, 16, 16)
     torch.testing.assert_close(captured["targets"][0][:, :, 0], torch.full((1, 3, 16, 16), -1.0))
     torch.testing.assert_close(captured["targets"][0][:, :, 1], torch.full((1, 3, 16, 16), 1.0))
     torch.testing.assert_close(captured["targets"][0][:, :, 2:], torch.full((1, 3, 3, 16, 16), 1.0))
@@ -1497,8 +1531,13 @@ class TestForwardRouting:
             )
         )
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
-        assert output.custom_output["action"].shape == (1, 2, 2)
-        assert "action_only_output" not in output.custom_output
+        assert output.output["payload"]["actions"].shape == (1, 2, 2)
+        assert output.output["metadata"]["actions"] == {
+            "raw_action_dim": 2,
+            "action_mode": "policy",
+            "domain_id": 7,
+        }
+        assert "common" not in output.output["metadata"]
 
     def test_forward_dispatches_robolab_policy_flow(
         self,
@@ -1582,12 +1621,15 @@ class TestForwardRouting:
         assert captured["prepare_action_video"]["kwargs"] == {"image_size": None}
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
         assert captured["diffuse_calls"][-1]["timesteps"].tolist() == [4, 3, 2, 1]
-        assert output.output == {}
-        assert output.custom_output["action_only_output"] is True
-        assert output.custom_output["action"].shape == (1, 2, 2)
-        assert "actions" not in output.custom_output
-        assert "robolab_action_postprocess" in output.custom_output
-        assert "robolab_policy_inputs" not in output.custom_output
+        assert output.output["payload"]["actions"].shape == (1, 2, 2)
+        assert output.output["metadata"]["actions"] == {
+            "raw_action_dim": 2,
+            "action_mode": "policy",
+            "domain_id": 7,
+        }
+        assert output.output["metadata"]["common"]["action_only_output"] is True
+        assert "robolab_action_postprocess" in output.output["metadata"]["internal"]
+        assert "robolab_policy_inputs" not in output.output["metadata"]
 
     @pytest.mark.parametrize(
         ("prompt", "sampling_params", "message"),

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -10,9 +10,19 @@ import pytest
 import torch
 from torch import nn
 
+from vllm_omni.platforms import current_omni_platform
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+@pytest.fixture
+def accelerator_device() -> torch.device:
+    """Provide an accelerator device, skipping when none is available."""
+    if current_omni_platform.get_device_count() == 0:
+        pytest.skip("Accelerator required for this test")
+    return current_omni_platform.get_torch_device(0)
+
+
 def _tiny_cosmos3_config(**overrides):
     config = {
         "hidden_size": 8,
@@ -88,7 +98,7 @@ def test_validate_supported_config_rejects_unsupported_flags(key: str, value) ->
 
 
 def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None:
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3LanguageModel, Cosmos3VFMTransformer
 
     model = object.__new__(Cosmos3VFMTransformer)
     nn.Module.__init__(model)
@@ -104,6 +114,7 @@ def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None:
     ]
     assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"]
     assert Cosmos3VFMTransformer._layerwise_offload_blocks_attrs == ["gen_layers"]
+    assert Cosmos3LanguageModel._layerwise_offload_blocks_attrs == ["layers"]
     assert Cosmos3VFMTransformer._repeated_blocks == ["Cosmos3GenDecoderLayer"]
 
     model.latent_patch_size = 2
@@ -131,6 +142,97 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
     assert tuple(output.shape) == (1, 2, 1, 2, 2)
 
 
+def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+
+    class ToyReasonerLayer(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weight = nn.Parameter(torch.ones(1))
+
+        def forward(self, hidden: torch.Tensor, freqs):
+            del freqs
+            kv = hidden.new_zeros(hidden.shape[0], hidden.shape[1], 1, 1)
+            return hidden + self.weight.to(hidden.dtype), kv, kv
+
+    class ToyGeneratorLayer(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weight = nn.Parameter(torch.ones(1))
+
+        def forward(self, hidden: torch.Tensor, **kwargs):
+            del kwargs
+            return hidden + self.weight.to(hidden.dtype)
+
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    model.language_model.layers = nn.ModuleList([ToyReasonerLayer()])
+    model.gen_layers = nn.ModuleList([ToyGeneratorLayer()])
+    model.enable_model_cpu_offload(device=torch.device("cpu"), pin_memory=False)
+
+    assert model._model_cpu_offload_enabled
+    assert set(model._model_cpu_offload_components()) == {"reasoner", "generator"}
+    assert model.device == torch.device("cpu")
+
+    output = model(
+        hidden_states=torch.zeros(1, 2, 1, 2, 2),
+        timestep=torch.tensor([1.0]),
+        text_ids=torch.tensor([[1, 2]], dtype=torch.long),
+        text_mask=torch.ones(1, 2, dtype=torch.long),
+        video_shape=(1, 2, 2),
+        fps=24.0,
+    )
+
+    assert tuple(output.shape) == (1, 2, 1, 2, 2)
+    # The reasoner runs once then the generator component stays resident for GEN.
+    assert model._active_model_cpu_offload_component == "generator"
+
+    model.disable_model_cpu_offload()
+    assert not model._model_cpu_offload_enabled
+
+
+def test_model_cpu_offload_moves_reasoner_and_generator_between_cpu_and_device(
+    monkeypatch: pytest.MonkeyPatch, accelerator_device: torch.device
+) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    model.language_model.layers = nn.ModuleList([nn.Linear(2, 2)])
+    model.gen_layers = nn.ModuleList([nn.Linear(2, 2)])
+    model.to(accelerator_device)
+
+    reasoner_param = model.language_model.layers[0].weight
+    generator_param = model.gen_layers[0].weight
+
+    model.enable_model_cpu_offload(device=accelerator_device, pin_memory=False)
+
+    assert model._model_cpu_offload_enabled
+    # On enable, every group is parked on CPU until a phase activates it.
+    assert reasoner_param.device.type == "cpu"
+    assert generator_param.device.type == "cpu"
+
+    model._activate_model_cpu_offload_component("reasoner")
+    assert reasoner_param.device == accelerator_device
+    assert generator_param.device.type == "cpu"
+
+    model._activate_model_cpu_offload_component("generator")
+    assert reasoner_param.device.type == "cpu"
+    assert generator_param.device == accelerator_device
+
+    model.disable_model_cpu_offload()
+    assert not model._model_cpu_offload_enabled
+    # Disable restores both components onto the device.
+    assert reasoner_param.device == accelerator_device
+    assert generator_param.device == accelerator_device
+
+
 def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 

tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image as qwen_image
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_qwen_image_postprocess_preserves_output_envelope(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path,
+) -> None:
+    model_path = tmp_path / "model"
+    vae_path = model_path / "vae"
+    vae_path.mkdir(parents=True)
+    (vae_path / "config.json").write_text(json.dumps({}), encoding="utf-8")
+
+    processed = object()
+    captured: dict[str, object] = {}
+
+    class FakeImageProcessor:
+        def __init__(self, *, vae_scale_factor: int) -> None:
+            captured["vae_scale_factor"] = vae_scale_factor
+
+        def postprocess(self, images: torch.Tensor) -> object:
+            captured["images"] = images
+            return processed
+
+    monkeypatch.setattr(qwen_image, "VaeImageProcessor", FakeImageProcessor)
+    postprocess = qwen_image.get_qwen_image_post_process_func(SimpleNamespace(model=str(model_path)))
+
+    image = torch.zeros(1, 3, 4, 4)
+    output = postprocess(
+        {
+            "payload": {"image": image},
+            "metadata": {"prompt_embeddings": {"present": True}},
+        }
+    )
+
+    assert captured["vae_scale_factor"] == 16
+    assert captured["images"] is image
+    assert output == {
+        "payload": {"image": processed},
+        "metadata": {"prompt_embeddings": {"present": True}},
+    }