lingbotworld2 ¶

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

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