284 lines
11 KiB
Python
284 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Generate a LingBot-World v2 video from an image and camera trajectory.
|
|
|
|
The official checkpoint is licensed separately under CC BY-NC-SA and is
|
|
restricted to non-commercial use. This vLLM-Omni integration remains
|
|
Apache-2.0 licensed.
|
|
|
|
Example:
|
|
python examples/offline_inference/diffusion/lingbot_world_v2.py \
|
|
--prompt "The camera moves slowly forward through the scene." \
|
|
--image /path/to/first_frame.png \
|
|
--action-dir /path/to/actions/forward \
|
|
--output lingbot_world_v2.mp4
|
|
|
|
``--action-dir`` must contain ``poses.npy`` and ``intrinsics.npy``. The
|
|
example configures its parent as the trusted action root and sends the
|
|
contained directory through ``extra_args["action_path"]``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import math
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
|
|
_MAX_SEQUENCE_LENGTH = 512
|
|
_MAX_PIXEL_AREA = 480 * 832
|
|
_MAX_SOURCE_CAMERA_FRAMES = 4096
|
|
_TEMPORAL_COMPRESSION = 4
|
|
_LATENT_FRAMES_PER_BLOCK = 3
|
|
_SPATIAL_ALIGNMENT = 16
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LingBotPaths:
|
|
image: Path
|
|
action_dir: Path
|
|
action_root: Path
|
|
action_relative: Path
|
|
camera_frames: int
|
|
output: Path
|
|
|
|
|
|
def _validate_camera_arrays(action_dir: Path) -> int:
|
|
"""Validate the bounded camera contract before starting the engine."""
|
|
|
|
try:
|
|
poses = np.load(action_dir / "poses.npy", mmap_mode="r", allow_pickle=False)
|
|
intrinsics = np.load(action_dir / "intrinsics.npy", mmap_mode="r", allow_pickle=False)
|
|
except (OSError, TypeError, ValueError):
|
|
raise ValueError("Camera files must be valid numeric poses.npy and intrinsics.npy arrays.") from None
|
|
if poses.ndim != 3 or poses.shape[1:] != (4, 4):
|
|
raise ValueError("poses.npy must have shape [frames, 4, 4].")
|
|
if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,):
|
|
raise ValueError("intrinsics.npy must have shape [frames, 4].")
|
|
if poses.shape[0] != intrinsics.shape[0]:
|
|
raise ValueError("poses.npy and intrinsics.npy must contain the same frame count.")
|
|
if not 1 <= poses.shape[0] <= _MAX_SOURCE_CAMERA_FRAMES:
|
|
raise ValueError("Camera arrays must contain between 1 and 4096 source frames.")
|
|
if poses.dtype.kind not in "fiu" or intrinsics.dtype.kind not in "fiu":
|
|
raise ValueError("Camera arrays must contain real numeric values.")
|
|
if not np.isfinite(poses).all() or not np.isfinite(intrinsics).all():
|
|
raise ValueError("Camera arrays must contain only finite values.")
|
|
return int(poses.shape[0])
|
|
|
|
|
|
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate a LingBot-World v2 video from a first-frame image and camera trajectory."
|
|
)
|
|
parser.add_argument("--model", default=_MODEL, help="Official Hugging Face model ID or local checkpoint path.")
|
|
parser.add_argument("--prompt", required=True, help="Text description of the scene and desired motion.")
|
|
parser.add_argument("--image", required=True, help="Path to the first-frame image.")
|
|
parser.add_argument(
|
|
"--action-dir",
|
|
required=True,
|
|
help="Camera action directory containing poses.npy and intrinsics.npy.",
|
|
)
|
|
parser.add_argument("--height", type=int, default=480, help="Output height; must be divisible by 16.")
|
|
parser.add_argument("--width", type=int, default=832, help="Output width; must be divisible by 16.")
|
|
parser.add_argument(
|
|
"--num-frames",
|
|
type=int,
|
|
default=9,
|
|
help="Raw output frame count. Nine frames are one three-latent-frame causal block.",
|
|
)
|
|
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
|
parser.add_argument(
|
|
"--tensor-parallel-size",
|
|
type=int,
|
|
default=1,
|
|
help="Number of GPUs used for tensor parallelism inside the DiT.",
|
|
)
|
|
parser.add_argument("--ulysses-degree", type=int, default=1, help="Pure Ulysses sequence parallel degree.")
|
|
parser.add_argument("--flow-shift", type=float, default=5.0, help="Positive FlowUniPC scheduler shift.")
|
|
parser.add_argument("--fps", type=int, default=16, help="Frames per second in the exported MP4.")
|
|
parser.add_argument("--output", default="lingbot_world_v2.mp4", help="Output MP4 path.")
|
|
parser.add_argument(
|
|
"--enforce-eager",
|
|
action="store_true",
|
|
help="Disable compilation and force eager execution.",
|
|
)
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def resolve_cli_paths(args: argparse.Namespace) -> LingBotPaths:
|
|
"""Canonicalize local inputs and construct a safe trusted-root mapping."""
|
|
|
|
image = Path(args.image).expanduser().resolve()
|
|
if not image.is_file():
|
|
raise ValueError("--image must point to an existing file.")
|
|
|
|
action_dir = Path(args.action_dir).expanduser().resolve()
|
|
if not action_dir.is_dir():
|
|
raise ValueError("--action-dir must point to an existing directory.")
|
|
if not all((action_dir / name).is_file() for name in ("poses.npy", "intrinsics.npy")):
|
|
raise ValueError("--action-dir must contain poses.npy and intrinsics.npy.")
|
|
camera_frames = _validate_camera_arrays(action_dir)
|
|
|
|
action_root = action_dir.parent.resolve()
|
|
try:
|
|
action_relative = action_dir.relative_to(action_root)
|
|
except ValueError:
|
|
raise ValueError("--action-dir must be contained by its canonical trusted action root.") from None
|
|
if not action_relative.parts:
|
|
raise ValueError("--action-dir must name a directory contained by a trusted action root.")
|
|
|
|
output = Path(args.output).expanduser().resolve()
|
|
return LingBotPaths(
|
|
image=image,
|
|
action_dir=action_dir,
|
|
action_root=action_root,
|
|
action_relative=action_relative,
|
|
camera_frames=camera_frames,
|
|
output=output,
|
|
)
|
|
|
|
|
|
def _positive_finite(value: float, flag: str) -> float:
|
|
result = float(value)
|
|
if not math.isfinite(result) or result <= 0:
|
|
raise ValueError(f"{flag} must be a positive finite number.")
|
|
return result
|
|
|
|
|
|
def build_omni_kwargs(
|
|
args: argparse.Namespace,
|
|
paths: LingBotPaths,
|
|
) -> dict[str, Any]:
|
|
"""Build Omni configuration without overriding checkpoint class discovery."""
|
|
|
|
if args.tensor_parallel_size <= 0:
|
|
raise ValueError("--tensor-parallel-size must be a positive integer.")
|
|
if args.ulysses_degree <= 0:
|
|
raise ValueError("--ulysses-degree must be a positive integer.")
|
|
flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
|
|
model_path = Path(args.model).expanduser()
|
|
model = str(model_path.resolve()) if model_path.exists() else args.model
|
|
return {
|
|
"model": model,
|
|
"flow_shift": flow_shift,
|
|
"tensor_parallel_size": args.tensor_parallel_size,
|
|
"ulysses_degree": args.ulysses_degree,
|
|
"enforce_eager": args.enforce_eager,
|
|
"model_config": {"lingbot_action_root": str(paths.action_root)},
|
|
}
|
|
|
|
|
|
def build_request(
|
|
args: argparse.Namespace,
|
|
paths: LingBotPaths,
|
|
) -> tuple[dict[str, Any], dict[str, Any]]:
|
|
"""Build the prompt and fixed v1 causal-DMD sampling contract."""
|
|
|
|
prompt_text = args.prompt.strip()
|
|
if not prompt_text:
|
|
raise ValueError("--prompt must contain non-whitespace text.")
|
|
if args.height <= 0 or args.width <= 0:
|
|
raise ValueError("--height and --width must be positive integers.")
|
|
if args.height % _SPATIAL_ALIGNMENT or args.width % _SPATIAL_ALIGNMENT:
|
|
raise ValueError("--height and --width must each be divisible by 16.")
|
|
if args.height * args.width > _MAX_PIXEL_AREA:
|
|
raise ValueError("--height * --width must not exceed 480 * 832 pixels.")
|
|
if args.num_frames <= 0:
|
|
raise ValueError("--num-frames must be a positive integer.")
|
|
if args.num_frames > paths.camera_frames:
|
|
raise ValueError("Camera arrays must contain at least --num-frames frames.")
|
|
if (args.num_frames - 1) % _TEMPORAL_COMPRESSION:
|
|
raise ValueError("--num-frames must satisfy (num_frames - 1) divisible by 4.")
|
|
latent_frames = (args.num_frames - 1) // _TEMPORAL_COMPRESSION + 1
|
|
if latent_frames % _LATENT_FRAMES_PER_BLOCK:
|
|
raise ValueError("--num-frames must map to whole three-frame latent blocks.")
|
|
if args.fps <= 0:
|
|
raise ValueError("--fps must be a positive integer.")
|
|
flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
|
|
|
|
prompt = {
|
|
"prompt": prompt_text,
|
|
"multi_modal_data": {"image": str(paths.image)},
|
|
}
|
|
sampling_kwargs = {
|
|
"height": args.height,
|
|
"width": args.width,
|
|
"num_frames": args.num_frames,
|
|
"max_sequence_length": _MAX_SEQUENCE_LENGTH,
|
|
"seed": args.seed,
|
|
"fps": args.fps,
|
|
"extra_args": {
|
|
"action_path": paths.action_relative.as_posix(),
|
|
"flow_shift": flow_shift,
|
|
},
|
|
}
|
|
return prompt, sampling_kwargs
|
|
|
|
|
|
def extract_video_array(outputs: Any) -> np.ndarray:
|
|
"""Unwrap one standard Omni diffusion video into ``[frames, H, W, C]``."""
|
|
|
|
if not isinstance(outputs, list) or len(outputs) != 1:
|
|
raise ValueError("Expected exactly one Omni output for the single LingBot request.")
|
|
output = outputs[0]
|
|
request_output = getattr(output, "request_output", None)
|
|
if request_output is not None and getattr(request_output, "images", None):
|
|
output = request_output
|
|
|
|
images = getattr(output, "images", None)
|
|
if not images:
|
|
multimodal_output = getattr(output, "multimodal_output", None) or {}
|
|
images = multimodal_output.get("video")
|
|
if images is None:
|
|
raise ValueError("No video frames found in Omni output.")
|
|
if isinstance(images, list) and len(images) == 1:
|
|
images = images[0]
|
|
|
|
video: np.ndarray = np.asarray(images)
|
|
if video.ndim == 5:
|
|
if video.shape[0] != 1:
|
|
raise ValueError(f"Expected one generated video, got batch shape {video.shape}.")
|
|
video = video[0]
|
|
if video.ndim != 4 or video.shape[-1] not in (3, 4):
|
|
raise ValueError(f"Expected video shape [frames, height, width, channels], got {video.shape}.")
|
|
return video
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> Path:
|
|
args = parse_args(argv)
|
|
paths = resolve_cli_paths(args)
|
|
prompt, sampling_kwargs = build_request(args, paths)
|
|
paths.output.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Keep vLLM imports below pure CLI/path validation so ``--help`` and the
|
|
# helper tests work on machines without the native vLLM runtime.
|
|
from diffusers.utils import export_to_video
|
|
|
|
from vllm_omni.entrypoints.omni import Omni
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
omni_kwargs = build_omni_kwargs(args, paths)
|
|
omni = Omni(**omni_kwargs)
|
|
try:
|
|
outputs = omni.generate(
|
|
prompt,
|
|
OmniDiffusionSamplingParams(**sampling_kwargs),
|
|
use_tqdm=False,
|
|
)
|
|
finally:
|
|
omni.close()
|
|
|
|
frames = extract_video_array(outputs)
|
|
export_to_video(frames, str(paths.output), fps=args.fps)
|
|
print(f"Saved {frames.shape[0]} frames to {paths.output}")
|
|
return paths.output
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|