model

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</code></pre></details>
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 = &quot;rms_norm_across_heads&quot;,
             cross_attn_norm: bool = False,
             eps: float = 1e-6,
             added_kv_proj_dim: int | None = None,
             supported_attention_backends: tuple&#91;AttentionBackendEnum, ...&#93;
             | None = None,
             prefix: str = &quot;&quot;):
    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&quot;&#123;prefix&#125;.attn1&quot;)
    self.hidden_dim = dim
    self.num_attention_heads = num_heads
    dim_head = dim // num_heads
    if qk_norm == &quot;rms_norm&quot;:
        self.norm_q = RMSNorm(dim_head, eps=eps)
        self.norm_k = RMSNorm(dim_head, eps=eps)
    elif qk_norm == &quot;rms_norm_across_heads&quot;:
        # 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&quot;QK Norm type &#x27;&#123;qk_norm&#125;&#x27; not supported&quot;)
    assert cross_attn_norm is True
    self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(dim,
                                                                    norm_type=&quot;layer&quot;,
                                                                    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=&quot;layer&quot;,
                                                                     eps=eps,
                                                                     elementwise_affine=False,
                                                                     dtype=torch.float32,
                                                                     compute_dtype=torch.float32)

    # 3. Feed-forward
    self.ffn = MLP(dim, ffn_dim, act_type=&quot;gelu_pytorch_tanh&quot;)
    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&#91;torch.Tensor, torch.Tensor&#93;,
    original_seq_len: int,
    c2ws_plucker_emb: torch.Tensor | None = None,
) -&gt; 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(&#91;0&#93;, 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</code></pre></details>
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(&#91;0&#93;, 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</code></pre></details>
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&#91;str, Any&#93;) -&gt; 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&quot;The number of attention heads (&#123;config.num_attention_heads&#125;) must be divisible by the sequence parallel size (&#123;get_sp_world_size()&#125;)&quot;

    # 1. Patch &amp; 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=&quot;silu&quot;)

    # 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(&#91;
        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&quot;&#123;config.prefix&#125;.blocks.&#123;i&#125;&quot;) for i in range(config.num_layers)
    &#93;)

    # 4. Output norm &amp; projection
    self.norm_out = LayerNormScaleShift(inner_dim,
                                        norm_type=&quot;layer&quot;,
                                        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&#91;torch.Tensor&#93;,
            timestep: torch.LongTensor,
            encoder_hidden_states_image: torch.Tensor | list&#91;torch.Tensor&#93;
            | None = None,
            guidance=None,
            c2ws_plucker_emb: torch.Tensor | None = None,
            **kwargs) -&gt; 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&#91;0&#93;
    if isinstance(encoder_hidden_states_image, list):
        encoder_hidden_states_image = (encoder_hidden_states_image&#91;0&#93;
                                       if len(encoder_hidden_states_image) &gt; 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 = &#91;d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)&#93;
    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&#91;1&#93;
        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(&#91;encoder_hidden_states_image, encoder_hidden_states&#93;, 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 &amp; 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 &amp; 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</code></pre></details>
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 = &#91;d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)&#93;
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&#91;1&#93;
    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(&#91;encoder_hidden_states_image, encoder_hidden_states&#93;, 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 &amp; 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 &amp; 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</code></pre></details>