Merge branch 'hunyuan-gamecraft-v2-port' of https://github.com/MihirJagtap/FastVideo into pr/1071

This commit is contained in:
Mihir Jagtap
2026-02-12 13:09:40 -08:00
25 changed files with 1278 additions and 2 deletions
@@ -0,0 +1,49 @@
from fastvideo import VideoGenerator
from fastvideo.models.dits.lingbotworld.cam_utils import prepare_camera_embedding
# from fastvideo.configs.sample import SamplingParam
OUTPUT_PATH = "video_samples_lingbotworld"
def main():
# FastVideo will automatically use the optimal default arguments for the
# model.
# If a local path is provided, FastVideo will make a best effort
# attempt to identify the optimal arguments.
generator = VideoGenerator.from_pretrained(
"FastVideo/LingBot-World-Base-Cam-Diffusers",
# FastVideo will automatically handle distributed setup
num_gpus=1,
use_fsdp_inference=False, # set to True if GPU is out of memory
dit_cpu_offload=True, # DiT need to be offloaded for MoE
vae_cpu_offload=False,
text_encoder_cpu_offload=True,
# Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
pin_cpu_memory=True,
# image_encoder_cpu_offload=False,
)
num_frames = 81
prompt = "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls."
image_path = "https://raw.githubusercontent.com/Robbyant/lingbot-world/main/examples/00/image.jpg"
action_path = "examples/inference/basic/lingbotworld_examples/00"
c2ws_plucker_emb, num_frames = prepare_camera_embedding(
action_path=action_path,
num_frames=num_frames,
height=480,
width=832,
spatial_scale=8,
)
generator.generate_video(
prompt,
image_path=image_path,
output_path=OUTPUT_PATH,
save_video=True,
num_frames=num_frames,
height=480,
width=832,
c2ws_plucker_emb=c2ws_plucker_emb,
)
if __name__ == "__main__":
main()
@@ -3,6 +3,7 @@ from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig
from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
from fastvideo.configs.models.dits.lingbotworld import LingBotWorldVideoConfig
from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
from fastvideo.configs.models.dits.stepvideo import StepVideoConfig
@@ -13,4 +14,8 @@ __all__ = [
"HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig",
"WanVideoConfig", "StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig",
"LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig"
"HunyuanVideoConfig", "HunyuanVideo15Config", "WanVideoConfig",
"StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig",
"LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig",
"LingBotWorldVideoConfig"
]
@@ -0,0 +1,110 @@
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
def is_blocks(n: str, m) -> bool:
return "blocks" in n and str.isdigit(n.split(".")[-1])
@dataclass
class LingBotWorldArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
param_names_mapping: dict = field(
default_factory=lambda: {
r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1",
r"^patch_embedding_wancamctrl\.(.*)$":
r"patch_embedding_wancamctrl.proj.\1",
r"^c2ws_hidden_states_layer1\.(.*)$": r"c2ws_mlp.fc_in.\1",
r"^c2ws_hidden_states_layer2\.(.*)$": r"c2ws_mlp.fc_out.\1",
r"^text_embedding\.0\.(.*)$":
r"condition_embedder.text_embedder.fc_in.\1",
r"^text_embedding\.2\.(.*)$":
r"condition_embedder.text_embedder.fc_out.\1",
r"^time_embedding\.0\.(.*)$":
r"condition_embedder.time_embedder.mlp.fc_in.\1",
r"^time_embedding\.2\.(.*)$":
r"condition_embedder.time_embedder.mlp.fc_out.\1",
r"^time_projection\.1\.(.*)$":
r"condition_embedder.time_modulation.linear.\1",
r"^blocks\.(\d+)\.modulation$": r"blocks.\1.scale_shift_table",
r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.to_q.\2",
r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.to_k.\2",
r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.to_v.\2",
r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.to_out.\2",
r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2",
r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2",
r"^blocks\.(\d+)\.norm3\.(.*)$":
r"blocks.\1.self_attn_residual_norm.norm.\2",
r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2",
r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2",
r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2",
r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$":
r"blocks.\1.attn2.to_out.\2",
r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$":
r"blocks.\1.attn2.norm_q.\2",
r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$":
r"blocks.\1.attn2.norm_k.\2",
r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
r"^blocks\.(\d+)\.cam_injector_layer1\.(.*)$":
r"blocks.\1.cam_conditioner.cam_injector.fc_in.\2",
r"^blocks\.(\d+)\.cam_injector_layer2\.(.*)$":
r"blocks.\1.cam_conditioner.cam_injector.fc_out.\2",
r"^blocks\.(\d+)\.cam_scale_layer\.(.*)$":
r"blocks.\1.cam_conditioner.cam_scale_layer.\2",
r"^blocks\.(\d+)\.cam_shift_layer\.(.*)$":
r"blocks.\1.cam_conditioner.cam_shift_layer.\2",
r"^head\.modulation$": r"scale_shift_table",
r"^head\.head\.(.*)$": r"proj_out.\1",
})
# Reverse mapping for saving checkpoints: custom -> hf
reverse_param_names_mapping: dict = field(default_factory=lambda: {})
# Some LoRA adapters use the original official layer names instead of hf layer names,
# so apply this before the param_names_mapping
lora_param_names_mapping: dict = field(default_factory=lambda: {})
patch_size: tuple[int, int, int] = (1, 2, 2)
text_len: int = 512
num_attention_heads: int = 40
attention_head_dim: int = 128
in_channels: int = 16
out_channels: int = 16
text_dim: int = 4096
freq_dim: int = 256
ffn_dim: int = 13824
num_layers: int = 40
cross_attn_norm: bool = True
qk_norm: str = "rms_norm_across_heads"
eps: float = 1e-6
image_dim: int | None = None
added_kv_proj_dim: int | None = None
rope_max_seq_len: int = 1024
pos_embed_seq_len: int | None = None
exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
# Wan MoE
boundary_ratio: float | None = None
# Causal Wan
local_attn_size: int = -1 # Window size for temporal local attention (-1 indicates global attention)
sink_size: int = 0 # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache
num_frames_per_block: int = 3
sliding_window_num_frames: int = 21
def __post_init__(self):
super().__post_init__()
self.out_channels = self.out_channels or self.in_channels
self.hidden_size = self.num_attention_heads * self.attention_head_dim
self.num_channels_latents = self.out_channels
@dataclass
class LingBotWorldVideoConfig(DiTConfig):
arch_config: DiTArchConfig = field(default_factory=LingBotWorldArchConfig)
prefix: str = "Wan"
+3 -1
View File
@@ -6,6 +6,7 @@ from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
from fastvideo.configs.pipelines.hyworld import HYWorldConfig
from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig
from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig
from fastvideo.configs.pipelines.sd35 import SD35Config
from fastvideo.registry import get_pipeline_config_cls_from_name
@@ -26,5 +27,6 @@ __all__ = [
"WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig",
"WanI2V720PConfig", "StepVideoT2VConfig", "SelfForcingWanT2V480PConfig",
"CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "HYWorldConfig",
"SD35Config", "get_pipeline_config_cls_from_name"
"SD35Config", "LingBotWorldI2V480PConfig",
"get_pipeline_config_cls_from_name"
]
@@ -0,0 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from fastvideo.configs.models import DiTConfig
from fastvideo.configs.pipelines.wan import Wan2_2_I2V_A14B_Config
from fastvideo.configs.models.dits.lingbotworld import LingBotWorldVideoConfig
@dataclass
class LingBotWorldI2V480PConfig(Wan2_2_I2V_A14B_Config):
dit_config: DiTConfig = field(default_factory=LingBotWorldVideoConfig)
flow_shift: float | None = 10.0
boundary_ratio: float | None = 0.947
+3
View File
@@ -31,6 +31,9 @@ class SamplingParam:
# Camera control inputs (HYWorld)
pose: str | None = None # Camera trajectory: pose string (e.g., 'w-31') or JSON file path
# Camera control inputs (LingBotWorld)
c2ws_plucker_emb: Any | None = None # Plucker embedding: [B, C, F_lat, H_lat, W_lat]
# Refine inputs (LongCat 480p->720p upscaling)
# Path-based refine (load stage1 video from disk, e.g. MP4)
refine_from: str | None = None # Path to stage1 video (480p output from distill)
+21
View File
@@ -0,0 +1,21 @@
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass
from fastvideo.configs.sample.wan import Wan2_2_I2V_A14B_SamplingParam
@dataclass
class LingBotWorld_SamplingParam(Wan2_2_I2V_A14B_SamplingParam):
guidance_scale: float = 5.0 # high_noise
guidance_scale_2: float = 5.0 # low_noise
num_inference_steps: int = 70
boundary_ratio: float | None = 0.947
negative_prompt: str | None = (
"画面突变,色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,"
"最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,"
"畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走,"
"镜头晃动,画面闪烁,模糊,噪点,水印,签名,文字,变形,扭曲,液化,不合逻辑的结构,卡顿,"
"PPT幻灯片感,过暗,欠曝,低对比度,霓虹灯光感,过度锐化,3D渲染感,人物,行人,游客,身体,"
"皮肤,肢体,面部特征,汽车,电线")
fps: int = 16
# NOTE(will): default boundary timestep is tracked by PipelineConfig, but
# can be overridden during sampling
+56 -1
View File
@@ -60,6 +60,61 @@ class PatchEmbed(nn.Module):
return x
class WanCamControlPatchEmbedding(nn.Module):
"""Lingbot World Patch embedding for Plucker features."""
def __init__(
self,
patch_size=(1, 2, 2),
in_chans=384, # 6 * 64
embed_dim=2048,
bias=True,
dtype=None,
prefix: str = ""):
super().__init__()
# must be 3-tuple
if isinstance(patch_size, list | tuple):
if len(patch_size) != 3:
raise ValueError(
f"patch_size must have length 3, got {len(patch_size)}")
else:
raise ValueError(f"Unsupported patch_size type: {type(patch_size)}")
self.patch_size = patch_size
pt, ph, pw = self.patch_size
self.in_features = in_chans * pt * ph * pw
self.proj = nn.Linear(self.in_features,
embed_dim,
bias=bias,
dtype=dtype)
def forward(self, x: torch.Tensor) -> torch.Tensor:
if x.dim() != 5:
raise ValueError(
f"Expected camera embedding shape [B, C, F, H, W], got {x.shape}"
)
bsz, channels, frames, height, width = x.shape
pt, ph, pw = self.patch_size
if (frames % pt) != 0 or (height % ph) != 0 or (width % pw) != 0:
raise ValueError(
f"Input shape {x.shape} must be divisible by patch_size {self.patch_size}"
)
# '1 c (f c1) (h c2) (w c3) -> 1 (f h w) (c c1 c2 c3)',
x = x.view(
bsz,
channels,
frames // pt,
pt,
height // ph,
ph,
width // pw,
pw,
)
x = x.permute(0, 2, 4, 6, 1, 3, 5, 7).reshape(bsz, -1, self.in_features)
return self.proj(x)
class TimestepEmbedder(nn.Module):
"""
Embeds scalar timesteps into vector representations.
@@ -252,4 +307,4 @@ class Timesteps(nn.Module):
downscale_freq_shift=self.downscale_freq_shift,
scale=self.scale,
)
return t_emb
return t_emb
@@ -0,0 +1,8 @@
from .model import LingBotWorldTransformer3DModel
__all__ = [
"LingBotWorldTransformer3DModel",
]
# Entry point for model registry
EntryClass = [LingBotWorldTransformer3DModel]
@@ -0,0 +1,203 @@
# SPDX-License-Identifier: Apache-2.0
# Adapted from LingBot World: https://github.com/Robbyant/lingbot-world/blob/main/wan/utils/cam_utils.py
import numpy as np
import os
import torch
from scipy.interpolate import interp1d
from scipy.spatial.transform import Rotation, Slerp
# --- Official Code (Leave Unchanged) ---
def interpolate_camera_poses(
src_indices: np.ndarray,
src_rot_mat: np.ndarray,
src_trans_vec: np.ndarray,
tgt_indices: np.ndarray,
) -> torch.Tensor:
# interpolate translation
interp_func_trans = interp1d(
src_indices,
src_trans_vec,
axis=0,
kind='linear',
bounds_error=False,
fill_value="extrapolate",
)
interpolated_trans_vec = interp_func_trans(tgt_indices)
# interpolate rotation
src_quat_vec = Rotation.from_matrix(src_rot_mat)
# ensure there is no sudden change in qw
quats = src_quat_vec.as_quat().copy() # [N, 4]
for i in range(1, len(quats)):
if np.dot(quats[i], quats[i-1]) < 0:
quats[i] = -quats[i]
src_quat_vec = Rotation.from_quat(quats)
slerp_func_rot = Slerp(src_indices, src_quat_vec)
interpolated_rot_quat = slerp_func_rot(tgt_indices)
interpolated_rot_mat = interpolated_rot_quat.as_matrix()
poses = np.zeros((len(tgt_indices), 4, 4))
poses[:, :3, :3] = interpolated_rot_mat
poses[:, :3, 3] = interpolated_trans_vec
poses[:, 3, 3] = 1.0
return torch.from_numpy(poses).float()
def SE3_inverse(T: torch.Tensor) -> torch.Tensor:
Rot = T[:, :3, :3] # [B,3,3]
trans = T[:, :3, 3:] # [B,3,1]
R_inv = Rot.transpose(-1, -2)
t_inv = -torch.bmm(R_inv, trans)
T_inv = torch.eye(4, device=T.device, dtype=T.dtype)[None, :, :].repeat(T.shape[0], 1, 1)
T_inv[:, :3, :3] = R_inv
T_inv[:, :3, 3:] = t_inv
return T_inv
def compute_relative_poses(
c2ws_mat: torch.Tensor,
framewise: bool = False,
normalize_trans: bool = True,
) -> torch.Tensor:
ref_w2cs = SE3_inverse(c2ws_mat[0:1])
relative_poses = torch.matmul(ref_w2cs, c2ws_mat)
# ensure identity matrix for 1st frame
relative_poses[0] = torch.eye(4, device=c2ws_mat.device, dtype=c2ws_mat.dtype)
if framewise:
# compute pose between i and i+1
relative_poses_framewise = torch.bmm(SE3_inverse(relative_poses[:-1]), relative_poses[1:])
relative_poses[1:] = relative_poses_framewise
if normalize_trans: # note refer to camctrl2: "we scale the coordinate inputs to roughly 1 standard deviation to simplify model learning."
translations = relative_poses[:, :3, 3] # [f, 3]
max_norm = torch.norm(translations, dim=-1).max()
# only normlaize when moving
if max_norm > 0:
relative_poses[:, :3, 3] = translations / max_norm
return relative_poses
@torch.no_grad()
def create_meshgrid(n_frames: int, height: int, width: int, bias: float = 0.5, device='cuda', dtype=torch.float32) -> torch.Tensor:
x_range = torch.arange(width, device=device, dtype=dtype)
y_range = torch.arange(height, device=device, dtype=dtype)
grid_y, grid_x = torch.meshgrid(y_range, x_range, indexing='ij')
grid_xy = torch.stack([grid_x, grid_y], dim=-1).view([-1, 2]) + bias # [h*w, 2]
grid_xy = grid_xy[None, ...].repeat(n_frames, 1, 1) # [f, h*w, 2]
return grid_xy
def get_plucker_embeddings(
c2ws_mat: torch.Tensor,
Ks: torch.Tensor,
height: int,
width: int,
):
n_frames = c2ws_mat.shape[0]
grid_xy = create_meshgrid(n_frames, height, width, device=c2ws_mat.device, dtype=c2ws_mat.dtype) # [f, h*w, 2]
fx, fy, cx, cy = Ks.chunk(4, dim=-1) # [f, 1]
i = grid_xy[..., 0] # [f, h*w]
j = grid_xy[..., 1] # [f, h*w]
zs = torch.ones_like(i) # [f, h*w]
xs = (i - cx) / fx * zs
ys = (j - cy) / fy * zs
directions = torch.stack([xs, ys, zs], dim=-1) # [f, h*w, 3]
directions = directions / directions.norm(dim=-1, keepdim=True) # [f, h*w, 3]
rays_d = directions @ c2ws_mat[:, :3, :3].transpose(-1, -2) # [f, h*w, 3]
rays_o = c2ws_mat[:, :3, 3] # [f, 3]
rays_o = rays_o[:, None, :].expand_as(rays_d) # [f, h*w, 3]
# rays_dxo = torch.cross(rays_o, rays_d, dim=-1) # [f, h*w, 3]
# note refer to: apt2
plucker_embeddings = torch.cat([rays_o, rays_d], dim=-1) # [f, h*w, 6]
plucker_embeddings = plucker_embeddings.view([n_frames, height, width, 6]) # [f*h*w, 6]
return plucker_embeddings
def get_Ks_transformed(
Ks: torch.Tensor,
height_org: int,
width_org: int,
height_resize: int,
width_resize: int,
height_final: int,
width_final: int,
):
fx, fy, cx, cy = Ks.chunk(4, dim=-1) # [f, 1]
scale_x = width_resize / width_org
scale_y = height_resize / height_org
fx_resize = fx * scale_x
fy_resize = fy * scale_y
cx_resize = cx * scale_x
cy_resize = cy * scale_y
crop_offset_x = (width_resize - width_final) / 2
crop_offset_y = (height_resize - height_final) / 2
cx_final = cx_resize - crop_offset_x
cy_final = cy_resize - crop_offset_y
Ks_transformed = torch.zeros_like(Ks)
Ks_transformed[:, 0:1] = fx_resize
Ks_transformed[:, 1:2] = fy_resize
Ks_transformed[:, 2:3] = cx_final
Ks_transformed[:, 3:4] = cy_final
return Ks_transformed
# --- Custom ---
def prepare_camera_embedding(
action_path: str,
num_frames: int,
height: int,
width: int,
spatial_scale: int = 8,
) -> tuple[torch.Tensor, int]:
c2ws = np.load(os.path.join(action_path, "poses.npy"))
len_c2ws = ((len(c2ws) - 1) // 4) * 4 + 1
num_frames = min(num_frames, len_c2ws)
c2ws = c2ws[:num_frames]
Ks = torch.from_numpy(
np.load(os.path.join(action_path, "intrinsics.npy"))
).float()
Ks = get_Ks_transformed(
Ks,
height_org=480,
width_org=832,
height_resize=height,
width_resize=width,
height_final=height,
width_final=width,
)
Ks = Ks[0] # use first frame
len_c2ws = len(c2ws)
num_latent_frames = (len_c2ws - 1) // 4 + 1
c2ws_infer = interpolate_camera_poses(
src_indices=np.linspace(0, len_c2ws - 1, len_c2ws),
src_rot_mat=c2ws[:, :3, :3],
src_trans_vec=c2ws[:, :3, 3],
tgt_indices=np.linspace(0, len_c2ws - 1, num_latent_frames),
)
c2ws_infer = compute_relative_poses(c2ws_infer, framewise=True)
Ks = Ks.repeat(num_latent_frames, 1)
plucker = get_plucker_embeddings(c2ws_infer, Ks, height, width) # [F, H, W, 6]
# reshpae
latent_height = height // spatial_scale
latent_width = width // spatial_scale
plucker = plucker.view(num_latent_frames, latent_height, spatial_scale, latent_width, spatial_scale, 6)
plucker = plucker.permute(0, 1, 3, 5, 2, 4).contiguous()
plucker = plucker.view(num_latent_frames, latent_height, latent_width, 6 * spatial_scale * spatial_scale)
c2ws_plucker_emb = plucker.permute(3, 0, 1, 2).contiguous().unsqueeze(0)
return c2ws_plucker_emb, num_frames
+569
View File
@@ -0,0 +1,569 @@
# SPDX-License-Identifier: Apache-2.0
import math
from contextlib import nullcontext
from typing import Any
import numpy as np
import torch
import torch.nn as nn
from fastvideo.attention import DistributedAttention
from fastvideo.configs.models.dits.lingbotworld import LingBotWorldVideoConfig
from fastvideo.configs.sample.wan import WanTeaCacheParams
from fastvideo.distributed.communication_op import (
sequence_model_parallel_all_gather_with_unpad,
sequence_model_parallel_shard)
from fastvideo.forward_context import get_forward_context
from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift,
RMSNorm, ScaleResidual,
ScaleResidualLayerNormScaleShift)
from fastvideo.layers.linear import ReplicatedLinear
# from torch.nn import RMSNorm
# TODO: RMSNorm ....
from fastvideo.layers.mlp import MLP
from fastvideo.layers.rotary_embedding import get_rotary_pos_embed
from fastvideo.layers.visual_embedding import (PatchEmbed, WanCamControlPatchEmbedding)
from fastvideo.logger import init_logger
from fastvideo.models.dits.base import CachableDiT
from fastvideo.models.dits.wanvideo import (
WanI2VCrossAttention,
WanT2VCrossAttention,
WanTimeTextImageEmbedding,
)
from fastvideo.platforms import AttentionBackendEnum, current_platform
from fastvideo.distributed.parallel_state import get_sp_world_size
from fastvideo.distributed.utils import create_attention_mask_for_padding
logger = init_logger(__name__)
class LingBotWorldCamConditioner(nn.Module):
def __init__(self, dim: int) -> None:
super().__init__()
self.cam_injector = MLP(dim, dim, dim, bias=True, act_type="silu")
self.cam_scale_layer = nn.Linear(dim, dim)
self.cam_shift_layer = nn.Linear(dim, dim)
def forward(
self,
hidden_states: torch.Tensor,
c2ws_plucker_emb: torch.Tensor | None,
) -> torch.Tensor:
if c2ws_plucker_emb is None:
return hidden_states
assert c2ws_plucker_emb.shape == hidden_states.shape, (
f"c2ws_plucker_emb shape must match hidden_states shape, got "
f"{tuple(c2ws_plucker_emb.shape)} vs {tuple(hidden_states.shape)}"
)
c2ws_hidden_states = self.cam_injector(c2ws_plucker_emb)
c2ws_hidden_states = c2ws_hidden_states + c2ws_plucker_emb
cam_scale = self.cam_scale_layer(c2ws_hidden_states)
cam_shift = self.cam_shift_layer(c2ws_hidden_states)
return (1.0 + cam_scale) * hidden_states + cam_shift
class LingBotWorldTransformerBlock(nn.Module):
def __init__(self,
dim: int,
ffn_dim: int,
num_heads: int,
qk_norm: str = "rms_norm_across_heads",
cross_attn_norm: bool = False,
eps: float = 1e-6,
added_kv_proj_dim: int | None = None,
supported_attention_backends: tuple[AttentionBackendEnum, ...]
| None = None,
prefix: str = ""):
super().__init__()
# 1. Self-attention
self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
self.to_q = ReplicatedLinear(dim, dim, bias=True)
self.to_k = ReplicatedLinear(dim, dim, bias=True)
self.to_v = ReplicatedLinear(dim, dim, bias=True)
self.to_out = ReplicatedLinear(dim, dim, bias=True)
self.attn1 = DistributedAttention(
num_heads=num_heads,
head_size=dim // num_heads,
causal=False,
supported_attention_backends=supported_attention_backends,
prefix=f"{prefix}.attn1")
self.hidden_dim = dim
self.num_attention_heads = num_heads
dim_head = dim // num_heads
if qk_norm == "rms_norm":
self.norm_q = RMSNorm(dim_head, eps=eps)
self.norm_k = RMSNorm(dim_head, eps=eps)
elif qk_norm == "rms_norm_across_heads":
# LTX applies qk norm across all heads
self.norm_q = RMSNorm(dim, eps=eps)
self.norm_k = RMSNorm(dim, eps=eps)
else:
raise NotImplementedError(
f"QK Norm type '{qk_norm}' not supported")
assert cross_attn_norm is True
self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
dim,
norm_type="layer",
eps=eps,
elementwise_affine=True,
dtype=torch.float32,
compute_dtype=torch.float32)
# 2. Cross-attention
if added_kv_proj_dim is not None:
# I2V
self.attn2 = WanI2VCrossAttention(dim,
num_heads,
qk_norm=qk_norm,
eps=eps)
else:
# T2V
self.attn2 = WanT2VCrossAttention(dim,
num_heads,
qk_norm=qk_norm,
eps=eps)
self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
dim,
norm_type="layer",
eps=eps,
elementwise_affine=False,
dtype=torch.float32,
compute_dtype=torch.float32)
# 3. Feed-forward
self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
self.mlp_residual = ScaleResidual()
self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
self.cam_conditioner = LingBotWorldCamConditioner(dim)
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
freqs_cis: tuple[torch.Tensor, torch.Tensor],
attention_mask: torch.Tensor | None = None,
c2ws_plucker_emb: torch.Tensor | None = None,
) -> torch.Tensor:
if hidden_states.dim() == 4:
hidden_states = hidden_states.squeeze(1)
bs, seq_length, _ = hidden_states.shape
orig_dtype = hidden_states.dtype
# assert orig_dtype != torch.float32
if temb.dim() == 4:
# temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
self.scale_shift_table.unsqueeze(0) + temb.float()
).chunk(6, dim=2)
# batch_size, seq_len, 1, inner_dim
shift_msa = shift_msa.squeeze(2)
scale_msa = scale_msa.squeeze(2)
gate_msa = gate_msa.squeeze(2)
c_shift_msa = c_shift_msa.squeeze(2)
c_scale_msa = c_scale_msa.squeeze(2)
c_gate_msa = c_gate_msa.squeeze(2)
else:
# temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B)
e = self.scale_shift_table + temb.float()
shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
6, dim=1)
assert shift_msa.dtype == torch.float32
# 1. Self-attention
norm_hidden_states = (self.norm1(hidden_states.float()) *
(1 + scale_msa) + shift_msa).to(orig_dtype)
query, _ = self.to_q(norm_hidden_states)
key, _ = self.to_k(norm_hidden_states)
value, _ = self.to_v(norm_hidden_states)
if self.norm_q is not None:
query = self.norm_q(query)
if self.norm_k is not None:
key = self.norm_k(key)
query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
attn_output, _ = self.attn1(query, key, value, freqs_cis=freqs_cis, attention_mask=attention_mask)
attn_output = attn_output.flatten(2)
attn_output, _ = self.to_out(attn_output)
attn_output = attn_output.squeeze(1)
null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
norm_hidden_states, hidden_states = self.self_attn_residual_norm(
hidden_states, attn_output, gate_msa, null_shift, null_scale)
norm_hidden_states, hidden_states = norm_hidden_states.to(
orig_dtype), hidden_states.to(orig_dtype)
# Inject camera condition
# must be applied after the self-attention residual update.
hidden_states = self.cam_conditioner(hidden_states, c2ws_plucker_emb)
norm_hidden_states = self.self_attn_residual_norm.norm(hidden_states)
norm_hidden_states = norm_hidden_states.to(orig_dtype)
# 2. Cross-attention
attn_output = self.attn2(norm_hidden_states,
context=encoder_hidden_states,
context_lens=None)
norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
norm_hidden_states, hidden_states = norm_hidden_states.to(
orig_dtype), hidden_states.to(orig_dtype)
# 3. Feed-forward
ff_output = self.ffn(norm_hidden_states)
hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
hidden_states = hidden_states.to(orig_dtype)
return hidden_states
class LingBotWorldTransformer3DModel(CachableDiT):
_fsdp_shard_conditions = LingBotWorldVideoConfig()._fsdp_shard_conditions
_compile_conditions = LingBotWorldVideoConfig()._compile_conditions
_supported_attention_backends = LingBotWorldVideoConfig(
)._supported_attention_backends
param_names_mapping = LingBotWorldVideoConfig().param_names_mapping
reverse_param_names_mapping = LingBotWorldVideoConfig().reverse_param_names_mapping
lora_param_names_mapping = LingBotWorldVideoConfig().lora_param_names_mapping
def __init__(self, config: LingBotWorldVideoConfig, hf_config: dict[str,
Any]) -> None:
super().__init__(config=config, hf_config=hf_config)
inner_dim = config.num_attention_heads * config.attention_head_dim
self.hidden_size = config.hidden_size
self.num_attention_heads = config.num_attention_heads
self.in_channels = config.in_channels
self.out_channels = config.out_channels
self.num_channels_latents = config.num_channels_latents
self.patch_size = config.patch_size
self.text_len = config.text_len
assert config.num_attention_heads % get_sp_world_size() == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})"
# 1. Patch & position embedding
self.patch_embedding = PatchEmbed(in_chans=config.in_channels,
embed_dim=inner_dim,
patch_size=config.patch_size,
flatten=False)
self.patch_embedding_wancamctrl = WanCamControlPatchEmbedding(in_chans=6 * 64,
embed_dim=inner_dim,
patch_size=config.patch_size)
self.c2ws_mlp = MLP(inner_dim, inner_dim, inner_dim, bias=True, act_type="silu")
# 2. Condition embeddings
self.condition_embedder = WanTimeTextImageEmbedding(
dim=inner_dim,
time_freq_dim=config.freq_dim,
text_embed_dim=config.text_dim,
image_embed_dim=config.image_dim,
)
# 3. Transformer blocks
transformer_block = LingBotWorldTransformerBlock
self.blocks = nn.ModuleList([
transformer_block(inner_dim,
config.ffn_dim,
config.num_attention_heads,
config.qk_norm,
config.cross_attn_norm,
config.eps,
config.added_kv_proj_dim,
self._supported_attention_backends,
prefix=f"{config.prefix}.blocks.{i}")
for i in range(config.num_layers)
])
# 4. Output norm & projection
self.norm_out = LayerNormScaleShift(inner_dim,
norm_type="layer",
eps=config.eps,
elementwise_affine=False,
dtype=torch.float32,
compute_dtype=torch.float32)
self.proj_out = nn.Linear(
inner_dim, config.out_channels * math.prod(config.patch_size))
self.scale_shift_table = nn.Parameter(
torch.randn(1, 2, inner_dim) / inner_dim**0.5)
self.gradient_checkpointing = False
self._logged_attention_mask = False
# For type checking
self.previous_e0_even = None
self.previous_e0_odd = None
self.previous_residual_even = None
self.previous_residual_odd = None
self.is_even = True
self.should_calc_even = True
self.should_calc_odd = True
self.accumulated_rel_l1_distance_even = 0
self.accumulated_rel_l1_distance_odd = 0
self.cnt = 0
self.__post_init__()
def forward(self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor | list[torch.Tensor],
timestep: torch.LongTensor,
encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
| None = None,
guidance=None,
c2ws_plucker_emb: torch.Tensor | None = None,
**kwargs) -> torch.Tensor:
forward_batch = get_forward_context().forward_batch
enable_teacache = forward_batch is not None and forward_batch.enable_teacache
orig_dtype = hidden_states.dtype
if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
encoder_hidden_states = encoder_hidden_states[0]
if isinstance(encoder_hidden_states_image,
list) and len(encoder_hidden_states_image) > 0:
encoder_hidden_states_image = encoder_hidden_states_image[0]
else:
encoder_hidden_states_image = None
batch_size, num_channels, num_frames, height, width = hidden_states.shape
p_t, p_h, p_w = self.patch_size
post_patch_num_frames = num_frames // p_t
post_patch_height = height // p_h
post_patch_width = width // p_w
# Get rotary embeddings
d = self.hidden_size // self.num_attention_heads
rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
freqs_cos, freqs_sin = get_rotary_pos_embed(
(post_patch_num_frames, post_patch_height,
post_patch_width),
self.hidden_size,
self.num_attention_heads,
rope_dim_list,
dtype=torch.float32 if current_platform.is_mps() else torch.float64,
rope_theta=10000)
freqs_cis = (freqs_cos.to(hidden_states.device).float(),
freqs_sin.to(hidden_states.device).float())
hidden_states = self.patch_embedding(hidden_states)
hidden_states = hidden_states.flatten(2).transpose(1, 2)
c2ws_hidden_states = None
if c2ws_plucker_emb is not None:
c2ws_plucker_emb = self.patch_embedding_wancamctrl(
c2ws_plucker_emb.to(device=hidden_states.device, dtype=hidden_states.dtype)
)
c2ws_hidden_states = self.c2ws_mlp(c2ws_plucker_emb)
c2ws_plucker_emb = c2ws_plucker_emb + c2ws_hidden_states
# Shard with padding support - returns (sharded_tensor, original_seq_len)
hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
# Shard c2ws_plucker_emb
if c2ws_plucker_emb is not None:
c2ws_plucker_emb, _ = sequence_model_parallel_shard(c2ws_plucker_emb, dim=1)
# Create attention mask for padded tokens if padding was applied
current_seq_len = hidden_states.shape[1]
sp_world_size = get_sp_world_size()
padded_seq_len = current_seq_len * sp_world_size
if padded_seq_len > original_seq_len:
if not self._logged_attention_mask:
logger.info(f"Padding applied, original seq len: {original_seq_len}, padded seq len: {padded_seq_len}")
self._logged_attention_mask = True
attention_mask = create_attention_mask_for_padding(
seq_len=original_seq_len,
padded_seq_len=padded_seq_len,
batch_size=batch_size,
device=hidden_states.device,
)
else:
if not self._logged_attention_mask:
logger.info(f"Padding not applied")
self._logged_attention_mask = True
attention_mask = None
# timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v)
if timestep.dim() == 2:
ts_seq_len = timestep.shape[1]
timestep = timestep.flatten() # batch_size * seq_len
else:
ts_seq_len = None
temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len)
if ts_seq_len is not None:
# batch_size, seq_len, 6, inner_dim
timestep_proj = timestep_proj.unflatten(2, (6, -1))
else:
# batch_size, 6, inner_dim
timestep_proj = timestep_proj.unflatten(1, (6, -1))
if encoder_hidden_states_image is not None:
if encoder_hidden_states is not None:
encoder_hidden_states = torch.concat(
[encoder_hidden_states_image, encoder_hidden_states], dim=1)
else:
encoder_hidden_states = encoder_hidden_states_image
if current_platform.is_mps() or current_platform.is_npu():
encoder_hidden_states = encoder_hidden_states.to(orig_dtype)
else:
encoder_hidden_states = encoder_hidden_states # cast to orig_dtype for MPS & NPU
assert encoder_hidden_states.dtype == orig_dtype
# 4. Transformer blocks
# if caching is enabled, we might be able to skip the forward pass
should_skip_forward = self.should_skip_forward_for_cached_states(
timestep_proj=timestep_proj, temb=temb)
if should_skip_forward:
print("skipping forward, cached")
hidden_states = self.retrieve_cached_states(hidden_states)
else:
# if teacache is enabled, we need to cache the original hidden states
if enable_teacache:
original_hidden_states = hidden_states.clone()
if torch.is_grad_enabled() and self.gradient_checkpointing:
for block in self.blocks:
hidden_states = self._gradient_checkpointing_func(
block, hidden_states, encoder_hidden_states,
timestep_proj, freqs_cis, attention_mask, c2ws_plucker_emb)
else:
for block in self.blocks:
hidden_states = block(hidden_states, encoder_hidden_states,
timestep_proj, freqs_cis, attention_mask, c2ws_plucker_emb)
# if teacache is enabled, we need to cache the original hidden states
if enable_teacache:
self.maybe_cache_states(hidden_states, original_hidden_states)
# 5. Output norm, projection & unpatchify
if temb.dim() == 3:
# batch_size, seq_len, inner_dim (wan 2.2 ti2v)
shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2)
shift = shift.squeeze(2)
scale = scale.squeeze(2)
else:
# batch_size, inner_dim
shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1)
hidden_states = self.norm_out(hidden_states, shift, scale)
# Gather and unpad in one operation
hidden_states = sequence_model_parallel_all_gather_with_unpad(
hidden_states, original_seq_len, dim=1)
hidden_states = self.proj_out(hidden_states)
hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
post_patch_height,
post_patch_width, p_t, p_h, p_w,
-1)
hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
return output
def maybe_cache_states(self, hidden_states: torch.Tensor,
original_hidden_states: torch.Tensor) -> None:
if self.is_even:
self.previous_residual_even = hidden_states.squeeze(
0) - original_hidden_states
else:
self.previous_residual_odd = hidden_states.squeeze(
0) - original_hidden_states
def should_skip_forward_for_cached_states(self, **kwargs) -> bool:
forward_context = get_forward_context()
forward_batch = forward_context.forward_batch
if forward_batch is None or not forward_batch.enable_teacache:
return False
teacache_params = forward_batch.teacache_params
assert teacache_params is not None, "teacache_params is not initialized"
assert isinstance(
teacache_params,
WanTeaCacheParams), "teacache_params is not a WanTeaCacheParams"
current_timestep = forward_context.current_timestep
num_inference_steps = forward_batch.num_inference_steps
# initialize the coefficients, cutoff_steps, and ret_steps
coefficients = teacache_params.coefficients
use_ret_steps = teacache_params.use_ret_steps
cutoff_steps = teacache_params.get_cutoff_steps(num_inference_steps)
ret_steps = teacache_params.ret_steps
teacache_thresh = teacache_params.teacache_thresh
if current_timestep == 0:
self.cnt = 0
timestep_proj = kwargs["timestep_proj"]
temb = kwargs["temb"]
modulated_inp = timestep_proj if use_ret_steps else temb
if self.cnt % 2 == 0: # even -> condition
self.is_even = True
if self.cnt < ret_steps or self.cnt >= cutoff_steps:
self.should_calc_even = True
self.accumulated_rel_l1_distance_even = 0
else:
assert self.previous_e0_even is not None, "previous_e0_even is not initialized"
assert self.accumulated_rel_l1_distance_even is not None, "accumulated_rel_l1_distance_even is not initialized"
rescale_func = np.poly1d(coefficients)
self.accumulated_rel_l1_distance_even += rescale_func(
((modulated_inp - self.previous_e0_even).abs().mean() /
self.previous_e0_even.abs().mean()).cpu().item())
if self.accumulated_rel_l1_distance_even < teacache_thresh:
self.should_calc_even = False
else:
self.should_calc_even = True
self.accumulated_rel_l1_distance_even = 0
self.previous_e0_even = modulated_inp.clone()
else: # odd -> unconditon
self.is_even = False
if self.cnt < ret_steps or self.cnt >= cutoff_steps:
self.should_calc_odd = True
self.accumulated_rel_l1_distance_odd = 0
else:
assert self.previous_e0_odd is not None, "previous_e0_odd is not initialized"
assert self.accumulated_rel_l1_distance_odd is not None, "accumulated_rel_l1_distance_odd is not initialized"
rescale_func = np.poly1d(coefficients)
self.accumulated_rel_l1_distance_odd += rescale_func(
((modulated_inp - self.previous_e0_odd).abs().mean() /
self.previous_e0_odd.abs().mean()).cpu().item())
if self.accumulated_rel_l1_distance_odd < teacache_thresh:
self.should_calc_odd = False
else:
self.should_calc_odd = True
self.accumulated_rel_l1_distance_odd = 0
self.previous_e0_odd = modulated_inp.clone()
self.cnt += 1
should_skip_forward = False
if self.is_even:
if not self.should_calc_even:
should_skip_forward = True
else:
if not self.should_calc_odd:
should_skip_forward = True
return should_skip_forward
def retrieve_cached_states(self,
hidden_states: torch.Tensor) -> torch.Tensor:
if self.is_even:
return hidden_states + self.previous_residual_even
else:
return hidden_states + self.previous_residual_odd
# Entry point for model registry
EntryClass = LingBotWorldTransformer3DModel
+1
View File
@@ -40,6 +40,7 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
"LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"), # Native (Phase 2)
"LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"),
"SD3Transformer2DModel": ("dits", "sd3", "SD3Transformer2DModel"),
"LingBotWorldTransformer3DModel": ("dits", "lingbotworld", "LingBotWorldTransformer3DModel"),
}
_IMAGE_TO_VIDEO_DIT_MODELS = {
@@ -0,0 +1 @@
@@ -0,0 +1,16 @@
# SPDX-License-Identifier: Apache-2.0
"""
Wan video diffusion pipeline implementation.
This module contains an implementation of the Wan video diffusion pipeline
using the modular pipeline architecture.
"""
from fastvideo.pipelines.basic.wan.wan_i2v_pipeline import WanImageToVideoPipeline
class LingBotWorldImageToVideoPipeline(WanImageToVideoPipeline):
pass
EntryClass = LingBotWorldImageToVideoPipeline
@@ -143,6 +143,8 @@ class ForwardBatch:
camera_trajectory: str | None = None # Camera trajectory file/identifier
action_list: list[str] | None = None # List of actions (e.g., ['forward', 'left'])
action_speed_list: list[float] | None = None # Speed for each action
# Camera control inputs (LingBotWorld)
c2ws_plucker_emb: torch.Tensor | None = None # Plucker embedding: [B, C, F_lat, H_lat, W_lat]
# Latent dimensions
height_latents: list[int] | int | None = None
+1
View File
@@ -173,6 +173,7 @@ class DenoisingStage(PipelineStage):
{
"mouse_cond": batch.mouse_cond,
"keyboard_cond": batch.keyboard_cond,
"c2ws_plucker_emb": batch.c2ws_plucker_emb,
},
)
+13
View File
@@ -23,6 +23,7 @@ from fastvideo.configs.pipelines.hunyuan15 import (
Hunyuan15T2V480PConfig, Hunyuan15I2V480PStepDistilledConfig,
Hunyuan15T2V720PConfig, Hunyuan15I2V720PConfig, Hunyuan15SR1080PConfig)
from fastvideo.configs.pipelines.hyworld import HYWorldConfig
from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig
from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig
from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
@@ -61,6 +62,7 @@ from fastvideo.configs.sample.hunyuan15 import (
from fastvideo.configs.sample.hyworld import HYWorld_SamplingParam
from fastvideo.configs.sample.hunyuangamecraft import HunyuanGameCraftSamplingParam
from fastvideo.configs.sample.ltx2 import LTX2SamplingParam
from fastvideo.configs.sample.lingbotworld import LingBotWorld_SamplingParam
from fastvideo.configs.sample.ltx2 import (LTX2BaseSamplingParam,
LTX2DistilledSamplingParam)
from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam
@@ -360,6 +362,17 @@ def _register_configs() -> None:
"tencent/Hunyuan-GameCraft-1.0",
],
model_detectors=[lambda path: "gamecraft" in path.lower()],
# LingBotWorld
register_configs(
sampling_param_cls=LingBotWorld_SamplingParam,
pipeline_config_cls=LingBotWorldI2V480PConfig,
hf_model_paths=[
"FastVideo/LingBot-World-Base-Cam-Diffusers",
],
model_detectors=[
lambda path:
("lingbotworld" in path.lower() or "lingbot-world" in path.lower())
],
)
# LongCat
@@ -0,0 +1,204 @@
# SPDX-License-Identifier: Apache-2.0
"""
SSIM-based similarity test for LingBotWorld I2V with camera control.
Camera trajectory is loaded from LingBot example npy files
(poses.npy/intrinsics.npy), matching the official script workflow.
Note: num_inference_steps is reduced to 4 for faster CI.
"""
import os
import pytest
import torch
from fastvideo import VideoGenerator
from fastvideo.logger import init_logger
from fastvideo.models.dits.lingbotworld.cam_utils import prepare_camera_embedding
from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results
logger = init_logger(__name__)
def _find_lingbotworld_examples_root() -> str | None:
script_dir = os.path.dirname(os.path.abspath(__file__))
repo_root = os.path.abspath(os.path.join(script_dir, "..", "..", ".."))
candidates = [
os.path.join(repo_root, "examples", "inference", "basic",
"lingbotworld_examples"),
os.path.join(repo_root, "..", "FastVideo", "examples", "inference",
"basic", "lingbotworld_examples"),
]
for candidate in candidates:
if (os.path.exists(os.path.join(candidate, "00", "poses.npy"))
and os.path.exists(
os.path.join(candidate, "00", "intrinsics.npy"))):
return os.path.abspath(candidate)
return None
device_name = torch.cuda.get_device_name()
device_reference_folder_suffix = "_reference_videos"
if "A40" in device_name:
device_reference_folder = "A40" + device_reference_folder_suffix
elif "L40S" in device_name:
device_reference_folder = "L40S" + device_reference_folder_suffix
elif "H100" in device_name:
device_reference_folder = "H100" + device_reference_folder_suffix
else:
device_reference_folder = None
logger.warning("Unsupported device for ssim tests: %s", device_name)
LINGBOT_PARAMS = {
"model_path": "FastVideo/LingBot-World-Base-Cam-Diffusers",
"num_gpus": 2,
"height": 256,
"width": 448,
"num_frames": 45, # must be 4k+1
"num_inference_steps": 4,
"guidance_scale": 5.0,
"guidance_scale_2": 5.0,
"embedded_cfg_scale": 6,
"flow_shift": 10.0,
"boundary_ratio": 0.947,
"seed": 42,
"fps": 16,
"spatial_scale": 8,
"example_case": "00",
"image_path": (
"https://raw.githubusercontent.com/Robbyant/lingbot-world/main/"
"examples/00/image.jpg"
),
"negative_prompt": (
"画面突变,色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,"
"最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,"
"畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走,"
"镜头晃动,画面闪烁,模糊,噪点,水印,签名,文字,变形,扭曲,液化,不合逻辑的结构,卡顿,"
"PPT幻灯片感,过暗,欠曝,低对比度,霓虹灯光感,过度锐化,3D渲染感,人物,行人,游客,身体,"
"皮肤,肢体,面部特征,汽车,电线"
),
}
TEST_PROMPTS = [
"The video presents a soaring journey through a fantasy jungle. The wind "
"whips past the rider's blue hands gripping the reins, causing the leather "
"straps to vibrate. The ancient gothic castle approaches steadily, its stone "
"details becoming clearer against the backdrop of floating islands and "
"distant waterfalls.",
]
@pytest.mark.parametrize("prompt", TEST_PROMPTS)
@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"])
def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str):
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
if device_reference_folder is None:
pytest.skip(f"Unsupported device for LingBot SSIM test: {device_name}")
if torch.cuda.device_count() < LINGBOT_PARAMS["num_gpus"]:
pytest.skip(
f"LingBot SSIM test requires {LINGBOT_PARAMS['num_gpus']} GPUs, "
f"but only {torch.cuda.device_count()} detected."
)
examples_root = _find_lingbotworld_examples_root()
if examples_root is None:
pytest.skip(
"lingbotworld_examples not found under examples/inference/basic.")
action_path = os.path.join(examples_root, LINGBOT_PARAMS["example_case"])
if not (os.path.exists(os.path.join(action_path, "poses.npy"))
and os.path.exists(os.path.join(action_path, "intrinsics.npy"))):
pytest.skip(f"Missing camera npy files under {action_path}")
c2ws_plucker_emb, aligned_num_frames = prepare_camera_embedding(
action_path=action_path,
num_frames=LINGBOT_PARAMS["num_frames"],
height=LINGBOT_PARAMS["height"],
width=LINGBOT_PARAMS["width"],
spatial_scale=LINGBOT_PARAMS["spatial_scale"],
)
script_dir = os.path.dirname(os.path.abspath(__file__))
model_id = "LingBot-World-Base-Cam-Diffusers"
output_dir = os.path.join(script_dir, "generated_videos", model_id,
ATTENTION_BACKEND)
output_video_name = f"{prompt[:100].strip()}.mp4"
os.makedirs(output_dir, exist_ok=True)
init_kwargs = {
"num_gpus": LINGBOT_PARAMS["num_gpus"],
"flow_shift": LINGBOT_PARAMS["flow_shift"],
"boundary_ratio": LINGBOT_PARAMS["boundary_ratio"],
"use_fsdp_inference": True,
"dit_cpu_offload": True,
"dit_layerwise_offload": False,
"text_encoder_cpu_offload": True,
"vae_cpu_offload": False,
"pin_cpu_memory": True,
}
generation_kwargs = {
"output_path": output_dir,
"image_path": LINGBOT_PARAMS["image_path"],
"height": LINGBOT_PARAMS["height"],
"width": LINGBOT_PARAMS["width"],
"num_frames": aligned_num_frames,
"num_inference_steps": LINGBOT_PARAMS["num_inference_steps"],
"guidance_scale": LINGBOT_PARAMS["guidance_scale"],
"guidance_scale_2": LINGBOT_PARAMS["guidance_scale_2"],
"embedded_cfg_scale": LINGBOT_PARAMS["embedded_cfg_scale"],
"seed": LINGBOT_PARAMS["seed"],
"fps": LINGBOT_PARAMS["fps"],
"negative_prompt": LINGBOT_PARAMS["negative_prompt"],
"c2ws_plucker_emb": c2ws_plucker_emb,
}
generator: VideoGenerator | None = None
try:
generator = VideoGenerator.from_pretrained(
model_path=LINGBOT_PARAMS["model_path"], **init_kwargs)
generator.generate_video(prompt, **generation_kwargs)
finally:
if generator is not None:
generator.shutdown()
generated_video_path = os.path.join(output_dir, output_video_name)
assert os.path.exists(generated_video_path), (
f"Output video was not generated at {generated_video_path}")
reference_folder = os.path.join(script_dir, device_reference_folder, model_id,
ATTENTION_BACKEND)
if not os.path.exists(reference_folder):
raise FileNotFoundError(
f"Reference video folder does not exist: {reference_folder}")
reference_video_name = None
for filename in os.listdir(reference_folder):
if filename.endswith(".mp4") and prompt[:100].strip() in filename:
reference_video_name = filename
break
if not reference_video_name:
raise FileNotFoundError(
f"Reference video missing for prompt/backend under {reference_folder}"
)
reference_video_path = os.path.join(reference_folder, reference_video_name)
logger.info("Computing SSIM between %s and %s", reference_video_path,
generated_video_path)
ssim_values = compute_video_ssim_torchvision(reference_video_path,
generated_video_path,
use_ms_ssim=True)
mean_ssim = ssim_values[0]
logger.info("SSIM mean value: %s", mean_ssim)
write_ssim_results(output_dir, ssim_values, reference_video_path,
generated_video_path,
LINGBOT_PARAMS["num_inference_steps"], prompt)
min_acceptable_ssim = 0.90
assert mean_ssim >= min_acceptable_ssim, (
f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} "
f"for {model_id} with backend {ATTENTION_BACKEND}")