Connect Lingbot world 2 registration and porting.

.gitignore

@@ -6,6 +6,7 @@ results/
 wandb/
 *.ipynb
 *.jpg
+!examples/dataset/lingbotworld2/image.jpg
 *.safetensors
 *.mp4
 *.png

examples/dataset/lingbotworld2/README.md

@@ -0,0 +1,17 @@
+# LingBot World 2 Example Dataset
+
+These files were copied unchanged from the LingBot World 2 repository for the
+FastVideo causal-fast inference example.
+
+- Repository: `https://github.com/Robbyant/lingbot-world-v2.git`
+- Source commit: `94f43115de8d4a4f9f282126528c300a0b232c5f`
+- Source directory: `examples/03`
+
+## Files
+
+- `image.jpg`: source image for image-to-video generation. SHA-256:
+  `6ee3dacfef32cfef504dd698adb8a660cf15f686535c52fed4903fef27c0edd0`
+- `poses.npy`: camera-to-world trajectory matrices. SHA-256:
+  `bd0a23a696e184b0b43e7767eb432bfe644690560fe327fa96961affc941c404`
+- `intrinsics.npy`: camera intrinsic parameters. SHA-256:
+  `821fca6cf957ae8fbb1181307f02479efb1705e04c9e05734cd02fb43462e082`

examples/inference/basic/basic_lingbotworld2_causal_fast.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run LingBot World 2 14B causal-fast I2V generation with FastVideo."""
+
+from pathlib import Path
+
+from fastvideo import VideoGenerator
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+DATASET_DIR = REPO_ROOT / "examples" / "dataset" / "lingbotworld2"
+OUTPUT_PATH = REPO_ROOT / "outputs" / "lingbotworld2_causal_fast.mp4"
+
+
+def main() -> None:
+    """Load the native FastVideo LingBot World 2 causal-fast pipeline and generate one video."""
+    generator = VideoGenerator.from_pretrained(
+        "/mnt/weka/shrd/wm/junda/fv-hub/lingbot-world-v2/ckpts/lingbot-world-v2-14b-causal-fast-fastvideo",
+        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,
+        override_pipeline_cls_name="LingBotWorld2CausalFastPipeline",
+    )
+
+    try:
+        generator.generate_video(
+            "A serene lakeside scene with a lone tree standing in calm water, surrounded by distant snow-capped mountains under a bright blue sky with drifting white clouds; gentle ripples reflect the tree and sky, creating a tranquil, meditative atmosphere.",
+            action_path=str(DATASET_DIR),
+            output_path=str(OUTPUT_PATH),
+            save_video=True,
+            height=480,
+            width=832,
+            num_frames=65,
+            num_inference_steps=4,
+            guidance_scale=1.0,
+            negative_prompt="",
+            fps=16,
+            seed=42,
+        )
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

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)

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
 

fastvideo/configs/models/dits/__init__.py

@@ -14,10 +14,11 @@ from fastvideo.configs.models.dits.stable_audio import StableAudioConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoConfig
 from fastvideo.configs.models.dits.hyworld import HYWorldConfig
 from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
+from fastvideo.configs.models.dits.lingbotworld2 import LingBotWorld2CausalFastVideoConfig
 
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
     "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config",
     "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig",
-    "StableAudioConfig", "GlmImageDiTConfig"
+    "StableAudioConfig", "GlmImageDiTConfig", "LingBotWorld2CausalFastVideoConfig"
 ]

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"

fastvideo/configs/models/encoders/__init__.py

@@ -2,6 +2,7 @@ from fastvideo.configs.models.encoders.base import (BaseEncoderOutput, EncoderCo
                                                     TextEncoderConfig)
 from fastvideo.configs.models.encoders.clip import (CLIPTextConfig, CLIPVisionConfig, WAN2_1ControlCLIPVisionConfig)
 from fastvideo.configs.models.encoders.llama import LlamaConfig
+from fastvideo.configs.models.encoders.lingbotworld2_t5 import LingBotWorld2UMT5ArchConfig, LingBotWorld2UMT5Config
 from fastvideo.configs.models.encoders.t5 import T5Config, T5LargeConfig
 from fastvideo.configs.models.encoders.qwen2_5 import Qwen2_5_VLConfig
 from fastvideo.configs.models.encoders.siglip import SiglipVisionConfig
@@ -17,5 +18,6 @@ __all__ = [
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig",
     "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig",
     "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig", "StableAudioConditionerArchConfig",
-    "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig", "Mistral3TextConfig"
+    "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig", "Mistral3TextConfig",
+    "LingBotWorld2UMT5ArchConfig", "LingBotWorld2UMT5Config"
 ]

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"

fastvideo/configs/pipelines/__init__.py

@@ -7,6 +7,7 @@ from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyua
 from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
 from fastvideo.configs.pipelines.hyworld import HYWorldConfig
 from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig
+from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
@@ -19,5 +20,6 @@ __all__ = [
     "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
     "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
     "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig",
-    "Kandinsky5I2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
+    "Kandinsky5I2VConfig", "LingBotWorld2CausalFastI2V480PConfig", "MatrixGame2I2V480PConfig",
+    "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
 ]

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

fastvideo/models/dits/lingbotworld2/__init__.py

@@ -0,0 +1,5 @@
+from .causal_fast import LingBotWorld2CausalFastTransformer3DModel
+
+__all__ = ["LingBotWorld2CausalFastTransformer3DModel"]
+
+EntryClass = LingBotWorld2CausalFastTransformer3DModel

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

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

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

fastvideo/models/loader/component_loader.py

@@ -820,6 +820,24 @@ class VAELoader(ComponentLoader):
                 vae.load_state_dict(sd, strict=False)
                 return vae.eval()
 
+            if class_name == "LingBotWorld2WanVAE":
+                dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
+                config.pop("_class_name", None)
+                vae_config = fastvideo_args.pipeline_config.vae_config
+                vae_config.update_model_arch(config)
+                vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+                weight_path = os.path.join(model_path, "Wan2.1_VAE.pth")
+                if not os.path.exists(weight_path):
+                    raise FileNotFoundError(
+                        f"Missing LingBot World 2 VAE weights: {weight_path}"
+                    )
+                vae = vae_cls(
+                    vae_config,
+                    checkpoint_path=weight_path,
+                    dtype=dtype,
+                ).to(target_device)
+                return vae.eval()
+
             # LTX-2 uses CausalVideoAutoencoder with nested "vae" config
             if class_name == "CausalVideoAutoencoder" and "vae" in config:
                 vae_cls, _ = ModelRegistry.resolve_model_cls(class_name)

fastvideo/models/registry.py

@@ -42,6 +42,11 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"),
     "SD3Transformer2DModel": ("dits", "sd3", "SD3Transformer2DModel"),
     "LingBotWorldTransformer3DModel": ("dits", "lingbotworld", "LingBotWorldTransformer3DModel"),
+    "LingBotWorld2CausalFastTransformer3DModel": (
+        "dits",
+        "lingbotworld2",
+        "LingBotWorld2CausalFastTransformer3DModel",
+    ),
     "Gen3CTransformer3DModel": ("dits", "gen3c", "Gen3CTransformer3DModel"),
     "Kandinsky5Transformer3DModel": ("dits", "kandinsky5", "Kandinsky5Transformer3DModel"),
     "Flux2Transformer2DModel": ("dits", "flux_2", "Flux2Transformer2DModel"),
@@ -53,6 +58,11 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
     "DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
     "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"),
     "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
+    "LingBotWorld2CausalFastTransformer3DModel": (
+        "dits",
+        "lingbotworld2",
+        "LingBotWorld2CausalFastTransformer3DModel",
+    ),
     "MatrixGame2WanModel": ("dits", "matrixgame2", "MatrixGame2WanModel"),
     "CausalMatrixGame2WanModel": ("dits", "matrixgame2", "CausalMatrixGame2WanModel"),
     # Legacy aliases for older HF model_index.json files
@@ -72,6 +82,7 @@ _TEXT_ENCODER_MODELS = {
     ("encoders", "clip", "CLIPTextModelWithProjection"),
     "LlamaModel": ("encoders", "llama", "LlamaModel"),
     "UMT5EncoderModel": ("encoders", "t5", "UMT5EncoderModel"),
+    "LingBotWorld2T5EncoderModel": ("encoders", "lingbotworld2_t5", "LingBotWorld2T5EncoderModel"),
     "T5EncoderModel": ("encoders", "t5_hf", "T5EncoderModel"),
     "BertModel": ("encoders", "clip", "CLIPTextModel"),
     "Qwen2_5_VLTextModel": ("encoders", "qwen2_5", "Qwen2_5_VLTextModel"),
@@ -98,6 +109,7 @@ _VAE_MODELS = {
     "AutoencoderKLHYWorld": ("vaes", "hyworldvae", "AutoencoderKLHYWorld"),
     "AutoencoderKLHunyuanVideo15": ("vaes", "hunyuan15vae", "AutoencoderKLHunyuanVideo15"),
     "AutoencoderKLWan": ("vaes", "wanvae", "AutoencoderKLWan"),
+    "LingBotWorld2WanVAE": ("vaes", "lingbotworld2_wanvae", "LingBotWorld2WanVAE"),
     "AutoencoderKL": ("vaes", "autoencoder_kl", "AutoencoderKL"),
     "AutoencoderKLGen3CTokenizer":
     ("vaes", "gen3c_tokenizer_vae", "AutoencoderKLGen3CTokenizer"),

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

fastvideo/pipelines/basic/lingbotworld2/__init__.py

@@ -0,0 +1,5 @@
+from .causal_fast_pipeline import LingBotWorld2CausalFastPipeline
+
+__all__ = ["LingBotWorld2CausalFastPipeline"]
+
+EntryClass = LingBotWorld2CausalFastPipeline

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

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, )

fastvideo/pipelines/pipeline_batch_info.py

@@ -147,8 +147,9 @@ class ForwardBatch:
     camera_trajectory: str | None = None  # Camera trajectory file/identifier
     action_list: list[str] | None = None  # List of actions (e.g., ['forward', 'left'])
     action_speed_list: list[float] | None = None  # Speed for each action
-    # Camera control inputs (LingBotWorld)
+    # Camera control inputs (LingBotWorld and LingBotWorld2)
     c2ws_plucker_emb: torch.Tensor | None = None  # Plucker embedding: [B, C, F_lat, H_lat, W_lat]
+    action_path: str | None = None  # Directory containing poses.npy and intrinsics.npy
 
     # Camera control inputs (GEN3C)
     trajectory_type: str | None = None

fastvideo/registry.py

@@ -30,6 +30,7 @@ from fastvideo.configs.pipelines.hunyuan15 import (Hunyuan15T2V480PConfig, Hunyu
 from fastvideo.configs.pipelines.hyworld import HYWorldConfig
 from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig
 from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig
+from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig
 from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
 from fastvideo.configs.pipelines.flux_2 import (
@@ -488,6 +489,20 @@ def _register_configs() -> None:
         model_family="gamecraft",
         default_preset="gamecraft_i2v",
     )
+    # LingBotWorld2 causal-fast
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=LingBotWorld2CausalFastI2V480PConfig,
+        workload_types=(WorkloadType.I2V, ),
+        hf_model_paths=[
+            "robbyant/lingbot-world-v2-14b-causal-fast",
+        ],
+        model_detectors=[lambda path: ("lingbot-world-v2-14b-causal-fast" in path.lower()
+                                       or "lingbotworld2causalfastpipeline" in path.lower())],
+        model_family="lingbotworld2",
+        default_preset="lingbotworld2_causal_fast_i2v",
+    )
+
     # LingBotWorld
     register_configs(
         sampling_param_cls=None,
@@ -496,7 +511,9 @@ def _register_configs() -> None:
         hf_model_paths=[
             "FastVideo/LingBot-World-Base-Cam-Diffusers",
         ],
-        model_detectors=[lambda path: ("lingbotworld" in path.lower() or "lingbot-world" in path.lower())],
+        model_detectors=[lambda path: (("lingbotworld" in path.lower() or "lingbot-world" in path.lower())
+                                       and "causal-fast" not in path.lower()
+                                       and "causalfast" not in path.lower())],
         model_family="lingbotworld",
         default_preset="lingbotworld_i2v",
     )
@@ -1210,6 +1227,8 @@ def _register_presets() -> None:
         ALL_PRESETS as KANDINSKY5_PRESETS, )
     from fastvideo.pipelines.basic.lingbotworld.presets import (
         ALL_PRESETS as LINGBOTWORLD_PRESETS, )
+    from fastvideo.pipelines.basic.lingbotworld2.presets import (
+        ALL_PRESETS as LINGBOTWORLD2_PRESETS, )
     from fastvideo.pipelines.basic.longcat.presets import (
         ALL_PRESETS as LONGCAT_PRESETS, )
     from fastvideo.pipelines.basic.ltx2.presets import (
@@ -1240,6 +1259,7 @@ def _register_presets() -> None:
         HYWORLD_PRESETS,
         KANDINSKY5_PRESETS,
         LINGBOTWORLD_PRESETS,
+        LINGBOTWORLD2_PRESETS,
         LONGCAT_PRESETS,
         LTX2_PRESETS,
         MATRIXGAME2_PRESETS,

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()

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.

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
+```

tests/local_tests/lingbotworld2/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

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

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 == []

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()