diff --git a/fastvideo/api/sampling_param.py b/fastvideo/api/sampling_param.py index 89d437cc7..c80310c1b 100644 --- a/fastvideo/api/sampling_param.py +++ b/fastvideo/api/sampling_param.py @@ -51,8 +51,9 @@ class SamplingParam: gt_latents: Any | None = None # Ground truth latents [B, 16, T, H, W] conditioning_mask: Any | None = None # Mask [B, 1, T, H, W] - # Camera control inputs (LingBotWorld) + # Camera control inputs (LingBotWorld and LingBotWorld2) c2ws_plucker_emb: Any | None = None # Plucker embedding: [B, C, F_lat, H_lat, W_lat] + action_path: str | None = None # Directory containing poses.npy and intrinsics.npy # Refine inputs (LongCat 480p->720p upscaling) # Path-based refine (load stage1 video from disk, e.g. MP4) diff --git a/fastvideo/api/schema.py b/fastvideo/api/schema.py index d3dbb0516..f676bfdcc 100644 --- a/fastvideo/api/schema.py +++ b/fastvideo/api/schema.py @@ -130,6 +130,7 @@ class InputConfig: keyboard_cond: Any | None = None grid_sizes: Any | None = None c2ws_plucker_emb: Any | None = None + action_path: str | None = None refine_from: str | None = None stage1_video: Any | None = None diff --git a/fastvideo/configs/models/dits/lingbotworld2.py b/fastvideo/configs/models/dits/lingbotworld2.py new file mode 100644 index 000000000..f57ee7035 --- /dev/null +++ b/fastvideo/configs/models/dits/lingbotworld2.py @@ -0,0 +1,55 @@ +# 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 LingBotWorld2CausalFastArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks]) + param_names_mapping: dict = field(default_factory=dict) + reverse_param_names_mapping: dict = field(default_factory=dict) + lora_param_names_mapping: dict = field(default_factory=dict) + + model_type: str = "i2v" + patch_size: tuple[int, int, int] = (1, 2, 2) + text_len: int = 512 + in_dim: int = 36 + dim: int = 5120 + ffn_dim: int = 13824 + freq_dim: int = 256 + text_dim: int = 4096 + out_dim: int = 16 + num_heads: int = 40 + num_layers: int = 40 + qk_norm: bool = True + cross_attn_norm: bool = True + eps: float = 1e-6 + + local_attn_size: int = 18 + sink_size: int = 6 + chunk_size: int = 4 + sample_shift: float = 10.0 + num_train_timesteps: int = 1000 + timesteps_index: tuple[int, int, int, int] = (0, 250, 500, 750) + max_area: int = 480 * 832 + + def __post_init__(self): + super().__post_init__() + self.hidden_size = self.dim + self.num_attention_heads = self.num_heads + self.attention_head_dim = self.dim // self.num_heads + self.in_channels = self.in_dim + self.out_channels = self.out_dim + self.num_channels_latents = self.out_dim + + +@dataclass +class LingBotWorld2CausalFastVideoConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=LingBotWorld2CausalFastArchConfig) + + prefix: str = "Wan" diff --git a/fastvideo/configs/models/encoders/lingbotworld2_t5.py b/fastvideo/configs/models/encoders/lingbotworld2_t5.py new file mode 100644 index 000000000..88443de5a --- /dev/null +++ b/fastvideo/configs/models/encoders/lingbotworld2_t5.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.encoders.base import ( + TextEncoderArchConfig, + TextEncoderConfig, +) + + +@dataclass +class LingBotWorld2UMT5ArchConfig(TextEncoderArchConfig): + architectures: list[str] = field(default_factory=lambda: ["LingBotWorld2T5EncoderModel"]) + vocab_size: int = 256384 + dim: int = 4096 + dim_attn: int = 4096 + dim_ffn: int = 10240 + num_heads: int = 64 + num_layers: int = 24 + num_buckets: int = 32 + text_len: int = 512 + hidden_size: int = 4096 + dropout: float = 0.1 + + def __post_init__(self) -> None: + super().__post_init__() + self.tokenizer_kwargs = { + "padding": "max_length", + "truncation": True, + "max_length": self.text_len, + "add_special_tokens": True, + "return_attention_mask": True, + "return_tensors": "pt", + } + + +@dataclass +class LingBotWorld2UMT5Config(TextEncoderConfig): + arch_config: TextEncoderArchConfig = field(default_factory=LingBotWorld2UMT5ArchConfig) + + prefix: str = "text_encoder" diff --git a/fastvideo/configs/pipelines/lingbotworld2.py b/fastvideo/configs/pipelines/lingbotworld2.py new file mode 100644 index 000000000..aa9d2cbd1 --- /dev/null +++ b/fastvideo/configs/pipelines/lingbotworld2.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +import html +from dataclasses import dataclass, field + +import ftfy +import torch + +from fastvideo.configs.models import DiTConfig +from fastvideo.configs.models.dits.lingbotworld2 import LingBotWorld2CausalFastVideoConfig +from fastvideo.configs.models.encoders import BaseEncoderOutput, LingBotWorld2UMT5Config +from fastvideo.configs.models.vaes import WanVAEConfig +from fastvideo.configs.pipelines.wan import Wan2_2_I2V_A14B_Config + + +def lingbotworld2_whitespace_preprocess(prompt: str) -> str: + text = ftfy.fix_text(prompt) + text = html.unescape(html.unescape(text)) + return " ".join(text.strip().split()) + + +def lingbotworld2_t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor: + assert outputs.last_hidden_state is not None + return outputs.last_hidden_state + + +@dataclass +class LingBotWorld2CausalFastI2V480PConfig(Wan2_2_I2V_A14B_Config): + dit_config: DiTConfig = field(default_factory=LingBotWorld2CausalFastVideoConfig) + vae_config: WanVAEConfig = field(default_factory=WanVAEConfig) + text_encoder_configs: tuple = field(default_factory=lambda: (LingBotWorld2UMT5Config(), )) + preprocess_text_funcs: tuple = field(default_factory=lambda: (lingbotworld2_whitespace_preprocess, )) + postprocess_text_funcs: tuple = field(default_factory=lambda: (lingbotworld2_t5_postprocess_text, )) + text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", )) + flow_shift: float | None = 10.0 + boundary_ratio: float | None = 0.947 + vae_precision: str = "fp32" + vae_decode_precision: str | None = "fp32" + vae_tiling: bool = False + vae_sp: bool = False + dit_precision: str = "bf16" + is_causal: bool = False + + def __post_init__(self) -> None: + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True diff --git a/fastvideo/models/dits/lingbotworld2/__init__.py b/fastvideo/models/dits/lingbotworld2/__init__.py new file mode 100644 index 000000000..c2aeafc2e --- /dev/null +++ b/fastvideo/models/dits/lingbotworld2/__init__.py @@ -0,0 +1,5 @@ +from .causal_fast import LingBotWorld2CausalFastTransformer3DModel + +__all__ = ["LingBotWorld2CausalFastTransformer3DModel"] + +EntryClass = LingBotWorld2CausalFastTransformer3DModel diff --git a/fastvideo/models/dits/lingbotworld2/cam_utils.py b/fastvideo/models/dits/lingbotworld2/cam_utils.py new file mode 100644 index 000000000..30521ddaa --- /dev/null +++ b/fastvideo/models/dits/lingbotworld2/cam_utils.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +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] + device = c2ws_mat.device + dtype = c2ws_mat.dtype + + grid_y, grid_x = torch.meshgrid( + torch.arange(height, device=device, dtype=dtype) + 0.5, + torch.arange(width, device=device, dtype=dtype) + 0.5, + indexing='ij', + ) + x_flat = grid_x.reshape(-1) + y_flat = grid_y.reshape(-1) + + fx, fy, cx, cy = Ks[0, 0], Ks[0, 1], Ks[0, 2], Ks[0, 3] + dirs = torch.stack([ + (x_flat - cx) / fx, + (y_flat - cy) / fy, + torch.ones_like(x_flat), + ], dim=-1) + dirs = dirs / dirs.norm(dim=-1, keepdim=True) + + rays_d = (c2ws_mat[:, :3, :3] @ dirs.T).transpose(1, 2) + rays_o = c2ws_mat[:, :3, 3].unsqueeze(1).expand_as(rays_d) + plucker_embeddings = torch.cat([rays_o, rays_d], dim=-1) + return plucker_embeddings.view(n_frames, height, width, 6) + + +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 diff --git a/fastvideo/models/dits/lingbotworld2/causal_fast.py b/fastvideo/models/dits/lingbotworld2/causal_fast.py new file mode 100644 index 000000000..c87b99f93 --- /dev/null +++ b/fastvideo/models/dits/lingbotworld2/causal_fast.py @@ -0,0 +1,776 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LingBot World 2 causal-fast DiT implemented inside FastVideo.""" + +import math +import warnings +from typing import Any + +from einops import rearrange +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.configs.models.dits.lingbotworld2 import ( + LingBotWorld2CausalFastVideoConfig, +) +from fastvideo.distributed.communication_op import ( + sequence_model_parallel_all_gather, + sequence_model_parallel_all_to_all_4D, +) +from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size +from fastvideo.models.dits.base import BaseDiT +from fastvideo.platforms import AttentionBackendEnum + +try: + import flash_attn_interface + + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + + +def is_blocks(n: str, m) -> bool: + return "blocks" in n and str.isdigit(n.split(".")[-1]) + + +def sinusoidal_embedding_1d(dim: int, position: torch.Tensor) -> torch.Tensor: + """Build Wan/LingBot World 2 sinusoidal timestep embeddings.""" + assert dim % 2 == 0 + half = dim // 2 + position = position.type(torch.float64) + sinusoid = torch.outer( + position, + torch.pow(10000, -torch.arange(half, device=position.device).to(position).div(half)), + ) + return torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + + +@torch.amp.autocast("cuda", enabled=False) +def rope_params(max_seq_len: int, dim: int, theta: int = 10000) -> torch.Tensor: + """Return complex RoPE frequencies used by the released LingBot World 2 model.""" + assert dim % 2 == 0 + freqs = torch.outer( + torch.arange(max_seq_len), + 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)), + ) + return torch.polar(torch.ones_like(freqs), freqs) + + +def flash_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_lens: torch.Tensor | None = None, + k_lens: torch.Tensor | None = None, + dropout_p: float = 0.0, + softmax_scale: float | None = None, + q_scale: float | None = None, + causal: bool = False, + window_size: tuple[int, int] = (-1, -1), + deterministic: bool = False, + dtype: torch.dtype = torch.bfloat16, + version: int | None = None, +) -> torch.Tensor: + """Run LingBot World 2-compatible FlashAttention on packed varlen inputs.""" + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == "cuda" and q.size(-1) <= 256 + + b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype + + def half(x: torch.Tensor) -> torch.Tensor: + return x if x.dtype in half_dtypes else x.to(dtype) + + if q_lens is None: + q = half(q.flatten(0, 1)) + q_lens = torch.tensor([lq] * b, dtype=torch.int32, device=q.device) + else: + q = half(torch.cat([u[:v] for u, v in zip(q, q_lens, strict=True)])) + + if k_lens is None: + k = half(k.flatten(0, 1)) + v = half(v.flatten(0, 1)) + k_lens = torch.tensor([lk] * b, dtype=torch.int32, device=k.device) + else: + k = half(torch.cat([u[:v] for u, v in zip(k, k_lens, strict=True)])) + v = half(torch.cat([u[:v] for u, v in zip(v, k_lens, strict=True)])) + + q = q.to(v.dtype) + k = k.to(v.dtype) + if q_scale is not None: + q = q * q_scale + + if version == 3 and not FLASH_ATTN_3_AVAILABLE: + warnings.warn("FlashAttention 3 is not available; using FlashAttention 2.") + + cu_q = torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum( + 0, dtype=torch.int32 + ).to(q.device, non_blocking=True) + cu_k = torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum( + 0, dtype=torch.int32 + ).to(k.device, non_blocking=True) + + if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE: + x = flash_attn_interface.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + seqused_q=None, + seqused_k=None, + max_seqlen_q=lq, + max_seqlen_k=lk, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic, + ).unflatten(0, (b, lq)) + else: + assert FLASH_ATTN_2_AVAILABLE + x = flash_attn.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=lq, + max_seqlen_k=lk, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + ).unflatten(0, (b, lq)) + return x.type(out_dtype) + + +def attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_lens: torch.Tensor | None = None, + k_lens: torch.Tensor | None = None, + dropout_p: float = 0.0, + softmax_scale: float | None = None, + q_scale: float | None = None, + causal: bool = False, + window_size: tuple[int, int] = (-1, -1), + deterministic: bool = False, + dtype: torch.dtype = torch.bfloat16, + fa_version: int | None = None, +) -> torch.Tensor: + """Dispatch LingBot World 2 attention to FlashAttention when available.""" + if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE: + return flash_attention( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + version=fa_version, + ) + + if q_lens is not None or k_lens is not None: + warnings.warn("Padding masks are disabled without FlashAttention.") + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + out = F.scaled_dot_product_attention( + q, k, v, attn_mask=None, is_causal=causal, dropout_p=dropout_p + ) + return out.transpose(1, 2).contiguous() + + +@torch.amp.autocast("cuda", enabled=False) +def causal_rope_apply( + x: torch.Tensor, + grid_sizes: torch.Tensor, + freqs: torch.Tensor, + start_frame: int = 0, +) -> torch.Tensor: + """Apply LingBot World 2 causal RoPE with the current chunk frame offset.""" + n, c = x.size(2), x.size(3) // 2 + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + output = [] + for i, (f, h, w) in enumerate(grid_sizes.tolist()): + seq_len = f * h * w + x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)) + freqs_i = torch.cat( + [ + freqs[0][start_frame : start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1), + freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(seq_len, 1, -1) + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + x_i = torch.cat([x_i, x[i, seq_len:]]) + output.append(x_i) + return torch.stack(output).type_as(x) + + +class WanRMSNorm(nn.Module): + """RMSNorm used by Wan/LingBot World 2 attention projections.""" + + def __init__(self, dim: int, eps: float = 1e-5): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Normalize the last dimension in fp32 and restore input dtype.""" + return self._norm(x.float()).type_as(x) * self.weight + + def _norm(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + +class WanLayerNorm(nn.LayerNorm): + """LayerNorm variant that computes in fp32 and returns the input dtype.""" + + def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = False): + super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply layer norm in fp32 for LingBot World 2 numerical parity.""" + return super().forward(x.float()).type_as(x) + + +class CausalWanSelfAttention(nn.Module): + """LingBot World 2 causal self-attention with rolling KV cache.""" + + def __init__( + self, + dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm: bool = True, + eps: float = 1e-6, + ): + super().__init__() + assert dim % num_heads == 0 + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.qk_norm = qk_norm + self.eps = eps + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + def forward( + self, + x: torch.Tensor, + seq_lens: torch.Tensor, + grid_sizes: torch.Tensor, + freqs: torch.Tensor, + kv_cache: dict, + current_start: int = 0, + max_attention_size: int = 1_000_000, + frame_seqlen: int | None = None, + seq_lens_int: int | None = None, + ) -> torch.Tensor: + """Project QKV, update the rolling cache, and attend to its active window.""" + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + + if frame_seqlen is None: + frame_seqlen = math.prod(grid_sizes[0][1:]).item() + current_start_frame = current_start // frame_seqlen + if seq_lens_int is None: + seq_lens_int = int(seq_lens[0].item() if seq_lens.dim() > 0 else seq_lens.item()) + + sp_size = get_sp_world_size() + if sp_size > 1: + q = sequence_model_parallel_all_to_all_4D(q, scatter_dim=2, gather_dim=1) + k = sequence_model_parallel_all_to_all_4D(k, scatter_dim=2, gather_dim=1) + v = sequence_model_parallel_all_to_all_4D(v, scatter_dim=2, gather_dim=1) + padded_seq_len = s * sp_size + roped_query = causal_rope_apply(q, grid_sizes, freqs, start_frame=current_start_frame).type_as(v) + roped_key = causal_rope_apply(k, grid_sizes, freqs, start_frame=current_start_frame).type_as(v) + roped_query = roped_query[:, :seq_lens_int] + roped_key = roped_key[:, :seq_lens_int] + v = v[:, :seq_lens_int] + num_new_tokens = seq_lens_int + else: + padded_seq_len = s + roped_query = causal_rope_apply(q, grid_sizes, freqs, start_frame=current_start_frame).type_as(v) + roped_key = causal_rope_apply(k, grid_sizes, freqs, start_frame=current_start_frame).type_as(v) + num_new_tokens = roped_query.shape[1] + + current_end = current_start + num_new_tokens + sink_tokens = self.sink_size * frame_seqlen + kv_cache_size = kv_cache["k"].shape[1] + + if self.local_attn_size == -1: + local_end_index = current_start + num_new_tokens + local_start_index = current_start + kv_cache["k"][:, local_start_index:local_end_index] = roped_key + kv_cache["v"][:, local_start_index:local_end_index] = v + elif (current_end > kv_cache["global_end_index"].item()) and ( + num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size + ): + num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size + num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens + kv_cache["k"][:, sink_tokens : sink_tokens + num_rolled_tokens] = kv_cache["k"][ + :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens + ].clone() + kv_cache["v"][:, sink_tokens : sink_tokens + num_rolled_tokens] = kv_cache["v"][ + :, sink_tokens + num_evicted_tokens : sink_tokens + num_evicted_tokens + num_rolled_tokens + ].clone() + local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() - num_evicted_tokens + local_start_index = local_end_index - num_new_tokens + kv_cache["k"][:, local_start_index:local_end_index] = roped_key + kv_cache["v"][:, local_start_index:local_end_index] = v + else: + local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() + local_start_index = local_end_index - num_new_tokens + kv_cache["k"][:, local_start_index:local_end_index] = roped_key + kv_cache["v"][:, local_start_index:local_end_index] = v + + k_cache = kv_cache["k"][:, max(0, local_end_index - max_attention_size) : local_end_index] + v_cache = kv_cache["v"][:, max(0, local_end_index - max_attention_size) : local_end_index] + x = attention(roped_query, k_cache, v_cache) + kv_cache["global_end_index"].fill_(current_end) + kv_cache["local_end_index"].fill_(local_end_index) + + if sp_size > 1: + sp_pad = padded_seq_len - seq_lens_int + if sp_pad > 0: + x = torch.cat([x, x.new_zeros(b, sp_pad, x.size(2), d)], dim=1) + x = sequence_model_parallel_all_to_all_4D(x, scatter_dim=1, gather_dim=2) + return self.o(x.flatten(2)) + +class WanCrossAttention(CausalWanSelfAttention): + """LingBot World 2 cross-attention with reusable text K/V cache.""" + + def forward( + self, + x: torch.Tensor, + context: torch.Tensor, + context_lens: torch.Tensor | None, + crossattn_cache: dict | None = None, + cross_attn_first_call: bool | None = None, + ) -> torch.Tensor: + """Attend hidden states to text context, populating cache on first use.""" + b, n, d = x.size(0), self.num_heads, self.head_dim + q = self.norm_q(self.q(x)).view(b, -1, n, d) + if crossattn_cache is not None: + is_first = crossattn_cache["is_init"].item() == 0 if cross_attn_first_call is None else cross_attn_first_call + if is_first: + crossattn_cache["is_init"].fill_(1) + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + crossattn_cache["k"].copy_(k) + crossattn_cache["v"].copy_(v) + else: + k = crossattn_cache["k"] + v = crossattn_cache["v"] + else: + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + x = flash_attention(q, k, v, k_lens=context_lens) + return self.o(x.flatten(2)) + + +class CausalWanAttentionBlock(nn.Module): + """One LingBot World 2 causal transformer block including camera injection.""" + + def __init__( + self, + dim: int, + ffn_dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm: bool = True, + cross_attn_norm: bool = False, + eps: float = 1e-6, + ): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.local_attn_size = local_attn_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = CausalWanSelfAttention(dim, num_heads, local_attn_size, sink_size, qk_norm, eps) + self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WanCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim)) + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + self.cam_injector_layer1 = nn.Linear(dim, dim) + self.cam_injector_layer2 = nn.Linear(dim, dim) + self.cam_scale_layer = nn.Linear(dim, dim) + self.cam_shift_layer = nn.Linear(dim, dim) + + def forward( + self, + x: torch.Tensor, + e: torch.Tensor, + seq_lens: torch.Tensor, + grid_sizes: torch.Tensor, + freqs: torch.Tensor, + context: torch.Tensor, + context_lens: torch.Tensor | None, + dit_cond_dict: dict[str, Any] | None = None, + kv_cache: dict | None = None, + crossattn_cache: dict | None = None, + current_start: int = 0, + max_attention_size: int = 1_000_000, + frame_seqlen: int | None = None, + cross_attn_first_call: bool | None = None, + seq_lens_int: int | None = None, + ) -> torch.Tensor: + """Apply self-attention, camera modulation, cross-attention, and FFN.""" + assert kv_cache is not None + assert e.dtype == torch.float32 + with torch.amp.autocast("cuda", dtype=torch.float32): + e = (self.modulation.unsqueeze(0) + e).chunk(6, dim=2) + + y = self.self_attn( + self.norm1(x).float() * (1 + e[1].squeeze(2)) + e[0].squeeze(2), + seq_lens, + grid_sizes, + freqs, + kv_cache, + current_start, + max_attention_size, + frame_seqlen=frame_seqlen, + seq_lens_int=seq_lens_int, + ) + with torch.amp.autocast("cuda", dtype=torch.float32): + x = x + y * e[2].squeeze(2) + + if dit_cond_dict is not None and "c2ws_plucker_emb" in dit_cond_dict: + c2ws_plucker_emb = dit_cond_dict["c2ws_plucker_emb"] + c2ws_hidden_states = self.cam_injector_layer2( + F.silu(self.cam_injector_layer1(c2ws_plucker_emb)) + ) + c2ws_hidden_states = c2ws_hidden_states + c2ws_plucker_emb + x = (1.0 + self.cam_scale_layer(c2ws_hidden_states)) * x + self.cam_shift_layer(c2ws_hidden_states) + + x = x + self.cross_attn( + self.norm3(x), + context, + context_lens, + crossattn_cache=crossattn_cache, + cross_attn_first_call=cross_attn_first_call, + ) + y = self.ffn(self.norm2(x).float() * (1 + e[4].squeeze(2)) + e[3].squeeze(2)) + with torch.amp.autocast("cuda", dtype=torch.float32): + x = x + y * e[5].squeeze(2) + return x + + +class CausalHead(nn.Module): + """Output projection head for LingBot World 2 causal-fast DiT.""" + + def __init__(self, dim: int, out_dim: int, patch_size: tuple[int, int, int], eps: float = 1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, math.prod(patch_size) * out_dim) + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x: torch.Tensor, e: torch.Tensor) -> torch.Tensor: + """Normalize, modulate, and project hidden states to latent patches.""" + assert e.dtype == torch.float32 + with torch.amp.autocast("cuda", dtype=torch.float32): + e = (self.modulation.unsqueeze(0) + e.unsqueeze(2)).chunk(2, dim=2) + x = self.head(self.norm(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2)) + return x + + +class LingBotWorld2CausalFastTransformer3DModel(BaseDiT): + """Released LingBot World 2 14B causal-fast model with native FastVideo loading.""" + + _fsdp_shard_conditions = [is_blocks] + _compile_conditions: list = [] + _supported_attention_backends = ( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ) + param_names_mapping: dict = {} + reverse_param_names_mapping: dict = {} + lora_param_names_mapping: dict = {} + + def __init__(self, config: LingBotWorld2CausalFastVideoConfig, hf_config: dict[str, Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + self.model_type = config.model_type + self.patch_size = tuple(config.patch_size) + self.text_len = config.text_len + self.in_dim = config.in_dim + self.dim = config.dim + self.hidden_size = config.dim + self.ffn_dim = config.ffn_dim + self.freq_dim = config.freq_dim + self.text_dim = config.text_dim + self.out_dim = config.out_dim + self.out_channels = config.out_dim + self.num_heads = config.num_heads + self.num_attention_heads = config.num_heads + self.attention_head_dim = config.dim // config.num_heads + self.num_layers = config.num_layers + self.local_attn_size = config.local_attn_size + self.sink_size = config.sink_size + self.qk_norm = config.qk_norm + self.cross_attn_norm = config.cross_attn_norm + self.eps = config.eps + self.num_channels_latents = config.out_dim + + control_dim = 6 + self.patch_embedding = nn.Conv3d(self.in_dim, self.dim, kernel_size=self.patch_size, stride=self.patch_size) + self.patch_embedding_wancamctrl = nn.Linear( + control_dim * 64 * self.patch_size[0] * self.patch_size[1] * self.patch_size[2], + self.dim, + ) + self.c2ws_hidden_states_layer1 = nn.Linear(self.dim, self.dim) + self.c2ws_hidden_states_layer2 = nn.Linear(self.dim, self.dim) + self.text_embedding = nn.Sequential( + nn.Linear(self.text_dim, self.dim), + nn.GELU(approximate="tanh"), + nn.Linear(self.dim, self.dim), + ) + self.time_embedding = nn.Sequential( + nn.Linear(self.freq_dim, self.dim), + nn.SiLU(), + nn.Linear(self.dim, self.dim), + ) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(self.dim, self.dim * 6)) + self.blocks = nn.ModuleList( + [ + CausalWanAttentionBlock( + self.dim, + self.ffn_dim, + self.num_heads, + self.local_attn_size, + self.sink_size, + self.qk_norm, + self.cross_attn_norm, + self.eps, + ) + for _ in range(self.num_layers) + ] + ) + self.head = CausalHead(self.dim, self.out_dim, self.patch_size, self.eps) + self.freqs: torch.Tensor | None = None + self.init_weights() + self.__post_init__() + + def _get_freqs(self, device: torch.device) -> torch.Tensor: + """Materialize the non-persistent RoPE frequency table outside meta init.""" + if self.freqs is None or self.freqs.is_meta or self.freqs.device != device: + d = self.dim // self.num_heads + self.freqs = torch.cat( + [ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + ], + dim=1, + ).to(device) + return self.freqs + + def forward( + self, + hidden_states: torch.Tensor | list[torch.Tensor] | None = None, + encoder_hidden_states: torch.Tensor | list[torch.Tensor] | None = None, + timestep: torch.Tensor | None = None, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + guidance=None, + *, + x: list[torch.Tensor] | None = None, + t: torch.Tensor | None = None, + context: list[torch.Tensor] | torch.Tensor | None = None, + seq_len: int | None = None, + y: list[torch.Tensor] | None = None, + dit_cond_dict: dict[str, Any] | None = None, + kv_cache: list[dict] | None = None, + crossattn_cache: list[dict] | None = None, + current_start: int = 0, + max_attention_size: int = 1_000_000, + frame_seqlen: int | None = None, + cross_attn_first_call: bool | None = None, + **kwargs, + ) -> list[torch.Tensor]: + """Run one cached causal-fast DiT forward using the released LingBot World 2 ABI.""" + del encoder_hidden_states_image, guidance, kwargs + if x is None: + assert isinstance(hidden_states, torch.Tensor) + x = [hidden_states[0]] + if t is None: + assert timestep is not None + t = timestep + if context is None: + context = encoder_hidden_states + if isinstance(context, torch.Tensor): + context = [u for u in context] + assert context is not None + assert seq_len is not None + assert kv_cache is not None + assert crossattn_cache is not None + if self.model_type == "i2v": + assert y is not None + + device = self.patch_embedding.weight.device + freqs = self._get_freqs(device) + if y is not None: + x = [torch.cat([u, v], dim=0) for u, v in zip(x, y, strict=True)] + + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long, device=u.device) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long, device=device) + assert seq_lens.max() <= seq_len + x = torch.cat(x) + seq_lens_int = int(seq_lens[0].item()) + sp_size = get_sp_world_size() + sp_rank = get_sp_parallel_rank() + padded_seq_len = ((seq_lens_int + sp_size - 1) // sp_size) * sp_size + sp_pad_len = padded_seq_len - seq_lens_int + if sp_pad_len > 0: + x = torch.cat([x, x.new_zeros(x.size(0), sp_pad_len, x.size(2))], dim=1) + + if t.dim() == 1: + t = t.expand(t.size(0), padded_seq_len) + with torch.amp.autocast("cuda", dtype=torch.float32): + bt = t.size(0) + t = t.flatten() + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, t).unflatten(0, (bt, padded_seq_len)).float() + ) + e0 = self.time_projection(e).unflatten(2, (6, self.dim)) + + context_lens = None + context = self.text_embedding( + torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context]) + ) + + if dit_cond_dict is not None and "c2ws_plucker_emb" in dit_cond_dict: + c2ws_plucker_emb = dit_cond_dict["c2ws_plucker_emb"] + c2ws_plucker_emb = [ + rearrange( + i, + "1 c (f c1) (h c2) (w c3) -> 1 (f h w) (c c1 c2 c3)", + c1=self.patch_size[0], + c2=self.patch_size[1], + c3=self.patch_size[2], + ) + for i in c2ws_plucker_emb + ] + c2ws_plucker_emb = torch.cat(c2ws_plucker_emb, dim=1) + c2ws_plucker_emb = self.patch_embedding_wancamctrl(c2ws_plucker_emb) + c2ws_hidden_states = self.c2ws_hidden_states_layer2( + F.silu(self.c2ws_hidden_states_layer1(c2ws_plucker_emb)) + ) + c2ws_plucker_emb = c2ws_plucker_emb + c2ws_hidden_states + cam_len = c2ws_plucker_emb.size(1) + if cam_len < padded_seq_len: + c2ws_plucker_emb = torch.cat( + [ + c2ws_plucker_emb, + c2ws_plucker_emb.new_zeros( + c2ws_plucker_emb.size(0), + padded_seq_len - cam_len, + c2ws_plucker_emb.size(2), + ), + ], + dim=1, + ) + elif cam_len > padded_seq_len: + c2ws_plucker_emb = c2ws_plucker_emb[:, :padded_seq_len, :] + if sp_size > 1: + c2ws_plucker_emb = torch.chunk(c2ws_plucker_emb, sp_size, dim=1)[sp_rank] + dit_cond_dict = dict(dit_cond_dict) + dit_cond_dict["c2ws_plucker_emb"] = c2ws_plucker_emb + + if sp_size > 1: + x = torch.chunk(x, sp_size, dim=1)[sp_rank] + e = torch.chunk(e, sp_size, dim=1)[sp_rank] + e0 = torch.chunk(e0, sp_size, dim=1)[sp_rank] + + for block_index, block in enumerate(self.blocks): + x = block( + x, + e=e0, + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=freqs, + context=context, + context_lens=context_lens, + dit_cond_dict=dit_cond_dict, + kv_cache=kv_cache[block_index], + crossattn_cache=crossattn_cache[block_index], + current_start=current_start, + max_attention_size=max_attention_size, + frame_seqlen=frame_seqlen, + cross_attn_first_call=cross_attn_first_call, + seq_lens_int=seq_lens_int, + ) + + x = self.head(x, e) + if sp_size > 1: + x = sequence_model_parallel_all_gather(x, dim=1) + return [u.float() for u in self.unpatchify(x, grid_sizes)] + + def unpatchify(self, x: torch.Tensor, grid_sizes: torch.Tensor) -> list[torch.Tensor]: + """Reconstruct latent videos from flattened patch tokens.""" + c = self.out_dim + out = [] + for u, v in zip(x, grid_sizes.tolist(), strict=True): + u = u[: math.prod(v)].view(*v, *self.patch_size, c) + u = torch.einsum("fhwpqrc->cfphqwr", u) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size, strict=True)]) + out.append(u) + return out + + def init_weights(self) -> None: + """Initialize modules for non-meta construction; checkpoint load overwrites them.""" + if self.patch_embedding.weight.is_meta: + return + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=0.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=0.02) + nn.init.zeros_(self.head.head.weight) + + +EntryClass = LingBotWorld2CausalFastTransformer3DModel diff --git a/fastvideo/models/encoders/lingbotworld2_t5.py b/fastvideo/models/encoders/lingbotworld2_t5.py new file mode 100644 index 000000000..6fa619ca8 --- /dev/null +++ b/fastvideo/models/encoders/lingbotworld2_t5.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LingBot World 2 UMT5 encoder with the released checkpoint's module names.""" + +from collections.abc import Iterable +import html +import math +import string + +import ftfy +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.configs.models.encoders import BaseEncoderOutput +from fastvideo.configs.models.encoders.lingbotworld2_t5 import LingBotWorld2UMT5Config +from fastvideo.models.encoders.base import TextEncoder +from fastvideo.models.loader.weight_utils import default_weight_loader + + +def basic_clean(text: str) -> str: + """Apply LingBot World 2's ftfy/html cleanup before tokenization.""" + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + + +def whitespace_clean(text: str) -> str: + """Collapse all whitespace runs to single spaces.""" + return " ".join(text.split()) + + +def canonicalize(text: str, keep_punctuation_exact_string: str | None = None) -> str: + """Normalize prompts with LingBot World 2's optional punctuation handling.""" + text = text.replace("_", " ") + if keep_punctuation_exact_string: + text = keep_punctuation_exact_string.join( + part.translate(str.maketrans("", "", string.punctuation)) + for part in text.split(keep_punctuation_exact_string) + ) + else: + text = text.translate(str.maketrans("", "", string.punctuation)) + return " ".join(text.lower().split()) + + +def lingbotworld2_whitespace_preprocess(prompt: str) -> str: + """Match the LingBot World 2 source tokenizer's `clean='whitespace'` behavior.""" + return whitespace_clean(basic_clean(prompt)) + + +class GELU(nn.Module): + """T5 gated GELU approximation used by the source checkpoint.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply the tanh GELU approximation.""" + return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) + + +class T5LayerNorm(nn.Module): + """T5 RMS-style layer norm with source-compatible parameter name.""" + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Normalize in fp32 and apply the learned scale.""" + x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps) + if self.weight.dtype in (torch.float16, torch.bfloat16): + x = x.type_as(self.weight) + return self.weight * x + + +class T5Attention(nn.Module): + """LingBot World 2 source T5 attention block.""" + + def __init__(self, dim: int, dim_attn: int, num_heads: int, dropout: float = 0.1): + super().__init__() + assert dim_attn % num_heads == 0 + self.dim = dim + self.dim_attn = dim_attn + self.num_heads = num_heads + self.head_dim = dim_attn // num_heads + self.q = nn.Linear(dim, dim_attn, bias=False) + self.k = nn.Linear(dim, dim_attn, bias=False) + self.v = nn.Linear(dim, dim_attn, bias=False) + self.o = nn.Linear(dim_attn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward( + self, + x: torch.Tensor, + context: torch.Tensor | None = None, + mask: torch.Tensor | None = None, + pos_bias: torch.Tensor | None = None, + ) -> torch.Tensor: + """Project QKV, add relative bias/mask, and return attended states.""" + context = x if context is None else context + b, n, c = x.size(0), self.num_heads, self.head_dim + q = self.q(x).view(b, -1, n, c) + k = self.k(context).view(b, -1, n, c) + v = self.v(context).view(b, -1, n, c) + attn_bias = x.new_zeros(b, n, q.size(1), k.size(1)) + if pos_bias is not None: + attn_bias += pos_bias + if mask is not None: + assert mask.ndim in (2, 3) + mask = mask.view(b, 1, 1, -1) if mask.ndim == 2 else mask.unsqueeze(1) + attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min) + attn = torch.einsum("binc,bjnc->bnij", q, k) + attn_bias + attn = F.softmax(attn.float(), dim=-1).type_as(attn) + x = torch.einsum("bnij,bjnc->binc", attn, v) + return self.dropout(self.o(x.reshape(b, -1, n * c))) + + +class T5FeedForward(nn.Module): + """LingBot World 2 source T5 gated feed-forward block.""" + + def __init__(self, dim: int, dim_ffn: int, dropout: float = 0.1): + super().__init__() + self.dim = dim + self.dim_ffn = dim_ffn + self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU()) + self.fc1 = nn.Linear(dim, dim_ffn, bias=False) + self.fc2 = nn.Linear(dim_ffn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply the gated feed-forward projection.""" + x = self.fc1(x) * self.gate(x) + x = self.dropout(x) + x = self.fc2(x) + return self.dropout(x) + + +class T5RelativeEmbedding(nn.Module): + """Per-block relative position embedding used by LingBot World 2 UMT5.""" + + def __init__(self, num_buckets: int, num_heads: int, bidirectional: bool, max_dist: int = 128): + super().__init__() + self.num_buckets = num_buckets + self.num_heads = num_heads + self.bidirectional = bidirectional + self.max_dist = max_dist + self.embedding = nn.Embedding(num_buckets, num_heads) + + def forward(self, lq: int, lk: int) -> torch.Tensor: + """Build a relative-position bias tensor for attention logits.""" + device = self.embedding.weight.device + rel_pos = torch.arange(lk, device=device).unsqueeze(0) - torch.arange(lq, device=device).unsqueeze(1) + rel_pos = self._relative_position_bucket(rel_pos) + rel_pos_embeds = self.embedding(rel_pos) + return rel_pos_embeds.permute(2, 0, 1).unsqueeze(0).contiguous() + + def _relative_position_bucket(self, rel_pos: torch.Tensor) -> torch.Tensor: + """Map token offsets to T5 relative-position buckets.""" + if self.bidirectional: + num_buckets = self.num_buckets // 2 + rel_buckets = (rel_pos > 0).long() * num_buckets + rel_pos = torch.abs(rel_pos) + else: + num_buckets = self.num_buckets + rel_buckets = 0 + rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos)) + max_exact = num_buckets // 2 + rel_pos_large = max_exact + ( + torch.log(rel_pos.float() / max_exact) / math.log(self.max_dist / max_exact) * (num_buckets - max_exact) + ).long() + rel_pos_large = torch.min(rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1)) + rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large) + return rel_buckets + + +class T5SelfAttention(nn.Module): + """One source-compatible UMT5 encoder block.""" + + def __init__( + self, + dim: int, + dim_attn: int, + dim_ffn: int, + num_heads: int, + num_buckets: int, + shared_pos: bool = False, + dropout: float = 0.1, + ): + super().__init__() + self.norm1 = T5LayerNorm(dim) + self.attn = T5Attention(dim, dim_attn, num_heads, dropout) + self.norm2 = T5LayerNorm(dim) + self.ffn = T5FeedForward(dim, dim_ffn, dropout) + self.pos_embedding = None if shared_pos else T5RelativeEmbedding(num_buckets, num_heads, bidirectional=True) + + def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None, pos_bias: torch.Tensor | None = None) -> torch.Tensor: + """Run self-attention and feed-forward residual updates.""" + e = pos_bias if self.pos_embedding is None else self.pos_embedding(x.size(1), x.size(1)) + x = self._fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e)) + return self._fp16_clamp(x + self.ffn(self.norm2(x))) + + @staticmethod + def _fp16_clamp(x: torch.Tensor) -> torch.Tensor: + if x.dtype == torch.float16 and torch.isinf(x).any(): + clamp = torch.finfo(x.dtype).max - 1000 + x = torch.clamp(x, min=-clamp, max=clamp) + return x + + +class LingBotWorld2T5EncoderModel(TextEncoder): + """FastVideo-native LingBot World 2 UMT5 encoder for the released `.pth` weights.""" + + fall_back_to_pt_during_load = True + allow_patterns_overrides = ["*.pt"] + + def __init__(self, config: LingBotWorld2UMT5Config, prefix: str = ""): + super().__init__(config) + del prefix + arch = config.arch_config + self.token_embedding = nn.Embedding(arch.vocab_size, arch.dim) + self.dropout = nn.Dropout(arch.dropout) + self.blocks = nn.ModuleList( + [ + T5SelfAttention( + arch.dim, + arch.dim_attn, + arch.dim_ffn, + arch.num_heads, + arch.num_buckets, + shared_pos=False, + dropout=arch.dropout, + ) + for _ in range(arch.num_layers) + ] + ) + self.norm = T5LayerNorm(arch.dim) + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> BaseEncoderOutput: + """Encode token ids and return source-compatible hidden states.""" + del position_ids, inputs_embeds, output_hidden_states, kwargs + assert input_ids is not None + x = self.dropout(self.token_embedding(input_ids)) + for block in self.blocks: + x = block(x, attention_mask) + x = self.dropout(self.norm(x)) + return BaseEncoderOutput(last_hidden_state=x, attention_mask=attention_mask) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load source `.pth` weights whose names already match this module.""" + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +EntryClass = LingBotWorld2T5EncoderModel diff --git a/fastvideo/models/vaes/lingbotworld2_wanvae.py b/fastvideo/models/vaes/lingbotworld2_wanvae.py new file mode 100644 index 000000000..0c8e9485c --- /dev/null +++ b/fastvideo/models/vaes/lingbotworld2_wanvae.py @@ -0,0 +1,722 @@ +import logging +from types import SimpleNamespace + +import torch +import torch.cuda.amp as amp +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +__all__ = [ + 'Wan2_1_VAE', + 'LingBotWorld2WanVAE', +] + +CACHE_T = 2 + + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = (self.padding[2], self.padding[2], self.padding[1], + self.padding[1], 2 * self.padding[0], 0) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + + return super().forward(x) + + +class RMS_norm(nn.Module): + + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. + + def forward(self, x): + return F.normalize( + x, dim=(1 if self.channel_first else + -1)) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + + def forward(self, x): + """ + Fix bfloat16 support for nearest neighbor interpolation. + """ + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + + def __init__(self, dim, mode): + assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d', + 'downsample3d') + super().__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == 'upsample2d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + elif mode == 'upsample3d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + self.time_conv = CausalConv3d( + dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + + elif mode == 'downsample2d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == 'downsample3d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d( + dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) + + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == 'upsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = 'Rep' + feat_idx[0] += 1 + else: + + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] != 'Rep': + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] == 'Rep': + cache_x = torch.cat([ + torch.zeros_like(cache_x).to(cache_x.device), + cache_x + ], + dim=2) + if feat_cache[idx] == 'Rep': + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), + 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.resample(x) + x = rearrange(x, '(b t) c h w -> b c t h w', t=t) + + if self.mode == 'downsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + + cache_x = x[:, :, -1:, :, :].clone() + # if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx]!='Rep': + # # cache last frame of last two chunk + # cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2) + + x = self.time_conv( + torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + def init_weight(self, conv): + conv_weight = conv.weight + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + one_matrix = torch.eye(c1, c2) + init_matrix = one_matrix + nn.init.zeros_(conv_weight) + #conv_weight.data[:,:,-1,1,1] = init_matrix * 0.5 + conv_weight.data[:, :, 1, 0, 0] = init_matrix #* 0.5 + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + def init_weight2(self, conv): + conv_weight = conv.weight.data + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1 // 2, c2) + #init_matrix = repeat(init_matrix, 'o ... -> (o 2) ...').permute(1,0,2).contiguous().reshape(c1,c2) + conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix + conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + +class ResidualBlock(nn.Module): + + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # layers + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1)) + self.shortcut = CausalConv3d(in_dim, out_dim, 1) \ + if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = self.shortcut(x) + for layer in self.residual: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class AttentionBlock(nn.Module): + """ + Causal self-attention with a single head. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + + # layers + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + # zero out the last layer params + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.norm(x) + # compute query, key, value + q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, + -1).permute(0, 1, 3, + 2).contiguous().chunk( + 3, dim=-1) + + # apply attention + x = F.scaled_dot_product_attention( + q, + k, + v, + ) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + + # output + x = self.proj(x) + x = rearrange(x, '(b t) c h w-> b c t h w', t=t) + return x + identity + + +class Encoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + for _ in range(num_res_blocks): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + downsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # downsample block + if i != len(dim_mult) - 1: + mode = 'downsample3d' if temperal_downsample[ + i] else 'downsample2d' + downsamples.append(Resample(out_dim, mode=mode)) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout)) + + # output blocks + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +class Decoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2**(len(dim_mult) - 2) + + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential( + ResidualBlock(dims[0], dims[0], dropout), AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if i == 1 or i == 2 or i == 3: + in_dim = in_dim // 2 + for _ in range(num_res_blocks + 1): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + upsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # upsample block + if i != len(dim_mult) - 1: + mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d' + upsamples.append(Resample(out_dim, mode=mode)) + scale *= 2.0 + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, 3, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + ## conv1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## middle + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def count_conv3d(model): + count = 0 + for m in model.modules(): + if isinstance(m, CausalConv3d): + count += 1 + return count + + +class WanVAE_(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout) + + def forward(self, x): + mu, log_var = self.encode(x) + z = self.reparameterize(mu, log_var) + x_recon = self.decode(z) + return x_recon, mu, log_var + + def encode(self, x, scale): + self.clear_cache() + ## cache + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder( + x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + else: + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + mu = (mu - scale[0]) * scale[1] + self.clear_cache() + return mu + + def decode(self, z, scale): + self.clear_cache() + # z: [b,c,t,h,w] + if isinstance(scale[0], torch.Tensor): + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + else: + out_ = self.decoder( + x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) + self.clear_cache() + return out + + def reparameterize(self, mu, log_var): + std = torch.exp(0.5 * log_var) + eps = torch.randn_like(std) + return eps * std + mu + + def sample(self, imgs, deterministic=False): + mu, log_var = self.encode(imgs) + if deterministic: + return mu + std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) + return mu + std * torch.randn_like(std) + + def clear_cache(self): + self._conv_num = count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + #cache encode + self._enc_conv_num = count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **kwargs): + """ + Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL. + """ + # params + cfg = dict( + dim=96, + z_dim=z_dim, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0) + cfg.update(**kwargs) + + # init model + with torch.device('meta'): + model = WanVAE_(**cfg) + + # load checkpoint + logging.info(f'loading {pretrained_path}') + model.load_state_dict( + torch.load(pretrained_path, map_location=device), assign=True) + + return model + + +class Wan2_1_VAE: + + def __init__(self, + z_dim=16, + vae_pth='cache/vae_step_411000.pth', + dtype=torch.float, + device="cuda"): + self.dtype = dtype + self.device = device + + mean = [ + -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, + 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 + ] + std = [ + 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, + 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 + ] + self.mean = torch.tensor(mean, dtype=dtype, device=device) + self.std = torch.tensor(std, dtype=dtype, device=device) + self.scale = [self.mean, 1.0 / self.std] + + # init model + self.model = _video_vae( + pretrained_path=vae_pth, + z_dim=z_dim, + ).eval().requires_grad_(False).to(device) + + def encode(self, videos): + """ + videos: A list of videos each with shape [C, T, H, W]. + """ + with amp.autocast(dtype=self.dtype): + return [ + self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0) + for u in videos + ] + + def decode(self, zs): + with amp.autocast(dtype=self.dtype): + return [ + self.model.decode(u.unsqueeze(0), + self.scale).float().clamp_(-1, 1).squeeze(0) + for u in zs + ] + + +class LingBotWorld2WanVAE(nn.Module): + """FastVideo-facing wrapper around the exact LingBot World 2 Wan2.1 VAE computation.""" + + handles_latent_denorm = True + + def __init__(self, config, checkpoint_path=None, dtype=torch.float): + """Load the official LingBot World 2 VAE weights and expose FastVideo VAE APIs.""" + super().__init__() + self.config = config + self.dtype = dtype + z_dim = int(getattr(config, "z_dim", 16)) + mean = torch.tensor(getattr(config, "latents_mean"), dtype=dtype) + std = torch.tensor(getattr(config, "latents_std"), dtype=dtype) + self.register_buffer("shift_factor", mean.view(1, z_dim, 1, 1, 1), persistent=False) + self.register_buffer("scaling_factor", (1.0 / std).view(1, z_dim, 1, 1, 1), persistent=False) + self.scale = [mean, 1.0 / std] + + if checkpoint_path is None: + self.model = WanVAE_(dim=96, z_dim=z_dim, dim_mult=[1, 2, 4, 4], + num_res_blocks=2, attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0) + else: + self.model = _video_vae(pretrained_path=checkpoint_path, z_dim=z_dim) + self.model.eval().requires_grad_(False) + + def _scale_for(self, device: torch.device) -> list[torch.Tensor]: + """Return source VAE scale tensors on the active device.""" + return [u.to(device) for u in self.scale] + + def encode(self, videos: torch.Tensor): + """Encode `[B,C,T,H,W]` videos and return a FastVideo-style mean tensor.""" + if videos.ndim != 5: + raise ValueError(f"LingBotWorld2WanVAE.encode expects 5D input, got {videos.shape}") + scale = self._scale_for(videos.device) + latents = [] + with amp.autocast(dtype=self.dtype): + for video in videos: + latent = self.model.encode(video.unsqueeze(0), scale).float().squeeze(0) + latents.append(latent) + normalized = torch.stack(latents, dim=0) + return SimpleNamespace(mean=normalized) + + def decode(self, latents: torch.Tensor) -> torch.Tensor: + """Decode normalized LingBot World 2 latents to clamped `[-1,1]` video tensors.""" + if latents.ndim != 5: + raise ValueError(f"LingBotWorld2WanVAE.decode expects 5D input, got {latents.shape}") + scale = self._scale_for(latents.device) + videos = [] + with amp.autocast(dtype=self.dtype): + for latent in latents: + video = self.model.decode(latent.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0) + videos.append(video) + return torch.stack(videos, dim=0) + + +EntryClass = LingBotWorld2WanVAE diff --git a/fastvideo/pipelines/basic/lingbotworld2/__init__.py b/fastvideo/pipelines/basic/lingbotworld2/__init__.py new file mode 100644 index 000000000..a1f16d26c --- /dev/null +++ b/fastvideo/pipelines/basic/lingbotworld2/__init__.py @@ -0,0 +1,5 @@ +from .causal_fast_pipeline import LingBotWorld2CausalFastPipeline + +__all__ = ["LingBotWorld2CausalFastPipeline"] + +EntryClass = LingBotWorld2CausalFastPipeline diff --git a/fastvideo/pipelines/basic/lingbotworld2/causal_fast_pipeline.py b/fastvideo/pipelines/basic/lingbotworld2/causal_fast_pipeline.py new file mode 100644 index 000000000..9cafb634b --- /dev/null +++ b/fastvideo/pipelines/basic/lingbotworld2/causal_fast_pipeline.py @@ -0,0 +1,370 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LingBot World 2 causal-fast image-to-video pipeline.""" + +import math +import os + +from einops import rearrange +import numpy as np +import torch +import torchvision.transforms.functional as TF + +from fastvideo.distributed import get_local_torch_device +from fastvideo.distributed.parallel_state import get_sp_world_size +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.dits.lingbotworld2.cam_utils import ( + compute_relative_poses, + get_Ks_transformed, + get_plucker_embeddings, + interpolate_camera_poses, +) +from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler, +) +from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages import ( + ConditioningStage, + DecodingStage, + InputValidationStage, + TextEncodingStage, +) +from fastvideo.pipelines.stages.base import PipelineStage + +logger = init_logger(__name__) + + +class LingBotWorld2TextEncodingStage(TextEncodingStage): + """Keep LingBot World 2 T5 attention masks so DiT context matches the source pipeline.""" + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + """Initialize mask storage before running the shared text-encoding stage.""" + if batch.prompt_attention_mask is None: + batch.prompt_attention_mask = [] + return super().forward(batch, fastvideo_args) + + +class LingBotWorld2CausalFastGenerationStage(PipelineStage): + """Prepare LingBot World 2 conditions and run the released causal-fast sampling loop.""" + + def __init__(self, transformer, scheduler, vae) -> None: + super().__init__() + self.transformer = transformer + self.scheduler = scheduler + self.vae = vae + self._cross_attn_initialized = False + + def _convert_flow_pred_to_x0( + self, + flow_pred: torch.Tensor, + xt: torch.Tensor, + timestep: torch.Tensor, + ) -> torch.Tensor: + """Convert LingBot World 2 flow prediction to x0 using the scheduler sigma.""" + original_dtype = flow_pred.dtype + flow_pred, xt, sigmas, timesteps = map( + lambda x: x.double().to(flow_pred.device), + [flow_pred, xt, self.scheduler.sigmas, self.scheduler.timesteps], + ) + timestep_id = torch.argmin((timesteps - timestep).abs()) + sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) + return (xt - sigma_t * flow_pred).to(original_dtype) + + def _initialize_self_kv_cache( + self, + batch_size: int, + kv_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> list[dict]: + """Allocate per-block self-attention KV cache tensors.""" + head_dim = self.transformer.dim // self.transformer.num_heads + num_heads = self.transformer.num_heads // get_sp_world_size() + shape = [batch_size, kv_size, num_heads, head_dim] + return [ + { + "k": torch.zeros(shape, dtype=dtype, device=device), + "v": torch.zeros(shape, dtype=dtype, device=device), + "global_end_index": torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": torch.tensor([0], dtype=torch.long, device=device), + } + for _ in range(self.transformer.num_layers) + ] + + def _initialize_crossattn_cache( + self, + batch_size: int, + max_sequence_length: int, + dtype: torch.dtype, + device: torch.device, + ) -> list[dict]: + """Allocate per-block text cross-attention KV cache tensors.""" + head_dim = self.transformer.dim // self.transformer.num_heads + shape = [batch_size, max_sequence_length, self.transformer.num_heads, head_dim] + return [ + { + "k": torch.zeros(shape, dtype=dtype, device=device), + "v": torch.zeros(shape, dtype=dtype, device=device), + "is_init": torch.tensor([0], dtype=torch.bool, device=device), + } + for _ in range(self.transformer.num_layers) + ] + + @staticmethod + def _prompt_context(batch: ForwardBatch, device: torch.device) -> list[torch.Tensor]: + """Slice padded text encoder states back to LingBot World 2's unpadded context list.""" + assert batch.prompt_embeds + context_tensor = batch.prompt_embeds[0].to(device) + if batch.prompt_attention_mask: + mask = batch.prompt_attention_mask[0].to(device) + seq_lens = mask.gt(0).sum(dim=1).long() + return [u[:v] for u, v in zip(context_tensor, seq_lens, strict=True)] + return [u for u in context_tensor] + + def _prepare_image_tensor(self, batch: ForwardBatch, device: torch.device) -> torch.Tensor: + """Return the source-style normalized image tensor `[C,H,W]`.""" + image = batch.pil_image + if image is None: + raise ValueError("LingBot World 2 causal-fast requires `image_path` or `pil_image`.") + if isinstance(image, torch.Tensor): + if image.ndim == 5: + return image[0, :, 0].to(device) + if image.ndim == 4: + return image[0].to(device) + return image.to(device) + return TF.to_tensor(image).sub_(0.5).div_(0.5).to(device) + + def _prepare_camera( + self, + action_path: str, + c2ws: np.ndarray, + h: int, + w: int, + lat_f: int, + lat_h: int, + lat_w: int, + chunk_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + """Build the LingBot World 2 camera Plucker tensor for latent chunks.""" + 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=h, + width_resize=w, + height_final=h, + width_final=w, + ) + Ks = Ks[0] + len_c2ws = len(c2ws) + len_c2ws_ = int((len_c2ws - 1) // 4) + 1 + len_c2ws_ = int(len_c2ws_ - (len_c2ws_ % chunk_size)) + 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, len_c2ws_), + ) + c2ws_infer = compute_relative_poses(c2ws_infer, framewise=True) + Ks = Ks.repeat(len(c2ws_infer), 1) + c2ws_plucker_emb = get_plucker_embeddings(c2ws_infer.to(device), Ks.to(device), h, w) + c2ws_plucker_emb = rearrange( + c2ws_plucker_emb, + "f (h c1) (w c2) c -> (f h w) (c c1 c2)", + c1=int(h // lat_h), + c2=int(w // lat_w), + ) + c2ws_plucker_emb = c2ws_plucker_emb[None, ...] + return rearrange( + c2ws_plucker_emb, + "b (f h w) c -> b c f h w", + f=lat_f, + h=lat_h, + w=lat_w, + ).to(device=device, dtype=dtype) + + def _encode_condition_video( + self, + img: torch.Tensor, + h: int, + w: int, + frames: int, + mask: torch.Tensor, + fastvideo_args: FastVideoArgs, + ) -> torch.Tensor: + """Encode the first-frame conditioning video and prepend mask channels.""" + device = get_local_torch_device() + self.vae = self.vae.to(device) + video_condition = torch.concat( + [ + torch.nn.functional.interpolate( + img[None].cpu(), + size=(h, w), + mode="bicubic", + ).transpose(0, 1), + torch.zeros(3, frames - 1, h, w), + ], + dim=1, + ).to(device) + vae_dtype = torch.float32 + with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=False): + encoder_output = self.vae.encode(video_condition.unsqueeze(0).to(torch.float32)) + latent_condition = encoder_output.mean + if not bool(getattr(self.vae, "handles_latent_denorm", False)): + if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None: + latent_condition -= self.vae.shift_factor.to(latent_condition.device, latent_condition.dtype) + latent_condition = latent_condition * self.vae.scaling_factor.to(latent_condition.device, latent_condition.dtype) + if fastvideo_args.vae_cpu_offload: + self.vae.to("cpu") + return torch.concat([mask, latent_condition[0]], dim=0) + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + """Execute LingBot World 2 causal-fast generation and store final latents on the batch.""" + device = get_local_torch_device() + cfg = fastvideo_args.pipeline_config.dit_config.arch_config + chunk_size = int(cfg.chunk_size) + max_sequence_length = int(batch.max_sequence_length or cfg.text_len) + action_path = batch.action_path + if action_path is None: + raise ValueError("LingBot World 2 causal-fast requires `action_path`.") + + c2ws = np.load(os.path.join(action_path, "poses.npy")) + len_c2ws = ((len(c2ws) - 1) // 4) * 4 + 1 + frame_num = ((int(batch.num_frames) - 1) // 4) * 4 + 1 + frame_num = min(frame_num, len_c2ws) + c2ws = c2ws[:frame_num] + + img = self._prepare_image_tensor(batch, device) + h0, w0 = img.shape[1:] + aspect_ratio = h0 / w0 + lat_h = round(np.sqrt(cfg.max_area * aspect_ratio) // 8 // cfg.patch_size[1] * cfg.patch_size[1]) + lat_w = round(np.sqrt(cfg.max_area / aspect_ratio) // 8 // cfg.patch_size[2] * cfg.patch_size[2]) + h = lat_h * 8 + w = lat_w * 8 + lat_f = (frame_num - 1) // 4 + 1 + lat_f = int(lat_f - (lat_f % chunk_size)) + frames = (lat_f - 1) * 4 + 1 + batch.height = h + batch.width = w + batch.num_frames = frames + + seed = int(batch.seed if batch.seed is not None else 42) + seed_g = torch.Generator(device=device) + seed_g.manual_seed(seed) + noise = torch.randn(16, lat_f, lat_h, lat_w, dtype=torch.float32, generator=seed_g, device=device) + + mask = torch.ones(1, frames, lat_h, lat_w, device=device) + mask[:, 1:] = 0 + mask = torch.concat([torch.repeat_interleave(mask[:, 0:1], repeats=4, dim=1), mask[:, 1:]], dim=1) + mask = mask.view(1, mask.shape[1] // 4, 4, lat_h, lat_w).transpose(1, 2)[0] + + self.scheduler.set_timesteps(cfg.num_train_timesteps, shift=cfg.sample_shift) + timesteps = self.scheduler.timesteps[list(cfg.timesteps_index)].to(device) + context = self._prompt_context(batch, device) + c2ws_plucker_emb = self._prepare_camera( + action_path, + c2ws, + h, + w, + lat_f, + lat_h, + lat_w, + chunk_size, + torch.bfloat16, + device, + ) + y = self._encode_condition_video(img, h, w, frames, mask, fastvideo_args).to(device=device, dtype=torch.bfloat16) + + transformer_dtype = torch.bfloat16 + frame_seqlen = int(noise.shape[-2] * noise.shape[-1] // 4) + kv_size = frame_seqlen * cfg.local_attn_size if cfg.local_attn_size > -1 else frame_seqlen * lat_f + self_kv_cache = self._initialize_self_kv_cache(1, kv_size, transformer_dtype, device) + cross_kv_cache = self._initialize_crossattn_cache(1, max_sequence_length, transformer_dtype, device) + + self.transformer = self.transformer.to(device) + self._cross_attn_initialized = False + pred_latent_chunks = [] + latents_chunk = noise.split(chunk_size, dim=1) + condition_chunk = y.split(chunk_size, dim=1) + c2ws_plucker_emb_chunk = c2ws_plucker_emb.split(chunk_size, dim=2) + max_seq_len = int(math.ceil((chunk_size * lat_h * lat_w // 4))) + + with torch.amp.autocast("cuda", dtype=transformer_dtype): + for chunk_id, current_latent in enumerate(latents_chunk): + current_condition = condition_chunk[chunk_id] + current_c2ws_plucker_emb = c2ws_plucker_emb_chunk[chunk_id] + dit_cond_dict = {"c2ws_plucker_emb": current_c2ws_plucker_emb.chunk(1, dim=0)} + kwargs = { + "context": [context[0]], + "seq_len": max_seq_len, + "y": [current_condition], + "dit_cond_dict": dit_cond_dict, + "kv_cache": self_kv_cache, + "crossattn_cache": cross_kv_cache, + "current_start": chunk_id * chunk_size * frame_seqlen, + "max_attention_size": kv_size, + "frame_seqlen": frame_seqlen, + } + x0 = current_latent + for timestep_idx, timestep_value in enumerate(timesteps): + timestep = torch.stack([timestep_value]).to(device) + noise_pred = self.transformer( + x=[current_latent.to(device)], + t=timestep, + cross_attn_first_call=not self._cross_attn_initialized, + **kwargs, + )[0] + self._cross_attn_initialized = True + x0 = self._convert_flow_pred_to_x0(noise_pred, current_latent, timestep_value) + if timestep_idx < len(timesteps) - 1: + next_timestep = timesteps[timestep_idx + 1].reshape(1) + current_latent = self.scheduler.add_noise( + x0, + torch.randn(x0.shape, generator=seed_g, device=x0.device, dtype=x0.dtype), + next_timestep, + ) + pred_latent_chunks.append(x0) + context_timestep = torch.stack([timesteps[-1] * 0.0]).to(device) + self.transformer(x=[x0], t=context_timestep, cross_attn_first_call=False, **kwargs) + + batch.latents = torch.cat(pred_latent_chunks, dim=1).unsqueeze(0) + return batch + + +class LingBotWorld2CausalFastPipeline(LoRAPipeline, ComposedPipelineBase): + """FastVideo pipeline for LingBot World 2 14B causal-fast I2V generation.""" + + _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + if "scheduler" not in self.modules: + self.modules["scheduler"] = FlowUniPCMultistepScheduler(shift=1.0, use_dynamic_shifting=False) + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + """Set up the LingBot World 2 causal-fast pipeline stages.""" + self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage()) + self.add_stage( + stage_name="prompt_encoding_stage", + stage=LingBotWorld2TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + ) + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + self.add_stage( + stage_name="lingbotworld2_causal_fast_generation_stage", + stage=LingBotWorld2CausalFastGenerationStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + vae=self.get_module("vae"), + ), + ) + self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))) + + +EntryClass = LingBotWorld2CausalFastPipeline diff --git a/fastvideo/pipelines/basic/lingbotworld2/presets.py b/fastvideo/pipelines/basic/lingbotworld2/presets.py new file mode 100644 index 000000000..109cce200 --- /dev/null +++ b/fastvideo/pipelines/basic/lingbotworld2/presets.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LingBotWorld2 causal-fast pipeline preset.""" + +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +_DENOISE_STAGE = PresetStageSpec( + name="denoise", + kind="denoising", + description="Causal-fast denoising pass", + allowed_overrides=frozenset({ + "num_inference_steps", + "guidance_scale", + }), +) + +LINGBOTWORLD2_CAUSAL_FAST_I2V = InferencePreset( + name="lingbotworld2_causal_fast_i2v", + version=1, + model_family="lingbotworld2", + description="LingBot World 2 14B causal-fast I2V", + workload_type="i2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "guidance_scale": 1.0, + "num_inference_steps": 4, + "fps": 16, + "seed": 42, + "num_frames": 65, + "height": 480, + "width": 832, + "negative_prompt": "", + }, +) + +ALL_PRESETS = (LINGBOTWORLD2_CAUSAL_FAST_I2V, ) diff --git a/scripts/checkpoint_conversion/convert_lingbotworld2_causal_fast.py b/scripts/checkpoint_conversion/convert_lingbotworld2_causal_fast.py new file mode 100644 index 000000000..8a232cffe --- /dev/null +++ b/scripts/checkpoint_conversion/convert_lingbotworld2_causal_fast.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Create a FastVideo-native LingBot World 2 causal-fast bundle.""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + + +LATENTS_MEAN = [ + -0.7571, + -0.7089, + -0.9113, + 0.1075, + -0.1745, + 0.9653, + -0.1517, + 1.5508, + 0.4134, + -0.0715, + 0.5517, + -0.3632, + -0.1922, + -0.9497, + 0.2503, + -0.2921, +] +LATENTS_STD = [ + 2.8184, + 1.4541, + 2.3275, + 2.6558, + 1.2196, + 1.7708, + 2.6052, + 2.0743, + 3.2687, + 2.1526, + 2.8652, + 1.5579, + 1.6382, + 1.1253, + 2.8251, + 1.9160, +] + + +def _write_json(path: Path, payload: dict) -> None: + """Write deterministic JSON config files for a local bundle.""" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + + +def _symlink_or_replace(source: Path, target: Path) -> None: + """Create a relative symlink, replacing stale symlinks only.""" + target.parent.mkdir(parents=True, exist_ok=True) + if target.is_symlink() or target.exists(): + if target.is_symlink() and target.resolve() == source.resolve(): + return + raise FileExistsError(f"Refusing to replace existing path: {target}") + target.symlink_to(os.path.relpath(source, target.parent)) + + +def convert_vae(source_dir: Path, output_dir: Path) -> None: + """Link the exact LingBot World 2 Wan VAE checkpoint consumed by LingBotWorld2WanVAE.""" + _symlink_or_replace(source_dir / "Wan2.1_VAE.pth", output_dir / "vae" / "Wan2.1_VAE.pth") + + +def build_bundle(source_dir: Path, output_dir: Path) -> None: + """Create configs and symlinks for FastVideo-native LingBot World 2 loading.""" + output_dir.mkdir(parents=True, exist_ok=True) + _write_json( + output_dir / "model_index.json", + { + "_class_name": "LingBotWorld2CausalFastPipeline", + "_diffusers_version": "0.30.0", + "scheduler": ["diffusers", "FlowUniPCMultistepScheduler"], + "text_encoder": ["transformers", "LingBotWorld2T5EncoderModel"], + "tokenizer": ["transformers", "AutoTokenizer"], + "transformer": ["diffusers", "LingBotWorld2CausalFastTransformer3DModel"], + "vae": ["diffusers", "LingBotWorld2WanVAE"], + }, + ) + _write_json( + output_dir / "transformer" / "config.json", + { + "_class_name": "LingBotWorld2CausalFastTransformer3DModel", + "_diffusers_version": "0.30.0", + "model_type": "i2v", + "patch_size": [1, 2, 2], + "text_len": 512, + "in_dim": 36, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "local_attn_size": 18, + "sink_size": 6, + "qk_norm": True, + "cross_attn_norm": True, + "eps": 1e-6, + }, + ) + _write_json( + output_dir / "text_encoder" / "config.json", + { + "architectures": ["LingBotWorld2T5EncoderModel"], + "vocab_size": 256384, + "dim": 4096, + "dim_attn": 4096, + "dim_ffn": 10240, + "num_heads": 64, + "num_layers": 24, + "num_buckets": 32, + "text_len": 512, + "hidden_size": 4096, + "dropout": 0.1, + }, + ) + _write_json( + output_dir / "vae" / "config.json", + { + "_class_name": "LingBotWorld2WanVAE", + "base_dim": 96, + "z_dim": 16, + "dim_mult": [1, 2, 4, 4], + "num_res_blocks": 2, + "attn_scales": [], + "temperal_downsample": [False, True, True], + "dropout": 0.0, + "latents_mean": LATENTS_MEAN, + "latents_std": LATENTS_STD, + }, + ) + _write_json( + output_dir / "scheduler" / "scheduler_config.json", + { + "_class_name": "FlowUniPCMultistepScheduler", + "num_train_timesteps": 1000, + "shift": 1.0, + "use_dynamic_shifting": False, + "prediction_type": "flow_prediction", + "solver_order": 2, + "solver_type": "bh2", + "final_sigmas_type": "zero", + }, + ) + _write_json(output_dir / "scheduler" / "config.json", json.loads((output_dir / "scheduler" / "scheduler_config.json").read_text())) + + _symlink_or_replace(source_dir / "models_t5_umt5-xxl-enc-bf16.pth", output_dir / "text_encoder" / "pytorch_model.pt") + _symlink_or_replace(source_dir / "google" / "umt5-xxl", output_dir / "tokenizer") + for shard in sorted((source_dir / "transformers").glob("*.safetensors")): + _symlink_or_replace(shard, output_dir / "transformer" / shard.name) + index_path = source_dir / "transformers" / "diffusion_pytorch_model.safetensors.index.json" + if index_path.exists(): + _symlink_or_replace(index_path, output_dir / "transformer" / "diffusion_pytorch_model.safetensors.index.json") + convert_vae(source_dir, output_dir) + + +def main() -> None: + """Parse CLI arguments and build the local converted bundle.""" + parser = argparse.ArgumentParser() + parser.add_argument("--source-dir", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + args = parser.parse_args() + build_bundle(args.source_dir, args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/tests/local_tests/lingbotworld2/PORT_STATUS.md b/tests/local_tests/lingbotworld2/PORT_STATUS.md new file mode 100644 index 000000000..bf745a7e2 --- /dev/null +++ b/tests/local_tests/lingbotworld2/PORT_STATUS.md @@ -0,0 +1,82 @@ +# LingBot World 2 Port Status + +## Summary + +- model_family: `lingbotworld2` +- workload_types: I2V causal-fast +- official_ref: `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2` +- official_ref_dir: `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2` +- hf_weights_path: `robbyant/lingbot-world-v2-14b-causal-fast` +- local_weights_dir: `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast` +- source_layout: `raw_official` +- local_tests_readme: `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/tests/local_tests/lingbotworld2/README.md` + +## Current Phase + +- phase: port verified +- status: complete +- owner: orchestrator +- last_updated: 2026-07-10 + +## Component Matrix + +| Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues | +| -------------- | --------- | ---------- | --------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------- | ---------------- | ------------------------------- | ------------------------- | ----------- | +| Transformer | DiT | ported | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/modules/model_fast.py` | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py` | `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/fastvideo/models/dits/lingbotworld2/causal_fast.py` | import/meta pass | symlink raw shards | exact generated parity | none | +| VAE | VAE | reused | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/modules/vae2_1.py` | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py` | `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/fastvideo/models/vaes/lingbotworld2_wanvae.py` | reuse proven | symlink raw `.pth` | exact encode/decode pass | none | +| Text encoder | encoder | ported | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/modules/t5.py` | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py` | `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/fastvideo/models/encoders/lingbotworld2_t5.py` | import pass | symlink official `.pth` as `.pt` | strict-load pass | none | +| Camera utility | generic | ported | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/utils/cam_utils.py` | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py` | `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/fastvideo/models/dits/lingbotworld2/cam_utils.py` | import pass | stateless | covered by pipeline smoke | none | +| Scheduler | scheduler | reused | `FlowUniPCMultistepScheduler` in official sampling loop | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py` | `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py` | reuse proven | config emitted | exact generated parity | none | +| Pipeline | pipeline | ported | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/generate.py`; `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py` | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/run_fast.sh` | `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/fastvideo/pipelines/basic/lingbotworld2/causal_fast_pipeline.py` | import pass | model_index emitted | exact generated parity | none | + +## Conversion State + +- conversion_script: `/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot/scripts/checkpoint_conversion/convert_lingbotworld2_causal_fast.py` +- converted_weights_dir: `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast-fastvideo` +- source_layout: `raw_official` +- strict_load_status: VAE pass; text encoder pass; transformer full 14B load pass +- passthrough_components: transformer shards, text encoder checkpoint, tokenizer files, raw VAE checkpoint +- retry_history: generic `AutoencoderKLWan` VAE conversion strict-loaded but was not numerically exact; final bundle uses `LingBotWorld2WanVAE`. + +## Parity Commands + +| Scope | Command | Last Result | Notes | +| ------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------- | ----------------------------------------------- | +| Preflight and parity | `/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python -m pytest tests/local_tests/lingbotworld2/test_lingbotworld2_causal_fast_preflight.py tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_smoke.py tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_parity.py -q -rs` | pass | 2026-07-10: 5 passed, 1 skipped; skip is opt-in 14B heavy smoke. | +| Reference matrix | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/scripts/run_reference_matrix.sh` | pass | `5` and `9` frame cases recorded as official_unsupported. | +| FastVideo matrix | `/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python /mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/scripts/run_fastvideo_matrix.py` | pass | Matches official-supported cases and records `5`/`9` as official_unsupported. | +| Generated video parity | `/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python /mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/scripts/compare_matrix_outputs.py` | pass | 2026-07-09: all 9 supported rows exact; max abs 0, mean abs 0.0. | +| Heavy FastVideo smoke | `LINGBOTWORLD2_RUN_HEAVY_SMOKE=1 /mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python -m pytest tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_smoke.py -q -rs` | pass | 2026-07-10: 2 passed; 8-GPU 14B latent generation passed in 3:15. | + +## Open Questions + +| ID | Question | Owner | Needed By Phase | Status | Resolution | +| ---- | ------------------------------------------------------------------------------------------ | ------------ | --------------- | -------- | ---------- | +| Q001 | Should final reported matrix include source-impossible requested `5` and `9` frame cases? | user/author | final parity | resolved | Record them as official-source failures; do not change `chunk_size=4`. | + +## Issues And Blockers + +| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution | +| ---- | ------------------ | --------- | -------- | --------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------ | ------------ | ------ | ---------- | +| I001 | pipeline parity | pipeline | high | Requested `5` and `9` frames cannot produce official reference videos. | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py` sets `lat_f=0`, then `F=-3`. | orchestrator | closed | Matrix scripts record unsupported cases; supported cases remain `17/33/65`. | + +## Escape Hatches + +| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution | +| -- | ----- | ------------- | -------- | ------------------ | ------ | ---------- | + +## Decisions + +| Date | Decision | Rationale | Impact | +| ---------- | ------------------------------------------------------------------------ | ------------------------------------------------------------------------------------------ | ------------------------------------------------------------ | +| 2026-07-09 | Keep the FastVideo causal-fast port native and do not import source repo. | User required a FastVideo-native model/pipeline implementation. | Production code owns DiT, T5 wrapper, camera utilities, and pipeline loop. | +| 2026-07-09 | Preserve official `chunk_size=4` behavior. | `run_fast.sh` does not override chunk size and `generate.py` defaults to `4`. | Requested `5` and `9` frame cases are recorded as official failures. | +| 2026-07-09 | Store generated outputs under the LingBot World 2 source tree. | User required outputs to stay within the named directories. | Matrix outputs live under `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/outputs`. | +| 2026-07-09 | Use the exact LingBot World 2 Wan VAE implementation in FastVideo. | Generic `AutoencoderKLWan` strict-loaded but decoded output was not numerically exact. | `LingBotWorld2WanVAE` loads the original `Wan2.1_VAE.pth` and restores exact video parity. | +| 2026-07-10 | Use `lingbotworld2` for the causal-fast model family and `LingBotWorld2` for its classes. | FastVideo already has a separate LingBot World v1 family under `lingbotworld`. | V2 modules, configs, presets, registry entries, tests, converter, dataset, and example no longer share the v1 namespace. | + +## Handoff Notes + +- The final comparison report is `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/outputs/verification/matrix_compare.md`. +- Exact generated-output parity is only meaningful for cases where the official reference writes a video. +- Do not loosen parity by changing `chunk_size` unless the user explicitly changes the reference config. diff --git a/tests/local_tests/lingbotworld2/README.md b/tests/local_tests/lingbotworld2/README.md new file mode 100644 index 000000000..cffc3e963 --- /dev/null +++ b/tests/local_tests/lingbotworld2/README.md @@ -0,0 +1,163 @@ +# LingBot World 2 Local Tests + +Local-only smoke, conversion, and generated-output parity notes for the +`lingbotworld2` causal-fast FastVideo port. + +## Reference Assets + +| Field | Value | +| ------------------- | ----------------------------------------------------------------------------------------------------------------- | +| Model family | `lingbotworld2` | +| First scope | `robbyant/lingbot-world-v2-14b-causal-fast`, I2V causal-fast inference | +| Workload types | I2V: text + first image + camera/action path to video | +| Official reference | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2` | +| HF weights | `robbyant/lingbot-world-v2-14b-causal-fast` | +| Raw weights | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast` | +| FastVideo bundle | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast-fastvideo` | +| Source layout | `raw_official` | +| Needs conversion | `yes` | +| Token env var | none recorded | + +Do not write token values in this file. + +## Shared Environment Setup + +Use the existing venv and do not rebuild FlashAttention: + +```bash +/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/pip install -r /mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/requirements.txt +/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/pip install --no-deps -e /mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot +``` + +The current environment already has the FastVideo editable install pointing at: + +```text +/mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot +``` + +Runtime cache and temporary paths used by the matrix scripts: + +```text +/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/.cache +``` + +## Conversion + +The converter writes a FastVideo-loadable bundle with native component configs, +the exact LingBot World 2 Wan VAE checkpoint, and symlinked raw LingBot World 2 transformer/T5 +assets: + +```bash +cd /mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot +/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python scripts/checkpoint_conversion/convert_lingbotworld2_causal_fast.py \ + --source-dir /mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast \ + --output-dir /mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast-fastvideo +``` + +Strict-load checks already run locally: + +| Component | Status | Notes | +| ------------ | ------ | --------------------------------------------------------------------- | +| VAE | pass | `LingBotWorld2WanVAE` loads the original `Wan2.1_VAE.pth` checkpoint. | +| Text encoder | pass | Native `LingBotWorld2T5EncoderModel` strict-load from symlinked `.pt`. | +| Transformer | pass | Full 14B load and 8-GPU causal-fast generation pass. | + +## Official Frame-Count Constraint + +The requested verification frame counts are `5, 9, 17, 33, 65`. The official +causal-fast source uses default `chunk_size=4`, then truncates latent frames to a +multiple of the chunk size. This makes `5` and `9` source-impossible and changes +the generated frame counts for the remaining values. + +| Requested frames | Official effective frames | Status | +| ---------------- | ------------------------- | ------------------------------ | +| `5` | `-3` | official source fails | +| `9` | `-3` | official source fails | +| `17` | `13` | runs with truncated output | +| `33` | `29` | runs with truncated output | +| `65` | `61` | runs with truncated output | + +Evidence: `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/wan/image2video.py` +rounds `frame_num` to `4n+1`, computes latent frames, truncates `lat_f` by +`chunk_size`, and resets `F = (lat_f - 1) * 4 + 1`. + +## Local Tests + +```bash +cd /mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot +/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python -m pytest \ + tests/local_tests/lingbotworld2/test_lingbotworld2_causal_fast_preflight.py \ + tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_smoke.py \ + tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_parity.py \ + -q -rs +``` + +Current result: + +```text +2026-07-10: 5 passed, 1 skipped. +Skipped test: opt-in 14B heavy smoke gated by LINGBOTWORLD2_RUN_HEAVY_SMOKE=1. +``` + +Heavy smoke is opt-in because it loads the 14B model: + +```bash +cd /mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot +LINGBOTWORLD2_RUN_HEAVY_SMOKE=1 /mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python -m pytest \ + tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_smoke.py \ + -q -rs +``` + +Current heavy result: + +```text +2026-07-10: 2 passed in 3:15 on 8 H200 GPUs. +``` + +## Matrix Verification + +Reference generation: + +```bash +/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/scripts/run_reference_matrix.sh +``` + +FastVideo generation: + +```bash +/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python \ + /mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/scripts/run_fastvideo_matrix.py +``` + +Decoded-frame comparison: + +```bash +/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python \ + /mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/scripts/compare_matrix_outputs.py +``` + +Current comparison result: + +```text +2026-07-09: all 9 official-supported FastVideo videos are exactly equal to the +reference videos after decoded-frame comparison: max abs 0, mean abs 0.0. +The requested 5-frame and 9-frame cases are recorded as official_unsupported for +both reference and FastVideo. +``` + +Expected output roots: + +| Output type | Path | +| ------------ | -------------------------------------------------------------------------- | +| Reference | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/outputs/reference` | +| FastVideo | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/outputs/fastvideo` | +| Verification | `/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/outputs/verification` | + +Generated-output parity test: + +```bash +cd /mnt/weka/shrd/wm/junda/fv-hub/fastvideo-port-lingbot +/mnt/weka/shrd/wm/junda/fv-hub/.venv/bin/python -m pytest \ + tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_parity.py \ + -q -rs +``` diff --git a/tests/local_tests/lingbotworld2/__init__.py b/tests/local_tests/lingbotworld2/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/tests/local_tests/lingbotworld2/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/local_tests/lingbotworld2/test_lingbotworld2_causal_fast_preflight.py b/tests/local_tests/lingbotworld2/test_lingbotworld2_causal_fast_preflight.py new file mode 100644 index 000000000..9002b9937 --- /dev/null +++ b/tests/local_tests/lingbotworld2/test_lingbotworld2_causal_fast_preflight.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Preflight checks for the LingBot World 2 causal-fast FastVideo port.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np + + +ROOT = Path("/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2") +FASTVIDEO_ROOT = Path(__file__).resolve().parents[3] +RAW_MODEL_DIR = ROOT / "ckpts" / "lingbot-world-v2-14b-causal-fast" +FASTVIDEO_MODEL_DIR = ROOT / "ckpts" / "lingbot-world-v2-14b-causal-fast-fastvideo" +ACTION_PATH = FASTVIDEO_ROOT / "examples" / "dataset" / "lingbotworld2" + + +def _source_effective_frames(frame_num: int, pose_count: int, chunk_size: int = 4) -> int: + """Mirror the released `wan/image2video.py` frame-count truncation.""" + rounded = ((frame_num - 1) // 4) * 4 + 1 + rounded = min(rounded, ((pose_count - 1) // 4) * 4 + 1) + latent_frames = (rounded - 1) // 4 + 1 + latent_frames = int(latent_frames - (latent_frames % chunk_size)) + return (latent_frames - 1) * 4 + 1 + + +def test_lingbotworld2_bundle_declares_native_components() -> None: + assert RAW_MODEL_DIR.exists() + assert FASTVIDEO_MODEL_DIR.exists() + assert (FASTVIDEO_MODEL_DIR / "model_index.json").exists() + assert (FASTVIDEO_MODEL_DIR / "transformer" / "config.json").exists() + assert (FASTVIDEO_MODEL_DIR / "text_encoder" / "config.json").exists() + assert (FASTVIDEO_MODEL_DIR / "text_encoder" / "pytorch_model.pt").exists() + assert (FASTVIDEO_MODEL_DIR / "vae" / "config.json").exists() + assert (FASTVIDEO_MODEL_DIR / "vae" / "Wan2.1_VAE.pth").exists() + assert (FASTVIDEO_MODEL_DIR / "scheduler" / "scheduler_config.json").exists() + + model_index = json.loads((FASTVIDEO_MODEL_DIR / "model_index.json").read_text(encoding="utf-8")) + assert model_index["_class_name"] == "LingBotWorld2CausalFastPipeline" + assert model_index["transformer"][1] == "LingBotWorld2CausalFastTransformer3DModel" + assert model_index["text_encoder"][1] == "LingBotWorld2T5EncoderModel" + assert model_index["vae"][1] == "LingBotWorld2WanVAE" + + +def test_lingbotworld2_source_frame_requests_truncate_as_expected() -> None: + pose_count = int(np.load(ACTION_PATH / "poses.npy").shape[0]) + effective = { + frame_num: _source_effective_frames(frame_num, pose_count) + for frame_num in (5, 9, 17, 33, 65) + } + assert effective == { + 5: -3, + 9: -3, + 17: 13, + 33: 29, + 65: 61, + } + + +def test_lingbotworld2_request_preserves_action_path() -> None: + from fastvideo.api.compat import normalize_generation_request, request_to_sampling_param + + request = normalize_generation_request({ + "prompt": "camera path schema check", + "inputs": { + "image_path": str(ACTION_PATH / "image.jpg"), + "action_path": str(ACTION_PATH), + }, + "sampling": { + "height": 480, + "width": 832, + "num_frames": 17, + "num_inference_steps": 4, + "guidance_scale": 1.0, + }, + "output": { + "save_video": False, + "return_frames": False, + }, + }) + sampling_param = request_to_sampling_param(request, model_path=str(FASTVIDEO_MODEL_DIR)) + assert sampling_param.action_path == str(ACTION_PATH) + assert sampling_param.image_path == str(ACTION_PATH / "image.jpg") + assert sampling_param.num_frames == 17 diff --git a/tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_parity.py b/tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_parity.py new file mode 100644 index 000000000..827d4ecbc --- /dev/null +++ b/tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_parity.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Generated-output parity for LingBot World 2 causal-fast matrix runs.""" + +from __future__ import annotations + +import json +from pathlib import Path +import subprocess +import sys + +import pytest + + +ROOT = Path("/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2") +REFERENCE_STATUS = ROOT / "outputs" / "reference" / "status.tsv" +FASTVIDEO_STATUS = ROOT / "outputs" / "fastvideo" / "status.tsv" +COMPARE_SCRIPT = ROOT / "scripts" / "compare_matrix_outputs.py" +COMPARE_JSON = ROOT / "outputs" / "verification" / "matrix_compare.json" + + +def test_lingbotworld2_generated_matrix_matches_reference_exactly() -> None: + """Compare generated matrix videos after the source-side launch scripts run.""" + if not REFERENCE_STATUS.exists() or not FASTVIDEO_STATUS.exists(): + pytest.skip("Run lingbot-world-v2/scripts/run_reference_matrix.sh and run_fastvideo_matrix.py first.") + + subprocess.run([sys.executable, str(COMPARE_SCRIPT)], check=True) + results = json.loads(COMPARE_JSON.read_text(encoding="utf-8")) + compared = [item for item in results if item.get("comparison") and "exact" in item["comparison"]] + if not compared: + pytest.fail("No LingBot World 2 matrix cases had both reference and FastVideo videos to compare.") + + failures = [ + item + for item in compared + if not item["comparison"].get("exact") + ] + assert failures == [] diff --git a/tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_smoke.py b/tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_smoke.py new file mode 100644 index 000000000..b891fd84d --- /dev/null +++ b/tests/local_tests/pipelines/test_lingbotworld2_causal_fast_pipeline_smoke.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Smoke tests for the LingBot World 2 causal-fast pipeline.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, cast + +import pytest +import torch + + +MODEL_DIR = Path("/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast-fastvideo") +FASTVIDEO_ROOT = Path(__file__).resolve().parents[3] +ACTION_PATH = FASTVIDEO_ROOT / "examples" / "dataset" / "lingbotworld2" + + +def test_lingbotworld2_checkpoint_selects_expected_pipeline_and_defaults() -> None: + import fastvideo.registry as registry + from fastvideo.api.presets import get_preset, get_presets_for_family + from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig + from fastvideo.fastvideo_args import WorkloadType + from fastvideo.pipelines.basic.lingbotworld2.causal_fast_pipeline import ( + EntryClass, + LingBotWorld2CausalFastPipeline, + ) + + assert EntryClass is LingBotWorld2CausalFastPipeline + assert LingBotWorld2CausalFastPipeline._required_config_modules == [ + "text_encoder", + "tokenizer", + "vae", + "transformer", + "scheduler", + ] + + default_preset, model_family = registry.get_preset_selection(str(MODEL_DIR)) + assert model_family == "lingbotworld2" + assert default_preset == "lingbotworld2_causal_fast_i2v" + + info = registry.get_model_info( + str(MODEL_DIR), + workload_type=WorkloadType.I2V, + override_pipeline_cls_name="LingBotWorld2CausalFastPipeline", + ) + assert info.pipeline_cls is LingBotWorld2CausalFastPipeline + assert info.pipeline_config_cls is LingBotWorld2CausalFastI2V480PConfig + + names = {preset.name for preset in get_presets_for_family("lingbotworld2")} + assert "lingbotworld2_causal_fast_i2v" in names + preset = get_preset("lingbotworld2_causal_fast_i2v", "lingbotworld2") + assert preset.defaults["num_inference_steps"] == 4 + assert preset.defaults["height"] == 480 + assert preset.defaults["width"] == 832 + assert preset.defaults["num_frames"] == 65 + assert preset.defaults["guidance_scale"] == 1.0 + + cfg = LingBotWorld2CausalFastI2V480PConfig() + arch = cfg.dit_config.arch_config + assert arch.model_type == "i2v" + assert arch.in_dim == 36 + assert arch.local_attn_size == 18 + assert arch.sink_size == 6 + assert arch.chunk_size == 4 + assert tuple(arch.timesteps_index) == (0, 250, 500, 750) + + +@pytest.mark.skipif( + os.getenv("LINGBOTWORLD2_RUN_HEAVY_SMOKE") != "1", + reason="Set LINGBOTWORLD2_RUN_HEAVY_SMOKE=1 to load the 14B causal-fast model.", +) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="LingBot World 2 heavy smoke requires CUDA") +def test_lingbotworld2_14b_generates_finite_latents_on_8_gpus() -> None: + from fastvideo import VideoGenerator + + generator = VideoGenerator.from_pretrained( + str(MODEL_DIR), + num_gpus=8, + sp_size=8, + hsdp_shard_dim=8, + use_fsdp_inference=True, + dit_layerwise_offload=False, + dit_cpu_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=False, + pin_cpu_memory=True, + output_type="latent", + override_pipeline_cls_name="LingBotWorld2CausalFastPipeline", + ) + try: + result = generator.generate_video( + "A serene lakeside scene with a lone tree standing in calm water.", + image_path=str(ACTION_PATH / "image.jpg"), + action_path=str(ACTION_PATH), + output_path="/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/outputs/fastvideo/heavy_smoke", + save_video=False, + return_frames=True, + height=480, + width=832, + num_frames=17, + num_inference_steps=4, + guidance_scale=1.0, + negative_prompt="", + fps=16, + seed=42, + ) + finally: + generator.shutdown() + + samples = cast(dict[str, Any], result)["samples"] + assert torch.is_tensor(samples) + assert samples.ndim == 5 + assert samples.shape[1] == 16 + assert torch.isfinite(samples).all()