Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
603a3ae67b | ||
|
|
ebd0e024e9 | ||
|
|
7d59b818d0 | ||
|
|
6f4e53553a | ||
|
|
0e6490d929 |
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()
|
||||
@@ -46,6 +46,7 @@ th {
|
||||
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||
| `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `LingBotVideoPipeline` | LingBot-Video dense T2V | `robbyant/lingbot-video-dense-1.3b` | ✅︎ | | | |
|
||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
|
||||
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
|
||||
|
||||
@@ -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()
|
||||
@@ -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) |
|
||||
| [`MammothModa2/MammothModa2-Preview.md`](./MammothModa2/MammothModa2-Preview.md) | Text-to-image with the shared offline image example (AR → DiT) | 1x L40S 48GB / 1x ≥40GB GPU |
|
||||
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
|
||||
| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense text-to-video serving | 1x L20X / A100 / H100 |
|
||||
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
|
||||
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
|
||||
| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
|
||||
|
||||
@@ -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",
|
||||
"HunyuanVideo15Pipeline",
|
||||
"HunyuanVideo15ImageToVideoPipeline",
|
||||
"LingBotVideoPipeline",
|
||||
"MagiHumanPipeline",
|
||||
"OmniVoicePipeline",
|
||||
"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)
|
||||
@@ -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",
|
||||
@@ -526,6 +531,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",
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
)
|
||||
@@ -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,
|
||||
@@ -165,6 +166,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,
|
||||
|
||||
Reference in new issue
Block a user