`fastvideo.models.dits.lingbotworld.model.DistributedAttention`

DistributedAttention(num_heads: int, head_size: int, num_kv_heads: int | None = None, softmax_scale: float | None = None, causal: bool = False, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, prefix: str = '', **extra_impl_args)

Distributed attention layer.

`fastvideo.models.dits.lingbotworld.model.DistributedAttention.forward`

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, original_seq_len: int | None = None, replicated_q: torch.Tensor | None = None, replicated_k: torch.Tensor | None = None, replicated_v: torch.Tensor | None = None, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, attn_metadata: AttentionMetadata | None = None) -> tuple[torch.Tensor, torch.Tensor | None]

Forward pass for distributed attention.

Parameters

q: torch.Tensor: Query tensor [batch_size, seq_len, num_heads, head_dim]
k: torch.Tensor: Key tensor [batch_size, seq_len, num_heads, head_dim]
v: torch.Tensor: Value tensor [batch_size, seq_len, num_heads, head_dim]
original_seq_len: int: Original (unpadded) full sequence length
replicated_q: Optional[torch.Tensor]: Replicated query tensor, typically for text tokens
replicated_k: Optional[torch.Tensor]: Replicated key tensor
replicated_v: Optional[torch.Tensor]: Replicated value tensor

Returns

tuple[torch.Tensor, torch.Tensor | None]: Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing: - o (torch.Tensor): Output tensor after attention for the main sequence - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided

`fastvideo.models.dits.lingbotworld.model.LingBotWorldVideoConfig`

LingBotWorldVideoConfig(arch_config: DiTArchConfig = LingBotWorldArchConfig(), prefix: str = 'Wan', quant_config: QuantizationConfig | None = None, *, _resolved_attention_backend: AttentionBackendEnum | None = None)

`fastvideo.models.dits.lingbotworld.model.sequence_model_parallel_all_gather_with_unpad`

sequence_model_parallel_all_gather_with_unpad(input_: torch.Tensor, original_seq_len: int, dim: int = -1) -> torch.Tensor

All-gather the input tensor and remove padding.

Parameters

input_: torch.Tensor: Sharded (and possibly padded) tensor to gather
original_seq_len: int: Original sequence length before padding
dim: int: Dimension to gather along (default: -1)

Returns

Tensor: torch.Tensor: Gathered and unpadded tensor

`fastvideo.models.dits.lingbotworld.model.sequence_model_parallel_shard`

sequence_model_parallel_shard(input_: torch.Tensor, dim: int = 1) -> tuple[torch.Tensor, int]

Shard the input tensor across model parallel group with optional padding.

Parameters

input_: torch.Tensor: Input tensor to shard
dim: int: Dimension to shard along (default: 1)

Returns

tuple: tuple[torch.Tensor, int]: (sharded_tensor, original_seq_len) - sharded_tensor: The sharded (and possibly padded) tensor - original_seq_len: Original sequence length before padding

`fastvideo.models.dits.lingbotworld.model.FP32LayerNorm`

FP32LayerNorm()

`fastvideo.models.dits.lingbotworld.model.FP32LayerNorm.forward`

forward(inputs: torch.Tensor) -> torch.Tensor

`fastvideo.models.dits.lingbotworld.model.LayerNormScaleShift`

LayerNormScaleShift(hidden_size: int, norm_type: str = 'rms', eps: float = 1e-06, elementwise_affine: bool = False, dtype: torch.dtype = torch.float32, compute_dtype: torch.dtype | None = None, prefix: str = '')

Fused operation that combines LayerNorm with scale and shift operations. This reduces memory bandwidth by combining memory-bound operations.

`fastvideo.models.dits.lingbotworld.model.LayerNormScaleShift.forward`

forward(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor, convert_modulation_dtype: bool = False) -> torch.Tensor

Apply ln followed by scale and shift in a single fused operation.

`fastvideo.models.dits.lingbotworld.model.RMSNorm`

RMSNorm(hidden_size: int, eps: float = 1e-06, dtype: torch.dtype = torch.float32, var_hidden_size: int | None = None, has_weight: bool = True)

Root mean square normalization.

Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight. Refer to https://arxiv.org/abs/1910.07467

`fastvideo.models.dits.lingbotworld.model.RMSNorm.forward_native`

forward_native(x: torch.Tensor, residual: torch.Tensor | None = None) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]

PyTorch-native implementation equivalent to forward().

`fastvideo.models.dits.lingbotworld.model.RMSNorm.extra_repr`

extra_repr() -> str

`fastvideo.models.dits.lingbotworld.model.ScaleResidual`

ScaleResidual(prefix: str = '')

Applies gated residual connection.

`fastvideo.models.dits.lingbotworld.model.ScaleResidual.forward`

forward(residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor

Apply gated residual connection.

`fastvideo.models.dits.lingbotworld.model.ScaleResidualLayerNormScaleShift`

ScaleResidualLayerNormScaleShift(hidden_size: int, norm_type: str = 'rms', eps: float = 1e-06, elementwise_affine: bool = False, dtype: torch.dtype = torch.float32, compute_dtype: torch.dtype | None = None, prefix: str = '', fuse_inference: bool = False)

Fused operation that combines: 1. Gated residual connection 2. LayerNorm 3. Scale and shift operations

This reduces memory bandwidth by combining memory-bound operations.

`fastvideo.models.dits.lingbotworld.model.ScaleResidualLayerNormScaleShift.forward`

forward(residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor | int, shift: torch.Tensor | None, scale: torch.Tensor | None, convert_modulation_dtype: bool = False) -> tuple[torch.Tensor, torch.Tensor]

Apply gated residual connection, followed by layernorm and scale/shift in a single fused operation.

Pass shift=None, scale=None (both, not just one) to skip the modulation entirely and return the normalized tensor as-is. Callers that need norm-only behavior (e.g. Wan's post-self-attn site) would otherwise pay two full-tensor elementwise passes computing normalized * (1.0 + 0) + 0.

Returns

torch.Tensor

Tuple containing:

torch.Tensor

normalized and modulated output

tuple[torch.Tensor, torch.Tensor]

residual value (value after residual connection but before normalization)

`fastvideo.models.dits.lingbotworld.model.ReplicatedLinear`

ReplicatedLinear(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 linear layer.

Parameters

input_size: int: input dimension of the linear layer.
output_size: int: output dimension of the linear layer.
bias: bool: If true, add bias.
skip_bias_add: bool: If true, skip adding bias but instead return it.
params_dtype: torch.dtype | None: Data type for the parameters.
quant_config: QuantizationConfig | None: Quantization configure.
prefix: str: The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj)

`fastvideo.models.dits.lingbotworld.model.ReplicatedLinear.weight_loader`

weight_loader(param: Parameter, loaded_weight: torch.Tensor) -> None

`fastvideo.models.dits.lingbotworld.model.ReplicatedLinear.forward`

forward(x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]

`fastvideo.models.dits.lingbotworld.model.ReplicatedLinear.extra_repr`

extra_repr() -> str

`fastvideo.models.dits.lingbotworld.model.ReplicatedLinear.get_shape_mapping`

get_shape_mapping() -> dict

Get the mapping from (input_shape, output_shape) to layer types.

`fastvideo.models.dits.lingbotworld.model.ReplicatedLinear.reset_shape_tracking`

reset_shape_tracking() -> None

Clear tracked shapes and layer type mappings.

`fastvideo.models.dits.lingbotworld.model.ReplicatedLinear.print_shape_summary`

print_shape_summary() -> None

Log a summary of all unique shapes and their layer types.

`fastvideo.models.dits.lingbotworld.model.MLP`

MLP(input_dim: int, mlp_hidden_dim: int, output_dim: int | None = None, bias: bool = True, act_type: str = 'gelu_pytorch_tanh', dtype: torch.dtype | None = None, prefix: str = '', quant_config: QuantizationConfig | None = None)

MLP for DiT blocks, NO gated linear units

`fastvideo.models.dits.lingbotworld.model.MLP.forward`

forward(x: torch.Tensor) -> torch.Tensor

`fastvideo.models.dits.lingbotworld.model.get_rotary_pos_embed`

get_rotary_pos_embed(rope_sizes, hidden_size, heads_num, rope_dim_list, rope_theta, theta_rescale_factor = 1.0, interpolation_factor = 1.0, shard_dim: int = 0, do_sp_sharding: bool = False, dtype: torch.dtype = torch.float32, start_frame: int = 0, use_real: bool = True) -> tuple[torch.Tensor, torch.Tensor]

Generate rotary positional embeddings for the given sizes.

Parameters

rope_sizes: Tuple of dimensions (t, h, w)
hidden_size: Hidden dimension size
heads_num: Number of attention heads
rope_dim_list: List of dimensions for each axis, or None
rope_theta: Base for frequency calculations
theta_rescale_factor: Rescale factor for theta. Defaults to 1.0
interpolation_factor: Factor to scale positions. Defaults to 1.0
shard_dim: int: Which dimension to shard for sequence parallelism. Defaults to 0.
do_sp_sharding: bool: Whether to shard the positional embeddings for sequence parallelism. Defaults to False.
use_real: bool: If True, output full head_dim for rotate_half style RoPE. Defaults to True.

Returns

tuple[torch.Tensor, torch.Tensor]: Tuple of (cos, sin) tensors for rotary embeddings. Shape [S, D] if use_real, [S, D/2] otherwise.

`fastvideo.models.dits.lingbotworld.model.PatchEmbed`

PatchEmbed(patch_size = 16, in_chans = 3, embed_dim = 768, norm_layer = None, flatten = True, bias = True, dtype = None, prefix: str = '')

2D Image to Patch Embedding

Image to Patch Embedding using Conv2d

A convolution based approach to patchifying a 2D image w/ embedding projection.

Based on the impl in https://github.com/google-research/vision_transformer

Remove the _assert function in forward function to be compatible with multi-resolution images.

`fastvideo.models.dits.lingbotworld.model.PatchEmbed.forward`

forward(x)

`fastvideo.models.dits.lingbotworld.model.WanCamControlPatchEmbedding`

WanCamControlPatchEmbedding(patch_size = (1, 2, 2), in_chans = 384, embed_dim = 2048, bias = True, dtype = None, prefix: str = '')

Lingbot World Patch embedding for Plucker features.

`fastvideo.models.dits.lingbotworld.model.WanCamControlPatchEmbedding.forward`

forward(x: torch.Tensor) -> torch.Tensor

`fastvideo.models.dits.lingbotworld.model.init_logger`

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

`fastvideo.models.dits.lingbotworld.model.BaseDiT`

BaseDiT(config: DiTConfig, hf_config: dict[str, Any], **kwargs)

`fastvideo.models.dits.lingbotworld.model.BaseDiT.forward`

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, guidance = None, **kwargs) -> torch.Tensor

`fastvideo.models.dits.lingbotworld.model.BaseDiT.preprocess_loaded_weight`

preprocess_loaded_weight(name: str, tensor: torch.Tensor) -> torch.Tensor

Adapt a checkpoint tensor before sharding; default models are unchanged.

`fastvideo.models.dits.lingbotworld.model.BaseDiT.inference_config_overrides`

inference_config_overrides() -> dict[str, Any]

Extra config.json keys a deployable checkpoint of this model needs.

`fastvideo.models.dits.lingbotworld.model.BaseDiT.inference_export_expected_shape`

inference_export_expected_shape(output_key: str, base_shape: tuple[int, ...]) -> tuple[int, ...]

The exported shape of output_key given the base checkpoint's shape.

`fastvideo.models.dits.lingbotworld.model.WanI2VCrossAttention`

WanI2VCrossAttention(dim: int, num_heads: int, window_size = (-1, -1), qk_norm = True, eps = 1e-06, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, quant_config: QuantizationConfig | None = None, prefix: str = '')

`fastvideo.models.dits.lingbotworld.model.WanI2VCrossAttention.forward`

forward(x, context, context_lens)

Parameters

x: Tensor: Shape [B, L1, C]
context: Tensor: Shape [B, L2, C]
context_lens: Tensor: Shape [B]

`fastvideo.models.dits.lingbotworld.model.WanT2VCrossAttention`

WanT2VCrossAttention(dim: int, num_heads: int, window_size = (-1, -1), qk_norm = True, eps = 1e-06, parallel_attention = False, quant_config: QuantizationConfig | None = None, prefix: str = '')

`fastvideo.models.dits.lingbotworld.model.WanT2VCrossAttention.forward`

forward(x, context, context_lens, crossattn_cache = None)

Parameters

x: Tensor: Shape [B, L1, C]
context: Tensor: Shape [B, L2, C]
context_lens: Tensor: Shape [B]

`fastvideo.models.dits.lingbotworld.model.WanTimeTextImageEmbedding`

WanTimeTextImageEmbedding(dim: int, time_freq_dim: int, text_embed_dim: int, image_embed_dim: int | None = None, *, r_embedder: bool = False, r_embedder_fusion: str = 'additive', r_embedder_gate_value: float = 0.25, r_embedder_deltatime_type: str = 'r')

`fastvideo.models.dits.lingbotworld.model.WanTimeTextImageEmbedding.forward`

forward(timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, encoder_hidden_states_image: torch.Tensor | None = None, timestep_seq_len: int | None = None, r_timestep: torch.Tensor | None = None)

`fastvideo.models.dits.lingbotworld.model.AttentionBackendEnum`

AttentionBackendEnum()

`fastvideo.models.dits.lingbotworld.model.get_sp_world_size`

get_sp_world_size() -> int

Return world size for the sequence model parallel group.

`fastvideo.models.dits.lingbotworld.model.LingBotWorldCamConditioner`

LingBotWorldCamConditioner(dim: int)

Source

class LingBotWorldCamConditioner(nn.Module):

    def __init__(self, dim: int) -> None:
        super().__init__()
        self.cam_injector = MLP(dim, dim, dim, bias=True, act_type="silu")
        self.cam_scale_layer = nn.Linear(dim, dim)
        self.cam_shift_layer = nn.Linear(dim, dim)

    def forward(
        self,
        hidden_states: torch.Tensor,
        c2ws_plucker_emb: torch.Tensor | None,
    ) -> torch.Tensor:
        if c2ws_plucker_emb is None:
            return hidden_states
        assert c2ws_plucker_emb.shape == hidden_states.shape, (
            f"c2ws_plucker_emb shape must match hidden_states shape, got "
            f"{tuple(c2ws_plucker_emb.shape)} vs {tuple(hidden_states.shape)}")
        c2ws_hidden_states = self.cam_injector(c2ws_plucker_emb)
        c2ws_hidden_states = c2ws_hidden_states + c2ws_plucker_emb
        cam_scale = self.cam_scale_layer(c2ws_hidden_states)
        cam_shift = self.cam_shift_layer(c2ws_hidden_states)
        return (1.0 + cam_scale) * hidden_states + cam_shift

`fastvideo.models.dits.lingbotworld.model.LingBotWorldCamConditioner.forward`

forward(hidden_states: torch.Tensor, c2ws_plucker_emb: torch.Tensor | None) -> torch.Tensor

Source

def forward(
    self,
    hidden_states: torch.Tensor,
    c2ws_plucker_emb: torch.Tensor | None,
) -> torch.Tensor:
    if c2ws_plucker_emb is None:
        return hidden_states
    assert c2ws_plucker_emb.shape == hidden_states.shape, (
        f"c2ws_plucker_emb shape must match hidden_states shape, got "
        f"{tuple(c2ws_plucker_emb.shape)} vs {tuple(hidden_states.shape)}")
    c2ws_hidden_states = self.cam_injector(c2ws_plucker_emb)
    c2ws_hidden_states = c2ws_hidden_states + c2ws_plucker_emb
    cam_scale = self.cam_scale_layer(c2ws_hidden_states)
    cam_shift = self.cam_shift_layer(c2ws_hidden_states)
    return (1.0 + cam_scale) * hidden_states + cam_shift

`fastvideo.models.dits.lingbotworld.model.LingBotWorldTransformerBlock`

LingBotWorldTransformerBlock(dim: int, ffn_dim: int, num_heads: int, qk_norm: str = 'rms_norm_across_heads', cross_attn_norm: bool = False, eps: float = 1e-06, added_kv_proj_dim: int | None = None, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, prefix: str = '')

Source

class LingBotWorldTransformerBlock(nn.Module):

    def __init__(self,
                 dim: int,
                 ffn_dim: int,
                 num_heads: int,
                 qk_norm: str = "rms_norm_across_heads",
                 cross_attn_norm: bool = False,
                 eps: float = 1e-6,
                 added_kv_proj_dim: int | None = None,
                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
                 | None = None,
                 prefix: str = ""):
        super().__init__()

        # 1. Self-attention
        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
        self.to_q = ReplicatedLinear(dim, dim, bias=True)
        self.to_k = ReplicatedLinear(dim, dim, bias=True)
        self.to_v = ReplicatedLinear(dim, dim, bias=True)

        self.to_out = ReplicatedLinear(dim, dim, bias=True)
        self.attn1 = DistributedAttention(num_heads=num_heads,
                                          head_size=dim // num_heads,
                                          causal=False,
                                          supported_attention_backends=supported_attention_backends,
                                          prefix=f"{prefix}.attn1")
        self.hidden_dim = dim
        self.num_attention_heads = num_heads
        dim_head = dim // num_heads
        if qk_norm == "rms_norm":
            self.norm_q = RMSNorm(dim_head, eps=eps)
            self.norm_k = RMSNorm(dim_head, eps=eps)
        elif qk_norm == "rms_norm_across_heads":
            # LTX applies qk norm across all heads
            self.norm_q = RMSNorm(dim, eps=eps)
            self.norm_k = RMSNorm(dim, eps=eps)
        else:
            raise NotImplementedError(f"QK Norm type '{qk_norm}' not supported")
        assert cross_attn_norm is True
        self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim,
                                                                        norm_type="layer",
                                                                        eps=eps,
                                                                        elementwise_affine=True,
                                                                        dtype=torch.float32,
                                                                        compute_dtype=torch.float32)

        # 2. Cross-attention
        if added_kv_proj_dim is not None:
            # I2V
            self.attn2 = WanI2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps)
        else:
            # T2V
            self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps)
        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim,
                                                                         norm_type="layer",
                                                                         eps=eps,
                                                                         elementwise_affine=False,
                                                                         dtype=torch.float32,
                                                                         compute_dtype=torch.float32)

        # 3. Feed-forward
        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
        self.mlp_residual = ScaleResidual()

        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
        self.cam_conditioner = LingBotWorldCamConditioner(dim)

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        temb: torch.Tensor,
        freqs_cis: tuple[torch.Tensor, torch.Tensor],
        original_seq_len: int,
        c2ws_plucker_emb: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if hidden_states.dim() == 4:
            hidden_states = hidden_states.squeeze(1)
        bs, seq_length, _ = hidden_states.shape
        orig_dtype = hidden_states.dtype
        # assert orig_dtype != torch.float32

        if temb.dim() == 4:
            # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
                self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2)
            # batch_size, seq_len, 1, inner_dim
            shift_msa = shift_msa.squeeze(2)
            scale_msa = scale_msa.squeeze(2)
            gate_msa = gate_msa.squeeze(2)
            c_shift_msa = c_shift_msa.squeeze(2)
            c_scale_msa = c_scale_msa.squeeze(2)
            c_gate_msa = c_gate_msa.squeeze(2)
        else:
            # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B)
            e = self.scale_shift_table + temb.float()
            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1)
        assert shift_msa.dtype == torch.float32

        # 1. Self-attention
        norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).to(orig_dtype)
        query, _ = self.to_q(norm_hidden_states)
        key, _ = self.to_k(norm_hidden_states)
        value, _ = self.to_v(norm_hidden_states)

        if self.norm_q is not None:
            query = self.norm_q(query)
        if self.norm_k is not None:
            key = self.norm_k(key)

        query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
        key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
        value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))

        attn_output, _ = self.attn1(
            query,
            key,
            value,
            original_seq_len,
            freqs_cis=freqs_cis,
        )
        attn_output = attn_output.flatten(2)
        attn_output, _ = self.to_out(attn_output)
        attn_output = attn_output.squeeze(1)

        null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
        norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa,
                                                                         null_shift, null_scale)
        norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype)
        # Inject camera condition
        # must be applied after the self-attention residual update.
        hidden_states = self.cam_conditioner(hidden_states, c2ws_plucker_emb)
        norm_hidden_states = self.self_attn_residual_norm.norm(hidden_states)
        norm_hidden_states = norm_hidden_states.to(orig_dtype)

        # 2. Cross-attention
        attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None)
        norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa,
                                                                          c_scale_msa)
        norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype)

        # 3. Feed-forward
        ff_output = self.ffn(norm_hidden_states)
        hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
        hidden_states = hidden_states.to(orig_dtype)

        return hidden_states

`fastvideo.models.dits.lingbotworld.model.LingBotWorldTransformerBlock.forward`

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], original_seq_len: int, c2ws_plucker_emb: torch.Tensor | None = None) -> torch.Tensor

Source

def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    temb: torch.Tensor,
    freqs_cis: tuple[torch.Tensor, torch.Tensor],
    original_seq_len: int,
    c2ws_plucker_emb: torch.Tensor | None = None,
) -> torch.Tensor:
    if hidden_states.dim() == 4:
        hidden_states = hidden_states.squeeze(1)
    bs, seq_length, _ = hidden_states.shape
    orig_dtype = hidden_states.dtype
    # assert orig_dtype != torch.float32

    if temb.dim() == 4:
        # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
            self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2)
        # batch_size, seq_len, 1, inner_dim
        shift_msa = shift_msa.squeeze(2)
        scale_msa = scale_msa.squeeze(2)
        gate_msa = gate_msa.squeeze(2)
        c_shift_msa = c_shift_msa.squeeze(2)
        c_scale_msa = c_scale_msa.squeeze(2)
        c_gate_msa = c_gate_msa.squeeze(2)
    else:
        # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B)
        e = self.scale_shift_table + temb.float()
        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1)
    assert shift_msa.dtype == torch.float32

    # 1. Self-attention
    norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).to(orig_dtype)
    query, _ = self.to_q(norm_hidden_states)
    key, _ = self.to_k(norm_hidden_states)
    value, _ = self.to_v(norm_hidden_states)

    if self.norm_q is not None:
        query = self.norm_q(query)
    if self.norm_k is not None:
        key = self.norm_k(key)

    query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
    key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
    value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))

    attn_output, _ = self.attn1(
        query,
        key,
        value,
        original_seq_len,
        freqs_cis=freqs_cis,
    )
    attn_output = attn_output.flatten(2)
    attn_output, _ = self.to_out(attn_output)
    attn_output = attn_output.squeeze(1)

    null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
    norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa,
                                                                     null_shift, null_scale)
    norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype)
    # Inject camera condition
    # must be applied after the self-attention residual update.
    hidden_states = self.cam_conditioner(hidden_states, c2ws_plucker_emb)
    norm_hidden_states = self.self_attn_residual_norm.norm(hidden_states)
    norm_hidden_states = norm_hidden_states.to(orig_dtype)

    # 2. Cross-attention
    attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None)
    norm_hidden_states, hidden_states = self.cross_attn_residual_norm(hidden_states, attn_output, 1, c_shift_msa,
                                                                      c_scale_msa)
    norm_hidden_states, hidden_states = norm_hidden_states.to(orig_dtype), hidden_states.to(orig_dtype)

    # 3. Feed-forward
    ff_output = self.ffn(norm_hidden_states)
    hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
    hidden_states = hidden_states.to(orig_dtype)

    return hidden_states

`fastvideo.models.dits.lingbotworld.model.LingBotWorldTransformer3DModel`

LingBotWorldTransformer3DModel(config: LingBotWorldVideoConfig, hf_config: dict[str, Any])

Source

class LingBotWorldTransformer3DModel(BaseDiT):
    _fsdp_shard_conditions = LingBotWorldVideoConfig()._fsdp_shard_conditions
    _compile_conditions = LingBotWorldVideoConfig()._compile_conditions
    _supported_attention_backends = LingBotWorldVideoConfig()._supported_attention_backends
    param_names_mapping = LingBotWorldVideoConfig().param_names_mapping
    reverse_param_names_mapping = LingBotWorldVideoConfig().reverse_param_names_mapping
    lora_param_names_mapping = LingBotWorldVideoConfig().lora_param_names_mapping

    def __init__(self, config: LingBotWorldVideoConfig, hf_config: dict[str, Any]) -> None:
        super().__init__(config=config, hf_config=hf_config)

        inner_dim = config.num_attention_heads * config.attention_head_dim
        self.hidden_size = config.hidden_size
        self.num_attention_heads = config.num_attention_heads
        self.in_channels = config.in_channels
        self.out_channels = config.out_channels
        self.num_channels_latents = config.num_channels_latents
        self.patch_size = config.patch_size
        self.text_len = config.text_len

        assert config.num_attention_heads % get_sp_world_size(
        ) == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})"

        # 1. Patch & position embedding
        self.patch_embedding = PatchEmbed(in_chans=config.in_channels,
                                          embed_dim=inner_dim,
                                          patch_size=config.patch_size,
                                          flatten=False)
        self.patch_embedding_wancamctrl = WanCamControlPatchEmbedding(in_chans=6 * 64,
                                                                      embed_dim=inner_dim,
                                                                      patch_size=config.patch_size)
        self.c2ws_mlp = MLP(inner_dim, inner_dim, inner_dim, bias=True, act_type="silu")

        # 2. Condition embeddings
        self.condition_embedder = WanTimeTextImageEmbedding(
            dim=inner_dim,
            time_freq_dim=config.freq_dim,
            text_embed_dim=config.text_dim,
            image_embed_dim=config.image_dim,
        )

        # 3. Transformer blocks
        transformer_block = LingBotWorldTransformerBlock
        self.blocks = nn.ModuleList([
            transformer_block(inner_dim,
                              config.ffn_dim,
                              config.num_attention_heads,
                              config.qk_norm,
                              config.cross_attn_norm,
                              config.eps,
                              config.added_kv_proj_dim,
                              self._supported_attention_backends,
                              prefix=f"{config.prefix}.blocks.{i}") for i in range(config.num_layers)
        ])

        # 4. Output norm & projection
        self.norm_out = LayerNormScaleShift(inner_dim,
                                            norm_type="layer",
                                            eps=config.eps,
                                            elementwise_affine=False,
                                            dtype=torch.float32,
                                            compute_dtype=torch.float32)
        self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size))
        self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5)

        self.gradient_checkpointing = False
        self.__post_init__()

    def forward(self,
                hidden_states: torch.Tensor,
                encoder_hidden_states: torch.Tensor | list[torch.Tensor],
                timestep: torch.LongTensor,
                encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
                | None = None,
                guidance=None,
                c2ws_plucker_emb: torch.Tensor | None = None,
                **kwargs) -> torch.Tensor:
        orig_dtype = hidden_states.dtype
        if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
            encoder_hidden_states = encoder_hidden_states[0]
        if isinstance(encoder_hidden_states_image, list):
            encoder_hidden_states_image = (encoder_hidden_states_image[0]
                                           if len(encoder_hidden_states_image) > 0 else None)

        batch_size, num_channels, num_frames, height, width = hidden_states.shape
        p_t, p_h, p_w = self.patch_size
        post_patch_num_frames = num_frames // p_t
        post_patch_height = height // p_h
        post_patch_width = width // p_w

        # Get rotary embeddings
        d = self.hidden_size // self.num_attention_heads
        rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
        freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width),
                                                    self.hidden_size,
                                                    self.num_attention_heads,
                                                    rope_dim_list,
                                                    dtype=torch.float32 if current_platform.is_mps() else torch.float64,
                                                    rope_theta=10000)
        freqs_cis = (freqs_cos.to(hidden_states.device).float(), freqs_sin.to(hidden_states.device).float())

        hidden_states = self.patch_embedding(hidden_states)
        hidden_states = hidden_states.flatten(2).transpose(1, 2)
        c2ws_hidden_states = None
        if c2ws_plucker_emb is not None:
            c2ws_plucker_emb = self.patch_embedding_wancamctrl(
                c2ws_plucker_emb.to(device=hidden_states.device, dtype=hidden_states.dtype))
            c2ws_hidden_states = self.c2ws_mlp(c2ws_plucker_emb)
            c2ws_plucker_emb = c2ws_plucker_emb + c2ws_hidden_states

        # Shard with padding support - returns (sharded_tensor, original_seq_len)
        hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)

        # Shard c2ws_plucker_emb
        if c2ws_plucker_emb is not None:
            c2ws_plucker_emb, _ = sequence_model_parallel_shard(c2ws_plucker_emb, dim=1)

        # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v)
        if timestep.dim() == 2:
            ts_seq_len = timestep.shape[1]
            timestep = timestep.flatten()  # batch_size * seq_len
        else:
            ts_seq_len = None

        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
            timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len)
        if ts_seq_len is not None:
            # batch_size, seq_len, 6, inner_dim
            timestep_proj = timestep_proj.unflatten(2, (6, -1))
        else:
            # batch_size, 6, inner_dim
            timestep_proj = timestep_proj.unflatten(1, (6, -1))

        if encoder_hidden_states_image is not None:
            if encoder_hidden_states is not None:
                encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1)
            else:
                encoder_hidden_states = encoder_hidden_states_image

        if current_platform.is_mps() or current_platform.is_npu():
            encoder_hidden_states = encoder_hidden_states.to(orig_dtype)
        else:
            encoder_hidden_states = encoder_hidden_states  # cast to orig_dtype for MPS & NPU

        assert encoder_hidden_states.dtype == orig_dtype

        # 4. Transformer blocks
        if torch.is_grad_enabled() and self.gradient_checkpointing:
            for block in self.blocks:
                hidden_states = self._gradient_checkpointing_func(block, hidden_states, encoder_hidden_states,
                                                                  timestep_proj, freqs_cis, original_seq_len,
                                                                  c2ws_plucker_emb)
        else:
            for block in self.blocks:
                hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, original_seq_len,
                                      c2ws_plucker_emb)
        # 5. Output norm, projection & unpatchify
        if temb.dim() == 3:
            # batch_size, seq_len, inner_dim (wan 2.2 ti2v)
            shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2)
            shift = shift.squeeze(2)
            scale = scale.squeeze(2)
        else:
            # batch_size, inner_dim
            shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1)

        hidden_states = self.norm_out(hidden_states, shift, scale)

        # Gather and unpad in one operation
        hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1)
        hidden_states = self.proj_out(hidden_states)

        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width,
                                              p_t, p_h, p_w, -1)
        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)

        return output

`fastvideo.models.dits.lingbotworld.model.LingBotWorldTransformer3DModel.forward`

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, guidance = None, c2ws_plucker_emb: torch.Tensor | None = None, **kwargs) -> torch.Tensor

Source

def forward(self,
            hidden_states: torch.Tensor,
            encoder_hidden_states: torch.Tensor | list[torch.Tensor],
            timestep: torch.LongTensor,
            encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
            | None = None,
            guidance=None,
            c2ws_plucker_emb: torch.Tensor | None = None,
            **kwargs) -> torch.Tensor:
    orig_dtype = hidden_states.dtype
    if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
        encoder_hidden_states = encoder_hidden_states[0]
    if isinstance(encoder_hidden_states_image, list):
        encoder_hidden_states_image = (encoder_hidden_states_image[0]
                                       if len(encoder_hidden_states_image) > 0 else None)

    batch_size, num_channels, num_frames, height, width = hidden_states.shape
    p_t, p_h, p_w = self.patch_size
    post_patch_num_frames = num_frames // p_t
    post_patch_height = height // p_h
    post_patch_width = width // p_w

    # Get rotary embeddings
    d = self.hidden_size // self.num_attention_heads
    rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
    freqs_cos, freqs_sin = get_rotary_pos_embed((post_patch_num_frames, post_patch_height, post_patch_width),
                                                self.hidden_size,
                                                self.num_attention_heads,
                                                rope_dim_list,
                                                dtype=torch.float32 if current_platform.is_mps() else torch.float64,
                                                rope_theta=10000)
    freqs_cis = (freqs_cos.to(hidden_states.device).float(), freqs_sin.to(hidden_states.device).float())

    hidden_states = self.patch_embedding(hidden_states)
    hidden_states = hidden_states.flatten(2).transpose(1, 2)
    c2ws_hidden_states = None
    if c2ws_plucker_emb is not None:
        c2ws_plucker_emb = self.patch_embedding_wancamctrl(
            c2ws_plucker_emb.to(device=hidden_states.device, dtype=hidden_states.dtype))
        c2ws_hidden_states = self.c2ws_mlp(c2ws_plucker_emb)
        c2ws_plucker_emb = c2ws_plucker_emb + c2ws_hidden_states

    # Shard with padding support - returns (sharded_tensor, original_seq_len)
    hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)

    # Shard c2ws_plucker_emb
    if c2ws_plucker_emb is not None:
        c2ws_plucker_emb, _ = sequence_model_parallel_shard(c2ws_plucker_emb, dim=1)

    # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v)
    if timestep.dim() == 2:
        ts_seq_len = timestep.shape[1]
        timestep = timestep.flatten()  # batch_size * seq_len
    else:
        ts_seq_len = None

    temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
        timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len)
    if ts_seq_len is not None:
        # batch_size, seq_len, 6, inner_dim
        timestep_proj = timestep_proj.unflatten(2, (6, -1))
    else:
        # batch_size, 6, inner_dim
        timestep_proj = timestep_proj.unflatten(1, (6, -1))

    if encoder_hidden_states_image is not None:
        if encoder_hidden_states is not None:
            encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1)
        else:
            encoder_hidden_states = encoder_hidden_states_image

    if current_platform.is_mps() or current_platform.is_npu():
        encoder_hidden_states = encoder_hidden_states.to(orig_dtype)
    else:
        encoder_hidden_states = encoder_hidden_states  # cast to orig_dtype for MPS & NPU

    assert encoder_hidden_states.dtype == orig_dtype

    # 4. Transformer blocks
    if torch.is_grad_enabled() and self.gradient_checkpointing:
        for block in self.blocks:
            hidden_states = self._gradient_checkpointing_func(block, hidden_states, encoder_hidden_states,
                                                              timestep_proj, freqs_cis, original_seq_len,
                                                              c2ws_plucker_emb)
    else:
        for block in self.blocks:
            hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, original_seq_len,
                                  c2ws_plucker_emb)
    # 5. Output norm, projection & unpatchify
    if temb.dim() == 3:
        # batch_size, seq_len, inner_dim (wan 2.2 ti2v)
        shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2)
        shift = shift.squeeze(2)
        scale = scale.squeeze(2)
    else:
        # batch_size, inner_dim
        shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1)

    hidden_states = self.norm_out(hidden_states, shift, scale)

    # Gather and unpad in one operation
    hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1)
    hidden_states = self.proj_out(hidden_states)

    hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width,
                                          p_t, p_h, p_w, -1)
    hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
    output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)

    return output