lingbot_video

@dataclass
class LingBotVideoTransformerOutput:
    """Output container matching the released transformer contract."""
sample: torch.Tensor</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoLinear
LingBotVideoLinear(input_size: int, output_size: int, bias: bool = True, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = '')
Replicated FastVideo linear with the tensor-only official call contract.
Source
class LingBotVideoLinear(ReplicatedLinear):
    """Replicated FastVideo linear with the tensor-only official call contract."""
def forward(self, hidden_states: torch.Tensor) -&gt; torch.Tensor:
    &quot;&quot;&quot;Return the projected tensor while preserving the normal weight surface.&quot;&quot;&quot;
    output, _ = super().forward(hidden_states)
    return output</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoLinear.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Return the projected tensor while preserving the normal weight surface.
Source
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
fastvideo.models.dits.lingbot_video.LingBotVideoRMSNorm
LingBotVideoRMSNorm(dim: int, eps: float = 1e-06)
RMSNorm with fp32 accumulation and input-dtype output.
Source
class LingBotVideoRMSNorm(nn.Module):
    """RMSNorm with fp32 accumulation and input-dtype output."""
def __init__(self, dim: int, eps: float = 1e-6) -&gt; None:
    super().__init__()
    self.weight = nn.Parameter(torch.ones(dim))
    self.variance_epsilon = eps

def forward(self, hidden_states: torch.Tensor) -&gt; torch.Tensor:
    &quot;&quot;&quot;Normalize the last dimension using the official accumulation order.&quot;&quot;&quot;
    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)</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoRMSNorm.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Normalize the last dimension using the official accumulation order.
Source
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)
fastvideo.models.dits.lingbot_video.LingBotVideoRotaryEmbedding
LingBotVideoRotaryEmbedding(axes_dims: tuple[int, ...], axes_lens: tuple[int, ...], theta: float)
Complex64 rotary table indexed by temporal and spatial positions.
Source
class LingBotVideoRotaryEmbedding(nn.Module):
    """Complex64 rotary table indexed by temporal and spatial positions."""
def __init__(self, axes_dims: tuple&#91;int, ...&#93;, axes_lens: tuple&#91;int, ...&#93;, theta: float) -&gt; None:
    super().__init__()
    self.axes_dims = tuple(axes_dims)
    self.axes_lens = list(axes_lens)
    self.theta = theta
    self.freqs_cis: list&#91;torch.Tensor&#93; | None = None

@staticmethod
def _precompute(dims: tuple&#91;int, ...&#93;, lengths: tuple&#91;int, ...&#93;, theta: float) -&gt; list&#91;torch.Tensor&#93;:
    &quot;&quot;&quot;Build the per-axis complex frequency tables on CPU.&quot;&quot;&quot;
    tables: list&#91;torch.Tensor&#93; = &#91;&#93;
    for dim, length in zip(dims, lengths, strict=True):
        frequencies = 1.0 / (theta**(torch.arange(0, dim, 2, dtype=torch.float64, device=&quot;cpu&quot;) / dim))
        positions = torch.arange(length, device=frequencies.device, dtype=torch.float64)
        phases = torch.outer(positions, frequencies).float()
        tables.append(torch.polar(torch.ones_like(phases), phases).to(torch.complex64))
    return tables

def forward(
    self,
    position_ids: torch.Tensor,
    maxima: tuple&#91;int, ...&#93; | None = None,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;Gather and concatenate rotary frequencies for `(S, 3)` positions.&quot;&quot;&quot;
    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 &gt;= length
                                            for maximum, length in zip(maxima, self.axes_lens, strict=True))
    if rebuild:
        for index, maximum in enumerate(maxima):
            if maximum &gt;= self.axes_lens&#91;index&#93;:
                self.axes_lens&#91;index&#93; = int(maximum * 1.5) + 1
        self.freqs_cis = self._precompute(self.axes_dims, tuple(self.axes_lens), self.theta)
        self.freqs_cis = &#91;table.to(device) for table in self.freqs_cis&#93;
    elif self.freqs_cis&#91;0&#93;.device != device:
        self.freqs_cis = &#91;table.to(device) for table in self.freqs_cis&#93;
    return torch.cat(
        &#91;self.freqs_cis&#91;index&#93;&#91;position_ids&#91;:, index&#93;&#93; for index in range(len(self.axes_dims))&#93;,
        dim=-1,
    )</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoRotaryEmbedding.forward
forward(position_ids: torch.Tensor, maxima: tuple[int, ...] | None = None) -> torch.Tensor
Gather and concatenate rotary frequencies for (S, 3) positions.
Source
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,
    )
fastvideo.models.dits.lingbot_video.LingBotVideoTextEmbedder
LingBotVideoTextEmbedder(text_dim: int, hidden_size: int)
Project Qwen3-VL hidden states into the DiT hidden dimension.
Source
class LingBotVideoTextEmbedder(nn.Module):
    """Project Qwen3-VL hidden states into the DiT hidden dimension."""
def __init__(self, text_dim: int, hidden_size: int) -&gt; None:
    super().__init__()
    self.norm = LingBotVideoRMSNorm(text_dim, eps=1e-6)
    self.linear_1 = LingBotVideoLinear(text_dim, hidden_size, bias=True)
    self.linear_2 = LingBotVideoLinear(hidden_size, hidden_size, bias=True)

def forward(self, hidden_states: torch.Tensor) -&gt; torch.Tensor:
    &quot;&quot;&quot;Apply RMSNorm followed by the released two-layer SiLU projection.&quot;&quot;&quot;
    hidden_states = self.norm(hidden_states)
    return self.linear_2(F.silu(self.linear_1(hidden_states)))</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoTextEmbedder.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Apply RMSNorm followed by the released two-layer SiLU projection.
Source
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)))
fastvideo.models.dits.lingbot_video.LingBotVideoAttention
LingBotVideoAttention(hidden_size: int, num_heads: int, norm_eps: float, qkv_bias: bool, out_bias: bool)
Joint video-text attention shared by the Dense and MoE variants.
Create released QKV projections, per-head norms, and output projection.
Source
class LingBotVideoAttention(nn.Module):
    """Joint video-text attention shared by the Dense and MoE variants."""
def __init__(
    self,
    hidden_size: int,
    num_heads: int,
    norm_eps: float,
    qkv_bias: bool,
    out_bias: bool,
) -&gt; None:
    &quot;&quot;&quot;Create released QKV projections, per-head norms, and output projection.&quot;&quot;&quot;
    super().__init__()
    self.num_heads = num_heads
    self.head_dim = hidden_size // num_heads
    self.to_q = LingBotVideoLinear(hidden_size, hidden_size, bias=qkv_bias)
    self.to_k = LingBotVideoLinear(hidden_size, hidden_size, bias=qkv_bias)
    self.to_v = LingBotVideoLinear(hidden_size, hidden_size, bias=qkv_bias)
    self.norm_q = LingBotVideoRMSNorm(self.head_dim, norm_eps)
    self.norm_k = LingBotVideoRMSNorm(self.head_dim, norm_eps)
    self.to_out = LingBotVideoLinear(hidden_size, hidden_size, bias=out_bias)

def forward(
    self,
    hidden_states: torch.Tensor,
    rotary_emb: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    original_seq_len: int | None = None,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;Project QKV, apply rotary embeddings, and run non-causal SDPA.&quot;&quot;&quot;
    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(&#91;query, key, value&#93;, dim=0)
        qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1)
        padded_seq_len = qkv.shape&#91;1&#93;
        query, key, value = qkv&#91;:, :original_seq_len&#93;.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))</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoAttention.forward
forward(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.
Source
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(&#91;query, key, value&#93;, dim=0)
    qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1)
    padded_seq_len = qkv.shape&#91;1&#93;
    query, key, value = qkv&#91;:, :original_seq_len&#93;.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))</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoMLP
LingBotVideoMLP(hidden_size: int, intermediate_size: int)
Dense SwiGLU feed-forward network.
Source
class LingBotVideoMLP(nn.Module):
    """Dense SwiGLU feed-forward network."""
def __init__(self, hidden_size: int, intermediate_size: int) -&gt; None:
    super().__init__()
    self.gate_proj = LingBotVideoLinear(hidden_size, intermediate_size, bias=False)
    self.up_proj = LingBotVideoLinear(hidden_size, intermediate_size, bias=False)
    self.down_proj = LingBotVideoLinear(intermediate_size, hidden_size, bias=False)

def forward(self, hidden_states: torch.Tensor) -&gt; torch.Tensor:
    &quot;&quot;&quot;Apply the released SiLU-gated MLP ordering.&quot;&quot;&quot;
    return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoMLP.forward
forward(hidden_states: torch.Tensor) -> torch.Tensor
Apply the released SiLU-gated MLP ordering.
Source
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))
fastvideo.models.dits.lingbot_video.LingBotVideoRouter
LingBotVideoRouter(hidden_size: int, num_experts: int, top_k: int, score_func: str, norm_topk_prob: bool, n_group: int | None, topk_group: int | None, route_scale: float)
Released token-choice router with bias-only expert selection correction.
Create fp32-routed expert scores with the released persistent bias.
Source
class LingBotVideoRouter(nn.Module):
    """Released token-choice router with bias-only expert selection correction."""
def __init__(
    self,
    hidden_size: int,
    num_experts: int,
    top_k: int,
    score_func: str,
    norm_topk_prob: bool,
    n_group: int | None,
    topk_group: int | None,
    route_scale: float,
) -&gt; None:
    &quot;&quot;&quot;Create fp32-routed expert scores with the released persistent bias.&quot;&quot;&quot;
    super().__init__()
    self.num_experts = num_experts
    self.top_k = top_k
    self.score_func = score_func
    self.norm_topk_prob = norm_topk_prob
    self.n_group = n_group
    self.topk_group = topk_group
    self.route_scale = route_scale
    self.weight = nn.Parameter(torch.empty(num_experts, hidden_size))
    self.register_buffer(&quot;e_score_correction_bias&quot;, torch.zeros(num_experts), persistent=True)

def _group_limited_topk(self, scores_for_choice: torch.Tensor) -&gt; torch.Tensor:
    &quot;&quot;&quot;Restrict token choices to groups with the two strongest expert scores.&quot;&quot;&quot;
    sequence_length = scores_for_choice.shape&#91;0&#93;
    experts_per_group = self.num_experts // self.n_group
    grouped = scores_for_choice.view(sequence_length, self.n_group, experts_per_group)
    group_scores = grouped.topk(2, dim=-1)&#91;0&#93;.sum(dim=-1)
    group_indices = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)&#91;1&#93;
    group_mask = torch.zeros_like(group_scores)
    group_mask.scatter_(1, group_indices, 1)
    score_mask = (group_mask.unsqueeze(-1).expand(sequence_length, self.n_group,
                                                  experts_per_group).reshape(sequence_length, -1))
    masked_scores = scores_for_choice.masked_fill(~score_mask.bool(), float(&quot;-inf&quot;))
    return torch.topk(masked_scores, k=self.top_k, dim=-1, sorted=False)&#91;1&#93;

def forward(self,
            tokens: torch.Tensor) -&gt; tuple&#91;torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor&#93;:
    &quot;&quot;&quot;Score in fp32, select with correction bias, and weight without it.&quot;&quot;&quot;
    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 == &quot;softmax&quot; 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 &gt; 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)&#91;1&#93;
    top_scores = scores.gather(1, top_indices)
    if self.top_k &gt; 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</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoRouter.forward
forward(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.
Source
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
fastvideo.models.dits.lingbot_video.LingBotVideoGroupedExperts
LingBotVideoGroupedExperts(num_experts: int, hidden_size: int, intermediate_size: int)
Released grouped-expert parameter layout: w1/w3 [E,I,H], w2 [E,H,I].
Source
class LingBotVideoGroupedExperts(nn.Module):
    """Released grouped-expert parameter layout: w1/w3 `[E,I,H]`, w2 `[E,H,I]`."""
def __init__(self, num_experts: int, hidden_size: int, intermediate_size: int) -&gt; None:
    super().__init__()
    self.num_experts = num_experts
    self.w1 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size))
    self.w2 = nn.Parameter(torch.empty(num_experts, hidden_size, intermediate_size))
    self.w3 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size))</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoSparseMoeBlock
LingBotVideoSparseMoeBlock(hidden_size: int, intermediate_size: int, num_experts: int, top_k: int, moe_intermediate_size: int, score_func: str, norm_topk_prob: bool, n_group: int | None, topk_group: int | None, routed_scaling_factor: float, n_shared_experts: int | None)
Token-choice sparse feed-forward block matching the released state surface.
Create routed and optional shared experts with released parameter names.
Source
class LingBotVideoSparseMoeBlock(nn.Module):
    """Token-choice sparse feed-forward block matching the released state surface."""
def __init__(
    self,
    hidden_size: int,
    intermediate_size: int,
    num_experts: int,
    top_k: int,
    moe_intermediate_size: int,
    score_func: str,
    norm_topk_prob: bool,
    n_group: int | None,
    topk_group: int | None,
    routed_scaling_factor: float,
    n_shared_experts: int | None,
) -&gt; None:
    &quot;&quot;&quot;Create routed and optional shared experts with released parameter names.&quot;&quot;&quot;
    super().__init__()
    del intermediate_size
    self.hidden_size = hidden_size
    self.num_experts = num_experts
    self.router = LingBotVideoRouter(
        hidden_size,
        num_experts,
        top_k,
        score_func,
        norm_topk_prob,
        n_group,
        topk_group,
        routed_scaling_factor,
    )
    self.experts = LingBotVideoGroupedExperts(num_experts, hidden_size, moe_intermediate_size)
    self.shared_experts = None
    if n_shared_experts is not None and n_shared_experts &gt; 0:
        self.shared_experts = LingBotVideoMLP(hidden_size, moe_intermediate_size * n_shared_experts)

@staticmethod
def _reorder_tokens(
    tokens: torch.Tensor,
    top_scores: torch.Tensor,
    top_indices: torch.Tensor,
    num_experts: int,
) -&gt; tuple&#91;torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int&#93;:
    &quot;&quot;&quot;Pack active token choices into stable expert-major order.&quot;&quot;&quot;
    num_tokens = tokens.shape&#91;0&#93;
    top_k = top_indices.shape&#91;1&#93;
    flat_scores = top_scores.reshape(-1)
    flat_indices = top_indices.reshape(-1)
    active_positions = torch.where(flat_scores != 0)&#91;0&#93;
    active_experts = flat_indices&#91;active_positions&#93;
    counts = torch.zeros(num_experts, device=tokens.device, dtype=torch.int64)
    counts.scatter_add_(0, active_experts, torch.ones_like(active_experts, dtype=torch.int64))
    sort_order = torch.argsort(active_experts, stable=True)
    sorted_positions = active_positions&#91;sort_order&#93;
    sorted_scores = flat_scores&#91;sorted_positions&#93;
    original_token_indices = sorted_positions // top_k
    permuted_tokens = tokens&#91;original_token_indices&#93;
    return permuted_tokens, counts, sorted_positions, sorted_scores, num_tokens, top_k

@staticmethod
def _pad_grouped_tokens(tokens: torch.Tensor,
                        counts: torch.Tensor,
                        align: int = 8) -&gt; tuple&#91;torch.Size, torch.Tensor, torch.Tensor, torch.Tensor&#93;:
    &quot;&quot;&quot;Align each expert segment for `torch._grouped_mm` and retain unpad indices.&quot;&quot;&quot;
    num_tokens = tokens.shape&#91;0&#93;
    num_experts = int(counts.shape&#91;0&#93;)
    max_length = _round_up_to_multiple(num_tokens + num_experts * align, align)
    # Build the small padding index on CPU after one synchronization instead
    # of reading three CUDA scalars for every expert.
    counts_cpu = counts.to(device=&quot;cpu&quot;, dtype=torch.int64)
    total_per_expert = torch.clamp_min(counts_cpu, align)
    aligned_counts_cpu = ((total_per_expert + align - 1) // align * align).to(torch.int32)
    write_offsets = torch.cumsum(aligned_counts_cpu, dim=0) - aligned_counts_cpu
    start_indices = torch.cumsum(counts_cpu, dim=0) - counts_cpu
    permuted_indices_cpu = torch.full((max_length, ), num_tokens, dtype=torch.int64, device=&quot;cpu&quot;)
    for expert_index in range(num_experts):
        length = int(counts_cpu&#91;expert_index&#93;)
        if length == 0:
            continue
        write_start = int(write_offsets&#91;expert_index&#93;)
        start = int(start_indices&#91;expert_index&#93;)
        permuted_indices_cpu&#91;write_start:write_start + length&#93; = torch.arange(start,
                                                                              start + length,
                                                                              device=&quot;cpu&quot;,
                                                                              dtype=torch.int64)
    permuted_indices = permuted_indices_cpu.to(tokens.device)
    aligned_counts = aligned_counts_cpu.to(tokens.device)
    tokens_with_pad = torch.vstack((tokens, tokens.new_zeros((tokens.shape&#91;-1&#93;, ))))
    input_shape = tokens_with_pad.shape
    return input_shape, tokens_with_pad&#91;permuted_indices&#93;, permuted_indices, aligned_counts

@staticmethod
def _unpad_grouped_tokens(output: torch.Tensor, input_shape: torch.Size,
                          permuted_indices: torch.Tensor) -&gt; torch.Tensor:
    &quot;&quot;&quot;Undo per-expert alignment while dropping the shared padding row.&quot;&quot;&quot;
    unpermuted = output.new_empty(input_shape)
    unpermuted&#91;permuted_indices, :&#93; = output
    return unpermuted&#91;:-1&#93;

def _run_grouped_experts(self, tokens: torch.Tensor, counts: torch.Tensor) -&gt; torch.Tensor:
    &quot;&quot;&quot;Use the released bf16 grouped matmuls on CUDA and an eager CPU fallback.&quot;&quot;&quot;
    if tokens.device.type == &quot;cpu&quot; or not hasattr(torch, &quot;_grouped_mm&quot;):
        return self._run_experts_for_loop(tokens, counts)
    input_shape, padded_tokens, permuted_indices, aligned_counts = self._pad_grouped_tokens(tokens, counts)
    offsets = torch.cumsum(aligned_counts, dim=0, dtype=torch.int32)
    hidden = F.silu(
        torch._grouped_mm(
            padded_tokens.bfloat16(),
            self.experts.w1.bfloat16().transpose(-2, -1),
            offs=offsets,
        ))
    hidden = hidden * torch._grouped_mm(
        padded_tokens.bfloat16(),
        self.experts.w3.bfloat16().transpose(-2, -1),
        offs=offsets,
    )
    output = torch._grouped_mm(
        hidden,
        self.experts.w2.bfloat16().transpose(-2, -1),
        offs=offsets,
    ).type_as(padded_tokens)
    return self._unpad_grouped_tokens(output, input_shape, permuted_indices)

def _run_experts_for_loop(self, tokens: torch.Tensor, counts: torch.Tensor) -&gt; torch.Tensor:
    &quot;&quot;&quot;Evaluate contiguous expert segments eagerly for CPU correctness tests.&quot;&quot;&quot;
    splits = torch.split(tokens, counts.tolist(), dim=0)
    outputs: list&#91;torch.Tensor&#93; = &#91;&#93;
    for expert_index, expert_tokens in enumerate(splits):
        if expert_tokens.numel() == 0:
            continue
        hidden = F.silu(expert_tokens @ self.experts.w1&#91;expert_index&#93;.transpose(-2, -1))
        hidden = hidden * (expert_tokens @ self.experts.w3&#91;expert_index&#93;.transpose(-2, -1))
        outputs.append(hidden @ self.experts.w2&#91;expert_index&#93;.transpose(-2, -1))
    if not outputs:
        return tokens.new_zeros(tokens.shape)
    return torch.cat(outputs, dim=0)

@staticmethod
def _restore_tokens(
    expert_output: torch.Tensor,
    sorted_positions: torch.Tensor,
    sorted_scores: torch.Tensor,
    num_tokens: int,
    top_k: int,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;Restore token order and combine expert outputs with fp32 weighted sums.&quot;&quot;&quot;
    hidden_size = expert_output.shape&#91;-1&#93;
    unsorted = torch.zeros(
        (num_tokens * top_k, hidden_size),
        dtype=expert_output.dtype,
        device=expert_output.device,
    )
    unsorted&#91;sorted_positions&#93; = expert_output
    unsorted = unsorted.reshape(num_tokens, top_k, hidden_size)
    scores_unsorted = torch.zeros(
        num_tokens * top_k,
        dtype=sorted_scores.dtype,
        device=sorted_scores.device,
    )
    scores_unsorted&#91;sorted_positions&#93; = sorted_scores
    scores_unsorted = scores_unsorted.reshape(num_tokens, top_k, 1)
    return (unsorted.float() * scores_unsorted).sum(dim=1).to(expert_output.dtype)

def _run_selected_experts(
    self,
    tokens: torch.Tensor,
    top_scores: torch.Tensor,
    top_indices: torch.Tensor,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;Dispatch routed choices, execute experts, and restore token-major order.&quot;&quot;&quot;
    permuted_tokens, counts, sorted_positions, sorted_scores, num_tokens, top_k = self._reorder_tokens(
        tokens, top_scores, top_indices, self.router.num_experts)
    expert_output = self._run_grouped_experts(permuted_tokens, counts)
    return self._restore_tokens(expert_output, sorted_positions, sorted_scores, num_tokens, top_k)

def forward(self, hidden_states: torch.Tensor, padding_mask: torch.Tensor | None = None) -&gt; torch.Tensor:
    &quot;&quot;&quot;Route token choices, zero padded routes, and add optional shared experts.&quot;&quot;&quot;
    batch = hidden_states.shape&#91;0&#93;
    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</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoSparseMoeBlock.forward
forward(hidden_states: torch.Tensor, padding_mask: torch.Tensor | None = None) -> torch.Tensor
Route token choices, zero padded routes, and add optional shared experts.
Source
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
fastvideo.models.dits.lingbot_video.LingBotVideoBlock
LingBotVideoBlock(hidden_size: int, num_attention_heads: int, intermediate_size: int, norm_eps: float, qkv_bias: bool, out_bias: bool, num_experts: int, num_experts_per_tok: int, moe_intermediate_size: int, decoder_sparse_step: int, mlp_only_layers: tuple[int, ...] | list[int], n_shared_experts: int | None, score_func: str, norm_topk_prob: bool, n_group: int | None, topk_group: int | None, routed_scaling_factor: float, layer_idx: int)
One Dense or sparse LingBot-Video transformer block.
Select the released Dense or sparse feed-forward structure for one layer.
Source
class LingBotVideoBlock(nn.Module):
    """One Dense or sparse LingBot-Video transformer block."""
def __init__(
    self,
    hidden_size: int,
    num_attention_heads: int,
    intermediate_size: int,
    norm_eps: float,
    qkv_bias: bool,
    out_bias: bool,
    num_experts: int,
    num_experts_per_tok: int,
    moe_intermediate_size: int,
    decoder_sparse_step: int,
    mlp_only_layers: tuple&#91;int, ...&#93; | list&#91;int&#93;,
    n_shared_experts: int | None,
    score_func: str,
    norm_topk_prob: bool,
    n_group: int | None,
    topk_group: int | None,
    routed_scaling_factor: float,
    layer_idx: int,
) -&gt; None:
    &quot;&quot;&quot;Select the released Dense or sparse feed-forward structure for one layer.&quot;&quot;&quot;
    super().__init__()
    self.layer_idx = layer_idx
    self.scale_shift_table = nn.Parameter(torch.zeros(1, 6 * hidden_size))
    self.norm1 = LingBotVideoRMSNorm(hidden_size, norm_eps)
    self.attn = LingBotVideoAttention(hidden_size, num_attention_heads, norm_eps, qkv_bias, out_bias)
    self.norm_post_attn = LingBotVideoRMSNorm(hidden_size, norm_eps)
    self.norm2 = LingBotVideoRMSNorm(hidden_size, norm_eps)
    if layer_idx not in mlp_only_layers and (num_experts &gt; 0 and (layer_idx + 1) % decoder_sparse_step == 0):
        self.ffn = LingBotVideoSparseMoeBlock(
            hidden_size,
            intermediate_size,
            num_experts,
            num_experts_per_tok,
            moe_intermediate_size,
            score_func,
            norm_topk_prob,
            n_group,
            topk_group,
            routed_scaling_factor,
            n_shared_experts,
        )
    else:
        self.ffn = LingBotVideoMLP(hidden_size, intermediate_size)
    self.norm_post_ffn = LingBotVideoRMSNorm(hidden_size, norm_eps)

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,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;Run attention and configured feed-forward residual branches with fp32 AdaLN.&quot;&quot;&quot;
    expected_tokens = hidden_states.shape&#91;0&#93; * hidden_states.shape&#91;1&#93;
    if temb6.ndim != 2 or temb6.shape&#91;0&#93; != expected_tokens:
        raise ValueError(&quot;LingBotVideoBlock expects token-level temb6 with shape &quot;
                         f&quot;(B*S, 6D); got &#123;tuple(temb6.shape)&#125; for &#123;tuple(hidden_states.shape)&#125;.&quot;)
    modulation = temb6.view(*hidden_states.shape&#91;:2&#93;, -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)</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoBlock.forward
forward(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.
Source
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)</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoTimestepEmbedding
LingBotVideoTimestepEmbedding(input_dim: int, hidden_size: int, bias: bool)
Two-layer timestep embedding with released parameter names.
Source
class LingBotVideoTimestepEmbedding(nn.Module):
    """Two-layer timestep embedding with released parameter names."""
def __init__(self, input_dim: int, hidden_size: int, bias: bool) -&gt; None:
    super().__init__()
    self.linear_1 = LingBotVideoLinear(input_dim, hidden_size, bias=bias)
    self.linear_2 = LingBotVideoLinear(hidden_size, hidden_size, bias=bias)

def forward(self, sample: torch.Tensor) -&gt; torch.Tensor:
    &quot;&quot;&quot;Apply the official linear-SiLU-linear timestep projection.&quot;&quot;&quot;
    return self.linear_2(F.silu(self.linear_1(sample)))</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoTimestepEmbedding.forward
forward(sample: torch.Tensor) -> torch.Tensor
Apply the official linear-SiLU-linear timestep projection.
Source
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)))
fastvideo.models.dits.lingbot_video.LingBotVideoTransformer3DModel
LingBotVideoTransformer3DModel(config: LingBotVideoConfig, hf_config: dict[str, Any])
LingBot-Video DiT with a source-compatible Dense or MoE state surface.
Construct the Dense or MoE variant from its released transformer config.
Source
class LingBotVideoTransformer3DModel(BaseDiT):
    """LingBot-Video DiT with a source-compatible Dense or MoE state surface."""
_fsdp_shard_conditions = &#91;_is_lingbot_video_block&#93;
_compile_conditions = &#91;_is_lingbot_video_block&#93;
_supported_attention_backends = (
    AttentionBackendEnum.TORCH_SDPA,
    AttentionBackendEnum.FLASH_ATTN,
)
param_names_mapping = &#123;r&quot;^(.*)$&quot;: r&quot;\1&quot;&#125;
reverse_param_names_mapping = &#123;r&quot;^(.*)$&quot;: r&quot;\1&quot;&#125;

def _get_parameter_dtype(self, name: str, default_dtype: torch.dtype) -&gt; torch.dtype:
    &quot;&quot;&quot;Select the released mixed-precision dtype while loading each parameter.&quot;&quot;&quot;
    return torch.float32 if _keep_in_fp32(name) else default_dtype

def __init__(self, config: LingBotVideoConfig, hf_config: dict&#91;str, Any&#93;) -&gt; None:
    &quot;&quot;&quot;Construct the Dense or MoE variant from its released transformer config.&quot;&quot;&quot;
    config.update_model_arch(hf_config)
    super().__init__(config, hf_config)
    head_dim = config.hidden_size // config.num_attention_heads
    if head_dim != sum(config.axes_dims):
        raise ValueError(f&quot;head_dim &#123;head_dim&#125; != sum(axes_dims) &#123;sum(config.axes_dims)&#125;&quot;)
    sp_world_size = _sequence_parallel_world_size()
    assert config.num_attention_heads % sp_world_size == 0, (
        f&quot;The number of attention heads (&#123;config.num_attention_heads&#125;) must be divisible by &quot;
        f&quot;the sequence parallel size (&#123;sp_world_size&#125;)&quot;)
    self.hidden_size = config.hidden_size
    self.num_attention_heads = config.num_attention_heads
    self.num_channels_latents = config.in_channels
    self.patch_embedder = LingBotVideoLinear(
        config.in_channels * math.prod(config.patch_size),
        config.hidden_size,
        bias=config.patch_embed_bias,
    )
    self.time_proj = Timesteps(config.freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
    self.time_embedder = LingBotVideoTimestepEmbedding(config.freq_dim, config.hidden_size,
                                                       config.timestep_mlp_bias)
    self.time_modulation = nn.Sequential(nn.SiLU(), LingBotVideoLinear(config.hidden_size, 6 * config.hidden_size))
    self.text_embedder = LingBotVideoTextEmbedder(config.text_dim, config.hidden_size)
    self.rope = LingBotVideoRotaryEmbedding(tuple(config.axes_dims), tuple(config.axes_lens), config.rope_theta)
    self.blocks = nn.ModuleList(&#91;
        LingBotVideoBlock(
            config.hidden_size,
            config.num_attention_heads,
            config.intermediate_size,
            config.norm_eps,
            config.qkv_bias,
            config.out_bias,
            config.num_experts,
            config.num_experts_per_tok,
            config.moe_intermediate_size,
            config.decoder_sparse_step,
            config.mlp_only_layers,
            config.n_shared_experts,
            config.score_func,
            config.norm_topk_prob,
            config.n_group,
            config.topk_group,
            config.routed_scaling_factor,
            layer_index,
        ) for layer_index in range(config.depth)
    &#93;)
    self.norm_out = nn.LayerNorm(config.hidden_size, elementwise_affine=False, eps=config.norm_eps)
    self.norm_out_modulation = nn.Sequential(nn.SiLU(),
                                             LingBotVideoLinear(config.hidden_size, 2 * config.hidden_size))
    self.proj_out = LingBotVideoLinear(config.hidden_size, math.prod(config.patch_size) * config.out_channels)
    self.__post_init__()

def to(self, *args: Any, **kwargs: Any) -&gt; LingBotVideoTransformer3DModel:
    &quot;&quot;&quot;Cast bulk weights while retaining the released fp32-sensitive modules.&quot;&quot;&quot;
    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

def forward(
    self,
    hidden_states: torch.Tensor,
    timestep: torch.Tensor,
    encoder_hidden_states: torch.Tensor | list&#91;torch.Tensor&#93;,
    encoder_hidden_states_image: torch.Tensor | list&#91;torch.Tensor&#93; | None = None,
    guidance: torch.Tensor | None = None,
    encoder_attention_mask: torch.Tensor | None = None,
    return_dict: bool = True,
    **kwargs: Any,
) -&gt; LingBotVideoTransformerOutput | tuple&#91;torch.Tensor&#93;:
    &quot;&quot;&quot;Denoise video latents with joint video-text attention.&quot;&quot;&quot;
    del encoder_hidden_states_image, guidance, kwargs
    if isinstance(encoder_hidden_states, list):
        if len(encoder_hidden_states) != 1:
            raise ValueError(&quot;LingBot-Video expects one text-encoder output tensor.&quot;)
        encoder_hidden_states = encoder_hidden_states&#91;0&#93;
    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&#91;1&#93;
    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(&#91;video_hidden, text_hidden&#93;, dim=1)

    rotary_parts: list&#91;torch.Tensor&#93; = &#91;&#93;
    for index in range(batch):
        real_text_length = int(text_lengths&#91;index&#93;.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 &lt; text_tokens:
            padding = torch.zeros(
                text_tokens - real_text_length,
                rotary.shape&#91;-1&#93;,
                device=device,
                dtype=rotary.dtype,
            )
            rotary = torch.cat(&#91;rotary, padding&#93;, 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 &lt; text_tokens).any()):
        key_mask = torch.cat(
            &#91;
                torch.ones(batch, video_tokens, dtype=torch.bool, device=device),
                encoder_attention_mask.bool(),
            &#93;,
            dim=1,
        )
        attention_mask = key_mask&#91;:, None, None, :&#93;
        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&#91;1&#93;, -1)
    original_joint_length: int | None = None
    if _sequence_parallel_world_size() &gt; 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&#91;1&#93;, -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&#91;:2&#93;, -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&#91;:, :video_tokens&#93;
    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)</code></pre></details>
fastvideo.models.dits.lingbot_video.LingBotVideoTransformer3DModel.to
to(*args: Any, **kwargs: Any) -> LingBotVideoTransformer3DModel
Cast bulk weights while retaining the released fp32-sensitive modules.
Source
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
fastvideo.models.dits.lingbot_video.LingBotVideoTransformer3DModel.forward
forward(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.
Source
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(&#91;video_hidden, text_hidden&#93;, dim=1)

rotary_parts: list&#91;torch.Tensor&#93; = &#91;&#93;
for index in range(batch):
    real_text_length = int(text_lengths&#91;index&#93;.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 &lt; text_tokens:
        padding = torch.zeros(
            text_tokens - real_text_length,
            rotary.shape&#91;-1&#93;,
            device=device,
            dtype=rotary.dtype,
        )
        rotary = torch.cat(&#91;rotary, padding&#93;, 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 &lt; text_tokens).any()):
    key_mask = torch.cat(
        &#91;
            torch.ones(batch, video_tokens, dtype=torch.bool, device=device),
            encoder_attention_mask.bool(),
        &#93;,
        dim=1,
    )
    attention_mask = key_mask&#91;:, None, None, :&#93;
    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&#91;1&#93;, -1)
original_joint_length: int | None = None
if _sequence_parallel_world_size() &gt; 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&#91;1&#93;, -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&#91;:2&#93;, -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&#91;:, :video_tokens&#93;
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)</code></pre></details>