diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index dd62ee68b..d0f866b5a 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -1322,6 +1322,44 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate + - label: ":full_moon: Diffusion X2V · LingBot Video MoE Function Test" + timeout_in_minutes: 90 + commands: + - pytest -s -v tests/e2e/online_serving/test_lingbot_video_moe.py -m "full_model and cuda" --run-level "full_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + - label: ":full_moon: Diffusion X2V · Doc Test" timeout_in_minutes: 60 commands: diff --git a/benchmarks/README.md b/benchmarks/README.md index 37a6264a9..0a8fb8058 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -28,6 +28,14 @@ Benchmarks for GLM-Image performance across HuggingFace baseline, vLLM-Omni offl - **Tasks**: text-to-image and image-to-image - **Key metrics**: request/image throughput, latency percentiles, optional per-stage pipeline timings +### [LingBot-Video](lingbot_video/README.md) — Dense and MoE Parity + +Manual cross-runtime validation for the dense LingBot-Video pipeline and the +LingBot-Video MoE transformer. + +- **Dense pipeline**: decoded-video MAE, MSE, PSNR, latency, and optional steady-state timings +- **MoE transformer**: bitwise router, sparse-block, shared-expert, and full-transformer parity + ### [Distributed](distributed/omni_connectors/README.md) — RDMA Connector Testing RDMA environment setup and transfer tests for `MooncakeTransferEngineConnector`, including pytest-based single-node checks and manual cross-node benchmarks. diff --git a/benchmarks/lingbot_video/README.md b/benchmarks/lingbot_video/README.md new file mode 100644 index 000000000..687973659 --- /dev/null +++ b/benchmarks/lingbot_video/README.md @@ -0,0 +1,62 @@ +# LingBot-Video Benchmarks + +These manual benchmarks compare the upstream LingBot-Video implementation with +vLLM-Omni. Run all commands from the repository root. + +| Benchmark | Script | Purpose | +|-----------|--------|---------| +| Dense pipeline parity | `dense_pipeline_parity.py` | Compare generated videos and report MAE, MSE, PSNR, latency, and optional steady-state timings | +| MoE numerical parity | `moe_transformer_parity.py` | Check bitwise parity for the sparse MoE block and the full MoE transformer | + +Both benchmarks require a local LingBot-Video checkout and local model files. +They do not run in CI. + +## Dense pipeline parity + +The dense benchmark launches the upstream Diffusers inference script and the +vLLM-Omni offline example with identical prompts and sampling parameters. It +then compares the decoded MP4 frames. + +```bash +CUDA_VISIBLE_DEVICES=0 \ +python benchmarks/lingbot_video/dense_pipeline_parity.py \ + --model /path/to/lingbot-video-dense-1.3b \ + --official-repo /path/to/lingbot-video \ + --output-dir /tmp/lingbot_dense_parity +``` + +Use `--runs N` to additionally measure repeated in-process vLLM-Omni +requests. The video comparison reports MAE, MSE, and PSNR; it is an end-to-end +pipeline diagnostic rather than a bitwise transformer check. + +## MoE numerical parity + +The lightweight block comparison uses deterministic weights and inputs to +cover correction-bias routing, group-limited top-k, routed experts, FP32 +scatter-weighted restore, padding masks, and the shared expert. + +```bash +CUDA_VISIBLE_DEVICES=0 \ +python benchmarks/lingbot_video/moe_transformer_parity.py \ + --scope block \ + --official-repo /path/to/lingbot-video \ + --output-json /tmp/lingbot_moe_block_parity.json +``` + +The full comparison loads the upstream and vLLM-Omni 30B transformers +sequentially and compares their output tensors. It requires a CUDA device with +`torch._grouped_mm` and enough memory for one 30B BF16 transformer. + +```bash +CUDA_VISIBLE_DEVICES=0 \ +python benchmarks/lingbot_video/moe_transformer_parity.py \ + --scope transformer \ + --official-repo /path/to/lingbot-video \ + --model /path/to/lingbot-video-moe-30b-a3b \ + --output-json /tmp/lingbot_moe_transformer_parity.json +``` + +For the correctness oracle, the upstream transformer uses +`diffusers:_native_math` and vLLM-Omni uses +`TORCH_SDPA + SDPBackend.MATH`. The command exits with a nonzero status unless +the selected comparisons are bitwise equal. diff --git a/benchmarks/lingbot_video/dense_pipeline_parity.py b/benchmarks/lingbot_video/dense_pipeline_parity.py new file mode 100644 index 000000000..0c0c1d7a0 --- /dev/null +++ b/benchmarks/lingbot_video/dense_pipeline_parity.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import argparse +import json +import re +import statistics +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +import imageio.v3 as iio +import numpy as np +import torch + +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.platforms import current_omni_platform + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="LingBot dense parity and timing harness.") + parser.add_argument("--model", default="/home/models/lingbot-video-dense-1.3b") + parser.add_argument("--official-repo", default="/tmp/lingbot-video") + parser.add_argument("--output-dir", default="/tmp/lingbot_dense_parity") + parser.add_argument("--prompt", default="a robotic arm picks up a red block") + parser.add_argument("--height", type=int, default=192) + parser.add_argument("--width", type=int, default=320) + parser.add_argument("--num-frames", type=int, default=9) + parser.add_argument("--steps", type=int, default=2) + parser.add_argument("--guidance-scale", type=float, default=3.0) + parser.add_argument("--shift", type=float, default=3.0) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--fps", type=int, default=24) + parser.add_argument("--runs", type=int, default=0, help="Optional in-process steady-state native runs.") + return parser.parse_args() + + +def _run(cmd: list[str], cwd: Path | None = None) -> tuple[float, str]: + start = time.perf_counter() + proc = subprocess.run(cmd, cwd=cwd, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + elapsed = time.perf_counter() - start + if proc.returncode != 0: + print(proc.stdout) + raise subprocess.CalledProcessError(proc.returncode, proc.args, output=proc.stdout) + return elapsed, proc.stdout + + +def _read_video(path: Path) -> np.ndarray: + frames = iio.imread(path, plugin="pyav") + if frames.dtype != np.float32: + frames = frames.astype(np.float32) / 255.0 + return frames + + +def _compare(a: np.ndarray, b: np.ndarray) -> dict[str, Any]: + if a.shape != b.shape: + return {"shape_a": list(a.shape), "shape_b": list(b.shape), "shape_match": False} + diff = a - b + mse = float(np.mean(diff * diff)) + mae = float(np.mean(np.abs(diff))) + psnr = float("inf") if mse == 0 else float(20.0 * np.log10(1.0 / np.sqrt(mse))) + return { + "shape": list(a.shape), + "shape_match": True, + "mae": mae, + "mse": mse, + "psnr": psnr, + } + + +def _extract_total_generation_seconds(log: str) -> float | None: + match = re.search(r"Total generation time:\s*([0-9.]+)\s*seconds", log) + return float(match.group(1)) if match else None + + +def _peak_memory_mb(output: Any) -> float: + result = output[0] if isinstance(output, list) and output else output + return float(getattr(result, "peak_memory_mb", 0.0) or 0.0) + + +def _benchmark_native_requests(args: argparse.Namespace) -> dict[str, Any]: + load_start = time.perf_counter() + omni = Omni( + model=args.model, + model_class_name="LingBotVideoPipeline", + flow_shift=args.shift, + parallel_config=DiffusionParallelConfig(), + ) + load_seconds = time.perf_counter() - load_start + prompt = {"prompt": args.prompt} + latencies: list[float] = [] + peak_mb = 0.0 + + for _ in range(args.runs): + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) + sampling_params = OmniDiffusionSamplingParams( + height=args.height, + width=args.width, + num_frames=args.num_frames, + num_inference_steps=args.steps, + guidance_scale=args.guidance_scale, + generator=generator, + output_type="pt", + extra_args={"shift": args.shift}, + ) + start = time.perf_counter() + output = omni.generate(prompt, sampling_params) + latencies.append(time.perf_counter() - start) + peak_mb = max(peak_mb, _peak_memory_mb(output)) + del output + omni.close() + + return { + "load_seconds": load_seconds, + "run_seconds": latencies, + "mean_seconds": statistics.mean(latencies) if latencies else None, + "median_seconds": statistics.median(latencies) if latencies else None, + "min_seconds": min(latencies) if latencies else None, + "max_seconds": max(latencies) if latencies else None, + "peak_mb": peak_mb, + } + + +def main() -> None: + args = parse_args() + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + official_out = output_dir / "official_diffusers.mp4" + native_out = output_dir / "vllm_omni_native.mp4" + + official_cmd = [ + sys.executable, + str(Path(args.official_repo) / "scripts" / "inference.py"), + "--backend", + "diffusers", + "--model_dir", + args.model, + "--mode", + "t2v", + "--prompt", + args.prompt, + "--output", + str(official_out), + "--height", + str(args.height), + "--width", + str(args.width), + "--num_frames", + str(args.num_frames), + "--steps", + str(args.steps), + "--guidance_scale", + str(args.guidance_scale), + "--shift", + str(args.shift), + "--seed", + str(args.seed), + "--fps", + str(args.fps), + "--quiet_progress", + ] + native_cmd = [ + sys.executable, + "examples/offline_inference/text_to_video/text_to_video_lingbot.py", + "--model", + args.model, + "--prompt", + args.prompt, + "--output", + str(native_out), + "--height", + str(args.height), + "--width", + str(args.width), + "--num-frames", + str(args.num_frames), + "--num-inference-steps", + str(args.steps), + "--guidance-scale", + str(args.guidance_scale), + "--flow-shift", + str(args.shift), + "--seed", + str(args.seed), + "--fps", + str(args.fps), + ] + + official_time, official_log = _run(official_cmd, cwd=Path(args.official_repo)) + native_time, native_log = _run(native_cmd, cwd=Path.cwd()) + + result = { + "official_seconds": official_time, + "native_process_seconds": native_time, + "native_request_seconds": _extract_total_generation_seconds(native_log), + "official_output": str(official_out), + "native_output": str(native_out), + "accuracy": _compare(_read_video(official_out), _read_video(native_out)), + "official_log_tail": official_log.splitlines()[-20:], + "native_log_tail": native_log.splitlines()[-20:], + } + if args.runs > 0: + result["native_steady_state"] = _benchmark_native_requests(args) + + result_path = output_dir / "comparison.json" + result_path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") + print(json.dumps(result, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/lingbot_video/moe_transformer_parity.py b/benchmarks/lingbot_video/moe_transformer_parity.py new file mode 100644 index 000000000..cdad31f19 --- /dev/null +++ b/benchmarks/lingbot_video/moe_transformer_parity.py @@ -0,0 +1,399 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Reproduce LingBot-Video MoE block and transformer numerical parity.""" + +from __future__ import annotations + +import argparse +import gc +import importlib +import json +import os +import sys +import time +from pathlib import Path +from typing import Any + +import torch + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--scope", + choices=("block", "transformer", "all"), + default="block", + help="Run the lightweight sparse-block check, the real transformer check, or both.", + ) + parser.add_argument( + "--official-repo", + type=Path, + required=True, + help="Local checkout of https://github.com/Robbyant/lingbot-video.", + ) + parser.add_argument( + "--model", + default=None, + help="Local MoE checkpoint path or a checkpoint already present in the Hugging Face cache.", + ) + parser.add_argument("--transformer-subfolder", default="transformer") + parser.add_argument("--output-json", type=Path, default=None) + parser.add_argument("--device", default="cuda") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--num-frames", type=int, default=1) + parser.add_argument("--height", type=int, default=8) + parser.add_argument("--width", type=int, default=8) + parser.add_argument("--text-length", type=int, default=16) + return parser.parse_args() + + +def _load_official_module(official_repo: Path): + transformer_file = official_repo / "lingbot_video" / "transformer_lingbot_video.py" + if not transformer_file.is_file(): + raise FileNotFoundError(f"Upstream transformer not found: {transformer_file}") + sys.path.insert(0, str(official_repo)) + module = importlib.import_module("lingbot_video.transformer_lingbot_video") + module_path = Path(module.__file__).resolve() + if official_repo.resolve() not in module_path.parents: + raise RuntimeError(f"Imported LingBot module from {module_path}, not {official_repo.resolve()}") + return module + + +def _load_native_module(): + return importlib.import_module("vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer") + + +def _tensor_metrics(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, Any]: + actual_float = actual.float() + expected_float = expected.float() + diff = (actual_float - expected_float).abs() + denominator = actual_float.norm() * expected_float.norm() + cosine = ( + float(torch.dot(actual_float.flatten(), expected_float.flatten()) / denominator) + if float(denominator) > 0 + else 1.0 + ) + return { + "shape": list(actual.shape), + "equal": bool(torch.equal(actual, expected)), + "max_abs": float(diff.max()) if diff.numel() else 0.0, + "mean_abs": float(diff.mean()) if diff.numel() else 0.0, + "rmse": float(torch.sqrt(torch.mean(diff.square()))) if diff.numel() else 0.0, + "cosine": cosine, + "finite": bool(torch.isfinite(actual).all() and torch.isfinite(expected).all()), + } + + +def _initialize_block(module: torch.nn.Module, seed: int) -> None: + generator = torch.Generator(device="cpu").manual_seed(seed) + with torch.no_grad(): + for parameter in module.parameters(): + values = torch.randn(parameter.shape, generator=generator, dtype=torch.float32) + parameter.copy_((values * 0.02).to(parameter.dtype)) + module.router.e_score_correction_bias.copy_(torch.tensor([0.9, 0.8, 1.0, 0.0, 0.7, 0.6, 0.5, 0.4])) + + +def _apply_padding( + scores: torch.Tensor, + padding_mask: torch.Tensor, + route_scale: float, +) -> torch.Tensor: + scores = scores * padding_mask.unsqueeze(-1).to(scores.dtype) + scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-9) + return scores * route_scale + + +def _run_block_parity( + official_module, + native_module, + *, + device: torch.device, + seed: int, +) -> dict[str, Any]: + if not hasattr(torch, "_grouped_mm"): + raise RuntimeError("The block parity check requires torch._grouped_mm.") + + common = { + "hidden_size": 16, + "num_experts": 8, + "top_k": 2, + "moe_intermediate_size": 8, + "score_func": "sigmoid", + "norm_topk_prob": True, + "n_group": 4, + "topk_group": 1, + "routed_scaling_factor": 1.5, + "n_shared_experts": 1, + } + official = official_module.LingBotVideoSparseMoeBlock( + intermediate_size=32, + **common, + ) + native = native_module.LingBotVideoSparseMoeBlock(**common) + _initialize_block(official, seed) + native.load_state_dict(official.state_dict(), strict=True) + official = official.to(device=device, dtype=torch.bfloat16).eval() + native = native.to(device=device, dtype=torch.bfloat16).eval() + + generator = torch.Generator(device=device).manual_seed(seed) + hidden_states = torch.randn( + 2, + 5, + common["hidden_size"], + generator=generator, + device=device, + dtype=torch.bfloat16, + ) + padding_mask = torch.tensor( + [1, 1, 1, 1, 1, 1, 1, 1, 0, 0], + device=device, + dtype=torch.float32, + ) + tokens = hidden_states.reshape(-1, common["hidden_size"]) + + with torch.inference_mode(): + official_router = official.router(tokens) + official_indices, official_scores = official_router[:2] + native_indices, native_scores = native.router(tokens) + + official_scores = _apply_padding( + official_scores, + padding_mask, + official.router.route_scale, + ) + native_scores = _apply_padding( + native_scores, + padding_mask, + native.router.route_scale, + ) + official_routed = official._run_selected_experts( + tokens, + official_scores, + official_indices, + ) + native_routed = native._run_selected_experts( + tokens, + native_scores, + native_indices, + ) + official_shared = official.shared_experts(hidden_states) + native_shared = native.shared_experts(hidden_states) + official_output = official(hidden_states, padding_mask=padding_mask) + native_output = native(hidden_states, padding_mask=padding_mask) + + result = { + "router_indices_equal": bool(torch.equal(native_indices, official_indices)), + "router_scores": _tensor_metrics(native_scores, official_scores), + "routed_output": _tensor_metrics(native_routed, official_routed), + "shared_output": _tensor_metrics(native_shared, official_shared), + "final_output": _tensor_metrics(native_output, official_output), + } + result["exact"] = bool( + result["router_indices_equal"] + and result["router_scores"]["equal"] + and result["routed_output"]["equal"] + and result["shared_output"]["equal"] + and result["final_output"]["equal"] + ) + return result + + +def _release_cuda() -> None: + gc.collect() + torch.accelerator.empty_cache() + torch.accelerator.synchronize() + + +def _make_transformer_inputs( + config, + args: argparse.Namespace, +) -> dict[str, torch.Tensor]: + generator = torch.Generator(device="cpu").manual_seed(args.seed) + hidden_states = torch.randn( + 1, + int(config.in_channels), + args.num_frames, + args.height, + args.width, + generator=generator, + dtype=torch.float32, + ).to(torch.bfloat16) + encoder_hidden_states = torch.randn( + 1, + args.text_length, + int(config.text_dim), + generator=generator, + dtype=torch.float32, + ).to(torch.bfloat16) + return { + "hidden_states": hidden_states, + "timestep": torch.tensor([500.0], dtype=torch.float32), + "encoder_hidden_states": encoder_hidden_states, + "encoder_attention_mask": torch.ones(1, args.text_length, dtype=torch.long), + } + + +def _load_transformer( + transformer_cls, + args: argparse.Namespace, + device: torch.device, +): + start = time.perf_counter() + model = transformer_cls.from_pretrained( + args.model, + subfolder=args.transformer_subfolder, + torch_dtype=torch.bfloat16, + local_files_only=True, + low_cpu_mem_usage=True, + ) + model = model.to(device=device, dtype=torch.bfloat16).eval() + torch.accelerator.synchronize() + return model, time.perf_counter() - start + + +def _forward_transformer( + model, + inputs: dict[str, torch.Tensor], + device: torch.device, + attention_context, +) -> tuple[torch.Tensor, float, float]: + torch.accelerator.reset_peak_memory_stats(device) + gpu_inputs = {name: value.to(device) for name, value in inputs.items()} + start = time.perf_counter() + with torch.inference_mode(), attention_context: + output = model(**gpu_inputs, return_dict=False)[0] + torch.accelerator.synchronize() + elapsed = time.perf_counter() - start + peak_gib = torch.accelerator.max_memory_reserved(device) / (1024**3) + return output.cpu(), elapsed, peak_gib + + +def _run_transformer_parity( + official_module, + native_module, + *, + args: argparse.Namespace, + device: torch.device, +) -> dict[str, Any]: + if args.model is None: + raise ValueError("--model is required for transformer parity.") + + from diffusers.models.attention_dispatch import attention_backend + from torch.nn.attention import SDPBackend, sdpa_kernel + + from vllm_omni.diffusion.config import set_current_diffusion_config + from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig + + official, official_load = _load_transformer( + official_module.LingBotVideoTransformer3DModel, + args, + device, + ) + inputs = _make_transformer_inputs(official.config, args) + official_output, official_forward, official_peak = _forward_transformer( + official, + inputs, + device, + attention_backend("_native_math"), + ) + del official + _release_cuda() + + native_config = OmniDiffusionConfig( + diffusion_attention_config=AttentionConfig( + default=AttentionSpec(backend="TORCH_SDPA"), + ), + ) + with set_current_diffusion_config(native_config): + native, native_load = _load_transformer( + native_module.LingBotVideoTransformer3DModel, + args, + device, + ) + native_backend_prefs = {block.attn.attn.backend_pref for block in native.blocks} + if native_backend_prefs != {"TORCH_SDPA"}: + raise RuntimeError( + f"Failed to force the native transformer to TORCH_SDPA: resolved preferences were {native_backend_prefs}." + ) + native_output, native_forward, native_peak = _forward_transformer( + native, + inputs, + device, + sdpa_kernel(SDPBackend.MATH), + ) + del native + _release_cuda() + + metrics = _tensor_metrics(native_output, official_output) + return { + "backends": { + "official": "diffusers:_native_math", + "native": "TORCH_SDPA+SDPBackend.MATH", + }, + "official": { + "load_seconds": official_load, + "forward_seconds": official_forward, + "peak_reserved_gib": official_peak, + }, + "native": { + "load_seconds": native_load, + "forward_seconds": native_forward, + "peak_reserved_gib": native_peak, + }, + "output": metrics, + "exact": metrics["equal"], + } + + +def main() -> int: + args = parse_args() + device = torch.device(args.device) + if device.type != "cuda" or not torch.cuda.is_available(): + raise RuntimeError("LingBot MoE parity requires an available CUDA device.") + + os.environ["DIFFUSERS_ATTN_BACKEND"] = "_native_math" + os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA" + os.environ["LINGBOT_MOE_EXPERT_BACKEND"] = "grouped_mm" + os.environ["LINGBOT_MOE_PAD_BACKEND"] = "loop" + os.environ["LINGBOT_MOE_REORDER_BACKEND"] = "sort" + os.environ["LINGBOT_MOE_RESTORE_BACKEND"] = "scatter" + + official_module = _load_official_module(args.official_repo) + native_module = _load_native_module() + + result: dict[str, Any] = { + "settings": { + "scope": args.scope, + "seed": args.seed, + "device": str(device), + "official_repo": str(args.official_repo.resolve()), + "model": args.model, + } + } + if args.scope in {"block", "all"}: + result["block"] = _run_block_parity( + official_module, + native_module, + device=device, + seed=args.seed, + ) + if args.scope in {"transformer", "all"}: + result["transformer"] = _run_transformer_parity( + official_module, + native_module, + args=args, + device=device, + ) + + result["exact"] = all(section["exact"] for name, section in result.items() if name in {"block", "transformer"}) + payload = json.dumps(result, indent=2, sort_keys=True) + print(payload) + if args.output_json is not None: + args.output_json.parent.mkdir(parents=True, exist_ok=True) + args.output_json.write_text(payload + "\n", encoding="utf-8") + return 0 if result["exact"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index c1dac6ca7..dda7c5f2d 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -47,6 +47,7 @@ th { | `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | | `LTX23Pipeline` | LTX-2.3-T2V | `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | | `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | +| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | 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 new file mode 100644 index 000000000..a8f8a2015 --- /dev/null +++ b/examples/offline_inference/text_to_video/text_to_video_lingbot.py @@ -0,0 +1,124 @@ +# 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="/home/models/lingbot-video-dense-1.3b") + parser.add_argument("--prompt", default="a robotic arm picks up a red block") + parser.add_argument("--negative-prompt", default="") + parser.add_argument("--output", default="lingbot_video_output.mp4") + parser.add_argument("--height", type=int, default=192) + parser.add_argument("--width", type=int, default=320) + parser.add_argument("--num-frames", type=int, default=9) + parser.add_argument("--num-inference-steps", type=int, default=2) + parser.add_argument("--guidance-scale", type=float, default=3.0) + parser.add_argument("--flow-shift", type=float, default=3.0) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--fps", type=int, default=24) + parser.add_argument("--output-type", choices=["pt", "np"], default="pt") + parser.add_argument("--batch-cfg", action="store_true") + return parser.parse_args() + + +def _peak_memory_mb(output: Any) -> float: + result = output[0] if isinstance(output, list) and output else output + return float(getattr(result, "peak_memory_mb", 0.0) or 0.0) + + +def _extract_video(output: Any) -> Any: + if isinstance(output, list): + output = output[0] if output else None + if isinstance(output, OmniRequestOutput): + if output.final_output_type != "image": + raise ValueError(f"Unexpected output type {output.final_output_type!r}.") + if not output.images: + raise ValueError("No video frames found in OmniRequestOutput.") + return output.images[0] if len(output.images) == 1 else output.images + return output + + +def _to_frame_array(video: Any) -> np.ndarray: + if isinstance(video, torch.Tensor): + tensor = video.detach().cpu() + if tensor.dim() == 5: + tensor = tensor[0] if tensor.shape[1] not in (3, 4) else tensor[0].permute(1, 2, 3, 0) + elif tensor.dim() == 4 and tensor.shape[0] in (3, 4): + tensor = tensor.permute(1, 2, 3, 0) + if tensor.is_floating_point(): + if float(tensor.min()) < 0.0: + tensor = tensor.clamp(-1, 1) * 0.5 + 0.5 + else: + tensor = tensor.clamp(0, 1) + return tensor.float().numpy() + if isinstance(video, np.ndarray): + array = video[0] if video.ndim == 5 else video + if np.issubdtype(array.dtype, np.integer): + return array.astype(np.float32) / 255.0 + return array.astype(np.float32, copy=False) + if isinstance(video, list): + return np.stack([_to_frame_array(frame) for frame in video], axis=0) + return np.asarray(video, dtype=np.float32) + + +def main() -> None: + args = parse_args() + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) + omni = Omni( + model=args.model, + model_class_name="LingBotVideoPipeline", + flow_shift=args.flow_shift, + parallel_config=DiffusionParallelConfig(), + ) + + prompt: dict[str, str] = {"prompt": args.prompt} + if args.negative_prompt: + prompt["negative_prompt"] = args.negative_prompt + + sampling_params = OmniDiffusionSamplingParams( + height=args.height, + width=args.width, + num_frames=args.num_frames, + num_inference_steps=args.num_inference_steps, + guidance_scale=args.guidance_scale, + generator=generator, + output_type=args.output_type, + extra_args={"batch_cfg": args.batch_cfg, "shift": args.flow_shift}, + ) + + start = time.perf_counter() + output = omni.generate(prompt, sampling_params) + elapsed = time.perf_counter() - start + print(f"Total generation time: {elapsed:.4f} seconds ({elapsed * 1000:.2f} ms)") + peak_mb = _peak_memory_mb(output) + if peak_mb: + print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)") + + video = _to_frame_array(_extract_video(output)) + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + from diffusers.utils import export_to_video + + export_to_video(list(video), str(output_path), fps=args.fps) + print(f"Saved generated video to {output_path}") + omni.close() + + +if __name__ == "__main__": + main() diff --git a/recipes/README.md b/recipes/README.md index d335183ca..be4c7cc08 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -49,6 +49,7 @@ recipes/ | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) | | [`MammothModa2/MammothModa2-Preview.md`](./MammothModa2/MammothModa2-Preview.md) | Text-to-image with the shared offline image example (AR → DiT) | 1x L40S 48GB / 1x ≥40GB GPU | | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB | +| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) | | [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 | | [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 | | [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB | diff --git a/recipes/Robbyant/LingBot-Video.md b/recipes/Robbyant/LingBot-Video.md new file mode 100644 index 000000000..289d36990 --- /dev/null +++ b/recipes/Robbyant/LingBot-Video.md @@ -0,0 +1,259 @@ +# LingBot-Video + +> Native dense and MoE text-to-video serving for LingBot-Video + +## Summary + +- Vendor: Robbyant +- Models: `robbyant/lingbot-video-dense-1.3b` and `robbyant/lingbot-video-moe-30b-a3b` +- Task: Text-to-video generation +- Mode: Offline generation and online serving with the OpenAI-compatible `/v1/videos` API +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want to run a dense or MoE LingBot-Video checkpoint +with vLLM-Omni's native pipeline. The runtime path does not import the upstream +`lingbot_video` Python package; it loads checkpoint components directly with +the in-tree `LingBotVideoPipeline`, dense or routed-MoE DiT blocks, shared +FlowUniPC scheduler, Qwen3-VL text encoder, and Wan VAE. + +This first MoE integration is intentionally text-to-video and single-GPU only. +It targets the base BF16 checkpoint without its optional refiner or additional +parallel, cache, quantized, or expert-kernel backends. + +## References + +- 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 online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md) + +## Hardware Support + +This recipe documents the CUDA single-GPU dense and BF16 MoE checkpoint paths. +Multi-GPU parallelism, Cache-DiT, quantization, and CPU offload are not +validated for LingBot-Video in this PR. + +## GPU + +### 1 x NVIDIA L20X + +Both checkpoints have been smoke-tested on one NVIDIA L20X at `192x320`, +9 frames, and 2 steps. The MoE smoke reserved approximately `67.70 GiB` of GPU +memory, so use a GPU with at least about 70 GiB of available memory for this +small validation shape. Larger resolutions, frame counts, or concurrent +requests require additional headroom. + +The MoE path is validated with BF16 expert weights. + +#### Dense offline T2V + +```bash +CUDA_VISIBLE_DEVICES=0 \ +python examples/offline_inference/text_to_video/text_to_video_lingbot.py \ + --model robbyant/lingbot-video-dense-1.3b \ + --prompt "a robotic arm picks up a red block" \ + --output lingbot_t2v.mp4 \ + --height 192 \ + --width 320 \ + --num-frames 9 \ + --num-inference-steps 2 \ + --guidance-scale 3.0 \ + --flow-shift 3.0 \ + --seed 42 \ + --fps 24 +``` + +#### MoE offline T2V + +```bash +CUDA_VISIBLE_DEVICES=0 \ +python examples/offline_inference/text_to_video/text_to_video_lingbot.py \ + --model robbyant/lingbot-video-moe-30b-a3b \ + --prompt "a robotic arm picks up a red block" \ + --output lingbot_moe_t2v.mp4 \ + --height 192 \ + --width 320 \ + --num-frames 9 \ + --num-inference-steps 2 \ + --guidance-scale 3.0 \ + --flow-shift 3.0 \ + --seed 42 \ + --fps 24 +``` + +#### Online serving + +```bash +CUDA_VISIBLE_DEVICES=0 \ +vllm serve robbyant/lingbot-video-dense-1.3b \ + --omni \ + --model-class-name LingBotVideoPipeline \ + --default-sampling-params \ + '{"0":{"num_frames":81,"num_inference_steps":40,"guidance_scale":6.0}}' \ + --port 8091 +``` + +For the MoE checkpoint, use the same single-GPU pipeline: + +```bash +CUDA_VISIBLE_DEVICES=0 \ +vllm serve robbyant/lingbot-video-moe-30b-a3b \ + --omni \ + --model-class-name LingBotVideoPipeline \ + --default-sampling-params \ + '{"0":{"num_frames":81,"num_inference_steps":40,"guidance_scale":6.0}}' \ + --port 8091 +``` + +These stage defaults match the LingBot reference pipeline. Request-level +values continue to override them, so the smaller smoke request below remains +unchanged. + +When serving MoE, replace the request's `model` form value below with +`robbyant/lingbot-video-moe-30b-a3b`. + +After the server is ready, submit a text-to-video job: + +```bash +create_response=$(curl -s http://localhost:8091/v1/videos \ + -F "model=robbyant/lingbot-video-dense-1.3b" \ + -F "prompt=a robotic arm picks up a red block" \ + -F "width=320" \ + -F "height=192" \ + -F "num_frames=9" \ + -F "fps=24" \ + -F "num_inference_steps=2" \ + -F "guidance_scale=3.0" \ + -F "flow_shift=3.0" \ + -F "seed=42") + +video_id=$(echo "${create_response}" | jq -r '.id') +while true; do + status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status') + if [ "${status}" = "completed" ]; then + break + fi + if [ "${status}" = "failed" ]; then + curl -s "http://localhost:8091/v1/videos/${video_id}" | jq . + exit 1 + fi + sleep 2 +done + +curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o lingbot_t2v.mp4 +``` + +## Key Parameters + +| Parameter | Suggested smoke value | Notes | +|-----------|-----------------------|-------| +| `height` | `192` | Must be a multiple of 16 | +| `width` | `320` | Must be a multiple of 16 | +| `num_frames` | `9` | Must be `1` or `4n + 1`; this PR validates T2V with video outputs | +| `num_inference_steps` | `2` | Use more steps for quality sweeps | +| `guidance_scale` | `3.0` | CFG is active when this is greater than `1.0` | +| `flow_shift` | `3.0` | Scheduler flow-shift; aliases the pipeline's internal `shift` | +| `negative_prompt` | model default | Optional text describing artifacts to avoid | +| `fps` | `24` | Output MP4 frame rate | + +## Validation + +Local dense smoke run: + +- Shape: 9 frames at `192x320` +- Steps: 2 +- Request generation time: `0.2923s` +- Peak reserved GPU memory: `14548 MiB` + +Local dense parity harness against the upstream repository: + +- Shape: `[9, 192, 320, 3]` +- MAE: `0.0065238` +- MSE: `0.00006650` +- PSNR: `41.77 dB` +- Native request time: `0.2875s` + +Local MoE validation on one NVIDIA L20X used checkpoint revision +`f2e538f64afe00cc4ae674db2aeb52e2945edfd5`: + +- Loaded all `977` transformer state keys and `30,084,506,176` parameters. +- Router weights and correction biases stayed in FP32; routed and shared + expert weights loaded in BF16. +- The complete 48-layer transformer forward returned finite + `[1, 16, 1, 8, 8]` BF16 output. +- With both implementations using matched math SDPA backends, the complete + native transformer matched the upstream implementation bitwise + (`max_abs=0`, `mean_abs=0`). +- An isolated sparse MoE block exercising the router, grouped experts, scatter + restore, and shared expert also matched the upstream implementation bitwise. +- The native pipeline generated a 9-frame `192x320` MP4 in 2 steps with + `69326 MiB` (`67.70 GiB`) peak reserved GPU memory. +- The online `/v1/videos` smoke created, polled, and downloaded the generated + MP4 successfully (`1 passed`). + +Do not treat these as production benchmarks; they are functional smoke plus +controlled numerical-parity evidence for small validation inputs. + +## Reproducing MoE numerical parity + +The in-tree parity harness compares the native MoE implementation with a local +checkout of the upstream LingBot-Video implementation. It uses local files +only and does not download checkpoints. + +Set paths for the upstream checkout and the cached MoE checkpoint: + +```bash +export LINGBOT_VIDEO_REPO=/path/to/lingbot-video +export LINGBOT_VIDEO_MOE_MODEL=/path/to/lingbot-video-moe-30b-a3b +``` + +Run the lightweight sparse-block comparison: + +```bash +CUDA_VISIBLE_DEVICES=0 \ +python benchmarks/lingbot_video/moe_transformer_parity.py \ + --scope block \ + --official-repo "${LINGBOT_VIDEO_REPO}" \ + --output-json /tmp/lingbot_moe_block_parity.json +``` + +This path covers bias-corrected router selection, group-limited top-k, routed +experts, FP32 scatter-weighted restore, padding masks, and the shared expert. +It exits with a nonzero status unless the upstream and native block outputs are +bitwise equal. + +The real-checkpoint transformer comparison requires one GPU with enough memory +for the 30B checkpoint. Models are loaded sequentially, so two copies are not +resident on the GPU at the same time: + +```bash +CUDA_VISIBLE_DEVICES=0 \ +python benchmarks/lingbot_video/moe_transformer_parity.py \ + --scope transformer \ + --official-repo "${LINGBOT_VIDEO_REPO}" \ + --model "${LINGBOT_VIDEO_MOE_MODEL}" \ + --output-json /tmp/lingbot_moe_transformer_parity.json +``` + +Transformer parity fixes the official implementation to +`diffusers:_native_math` and the native implementation to +`TORCH_SDPA + SDPBackend.MATH`. The expected result is `exact=true`, +`output.equal=true`, and zero max/mean/RMSE error. Different fused attention +kernels can introduce BF16 rounding differences before MoE routing, so +automatic attention-backend comparisons are diagnostic only and are not the +bitwise correctness oracle. + +## Known Limitations + +- T2V base-transformer inference only. T2I, I2V, TI2V, and the checkpoint's + optional `refiner/` transformer are not supported by this PR. +- Only the BF16 MoE checkpoint is validated in this first integration. +- No HSDP, tensor, sequence, expert, or CFG parallelism is claimed. +- No Cache-DiT, TeaCache, CPU offload, VAE patch parallelism, or quantized + inference is claimed. +- Optional Triton, SGLang, FP8, and alternative fused-expert backends from the + upstream project are not included. +- Only one request per LingBot pipeline batch is currently supported. diff --git a/tests/dfx/perf/tests/test_lingbot_video_vllm_omni.json b/tests/dfx/perf/tests/test_lingbot_video_vllm_omni.json new file mode 100644 index 000000000..a6751bda1 --- /dev/null +++ b/tests/dfx/perf/tests/test_lingbot_video_vllm_omni.json @@ -0,0 +1,34 @@ +[ + { + "test_name": "test_lingbot_video_single_device", + "description": "Single-device LingBot dense T2V baseline at 320x192, 9 frames, 2 steps.", + "server_type": "vllm-omni", + "server_params": { + "model": "robbyant/lingbot-video-dense-1.3b", + "serve_args": { + "model-class-name": "LingBotVideoPipeline", + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "320x192_frames9_steps2", + "dataset": "random", + "task": "t2v", + "backend": "v1/videos", + "width": 320, + "height": 192, + "num-frames": 9, + "fps": 24, + "num-inference-steps": 2, + "num-prompts": 3, + "max-concurrency": 1, + "warmup-requests": 1, + "warmup-concurrency": 1, + "warmup-num-inference-steps": 2, + "enable-negative-prompt": true, + "skip-performance-assertion": true + } + ] + } +] diff --git a/tests/diffusion/models/lingbot_video/__init__.py b/tests/diffusion/models/lingbot_video/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/tests/diffusion/models/lingbot_video/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/diffusion/models/lingbot_video/test_lingbot_video_moe_cuda.py b/tests/diffusion/models/lingbot_video/test_lingbot_video_moe_cuda.py new file mode 100644 index 000000000..5ea19948b --- /dev/null +++ b/tests/diffusion/models/lingbot_video/test_lingbot_video_moe_cuda.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +from tests.helpers.mark import hardware_test + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + + +def _upstream_sparse_moe_reference(block, hidden_states, padding_mask=None): + """Independent eager reference for the upstream LingBot sparse MoE path.""" + batch_size, _, hidden_size = hidden_states.shape + tokens = hidden_states.reshape(-1, hidden_size) + + logits = F.linear(tokens.float(), block.router.weight.float()) + if block.router.score_func == "softmax": + scores = F.softmax(logits, dim=-1) + else: + scores = logits.sigmoid() + scores_for_choice = scores + block.router.e_score_correction_bias.unsqueeze(0) + + if block.router.n_group is not None and block.router.n_group > 1: + experts_per_group = block.router.num_experts // block.router.n_group + grouped = scores_for_choice.view(-1, block.router.n_group, experts_per_group) + group_scores = grouped.topk(2, dim=-1)[0].sum(dim=-1) + group_idx = torch.topk( + group_scores, + k=block.router.topk_group, + dim=-1, + sorted=False, + )[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = group_mask.unsqueeze(-1).expand_as(grouped).reshape_as(scores_for_choice) + scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf")) + + top_indices = torch.topk( + scores_for_choice, + k=block.router.top_k, + dim=-1, + sorted=False, + )[1] + top_scores = scores.gather(1, top_indices) + if block.router.top_k > 1 and block.router.norm_topk_prob: + top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20) + top_scores = top_scores.to(tokens.dtype) * block.router.route_scale + + if padding_mask is not None: + mask = padding_mask.unsqueeze(-1).to(top_scores.dtype) + top_scores = top_scores * mask + top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-9) + top_scores = top_scores * block.router.route_scale + + routed = torch.zeros( + tokens.shape, + dtype=torch.float32, + device=tokens.device, + ) + for expert_idx in range(block.router.num_experts): + token_idx, route_idx = torch.where(top_indices == expert_idx) + if token_idx.numel() == 0: + continue + expert_tokens = tokens[token_idx] + h = F.silu(expert_tokens @ block.experts.w1[expert_idx].transpose(-2, -1)) + h = h * (expert_tokens @ block.experts.w3[expert_idx].transpose(-2, -1)) + expert_output = h @ block.experts.w2[expert_idx].transpose(-2, -1) + weighted = expert_output.float() * top_scores[token_idx, route_idx, None].float() + routed.index_add_(0, token_idx, weighted) + + output = routed.to(tokens.dtype).reshape(batch_size, -1, hidden_size) + if block.shared_experts is not None: + shared = F.linear(hidden_states, block.shared_experts.gate_proj.weight) + shared = F.silu(shared) * F.linear(hidden_states, block.shared_experts.up_proj.weight) + shared = F.linear(shared, block.shared_experts.down_proj.weight) + output = output + shared + return output, top_indices, top_scores + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_grouped_mm_matches_per_expert_reference() -> None: + from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module + + if not hasattr(torch, "_grouped_mm"): + pytest.skip("torch._grouped_mm is unavailable") + + torch.manual_seed(42) + block = module.LingBotVideoSparseMoeBlock( + hidden_size=16, + num_experts=4, + top_k=2, + moe_intermediate_size=8, + score_func="sigmoid", + norm_topk_prob=True, + n_group=None, + topk_group=None, + routed_scaling_factor=1.0, + n_shared_experts=None, + ).to(device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + block.experts.w1.normal_(mean=0.0, std=0.02) + block.experts.w2.normal_(mean=0.0, std=0.02) + block.experts.w3.normal_(mean=0.0, std=0.02) + tokens = torch.randn(23, 16, device="cuda", dtype=torch.bfloat16) + counts = torch.tensor([5, 0, 9, 9], device="cuda", dtype=torch.int64) + + expected = block._run_experts_for_loop(tokens, counts) + actual = block._run_grouped_experts(tokens, counts) + + torch.testing.assert_close(actual, expected) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_sparse_moe_block_matches_upstream_reference() -> None: + from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module + + if not hasattr(torch, "_grouped_mm"): + pytest.skip("torch._grouped_mm is unavailable") + + torch.manual_seed(42) + block = module.LingBotVideoSparseMoeBlock( + hidden_size=16, + num_experts=8, + top_k=2, + moe_intermediate_size=8, + score_func="sigmoid", + norm_topk_prob=True, + n_group=4, + topk_group=1, + routed_scaling_factor=1.5, + n_shared_experts=1, + ) + with torch.no_grad(): + block.router.weight.zero_() + block.router.e_score_correction_bias.copy_(torch.tensor([0.9, 0.8, 1.0, 0.0, 0.7, 0.6, 0.5, 0.4])) + for parameter in ( + block.experts.w1, + block.experts.w2, + block.experts.w3, + block.shared_experts.gate_proj.weight, + block.shared_experts.up_proj.weight, + block.shared_experts.down_proj.weight, + ): + parameter.normal_(mean=0.0, std=0.02) + block = block.to(device="cuda", dtype=torch.bfloat16) + + hidden_states = torch.randn(2, 5, 16, device="cuda", dtype=torch.bfloat16) + padding_mask = torch.tensor( + [1, 1, 1, 1, 1, 1, 1, 1, 0, 0], + device="cuda", + dtype=torch.float32, + ) + + with torch.no_grad(): + expected, expected_indices, expected_scores = _upstream_sparse_moe_reference( + block, + hidden_states, + padding_mask, + ) + actual_indices, actual_scores = block.router(hidden_states.reshape(-1, 16)) + actual_scores = actual_scores * padding_mask.unsqueeze(-1).to(actual_scores.dtype) + actual_scores = actual_scores / (actual_scores.sum(dim=-1, keepdim=True) + 1e-9) + actual_scores = actual_scores * block.router.route_scale + actual = block(hidden_states, padding_mask=padding_mask) + + corrected_scores = ( + torch.full( + (8,), + 0.5, + device="cuda", + ) + + block.router.e_score_correction_bias + ) + unrestricted_indices = torch.topk(corrected_scores, k=2, sorted=False)[1] + + assert set(actual_indices[0].tolist()) == {0, 1} + assert set(unrestricted_indices.tolist()) == {0, 2} + assert torch.equal(actual_indices, expected_indices) + torch.testing.assert_close(actual_scores, expected_scores, rtol=0, atol=0) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert torch.isfinite(actual).all() diff --git a/tests/diffusion/models/lingbot_video/test_lingbot_video_transformer.py b/tests/diffusion/models/lingbot_video/test_lingbot_video_transformer.py new file mode 100644 index 000000000..e9fad3825 --- /dev/null +++ b/tests/diffusion/models/lingbot_video/test_lingbot_video_transformer.py @@ -0,0 +1,266 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _tiny_transformer(**overrides): + from vllm_omni.diffusion.models.lingbot_video import LingBotVideoTransformer3DModel + + config = { + "patch_size": (1, 1, 1), + "in_channels": 2, + "out_channels": 2, + "hidden_size": 16, + "num_attention_heads": 1, + "depth": 0, + "intermediate_size": 32, + "text_dim": 8, + "freq_dim": 8, + "axes_dims": (4, 4, 8), + "axes_lens": (32, 32, 32), + } + config.update(overrides) + return LingBotVideoTransformer3DModel(**config) + + +def test_joint_position_ids_video_then_text_order(): + from vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer import make_joint_position_ids + + positions = make_joint_position_ids(text_len=3, grid_t=1, grid_h=2, grid_w=2, device=torch.device("cpu")) + + assert positions.shape == (7, 3) + assert positions[:4, 0].tolist() == [4, 4, 4, 4] + assert positions[:4, 1:].tolist() == [[0, 0], [0, 1], [1, 0], [1, 1]] + assert positions[4:].tolist() == [[1, 0, 0], [2, 0, 0], [3, 0, 0]] + + +def test_tiny_transformer_depth_zero_forward_shape(): + model = _tiny_transformer() + hidden_states = torch.randn(1, 2, 1, 2, 2) + timestep = torch.tensor([300.0]) + encoder_hidden_states = torch.randn(1, 3, 8) + encoder_attention_mask = torch.ones(1, 3, dtype=torch.long) + + with torch.no_grad(): + out = model( + hidden_states, + timestep, + encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + )[0] + + assert out.shape == hidden_states.shape + assert torch.isfinite(out).all() + + +def test_packed_attention_uses_sdpa_fallback_without_flash_varlen(monkeypatch): + from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module + + monkeypatch.setattr(module, "flash_attn_varlen_func_v3", None) + attn = module.LingBotVideoAttention( + hidden_size=8, + num_heads=2, + norm_eps=1e-6, + qkv_bias=False, + out_bias=False, + ) + captured = {} + + def fake_sdpa_forward(query, key, value, attn_metadata): + captured["mask"] = attn_metadata.attn_mask + return torch.zeros_like(query) + + monkeypatch.setattr(attn.attn.sdpa_fallback, "forward", fake_sdpa_forward) + x = torch.randn(1, 5, 8) + rotary = torch.ones(1, 5, 2, dtype=torch.complex64) + packed_indices = { + "cu_seqlens_kv": torch.tensor([0, 2, 5], dtype=torch.int32), + "max_seqlen_in_batch_kv": 3, + "attention_mask": module._packed_block_attention_mask([2, 3], x.device), + } + + out = attn(x, rotary, packed_indices=packed_indices) + + assert out.shape == x.shape + mask = captured["mask"] + assert mask.shape == (1, 1, 5, 5) + assert mask[0, 0, :2, :2].all() + assert mask[0, 0, 2:, 2:].all() + assert not mask[0, 0, :2, 2:].any() + assert not mask[0, 0, 2:, :2].any() + + +def test_tiny_transformer_rejects_invalid_rope_dims(): + from vllm_omni.diffusion.models.lingbot_video import LingBotVideoTransformer3DModel + + with pytest.raises(AssertionError, match="head_dim"): + LingBotVideoTransformer3DModel( + hidden_size=16, + num_attention_heads=1, + axes_dims=(4, 4, 4), + depth=0, + ) + + +def test_transformer_to_keeps_sensitive_modules_in_fp32(): + model = _tiny_transformer() + + model.to(dtype=torch.bfloat16) + + assert model.patch_embedder.weight.dtype == torch.bfloat16 + assert model.time_embedder.linear_1.weight.dtype == torch.float32 + assert model.norm_out_modulation[1].weight.dtype == torch.float32 + + +def test_router_group_limited_topk_uses_bias_corrected_choice(): + from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module + + router = module.LingBotVideoRouter( + hidden_size=2, + num_experts=4, + top_k=2, + score_func="sigmoid", + norm_topk_prob=True, + n_group=2, + topk_group=1, + route_scale=2.5, + ) + with torch.no_grad(): + router.weight.copy_( + torch.tensor( + [ + [8.0, 0.0], + [7.0, 0.0], + [-8.0, 0.0], + [-7.0, 0.0], + ] + ) + ) + router.e_score_correction_bias.copy_(torch.tensor([0.0, 0.0, 2.0, 2.0])) + + top_indices, top_scores = router(torch.tensor([[1.0, 0.0]])) + + assert set(top_indices[0].tolist()) == {2, 3} + assert torch.allclose(top_scores.sum(dim=-1), torch.tensor([2.5])) + low_score, high_score = sorted(top_scores[0].tolist()) + assert low_score < 0.8 + assert high_score > 1.7 + + +def test_sparse_moe_block_masks_padding_tokens(): + from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module + + block = module.LingBotVideoSparseMoeBlock( + hidden_size=4, + num_experts=2, + top_k=1, + moe_intermediate_size=3, + score_func="sigmoid", + norm_topk_prob=True, + n_group=None, + topk_group=None, + routed_scaling_factor=1.0, + n_shared_experts=None, + ) + with torch.no_grad(): + block.router.weight.copy_( + torch.tensor( + [ + [8.0, 0.0, 0.0, 0.0], + [-8.0, 0.0, 0.0, 0.0], + ] + ) + ) + block.router.e_score_correction_bias.zero_() + block.experts.w1.fill_(0.5) + block.experts.w2.fill_(0.5) + block.experts.w3.fill_(0.5) + + hidden_states = torch.tensor([[[1.0, 0.0, 0.0, 0.0], [-1.0, 0.0, 0.0, 0.0]]]) + padding_mask = torch.tensor([1.0, 0.0]) + + out = block(hidden_states, padding_mask=padding_mask) + + assert out.shape == hidden_states.shape + assert torch.isfinite(out).all() + assert not torch.allclose(out[0, 0], torch.zeros_like(out[0, 0])) + assert torch.allclose(out[0, 1], torch.zeros_like(out[0, 1])) + + +@pytest.mark.parametrize( + "counts_list", + [ + [0, 0, 0, 0], + [1, 0, 7, 8, 9], + [16, 0, 1, 0], + [0, 0, 33, 0], + ], +) +def test_grouped_padding_matches_reference(counts_list): + from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module + + align = 8 + counts = torch.tensor(counts_list, dtype=torch.int64) + num_tokens = int(counts.sum()) + tokens = torch.arange(num_tokens * 4, dtype=torch.float32).reshape(num_tokens, 4) + + actual = module.LingBotVideoSparseMoeBlock._pad_grouped_tokens( + tokens, + counts, + align, + ) + + num_experts = len(counts_list) + max_len = ((num_tokens + num_experts * align + align - 1) // align) * align + aligned_counts = [((max(count, align) + align - 1) // align) * align for count in counts_list] + expected_indices = [num_tokens] * max_len + source_start = 0 + write_start = 0 + for count, aligned_count in zip(counts_list, aligned_counts): + expected_indices[write_start : write_start + count] = range( + source_start, + source_start + count, + ) + source_start += count + write_start += aligned_count + + expected_indices_tensor = torch.tensor(expected_indices, dtype=torch.int64) + tokens_with_pad = torch.vstack((tokens, tokens.new_zeros((tokens.shape[-1],)))) + expected_aligned_counts = torch.tensor(aligned_counts, dtype=torch.int32) + + assert actual[0] == tokens_with_pad.shape + torch.testing.assert_close(actual[1], tokens_with_pad[expected_indices_tensor]) + torch.testing.assert_close(actual[2], expected_indices_tensor) + torch.testing.assert_close(actual[3], expected_aligned_counts) + + +def test_tiny_transformer_constructs_moe_and_dense_layers(): + from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module + + model = _tiny_transformer( + depth=2, + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=8, + decoder_sparse_step=1, + mlp_only_layers=(1,), + n_shared_experts=1, + n_group=2, + topk_group=1, + routed_scaling_factor=2.5, + ) + + assert isinstance(model.blocks[0].ffn, module.LingBotVideoSparseMoeBlock) + assert isinstance(model.blocks[1].ffn, module.LingBotVideoMLP) + assert "blocks.0.ffn.experts.w1" in model.state_dict() + assert "blocks.0.ffn.shared_experts.gate_proj.weight" in model.state_dict() + + model.to(dtype=torch.bfloat16) + + assert model.blocks[0].ffn.router.weight.dtype == torch.float32 + assert model.blocks[0].ffn.experts.w1.dtype == torch.bfloat16 diff --git a/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py b/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py new file mode 100644 index 000000000..1322a542f --- /dev/null +++ b/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _make_pipeline(): + from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import LingBotVideoPipeline + + pipeline = object.__new__(LingBotVideoPipeline) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.default_negative_prompt = "default negative" + pipeline.od_config = SimpleNamespace(flow_shift=None) + return pipeline + + +def _make_request_batch(prompt, **sampling_overrides): + sampling = OmniDiffusionSamplingParams(**sampling_overrides) + request = OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id="lingbot-test") + return DiffusionRequestBatch([request]) + + +def test_lingbot_video_pipeline_import_and_registry(): + from vllm_omni.diffusion.models.lingbot_video import ( + LingBotVideoPipeline, + LingBotVideoTransformer3DModel, + get_lingbot_video_post_process_func, + ) + from vllm_omni.diffusion.registry import _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS + + assert LingBotVideoPipeline is not None + assert LingBotVideoTransformer3DModel is not None + assert get_lingbot_video_post_process_func is not None + assert _DIFFUSION_MODELS["LingBotVideoPipeline"] == ( + "lingbot_video", + "pipeline_lingbot_video", + "LingBotVideoPipeline", + ) + assert _DIFFUSION_POST_PROCESS_FUNCS["LingBotVideoPipeline"] == "get_lingbot_video_post_process_func" + + +def test_component_discovery_declarations(): + from vllm_omni.diffusion.models.lingbot_video import LingBotVideoPipeline + + assert LingBotVideoPipeline._dit_modules == ["transformer"] + assert LingBotVideoPipeline._encoder_modules == ["text_encoder"] + assert LingBotVideoPipeline._vae_modules == ["vae"] + assert LingBotVideoPipeline.supports_step_execution is False + + +def test_extra_body_params_include_video_flow_shift_alias(): + from vllm_omni.model_extras import get_extra_body_params + + params = get_extra_body_params("LingBotVideoPipeline") + assert "shift" in params + assert "flow_shift" in params + assert "batch_cfg" in params + + +def test_extract_prompt_string_and_mapping(): + from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import _extract_prompt + + string_req = SimpleNamespace(prompt="a robot waves") + assert _extract_prompt(string_req) == ("a robot waves", None) + + mapping_req = SimpleNamespace(prompt={"prompt": "a robot walks", "negative_prompt": "low quality"}) + assert _extract_prompt(mapping_req) == ("a robot walks", "low quality") + + with pytest.raises(TypeError, match="Unsupported LingBot prompt type"): + _extract_prompt(SimpleNamespace(prompt=["not", "supported"])) + + +def test_check_inputs_accepts_t2v_shapes_and_rejects_invalid_values(): + from vllm_omni.diffusion.models.lingbot_video import LingBotVideoPipeline + + LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=1) + LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=9) + + with pytest.raises(ValueError, match="num_frames"): + LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=8) + with pytest.raises(ValueError, match="height"): + LingBotVideoPipeline.check_inputs(height=190, width=320, num_frames=9) + + +def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested(): + from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func + + frames = torch.ones(2, 4, 4, 3) + post = get_lingbot_video_post_process_func(SimpleNamespace()) + + assert post(frames, SimpleNamespace(output_type="pt")) is frames + array = post(frames, SimpleNamespace(output_type="np")) + assert isinstance(array, np.ndarray) + assert array.shape == (2, 4, 4, 3) + + +def test_forward_resolves_t2v_sampling_and_flow_shift_alias(): + pipeline = _make_pipeline() + calls = [] + + def fake_generate(**kwargs): + calls.append(kwargs) + return torch.zeros(9, 192, 320, 3) + + pipeline._generate = fake_generate + req = _make_request_batch( + {"prompt": "a robot arm", "negative_prompt": "bad hands"}, + height=192, + width=320, + num_frames=9, + num_inference_steps=2, + guidance_scale=3.0, + seed=42, + output_type="pt", + extra_args={"flow_shift": 4.0, "batch_cfg": True}, + ) + + out = pipeline.forward(req) + + assert torch.equal(out.output, torch.zeros(9, 192, 320, 3)) + assert len(calls) == 1 + call = calls[0] + assert call["prompt"] == "a robot arm" + assert call["negative_prompt"] == "bad hands" + assert call["height"] == 192 + assert call["width"] == 320 + assert call["num_frames"] == 9 + assert call["num_inference_steps"] == 2 + assert call["guidance_scale"] == 3.0 + assert call["shift"] == 4.0 + assert call["batch_cfg"] is True + assert call["output_type"] == "pt" + + +def test_forward_uses_lingbot_guidance_default_when_omitted(): + pipeline = _make_pipeline() + calls = [] + + def fake_generate(**kwargs): + calls.append(kwargs) + return torch.zeros(1) + + pipeline._generate = fake_generate + req = _make_request_batch( + "a robot arm", + height=192, + width=320, + num_frames=81, + num_inference_steps=2, + seed=42, + ) + + pipeline.forward(req) + + assert calls[0]["num_frames"] == 81 + assert calls[0]["guidance_scale"] == 6.0 + + +def test_forward_prefers_shift_over_flow_shift_and_defaults_negative_prompt(): + pipeline = _make_pipeline() + calls = [] + + def fake_generate(**kwargs): + calls.append(kwargs) + return torch.zeros(5, 192, 192, 3) + + pipeline._generate = fake_generate + req = _make_request_batch( + "a robot arm", + height=192, + width=192, + num_frames=5, + num_inference_steps=2, + guidance_scale=1.0, + seed=42, + extra_args={"shift": 5.0, "flow_shift": 4.0}, + ) + + pipeline.forward(req) + + assert calls[0]["shift"] == 5.0 + assert calls[0]["negative_prompt"] == "default negative" + + +def test_forward_rejects_multi_request_batches(): + pipeline = _make_pipeline() + first = OmniDiffusionRequest( + prompt="first", + sampling_params=OmniDiffusionSamplingParams(seed=1), + request_id="lingbot-first", + ) + second = OmniDiffusionRequest( + prompt="second", + sampling_params=OmniDiffusionSamplingParams(seed=2), + request_id="lingbot-second", + ) + + with pytest.raises(ValueError, match="only supports one request"): + pipeline.forward(DiffusionRequestBatch([first, second])) + + +def test_load_weights_rejects_external_weight_stream(): + pipeline = _make_pipeline() + + assert pipeline.load_weights([]) == set() + with pytest.raises(RuntimeError, match="components are loaded directly"): + pipeline.load_weights([("transformer.weight", torch.zeros(1))]) diff --git a/tests/e2e/online_serving/test_lingbot_video.py b/tests/e2e/online_serving/test_lingbot_video.py new file mode 100644 index 000000000..e7af994cb --- /dev/null +++ b/tests/e2e/online_serving/test_lingbot_video.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Online serving smoke for ``robbyant/lingbot-video-dense-1.3b``. + +The LingBot dense PR supports text-to-video only, so this module exercises the +native ``/v1/videos`` path with a small 9-frame request. +""" + +import os + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +MODEL = "robbyant/lingbot-video-dense-1.3b" +PROMPT = "a robotic arm picks up a red block" +NEGATIVE_PROMPT = "low quality, blurry, watermark, text" +DEFAULT_SAMPLING_PARAMS = '{"0":{"num_frames":81,"num_inference_steps":40,"guidance_scale":6.0}}' + +SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) + + +def _get_diffusion_feature_cases(model: str): + return [ + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--model-class-name", + "LingBotVideoPipeline", + "--default-sampling-params", + DEFAULT_SAMPLING_PARAMS, + ], + ), + id="default", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + ] + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) +def test_text_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + request_config = { + "model": omni_server.model, + "form_data": { + "model": omni_server.model, + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "height": 192, + "width": 320, + "num_frames": 9, + "fps": 24, + "num_inference_steps": 2, + "guidance_scale": 3.0, + "flow_shift": 3.0, + "seed": 42, + }, + } + openai_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_lingbot_video_expansion.py b/tests/e2e/online_serving/test_lingbot_video_expansion.py new file mode 100644 index 000000000..cb5a314ff --- /dev/null +++ b/tests/e2e/online_serving/test_lingbot_video_expansion.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +L4 expansion coverage for ``robbyant/lingbot-video-dense-1.3b``. + +This file remains dense-only. Basic single-GPU MoE serving is covered by +``test_lingbot_video_moe.py``; T2I, I2V, TI2V, and multi-GPU feature rows +belong in follow-up PRs. +""" + +import json +import os + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + +MODEL = "robbyant/lingbot-video-dense-1.3b" +PROMPT = "a robotic arm picks up a red block" +NEGATIVE_PROMPT = "low quality, blurry, watermark, text" + +SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) + + +def _get_diffusion_feature_cases(model: str): + return [ + pytest.param( + OmniServerParams( + model=model, + server_args=["--model-class-name", "LingBotVideoPipeline"], + ), + id="default", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + ] + + +@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) +def test_cfg_off(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + request_config = { + "model": omni_server.model, + "form_data": { + "model": omni_server.model, + "prompt": PROMPT, + "height": 192, + "width": 192, + "num_frames": 9, + "fps": 24, + "num_inference_steps": 2, + "guidance_scale": 1.0, + "flow_shift": 3.0, + "seed": 42, + }, + } + openai_client.send_video_diffusion_request(request_config) + + +@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) +def test_batch_cfg_extra_params(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + request_config = { + "model": omni_server.model, + "form_data": { + "model": omni_server.model, + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "height": 192, + "width": 320, + "num_frames": 9, + "fps": 24, + "num_inference_steps": 2, + "guidance_scale": 3.0, + "flow_shift": 3.0, + "seed": 42, + "extra_params": json.dumps({"batch_cfg": True}, separators=(",", ":")), + }, + } + openai_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_lingbot_video_moe.py b/tests/e2e/online_serving/test_lingbot_video_moe.py new file mode 100644 index 000000000..32a6cbf2f --- /dev/null +++ b/tests/e2e/online_serving/test_lingbot_video_moe.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Single-GPU serving smoke for ``robbyant/lingbot-video-moe-30b-a3b``.""" + +import os + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +MODEL = "robbyant/lingbot-video-moe-30b-a3b" +PROMPT = "a robotic arm picks up a red block" +NEGATIVE_PROMPT = "low quality, blurry, watermark, text" +SMOKE_DEFAULT_SAMPLING_PARAMS = '{"0":{"num_frames":81,"num_inference_steps":2,"guidance_scale":3.0}}' + +SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) + + +def _get_server_cases(model: str): + return [ + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--model-class-name", + "LingBotVideoPipeline", + "--default-sampling-params", + SMOKE_DEFAULT_SAMPLING_PARAMS, + ], + ), + id="default", + marks=SINGLE_CARD_MARKS, + ), + ] + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parametrize("omni_server", _get_server_cases(MODEL), indirect=True) +def test_text_to_video_moe(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + request_config = { + "model": omni_server.model, + "form_data": { + "model": omni_server.model, + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "height": 192, + "width": 320, + "num_frames": 9, + "fps": 24, + "flow_shift": 3.0, + "seed": 42, + }, + } + openai_client.send_video_diffusion_request(request_config) diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index c95644143..4b51220c7 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -67,6 +67,7 @@ EXCLUDED_MODELS = [ "AudioXPipeline", "HunyuanVideo15Pipeline", "HunyuanVideo15ImageToVideoPipeline", + "LingBotVideoPipeline", "MagiHumanPipeline", "OmniVoicePipeline", "OmniVoice", diff --git a/vllm_omni/diffusion/models/lingbot_video/__init__.py b/vllm_omni/diffusion/models/lingbot_video/__init__.py new file mode 100644 index 000000000..013f7b2c3 --- /dev/null +++ b/vllm_omni/diffusion/models/lingbot_video/__init__.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer import ( + LingBotVideoTransformer3DModel, +) +from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import ( + LingBotVideoPipeline, + get_lingbot_video_post_process_func, +) + +__all__ = [ + "LingBotVideoPipeline", + "LingBotVideoTransformer3DModel", + "get_lingbot_video_post_process_func", +] diff --git a/vllm_omni/diffusion/models/lingbot_video/lingbot_video_transformer.py b/vllm_omni/diffusion/models/lingbot_video/lingbot_video_transformer.py new file mode 100644 index 000000000..ad820eba3 --- /dev/null +++ b/vllm_omni/diffusion/models/lingbot_video/lingbot_video_transformer.py @@ -0,0 +1,1005 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from LingBot-Video (https://github.com/Robbyant/lingbot-video). + +import math + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.embeddings import TimestepEmbedding, Timesteps +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.modeling_utils import ModelMixin + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.utils.fa import ( + flash_attn_varlen_func as flash_attn_varlen_func_v3, +) +from vllm_omni.diffusion.attention.layer import Attention + +LINGBOT_VIDEO_FP32_MODULES = ( + "time_embedder", + "time_modulation", + "scale_shift_table", + "norm", + "norm1", + "norm2", + "norm_q", + "norm_k", + "norm_post_attn", + "norm_post_ffn", + "norm_out", + "norm_out_modulation", + "router", +) + + +def should_keep_in_fp32(name: str) -> bool: + return any(module_name in name.split(".") for module_name in LINGBOT_VIDEO_FP32_MODULES) + + +def _all_to_all_split_cat( + local_input: torch.Tensor, + scatter_dim: int, + gather_dim: int, + group: dist.ProcessGroup, +) -> torch.Tensor: + world_size = dist.get_world_size(group) + input_list = [tensor.contiguous() for tensor in torch.tensor_split(local_input, world_size, scatter_dim)] + output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)] + dist.all_to_all(output_list, input_list, group=group) + return torch.cat(output_list, dim=gather_dim).contiguous() + + +class LingBotVideoRMSNorm(nn.Module): + """RMSNorm with fp32 accumulation.""" + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return (self.weight * hidden_states).to(input_dtype) + + +def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: + """Apply complex RoPE to `(B, S, H, D)` attention tensors.""" + with torch.amp.autocast("cuda", enabled=False): + x_c = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + out = torch.view_as_real(x_c * freqs_cis.unsqueeze(2)).flatten(3) + return out.type_as(x) + + +class LingBotVideoRotaryEmbedding(nn.Module): + """Complex64 RoPE table indexed by position ids.""" + + def __init__(self, axes_dims: tuple[int, ...], axes_lens: tuple[int, ...], theta: float): + super().__init__() + self.axes_dims = tuple(axes_dims) + self.axes_lens = list(axes_lens) + self.theta = theta + self.freqs_cis = None + + @staticmethod + def precompute_freqs_cis(dim: tuple[int, ...], end: tuple[int, ...], theta: float): + freqs_cis = [] + for d, e in zip(dim, end): + freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d)) + timestep = torch.arange(e, device=freqs.device, dtype=torch.float64) + freqs = torch.outer(timestep, freqs).float() + freqs_cis.append(torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64)) + return freqs_cis + + def forward(self, position_ids: torch.Tensor) -> torch.Tensor: + # position_ids: (S, 3) int -> (S, head_dim/2) complex64 + device = position_ids.device + max_vals = position_ids.max(dim=0).values.tolist() + needs_rebuild = self.freqs_cis is None or any( + max_val >= axis_len for max_val, axis_len in zip(max_vals, self.axes_lens) + ) + if needs_rebuild: + for i in range(len(self.axes_lens)): + if max_vals[i] >= self.axes_lens[i]: + self.axes_lens[i] = int(max_vals[i] * 1.5) + 1 + self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, tuple(self.axes_lens), theta=self.theta) + self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] + elif self.freqs_cis[0].device != device: + self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis] + + return torch.cat([self.freqs_cis[i][position_ids[:, i]] for i in range(len(self.axes_dims))], dim=-1) + + +def make_joint_position_ids(text_len: int, grid_t: int, grid_h: int, grid_w: int, device: torch.device) -> torch.Tensor: + """Return ``(t, h, w)`` positions with video rows before text rows. + + Text t-axis positions are 1..text_len, and video t-axis positions start at + text_len + 1. This matches the token order produced by ``_cat_interleave``. + """ + tt = torch.arange(grid_t, device=device, dtype=torch.int32) + (text_len + 1) + hh = torch.arange(grid_h, device=device, dtype=torch.int32) + ww = torch.arange(grid_w, device=device, dtype=torch.int32) + grid = torch.stack(torch.meshgrid(tt, hh, ww, indexing="ij"), dim=-1).flatten(0, 2) + text_t = torch.arange(text_len, device=device, dtype=torch.int32) + 1 + text_pos = torch.stack([text_t, torch.zeros_like(text_t), torch.zeros_like(text_t)], dim=-1) + return torch.cat([grid, text_pos], dim=0) # (Nx + L, 3) + + +def _cat_interleave( + a: torch.Tensor, + len_a: list[int], + b: torch.Tensor, + len_b: list[int], +) -> torch.Tensor: + a_split = torch.split(a, len_a, dim=1) + b_split = torch.split(b, len_b, dim=1) + blocks: list[torch.Tensor] = [] + for x_part, text_part in zip(a_split, b_split): + blocks.extend([x_part, text_part]) + return torch.cat(blocks, dim=1) + + +def _packed_block_attention_mask(sample_seq_lens: list[int], device: torch.device) -> torch.Tensor: + total_seq_len = sum(sample_seq_lens) + mask = torch.zeros( + 1, + 1, + total_seq_len, + total_seq_len, + dtype=torch.bool, + device=device, + ) + start = 0 + for seq_len in sample_seq_lens: + end = start + seq_len + mask[:, :, start:end, start:end] = True + start = end + return mask + + +class LingBotVideoTextEmbedder(nn.Module): + """Matches CondProjection: RMSNorm(text_dim, eps=1e-6 fixed) -> Linear-SiLU-Linear.""" + + def __init__(self, text_dim: int, hidden_size: int): + super().__init__() + self.norm = LingBotVideoRMSNorm(text_dim, eps=1e-6) + self.linear_1 = nn.Linear(text_dim, hidden_size, bias=True) + self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.norm(x) + return self.linear_2(F.silu(self.linear_1(x))) + + +class LingBotVideoAttention(nn.Module): + def __init__(self, hidden_size, num_heads, norm_eps, qkv_bias, out_bias): + super().__init__() + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.to_q = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + self.to_k = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + self.to_v = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + self.norm_q = LingBotVideoRMSNorm(self.head_dim, norm_eps) + self.norm_k = LingBotVideoRMSNorm(self.head_dim, norm_eps) + self.to_out = nn.Linear(hidden_size, hidden_size, bias=out_bias) + self.attn = Attention( + num_heads=num_heads, + head_size=self.head_dim, + softmax_scale=1.0 / math.sqrt(self.head_dim), + causal=False, + num_kv_heads=num_heads, + ) + + def forward( + self, + x, + rotary_emb, + attention_mask=None, + packed_indices: dict[str, torch.Tensor] | None = None, + parallel_config=None, + ): + B, S, _ = x.shape + q = self.to_q(x).unflatten(2, (self.num_heads, self.head_dim)) + k = self.to_k(x).unflatten(2, (self.num_heads, self.head_dim)) + v = self.to_v(x).unflatten(2, (self.num_heads, self.head_dim)) + q = apply_rotary_emb(self.norm_q(q), rotary_emb) + k = apply_rotary_emb(self.norm_k(k), rotary_emb) + if packed_indices is None: + out = self.attn(q, k, v, attn_metadata=AttentionMetadata(attn_mask=attention_mask)) + else: + packed_attention_mask = packed_indices.get("attention_mask") + if packed_attention_mask is not None and parallel_config is None: + out = self.attn.sdpa_fallback.forward( + q, + k, + v, + AttentionMetadata(attn_mask=packed_attention_mask), + ) + return self.to_out(out.flatten(2, 3).type_as(x)) + if flash_attn_varlen_func_v3 is None: + raise RuntimeError( + "A flash attention varlen function is required for packed context parallel attention." + ) + if parallel_config is None: + result = flash_attn_varlen_func_v3( + q=q.reshape(-1, self.num_heads, self.head_dim), + k=k.reshape(-1, self.num_heads, self.head_dim), + v=v.reshape(-1, self.num_heads, self.head_dim), + cu_seqlens_q=packed_indices["cu_seqlens_kv"], + cu_seqlens_k=packed_indices["cu_seqlens_kv"], + max_seqlen_q=packed_indices["max_seqlen_in_batch_kv"], + max_seqlen_k=packed_indices["max_seqlen_in_batch_kv"], + causal=False, + ) + out = result[0] if isinstance(result, tuple) else result + out = out.reshape(B, S, self.num_heads, self.head_dim) + else: + group = parallel_config.context_parallel_config._ulysses_mesh.get_group() + world_size = dist.get_world_size(group) + local_heads = self.num_heads // world_size + q_global = _all_to_all_split_cat( + q.reshape(B, S, self.num_heads * self.head_dim), + scatter_dim=2, + gather_dim=1, + group=group, + ).view(B, S * world_size, local_heads, self.head_dim) + k_global = _all_to_all_split_cat( + k.reshape(B, S, self.num_heads * self.head_dim), + scatter_dim=2, + gather_dim=1, + group=group, + ).view(B, S * world_size, local_heads, self.head_dim) + v_global = _all_to_all_split_cat( + v.reshape(B, S, self.num_heads * self.head_dim), + scatter_dim=2, + gather_dim=1, + group=group, + ).view(B, S * world_size, local_heads, self.head_dim) + q_flat = q_global.reshape(-1, local_heads, self.head_dim) + k_flat = k_global.reshape(-1, local_heads, self.head_dim) + v_flat = v_global.reshape(-1, local_heads, self.head_dim) + result = flash_attn_varlen_func_v3( + q=q_flat, + k=k_flat, + v=v_flat, + cu_seqlens_q=packed_indices["cu_seqlens_kv"], + cu_seqlens_k=packed_indices["cu_seqlens_kv"], + max_seqlen_q=packed_indices["max_seqlen_in_batch_kv"], + max_seqlen_k=packed_indices["max_seqlen_in_batch_kv"], + causal=False, + ) + out_global = result[0] if isinstance(result, tuple) else result + out_global = out_global.reshape(B, S * world_size, local_heads * self.head_dim) + out = _all_to_all_split_cat( + out_global, + scatter_dim=1, + gather_dim=2, + group=group, + ).view(B, S, self.num_heads, self.head_dim) + return self.to_out(out.flatten(2, 3).type_as(x)) + + +class LingBotVideoMLP(nn.Module): + def __init__(self, hidden_size, intermediate_size): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + + def forward(self, x): + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class LingBotVideoRouter(nn.Module): + """Token-choice top-k router used by the LingBot MoE checkpoints. + + Selection uses the bias-corrected score, while the gating weights use the + original score. This asymmetry matches the reference LingBot inference path. + """ + + def __init__( + self, + hidden_size: int, + num_experts: int, + top_k: int, + score_func: str, + norm_topk_prob: bool, + n_group: int | None, + topk_group: int | None, + route_scale: float, + ): + super().__init__() + self.num_experts = num_experts + self.top_k = top_k + self.score_func = score_func + self.norm_topk_prob = norm_topk_prob + self.n_group = n_group + self.topk_group = topk_group + self.route_scale = route_scale + self.weight = nn.Parameter(torch.empty(num_experts, hidden_size)) + self.register_buffer( + "e_score_correction_bias", + torch.zeros(num_experts), + persistent=True, + ) + + def _group_limited_topk(self, scores_for_choice: torch.Tensor) -> torch.Tensor: + if self.n_group is None or self.topk_group is None: + raise ValueError("group-limited top-k requires n_group and topk_group.") + seq_len = scores_for_choice.shape[0] + experts_per_group = self.num_experts // self.n_group + grouped = scores_for_choice.view(seq_len, self.n_group, experts_per_group) + group_scores = grouped.topk(2, dim=-1)[0].sum(dim=-1) + group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = group_mask.unsqueeze(-1).expand(seq_len, self.n_group, experts_per_group).reshape(seq_len, -1) + masked = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf")) + return torch.topk(masked, k=self.top_k, dim=-1, sorted=False)[1] + + def forward(self, tokens: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + with torch.amp.autocast(tokens.device.type, enabled=False): + logits = F.linear(tokens.float(), self.weight.float()) + if self.score_func == "softmax": + scores = F.softmax(logits, dim=-1) + elif self.score_func == "sigmoid": + scores = logits.sigmoid() + else: + raise ValueError(f"Unsupported LingBot router score_func: {self.score_func!r}.") + + scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0) + if self.n_group is not None and self.n_group > 1: + top_indices = self._group_limited_topk(scores_for_choice) + else: + top_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1] + top_scores = scores.gather(1, top_indices) + if self.top_k > 1 and self.norm_topk_prob: + top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20) + top_scores = top_scores * self.route_scale + return top_indices, top_scores.to(tokens.dtype) + + +class LingBotVideoGroupedExperts(nn.Module): + """Grouped expert weights. + + Weight layout matches the reference checkpoint: + w1 [E, I, H], w2 [E, H, I], w3 [E, I, H]. + """ + + def __init__(self, num_experts: int, hidden_size: int, intermediate_size: int): + super().__init__() + self.num_experts = num_experts + self.w1 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size)) + self.w2 = nn.Parameter(torch.empty(num_experts, hidden_size, intermediate_size)) + self.w3 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size)) + + +def _round_up_to_multiple(value: int, multiple: int) -> int: + return ((value + multiple - 1) // multiple) * multiple + + +class LingBotVideoSparseMoeBlock(nn.Module): + def __init__( + self, + hidden_size: int, + num_experts: int, + top_k: int, + moe_intermediate_size: int, + score_func: str, + norm_topk_prob: bool, + n_group: int | None, + topk_group: int | None, + routed_scaling_factor: float, + n_shared_experts: int | None, + ): + super().__init__() + self.hidden_size = hidden_size + self.num_experts = num_experts + self.router = LingBotVideoRouter( + hidden_size, + num_experts, + top_k, + score_func, + norm_topk_prob, + n_group, + topk_group, + routed_scaling_factor, + ) + self.experts = LingBotVideoGroupedExperts(num_experts, hidden_size, moe_intermediate_size) + self.shared_experts = None + if n_shared_experts is not None and n_shared_experts > 0: + self.shared_experts = LingBotVideoMLP( + hidden_size, + moe_intermediate_size * n_shared_experts, + ) + + @staticmethod + def _reorder_tokens( + tokens: torch.Tensor, + top_scores: torch.Tensor, + top_indices: torch.Tensor, + num_experts: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]: + num_tokens = tokens.shape[0] + top_k = top_indices.shape[1] + flat_scores = top_scores.reshape(-1) + flat_indices = top_indices.reshape(-1) + active_positions = torch.where(flat_scores != 0)[0] + active_experts = flat_indices[active_positions] + + counts = torch.zeros(num_experts, device=tokens.device, dtype=torch.int64) + counts.scatter_add_( + 0, + active_experts, + torch.ones_like(active_experts, dtype=torch.int64), + ) + + sort_order = torch.argsort(active_experts, stable=True) + sorted_positions = active_positions[sort_order] + sorted_scores = flat_scores[sorted_positions] + original_token_idx = sorted_positions // top_k + permuted_tokens = tokens[original_token_idx] + return permuted_tokens, counts, sorted_positions, sorted_scores, num_tokens, top_k + + @staticmethod + def _pad_grouped_tokens( + tokens: torch.Tensor, + counts: torch.Tensor, + align: int = 8, + ) -> tuple[torch.Size, torch.Tensor, torch.Tensor, torch.Tensor]: + num_tokens = tokens.shape[0] + num_experts = int(counts.shape[0]) + max_len = _round_up_to_multiple(num_tokens + num_experts * align, align) + counts_i64 = counts.to(torch.int64) + total_per_expert = torch.clamp_min(counts_i64, align) + aligned_counts_i64 = (total_per_expert + align - 1) // align * align + write_offsets = torch.cumsum(aligned_counts_i64, dim=0) - aligned_counts_i64 + end_offsets = torch.cumsum(aligned_counts_i64, dim=0) + start_indices = torch.cumsum(counts_i64, dim=0) - counts_i64 + + slots = torch.arange(max_len, dtype=torch.int64, device=tokens.device) + expert_idx = torch.bucketize(slots, end_offsets, right=True) + valid_expert = expert_idx < num_experts + safe_expert_idx = expert_idx.clamp(max=num_experts - 1) + local_idx = slots - write_offsets[safe_expert_idx] + source_idx = start_indices[safe_expert_idx] + local_idx + valid = valid_expert & (local_idx < counts_i64[safe_expert_idx]) + fill = torch.full_like(source_idx, num_tokens) + permuted_indices = torch.where(valid, source_idx, fill) + + tokens_with_pad = torch.vstack((tokens, tokens.new_zeros((tokens.shape[-1],)))) + input_shape = tokens_with_pad.shape + return ( + input_shape, + tokens_with_pad[permuted_indices], + permuted_indices, + aligned_counts_i64.to(torch.int32), + ) + + @staticmethod + def _unpad_grouped_tokens( + output: torch.Tensor, + input_shape: torch.Size, + permuted_indices: torch.Tensor, + ) -> torch.Tensor: + unpermuted = output.new_empty(input_shape) + unpermuted[permuted_indices, :] = output + return unpermuted[:-1] + + def _run_experts_for_loop(self, tokens: torch.Tensor, counts: torch.Tensor) -> torch.Tensor: + count_list = counts.tolist() + splits = torch.split(tokens, count_list, dim=0) + outputs = [] + for expert_idx, expert_tokens in enumerate(splits): + if expert_tokens.numel() == 0: + continue + h = F.silu(expert_tokens @ self.experts.w1[expert_idx].transpose(-2, -1)) + h = h * (expert_tokens @ self.experts.w3[expert_idx].transpose(-2, -1)) + h = h @ self.experts.w2[expert_idx].transpose(-2, -1) + outputs.append(h) + if not outputs: + return tokens.new_zeros(tokens.shape) + return torch.cat(outputs, dim=0) + + def _run_grouped_experts(self, tokens: torch.Tensor, counts: torch.Tensor) -> torch.Tensor: + if not hasattr(torch, "_grouped_mm") or tokens.device.type != "cuda": + return self._run_experts_for_loop(tokens, counts) + input_shape, padded_tokens, permuted_indices, aligned_counts = self._pad_grouped_tokens(tokens, counts) + offsets = torch.cumsum(aligned_counts, dim=0, dtype=torch.int32) + h = F.silu( + torch._grouped_mm( + padded_tokens.bfloat16(), + self.experts.w1.bfloat16().transpose(-2, -1), + offs=offsets, + ) + ) + h = h * torch._grouped_mm( + padded_tokens.bfloat16(), + self.experts.w3.bfloat16().transpose(-2, -1), + offs=offsets, + ) + out = torch._grouped_mm( + h, + self.experts.w2.bfloat16().transpose(-2, -1), + offs=offsets, + ).type_as(padded_tokens) + return self._unpad_grouped_tokens(out, input_shape, permuted_indices) + + @staticmethod + def _restore_tokens( + expert_output: torch.Tensor, + sorted_positions: torch.Tensor, + sorted_scores: torch.Tensor, + num_tokens: int, + top_k: int, + ) -> torch.Tensor: + hidden_size = expert_output.shape[-1] + unsorted = torch.zeros( + (num_tokens * top_k, hidden_size), + dtype=expert_output.dtype, + device=expert_output.device, + ) + unsorted[sorted_positions] = expert_output + unsorted = unsorted.reshape(num_tokens, top_k, hidden_size) + + scores_unsorted = torch.zeros( + num_tokens * top_k, + dtype=sorted_scores.dtype, + device=sorted_scores.device, + ) + scores_unsorted[sorted_positions] = sorted_scores + scores_unsorted = scores_unsorted.reshape(num_tokens, top_k, 1) + return (unsorted.float() * scores_unsorted).sum(dim=1).to(expert_output.dtype) + + def _run_selected_experts( + self, + tokens: torch.Tensor, + top_scores: torch.Tensor, + top_indices: torch.Tensor, + ) -> torch.Tensor: + ( + permuted_tokens, + counts, + sorted_positions, + sorted_scores, + num_tokens, + top_k, + ) = self._reorder_tokens(tokens, top_scores, top_indices, self.router.num_experts) + expert_output = self._run_grouped_experts(permuted_tokens, counts) + return self._restore_tokens( + expert_output, + sorted_positions, + sorted_scores, + num_tokens, + top_k, + ) + + def forward( + self, + hidden_states: torch.Tensor, + padding_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + batch_size = hidden_states.shape[0] + tokens = hidden_states.reshape(-1, self.hidden_size) + top_indices, top_scores = self.router(tokens) + if padding_mask is not None: + mask = padding_mask.unsqueeze(-1).to(top_scores.dtype) + top_scores = top_scores * mask + top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-9) + top_scores = top_scores * self.router.route_scale + + out = self._run_selected_experts(tokens, top_scores, top_indices) + out = out.reshape(batch_size, -1, self.hidden_size) + if self.shared_experts is not None: + out = out + self.shared_experts(hidden_states) + return out + + +class LingBotVideoBlock(nn.Module): + def __init__( + self, + hidden_size, + num_attention_heads, + intermediate_size, + norm_eps, + qkv_bias, + out_bias, + num_experts, + num_experts_per_tok, + moe_intermediate_size, + decoder_sparse_step, + mlp_only_layers, + n_shared_experts, + score_func, + norm_topk_prob, + n_group, + topk_group, + routed_scaling_factor, + layer_idx: int, + ): + super().__init__() + self.layer_idx = layer_idx + h = hidden_size + self.scale_shift_table = nn.Parameter(torch.zeros(1, 6 * h)) + self.norm1 = LingBotVideoRMSNorm(h, norm_eps) + self.attn = LingBotVideoAttention(h, num_attention_heads, norm_eps, qkv_bias, out_bias) + self.norm_post_attn = LingBotVideoRMSNorm(h, norm_eps) + self.norm2 = LingBotVideoRMSNorm(h, norm_eps) + use_sparse_moe = ( + num_experts > 0 and layer_idx not in mlp_only_layers and (layer_idx + 1) % decoder_sparse_step == 0 + ) + if use_sparse_moe: + self.ffn = LingBotVideoSparseMoeBlock( + hidden_size=h, + num_experts=num_experts, + top_k=num_experts_per_tok, + moe_intermediate_size=moe_intermediate_size, + score_func=score_func, + norm_topk_prob=norm_topk_prob, + n_group=n_group, + topk_group=topk_group, + routed_scaling_factor=routed_scaling_factor, + n_shared_experts=n_shared_experts, + ) + else: + self.ffn = LingBotVideoMLP(h, intermediate_size) + self.norm_post_ffn = LingBotVideoRMSNorm(h, norm_eps) + + def forward( + self, + x, + temb6, + rotary_emb, + attention_mask=None, + moe_padding_mask=None, + packed_indices: dict[str, torch.Tensor] | None = None, + parallel_config=None, + ): + expected_tokens = x.shape[0] * x.shape[1] + if temb6.ndim != 2 or temb6.shape[0] != expected_tokens: + raise ValueError( + "LingBotVideoBlock expects token-level temb6 with shape " + f"(B*S, 6D); got {tuple(temb6.shape)} for hidden states {tuple(x.shape)}." + ) + # AdaLN modulation and normalization stay in fp32 for the sensitive path. + mod = temb6.view(x.shape[0], x.shape[1], -1) + self.scale_shift_table.unsqueeze(0) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=-1) + gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh() + scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp + + # AdaLN modulation and norms stay in fp32; cast to the transformer + # compute dtype only at Linear boundaries. + bulk_dtype = self.attn.to_q.weight.dtype + attn_in = (self.norm1(x) * scale_msa + shift_msa).to(bulk_dtype) + attn_out = self.attn( + attn_in, + rotary_emb, + attention_mask, + packed_indices=packed_indices, + parallel_config=parallel_config, + ) + x = x + (gate_msa * self.norm_post_attn(attn_out)).to(x.dtype) + + ffn_in = (self.norm2(x) * scale_mlp + shift_mlp).to(bulk_dtype) + if isinstance(self.ffn, LingBotVideoSparseMoeBlock): + ffn_out = self.ffn(ffn_in, padding_mask=moe_padding_mask) + else: + ffn_out = self.ffn(ffn_in) + ffn_normed = self.norm_post_ffn(ffn_out) + x = x + (gate_mlp * ffn_normed).to(x.dtype) + return x + + +class LingBotVideoTransformer3DModel(ModelMixin, ConfigMixin): + _supports_gradient_checkpointing = False + _repeated_blocks = ["LingBotVideoBlock"] + _layerwise_offload_blocks_attr = "blocks" + _no_split_modules = ["LingBotVideoBlock"] + _keep_in_fp32_modules = list(LINGBOT_VIDEO_FP32_MODULES) + + def to(self, *args, **kwargs): + device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs) + if dtype is None or dtype == torch.float32: + return super().to(*args, **kwargs) + + dtype_is_floating = torch.is_floating_point(torch.empty((), dtype=dtype)) + if not dtype_is_floating: + return super().to(*args, **kwargs) + + if device is not None: + super().to(device=device, non_blocking=non_blocking) + + for name, param in self.named_parameters(): + if not torch.is_floating_point(param): + continue + target_dtype = torch.float32 if should_keep_in_fp32(name) else dtype + param.data = param.data.to(dtype=target_dtype, non_blocking=non_blocking) + if param.grad is not None: + param.grad.data = param.grad.data.to(dtype=target_dtype, non_blocking=non_blocking) + + for name, buffer in self.named_buffers(): + if not torch.is_floating_point(buffer): + continue + target_dtype = torch.float32 if should_keep_in_fp32(name) else dtype + buffer.data = buffer.data.to(dtype=target_dtype, non_blocking=non_blocking) + + return self + + @register_to_config + def __init__( + self, + patch_size: tuple[int, int, int] = (1, 2, 2), + in_channels: int = 16, + out_channels: int = 16, + hidden_size: int = 2048, + num_attention_heads: int = 16, + depth: int = 24, + intermediate_size: int = 6144, + text_dim: int = 2560, + freq_dim: int = 256, + norm_eps: float = 1e-6, + rope_theta: float = 256.0, + axes_dims: tuple[int, int, int] = (32, 48, 48), + axes_lens: tuple[int, int, int] = (8192, 1024, 1024), + qkv_bias: bool = False, + out_bias: bool = True, + patch_embed_bias: bool = True, + timestep_mlp_bias: bool = True, + num_experts: int = 0, + num_experts_per_tok: int = 8, + moe_intermediate_size: int = 512, + decoder_sparse_step: int = 1, + mlp_only_layers: tuple[int, ...] = (), + n_shared_experts: int | None = None, + score_func: str = "sigmoid", + norm_topk_prob: bool = True, + n_group: int | None = None, + topk_group: int | None = None, + routed_scaling_factor: float = 1.0, + ): + super().__init__() + head_dim = hidden_size // num_attention_heads + assert head_dim == sum(axes_dims), f"head_dim {head_dim} != sum(axes_dims) {sum(axes_dims)}" + mlp_only_layers = tuple(mlp_only_layers) + + self.patch_embedder = nn.Linear(in_channels * math.prod(patch_size), hidden_size, bias=patch_embed_bias) + self.time_proj = Timesteps(freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.time_embedder = TimestepEmbedding(freq_dim, hidden_size, act_fn="silu", sample_proj_bias=timestep_mlp_bias) + self.time_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size)) + self.text_embedder = LingBotVideoTextEmbedder(text_dim, hidden_size) + self.rope = LingBotVideoRotaryEmbedding(axes_dims, axes_lens, rope_theta) + self.blocks = nn.ModuleList( + [ + LingBotVideoBlock( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + intermediate_size=intermediate_size, + norm_eps=norm_eps, + qkv_bias=qkv_bias, + out_bias=out_bias, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + moe_intermediate_size=moe_intermediate_size, + decoder_sparse_step=decoder_sparse_step, + mlp_only_layers=mlp_only_layers, + n_shared_experts=n_shared_experts, + score_func=score_func, + norm_topk_prob=norm_topk_prob, + n_group=n_group, + topk_group=topk_group, + routed_scaling_factor=routed_scaling_factor, + layer_idx=i, + ) + for i in range(depth) + ] + ) + self.norm_out = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps) + self.norm_out_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size)) + self.proj_out = nn.Linear(hidden_size, math.prod(patch_size) * out_channels) + + def forward( + self, + hidden_states: torch.Tensor, # (B, C, T, H, W) + timestep: torch.Tensor, # (B,) in [0, 1000](= sigma*1000) + encoder_hidden_states: torch.Tensor, # (B, L, text_dim) + encoder_attention_mask: torch.Tensor | None = None, # (B, L) 1=valid + return_dict: bool = True, + ): + B, C, T, H, W = hidden_states.shape + pF, pH, pW = self.config.patch_size + gt, gh, gw = T // pF, H // pH, W // pW + n_video = gt * gh * gw + L = encoder_hidden_states.shape[1] + device = hidden_states.device + if encoder_attention_mask is not None: + text_lens = encoder_attention_mask.sum(dim=-1).long() + else: + text_lens = torch.full((B,), L, dtype=torch.long, device=device) + text_lens_list = [int(v) for v in text_lens.detach().cpu().tolist()] + packed_batch = B > 1 + + # patchify: token order (f h w), feature order (pf ph pw c) -- matches patchify_and_embed + patch_tokens = hidden_states.reshape(B, C, gt, pF, gh, pH, gw, pW) + patch_tokens = patch_tokens.permute(0, 2, 4, 6, 3, 5, 7, 1).reshape( + B, + n_video, + pF * pH * pW * C, + ) + if packed_batch: + x = torch.cat( + [self.patch_embedder(patch_tokens[i : i + 1]) for i in range(B)], + dim=1, + ) + else: + x = self.patch_embedder(patch_tokens) + + if packed_batch: + text_parts = [ + self.text_embedder(encoder_hidden_states[i : i + 1, : text_lens_list[i], :]) for i in range(B) + ] + text = torch.cat(text_parts, dim=1) + joint = _cat_interleave( + x, + [n_video] * B, + text, + text_lens_list, + ) + else: + text = self.text_embedder(encoder_hidden_states) + joint = torch.cat([x, text], dim=1) # [video; text] + joint_seq_len = joint.shape[1] + + # Per-sample RoPE: video t-axis start = real text length of this sample + 1 + rotary_parts = [self.rope(make_joint_position_ids(text_lens_list[i], gt, gh, gw, device)) for i in range(B)] + if packed_batch: + rotary = torch.cat(rotary_parts, dim=0).unsqueeze(0) + else: + rotary = torch.stack(rotary_parts, dim=0) # (B, S, head_dim/2) complex64 + + parallel_config = getattr(self, "_parallel_config", None) + use_packed_attention = parallel_config is not None + + attention_mask = None + moe_padding_mask = None + packed_indices = None + has_padding = encoder_attention_mask is not None and bool((text_lens < L).any()) + if packed_batch or use_packed_attention: + sample_seq_lens = [n_video + text_len for text_len in text_lens_list] + cu_seqlens = torch.zeros(B + 1, device=device, dtype=torch.int32) + cu_seqlens[1:] = torch.cumsum( + torch.tensor(sample_seq_lens, device=device, dtype=torch.int32), + dim=0, + ) + packed_indices = { + "cu_seqlens_kv": cu_seqlens, + "max_seqlen_in_batch_kv": max(sample_seq_lens), + } + if packed_batch and not use_packed_attention: + packed_indices["attention_mask"] = _packed_block_attention_mask(sample_seq_lens, device) + has_padding = False + if has_padding: + key_mask = torch.cat( + [torch.ones(B, n_video, dtype=torch.bool, device=device), encoder_attention_mask.bool()], + dim=1, + ) + attention_mask = key_mask[:, None, None, :] # (B,1,1,S) -> SDPA broadcast + moe_padding_mask = key_mask.reshape(-1).float() # (B*S,) + packed_cp = packed_indices is not None and parallel_config is not None + padding_size = 0 + if packed_cp: + cp_config = parallel_config.context_parallel_config + cp_world_size = int(getattr(cp_config, "ulysses_degree", getattr(cp_config, "_world_size", 1))) + padding_size = (cp_world_size - (joint_seq_len % cp_world_size)) % cp_world_size + if padding_size: + joint = torch.cat( + [ + joint, + torch.zeros( + joint.shape[0], + padding_size, + joint.shape[2], + device=joint.device, + dtype=joint.dtype, + ), + ], + dim=1, + ) + rotary = torch.cat( + [ + rotary, + torch.zeros( + rotary.shape[0], + padding_size, + rotary.shape[2], + device=rotary.device, + dtype=rotary.dtype, + ), + ], + dim=1, + ) + if packed_indices is None: + raise RuntimeError("packed_indices must be initialized for packed context parallel.") + packed_indices["cu_seqlens_kv"] = torch.cat( + [ + packed_indices["cu_seqlens_kv"], + packed_indices["cu_seqlens_kv"][-1:] + padding_size, + ], + dim=0, + ) + packed_indices["max_seqlen_in_batch_kv"] = max( + int(packed_indices["max_seqlen_in_batch_kv"]), + int(padding_size), + ) + joint_seq_len = joint.shape[1] + + timestep_for_embed = timestep.float() + timestep_proj = self.time_proj(timestep_for_embed) + t_emb = self.time_embedder(timestep_proj) # (B, D) + if packed_batch: + temb_input = torch.cat( + [t_emb[i : i + 1].unsqueeze(1).expand(1, n_video + text_lens_list[i], -1) for i in range(B)], + dim=1, + ) + if padding_size: + temb_input = torch.cat( + [ + temb_input, + torch.zeros( + temb_input.shape[0], + padding_size, + temb_input.shape[2], + device=temb_input.device, + dtype=temb_input.dtype, + ), + ], + dim=1, + ) + temb6 = self.time_modulation(temb_input.reshape(joint_seq_len, -1)) + temb6 = temb6.reshape(1, joint_seq_len, -1) + else: + temb_input = t_emb.unsqueeze(1).expand(B, joint_seq_len, -1) # (B, S, D) + temb6 = self.time_modulation(temb_input.reshape(B * joint_seq_len, -1)) + temb6 = temb6.reshape(B, joint_seq_len, -1) # (B, S, 6D) + + temb6 = temb6.reshape(temb6.shape[0] * temb6.shape[1], -1) + + for block in self.blocks: + joint = block( + joint, + temb6, + rotary, + attention_mask, + moe_padding_mask, + packed_indices=packed_indices, + parallel_config=parallel_config, + ) + final_mod = self.norm_out_modulation(temb_input.reshape(joint.shape[0] * joint.shape[1], -1)) + shift, scale = final_mod.reshape(joint.shape[0], joint.shape[1], -1).chunk(2, dim=-1) + final_hidden = self.norm_out(joint) * (1.0 + scale) + shift + projected = self.proj_out(final_hidden.to(self.proj_out.weight.dtype)) + if packed_cp: + if padding_size: + projected = projected[:, :-padding_size, :] + if packed_batch: + split_lengths: list[int] = [] + for text_len in text_lens_list: + split_lengths.extend([n_video, text_len]) + parts = torch.split(projected, split_lengths, dim=1) + x = torch.cat(parts[::2], dim=1).reshape(B, n_video, -1) + else: + x = projected[:, :n_video] + + # unpatchify (matches the rearrange in postprocess) + Cout = self.config.out_channels + x = x.reshape(B, gt, gh, gw, pF, pH, pW, Cout) + x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(B, Cout, T, H, W) + + if not return_dict: + return (x,) + return Transformer2DModelOutput(sample=x) diff --git a/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py b/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py new file mode 100644 index 000000000..f54fa387c --- /dev/null +++ b/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py @@ -0,0 +1,695 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from LingBot-Video (https://github.com/Robbyant/lingbot-video). + +from __future__ import annotations + +import os +from collections.abc import Iterable, Mapping +from contextlib import nullcontext +from typing import Any, ClassVar + +import numpy as np +import torch +import torch.distributed as dist +from diffusers import AutoencoderKLWan +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery +from vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer import LingBotVideoTransformer3DModel +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + +TOKEN_LENGTH = 37698 +HIDDEN_STATE_SKIP_LAYER = 0 +LOW_NOISE_TAIL_V1_DEFAULT_STEPS = 2 + +PROMPT_TEMPLATE = ( + "<|im_start|>system\nGiven a user input that may include a text prompt alone, " + "a text prompt with an image reference, or a text prompt with a video reference " + 'or a video reference alone, generate an "Enhanced prompt" that provides detailed ' + "visual descriptions suitable for video generation. Evaluate the level of detail " + "in the user's input: if it is simple, enrich it by adding specifics about colors, " + "shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal " + "progression, and spatial relationships to create vivid, concrete, and temporally " + "coherent scenes to create vivid and concrete scenes. Please generate only the " + "enhanced description for the prompt below and avoid including any additional " + "commentary or evaluations:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n" + "<|im_start|>assistant\n" +) +DEFAULT_NEGATIVE_PROMPT = ( + '{"universal_negative": {"visual_quality": ["low quality", "worst quality", "blurry", ' + '"pixelated", "jpeg artifacts", "low resolution", "unstable color", "color flicker", ' + '"underexposed", "overexposed", "invisible subject", "subject hidden in darkness"], ' + '"artistic_style": ["painting", "illustration", "drawing", "cartoon", "3d render", ' + '"cgi", "sketch", "digital art"], "composition_and_content": ["text", "watermark", ' + '"signature", "logo", "subtitles", "pillarboxed", "side bars", ' + '"portrait image in landscape frame"], "temporal_and_motion_stability": ["flickering", ' + '"jittery", "motion blur", "temporal inconsistency", "warping", "morphing", ' + '"incoherent motion", "unnatural movement", "static object with sudden jump", ' + '"frame-to-frame inconsistency"], "material_and_structure": ["plastic-like glass", ' + '"unrealistic texture", "deformed bottle", "liquid freezing improperly", ' + '"distorted reflections"]}}' +) + + +def _dtype_from_name(value: Any, default: torch.dtype) -> torch.dtype: + if isinstance(value, torch.dtype): + return value + if value is None: + return default + normalized = str(value).lower() + if normalized in {"bf16", "bfloat16", "torch.bfloat16"}: + return torch.bfloat16 + if normalized in {"fp16", "float16", "torch.float16"}: + return torch.float16 + if normalized in {"fp32", "float32", "torch.float32"}: + return torch.float32 + raise ValueError(f"Unsupported LingBot dtype: {value!r}.") + + +def _extract_prompt(req: OmniDiffusionRequest) -> tuple[str, str | None]: + prompt_obj = req.prompt + if isinstance(prompt_obj, str): + return prompt_obj, None + if isinstance(prompt_obj, Mapping): + prompt = prompt_obj.get("prompt", "") + negative_prompt = prompt_obj.get("negative_prompt") + return str(prompt), None if negative_prompt is None else str(negative_prompt) + raise TypeError(f"Unsupported LingBot prompt type: {type(prompt_obj)!r}.") + + +def _module_dtype(module: torch.nn.Module) -> torch.dtype: + try: + return next(module.parameters()).dtype + except StopIteration: + return torch.float32 + + +def _module_device(module: torch.nn.Module) -> torch.device: + try: + return next(module.parameters()).device + except StopIteration: + return torch.device("cpu") + + +def _transformer_timestep(timestep: torch.Tensor, transformer_dtype: torch.dtype) -> torch.Tensor: + sigma = timestep.float() / 1000.0 + if transformer_dtype in {torch.bfloat16, torch.float16}: + sigma = sigma.to(transformer_dtype) + return (sigma * 1000.0).float() + + +def _transformer_autocast(device: torch.device, transformer_dtype: torch.dtype): + if device.type != "cuda" or transformer_dtype not in {torch.bfloat16, torch.float16}: + return nullcontext() + return torch.autocast(device_type="cuda", dtype=transformer_dtype) + + +def _group_global_rank(group: Any | None, group_rank: int) -> int: + if group is None: + return group_rank + get_global_rank = getattr(dist, "get_global_rank", None) + if get_global_rank is None: + return group_rank + return int(get_global_rank(group, group_rank)) + + +def _validate_refiner_sigmas(sigmas: np.ndarray, t_thresh: float | None = None) -> np.ndarray: + arr = np.asarray(list(sigmas), dtype=np.float64) + if arr.ndim != 1 or arr.size == 0: + raise ValueError("refiner sigma schedule must be a non-empty 1D list") + if not np.all(np.isfinite(arr)): + raise ValueError("refiner sigma schedule contains non-finite values") + if np.any(arr < 0.0) or np.any(arr > 1.0): + raise ValueError(f"refiner sigma schedule values must be in [0, 1], got {arr.tolist()}") + if arr.size > 1 and not np.all(np.diff(arr) < 0.0): + raise ValueError(f"refiner sigma schedule must be strictly descending, got {arr.tolist()}") + if t_thresh is not None and abs(float(arr[0]) - float(t_thresh)) > 1e-6: + raise ValueError(f"refiner sigma schedule must start at t_thresh={float(t_thresh)}, got {float(arr[0])}") + return arr + + +def _compute_refiner_sigmas( + *, + sigma_max: float, + sigma_min: float, + num_inference_steps: int, + shift: float, + t_thresh: float | None, + tail_steps: int = 0, +) -> np.ndarray | None: + if t_thresh is None: + return None + t_value = float(t_thresh) + if not (0.0 < t_value <= 1.0): + raise ValueError(f"refiner t_thresh must lie in (0, 1], got {t_value}") + steps = int(num_inference_steps) + if steps < 1: + raise ValueError(f"num_inference_steps must be >= 1, got {steps}") + tail = int(tail_steps or 0) + if tail < 0: + raise ValueError(f"refiner_sigma_tail_steps must be >= 0, got {tail}") + + base = np.linspace(float(sigma_max), float(sigma_min), steps + 1).copy()[:-1] + shift_value = float(shift) + shifted = shift_value * base / (1.0 + (shift_value - 1.0) * base) + eps = 1e-6 + sigmas = shifted[shifted <= t_value + eps] + if sigmas.size == 0 or abs(float(sigmas[0]) - t_value) > eps: + sigmas = np.concatenate([[t_value], sigmas]) + if tail > 0: + start = float(sigmas[-1]) + stop = min(float(sigma_min), start) + extra = np.linspace(start, stop, tail + 2, dtype=np.float64)[1:-1] + sigmas = np.concatenate([sigmas, extra]) + return _validate_refiner_sigmas(sigmas, t_value).astype(np.float32) + + +def _pad_prompt_embeds( + embeds: torch.Tensor, + mask: torch.Tensor, + target_length: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if embeds.shape[0] != 1: + raise ValueError(f"batched CFG helper expects batch=1 inputs, got {embeds.shape[0]}") + if embeds.shape[1] > target_length: + raise ValueError(f"cannot pad length {embeds.shape[1]} down to {target_length}") + pad_len = target_length - embeds.shape[1] + if pad_len == 0: + return embeds, mask + embed_pad = torch.zeros(embeds.shape[0], pad_len, embeds.shape[2], dtype=embeds.dtype, device=embeds.device) + mask_pad = torch.zeros(mask.shape[0], pad_len, dtype=mask.dtype, device=mask.device) + return torch.cat([embeds, embed_pad], dim=1), torch.cat([mask, mask_pad], dim=1) + + +def _batch_cfg_prompt_inputs( + prompt_embeds: torch.Tensor, + prompt_mask: torch.Tensor, + negative_embeds: torch.Tensor, + negative_mask: torch.Tensor, + *, + null_cond_clone_zero: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + if null_cond_clone_zero: + zero_negative = torch.zeros_like(prompt_embeds) + return ( + torch.cat([prompt_embeds, zero_negative], dim=0), + torch.cat([prompt_mask, prompt_mask.clone()], dim=0), + ) + + target_length = max(int(prompt_embeds.shape[1]), int(negative_embeds.shape[1])) + prompt_padded, prompt_mask_padded = _pad_prompt_embeds(prompt_embeds, prompt_mask, target_length) + negative_padded, negative_mask_padded = _pad_prompt_embeds(negative_embeds, negative_mask, target_length) + return ( + torch.cat([prompt_padded, negative_padded], dim=0), + torch.cat([prompt_mask_padded, negative_mask_padded], dim=0), + ) + + +def get_lingbot_video_post_process_func(od_config: OmniDiffusionConfig): + del od_config + + def post_process_func(frames: torch.Tensor, sampling_params=None): + output_type = getattr(sampling_params, "output_type", None) or "pt" + if output_type == "np" and isinstance(frames, torch.Tensor): + return frames.float().cpu().numpy() + return frames + + return post_process_func + + +class LingBotVideoPipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery): + """Native vLLM-Omni entry for LingBot-Video checkpoints. + + The in-tree transformer supports both dense MLP blocks and routed MoE blocks. + Fused expert kernels and the optional ``refiner/`` transformer are not loaded + or executed. ``t_thresh`` only selects a low-noise sigma schedule for the + primary transformer; it does not enable automatic refiner orchestration. + """ + + supports_step_execution: ClassVar[bool] = False + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__() + del prefix + self.od_config = od_config + self.device = get_local_device() + self.vae_scale_factor_temporal = 4 + self.vae_scale_factor_spatial = 8 + self.token_length = TOKEN_LENGTH + self.hidden_state_skip_layer = HIDDEN_STATE_SKIP_LAYER + self.prompt_template = PROMPT_TEMPLATE + self._crop_start: int | None = None + + model = od_config.model + local_files_only = os.path.exists(model) + dtype = getattr(od_config, "dtype", torch.bfloat16) + model_config = getattr(od_config, "model_config", None) or {} + transformer_dtype = _dtype_from_name(model_config.get("transformer_dtype"), dtype) + text_encoder_dtype = _dtype_from_name(model_config.get("text_encoder_dtype"), dtype) + vae_dtype = _dtype_from_name(model_config.get("vae_dtype"), torch.float32) + + transformer_subfolder = str(model_config.get("transformer_subfolder", "transformer")) + text_encoder_subfolder = str(model_config.get("text_encoder_subfolder", "text_encoder")) + processor_subfolder = str(model_config.get("processor_subfolder", "processor")) + vae_subfolder = str(model_config.get("vae_subfolder", "vae")) + scheduler_subfolder = str(model_config.get("scheduler_subfolder", "scheduler")) + + self.transformer = LingBotVideoTransformer3DModel.from_pretrained( + model, + subfolder=transformer_subfolder, + torch_dtype=transformer_dtype, + local_files_only=local_files_only, + ).to(self.device) + text_encoder_kwargs: dict[str, Any] = { + "dtype": text_encoder_dtype, + "local_files_only": local_files_only, + } + self.text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( + model, + subfolder=text_encoder_subfolder, + **text_encoder_kwargs, + ).to(self.device) + self.processor = Qwen3VLProcessor.from_pretrained( + model, + subfolder=processor_subfolder, + local_files_only=local_files_only, + ) + self.vae = AutoencoderKLWan.from_pretrained( + model, + subfolder=vae_subfolder, + torch_dtype=vae_dtype, + local_files_only=local_files_only, + ).to(self.device) + self.scheduler = FlowUniPCMultistepScheduler.from_pretrained( + model, + subfolder=scheduler_subfolder, + local_files_only=local_files_only, + ) + self.set_progress_bar_config(disable=bool(model_config.get("quiet_progress", True))) + self.default_negative_prompt = DEFAULT_NEGATIVE_PROMPT + + def to(self, *args, **kwargs): + device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs) + super().to(*args, **kwargs) + if device is not None: + self.device = torch.device(device) + self.transformer.to(device=self.device, non_blocking=non_blocking) + self.text_encoder.to(device=self.device, non_blocking=non_blocking) + self.vae.to(device=self.device, non_blocking=non_blocking) + if dtype is not None: + self.transformer.to(dtype=dtype, non_blocking=non_blocking) + return self + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + consumed = list(weights) + if consumed: + raise RuntimeError( + f"{self.__class__.__name__}.load_weights received {len(consumed)} weight tensors; " + "LingBot components are loaded directly from subfolders during __init__." + ) + return {name for name, _ in self.named_parameters()} + + @staticmethod + def check_inputs(height: int, width: int, num_frames: int) -> None: + if num_frames != 1 and (num_frames - 1) % 4 != 0: + raise ValueError(f"`num_frames` must be 1 or 4n+1, got {num_frames}.") + if height % 16 != 0 or width % 16 != 0: + raise ValueError(f"`height` and `width` must be multiples of 16, got {height}x{width}.") + + @staticmethod + def apply_text_to_template(text: str, template: str = PROMPT_TEMPLATE) -> str: + return template.format(text) + + def _compute_crop_start(self) -> int: + if self._crop_start is None: + marker = "<|USER_INPUT_MARKER|>" + marked = self.prompt_template.format(marker) + marker_pos = marked.find(marker) + if marker_pos < 0: + self._crop_start = 0 + else: + prefix = self.processor( + text=marked[:marker_pos], + images=None, + videos=None, + return_tensors="pt", + ) + self._crop_start = int(prefix["input_ids"].shape[1]) + return self._crop_start + + def _build_prompt_inputs(self, prompt: str | list[str]): + prompts = [prompt] if isinstance(prompt, str) else list(prompt) + texts = [self.apply_text_to_template(text, self.prompt_template) for text in prompts] + return self.processor( + text=texts, + images=None, + videos=None, + do_resize=False, + truncation=True, + max_length=self.token_length, + padding="longest", + return_tensors="pt", + ) + + @torch.no_grad() + def encode_prompt( + self, + prompt: str | list[str], + *, + device: str | torch.device | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + device = torch.device(device) if device is not None else self.device + inputs = self._build_prompt_inputs(prompt).to(device) + outputs = self.text_encoder( + **inputs, + output_hidden_states=self.hidden_state_skip_layer is not None, + ) + if self.hidden_state_skip_layer is not None: + prompt_embeds = outputs.hidden_states[-(self.hidden_state_skip_layer + 1)] + else: + prompt_embeds = outputs.last_hidden_state + + prompt_mask = inputs["attention_mask"] + crop_start = self._compute_crop_start() + if crop_start > 0: + prompt_embeds = prompt_embeds[:, crop_start:] + prompt_mask = prompt_mask[:, crop_start:] + + if prompt_embeds.shape[0] == 1: + true_len = int(prompt_mask[0].sum().item()) + prompt_embeds = prompt_embeds[:, :true_len] + prompt_mask = prompt_mask[:, :true_len] + return prompt_embeds, prompt_mask + + def prepare_latents( + self, + num_frames: int, + height: int, + width: int, + generator: torch.Generator | None, + latents: torch.Tensor | None, + device: torch.device, + ) -> torch.Tensor: + latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + latent_height = height // self.vae_scale_factor_spatial + latent_width = width // self.vae_scale_factor_spatial + shape = (1, self.transformer.config.in_channels, latent_frames, latent_height, latent_width) + if latents is None: + return randn_tensor(shape, generator=generator, device=device, dtype=torch.float32) + if tuple(latents.shape) != shape: + raise ValueError(f"`latents` shape must be {shape}, got {tuple(latents.shape)}.") + return latents.to(device=device, dtype=torch.float32) + + def _dit_latent_to_vae(self, latents: torch.Tensor) -> torch.Tensor: + mean = torch.tensor(self.vae.config.latents_mean, device=latents.device, dtype=torch.float32) + std_inv = 1.0 / torch.tensor(self.vae.config.latents_std, device=latents.device, dtype=torch.float32) + mean = mean.view(1, -1, 1, 1, 1) + std_inv = std_inv.view(1, -1, 1, 1, 1) + return latents.float() / std_inv + mean + + @torch.no_grad() + def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: + vae_device = _module_device(self.vae) + vae_dtype = _module_dtype(self.vae) + vae_latents = self._dit_latent_to_vae(latents).to(device=vae_device, dtype=torch.float32) + if vae_latents.ndim == 5: + vae_latents = vae_latents.contiguous(memory_format=torch.channels_last_3d) + autocast_dtype = ( + vae_dtype if vae_device.type == "cuda" and vae_dtype in {torch.bfloat16, torch.float16} else None + ) + with torch.autocast("cuda", dtype=autocast_dtype or torch.bfloat16, enabled=autocast_dtype is not None): + decoded = self.vae.decode(vae_latents) + frames = decoded[0] if isinstance(decoded, tuple) else decoded.sample + frames = frames.float().clamp_(-1, 1) + frames = (frames + 1.0) / 2.0 + frames = frames.permute(0, 2, 3, 4, 1).cpu() + return frames[0] + + @torch.no_grad() + def _generate( + self, + *, + prompt: str, + negative_prompt: str = DEFAULT_NEGATIVE_PROMPT, + height: int = 480, + width: int = 480, + num_frames: int = 81, + num_inference_steps: int = 40, + guidance_scale: float = 6.0, + shift: float = 3.0, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_mask: torch.Tensor | None = None, + output_type: str = "pt", + cfg_parallel_group: Any | None = None, + batch_cfg: bool = False, + null_cond_clone_zero: bool = False, + t_thresh: float | None = None, + refiner_sigma_tail_steps: int = LOW_NOISE_TAIL_V1_DEFAULT_STEPS, + offload_vae_during_denoise: bool = False, + **extra_args, + ) -> torch.Tensor: + del extra_args + self.check_inputs(height, width, num_frames) + device = self.device + do_cfg = guidance_scale > 1.0 + effective_batch_cfg = bool(batch_cfg) + cfg_parallel = cfg_parallel_group is not None + cfg_parallel_rank = 0 + if cfg_parallel: + if not dist.is_available() or not dist.is_initialized(): + raise ValueError("`cfg_parallel_group` requires an initialized process group.") + if effective_batch_cfg: + raise ValueError("`cfg_parallel_group` and `batch_cfg` are mutually exclusive.") + if not do_cfg: + raise ValueError("CFG parallel requires `guidance_scale > 1.0`.") + cfg_parallel_rank = dist.get_rank(cfg_parallel_group) + cfg_parallel_world_size = dist.get_world_size(cfg_parallel_group) + if cfg_parallel_world_size != 2: + raise ValueError(f"CFG parallel currently requires exactly 2 ranks, got {cfg_parallel_world_size}.") + + if prompt_embeds is not None: + if prompt_mask is None: + raise ValueError("`prompt_mask` is required when `prompt_embeds` is provided.") + prompt_embeds = prompt_embeds.to(device=device) + prompt_mask = prompt_mask.to(device=device) + if negative_prompt_embeds is not None: + if negative_prompt_mask is None: + raise ValueError("`negative_prompt_mask` is required when `negative_prompt_embeds` is provided.") + negative_prompt_embeds = negative_prompt_embeds.to(device=device) + negative_prompt_mask = negative_prompt_mask.to(device=device) + + negative_embeds = None + negative_mask = None + if cfg_parallel and cfg_parallel_rank == 1: + if negative_prompt_embeds is not None: + negative_embeds, negative_mask = negative_prompt_embeds, negative_prompt_mask + else: + negative_embeds, negative_mask = self.encode_prompt(negative_prompt, device=device) + prompt_embeds = prompt_mask = None + else: + if prompt_embeds is None: + prompt_embeds, prompt_mask = self.encode_prompt(prompt, device=device) + if do_cfg and not cfg_parallel: + if null_cond_clone_zero: + negative_embeds = torch.zeros_like(prompt_embeds) + negative_mask = prompt_mask.clone() + elif negative_prompt_embeds is not None: + negative_embeds, negative_mask = negative_prompt_embeds, negative_prompt_mask + else: + negative_embeds, negative_mask = self.encode_prompt(negative_prompt, device=device) + + latents = self.prepare_latents(num_frames, height, width, generator, latents, device) + sigmas = _compute_refiner_sigmas( + sigma_max=float(self.scheduler.sigma_max), + sigma_min=float(self.scheduler.sigma_min), + num_inference_steps=num_inference_steps, + shift=shift, + t_thresh=t_thresh, + tail_steps=refiner_sigma_tail_steps, + ) + if sigmas is None: + self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift) + else: + self.scheduler.set_timesteps(int(sigmas.shape[0]), device=device, sigmas=sigmas, shift=1.0) + + transformer_dtype = _module_dtype(self.transformer) + vae_restore_device: torch.device | None = None + vae_offloaded = False + if offload_vae_during_denoise and output_type != "latent": + vae_device = _module_device(self.vae) + if vae_device.type == "cuda": + self.vae.to("cpu") + torch.accelerator.empty_cache() + vae_restore_device = vae_device + vae_offloaded = True + + cfg_latent_src = _group_global_rank(cfg_parallel_group, 0) + cfg_uncond_src = _group_global_rank(cfg_parallel_group, 1) + for timestep in self.progress_bar(self.scheduler.timesteps): + if cfg_parallel: + dist.broadcast(latents, src=cfg_latent_src, group=cfg_parallel_group) + timestep_batch = _transformer_timestep(timestep, transformer_dtype).expand(1).to(device) + latent_model_input = latents + if cfg_parallel: + if cfg_parallel_rank == 0: + branch_embeds = prompt_embeds + branch_mask = prompt_mask + else: + branch_embeds = negative_embeds + branch_mask = negative_mask + if branch_embeds is None: + raise RuntimeError("CFG branch embeddings were not initialized.") + with _transformer_autocast(device, transformer_dtype): + branch_noise_pred = self.transformer( + latent_model_input, + timestep_batch, + branch_embeds.to(transformer_dtype), + encoder_attention_mask=branch_mask, + return_dict=False, + )[0].float() + if cfg_parallel_rank == 0: + noise_pred = branch_noise_pred + noise_pred_uncond = torch.empty_like(noise_pred) + else: + noise_pred_uncond = branch_noise_pred + dist.broadcast(noise_pred_uncond, src=cfg_uncond_src, group=cfg_parallel_group) + if cfg_parallel_rank != 0: + continue + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond) + else: + if prompt_embeds is None: + raise RuntimeError("Prompt embeddings were not initialized.") + prompt_model_input = prompt_embeds.to(transformer_dtype) + if do_cfg and effective_batch_cfg: + if negative_embeds is None or negative_mask is None: + raise RuntimeError("Negative embeddings were not initialized for CFG.") + cfg_embeds, cfg_mask = _batch_cfg_prompt_inputs( + prompt_model_input, + prompt_mask, + negative_embeds.to(transformer_dtype), + negative_mask, + null_cond_clone_zero=False, + ) + cfg_latents = torch.cat([latent_model_input, latent_model_input], dim=0) + cfg_timesteps = torch.cat([timestep_batch, timestep_batch], dim=0) + with _transformer_autocast(device, transformer_dtype): + noise_batched = self.transformer( + cfg_latents, + cfg_timesteps, + cfg_embeds, + encoder_attention_mask=cfg_mask, + return_dict=False, + )[0].float() + noise_pred, noise_pred_uncond = noise_batched.chunk(2, dim=0) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond) + else: + with _transformer_autocast(device, transformer_dtype): + noise_pred = self.transformer( + latent_model_input, + timestep_batch, + prompt_model_input, + encoder_attention_mask=prompt_mask, + return_dict=False, + )[0].float() + + if do_cfg and not effective_batch_cfg: + if negative_embeds is None or negative_mask is None: + raise RuntimeError("Negative embeddings were not initialized for CFG.") + with _transformer_autocast(device, transformer_dtype): + noise_pred_uncond = self.transformer( + latent_model_input, + timestep_batch, + negative_embeds.to(transformer_dtype), + encoder_attention_mask=negative_mask, + return_dict=False, + )[0].float() + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond) + + latents = self.scheduler.step(noise_pred, timestep, latents, return_dict=False, generator=generator)[0] + + if cfg_parallel: + dist.barrier(group=cfg_parallel_group) + if cfg_parallel_rank != 0: + return latents if output_type == "latent" else [] + + if output_type == "latent": + return latents + if output_type in {"pt", "np"}: + if vae_offloaded and vae_restore_device is not None: + self.vae.to(device=vae_restore_device) + torch.accelerator.empty_cache() + return self._decode_latents(latents) + raise ValueError(f"Unsupported output_type: {output_type}") + + @torch.inference_mode() + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + if req.num_reqs != 1: + raise ValueError(f"LingBotVideoPipeline only supports one request per batch, got {req.num_reqs}.") + request = req.requests[0] + prompt, prompt_negative = _extract_prompt(request) + sampling = request.sampling_params + extra_args = dict(sampling.extra_args or {}) + + generator = sampling.generator + if isinstance(generator, list): + generator = generator[0] if generator else None + if generator is None and sampling.seed is not None: + generator = torch.Generator(device=self.device).manual_seed(int(sampling.seed)) + + height = sampling.height if sampling.height is not None else extra_args.pop("height", 480) + width = sampling.width if sampling.width is not None else extra_args.pop("width", 480) + num_frames = sampling.num_frames or extra_args.pop("num_frames", 81) + num_inference_steps = ( + sampling.num_inference_steps + if sampling.num_inference_steps is not None + else extra_args.pop("num_inference_steps", 40) + ) + guidance_scale = ( + sampling.guidance_scale if sampling.guidance_scale_provided else extra_args.pop("guidance_scale", 6.0) + ) + shift = extra_args.pop( + "shift", + extra_args.pop("flow_shift", getattr(self.od_config, "flow_shift", None) or 3.0), + ) + negative_prompt = extra_args.pop("negative_prompt", prompt_negative or self.default_negative_prompt) + output_type = ( + sampling.output_type or getattr(self.od_config, "output_type", None) or extra_args.pop("output_type", "pt") + ) + if output_type not in {"pt", "np", "latent"}: + output_type = "pt" + sampling.output_type = output_type + + frames = self._generate( + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + shift=shift, + generator=generator, + latents=sampling.latents, + output_type=output_type, + batch_cfg=bool(extra_args.pop("batch_cfg", False)), + null_cond_clone_zero=bool(extra_args.pop("null_cond_clone_zero", False)), + offload_vae_during_denoise=bool(extra_args.pop("offload_vae_during_denoise", False)), + t_thresh=extra_args.pop("t_thresh", None), + refiner_sigma_tail_steps=int(extra_args.pop("refiner_sigma_tail_steps", LOW_NOISE_TAIL_V1_DEFAULT_STEPS)), + ) + return DiffusionOutput(output=frames) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 22c36f32b..aea257dd6 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -256,6 +256,11 @@ _DIFFUSION_MODELS = { "pipeline_hunyuan_video_1_5_i2v", "HunyuanVideo15I2VPipeline", ), + "LingBotVideoPipeline": ( + "lingbot_video", + "pipeline_lingbot_video", + "LingBotVideoPipeline", + ), "MagiHumanPipeline": ( "magi_human", "pipeline_magi_human", @@ -536,6 +541,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "Flux2Pipeline": "get_flux2_post_process_func", "HunyuanVideo15Pipeline": "get_hunyuan_video_15_post_process_func", "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_post_process_func", + "LingBotVideoPipeline": "get_lingbot_video_post_process_func", "MagiHumanPipeline": "get_magi_human_post_process_func", "OmniVoicePipeline": "get_omnivoice_post_process_func", "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func", diff --git a/vllm_omni/model_extras/lingbot_video.py b/vllm_omni/model_extras/lingbot_video.py new file mode 100644 index 000000000..c7aae3763 --- /dev/null +++ b/vllm_omni/model_extras/lingbot_video.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +LINGBOT_VIDEO_EXTRA_BODY_PARAMS = frozenset( + { + "batch_cfg", + "flow_shift", + "negative_prompt", + "null_cond_clone_zero", + "offload_vae_during_denoise", + "output_type", + "refiner_sigma_tail_steps", + "shift", + "t_thresh", + } +) diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 79f8b4360..ca13a35c5 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -34,6 +34,7 @@ from vllm_omni.model_extras.helios import ( HELIOS_EXTRA_BODY_PARAMS, HELIOS_EXTRA_OUTPUT_PARAMS, ) +from vllm_omni.model_extras.lingbot_video import LINGBOT_VIDEO_EXTRA_BODY_PARAMS from vllm_omni.model_extras.magi_human import ( MAGI_HUMAN_EXTRA_BODY_PARAMS, MAGI_HUMAN_EXTRA_OUTPUT_PARAMS, @@ -170,6 +171,9 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { "extra_body_params": HELIOS_EXTRA_BODY_PARAMS, "extra_output_params": HELIOS_EXTRA_OUTPUT_PARAMS, }, + "LingBotVideoPipeline": { + "extra_body_params": LINGBOT_VIDEO_EXTRA_BODY_PARAMS, + }, "WanVACEPipeline": { "extra_body_params": VACE_EXTRA_BODY_PARAMS, "extra_output_params": VACE_EXTRA_OUTPUT_PARAMS,