diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index a1c0c4862..d3246a98f 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -345,7 +345,7 @@ steps: commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - - export CACHE_DIT_VERSION=1.3.0 + - export CACHE_DIT_VERSION=1.5.0 - | set +e pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json @@ -361,7 +361,7 @@ steps: commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - - export CACHE_DIT_VERSION=1.3.0 + - export CACHE_DIT_VERSION=1.5.0 - | set +e pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json @@ -377,7 +377,7 @@ steps: commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - - export CACHE_DIT_VERSION=1.3.0 + - export CACHE_DIT_VERSION=1.5.0 - | set +e pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json @@ -393,7 +393,7 @@ steps: commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - - export CACHE_DIT_VERSION=1.3.0 + - export CACHE_DIT_VERSION=1.5.0 - | set +e pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json @@ -409,7 +409,7 @@ steps: commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - - export CACHE_DIT_VERSION=1.3.0 + - export CACHE_DIT_VERSION=1.5.0 - | set +e pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json diff --git a/.buildkite/npu/test-npu-nightly.yml b/.buildkite/npu/test-npu-nightly.yml index 7f277258c..050cef4f7 100644 --- a/.buildkite/npu/test-npu-nightly.yml +++ b/.buildkite/npu/test-npu-nightly.yml @@ -115,7 +115,7 @@ steps: DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - - export CACHE_DIT_VERSION=1.3.0 + - export CACHE_DIT_VERSION=1.5.0 - | set +e pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json diff --git a/.readthedocs.yml b/.readthedocs.yml index 2b5b81bbc..be2996355 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -11,10 +11,15 @@ build: post_checkout: - git fetch --unshallow || true pre_install: + # `pre_install` runs before RTD upgrades pip, and the seeded pip 23.1 + # rejects the `typing_extensions` wheel torch needs (fixed in pip 23.2). + - python -m pip install --upgrade pip # Pre-install CPU torch so the subsequent `pip install .[docs]` does not # pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's - # 1 GB memory cap and time out the mkdocs build. - - python -m pip install --index-url https://download.pytorch.org/whl/cpu torch + # 1 GB memory cap and time out the mkdocs build. Wheel-only: `--index-url` + # replaces PyPI, so an sdist fallback cannot reach its build deps. The + # quotes keep YAML from reading `:all:` as a mapping key. + - "python -m pip install --only-binary=:all: --index-url https://download.pytorch.org/whl/cpu torch" mkdocs: configuration: mkdocs.yml diff --git a/benchmarks/lingbot_video/dense_pipeline_parity.py b/benchmarks/lingbot_video/dense_pipeline_parity.py index 0c0c1d7a0..cc781a95b 100644 --- a/benchmarks/lingbot_video/dense_pipeline_parity.py +++ b/benchmarks/lingbot_video/dense_pipeline_parity.py @@ -167,9 +167,11 @@ def main() -> None: ] native_cmd = [ sys.executable, - "examples/offline_inference/text_to_video/text_to_video_lingbot.py", + "examples/offline_inference/text_to_video/text_to_video.py", "--model", args.model, + "--model-class-name", + "LingBotVideoPipeline", "--prompt", args.prompt, "--output", diff --git a/docs/user_guide/examples/offline_inference/text_to_video.md b/docs/user_guide/examples/offline_inference/text_to_video.md index d32a1f562..1a2082f82 100644 --- a/docs/user_guide/examples/offline_inference/text_to_video.md +++ b/docs/user_guide/examples/offline_inference/text_to_video.md @@ -12,6 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back | Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) | |---|---|---|---|---|---| | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB | +| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) | | `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | @@ -35,6 +36,25 @@ python text_to_video.py \ --output t2v_out.mp4 ``` +### LingBot-Video + +The shared runner recognizes both official LingBot checkpoint IDs and selects +the native `LingBotVideoPipeline` automatically: + +```bash +python text_to_video.py \ + --model robbyant/lingbot-video-dense-1.3b \ + --prompt "a robotic arm picks up a red block" \ + --height 192 --width 320 --num-frames 9 --num-inference-steps 2 \ + --guidance-scale 3.0 --flow-shift 3.0 --fps 24 \ + --output lingbot_t2v.mp4 +``` + +Use `robbyant/lingbot-video-moe-30b-a3b` for the MoE checkpoint. For local +paths without `lingbot` in the directory name, add +`--model-class-name LingBotVideoPipeline`. Model-only options use +`--extra-body`, for example `'{"batch_cfg": true, "output_type": "np"}'`. + ### LTX-2 ```bash diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index 50b207526..339c7dac9 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -8,6 +8,7 @@ A unified script for text-to-video generation. Supports multiple models with mod |---|---|---|---|---|---| | `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) | | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB | +| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) | | `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | @@ -57,18 +58,31 @@ python text_to_video.py \ ### LingBot-Video -The dedicated LingBot example explicitly selects the video output modality: +The shared runner recognizes the official dense and MoE checkpoint IDs, selects +`LingBotVideoPipeline`, and builds the canonical video request envelope: ```bash -python examples/offline_inference/text_to_video/text_to_video_lingbot.py \ +python text_to_video.py \ --model robbyant/lingbot-video-dense-1.3b \ + --prompt "a robotic arm picks up a red block" \ --height 192 --width 320 --num-frames 9 --num-inference-steps 2 \ + --guidance-scale 3.0 --flow-shift 3.0 --fps 24 \ --output lingbot_t2v.mp4 ``` Use the same command with `robbyant/lingbot-video-moe-30b-a3b` for the MoE checkpoint. Requested frame counts are rounded upward to the causal -VAE's `4n+1` grid. +VAE's `4n+1` grid. For a local checkpoint path whose name does not contain +`lingbot`, pass `--model-class-name LingBotVideoPipeline`. + +LingBot-only options such as `batch_cfg` and `output_type` use the shared +model-extra channel, for example: + +```bash +python text_to_video.py \ + --model robbyant/lingbot-video-dense-1.3b \ + --extra-body '{"batch_cfg": true, "output_type": "np"}' +``` ### LTX-2 diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 2e272b7a7..129ca860a 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -16,7 +16,7 @@ from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.lora.request import LoRARequest from vllm_omni.lora.utils import stable_lora_int_id -from vllm_omni.model_extras import get_extra_body_params, get_model_class_name +from vllm_omni.model_extras import get_extra_body_params, get_model_class_name, get_output_tensor_range from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform @@ -80,6 +80,17 @@ _MODEL_PRESETS = { "fps": 16, "output": "helios_output.mp4", }, + "lingbot": { + "model_class_name": "LingBotVideoPipeline", + "height": 192, + "width": 320, + "num_frames": 9, + "num_inference_steps": 2, + "guidance_scale": 3.0, + "fps": 24, + "flow_shift": 3.0, + "output": "lingbot_video_output.mp4", + }, "ltx2": { "height": 512, "width": 768, @@ -110,27 +121,51 @@ _MODEL_PRESETS = { def _detect_preset(model: str, model_class_name: str | None = None) -> dict: model_lower = model.lower() class_lower = (model_class_name or "").lower() - if "distilled" in class_lower or "distilled" in model_lower: - return _MODEL_PRESETS["ltx2_distilled"] - if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower: - return _MODEL_PRESETS["ltx23"] - if "ltx2" in class_lower or "ltx-2" in model_lower or "ltx_2" in model_lower: + if "lingbot" in model_lower or "lingbotvideo" in class_lower: + return _MODEL_PRESETS["lingbot"] + if "ltx" in class_lower or "ltx" in model_lower: + if "distilled" in class_lower or "distilled" in model_lower: + return _MODEL_PRESETS["ltx2_distilled"] + if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower: + return _MODEL_PRESETS["ltx23"] return _MODEL_PRESETS["ltx2"] - if "vace" in model_lower: + if "vace" in model_lower or "vace" in class_lower: return _MODEL_PRESETS["vace"] # Edge must be matched before the generic cosmos branch (its "cosmos" substring would # otherwise pick up the Nano/Super 720p / gs 6.0 / flow_shift 10.0 preset). if ("cosmos" in model_lower or "cosmos" in class_lower) and ("edge" in model_lower or "edge" in class_lower): return _MODEL_PRESETS["cosmos3_edge"] - if "cosmos" in model_lower: + if "cosmos" in model_lower or "cosmos" in class_lower: return _MODEL_PRESETS["cosmos"] - if "hunyuan" in model_lower: + if "hunyuan" in model_lower or "hunyuan" in class_lower: return _MODEL_PRESETS["hunyuan"] - if "helios" in model_lower: + if "helios" in model_lower or "helios" in class_lower: return _MODEL_PRESETS["helios"] return _MODEL_PRESETS["wan"] +def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]: + """Build the canonical request envelope for the shared T2V example.""" + result: dict[str, Any] = { + "prompt": prompt, + "modalities": ["video"], + } + if negative_prompt is not None: + result["negative_prompt"] = negative_prompt + return result + + +def _normalize_float_tensor(tensor: torch.Tensor, source_range: str) -> torch.Tensor: + """Normalize decoded frames according to the pipeline's output contract.""" + if not tensor.is_floating_point(): + return tensor + if source_range == "negative_one_to_one": + return tensor.clamp(-1, 1) * 0.5 + 0.5 + if source_range == "zero_to_one": + return tensor.clamp(0, 1) + raise ValueError(f"Unsupported floating-point tensor range: {source_range!r}") + + def parse_profiler_config(value: str) -> dict[str, Any]: try: config = json.loads(value) @@ -160,7 +195,7 @@ def parse_extra_body(value: str) -> dict[str, Any]: def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Generate a video from a text prompt. " - "Supports Wan2.2, HunyuanVideo-1.5, Helios, and other text-to-video models." + "Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, and other text-to-video models." ) parser.add_argument( "--model", @@ -408,6 +443,7 @@ def main(): for key, default_val in preset.items(): if getattr(args, key.replace("-", "_"), None) is None: setattr(args, key.replace("-", "_"), default_val) + model_class_name = args.model_class_name generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) # Cache-dit config (Wan2.2 only) @@ -480,6 +516,7 @@ def main(): omni = Omni(**omni_kwargs) model_class_name = get_model_class_name(omni) or model_class_name declared_extra_body_params = get_extra_body_params(model_class_name) + output_tensor_range = get_output_tensor_range(model_class_name) if profiler_enabled: print("[Profiler] Starting profiling...") @@ -513,12 +550,11 @@ def main(): lora_path=lora_path, ) - prompt_dict = {"prompt": args.prompt} - if args.negative_prompt is not None: - prompt_dict["negative_prompt"] = args.negative_prompt - elif preset not in (_MODEL_PRESETS["ltx2"], _MODEL_PRESETS["ltx23"]): + negative_prompt = args.negative_prompt + if negative_prompt is None and all(preset is not _MODEL_PRESETS[name] for name in ("lingbot", "ltx2", "ltx23")): # Preserve the historical empty-prompt behavior for non-LTX examples. - prompt_dict["negative_prompt"] = "" + negative_prompt = "" + prompt_dict = build_text_to_video_prompt(args.prompt, negative_prompt) extra_args = {} if lora_request: @@ -630,8 +666,7 @@ def main(): frame_tensor = frame_tensor[0] if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4): frame_tensor = frame_tensor.permute(1, 2, 0) - if frame_tensor.is_floating_point(): - frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5 + frame_tensor = _normalize_float_tensor(frame_tensor, output_tensor_range) return frame_tensor.float().numpy() if isinstance(frame, np.ndarray): frame_array = frame @@ -683,8 +718,7 @@ def main(): video_tensor = video_tensor[0] elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4): video_tensor = video_tensor.permute(1, 2, 3, 0) - if video_tensor.is_floating_point(): - video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5 + video_tensor = _normalize_float_tensor(video_tensor, output_tensor_range) video_array = video_tensor.float().numpy() elif isinstance(frames, np.ndarray): video_array = frames diff --git a/examples/offline_inference/text_to_video/text_to_video_lingbot.py b/examples/offline_inference/text_to_video/text_to_video_lingbot.py deleted file mode 100644 index 108e2dc1f..000000000 --- a/examples/offline_inference/text_to_video/text_to_video_lingbot.py +++ /dev/null @@ -1,131 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import argparse -import time -from pathlib import Path -from typing import Any - -import numpy as np -import torch - -from vllm_omni.diffusion.data import DiffusionParallelConfig -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.platforms import current_omni_platform - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Generate text-to-video with a LingBot-Video checkpoint.") - parser.add_argument("--model", default="robbyant/lingbot-video-dense-1.3b") - parser.add_argument("--prompt", default="a robotic arm picks up a red block") - parser.add_argument("--negative-prompt", default=None) - parser.add_argument("--output", default="lingbot_video_output.mp4") - parser.add_argument("--height", type=int, default=192) - parser.add_argument("--width", type=int, default=320) - parser.add_argument("--num-frames", type=int, default=9) - parser.add_argument("--num-inference-steps", type=int, default=2) - parser.add_argument("--guidance-scale", type=float, default=3.0) - parser.add_argument("--flow-shift", type=float, default=3.0) - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--fps", type=int, default=24) - parser.add_argument("--output-type", choices=["pt", "np"], default="pt") - parser.add_argument("--batch-cfg", action="store_true") - return parser.parse_args() - - -def _peak_memory_mb(output: Any) -> float: - result = output[0] if isinstance(output, list) and output else output - return float(getattr(result, "peak_memory_mb", 0.0) or 0.0) - - -def _extract_video(output: Any) -> Any: - if isinstance(output, list): - output = output[0] if output else None - if isinstance(output, OmniRequestOutput): - if output.final_output_type != "image": - raise ValueError(f"Unexpected output type {output.final_output_type!r}.") - if not output.images: - raise ValueError("No video frames found in OmniRequestOutput.") - return output.images[0] if len(output.images) == 1 else output.images - return output - - -def _to_frame_array(video: Any) -> np.ndarray: - if isinstance(video, torch.Tensor): - tensor = video.detach().cpu() - if tensor.dim() == 5: - tensor = tensor[0] - if tensor.dim() == 4 and tensor.shape[-1] not in (3, 4): - if tensor.shape[0] in (3, 4): - tensor = tensor.permute(1, 2, 3, 0) - elif tensor.shape[1] in (3, 4): - tensor = tensor.permute(0, 2, 3, 1) - if tensor.is_floating_point(): - if float(tensor.min()) < 0.0: - tensor = tensor.clamp(-1, 1) * 0.5 + 0.5 - else: - tensor = tensor.clamp(0, 1) - return tensor.float().numpy() - if isinstance(video, np.ndarray): - array = video[0] if video.ndim == 5 else video - if np.issubdtype(array.dtype, np.integer): - return array.astype(np.float32) / 255.0 - return array.astype(np.float32, copy=False) - if isinstance(video, list): - return np.stack([_to_frame_array(frame) for frame in video], axis=0) - return np.asarray(video, dtype=np.float32) - - -def main() -> None: - args = parse_args() - generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) - omni = Omni( - model=args.model, - model_class_name="LingBotVideoPipeline", - flow_shift=args.flow_shift, - parallel_config=DiffusionParallelConfig(), - ) - - prompt: dict[str, Any] = { - "prompt": args.prompt, - "modalities": ["video"], - "num_frames": args.num_frames, - } - if args.negative_prompt is not None: - prompt["negative_prompt"] = args.negative_prompt - - sampling_params = OmniDiffusionSamplingParams( - height=args.height, - width=args.width, - num_frames=args.num_frames, - num_inference_steps=args.num_inference_steps, - guidance_scale=args.guidance_scale, - generator=generator, - output_type=args.output_type, - extra_args={"batch_cfg": args.batch_cfg, "shift": args.flow_shift}, - ) - - start = time.perf_counter() - output = omni.generate(prompt, sampling_params) - elapsed = time.perf_counter() - start - print(f"Total generation time: {elapsed:.4f} seconds ({elapsed * 1000:.2f} ms)") - peak_mb = _peak_memory_mb(output) - if peak_mb: - print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)") - - video = _to_frame_array(_extract_video(output)) - output_path = Path(args.output) - output_path.parent.mkdir(parents=True, exist_ok=True) - from diffusers.utils import export_to_video - - export_to_video(list(video), str(output_path), fps=args.fps) - print(f"Saved generated video to {output_path}") - omni.close() - - -if __name__ == "__main__": - main() diff --git a/recipes/Robbyant/LingBot-Video.md b/recipes/Robbyant/LingBot-Video.md index 289d36990..29d112217 100644 --- a/recipes/Robbyant/LingBot-Video.md +++ b/recipes/Robbyant/LingBot-Video.md @@ -27,7 +27,7 @@ parallel, cache, quantized, or expert-kernel backends. - Dense checkpoint: - MoE checkpoint: - Upstream project: -- Related offline example: [`examples/offline_inference/text_to_video/text_to_video_lingbot.py`](../../examples/offline_inference/text_to_video/text_to_video_lingbot.py) +- Related offline example: [`examples/offline_inference/text_to_video/text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py) - Related online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md) ## Hardware Support @@ -52,7 +52,7 @@ The MoE path is validated with BF16 expert weights. ```bash CUDA_VISIBLE_DEVICES=0 \ -python examples/offline_inference/text_to_video/text_to_video_lingbot.py \ +python examples/offline_inference/text_to_video/text_to_video.py \ --model robbyant/lingbot-video-dense-1.3b \ --prompt "a robotic arm picks up a red block" \ --output lingbot_t2v.mp4 \ @@ -70,7 +70,7 @@ python examples/offline_inference/text_to_video/text_to_video_lingbot.py \ ```bash CUDA_VISIBLE_DEVICES=0 \ -python examples/offline_inference/text_to_video/text_to_video_lingbot.py \ +python examples/offline_inference/text_to_video/text_to_video.py \ --model robbyant/lingbot-video-moe-30b-a3b \ --prompt "a robotic arm picks up a red block" \ --output lingbot_moe_t2v.mp4 \ diff --git a/recipes/Tencent/HunyuanImage-3.0-Instruct.md b/recipes/Tencent/HunyuanImage-3.0-Instruct.md index a9a63d5ce..4614df382 100644 --- a/recipes/Tencent/HunyuanImage-3.0-Instruct.md +++ b/recipes/Tencent/HunyuanImage-3.0-Instruct.md @@ -82,7 +82,7 @@ example and the DiT sections below. - Optional environment variables: ```bash -export CACHE_DIT_VERSION=1.3.0 +export CACHE_DIT_VERSION=1.5.0 ``` HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA` diff --git a/requirements/common.txt b/requirements/common.txt index 13c7d35a9..e0233bd8d 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -16,7 +16,7 @@ diffusers==0.38.0 safetensors>=0.8.0 accelerate==1.12.0 soundfile>=0.13.1 -cache-dit==1.3.0 +cache-dit==1.5.0 tqdm>=4.66.0 torchsde>=0.2.6 openai-whisper>=20250625 diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 581a7498f..617083398 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/tests/diffusion/quantization/test_quantization_fp8.py @@ -189,7 +189,7 @@ def _generate_single_stage_video( def _generate_bagel_image( - quantization_config: str | None = None, + diffusion_quantization_config: str | None = None, num_inference_steps: int = 15, ) -> tuple[Any, float]: """Generate an image with BAGEL (multi-stage: LLM + Diffusion). @@ -202,8 +202,8 @@ def _generate_bagel_image( "stage_configs_path": config_path, "stage_init_timeout": 300, } - if quantization_config: - omni_kwargs["quantization_config"] = quantization_config + if diffusion_quantization_config: + omni_kwargs["diffusion_quantization_config"] = diffusion_quantization_config model_name = omni_kwargs.pop("model") with OmniRunner(model_name, **omni_kwargs) as runner: @@ -380,21 +380,21 @@ def test_single_stage_ltx2_fp8_uses_less_memory(): @hardware_test(res={"cuda": "H100"}) def test_bagel_fp8_generates_image(): - """BAGEL with FP8 quantization_config generates a valid image. + """BAGEL with diffusion-stage FP8 generates a valid image. FP8 should only apply to the diffusion stage (Stage-1), not the LLM stage (Stage-0). We verify this by checking: 1. Image is generated successfully 2. LLM stage finish_reason is 'stop' (not 'length' from garbled output) """ - image, _ = _generate_bagel_image(quantization_config="fp8") + image, _ = _generate_bagel_image(diffusion_quantization_config="fp8") image.save("test_bagel_fp8.png") @hardware_test(res={"cuda": "H100"}) def test_bagel_bf16_generates_image(): """BAGEL without quantization generates a valid image (baseline).""" - image, _ = _generate_bagel_image(quantization_config=None) + image, _ = _generate_bagel_image(diffusion_quantization_config=None) image.save("test_bagel_bf16.png") diff --git a/tests/examples/offline_inference/test_text_to_video_prompts.py b/tests/examples/offline_inference/test_text_to_video_prompts.py new file mode 100644 index 000000000..ff7244d65 --- /dev/null +++ b/tests/examples/offline_inference/test_text_to_video_prompts.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import pytest +import torch + +from examples.offline_inference.text_to_video.text_to_video import ( + _MODEL_PRESETS, + _detect_preset, + _normalize_float_tensor, + build_text_to_video_prompt, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + ("negative_prompt", "expected_negative_prompt"), + [(None, None), ("", ""), ("blurry", "blurry")], +) +def test_text_to_video_builds_canonical_prompt( + negative_prompt: str | None, + expected_negative_prompt: str | None, +) -> None: + result = build_text_to_video_prompt("a robot waves", negative_prompt) + + assert result["prompt"] == "a robot waves" + assert result["modalities"] == ["video"] + if expected_negative_prompt is None: + assert "negative_prompt" not in result + else: + assert result["negative_prompt"] == expected_negative_prompt + + +@pytest.mark.parametrize( + ("model", "model_class_name", "preset_name"), + [ + ("robbyant/lingbot-video-dense-1.3b", None, "lingbot"), + ("/models/custom-checkpoint", "LingBotVideoPipeline", "lingbot"), + ("Lightricks/LTX-2", "LTX2DistilledPipeline", "ltx2_distilled"), + ("BestWishYsh/Helios-Distilled", None, "helios"), + ("/models/cosmos-edge", "Cosmos3EdgeVFMTransformer", "cosmos3_edge"), + ("/models/vace", "WanVACEPipeline", "vace"), + ], +) +def test_detect_preset_is_scoped_to_model_family( + model: str, + model_class_name: str | None, + preset_name: str, +) -> None: + assert _detect_preset(model, model_class_name) is _MODEL_PRESETS[preset_name] + + +def test_lingbot_preset_matches_lingbot_defaults() -> None: + assert _MODEL_PRESETS["lingbot"] == { + "model_class_name": "LingBotVideoPipeline", + "height": 192, + "width": 320, + "num_frames": 9, + "num_inference_steps": 2, + "guidance_scale": 3.0, + "fps": 24, + "flow_shift": 3.0, + "output": "lingbot_video_output.mp4", + } + + +def test_normalize_float_tensor_preserves_zero_to_one_frames() -> None: + frames = torch.tensor([0.0, 0.25, 1.0, 1.5]) + + assert torch.equal(_normalize_float_tensor(frames, "zero_to_one"), torch.tensor([0.0, 0.25, 1.0, 1.0])) + + +def test_normalize_float_tensor_maps_negative_one_to_one_frames() -> None: + frames = torch.tensor([-1.5, -1.0, 0.0, 1.0, 1.5]) + + assert torch.equal( + _normalize_float_tensor(frames, "negative_one_to_one"), + torch.tensor([0.0, 0.0, 0.5, 1.0, 1.0]), + ) + + +def test_normalize_float_tensor_uses_one_contract_for_ambiguous_frames() -> None: + frames = [torch.tensor([0.0, 0.5, 1.0]), torch.tensor([-1.0, 0.0, 1.0])] + + normalized = [_normalize_float_tensor(frame, "negative_one_to_one") for frame in frames] + + assert torch.equal(normalized[0], torch.tensor([0.5, 0.75, 1.0])) + assert torch.equal(normalized[1], torch.tensor([0.0, 0.5, 1.0])) diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py index 16e04fc09..7ad35b86d 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py +++ b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py @@ -1,21 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from functools import lru_cache from types import SimpleNamespace from unittest.mock import Mock import pytest import torch import torch.nn as nn +from vllm.platforms import current_platform from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( HiFTGenerator, ) from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import ( + CFMGraphWrapper, HiFTGraphWrapper, ) -pytestmark = [pytest.mark.core_model] +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] class _F0Predictor(nn.Module): @@ -142,3 +145,248 @@ def test_lazy_capture_limit_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) wrapper._capture.assert_not_called() wrapper.decode_fn.assert_called_once_with(speech_feat, cache_source) assert result is wrapper.decode_fn.return_value + + +# --------------------------------------------------------------------------- +# CFMGraphWrapper tests +# --------------------------------------------------------------------------- + + +class _MiniDiT(nn.Module): + """Minimal DiT-like module with blocks_forward_chunk for CFM graph testing. + + Mimics the upstream cosyvoice2 DiT's cache semantics: + - CausalConv1d.forward_chunk: cat([cnn_cache, x], dim=time) when cache is not None + - Attention.forward_chunk: cat([k, k_cache], dim=seq) when att_cache is not None + """ + + def __init__(self, in_dim: int = 16, hidden: int = 8, depth: int = 2) -> None: + super().__init__() + self.in_proj = nn.Linear(in_dim, hidden) + self.blocks = nn.ModuleList([nn.LayerNorm(hidden) for _ in range(depth)]) + self.final_layer = nn.Linear(hidden, in_dim) + + def blocks_forward_chunk( + self, + x: torch.Tensor, + t: torch.Tensor, + mask: torch.Tensor | None, + cnn_cache: torch.Tensor | None = None, + att_cache: torch.Tensor | None = None, + cnn_cache_buffer: torch.Tensor | None = None, + att_cache_buffer: torch.Tensor | None = None, + ) -> torch.Tensor: + x = x.transpose(1, 2) + x = self.in_proj(x) + for b_idx in range(len(self.blocks)): + # Simulate CausalConv1d: cnn_cache contributes to the first frames. + # Upstream: if cnn_cache[b_idx] is None, creates zeros internally (same as zeros). + cnn_b = cnn_cache[b_idx] if cnn_cache is not None else None + if cnn_b is not None: + x[:, : cnn_b.shape[2], :] += cnn_b.transpose(1, 2) + # Simulate Attention: att_cache contributes a bias when non-empty. + # Upstream: if att_cache[b_idx] is None, skips cat entirely (different path). + att_b = att_cache[b_idx] if att_cache is not None else None + if att_b is not None and att_b.shape[3] > 0: + x += att_b.sum(dim=(1, 2), keepdim=False).unsqueeze(1) + x = self.blocks[b_idx](x) + x = x + t + cnn_cache_buffer[b_idx] = x[:, -2:, :].transpose(1, 2).contiguous() + dt = x.shape[1] + att_cache_buffer[b_idx][:, :, :dt, :] = x.unsqueeze(1) + x = self.final_layer(x) + x = x.transpose(1, 2) + return x + + +def _cfm_inputs( + batch_size: int, chunk_size: int, old_att_len: int, *, device: str = "cuda" +) -> tuple[torch.Tensor, ...]: + depth = 2 + hidden = 8 + estimator_input = torch.randn(batch_size, 16, chunk_size, device=device) + time_emb = torch.randn(batch_size, 1, hidden, device=device) + cnn_cache = torch.randn(depth, batch_size, hidden, 2, device=device) + att_cache = torch.randn(depth, batch_size, 1, old_att_len, hidden, device=device) + cnn_out = torch.empty(depth, batch_size, hidden, 2, device=device) + att_out = torch.empty(depth, batch_size, 1, old_att_len + chunk_size, hidden, device=device) + return estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _MiniDiT().eval().cuda() + wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=32) + + with torch.inference_mode(): + for _, chunk_size, old_att_len in ((2, 10, 0), (2, 10, 5)): + inputs = _cfm_inputs(2, chunk_size, old_att_len) + + eager_inputs = tuple(v.clone() for v in inputs) + with torch.no_grad(): + eager_result = estimator.blocks_forward_chunk( + eager_inputs[0], + eager_inputs[1], + None, + eager_inputs[2], + eager_inputs[3], + eager_inputs[4], + eager_inputs[5], + ) + + wrapper.replay(*inputs) + replay_inputs = tuple(v.clone() for v in inputs) + graph_result, graph_cnn, graph_att = wrapper.replay(*replay_inputs) + + torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_cnn, eager_inputs[4], rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_att, eager_inputs[5], rtol=1e-4, atol=1e-5) + + wrapper._capture_graph.cache_clear() + + +def _cfm_mock_wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 1) -> CFMGraphWrapper: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + wrapper = object.__new__(CFMGraphWrapper) + wrapper.max_graphs = max_graphs + wrapper.graph_fn = Mock(return_value=torch.tensor([42.0])) + wrapper.device = torch.device("cuda") + + @lru_cache(maxsize=max_graphs) + def _capture_graph(key: tuple): + return wrapper._do_capture(key) + + wrapper._capture_graph = _capture_graph + wrapper._do_capture = Mock(return_value=None) + return wrapper + + +def test_cfm_unseen_shape_is_lazily_captured(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _cfm_mock_wrapper(monkeypatch) + wrapper._capture_graph = Mock(return_value=None) + + inputs = _cfm_inputs(2, 10, 0) + wrapper.replay(*inputs) + + wrapper._capture_graph.assert_called_once() + wrapper.graph_fn.assert_called_once() + + +def test_cfm_capture_failure_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _cfm_mock_wrapper(monkeypatch) + wrapper._capture_graph = lru_cache(maxsize=1)(lambda key: None) + + inputs = _cfm_inputs(2, 10, 0) + result = wrapper.replay(*inputs) + + wrapper.graph_fn.assert_called_once() + assert result[0] is wrapper.graph_fn.return_value + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None: + """Verify that LRU cache evicts oldest entries when full.""" + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _MiniDiT().eval().cuda() + wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=2) + + with torch.inference_mode(): + inputs_a = _cfm_inputs(2, 10, 0) + wrapper.replay(*inputs_a) + assert wrapper._capture_graph.cache_info().currsize == 1 + + inputs_b = _cfm_inputs(2, 12, 0) + wrapper.replay(*inputs_b) + assert wrapper._capture_graph.cache_info().currsize == 2 + + wrapper.replay(*inputs_a) + assert wrapper._capture_graph.cache_info().currsize == 2 + + inputs_c = _cfm_inputs(2, 14, 0) + wrapper.replay(*inputs_c) + assert wrapper._capture_graph.cache_info().currsize == 2 + + wrapper.replay(*inputs_b) + hits_before = wrapper._capture_graph.cache_info().hits + wrapper.replay(*inputs_b) + hits_after = wrapper._capture_graph.cache_info().hits + assert hits_after == hits_before + 1 + + wrapper._capture_graph.cache_clear() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cfm_none_cache_parity_between_graph_and_eager( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verify that None→zeros (graph path) matches [None]*depth (eager path). + + setup_batch calls _decode_cfm with cnn_cache=None on every request's + prompt pass. The graph path replaces None with zeros; the eager path + passes [None]*depth. Both must produce identical outputs. + """ + torch.accelerator.synchronize() + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _MiniDiT().eval().cuda() + wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=32) + + depth = len(estimator.blocks) + hidden = 8 + batch_size = 2 + chunk_size = 20 + + estimator_input = torch.randn(batch_size, 16, chunk_size, device="cuda") + time_emb = torch.randn(batch_size, 1, hidden, device="cuda") + cnn_out = torch.empty(depth, batch_size, hidden, 2, device="cuda") + att_out = torch.empty(depth, batch_size, 1, chunk_size, hidden, device="cuda") + + # Eager path: cnn_cache=[None]*depth, att_cache=[None]*depth + eager_input = estimator_input.clone() + eager_time = time_emb.clone() + eager_cnn_out = torch.empty_like(cnn_out) + eager_att_out = torch.empty_like(att_out) + with torch.no_grad(): + eager_result = estimator.blocks_forward_chunk( + eager_input, + eager_time, + None, + [None] * depth, + [None] * depth, + eager_cnn_out, + eager_att_out, + ) + + # Graph path: cnn_cache=zeros, att_cache=zero-length (simulating _estimator_step's None→zeros) + zero_cnn = torch.zeros_like(cnn_out) + zero_att = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:]) + wrapper.replay(estimator_input, time_emb, zero_cnn, zero_att, cnn_out, att_out) + graph_input = estimator_input.clone() + graph_time = time_emb.clone() + graph_cnn_out = torch.empty_like(cnn_out) + graph_att_out = torch.empty_like(att_out) + graph_zero_cnn = torch.zeros_like(cnn_out) + graph_zero_att = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:]) + graph_result, graph_cnn, graph_att = wrapper.replay( + graph_input, + graph_time, + graph_zero_cnn, + graph_zero_att, + graph_cnn_out, + graph_att_out, + ) + + torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_cnn, eager_cnn_out, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_att, eager_att_out, rtol=1e-4, atol=1e-5) diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index fde3ea878..d57671adc 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -17,6 +17,7 @@ from vllm_omni.model_extras import ( build_x_to_text_prompt, get_extra_body_params, get_extra_output_params, + get_output_tensor_range, get_x_to_text_model_family, should_init_extra_args_for_non_diffusion_stages, ) @@ -134,6 +135,8 @@ def test_lingbot_extra_registry_declares_request_params() -> None: } ) assert get_extra_output_params("LingBotVideoPipeline") == frozenset() + assert get_output_tensor_range("LingBotVideoPipeline") == "zero_to_one" + assert get_output_tensor_range("WanPipeline") == "negative_one_to_one" assert should_init_extra_args_for_non_diffusion_stages("LingBotVideoPipeline") is False diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 668954f5d..bc34bfe9f 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -19,6 +19,8 @@ connectors: # A future initial_codec_chunk_frames shape will use lazy capture unless # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. enable_hift_graph: true + enable_cfm_graph: true + cfm_max_graphs: 32 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index e44140b2d..2a4abed42 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml @@ -22,6 +22,8 @@ connectors: # A future initial_codec_chunk_frames shape will use lazy capture unless # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. enable_hift_graph: true + enable_cfm_graph: true + cfm_max_graphs: 32 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index 4933d0c74..b34673910 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml @@ -22,6 +22,8 @@ connectors: # A future initial_codec_chunk_frames shape will use lazy capture unless # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. enable_hift_graph: true + enable_cfm_graph: true + cfm_max_graphs: 32 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml index 69fa6d5b6..813a2e2eb 100644 --- a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml @@ -33,6 +33,8 @@ connectors: # A future initial_codec_chunk_frames shape will use lazy capture unless # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. enable_hift_graph: true + enable_cfm_graph: true + cfm_max_graphs: 32 stages: - stage_id: 0 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py index 651284fe0..41223e864 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py @@ -14,7 +14,7 @@ import torch.nn as nn import torch.nn.functional as F from vllm.logger import init_logger -from .cuda_graph_wrapper import HiFTGraphWrapper +from .cuda_graph_wrapper import CFMGraphWrapper, HiFTGraphWrapper logger = init_logger(__name__) @@ -73,6 +73,7 @@ class BatchedToken2Wav(nn.Module): *, connector_config: Mapping[str, int] | None = None, hift_graph_config: Mapping[str, Any] | None = None, + cfm_graph_config: Mapping[str, Any] | None = None, ): super().__init__() self._token2wav = token2wav @@ -135,6 +136,22 @@ class BatchedToken2Wav(nn.Module): with torch.inference_mode(), _autocast_disabled(hift_parameter.device): self.hift_graph_wrapper.capture() logger.info("HiFT CUDA Graph captured successfully") + self._cfm_graph_wrapper: CFMGraphWrapper | None = None + cfm_graph_cfg = dict(cfm_graph_config or {}) + if bool(cfm_graph_cfg.get("enabled", False)): + flow_parameter = next(self.flow.parameters(), None) + if flow_parameter is not None and flow_parameter.device.type == "cuda": + estimator = self.flow.decoder.estimator + self._cfm_graph_wrapper = CFMGraphWrapper( + graph_fn=estimator.blocks_forward_chunk, + max_graphs=int(cfm_graph_cfg.get("max_graphs", 32)), + ) + logger.info("CFM CUDA Graph enabled (max_graphs=%d)", int(cfm_graph_cfg.get("max_graphs", 32))) + else: + logger.info( + "CFM CUDA Graph is disabled on device type %s", + flow_parameter.device.type if flow_parameter is not None else "unknown", + ) self._prompt_features: dict[tuple[str, str], PromptFeatures] = {} def _hift_inference( @@ -265,6 +282,19 @@ class BatchedToken2Wav(nn.Module): speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width) estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1) cnn_out, att_out = self._estimator_buffers(estimator, estimator_input, att_cache) + if self._cfm_graph_wrapper is not None: + if cnn_cache is None: + cnn_cache = torch.zeros_like(cnn_out) + if att_cache is None: + att_cache = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:]) + return self._cfm_graph_wrapper.replay( + estimator_input, + time_embedding, + cnn_cache, + att_cache, + cnn_out, + att_out, + ) old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks) old_att: Any = att_cache if att_cache is not None else [None] * len(estimator.blocks) result = estimator.blocks_forward_chunk( diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py index 5f6c965d5..97015855e 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from functools import lru_cache + import torch from torch.cuda import CUDAGraph from vllm.logger import init_logger @@ -133,3 +135,126 @@ class HiFTGraphWrapper: cache_source = static_cache_source_output[:batch_size].clone() speech = self.finalize_fn(static_magnitude_output[:batch_size], static_phase_output[:batch_size]).clone() return speech, cache_source + + +def _tensor_signature(value: torch.Tensor) -> tuple: + return tuple(value.shape), str(value.dtype), str(value.device) + + +_DTYPE_MAP = { + "torch.float32": torch.float32, + "torch.float16": torch.float16, + "torch.bfloat16": torch.bfloat16, + "torch.float64": torch.float64, +} + + +def _tensors_from_key(key: tuple) -> tuple[torch.Tensor, ...]: + """Rebuild zero tensors from a cache key (shape, dtype, device tuples).""" + tensors = [] + for shape, dtype_str, device_str in key[1:]: + dtype = _DTYPE_MAP.get(dtype_str, torch.float32) + device = torch.device(device_str) + tensors.append(torch.zeros(shape, dtype=dtype, device=device)) + return tuple(tensors) + + +class CFMGraphWrapper: + """Per-shape CUDA graph capture/replay for the CFM DiT estimator. + + Captures one blocks_forward_chunk call (in_proj -> DiT blocks -> final_layer) + as the graph target. The 10-step Euler loop stays in Python, replaying + the graph 10 times per decode. + + Uses functools.lru_cache for automatic LRU eviction. Cache misses + trigger capture; capture failures fall back to eager. Outputs are + cloned after replay to prevent streaming cache corruption. + """ + + def __init__(self, graph_fn, *, max_graphs: int = 32) -> None: + self.graph_fn = graph_fn + self.max_graphs = int(max_graphs) + self.device = next(graph_fn.__self__.parameters()).device + + @lru_cache(maxsize=self.max_graphs) + def _capture_graph(key: tuple): + return self._capture(key) + + self._capture_graph = _capture_graph + + def _call_graph_fn(self, args: tuple[torch.Tensor, ...]) -> torch.Tensor: + return self.graph_fn(args[0], args[1], None, args[2], args[3], args[4], args[5]) + + def _capture(self, key: tuple) -> tuple | None: + """Capture a CUDA graph for the given key. Returns None on failure.""" + static_inputs = _tensors_from_key(key) + + current_stream = torch.cuda.current_stream(self.device) + warmup_stream = torch.cuda.Stream(device=self.device) + warmup_stream.wait_stream(current_stream) + with torch.cuda.stream(warmup_stream), torch.no_grad(): + for _ in range(3): + self._call_graph_fn(static_inputs) + current_stream.wait_stream(warmup_stream) + + try: + graph = CUDAGraph() + with torch.no_grad(), torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + static_output = self._call_graph_fn(static_inputs) + except Exception: + logger.warning( + "CFM graph capture failed for shape=%s; using eager", + key, + exc_info=True, + ) + return None + + logger.info( + "Captured CFM CUDA Graph for shape %s (cache=%d/%d, hits=%d, misses=%d)", + key, + self._capture_graph.cache_info().currsize, + self.max_graphs, + self._capture_graph.cache_info().hits, + self._capture_graph.cache_info().misses, + ) + return (static_inputs, static_output, graph) + + def replay( + self, + estimator_input: torch.Tensor, + time_emb: torch.Tensor, + cnn_cache: torch.Tensor, + att_cache: torch.Tensor, + cnn_out: torch.Tensor, + att_out: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + inputs = (estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out) + + if torch.cuda.is_current_stream_capturing() or estimator_input.device.type != "cuda": + with torch.no_grad(): + result = self._call_graph_fn(inputs) + return result, inputs[4], inputs[5] + + key = ("estimator_step",) + tuple(_tensor_signature(v) for v in inputs) + entry = self._capture_graph(key) + + if entry is None: + logger.debug( + "CFM graph eager fallback for shape=%s (hits=%d, misses=%d)", + key, + self._capture_graph.cache_info().hits, + self._capture_graph.cache_info().misses, + ) + with torch.no_grad(): + result = self._call_graph_fn(inputs) + return result, inputs[4], inputs[5] + + static_inputs, static_output, graph = entry + for static, current in zip(static_inputs, inputs, strict=True): + static.copy_(current) + graph.replay() + return ( + static_output.detach().clone(), + static_inputs[4].detach().clone(), + static_inputs[5].detach().clone(), + ) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index d6d3e3d47..686b39c5b 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -162,6 +162,10 @@ class MiniCPMO45Code2Wav(nn.Module): "enabled": bool(extra.get("enable_hift_graph", False)), "capture_batch_sizes": capture_batch_sizes, } + self._cfm_graph_config = { + "enabled": bool(extra.get("enable_cfm_graph", False)), + "max_graphs": int(extra.get("cfm_max_graphs", 32)), + } self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1)) if self._min_batch_size < 1: raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1") @@ -802,4 +806,5 @@ class MiniCPMO45Code2Wav(nn.Module): trt_stepper=trt_stepper, connector_config=self._connector_config, hift_graph_config=self._hift_graph_config, + cfm_graph_config=self._cfm_graph_config, ) diff --git a/vllm_omni/model_extras/__init__.py b/vllm_omni/model_extras/__init__.py index 1a3d71870..b11ceadec 100644 --- a/vllm_omni/model_extras/__init__.py +++ b/vllm_omni/model_extras/__init__.py @@ -9,6 +9,7 @@ from vllm_omni.model_extras.registry import ( get_extra_body_params, get_extra_output_params, get_model_class_name, + get_output_tensor_range, get_x_to_text_model_family, should_init_extra_args_for_non_diffusion_stages, ) @@ -21,6 +22,7 @@ __all__ = [ "get_extra_body_params", "get_extra_output_params", "get_model_class_name", + "get_output_tensor_range", "get_x_to_text_model_family", "should_init_extra_args_for_non_diffusion_stages", ] diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 4c41c5920..8d7ee1e29 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -4,7 +4,7 @@ from __future__ import annotations from collections.abc import Callable, Mapping -from typing import Any +from typing import Any, Literal from PIL import Image @@ -93,6 +93,7 @@ ImageToVideoPromptBuilder = Callable[ dict[str, Any], ] XToTextPromptBuilder = Callable[[str, str, bool], tuple[dict[str, Any], list[int] | None]] +OutputTensorRange = Literal["negative_one_to_one", "zero_to_one"] def default_x_to_text_prompt( @@ -193,6 +194,7 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { }, "LingBotVideoPipeline": { "extra_body_params": LINGBOT_VIDEO_EXTRA_BODY_PARAMS, + "output_tensor_range": "zero_to_one", # Shared T2I/I2V envelopes select the output modality. LingBot's # pipeline owns model-specific validation and normalization. }, @@ -262,6 +264,18 @@ def get_extra_output_params(model_class_name: str | None) -> frozenset[str]: return spec.get("extra_output_params", frozenset()) if spec is not None else frozenset() +def get_output_tensor_range(model_class_name: str | None) -> OutputTensorRange: + """Return the declared range for floating-point tensor outputs. + + The default preserves the shared examples' historical handling. Pipelines + that already return normalized tensors declare ``zero_to_one`` explicitly. + """ + spec = _get_spec(model_class_name) + if spec is None: + return "negative_one_to_one" + return spec.get("output_tensor_range", "negative_one_to_one") + + def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None) -> bool: spec = _get_spec(model_class_name) return bool(spec and spec.get("init_extra_args_for_non_diffusion_stages", False))