Merge branch 'hunyuan-gamecraft-v2-port' of https://github.com/MihirJagtap/FastVideo into pr/1071
This commit is contained in:
@@ -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()
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -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"
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -173,6 +173,7 @@ class DenoisingStage(PipelineStage):
|
||||
{
|
||||
"mouse_cond": batch.mouse_cond,
|
||||
"keyboard_cond": batch.keyboard_cond,
|
||||
"c2ws_plucker_emb": batch.c2ws_plucker_emb,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
BIN
Binary file not shown.
@@ -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}")
|
||||
Reference in New Issue
Block a user