Compare commits

...
Author SHA1 Message Date
hsliuustc0106 603a3ae67b Fix LingBot diffusion alignment CI
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-07-10 08:31:38 +08:00
hsliuustc0106 ebd0e024e9 Fix LingBot packed CFG attention fallback
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-07-10 00:43:34 +08:00
hsliuustc0106 7d59b818d0 Remove LingBot-specific attention env
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-07-10 00:29:02 +08:00
hsliuustc0106 6f4e53553a Add LingBot T2V recipe and tests
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-07-10 00:29:02 +08:00
hsliuustc0106 0e6490d929 [diffusion][feature] Add LingBot dense video support
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-07-10 00:29:02 +08:00
18 changed files with 2417 additions and 0 deletions

No files matched your search

@@ -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()
+1
View File
@@ -46,6 +46,7 @@ th {
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | | `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | | `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | | `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `LingBotVideoPipeline` | LingBot-Video dense T2V | `robbyant/lingbot-video-dense-1.3b` | ✅︎ | | | |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
@@ -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 the LingBot dense preset.")
parser.add_argument("--model", default="/home/models/lingbot-video-dense-1.3b")
parser.add_argument("--prompt", default="a robotic arm picks up a red block")
parser.add_argument("--negative-prompt", default="")
parser.add_argument("--output", default="lingbot_video_output.mp4")
parser.add_argument("--height", type=int, default=192)
parser.add_argument("--width", type=int, default=320)
parser.add_argument("--num-frames", type=int, default=9)
parser.add_argument("--num-inference-steps", type=int, default=2)
parser.add_argument("--guidance-scale", type=float, default=3.0)
parser.add_argument("--flow-shift", type=float, default=3.0)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--fps", type=int, default=24)
parser.add_argument("--output-type", choices=["pt", "np"], default="pt")
parser.add_argument("--batch-cfg", action="store_true")
return parser.parse_args()
def _peak_memory_mb(output: Any) -> float:
result = output[0] if isinstance(output, list) and output else output
return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
def _extract_video(output: Any) -> Any:
if isinstance(output, list):
output = output[0] if output else None
if isinstance(output, OmniRequestOutput):
if output.final_output_type != "image":
raise ValueError(f"Unexpected output type {output.final_output_type!r}.")
if not output.images:
raise ValueError("No video frames found in OmniRequestOutput.")
return output.images[0] if len(output.images) == 1 else output.images
return output
def _to_frame_array(video: Any) -> np.ndarray:
if isinstance(video, torch.Tensor):
tensor = video.detach().cpu()
if tensor.dim() == 5:
tensor = tensor[0] if tensor.shape[1] not in (3, 4) else tensor[0].permute(1, 2, 3, 0)
elif tensor.dim() == 4 and tensor.shape[0] in (3, 4):
tensor = tensor.permute(1, 2, 3, 0)
if tensor.is_floating_point():
if float(tensor.min()) < 0.0:
tensor = tensor.clamp(-1, 1) * 0.5 + 0.5
else:
tensor = tensor.clamp(0, 1)
return tensor.float().numpy()
if isinstance(video, np.ndarray):
array = video[0] if video.ndim == 5 else video
if np.issubdtype(array.dtype, np.integer):
return array.astype(np.float32) / 255.0
return array.astype(np.float32, copy=False)
if isinstance(video, list):
return np.stack([_to_frame_array(frame) for frame in video], axis=0)
return np.asarray(video, dtype=np.float32)
def main() -> None:
args = parse_args()
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
omni = Omni(
model=args.model,
model_class_name="LingBotVideoPipeline",
flow_shift=args.flow_shift,
parallel_config=DiffusionParallelConfig(),
)
prompt: dict[str, str] = {"prompt": args.prompt}
if args.negative_prompt:
prompt["negative_prompt"] = args.negative_prompt
sampling_params = OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
num_frames=args.num_frames,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
generator=generator,
output_type=args.output_type,
extra_args={"batch_cfg": args.batch_cfg, "shift": args.flow_shift},
)
start = time.perf_counter()
output = omni.generate(prompt, sampling_params)
elapsed = time.perf_counter() - start
print(f"Total generation time: {elapsed:.4f} seconds ({elapsed * 1000:.2f} ms)")
peak_mb = _peak_memory_mb(output)
if peak_mb:
print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
video = _to_frame_array(_extract_video(output))
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
from diffusers.utils import export_to_video
export_to_video(list(video), str(output_path), fps=args.fps)
print(f"Saved generated video to {output_path}")
omni.close()
if __name__ == "__main__":
main()
+1
View File
@@ -48,6 +48,7 @@ recipes/
| [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) | | [`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 | | [`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 | | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense text-to-video serving | 1x L20X / A100 / H100 |
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 | | [`cosmos3/Cosmos3-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 | | [`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 | | [`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 |
+143
View File
@@ -0,0 +1,143 @@
# LingBot-Video
> Native dense text-to-video serving for `robbyant/lingbot-video-dense-1.3b`
## Summary
- Vendor: Robbyant
- Model: `robbyant/lingbot-video-dense-1.3b`
- Task: Text-to-video generation
- Mode: Offline generation and online serving with the OpenAI-compatible `/v1/videos` API
- Maintainer: Community
## When to use this recipe
Use this recipe when you want to run the dense LingBot-Video checkpoint with
vLLM-Omni's native pipeline. The runtime path does not import the upstream
`lingbot_video` Python package; it loads the checkpoint components directly
with the in-tree `LingBotVideoPipeline`, dense DiT transformer, shared FlowUniPC
scheduler, Qwen3-VL text encoder, and Wan VAE.
This recipe is intentionally text-to-video only. Text-to-image, image-conditioned
video, TI2V, MoE checkpoints, and fused-expert kernels are not covered here.
## References
- Dense checkpoint: <https://huggingface.co/robbyant/lingbot-video-dense-1.3b>
- Upstream project: <https://github.com/Robbyant/lingbot-video>
- Related offline example: [`examples/offline_inference/text_to_video/text_to_video_lingbot.py`](../../examples/offline_inference/text_to_video/text_to_video_lingbot.py)
- Related online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
## Hardware Support
This recipe documents the CUDA single-GPU dense checkpoint path. Multi-GPU
parallelism, Cache-DiT, CPU offload, and the LingBot MoE path are not validated
for this model in this PR.
## GPU
### 1 x NVIDIA L20X / A100 / H100
The dense 1.3B checkpoint has been smoke-tested on a single NVIDIA L20X at a
small validation shape (`192x320`, 9 frames, 2 steps). Use a larger GPU or lower
resolution/frame count if your target workload runs out of memory.
#### Offline T2V
```bash
CUDA_VISIBLE_DEVICES=0 \
python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
--model robbyant/lingbot-video-dense-1.3b \
--prompt "a robotic arm picks up a red block" \
--output lingbot_t2v.mp4 \
--height 192 \
--width 320 \
--num-frames 9 \
--num-inference-steps 2 \
--guidance-scale 3.0 \
--flow-shift 3.0 \
--seed 42 \
--fps 24
```
#### Online Serving
```bash
CUDA_VISIBLE_DEVICES=0 \
vllm serve robbyant/lingbot-video-dense-1.3b \
--omni \
--model-class-name LingBotVideoPipeline \
--port 8091
```
After the server is ready, submit a text-to-video job:
```bash
create_response=$(curl -s http://localhost:8091/v1/videos \
-F "model=robbyant/lingbot-video-dense-1.3b" \
-F "prompt=a robotic arm picks up a red block" \
-F "width=320" \
-F "height=192" \
-F "num_frames=9" \
-F "fps=24" \
-F "num_inference_steps=2" \
-F "guidance_scale=3.0" \
-F "flow_shift=3.0" \
-F "seed=42")
video_id=$(echo "${create_response}" | jq -r '.id')
while true; do
status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
if [ "${status}" = "completed" ]; then
break
fi
if [ "${status}" = "failed" ]; then
curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
exit 1
fi
sleep 2
done
curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o lingbot_t2v.mp4
```
## Key Parameters
| Parameter | Suggested smoke value | Notes |
|-----------|-----------------------|-------|
| `height` | `192` | Must be a multiple of 16 |
| `width` | `320` | Must be a multiple of 16 |
| `num_frames` | `9` | Must be `1` or `4n + 1`; this PR validates T2V with video outputs |
| `num_inference_steps` | `2` | Use more steps for quality sweeps |
| `guidance_scale` | `3.0` | CFG is active when this is greater than `1.0` |
| `flow_shift` | `3.0` | Scheduler flow-shift; aliases the pipeline's internal `shift` |
| `negative_prompt` | model default | Optional text describing artifacts to avoid |
| `fps` | `24` | Output MP4 frame rate |
## Validation
Local dense smoke run:
- Shape: 9 frames at `192x320`
- Steps: 2
- Request generation time: `0.2923s`
- Peak reserved GPU memory: `14548 MiB`
Local parity harness against the upstream repository:
- Shape: `[9, 192, 320, 3]`
- MAE: `0.0065238`
- MSE: `0.00006650`
- PSNR: `41.77 dB`
- Native request time: `0.2875s`
Do not treat these as production benchmarks; they are draft-PR validation
numbers for the small dense smoke configuration.
## Known Limitations
- T2V only. T2I, I2V, and TI2V are not claimed by this PR.
- Dense checkpoint only. MoE/fused-expert support belongs in a separate PR.
- No Cache-DiT, TeaCache, tensor parallelism, sequence parallelism, CFG
parallelism, HSDP, CPU offload, or VAE patch parallelism validation yet.
- Only one request per LingBot pipeline batch is currently supported.
@@ -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
}
]
}
]
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
@@ -0,0 +1,117 @@
# 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
@@ -0,0 +1,195 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import numpy as np
import pytest
import torch
from torch import nn
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def _make_pipeline():
from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import LingBotVideoPipeline
pipeline = object.__new__(LingBotVideoPipeline)
nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
pipeline.default_negative_prompt = "default negative"
pipeline.od_config = SimpleNamespace(flow_shift=None)
return pipeline
def _make_request_batch(prompt, **sampling_overrides):
sampling = OmniDiffusionSamplingParams(**sampling_overrides)
request = OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id="lingbot-test")
return DiffusionRequestBatch([request])
def test_lingbot_video_pipeline_import_and_registry():
from vllm_omni.diffusion.models.lingbot_video import (
LingBotVideoPipeline,
LingBotVideoTransformer3DModel,
get_lingbot_video_post_process_func,
)
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS
assert LingBotVideoPipeline is not None
assert LingBotVideoTransformer3DModel is not None
assert get_lingbot_video_post_process_func is not None
assert _DIFFUSION_MODELS["LingBotVideoPipeline"] == (
"lingbot_video",
"pipeline_lingbot_video",
"LingBotVideoPipeline",
)
assert _DIFFUSION_POST_PROCESS_FUNCS["LingBotVideoPipeline"] == "get_lingbot_video_post_process_func"
def test_component_discovery_declarations():
from vllm_omni.diffusion.models.lingbot_video import LingBotVideoPipeline
assert LingBotVideoPipeline._dit_modules == ["transformer"]
assert LingBotVideoPipeline._encoder_modules == ["text_encoder"]
assert LingBotVideoPipeline._vae_modules == ["vae"]
assert LingBotVideoPipeline.supports_step_execution is False
def test_extra_body_params_include_video_flow_shift_alias():
from vllm_omni.model_extras import get_extra_body_params
params = get_extra_body_params("LingBotVideoPipeline")
assert "shift" in params
assert "flow_shift" in params
assert "batch_cfg" in params
def test_extract_prompt_string_and_mapping():
from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import _extract_prompt
string_req = SimpleNamespace(prompt="a robot waves")
assert _extract_prompt(string_req) == ("a robot waves", None)
mapping_req = SimpleNamespace(prompt={"prompt": "a robot walks", "negative_prompt": "low quality"})
assert _extract_prompt(mapping_req) == ("a robot walks", "low quality")
with pytest.raises(TypeError, match="Unsupported LingBot prompt type"):
_extract_prompt(SimpleNamespace(prompt=["not", "supported"]))
def test_check_inputs_accepts_t2v_shapes_and_rejects_invalid_values():
from vllm_omni.diffusion.models.lingbot_video import LingBotVideoPipeline
LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=1)
LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=9)
with pytest.raises(ValueError, match="num_frames"):
LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=8)
with pytest.raises(ValueError, match="height"):
LingBotVideoPipeline.check_inputs(height=190, width=320, num_frames=9)
def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
frames = torch.ones(2, 4, 4, 3)
post = get_lingbot_video_post_process_func(SimpleNamespace())
assert post(frames, SimpleNamespace(output_type="pt")) is frames
array = post(frames, SimpleNamespace(output_type="np"))
assert isinstance(array, np.ndarray)
assert array.shape == (2, 4, 4, 3)
def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
pipeline = _make_pipeline()
calls = []
def fake_generate(**kwargs):
calls.append(kwargs)
return torch.zeros(9, 192, 320, 3)
pipeline._generate = fake_generate
req = _make_request_batch(
{"prompt": "a robot arm", "negative_prompt": "bad hands"},
height=192,
width=320,
num_frames=9,
num_inference_steps=2,
guidance_scale=3.0,
seed=42,
output_type="pt",
extra_args={"flow_shift": 4.0, "batch_cfg": True},
)
out = pipeline.forward(req)
assert torch.equal(out.output, torch.zeros(9, 192, 320, 3))
assert len(calls) == 1
call = calls[0]
assert call["prompt"] == "a robot arm"
assert call["negative_prompt"] == "bad hands"
assert call["height"] == 192
assert call["width"] == 320
assert call["num_frames"] == 9
assert call["num_inference_steps"] == 2
assert call["guidance_scale"] == 3.0
assert call["shift"] == 4.0
assert call["batch_cfg"] is True
assert call["output_type"] == "pt"
def test_forward_prefers_shift_over_flow_shift_and_defaults_negative_prompt():
pipeline = _make_pipeline()
calls = []
def fake_generate(**kwargs):
calls.append(kwargs)
return torch.zeros(5, 192, 192, 3)
pipeline._generate = fake_generate
req = _make_request_batch(
"a robot arm",
height=192,
width=192,
num_frames=5,
num_inference_steps=2,
guidance_scale=1.0,
seed=42,
extra_args={"shift": 5.0, "flow_shift": 4.0},
)
pipeline.forward(req)
assert calls[0]["shift"] == 5.0
assert calls[0]["negative_prompt"] == "default negative"
def test_forward_rejects_multi_request_batches():
pipeline = _make_pipeline()
first = OmniDiffusionRequest(
prompt="first",
sampling_params=OmniDiffusionSamplingParams(seed=1),
request_id="lingbot-first",
)
second = OmniDiffusionRequest(
prompt="second",
sampling_params=OmniDiffusionSamplingParams(seed=2),
request_id="lingbot-second",
)
with pytest.raises(ValueError, match="only supports one request"):
pipeline.forward(DiffusionRequestBatch([first, second]))
def test_load_weights_rejects_external_weight_stream():
pipeline = _make_pipeline()
assert pipeline.load_weights([]) == set()
with pytest.raises(RuntimeError, match="components are loaded directly"):
pipeline.load_weights([("transformer.weight", torch.zeros(1))])
@@ -0,0 +1,61 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Online serving smoke for ``robbyant/lingbot-video-dense-1.3b``.
The LingBot dense PR supports text-to-video only, so this module exercises the
native ``/v1/videos`` path with a small 9-frame request.
"""
import os
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
MODEL = "robbyant/lingbot-video-dense-1.3b"
PROMPT = "a robotic arm picks up a red block"
NEGATIVE_PROMPT = "low quality, blurry, watermark, text"
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
def _get_diffusion_feature_cases(model: str):
return [
pytest.param(
OmniServerParams(
model=model,
server_args=["--model-class-name", "LingBotVideoPipeline"],
),
id="default",
marks=SINGLE_CARD_FEATURE_MARKS,
),
]
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
def test_text_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
request_config = {
"model": omni_server.model,
"form_data": {
"model": omni_server.model,
"prompt": PROMPT,
"negative_prompt": NEGATIVE_PROMPT,
"height": 192,
"width": 320,
"num_frames": 9,
"fps": 24,
"num_inference_steps": 2,
"guidance_scale": 3.0,
"flow_shift": 3.0,
"seed": 42,
},
}
openai_client.send_video_diffusion_request(request_config)
@@ -0,0 +1,82 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
L4 expansion coverage for ``robbyant/lingbot-video-dense-1.3b``.
This is intentionally T2V-only. T2I, I2V, TI2V, MoE, and multi-GPU feature
rows should be added in follow-up PRs when those paths are implemented.
"""
import json
import os
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
MODEL = "robbyant/lingbot-video-dense-1.3b"
PROMPT = "a robotic arm picks up a red block"
NEGATIVE_PROMPT = "low quality, blurry, watermark, text"
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
def _get_diffusion_feature_cases(model: str):
return [
pytest.param(
OmniServerParams(
model=model,
server_args=["--model-class-name", "LingBotVideoPipeline"],
),
id="default",
marks=SINGLE_CARD_FEATURE_MARKS,
),
]
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
def test_cfg_off(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
request_config = {
"model": omni_server.model,
"form_data": {
"model": omni_server.model,
"prompt": PROMPT,
"height": 192,
"width": 192,
"num_frames": 9,
"fps": 24,
"num_inference_steps": 2,
"guidance_scale": 1.0,
"flow_shift": 3.0,
"seed": 42,
},
}
openai_client.send_video_diffusion_request(request_config)
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
def test_batch_cfg_extra_params(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
request_config = {
"model": omni_server.model,
"form_data": {
"model": omni_server.model,
"prompt": PROMPT,
"negative_prompt": NEGATIVE_PROMPT,
"height": 192,
"width": 320,
"num_frames": 9,
"fps": 24,
"num_inference_steps": 2,
"guidance_scale": 3.0,
"flow_shift": 3.0,
"seed": 42,
"extra_params": json.dumps({"batch_cfg": True}, separators=(",", ":")),
},
}
openai_client.send_video_diffusion_request(request_config)
@@ -66,6 +66,7 @@ EXCLUDED_MODELS = [
"AudioXPipeline", "AudioXPipeline",
"HunyuanVideo15Pipeline", "HunyuanVideo15Pipeline",
"HunyuanVideo15ImageToVideoPipeline", "HunyuanVideo15ImageToVideoPipeline",
"LingBotVideoPipeline",
"MagiHumanPipeline", "MagiHumanPipeline",
"OmniVoicePipeline", "OmniVoicePipeline",
"OmniVoice", "OmniVoice",
@@ -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",
]
@@ -0,0 +1,702 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from LingBot-Video (https://github.com/Robbyant/lingbot-video).
import math
import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.models.modeling_utils import ModelMixin
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.backends.utils.fa import (
flash_attn_varlen_func as flash_attn_varlen_func_v3,
)
from vllm_omni.diffusion.attention.layer import Attention
LINGBOT_VIDEO_FP32_MODULES = (
"time_embedder",
"time_modulation",
"scale_shift_table",
"norm",
"norm1",
"norm2",
"norm_q",
"norm_k",
"norm_post_attn",
"norm_post_ffn",
"norm_out",
"norm_out_modulation",
"router",
)
def should_keep_in_fp32(name: str) -> bool:
return any(module_name in name.split(".") for module_name in LINGBOT_VIDEO_FP32_MODULES)
def _all_to_all_split_cat(
local_input: torch.Tensor,
scatter_dim: int,
gather_dim: int,
group: dist.ProcessGroup,
) -> torch.Tensor:
world_size = dist.get_world_size(group)
input_list = [tensor.contiguous() for tensor in torch.tensor_split(local_input, world_size, scatter_dim)]
output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)]
dist.all_to_all(output_list, input_list, group=group)
return torch.cat(output_list, dim=gather_dim).contiguous()
class LingBotVideoRMSNorm(nn.Module):
"""RMSNorm with fp32 accumulation."""
def __init__(self, dim: int, eps: float = 1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(dim))
self.variance_epsilon = eps
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return (self.weight * hidden_states).to(input_dtype)
def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
"""Apply complex RoPE to `(B, S, H, D)` attention tensors."""
with torch.amp.autocast("cuda", enabled=False):
x_c = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
out = torch.view_as_real(x_c * freqs_cis.unsqueeze(2)).flatten(3)
return out.type_as(x)
class LingBotVideoRotaryEmbedding(nn.Module):
"""Complex64 RoPE table indexed by position ids."""
def __init__(self, axes_dims: tuple[int, ...], axes_lens: tuple[int, ...], theta: float):
super().__init__()
self.axes_dims = tuple(axes_dims)
self.axes_lens = list(axes_lens)
self.theta = theta
self.freqs_cis = None
@staticmethod
def precompute_freqs_cis(dim: tuple[int, ...], end: tuple[int, ...], theta: float):
freqs_cis = []
for d, e in zip(dim, end):
freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d))
timestep = torch.arange(e, device=freqs.device, dtype=torch.float64)
freqs = torch.outer(timestep, freqs).float()
freqs_cis.append(torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64))
return freqs_cis
def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
# position_ids: (S, 3) int -> (S, head_dim/2) complex64
device = position_ids.device
max_vals = position_ids.max(dim=0).values.tolist()
needs_rebuild = self.freqs_cis is None or any(
max_val >= axis_len for max_val, axis_len in zip(max_vals, self.axes_lens)
)
if needs_rebuild:
for i in range(len(self.axes_lens)):
if max_vals[i] >= self.axes_lens[i]:
self.axes_lens[i] = int(max_vals[i] * 1.5) + 1
self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, tuple(self.axes_lens), theta=self.theta)
self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
elif self.freqs_cis[0].device != device:
self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
return torch.cat([self.freqs_cis[i][position_ids[:, i]] for i in range(len(self.axes_dims))], dim=-1)
def make_joint_position_ids(text_len: int, grid_t: int, grid_h: int, grid_w: int, device: torch.device) -> torch.Tensor:
"""3D positions in [video; text] order. Text t-axis is 1..text_len; video t-axis starts at text_len+1.
Matches patchify_and_embed: cap start (1,0,0); vision start (cap_len+1,0,0);
freqs ordered with x first and cap second (same order as cat_interleave).
"""
tt = torch.arange(grid_t, device=device, dtype=torch.int32) + (text_len + 1)
hh = torch.arange(grid_h, device=device, dtype=torch.int32)
ww = torch.arange(grid_w, device=device, dtype=torch.int32)
grid = torch.stack(torch.meshgrid(tt, hh, ww, indexing="ij"), dim=-1).flatten(0, 2)
text_t = torch.arange(text_len, device=device, dtype=torch.int32) + 1
text_pos = torch.stack([text_t, torch.zeros_like(text_t), torch.zeros_like(text_t)], dim=-1)
return torch.cat([grid, text_pos], dim=0) # (Nx + L, 3)
def _cat_interleave(
a: torch.Tensor,
len_a: list[int],
b: torch.Tensor,
len_b: list[int],
) -> torch.Tensor:
a_split = torch.split(a, len_a, dim=1)
b_split = torch.split(b, len_b, dim=1)
blocks: list[torch.Tensor] = []
for x_part, text_part in zip(a_split, b_split):
blocks.extend([x_part, text_part])
return torch.cat(blocks, dim=1)
def _packed_block_attention_mask(sample_seq_lens: list[int], device: torch.device) -> torch.Tensor:
total_seq_len = sum(sample_seq_lens)
mask = torch.zeros(
1,
1,
total_seq_len,
total_seq_len,
dtype=torch.bool,
device=device,
)
start = 0
for seq_len in sample_seq_lens:
end = start + seq_len
mask[:, :, start:end, start:end] = True
start = end
return mask
class LingBotVideoTextEmbedder(nn.Module):
"""Matches CondProjection: RMSNorm(text_dim, eps=1e-6 fixed) -> Linear-SiLU-Linear."""
def __init__(self, text_dim: int, hidden_size: int):
super().__init__()
self.norm = LingBotVideoRMSNorm(text_dim, eps=1e-6)
self.linear_1 = nn.Linear(text_dim, hidden_size, bias=True)
self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.norm(x)
return self.linear_2(F.silu(self.linear_1(x)))
class LingBotVideoAttention(nn.Module):
def __init__(self, hidden_size, num_heads, norm_eps, qkv_bias, out_bias):
super().__init__()
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.to_q = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
self.to_k = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
self.to_v = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
self.norm_q = LingBotVideoRMSNorm(self.head_dim, norm_eps)
self.norm_k = LingBotVideoRMSNorm(self.head_dim, norm_eps)
self.to_out = nn.Linear(hidden_size, hidden_size, bias=out_bias)
self.attn = Attention(
num_heads=num_heads,
head_size=self.head_dim,
softmax_scale=1.0 / math.sqrt(self.head_dim),
causal=False,
num_kv_heads=num_heads,
)
def forward(
self,
x,
rotary_emb,
attention_mask=None,
packed_indices: dict[str, torch.Tensor] | None = None,
parallel_config=None,
):
B, S, _ = x.shape
if os.environ.get("LINGBOT_FUSED_QKV_LINEAR") == "1":
weight = torch.cat(
(self.to_q.weight, self.to_k.weight, self.to_v.weight),
dim=0,
)
bias = None
if self.to_q.bias is not None:
bias = torch.cat(
(self.to_q.bias, self.to_k.bias, self.to_v.bias),
dim=0,
)
qkv = F.linear(x, weight, bias)
q, k, v = qkv.view(B, S, 3, self.num_heads, self.head_dim).unbind(2)
else:
q = self.to_q(x).unflatten(2, (self.num_heads, self.head_dim))
k = self.to_k(x).unflatten(2, (self.num_heads, self.head_dim))
v = self.to_v(x).unflatten(2, (self.num_heads, self.head_dim))
q = apply_rotary_emb(self.norm_q(q), rotary_emb)
k = apply_rotary_emb(self.norm_k(k), rotary_emb)
if packed_indices is None:
out = self.attn(q, k, v, attn_metadata=AttentionMetadata(attn_mask=attention_mask))
else:
packed_attention_mask = packed_indices.get("attention_mask")
if packed_attention_mask is not None and parallel_config is None:
out = self.attn.sdpa_fallback.forward(
q,
k,
v,
AttentionMetadata(attn_mask=packed_attention_mask),
)
return self.to_out(out.flatten(2, 3).type_as(x))
if flash_attn_varlen_func_v3 is None:
raise RuntimeError(
"A flash attention varlen function is required for packed context parallel attention."
)
if parallel_config is None:
result = flash_attn_varlen_func_v3(
q=q.reshape(-1, self.num_heads, self.head_dim),
k=k.reshape(-1, self.num_heads, self.head_dim),
v=v.reshape(-1, self.num_heads, self.head_dim),
cu_seqlens_q=packed_indices["cu_seqlens_kv"],
cu_seqlens_k=packed_indices["cu_seqlens_kv"],
max_seqlen_q=packed_indices["max_seqlen_in_batch_kv"],
max_seqlen_k=packed_indices["max_seqlen_in_batch_kv"],
causal=False,
)
out = result[0] if isinstance(result, tuple) else result
out = out.reshape(B, S, self.num_heads, self.head_dim)
else:
group = parallel_config.context_parallel_config._ulysses_mesh.get_group()
world_size = dist.get_world_size(group)
local_heads = self.num_heads // world_size
q_global = _all_to_all_split_cat(
q.reshape(B, S, self.num_heads * self.head_dim),
scatter_dim=2,
gather_dim=1,
group=group,
).view(B, S * world_size, local_heads, self.head_dim)
k_global = _all_to_all_split_cat(
k.reshape(B, S, self.num_heads * self.head_dim),
scatter_dim=2,
gather_dim=1,
group=group,
).view(B, S * world_size, local_heads, self.head_dim)
v_global = _all_to_all_split_cat(
v.reshape(B, S, self.num_heads * self.head_dim),
scatter_dim=2,
gather_dim=1,
group=group,
).view(B, S * world_size, local_heads, self.head_dim)
q_flat = q_global.reshape(-1, local_heads, self.head_dim)
k_flat = k_global.reshape(-1, local_heads, self.head_dim)
v_flat = v_global.reshape(-1, local_heads, self.head_dim)
result = flash_attn_varlen_func_v3(
q=q_flat,
k=k_flat,
v=v_flat,
cu_seqlens_q=packed_indices["cu_seqlens_kv"],
cu_seqlens_k=packed_indices["cu_seqlens_kv"],
max_seqlen_q=packed_indices["max_seqlen_in_batch_kv"],
max_seqlen_k=packed_indices["max_seqlen_in_batch_kv"],
causal=False,
)
out_global = result[0] if isinstance(result, tuple) else result
out_global = out_global.reshape(B, S * world_size, local_heads * self.head_dim)
out = _all_to_all_split_cat(
out_global,
scatter_dim=1,
gather_dim=2,
group=group,
).view(B, S, self.num_heads, self.head_dim)
return self.to_out(out.flatten(2, 3).type_as(x))
class LingBotVideoMLP(nn.Module):
def __init__(self, hidden_size, intermediate_size):
super().__init__()
self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
def forward(self, x):
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
class LingBotVideoBlock(nn.Module):
def __init__(
self,
hidden_size,
num_attention_heads,
intermediate_size,
norm_eps,
qkv_bias,
out_bias,
num_experts,
num_experts_per_tok,
moe_intermediate_size,
decoder_sparse_step,
mlp_only_layers,
n_shared_experts,
score_func,
norm_topk_prob,
n_group,
topk_group,
routed_scaling_factor,
layer_idx: int,
):
super().__init__()
self.layer_idx = layer_idx
h = hidden_size
self.scale_shift_table = nn.Parameter(torch.zeros(1, 6 * h))
self.norm1 = LingBotVideoRMSNorm(h, norm_eps)
self.attn = LingBotVideoAttention(h, num_attention_heads, norm_eps, qkv_bias, out_bias)
self.norm_post_attn = LingBotVideoRMSNorm(h, norm_eps)
self.norm2 = LingBotVideoRMSNorm(h, norm_eps)
if num_experts > 0:
raise NotImplementedError(
"LingBotVideoTransformer3DModel in vLLM-Omni currently supports "
"the dense checkpoint only. MoE/fused-expert support belongs in "
"the follow-up MoE PR."
)
self.ffn = LingBotVideoMLP(h, intermediate_size)
self.norm_post_ffn = LingBotVideoRMSNorm(h, norm_eps)
def forward(
self,
x,
temb6,
rotary_emb,
attention_mask=None,
moe_padding_mask=None,
packed_indices: dict[str, torch.Tensor] | None = None,
parallel_config=None,
):
expected_tokens = x.shape[0] * x.shape[1]
if temb6.ndim != 2 or temb6.shape[0] != expected_tokens:
raise ValueError(
"LingBotVideoBlock expects token-level temb6 with shape "
f"(B*S, 6D); got {tuple(temb6.shape)} for hidden states {tuple(x.shape)}."
)
# AdaLN modulation and normalization stay in fp32 for the sensitive path.
mod = temb6.view(x.shape[0], x.shape[1], -1) + self.scale_shift_table.unsqueeze(0)
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=-1)
gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
# AdaLN modulation / norms run in fp32 (sensitive path); cast to the bulk
# compute dtype only at the bf16 Linear boundary. This replaces the old
# ambient autocast, which rounded Linear inputs to bf16 at the same point.
bulk_dtype = self.attn.to_q.weight.dtype
attn_in = (self.norm1(x) * scale_msa + shift_msa).to(bulk_dtype)
attn_out = self.attn(
attn_in,
rotary_emb,
attention_mask,
packed_indices=packed_indices,
parallel_config=parallel_config,
)
x = x + (gate_msa * self.norm_post_attn(attn_out)).to(x.dtype)
ffn_in = (self.norm2(x) * scale_mlp + shift_mlp).to(bulk_dtype)
ffn_out = self.ffn(ffn_in)
ffn_normed = self.norm_post_ffn(ffn_out)
x = x + (gate_mlp * ffn_normed).to(x.dtype)
return x
class LingBotVideoTransformer3DModel(ModelMixin, ConfigMixin):
_supports_gradient_checkpointing = False
_repeated_blocks = ["LingBotVideoBlock"]
_layerwise_offload_blocks_attr = "blocks"
_no_split_modules = ["LingBotVideoBlock"]
_keep_in_fp32_modules = list(LINGBOT_VIDEO_FP32_MODULES)
def to(self, *args, **kwargs):
device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
if dtype is None or dtype == torch.float32:
return super().to(*args, **kwargs)
dtype_is_floating = torch.is_floating_point(torch.empty((), dtype=dtype))
if not dtype_is_floating:
return super().to(*args, **kwargs)
if device is not None:
super().to(device=device, non_blocking=non_blocking)
for name, param in self.named_parameters():
if not torch.is_floating_point(param):
continue
target_dtype = torch.float32 if should_keep_in_fp32(name) else dtype
param.data = param.data.to(dtype=target_dtype, non_blocking=non_blocking)
if param.grad is not None:
param.grad.data = param.grad.data.to(dtype=target_dtype, non_blocking=non_blocking)
for name, buffer in self.named_buffers():
if not torch.is_floating_point(buffer):
continue
target_dtype = torch.float32 if should_keep_in_fp32(name) else dtype
buffer.data = buffer.data.to(dtype=target_dtype, non_blocking=non_blocking)
return self
@register_to_config
def __init__(
self,
patch_size: tuple[int, int, int] = (1, 2, 2),
in_channels: int = 16,
out_channels: int = 16,
hidden_size: int = 2048,
num_attention_heads: int = 16,
depth: int = 24,
intermediate_size: int = 6144,
text_dim: int = 2560,
freq_dim: int = 256,
norm_eps: float = 1e-6,
rope_theta: float = 256.0,
axes_dims: tuple[int, int, int] = (32, 48, 48),
axes_lens: tuple[int, int, int] = (8192, 1024, 1024),
qkv_bias: bool = False,
out_bias: bool = True,
patch_embed_bias: bool = True,
timestep_mlp_bias: bool = True,
num_experts: int = 0,
num_experts_per_tok: int = 8,
moe_intermediate_size: int = 512,
decoder_sparse_step: int = 1,
mlp_only_layers: tuple[int, ...] = (),
n_shared_experts: int | None = None,
score_func: str = "sigmoid",
norm_topk_prob: bool = True,
n_group: int | None = None,
topk_group: int | None = None,
routed_scaling_factor: float = 1.0,
):
super().__init__()
head_dim = hidden_size // num_attention_heads
assert head_dim == sum(axes_dims), f"head_dim {head_dim} != sum(axes_dims) {sum(axes_dims)}"
mlp_only_layers = tuple(mlp_only_layers)
self.patch_embedder = nn.Linear(in_channels * math.prod(patch_size), hidden_size, bias=patch_embed_bias)
self.time_proj = Timesteps(freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
self.time_embedder = TimestepEmbedding(freq_dim, hidden_size, act_fn="silu", sample_proj_bias=timestep_mlp_bias)
self.time_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size))
self.text_embedder = LingBotVideoTextEmbedder(text_dim, hidden_size)
self.rope = LingBotVideoRotaryEmbedding(axes_dims, axes_lens, rope_theta)
self.blocks = nn.ModuleList(
[
LingBotVideoBlock(
hidden_size=hidden_size,
num_attention_heads=num_attention_heads,
intermediate_size=intermediate_size,
norm_eps=norm_eps,
qkv_bias=qkv_bias,
out_bias=out_bias,
num_experts=num_experts,
num_experts_per_tok=num_experts_per_tok,
moe_intermediate_size=moe_intermediate_size,
decoder_sparse_step=decoder_sparse_step,
mlp_only_layers=mlp_only_layers,
n_shared_experts=n_shared_experts,
score_func=score_func,
norm_topk_prob=norm_topk_prob,
n_group=n_group,
topk_group=topk_group,
routed_scaling_factor=routed_scaling_factor,
layer_idx=i,
)
for i in range(depth)
]
)
self.norm_out = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps)
self.norm_out_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size))
self.proj_out = nn.Linear(hidden_size, math.prod(patch_size) * out_channels)
def forward(
self,
hidden_states: torch.Tensor, # (B, C, T, H, W)
timestep: torch.Tensor, # (B,) in [0, 1000](= sigma*1000)
encoder_hidden_states: torch.Tensor, # (B, L, text_dim)
encoder_attention_mask: torch.Tensor | None = None, # (B, L) 1=valid
return_dict: bool = True,
):
B, C, T, H, W = hidden_states.shape
pF, pH, pW = self.config.patch_size
gt, gh, gw = T // pF, H // pH, W // pW
n_video = gt * gh * gw
L = encoder_hidden_states.shape[1]
device = hidden_states.device
if encoder_attention_mask is not None:
text_lens = encoder_attention_mask.sum(dim=-1).long()
else:
text_lens = torch.full((B,), L, dtype=torch.long, device=device)
text_lens_list = [int(v) for v in text_lens.detach().cpu().tolist()]
packed_batch = B > 1
# patchify: token order (f h w), feature order (pf ph pw c) -- matches patchify_and_embed
patch_tokens = hidden_states.reshape(B, C, gt, pF, gh, pH, gw, pW)
patch_tokens = patch_tokens.permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(
B,
n_video,
pF * pH * pW * C,
)
if packed_batch:
x = torch.cat(
[self.patch_embedder(patch_tokens[i : i + 1]) for i in range(B)],
dim=1,
)
else:
x = self.patch_embedder(patch_tokens)
if packed_batch:
text_parts = [
self.text_embedder(encoder_hidden_states[i : i + 1, : text_lens_list[i], :]) for i in range(B)
]
text = torch.cat(text_parts, dim=1)
joint = _cat_interleave(
x,
[n_video] * B,
text,
text_lens_list,
)
else:
text = self.text_embedder(encoder_hidden_states)
joint = torch.cat([x, text], dim=1) # [video; text]
joint_seq_len = joint.shape[1]
# Per-sample RoPE: video t-axis start = real text length of this sample + 1
rotary_parts = [self.rope(make_joint_position_ids(text_lens_list[i], gt, gh, gw, device)) for i in range(B)]
if packed_batch:
rotary = torch.cat(rotary_parts, dim=0).unsqueeze(0)
else:
rotary = torch.stack(rotary_parts, dim=0) # (B, S, head_dim/2) complex64
parallel_config = getattr(self, "_parallel_config", None)
use_packed_attention = parallel_config is not None
attention_mask = None
moe_padding_mask = None
packed_indices = None
has_padding = encoder_attention_mask is not None and bool((text_lens < L).any())
if packed_batch or use_packed_attention:
sample_seq_lens = [n_video + text_len for text_len in text_lens_list]
cu_seqlens = torch.zeros(B + 1, device=device, dtype=torch.int32)
cu_seqlens[1:] = torch.cumsum(
torch.tensor(sample_seq_lens, device=device, dtype=torch.int32),
dim=0,
)
packed_indices = {
"cu_seqlens_kv": cu_seqlens,
"max_seqlen_in_batch_kv": max(sample_seq_lens),
}
if packed_batch and not use_packed_attention:
packed_indices["attention_mask"] = _packed_block_attention_mask(sample_seq_lens, device)
has_padding = False
if has_padding:
key_mask = torch.cat(
[torch.ones(B, n_video, dtype=torch.bool, device=device), encoder_attention_mask.bool()],
dim=1,
)
attention_mask = key_mask[:, None, None, :] # (B,1,1,S) -> SDPA broadcast
moe_padding_mask = key_mask.reshape(-1).float() # (B*S,)
packed_cp = packed_indices is not None and parallel_config is not None
padding_size = 0
if packed_cp:
cp_config = parallel_config.context_parallel_config
cp_world_size = int(getattr(cp_config, "ulysses_degree", getattr(cp_config, "_world_size", 1)))
padding_size = (cp_world_size - (joint_seq_len % cp_world_size)) % cp_world_size
if padding_size:
joint = torch.cat(
[
joint,
torch.zeros(
joint.shape[0],
padding_size,
joint.shape[2],
device=joint.device,
dtype=joint.dtype,
),
],
dim=1,
)
rotary = torch.cat(
[
rotary,
torch.zeros(
rotary.shape[0],
padding_size,
rotary.shape[2],
device=rotary.device,
dtype=rotary.dtype,
),
],
dim=1,
)
if packed_indices is None:
raise RuntimeError("packed_indices must be initialized for packed context parallel.")
packed_indices["cu_seqlens_kv"] = torch.cat(
[
packed_indices["cu_seqlens_kv"],
packed_indices["cu_seqlens_kv"][-1:] + padding_size,
],
dim=0,
)
packed_indices["max_seqlen_in_batch_kv"] = max(
int(packed_indices["max_seqlen_in_batch_kv"]),
int(padding_size),
)
joint_seq_len = joint.shape[1]
timestep_for_embed = timestep.float()
timestep_proj = self.time_proj(timestep_for_embed)
t_emb = self.time_embedder(timestep_proj) # (B, D)
if packed_batch:
temb_input = torch.cat(
[t_emb[i : i + 1].unsqueeze(1).expand(1, n_video + text_lens_list[i], -1) for i in range(B)],
dim=1,
)
if padding_size:
temb_input = torch.cat(
[
temb_input,
torch.zeros(
temb_input.shape[0],
padding_size,
temb_input.shape[2],
device=temb_input.device,
dtype=temb_input.dtype,
),
],
dim=1,
)
temb6 = self.time_modulation(temb_input.reshape(joint_seq_len, -1))
temb6 = temb6.reshape(1, joint_seq_len, -1)
else:
temb_input = t_emb.unsqueeze(1).expand(B, joint_seq_len, -1) # (B, S, D)
temb6 = self.time_modulation(temb_input.reshape(B * joint_seq_len, -1))
temb6 = temb6.reshape(B, joint_seq_len, -1) # (B, S, 6D)
temb6 = temb6.reshape(temb6.shape[0] * temb6.shape[1], -1)
for block in self.blocks:
joint = block(
joint,
temb6,
rotary,
attention_mask,
moe_padding_mask,
packed_indices=packed_indices,
parallel_config=parallel_config,
)
final_mod = self.norm_out_modulation(temb_input.reshape(joint.shape[0] * joint.shape[1], -1))
shift, scale = final_mod.reshape(joint.shape[0], joint.shape[1], -1).chunk(2, dim=-1)
final_hidden = self.norm_out(joint) * (1.0 + scale) + shift
projected = self.proj_out(final_hidden.to(self.proj_out.weight.dtype))
if packed_cp:
if padding_size:
projected = projected[:, :-padding_size, :]
if packed_batch:
split_lengths: list[int] = []
for text_len in text_lens_list:
split_lengths.extend([n_video, text_len])
parts = torch.split(projected, split_lengths, dim=1)
x = torch.cat(parts[::2], dim=1).reshape(B, n_video, -1)
else:
x = projected[:, :n_video]
# unpatchify (matches the rearrange in postprocess)
Cout = self.config.out_channels
x = x.reshape(B, gt, gh, gw, pF, pH, pW, Cout)
x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(B, Cout, T, H, W)
if not return_dict:
return (x,)
return Transformer2DModelOutput(sample=x)
@@ -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 the dense LingBot-Video checkpoint.
This PR supports the dense checkpoint only. MoE and fused-expert kernels are
intentionally left for the follow-up MoE PR.
"""
supports_step_execution: ClassVar[bool] = False
_dit_modules: ClassVar[list[str]] = ["transformer"]
_encoder_modules: ClassVar[list[str]] = ["text_encoder"]
_vae_modules: ClassVar[list[str]] = ["vae"]
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__()
del prefix
self.od_config = od_config
self.device = get_local_device()
self.vae_scale_factor_temporal = 4
self.vae_scale_factor_spatial = 8
self.token_length = TOKEN_LENGTH
self.hidden_state_skip_layer = HIDDEN_STATE_SKIP_LAYER
self.prompt_template = PROMPT_TEMPLATE
self._crop_start: int | None = None
model = od_config.model
local_files_only = os.path.exists(model)
dtype = getattr(od_config, "dtype", torch.bfloat16)
model_config = getattr(od_config, "model_config", None) or {}
transformer_dtype = _dtype_from_name(model_config.get("transformer_dtype"), dtype)
text_encoder_dtype = _dtype_from_name(model_config.get("text_encoder_dtype"), dtype)
vae_dtype = _dtype_from_name(model_config.get("vae_dtype"), torch.float32)
transformer_subfolder = str(model_config.get("transformer_subfolder", "transformer"))
text_encoder_subfolder = str(model_config.get("text_encoder_subfolder", "text_encoder"))
processor_subfolder = str(model_config.get("processor_subfolder", "processor"))
vae_subfolder = str(model_config.get("vae_subfolder", "vae"))
scheduler_subfolder = str(model_config.get("scheduler_subfolder", "scheduler"))
self.transformer = LingBotVideoTransformer3DModel.from_pretrained(
model,
subfolder=transformer_subfolder,
torch_dtype=transformer_dtype,
local_files_only=local_files_only,
).to(self.device)
text_encoder_kwargs: dict[str, Any] = {
"dtype": text_encoder_dtype,
"local_files_only": local_files_only,
}
self.text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
model,
subfolder=text_encoder_subfolder,
**text_encoder_kwargs,
).to(self.device)
self.processor = Qwen3VLProcessor.from_pretrained(
model,
subfolder=processor_subfolder,
local_files_only=local_files_only,
)
self.vae = AutoencoderKLWan.from_pretrained(
model,
subfolder=vae_subfolder,
torch_dtype=vae_dtype,
local_files_only=local_files_only,
).to(self.device)
self.scheduler = FlowUniPCMultistepScheduler.from_pretrained(
model,
subfolder=scheduler_subfolder,
local_files_only=local_files_only,
)
self.set_progress_bar_config(disable=bool(model_config.get("quiet_progress", True)))
self.default_negative_prompt = DEFAULT_NEGATIVE_PROMPT
def to(self, *args, **kwargs):
device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
super().to(*args, **kwargs)
if device is not None:
self.device = torch.device(device)
self.transformer.to(device=self.device, non_blocking=non_blocking)
self.text_encoder.to(device=self.device, non_blocking=non_blocking)
self.vae.to(device=self.device, non_blocking=non_blocking)
if dtype is not None:
self.transformer.to(dtype=dtype, non_blocking=non_blocking)
return self
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
consumed = list(weights)
if consumed:
raise RuntimeError(
f"{self.__class__.__name__}.load_weights received {len(consumed)} weight tensors; "
"LingBot components are loaded directly from subfolders during __init__."
)
return {name for name, _ in self.named_parameters()}
@staticmethod
def check_inputs(height: int, width: int, num_frames: int) -> None:
if num_frames != 1 and (num_frames - 1) % 4 != 0:
raise ValueError(f"`num_frames` must be 1 or 4n+1, got {num_frames}.")
if height % 16 != 0 or width % 16 != 0:
raise ValueError(f"`height` and `width` must be multiples of 16, got {height}x{width}.")
@staticmethod
def apply_text_to_template(text: str, template: str = PROMPT_TEMPLATE) -> str:
return template.format(text)
def _compute_crop_start(self) -> int:
if self._crop_start is None:
marker = "<|USER_INPUT_MARKER|>"
marked = self.prompt_template.format(marker)
marker_pos = marked.find(marker)
if marker_pos < 0:
self._crop_start = 0
else:
prefix = self.processor(
text=marked[:marker_pos],
images=None,
videos=None,
return_tensors="pt",
)
self._crop_start = int(prefix["input_ids"].shape[1])
return self._crop_start
def _build_prompt_inputs(self, prompt: str | list[str]):
prompts = [prompt] if isinstance(prompt, str) else list(prompt)
texts = [self.apply_text_to_template(text, self.prompt_template) for text in prompts]
return self.processor(
text=texts,
images=None,
videos=None,
do_resize=False,
truncation=True,
max_length=self.token_length,
padding="longest",
return_tensors="pt",
)
@torch.no_grad()
def encode_prompt(
self,
prompt: str | list[str],
*,
device: str | torch.device | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
device = torch.device(device) if device is not None else self.device
inputs = self._build_prompt_inputs(prompt).to(device)
outputs = self.text_encoder(
**inputs,
output_hidden_states=self.hidden_state_skip_layer is not None,
)
if self.hidden_state_skip_layer is not None:
prompt_embeds = outputs.hidden_states[-(self.hidden_state_skip_layer + 1)]
else:
prompt_embeds = outputs.last_hidden_state
prompt_mask = inputs["attention_mask"]
crop_start = self._compute_crop_start()
if crop_start > 0:
prompt_embeds = prompt_embeds[:, crop_start:]
prompt_mask = prompt_mask[:, crop_start:]
if prompt_embeds.shape[0] == 1:
true_len = int(prompt_mask[0].sum().item())
prompt_embeds = prompt_embeds[:, :true_len]
prompt_mask = prompt_mask[:, :true_len]
return prompt_embeds, prompt_mask
def prepare_latents(
self,
num_frames: int,
height: int,
width: int,
generator: torch.Generator | None,
latents: torch.Tensor | None,
device: torch.device,
) -> torch.Tensor:
latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
latent_height = height // self.vae_scale_factor_spatial
latent_width = width // self.vae_scale_factor_spatial
shape = (1, self.transformer.config.in_channels, latent_frames, latent_height, latent_width)
if latents is None:
return randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
if tuple(latents.shape) != shape:
raise ValueError(f"`latents` shape must be {shape}, got {tuple(latents.shape)}.")
return latents.to(device=device, dtype=torch.float32)
def _dit_latent_to_vae(self, latents: torch.Tensor) -> torch.Tensor:
mean = torch.tensor(self.vae.config.latents_mean, device=latents.device, dtype=torch.float32)
std_inv = 1.0 / torch.tensor(self.vae.config.latents_std, device=latents.device, dtype=torch.float32)
mean = mean.view(1, -1, 1, 1, 1)
std_inv = std_inv.view(1, -1, 1, 1, 1)
return latents.float() / std_inv + mean
@torch.no_grad()
def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
vae_device = _module_device(self.vae)
vae_dtype = _module_dtype(self.vae)
vae_latents = self._dit_latent_to_vae(latents).to(device=vae_device, dtype=torch.float32)
if vae_latents.ndim == 5:
vae_latents = vae_latents.contiguous(memory_format=torch.channels_last_3d)
autocast_dtype = (
vae_dtype if vae_device.type == "cuda" and vae_dtype in {torch.bfloat16, torch.float16} else None
)
with torch.autocast("cuda", dtype=autocast_dtype or torch.bfloat16, enabled=autocast_dtype is not None):
decoded = self.vae.decode(vae_latents)
frames = decoded[0] if isinstance(decoded, tuple) else decoded.sample
frames = frames.float().clamp_(-1, 1)
frames = (frames + 1.0) / 2.0
frames = frames.permute(0, 2, 3, 4, 1).cpu()
return frames[0]
@torch.no_grad()
def _generate(
self,
*,
prompt: str,
negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
height: int = 480,
width: int = 480,
num_frames: int = 81,
num_inference_steps: int = 40,
guidance_scale: float = 6.0,
shift: float = 3.0,
generator: torch.Generator | None = None,
latents: torch.Tensor | None = None,
prompt_embeds: torch.Tensor | None = None,
prompt_mask: torch.Tensor | None = None,
negative_prompt_embeds: torch.Tensor | None = None,
negative_prompt_mask: torch.Tensor | None = None,
output_type: str = "pt",
cfg_parallel_group: Any | None = None,
batch_cfg: bool = False,
null_cond_clone_zero: bool = False,
t_thresh: float | None = None,
refiner_sigma_tail_steps: int = LOW_NOISE_TAIL_V1_DEFAULT_STEPS,
offload_vae_during_denoise: bool = False,
**extra_args,
) -> torch.Tensor:
del extra_args
self.check_inputs(height, width, num_frames)
device = self.device
do_cfg = guidance_scale > 1.0
effective_batch_cfg = bool(batch_cfg)
cfg_parallel = cfg_parallel_group is not None
cfg_parallel_rank = 0
if cfg_parallel:
if not dist.is_available() or not dist.is_initialized():
raise ValueError("`cfg_parallel_group` requires an initialized process group.")
if effective_batch_cfg:
raise ValueError("`cfg_parallel_group` and `batch_cfg` are mutually exclusive.")
if not do_cfg:
raise ValueError("CFG parallel requires `guidance_scale > 1.0`.")
cfg_parallel_rank = dist.get_rank(cfg_parallel_group)
cfg_parallel_world_size = dist.get_world_size(cfg_parallel_group)
if cfg_parallel_world_size != 2:
raise ValueError(f"CFG parallel currently requires exactly 2 ranks, got {cfg_parallel_world_size}.")
if prompt_embeds is not None:
if prompt_mask is None:
raise ValueError("`prompt_mask` is required when `prompt_embeds` is provided.")
prompt_embeds = prompt_embeds.to(device=device)
prompt_mask = prompt_mask.to(device=device)
if negative_prompt_embeds is not None:
if negative_prompt_mask is None:
raise ValueError("`negative_prompt_mask` is required when `negative_prompt_embeds` is provided.")
negative_prompt_embeds = negative_prompt_embeds.to(device=device)
negative_prompt_mask = negative_prompt_mask.to(device=device)
negative_embeds = None
negative_mask = None
if cfg_parallel and cfg_parallel_rank == 1:
if negative_prompt_embeds is not None:
negative_embeds, negative_mask = negative_prompt_embeds, negative_prompt_mask
else:
negative_embeds, negative_mask = self.encode_prompt(negative_prompt, device=device)
prompt_embeds = prompt_mask = None
else:
if prompt_embeds is None:
prompt_embeds, prompt_mask = self.encode_prompt(prompt, device=device)
if do_cfg and not cfg_parallel:
if null_cond_clone_zero:
negative_embeds = torch.zeros_like(prompt_embeds)
negative_mask = prompt_mask.clone()
elif negative_prompt_embeds is not None:
negative_embeds, negative_mask = negative_prompt_embeds, negative_prompt_mask
else:
negative_embeds, negative_mask = self.encode_prompt(negative_prompt, device=device)
latents = self.prepare_latents(num_frames, height, width, generator, latents, device)
sigmas = _compute_refiner_sigmas(
sigma_max=float(self.scheduler.sigma_max),
sigma_min=float(self.scheduler.sigma_min),
num_inference_steps=num_inference_steps,
shift=shift,
t_thresh=t_thresh,
tail_steps=refiner_sigma_tail_steps,
)
if sigmas is None:
self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
else:
self.scheduler.set_timesteps(int(sigmas.shape[0]), device=device, sigmas=sigmas, shift=1.0)
transformer_dtype = _module_dtype(self.transformer)
vae_restore_device: torch.device | None = None
vae_offloaded = False
if offload_vae_during_denoise and output_type != "latent":
vae_device = _module_device(self.vae)
if vae_device.type == "cuda":
self.vae.to("cpu")
torch.accelerator.empty_cache()
vae_restore_device = vae_device
vae_offloaded = True
cfg_latent_src = _group_global_rank(cfg_parallel_group, 0)
cfg_uncond_src = _group_global_rank(cfg_parallel_group, 1)
for timestep in self.progress_bar(self.scheduler.timesteps):
if cfg_parallel:
dist.broadcast(latents, src=cfg_latent_src, group=cfg_parallel_group)
timestep_batch = _transformer_timestep(timestep, transformer_dtype).expand(1).to(device)
latent_model_input = latents
if cfg_parallel:
if cfg_parallel_rank == 0:
branch_embeds = prompt_embeds
branch_mask = prompt_mask
else:
branch_embeds = negative_embeds
branch_mask = negative_mask
if branch_embeds is None:
raise RuntimeError("CFG branch embeddings were not initialized.")
with _transformer_autocast(device, transformer_dtype):
branch_noise_pred = self.transformer(
latent_model_input,
timestep_batch,
branch_embeds.to(transformer_dtype),
encoder_attention_mask=branch_mask,
return_dict=False,
)[0].float()
if cfg_parallel_rank == 0:
noise_pred = branch_noise_pred
noise_pred_uncond = torch.empty_like(noise_pred)
else:
noise_pred_uncond = branch_noise_pred
dist.broadcast(noise_pred_uncond, src=cfg_uncond_src, group=cfg_parallel_group)
if cfg_parallel_rank != 0:
continue
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
else:
if prompt_embeds is None:
raise RuntimeError("Prompt embeddings were not initialized.")
prompt_model_input = prompt_embeds.to(transformer_dtype)
if do_cfg and effective_batch_cfg:
if negative_embeds is None or negative_mask is None:
raise RuntimeError("Negative embeddings were not initialized for CFG.")
cfg_embeds, cfg_mask = _batch_cfg_prompt_inputs(
prompt_model_input,
prompt_mask,
negative_embeds.to(transformer_dtype),
negative_mask,
null_cond_clone_zero=False,
)
cfg_latents = torch.cat([latent_model_input, latent_model_input], dim=0)
cfg_timesteps = torch.cat([timestep_batch, timestep_batch], dim=0)
with _transformer_autocast(device, transformer_dtype):
noise_batched = self.transformer(
cfg_latents,
cfg_timesteps,
cfg_embeds,
encoder_attention_mask=cfg_mask,
return_dict=False,
)[0].float()
noise_pred, noise_pred_uncond = noise_batched.chunk(2, dim=0)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
else:
with _transformer_autocast(device, transformer_dtype):
noise_pred = self.transformer(
latent_model_input,
timestep_batch,
prompt_model_input,
encoder_attention_mask=prompt_mask,
return_dict=False,
)[0].float()
if do_cfg and not effective_batch_cfg:
if negative_embeds is None or negative_mask is None:
raise RuntimeError("Negative embeddings were not initialized for CFG.")
with _transformer_autocast(device, transformer_dtype):
noise_pred_uncond = self.transformer(
latent_model_input,
timestep_batch,
negative_embeds.to(transformer_dtype),
encoder_attention_mask=negative_mask,
return_dict=False,
)[0].float()
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
latents = self.scheduler.step(noise_pred, timestep, latents, return_dict=False, generator=generator)[0]
if cfg_parallel:
dist.barrier(group=cfg_parallel_group)
if cfg_parallel_rank != 0:
return latents if output_type == "latent" else []
if output_type == "latent":
return latents
if output_type in {"pt", "np"}:
if vae_offloaded and vae_restore_device is not None:
self.vae.to(device=vae_restore_device)
torch.accelerator.empty_cache()
return self._decode_latents(latents)
raise ValueError(f"Unsupported output_type: {output_type}")
@torch.inference_mode()
def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
if req.num_reqs != 1:
raise ValueError(f"LingBotVideoPipeline only supports one request per batch, got {req.num_reqs}.")
request = req.requests[0]
prompt, prompt_negative = _extract_prompt(request)
sampling = request.sampling_params
extra_args = dict(sampling.extra_args or {})
generator = sampling.generator
if isinstance(generator, list):
generator = generator[0] if generator else None
if generator is None and sampling.seed is not None:
generator = torch.Generator(device=self.device).manual_seed(int(sampling.seed))
height = sampling.height if sampling.height is not None else extra_args.pop("height", 480)
width = sampling.width if sampling.width is not None else extra_args.pop("width", 480)
num_frames = sampling.num_frames or extra_args.pop("num_frames", 81)
num_inference_steps = (
sampling.num_inference_steps
if sampling.num_inference_steps is not None
else extra_args.pop("num_inference_steps", 40)
)
guidance_scale = (
sampling.guidance_scale
if sampling.guidance_scale_provided or sampling.guidance_scale > 0
else extra_args.pop("guidance_scale", 6.0)
)
shift = extra_args.pop(
"shift",
extra_args.pop("flow_shift", getattr(self.od_config, "flow_shift", None) or 3.0),
)
negative_prompt = extra_args.pop("negative_prompt", prompt_negative or self.default_negative_prompt)
output_type = (
sampling.output_type or getattr(self.od_config, "output_type", None) or extra_args.pop("output_type", "pt")
)
if output_type not in {"pt", "np", "latent"}:
output_type = "pt"
sampling.output_type = output_type
frames = self._generate(
prompt=prompt,
negative_prompt=negative_prompt,
height=height,
width=width,
num_frames=num_frames,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
shift=shift,
generator=generator,
latents=sampling.latents,
output_type=output_type,
batch_cfg=bool(extra_args.pop("batch_cfg", False)),
null_cond_clone_zero=bool(extra_args.pop("null_cond_clone_zero", False)),
offload_vae_during_denoise=bool(extra_args.pop("offload_vae_during_denoise", False)),
t_thresh=extra_args.pop("t_thresh", None),
refiner_sigma_tail_steps=int(extra_args.pop("refiner_sigma_tail_steps", LOW_NOISE_TAIL_V1_DEFAULT_STEPS)),
)
return DiffusionOutput(output=frames)
+6
View File
@@ -256,6 +256,11 @@ _DIFFUSION_MODELS = {
"pipeline_hunyuan_video_1_5_i2v", "pipeline_hunyuan_video_1_5_i2v",
"HunyuanVideo15I2VPipeline", "HunyuanVideo15I2VPipeline",
), ),
"LingBotVideoPipeline": (
"lingbot_video",
"pipeline_lingbot_video",
"LingBotVideoPipeline",
),
"MagiHumanPipeline": ( "MagiHumanPipeline": (
"magi_human", "magi_human",
"pipeline_magi_human", "pipeline_magi_human",
@@ -526,6 +531,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
"Flux2Pipeline": "get_flux2_post_process_func", "Flux2Pipeline": "get_flux2_post_process_func",
"HunyuanVideo15Pipeline": "get_hunyuan_video_15_post_process_func", "HunyuanVideo15Pipeline": "get_hunyuan_video_15_post_process_func",
"HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_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", "MagiHumanPipeline": "get_magi_human_post_process_func",
"OmniVoicePipeline": "get_omnivoice_post_process_func", "OmniVoicePipeline": "get_omnivoice_post_process_func",
"DreamIDOmniPipeline": "get_dreamid_omni_post_process_func", "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func",
+16
View File
@@ -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",
}
)
+4
View File
@@ -34,6 +34,7 @@ from vllm_omni.model_extras.helios import (
HELIOS_EXTRA_BODY_PARAMS, HELIOS_EXTRA_BODY_PARAMS,
HELIOS_EXTRA_OUTPUT_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 ( from vllm_omni.model_extras.magi_human import (
MAGI_HUMAN_EXTRA_BODY_PARAMS, MAGI_HUMAN_EXTRA_BODY_PARAMS,
MAGI_HUMAN_EXTRA_OUTPUT_PARAMS, MAGI_HUMAN_EXTRA_OUTPUT_PARAMS,
@@ -165,6 +166,9 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
"extra_body_params": HELIOS_EXTRA_BODY_PARAMS, "extra_body_params": HELIOS_EXTRA_BODY_PARAMS,
"extra_output_params": HELIOS_EXTRA_OUTPUT_PARAMS, "extra_output_params": HELIOS_EXTRA_OUTPUT_PARAMS,
}, },
"LingBotVideoPipeline": {
"extra_body_params": LINGBOT_VIDEO_EXTRA_BODY_PARAMS,
},
"WanVACEPipeline": { "WanVACEPipeline": {
"extra_body_params": VACE_EXTRA_BODY_PARAMS, "extra_body_params": VACE_EXTRA_BODY_PARAMS,
"extra_output_params": VACE_EXTRA_OUTPUT_PARAMS, "extra_output_params": VACE_EXTRA_OUTPUT_PARAMS,