lingbotworld2_wanvae ¶

Classes¶

fastvideo.models.vaes.lingbotworld2_wanvae.AttentionBlock ¶

AttentionBlock(dim)

Bases: Module

Causal self-attention with a single head.

Source code in fastvideo/models/vaes/lingbotworld2_wanvae.py

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)

fastvideo.models.vaes.lingbotworld2_wanvae.CausalConv3d ¶

CausalConv3d(*args, **kwargs)

Bases: Conv3d

Causal 3d convolusion.

Source code in fastvideo/models/vaes/lingbotworld2_wanvae.py

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)

fastvideo.models.vaes.lingbotworld2_wanvae.LingBotWorld2WanVAE ¶

LingBotWorld2WanVAE(config, checkpoint_path=None, dtype=float)

Bases: Module

FastVideo-facing wrapper around the exact LingBot World 2 Wan2.1 VAE computation.

Load the official LingBot World 2 VAE weights and expose FastVideo VAE APIs.

Source code in fastvideo/models/vaes/lingbotworld2_wanvae.py

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)

Methods:¶

fastvideo.models.vaes.lingbotworld2_wanvae.LingBotWorld2WanVAE.decode ¶

decode(latents: Tensor) -> Tensor

Decode normalized LingBot World 2 latents to clamped [-1,1] video tensors.

Source code in fastvideo/models/vaes/lingbotworld2_wanvae.py

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)

fastvideo.models.vaes.lingbotworld2_wanvae.LingBotWorld2WanVAE.encode ¶

encode(videos: Tensor)

Encode [B,C,T,H,W] videos and return a FastVideo-style mean tensor.

Source code in fastvideo/models/vaes/lingbotworld2_wanvae.py

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)

fastvideo.models.vaes.lingbotworld2_wanvae.Upsample ¶

Bases: Upsample

Methods:¶

fastvideo.models.vaes.lingbotworld2_wanvae.Upsample.forward ¶

forward(x)

Fix bfloat16 support for nearest neighbor interpolation.

Source code in fastvideo/models/vaes/lingbotworld2_wanvae.py

def forward(self, x):
    """
    Fix bfloat16 support for nearest neighbor interpolation.
    """
    return super().forward(x.float()).type_as(x)

fastvideo.models.vaes.lingbotworld2_wanvae.Wan2_1_VAE ¶

Wan2_1_VAE(z_dim=16, vae_pth='cache/vae_step_411000.pth', dtype=float, device='cuda')

Source code in fastvideo/models/vaes/lingbotworld2_wanvae.py

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)

Methods:¶

fastvideo.models.vaes.lingbotworld2_wanvae.Wan2_1_VAE.encode ¶

encode(videos)

videos: A list of videos each with shape [C, T, H, W].

Source code in fastvideo/models/vaes/lingbotworld2_wanvae.py

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
        ]