lingbot_video ¶

def forward(
    self,
    hidden_states: torch.Tensor,
    rotary_emb: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    original_seq_len: int | None = None,
) -> torch.Tensor:
    """Project QKV, apply rotary embeddings, and run non-causal SDPA."""
    batch, sequence, _ = hidden_states.shape
    query = self.to_q(hidden_states).unflatten(2, (self.num_heads, self.head_dim))
    key = self.to_k(hidden_states).unflatten(2, (self.num_heads, self.head_dim))
    value = self.to_v(hidden_states).unflatten(2, (self.num_heads, self.head_dim))
    query = _apply_rotary_emb(self.norm_q(query), rotary_emb)
    key = _apply_rotary_emb(self.norm_k(key), rotary_emb)

    if original_seq_len is not None:
        # Attention needs the full unpadded joint sequence while projections
        # and residual blocks stay sharded over tokens.
        qkv = torch.cat([query, key, value], dim=0)
        qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1)
        padded_seq_len = qkv.shape[1]
        query, key, value = qkv[:, :original_seq_len].chunk(3, dim=0)
    output = F.scaled_dot_product_attention(
        query.transpose(1, 2),
        key.transpose(1, 2),
        value.transpose(1, 2),
        attn_mask=attention_mask,
        dropout_p=0.0,
        is_causal=False,
    ).transpose(1, 2)
    if original_seq_len is not None:
        output = F.pad(output, (0, 0, 0, 0, 0, padded_seq_len - original_seq_len))
        output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2)
    return self.to_out(output.reshape(batch, sequence, -1).type_as(hidden_states))

def forward(
    self,
    hidden_states: torch.Tensor,
    temb6: torch.Tensor,
    rotary_emb: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    moe_padding_mask: torch.Tensor | None = None,
    original_seq_len: int | None = None,
) -> torch.Tensor:
    """Run attention and configured feed-forward residual branches with fp32 AdaLN."""
    expected_tokens = hidden_states.shape[0] * hidden_states.shape[1]
    if temb6.ndim != 2 or temb6.shape[0] != expected_tokens:
        raise ValueError("LingBotVideoBlock expects token-level temb6 with shape "
                         f"(B*S, 6D); got {tuple(temb6.shape)} for {tuple(hidden_states.shape)}.")
    modulation = temb6.view(*hidden_states.shape[:2], -1) + self.scale_shift_table.unsqueeze(0)
    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = modulation.chunk(6, dim=-1)
    gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
    scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
    bulk_dtype = self.attn.to_q.weight.dtype

    attention_input = (self.norm1(hidden_states) * scale_msa + shift_msa).to(bulk_dtype)
    attention_output = self.attn(attention_input, rotary_emb, attention_mask, original_seq_len)
    hidden_states = hidden_states + (gate_msa * self.norm_post_attn(attention_output)).to(hidden_states.dtype)
    mlp_input = (self.norm2(hidden_states) * scale_mlp + shift_mlp).to(bulk_dtype)
    if isinstance(self.ffn, LingBotVideoSparseMoeBlock):
        mlp_output = self.ffn(mlp_input, padding_mask=moe_padding_mask)
    else:
        mlp_output = self.ffn(mlp_input)
    mlp_output = self.norm_post_ffn(mlp_output)
    return hidden_states + (gate_mlp * mlp_output).to(hidden_states.dtype)

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """Return the projected tensor while preserving the normal weight surface."""
    output, _ = super().forward(hidden_states)
    return output

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """Apply the released SiLU-gated MLP ordering."""
    return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """Normalize the last dimension using the official accumulation order."""
    input_dtype = hidden_states.dtype
    hidden_states = hidden_states.to(torch.float32)
    variance = hidden_states.pow(2).mean(-1, keepdim=True)
    hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
    return (self.weight * hidden_states).to(input_dtype)

def forward(
    self,
    position_ids: torch.Tensor,
    maxima: tuple[int, ...] | None = None,
) -> torch.Tensor:
    """Gather and concatenate rotary frequencies for `(S, 3)` positions."""
    device = position_ids.device
    if maxima is None:
        maxima = tuple(int(value) for value in position_ids.max(dim=0).values.tolist())
    rebuild = self.freqs_cis is None or any(maximum >= length
                                            for maximum, length in zip(maxima, self.axes_lens, strict=True))
    if rebuild:
        for index, maximum in enumerate(maxima):
            if maximum >= self.axes_lens[index]:
                self.axes_lens[index] = int(maximum * 1.5) + 1
        self.freqs_cis = self._precompute(self.axes_dims, tuple(self.axes_lens), self.theta)
        self.freqs_cis = [table.to(device) for table in self.freqs_cis]
    elif self.freqs_cis[0].device != device:
        self.freqs_cis = [table.to(device) for table in self.freqs_cis]
    return torch.cat(
        [self.freqs_cis[index][position_ids[:, index]] for index in range(len(self.axes_dims))],
        dim=-1,
    )

def forward(self,
            tokens: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Score in fp32, select with correction bias, and weight without it."""
    with torch.amp.autocast(tokens.device.type, enabled=False):
        logits = F.linear(tokens.float(), self.weight.float())
    scores = F.softmax(logits, dim=-1) if self.score_func == "softmax" else logits.sigmoid()
    scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0)
    if self.n_group is not None and self.n_group > 1:
        top_indices = self._group_limited_topk(scores_for_choice)
    else:
        top_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1]
    top_scores = scores.gather(1, top_indices)
    if self.top_k > 1 and self.norm_topk_prob:
        top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20)
    top_scores = top_scores * self.route_scale
    return top_indices, top_scores.to(tokens.dtype), logits, scores, scores_for_choice

def forward(self, hidden_states: torch.Tensor, padding_mask: torch.Tensor | None = None) -> torch.Tensor:
    """Route token choices, zero padded routes, and add optional shared experts."""
    batch = hidden_states.shape[0]
    tokens = hidden_states.view(-1, self.hidden_size)
    top_indices, top_scores, logits, scores, scores_for_choice = self.router(tokens)
    del logits, scores, scores_for_choice
    if padding_mask is not None:
        mask = padding_mask.unsqueeze(-1).to(top_scores.dtype)
        top_scores = top_scores * mask
        top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-9)
        top_scores = top_scores * self.router.route_scale
    output = self._run_selected_experts(tokens, top_scores, top_indices)
    output = output.view(batch, -1, self.hidden_size)
    if self.shared_experts is not None:
        output = output + self.shared_experts(hidden_states)
    return output

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """Apply RMSNorm followed by the released two-layer SiLU projection."""
    hidden_states = self.norm(hidden_states)
    return self.linear_2(F.silu(self.linear_1(hidden_states)))

def forward(self, sample: torch.Tensor) -> torch.Tensor:
    """Apply the official linear-SiLU-linear timestep projection."""
    return self.linear_2(F.silu(self.linear_1(sample)))

def forward(
    self,
    hidden_states: torch.Tensor,
    timestep: torch.Tensor,
    encoder_hidden_states: torch.Tensor | list[torch.Tensor],
    encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
    guidance: torch.Tensor | None = None,
    encoder_attention_mask: torch.Tensor | None = None,
    return_dict: bool = True,
    **kwargs: Any,
) -> LingBotVideoTransformerOutput | tuple[torch.Tensor]:
    """Denoise video latents with joint video-text attention."""
    del encoder_hidden_states_image, guidance, kwargs
    if isinstance(encoder_hidden_states, list):
        if len(encoder_hidden_states) != 1:
            raise ValueError("LingBot-Video expects one text-encoder output tensor.")
        encoder_hidden_states = encoder_hidden_states[0]
    batch, channels, frames, height, width = hidden_states.shape
    patch_t, patch_h, patch_w = self.config.patch_size
    grid_t, grid_h, grid_w = frames // patch_t, height // patch_h, width // patch_w
    video_tokens = grid_t * grid_h * grid_w
    text_tokens = encoder_hidden_states.shape[1]
    device = hidden_states.device
    if encoder_attention_mask is None:
        encoder_attention_mask = torch.ones((batch, text_tokens), device=device, dtype=torch.bool)
    text_lengths = encoder_attention_mask.sum(dim=-1).long()

    patches = hidden_states.reshape(batch, channels, grid_t, patch_t, grid_h, patch_h, grid_w, patch_w)
    patches = patches.permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(batch, video_tokens,
                                                              patch_t * patch_h * patch_w * channels)
    video_hidden = self.patch_embedder(patches)
    text_hidden = self.text_embedder(encoder_hidden_states)
    joint = torch.cat([video_hidden, text_hidden], dim=1)

    rotary_parts: list[torch.Tensor] = []
    for index in range(batch):
        real_text_length = int(text_lengths[index].item())
        positions = _make_joint_position_ids(real_text_length, grid_t, grid_h, grid_w, device)
        maxima = (real_text_length + grid_t, grid_h - 1, grid_w - 1)
        rotary = self.rope(positions, maxima=maxima)
        if real_text_length < text_tokens:
            padding = torch.zeros(
                text_tokens - real_text_length,
                rotary.shape[-1],
                device=device,
                dtype=rotary.dtype,
            )
            rotary = torch.cat([rotary, padding], dim=0)
        rotary_parts.append(rotary)
    rotary_emb = torch.stack(rotary_parts, dim=0)

    attention_mask = None
    moe_padding_mask = None
    if bool((text_lengths < text_tokens).any()):
        key_mask = torch.cat(
            [
                torch.ones(batch, video_tokens, dtype=torch.bool, device=device),
                encoder_attention_mask.bool(),
            ],
            dim=1,
        )
        attention_mask = key_mask[:, None, None, :]
        moe_padding_mask = key_mask
    timestep_projection = self.time_proj(timestep.float())
    timestep_embedding = self.time_embedder(timestep_projection)
    token_embedding = timestep_embedding.unsqueeze(1).expand(batch, joint.shape[1], -1)
    original_joint_length: int | None = None
    if _sequence_parallel_world_size() > 1:
        # Match the official CP order: project token modulation before
        # placing every token-aligned tensor on the same padded shard.
        temb6 = self.time_modulation(token_embedding.reshape(-1,
                                                             self.hidden_size)).reshape(batch, joint.shape[1], -1)
        joint, original_joint_length = sequence_model_parallel_shard(joint, dim=1)
        rotary_emb, _ = sequence_model_parallel_shard(rotary_emb, dim=1)
        token_embedding, _ = sequence_model_parallel_shard(token_embedding, dim=1)
        temb6, _ = sequence_model_parallel_shard(temb6, dim=1)
        if moe_padding_mask is None:
            moe_padding_mask = torch.ones(batch, original_joint_length, dtype=torch.bool, device=device)
        moe_padding_mask, _ = sequence_model_parallel_shard(moe_padding_mask, dim=1)
        moe_padding_mask = moe_padding_mask.reshape(-1)
        temb6 = temb6.reshape(-1, 6 * self.hidden_size)
    else:
        temb6 = self.time_modulation(token_embedding.reshape(-1, self.hidden_size))
        if moe_padding_mask is not None:
            moe_padding_mask = moe_padding_mask.reshape(-1)

    for block in self.blocks:
        joint = block(
            joint,
            temb6,
            rotary_emb,
            attention_mask,
            moe_padding_mask,
            original_joint_length,
        )

    final_modulation = self.norm_out_modulation(token_embedding.reshape(-1, self.hidden_size))
    shift, scale = final_modulation.reshape(*joint.shape[:2], -1).chunk(2, dim=-1)
    final_hidden = self.norm_out(joint) * (1.0 + scale) + shift
    projected = self.proj_out(final_hidden.to(self.proj_out.weight.dtype))
    if original_joint_length is not None:
        projected = sequence_model_parallel_all_gather_with_unpad(projected, original_joint_length, dim=1)
    projected = projected[:, :video_tokens]
    output_channels = self.config.out_channels
    output = projected.reshape(batch, grid_t, grid_h, grid_w, patch_t, patch_h, patch_w, output_channels)
    output = output.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(batch, output_channels, frames, height, width)
    if not return_dict:
        return (output, )
    return LingBotVideoTransformerOutput(sample=output)

def to(self, *args: Any, **kwargs: Any) -> LingBotVideoTransformer3DModel:
    """Cast bulk weights while retaining the released fp32-sensitive modules."""
    device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
    if dtype is None or dtype == torch.float32:
        return super().to(*args, **kwargs)
    if not torch.is_floating_point(torch.empty((), dtype=dtype)):
        return super().to(*args, **kwargs)
    if device is not None:
        super().to(device=device, non_blocking=non_blocking)
    for name, parameter in self.named_parameters():
        if torch.is_floating_point(parameter):
            target_dtype = torch.float32 if _keep_in_fp32(name) else dtype
            parameter.data = parameter.data.to(target_dtype, non_blocking=non_blocking)
            if parameter.grad is not None:
                parameter.grad.data = parameter.grad.data.to(target_dtype, non_blocking=non_blocking)
    for name, buffer in self.named_buffers():
        if torch.is_floating_point(buffer):
            target_dtype = torch.float32 if _keep_in_fp32(name) else dtype
            buffer.data = buffer.data.to(target_dtype, non_blocking=non_blocking)
    return self