videox_fun/models/creator_gating.py

"""Configurable gated cross-attention joint audio-video model.

This experimental variant follows the AV cross-attention design but allows
audio-to-video (A2V) and video-to-audio (V2A) cross attention to be enabled
independently for each layer. Cross-modal attention can optionally apply a
fixed per-layer alpha times a per-head sigmoid gate to the attention context
before the output projection.
"""

import logging
import math
import os
from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union

import torch
import torch.nn as nn
from einops import rearrange

from .attention_utils import attention
from .creator_audio import CreatorAudioModel
from .wan_transformer3d_prope import (
    Wan2_2Transformer3DModel,
    WanRMSNorm,
    WanTransformer3DModel,
    rope_apply_qk,
)
from .creator.creator_video_dit import sinusoidal_embedding_1d
from .creator.creator_video_dit import rope_apply_head_dim
from ..dist.sequence_parallel import all_gather_sequence, ulysses_attention


LayerSelection = Optional[Union[bool, str, Sequence[bool], Sequence[int], torch.Tensor]]
LayerAlphas = Optional[
    Union[float, Sequence[float], Mapping[Union[int, str], float], torch.Tensor]
]


@torch.amp.autocast("cuda", enabled=False)
def temporal_rope_1d(
    x: torch.Tensor,
    temporal_positions: torch.Tensor,
    inv_freqs_1d: torch.Tensor,
) -> torch.Tensor:
    """Apply 1D temporal RoPE to [B, L, num_heads, head_dim] tensors."""
    dtype = x.dtype
    batch_size, seq_len, num_heads, head_dim = x.shape
    half_dim = head_dim // 2

    freqs = torch.einsum(
        "bl,d->bld",
        temporal_positions.to(torch.float64),
        inv_freqs_1d.to(x.device, torch.float64),
    )
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)

    x_complex = torch.view_as_complex(
        x.to(torch.float64).reshape(batch_size, seq_len, num_heads, half_dim, 2)
    )
    x_out = torch.view_as_real(x_complex * freqs_cis.unsqueeze(2)).flatten(3)
    return x_out.to(dtype)


def compute_video_temporal_positions(
    grid_sizes: torch.Tensor,
    seq_len: int,
    device: torch.device,
    audio_fps: float = 48000.0 / 960.0,
    video_fps: float = 16.0,
    vae_temporal_stride: int = 4,
) -> torch.Tensor:
    """Compute video token temporal positions in audio-token time units."""
    batch_size = grid_sizes.size(0)
    positions = torch.zeros(batch_size, seq_len, device=device, dtype=torch.float64)
    video_latent_fps = video_fps / vae_temporal_stride
    scale = audio_fps / video_latent_fps

    for sample_idx, (num_frames, height, width) in enumerate(grid_sizes.tolist()):
        spatial_size = int(height * width)
        num_tokens = int(num_frames * spatial_size)
        frame_indices = torch.arange(num_tokens, device=device, dtype=torch.float64) // spatial_size
        positions[sample_idx, :num_tokens] = frame_indices * scale
    return positions


def compute_audio_temporal_positions(
    seq_lens: torch.Tensor,
    seq_len: int,
    device: torch.device,
) -> torch.Tensor:
    """Compute sequential temporal positions for audio tokens."""
    batch_size = seq_lens.size(0)
    positions = torch.zeros(batch_size, seq_len, device=device, dtype=torch.float64)
    base_positions = torch.arange(seq_len, device=device, dtype=torch.float64)
    for sample_idx in range(batch_size):
        valid_len = int(seq_lens[sample_idx].item())
        positions[sample_idx, :valid_len] = base_positions[:valid_len]
    return positions


def _apply_video_rope_local(
    x: torch.Tensor,
    grid_sizes: torch.Tensor,
    freqs: torch.Tensor,
    sp_rank: int,
    sp_world_size: int,
) -> torch.Tensor:
    """Apply the Wan 3D RoPE slice belonging to this sequence-parallel rank."""
    if sp_world_size <= 1:
        return rope_apply_qk(x, x, grid_sizes, freqs)[0]

    local_len, num_heads, complex_dim = x.size(1), x.size(2), x.size(3) // 2
    freq_parts = freqs.split(
        [complex_dim - 2 * (complex_dim // 3), complex_dim // 3, complex_dim // 3],
        dim=1,
    )
    output = []
    for sample_idx, (frames, height, width) in enumerate(grid_sizes.tolist()):
        full_len = int(frames * height * width)
        sample = x[sample_idx, :local_len].to(torch.float64)
        sample_complex = torch.view_as_complex(
            sample.reshape(local_len, num_heads, -1, 2)
        )
        full_freqs = torch.cat(
            [
                freq_parts[0][:frames].view(frames, 1, 1, -1).expand(frames, height, width, -1),
                freq_parts[1][:height].view(1, height, 1, -1).expand(frames, height, width, -1),
                freq_parts[2][:width].view(1, 1, width, -1).expand(frames, height, width, -1),
            ],
            dim=-1,
        ).reshape(full_len, 1, -1)
        if full_freqs.size(0) < local_len * sp_world_size:
            full_freqs = torch.cat(
                [
                    full_freqs,
                    torch.ones(
                        local_len * sp_world_size - full_freqs.size(0),
                        full_freqs.size(1),
                        full_freqs.size(2),
                        dtype=full_freqs.dtype,
                        device=full_freqs.device,
                    ),
                ],
                dim=0,
            )
        start = sp_rank * local_len
        local_freqs = full_freqs[start : start + local_len]
        rotated = torch.view_as_real(sample_complex * local_freqs).flatten(2)
        if x.size(1) > local_len:
            rotated = torch.cat([rotated, x[sample_idx, local_len:]], dim=0)
        output.append(rotated)
    return torch.stack(output).to(x.dtype)


def _to_list(value) -> List[torch.Tensor]:
    if isinstance(value, torch.Tensor):
        return [sample for sample in value]
    return list(value)


def _build_time_embeddings(
    time_embedding: nn.Module,
    time_projection: nn.Module,
    freq_dim: int,
    dim: int,
    timesteps: torch.Tensor,
    seq_len: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Build time embeddings and projected modulations."""

    with torch.amp.autocast("cuda", dtype=torch.float32):
        if timesteps.dim() != 1:
            if timesteps.size(1) < seq_len:
                pad_size = seq_len - timesteps.size(1)
                timesteps = torch.cat(
                    [timesteps, timesteps[:, -1:].repeat(1, pad_size)], dim=1
                )
            batch_size = timesteps.size(0)
            embedding = time_embedding(
                sinusoidal_embedding_1d(freq_dim, timesteps.flatten())
                .unflatten(0, (batch_size, seq_len))
                .float()
            )
            modulation = time_projection(embedding).unflatten(2, (6, dim))
        else:
            embedding = time_embedding(sinusoidal_embedding_1d(freq_dim, timesteps).float())
            modulation = time_projection(embedding).unflatten(1, (6, dim))
    return embedding, modulation


def _embed_context(text_embedding: nn.Module, context, text_len: int) -> torch.Tensor:
    """Embed and right-pad text context to the model fixed text length."""
    if isinstance(context, torch.Tensor):
        samples = [sample for sample in context]
    else:
        samples = list(context)
    return text_embedding(
        torch.stack([
            torch.cat([sample, sample.new_zeros(text_len - sample.size(0), sample.size(1))])
            for sample in samples
        ])
    )


def _logit_from_gate_value(gate_init_value: Optional[float]) -> float:
    """Convert an initial gate value in [0, 1] to a sigmoid bias."""
    if gate_init_value is None:
        return 0.0
    gate_value = min(max(float(gate_init_value), 1e-6), 1.0 - 1e-6)
    return math.log(gate_value / (1.0 - gate_value))


def _expand_layer_selection(
    selection: LayerSelection,
    num_layers: int,
    name: str,
) -> List[bool]:
    """Expand a layer-selection config to a bool mask of length ``num_layers``.

    Accepted forms:
      - ``None`` or ``True``: enable every layer.
      - ``False``: disable every layer.
      - bool mask with length ``num_layers``.
      - 0/1 mask with length ``num_layers``.
      - list/tuple/tensor of layer indices to enable.
      - strings: ``"all"``, ``"none"``, ``"0,2,5"``.
    """
    if selection is None:
        return [False] * num_layers
    if isinstance(selection, bool):
        return [selection] * num_layers
    if isinstance(selection, torch.Tensor):
        selection = selection.cpu().tolist()
    if isinstance(selection, str):
        normalized = selection.strip().lower()
        if normalized in {"", "none", "false", "off", "0"}:
            return [False] * num_layers
        if normalized in {"all", "true", "on", "1"}:
            return [True] * num_layers
        indices = [int(part.strip()) for part in selection.split(",") if part.strip()]
        mask = [False] * num_layers
        for layer_idx in indices:
            if layer_idx < 0 or layer_idx >= num_layers:
                raise ValueError(f"{name} layer index {layer_idx} out of range [0, {num_layers})")
            mask[layer_idx] = True
        return mask

    values = list(selection)
    if not values:
        return [False] * num_layers

    if all(isinstance(value, bool) for value in values):
        if len(values) != num_layers:
            raise ValueError(f"{name} bool mask must have length {num_layers}, got {len(values)}")
        return [bool(value) for value in values]

    if all(isinstance(value, int) for value in values):
        if len(values) == num_layers and all(int(value) in {0, 1} for value in values):
            return [bool(value) for value in values]
        mask = [False] * num_layers
        for layer_idx in values:
            if layer_idx < 0 or layer_idx >= num_layers:
                raise ValueError(f"{name} layer index {layer_idx} out of range [0, {num_layers})")
            mask[int(layer_idx)] = True
        return mask

    raise TypeError(
        f"{name} must be None, bool, string, bool mask, 0/1 mask, or layer-index sequence"
    )


def _expand_layer_alphas(
    values: LayerAlphas,
    num_layers: int,
    name: str,
) -> List[float]:
    """Expand fixed per-layer gate multipliers, defaulting each layer to 1.0.

    Accepted forms are a scalar shared by all layers, a full sequence with
    ``num_layers`` entries, or a mapping of layer index to alpha. Unspecified
    mapping entries retain the default value 1.0.
    """
    if values is None:
        return [1.0] * num_layers
    if isinstance(values, torch.Tensor):
        values = values.cpu().tolist()

    def validate(value, layer_label: str) -> float:
        alpha = float(value)
        if not math.isfinite(alpha) or alpha < 0.0:
            raise ValueError(f"{name} {layer_label} must be finite and non-negative, got {value}")
        return alpha

    if isinstance(values, (int, float)) and not isinstance(values, bool):
        return [validate(values, "scalar")] * num_layers

    if isinstance(values, Mapping):
        alphas = [1.0] * num_layers
        for raw_layer_idx, value in values.items():
            try:
                layer_idx = int(raw_layer_idx)
            except (TypeError, ValueError) as exc:
                raise ValueError(
                    f"{name} mapping key must be a layer index, got {raw_layer_idx!r}"
                ) from exc
            if layer_idx < 0 or layer_idx >= num_layers:
                raise ValueError(f"{name} layer index {layer_idx} out of range [0, {num_layers})")
            alphas[layer_idx] = validate(value, f"layer {layer_idx}")
        return alphas

    if isinstance(values, Sequence) and not isinstance(values, (str, bytes)):
        if len(values) != num_layers:
            raise ValueError(f"{name} sequence must have length {num_layers}, got {len(values)}")
        return [validate(value, f"layer {layer_idx}") for layer_idx, value in enumerate(values)]

    raise TypeError(f"{name} must be None, a scalar, a mapping, or a full layer sequence")


class GatedCrossModalAttention(nn.Module):
    """Cross-modal attention with optional alpha-scaled sigmoid context gating.

    For A2V, ``x`` is video hidden states and ``y`` is audio hidden states.
    For V2A, ``x`` is audio hidden states and ``y`` is video hidden states.
    When ``use_gating=False``, this module follows the reference cross-attn
    behavior: query from raw ``x`` and key/value from normalized ``y``.
    """

    def __init__(
        self,
        q_dim: int,
        kv_dim: int,
        num_heads: int,
        eps: float = 1e-6,
        zero_init_output: bool = False,
        use_gating: bool = True,
        zero_init_gating: bool = False,
        gate_init_value: Optional[float] = None,
        gate_alpha: float = 1.0,
    ):
        super().__init__()
        assert q_dim % num_heads == 0
        self.q_dim = q_dim
        self.kv_dim = kv_dim
        self.num_heads = num_heads
        self.head_dim = q_dim // num_heads
        self.use_gating = bool(use_gating)
        self.gate_alpha = _expand_layer_alphas(gate_alpha, 1, "gate_alpha")[0]

        self.norm = nn.LayerNorm(kv_dim, eps=eps)
        # if self.use_gating:
        self.norm_x = nn.LayerNorm(q_dim, eps=eps)

        self.q = nn.Linear(q_dim, q_dim)
        self.k = nn.Linear(kv_dim, q_dim)
        self.v = nn.Linear(kv_dim, q_dim)
        self.o = nn.Linear(q_dim, q_dim)

        self.norm_q = WanRMSNorm(q_dim, eps=eps)
        self.norm_k = WanRMSNorm(q_dim, eps=eps)

        if self.use_gating:
            self.gate_hidden = nn.Linear(q_dim, num_heads, bias=False)
            self.gate_context_norm = nn.LayerNorm(self.head_dim, eps=eps)
            self.gate_context = nn.Linear(self.head_dim, 1, bias=False)
            self.gate_bias = nn.Parameter(torch.empty(num_heads))

        self._init_weights(
            zero_init_output=zero_init_output,
            zero_init_gating=zero_init_gating,
            gate_init_value=gate_init_value,
        )

    def _init_weights(
        self,
        zero_init_output: bool,
        zero_init_gating: bool,
        gate_init_value: Optional[float],
    ):
        for module in [self.q, self.k, self.v, self.o]:
            nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)

        if zero_init_output:
            nn.init.zeros_(self.o.weight)
            nn.init.zeros_(self.o.bias)

        if self.use_gating:
            if zero_init_gating:
                nn.init.zeros_(self.gate_hidden.weight)
                nn.init.zeros_(self.gate_context.weight)
            else:
                nn.init.xavier_uniform_(self.gate_hidden.weight)
                nn.init.xavier_uniform_(self.gate_context.weight)
            nn.init.constant_(self.gate_bias, _logit_from_gate_value(gate_init_value))

    def forward(
        self,
        x: torch.Tensor,
        y: torch.Tensor,
        y_lens: Optional[torch.Tensor] = None,
        dtype: torch.dtype = torch.bfloat16,
        q_temporal_pos: Optional[torch.Tensor] = None,
        k_temporal_pos: Optional[torch.Tensor] = None,
        temporal_rope_inv_freq: Optional[torch.Tensor] = None,
        sequence_parallel: bool = False,
        sp_group=None,
    ) -> torch.Tensor:
        """
        Args:
            x: primary hidden states [B, Lq, q_dim]
            y: conditioning hidden states [B, Lk, kv_dim]
            y_lens: valid lengths of y per sample [B]
            q_temporal_pos: [B, Lq] temporal positions for query
            k_temporal_pos: [B, Lk] temporal positions for key
            temporal_rope_inv_freq: [head_dim // 2] inv frequencies for temporal RoPE
        Returns:
            Cross-attention output [B, Lq, q_dim].
        """
        batch_size = x.size(0)
        num_heads = self.num_heads
        head_dim = self.head_dim

        # In SP mode x is a local query chunk while y is a local conditioning
        # chunk. Cross-modal attention needs the complete conditioning sequence.
        if sequence_parallel:
            y = all_gather_sequence(y, group=sp_group)
            if k_temporal_pos is not None:
                k_temporal_pos = all_gather_sequence(
                    k_temporal_pos.unsqueeze(-1), group=sp_group
                ).squeeze(-1)

        x_for_q = self.norm_x(x)
        y_norm = self.norm(y)

        query = self.norm_q(self.q(x_for_q.to(dtype))).view(batch_size, -1, num_heads, head_dim)
        key = self.norm_k(self.k(y_norm.to(dtype))).view(batch_size, -1, num_heads, head_dim)
        value = self.v(y_norm.to(dtype)).view(batch_size, -1, num_heads, head_dim)

        if q_temporal_pos is not None and k_temporal_pos is not None and temporal_rope_inv_freq is not None:
            query = temporal_rope_1d(query, q_temporal_pos, temporal_rope_inv_freq)
            key = temporal_rope_1d(key, k_temporal_pos, temporal_rope_inv_freq)

        context = attention(query.to(dtype), key.to(dtype), value.to(dtype), k_lens=y_lens)
        context = context.to(dtype)

        if self.use_gating:
            hidden_gate = self.gate_hidden(x_for_q.to(dtype)).view(batch_size, -1, num_heads, 1)
            context_gate = self.gate_context(
                self.gate_context_norm(context)
            )
            gate = self.gate_alpha * torch.sigmoid(
                hidden_gate + context_gate + self.gate_bias.view(1, 1, num_heads, 1)
            )
            context = gate * context
        return self.o(context.flatten(2))


class GatedJointBlock(nn.Module):
    """One joint block with independently configurable A2V and V2A attention."""

    def __init__(
        self,
        video_block: nn.Module,
        audio_block: nn.Module,
        video_dim: int,
        audio_dim: int,
        video_num_heads: int,
        audio_num_heads: int,
        enable_a2v_cross_attn: bool = True,
        enable_v2a_cross_attn: bool = True,
        zero_init_output: bool = False,
        zero_init_video_cross_attn: bool | None = None,
        zero_init_audio_cross_attn: bool | None = None,
        use_a2v_gating: bool = True,
        use_v2a_gating: bool = True,
        zero_init_a2v_gating: bool = False,
        zero_init_v2a_gating: bool = False,
        a2v_gate_init_value: Optional[float] = None,
        v2a_gate_init_value: Optional[float] = None,
        a2v_gate_alpha: float = 1.0,
        v2a_gate_alpha: float = 1.0,
    ):
        super().__init__()
        self.video_block = video_block
        self.audio_block = audio_block
        self.enable_a2v_cross_attn = bool(enable_a2v_cross_attn)
        self.enable_v2a_cross_attn = bool(enable_v2a_cross_attn)

        zero_init_video = zero_init_video_cross_attn if zero_init_video_cross_attn is not None else zero_init_output
        zero_init_audio = zero_init_audio_cross_attn if zero_init_audio_cross_attn is not None else zero_init_output

        if self.enable_a2v_cross_attn:
            self.video_cross_attn_audio = GatedCrossModalAttention(
                q_dim=video_dim,
                kv_dim=audio_dim,
                num_heads=video_num_heads,
                zero_init_output=zero_init_video,
                use_gating=use_a2v_gating,
                zero_init_gating=zero_init_a2v_gating,
                gate_init_value=a2v_gate_init_value,
                gate_alpha=a2v_gate_alpha,
            )
        else:
            self.video_cross_attn_audio = None

        if self.enable_v2a_cross_attn:
            self.audio_cross_attn_video = GatedCrossModalAttention(
                q_dim=audio_dim,
                kv_dim=video_dim,
                num_heads=audio_num_heads,
                zero_init_output=zero_init_audio,
                use_gating=use_v2a_gating,
                zero_init_gating=zero_init_v2a_gating,
                gate_init_value=v2a_gate_init_value,
                gate_alpha=v2a_gate_alpha,
            )
        else:
            self.audio_cross_attn_video = None

    def forward(
        self,
        video_x: torch.Tensor,
        audio_x: torch.Tensor,
        video_kwargs: Dict[str, Any],
        audio_kwargs: Dict[str, Any],
        dtype: torch.dtype = torch.bfloat16,
        enable_a2v: Optional[Union[bool, torch.Tensor]] = None,
        enable_v2a: Optional[Union[bool, torch.Tensor]] = None,
        sequence_parallel: bool = False,
        sp_group=None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        video_block = self.video_block
        audio_block = self.audio_block

        video_x = self._video_selfattn_and_text(
            video_x, video_block, video_kwargs, dtype,
            sequence_parallel=sequence_parallel,
            sp_group=sp_group,
        )
        audio_x = self._audio_selfattn_and_text(
            audio_x, audio_block, audio_kwargs, dtype,
            sequence_parallel=sequence_parallel,
            sp_group=sp_group,
        )

        temporal_rope_inv_freq = video_kwargs.get("temporal_rope_inv_freq")
        video_temporal_pos = video_kwargs.get("temporal_positions")
        audio_temporal_pos = audio_kwargs.get("temporal_positions")

        # Runtime switches can be batch-wide booleans or per-branch CFG masks.
        a2v_is_tensor = isinstance(enable_a2v, torch.Tensor)
        v2a_is_tensor = isinstance(enable_v2a, torch.Tensor)
        run_a2v = self.video_cross_attn_audio is not None and (
            a2v_is_tensor or enable_a2v is None or enable_a2v
        )
        run_v2a = self.audio_cross_attn_video is not None and (
            v2a_is_tensor or enable_v2a is None or enable_v2a
        )

        # Cache pre-cross-attention states so A2V and V2A are updated jointly:
        # both directions must attend to the *same* pre-update snapshot, otherwise
        # V2A would condition on the already-A2V-updated video (serial dependency).
        video_x_pre = video_x
        audio_x_pre = audio_x

        if run_a2v:
            a2v_result = self.video_cross_attn_audio(
                x=video_x_pre,
                y=audio_x_pre,
                y_lens=audio_kwargs.get("seq_lens"),
                dtype=dtype,
                q_temporal_pos=video_temporal_pos,
                k_temporal_pos=audio_temporal_pos,
                temporal_rope_inv_freq=temporal_rope_inv_freq,
                sequence_parallel=sequence_parallel,
                sp_group=sp_group,
            )
            a2v_out = a2v_result
            if a2v_is_tensor:
                a2v_mask = enable_a2v.view(-1, *([1] * (a2v_out.dim() - 1))).to(
                    device=a2v_out.device, dtype=a2v_out.dtype
                )
                a2v_out = a2v_out * a2v_mask
            video_x = video_x + a2v_out

        if run_v2a:
            v2a_result = self.audio_cross_attn_video(
                x=audio_x_pre,
                y=video_x_pre,
                y_lens=video_kwargs.get("seq_lens"),
                dtype=dtype,
                q_temporal_pos=audio_temporal_pos,
                k_temporal_pos=video_temporal_pos,
                temporal_rope_inv_freq=temporal_rope_inv_freq,
                sequence_parallel=sequence_parallel,
                sp_group=sp_group,
            )
            v2a_out = v2a_result
            if v2a_is_tensor:
                v2a_mask = enable_v2a.view(-1, *([1] * (v2a_out.dim() - 1))).to(
                    device=v2a_out.device, dtype=v2a_out.dtype
                )
                v2a_out = v2a_out * v2a_mask
            audio_x = audio_x + v2a_out

        video_x = self._video_ffn(video_x, video_block, video_kwargs, dtype)
        audio_x = self._audio_ffn(audio_x, audio_block, audio_kwargs, dtype)
        return video_x, audio_x

    def _video_selfattn_and_text(
        self,
        x: torch.Tensor,
        block,
        kwargs: Dict[str, Any],
        dtype: torch.dtype,
        sequence_parallel: bool = False,
        sp_group=None,
    ) -> torch.Tensor:
        e0 = kwargs["e0"]
        seq_lens = kwargs["seq_lens"]
        grid_sizes = kwargs["grid_sizes"]
        freqs = kwargs["freqs"]
        context = kwargs["context"]
        context_lens = kwargs.get("context_lens")

        if e0.dim() > 3:
            modulation = (block.modulation.unsqueeze(0) + e0).chunk(6, dim=2)
            modulation = [part.squeeze(2) for part in modulation]
        else:
            modulation = (block.modulation + e0).chunk(6, dim=1)

        kwargs["_video_e"] = modulation

        temp_x = block.norm1(x) * (1 + modulation[1]) + modulation[0]
        temp_x = temp_x.to(dtype)

        self_attn = block.self_attn
        batch_size, seq_len = temp_x.shape[:2]
        num_heads, head_dim = self_attn.num_heads, self_attn.head_dim
        query = self_attn.norm_q(self_attn.q(temp_x)).view(batch_size, seq_len, num_heads, head_dim)
        key = self_attn.norm_k(self_attn.k(temp_x)).view(batch_size, seq_len, num_heads, head_dim)
        value = self_attn.v(temp_x).view(batch_size, seq_len, num_heads, head_dim)
        if sequence_parallel:
            sp_rank = int(kwargs["sp_rank"])
            sp_world_size = int(kwargs["sp_world_size"])
            query = _apply_video_rope_local(query, grid_sizes, freqs, sp_rank, sp_world_size)
            key = _apply_video_rope_local(key, grid_sizes, freqs, sp_rank, sp_world_size)
            attn_output = ulysses_attention(
                query.to(dtype),
                key.to(dtype),
                value.to(dtype),
                attention,
                k_lens=seq_lens,
                window_size=getattr(self_attn, "window_size", (-1, -1)),
                group=sp_group,
            )
        else:
            query, key = rope_apply_qk(query, key, grid_sizes, freqs)
            attn_output = attention(
                query.to(dtype),
                key.to(dtype),
                v=value.to(dtype),
                k_lens=seq_lens,
                window_size=getattr(self_attn, "window_size", (-1, -1)),
            )
        attn_output = attn_output.to(dtype).flatten(2)
        attn_output = self_attn.o(attn_output)

        x = x + attn_output * modulation[2]
        x = x + block.cross_attn(block.norm3(x), context, context_lens, dtype)
        return x

    def _audio_selfattn_and_text(
        self,
        x: torch.Tensor,
        block,
        kwargs: Dict[str, Any],
        dtype: torch.dtype,
        sequence_parallel: bool = False,
        sp_group=None,
    ) -> torch.Tensor:
        time_mod = kwargs["e0"]
        freqs = kwargs["freqs"]
        context = kwargs["context"]

        has_seq_mod = len(time_mod.shape) == 4
        chunk_dim = 2 if has_seq_mod else 1
        modulation = (
            block.modulation.to(dtype=time_mod.dtype, device=time_mod.device) + time_mod
        ).chunk(6, dim=chunk_dim)
        if has_seq_mod:
            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = [
                part.squeeze(2) for part in modulation
            ]
        else:
            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = modulation

        kwargs["_audio_mod"] = (shift_mlp, scale_mlp, gate_mlp)

        # norm1 (LayerNorm) upcasts to fp32 and the modulation terms are fp32,
        # so input_x is fp32 here. block.self_attn calls flash-attn, which only
        # supports fp16/bf16 -> cast down first (mirrors the video self-attn path).
        input_x = (block.norm1(x) * (1 + scale_msa) + shift_msa).to(dtype)
        if sequence_parallel:
            self_attn = block.self_attn
            batch_size, seq_len = input_x.shape[:2]
            num_heads, head_dim = self_attn.num_heads, self_attn.head_dim
            query = self_attn.norm_q(self_attn.q(input_x)).view(batch_size, seq_len, num_heads, head_dim)
            key = self_attn.norm_k(self_attn.k(input_x)).view(batch_size, seq_len, num_heads, head_dim)
            value = self_attn.v(input_x).view(batch_size, seq_len, num_heads, head_dim)
            query = rope_apply_head_dim(
                query.flatten(2), freqs, head_dim
            ).view(batch_size, seq_len, num_heads, head_dim)
            key = rope_apply_head_dim(
                key.flatten(2), freqs, head_dim
            ).view(batch_size, seq_len, num_heads, head_dim)
            self_attn_output = ulysses_attention(
                query,
                key,
                value,
                attention,
                k_lens=kwargs.get("seq_lens"),
                window_size=getattr(self_attn, "window_size", (-1, -1)),
                group=sp_group,
            ).flatten(2)
            self_attn_output = self_attn.o(self_attn_output)
        else:
            self_attn_output = block.self_attn(input_x, freqs, seq_lens=kwargs.get("seq_lens"))
        x = block.gate(x, gate_msa, self_attn_output)
        # norm3 (LayerNorm) upcasts to fp32 and context (text embeds) may be fp32;
        # audio cross_attn calls flash-attn (fp16/bf16 only) with no internal cast,
        # so cast both q-source and kv-source down first.
        x = x + block.cross_attn(block.norm3(x).to(dtype), context.to(dtype))
        return x

    def _video_ffn(
        self, x: torch.Tensor, block, kwargs: Dict[str, Any], dtype: torch.dtype
    ) -> torch.Tensor:
        modulation = kwargs["_video_e"]
        temp_x = block.norm2(x) * (1 + modulation[4]) + modulation[3]
        temp_x = temp_x.to(dtype)
        return x + block.ffn(temp_x) * modulation[5]

    def _audio_ffn(
        self, x: torch.Tensor, block, kwargs: Dict[str, Any], dtype: torch.dtype
    ) -> torch.Tensor:
        shift_mlp, scale_mlp, gate_mlp = kwargs["_audio_mod"]
        input_x = block.norm2(x) * (1 + scale_mlp) + shift_mlp
        return block.gate(x, gate_mlp, block.ffn(input_x))


class WanCreatorGatingAVModel(nn.Module):
    """Wan+Creator audio-video model with configurable gated cross attention.

    A2V means the video branch attends to audio and updates video tokens.
    V2A means the audio branch attends to video and updates audio tokens.
    """

    def __init__(
        self,
        video_model: WanTransformer3DModel,
        audio_model: CreatorAudioModel,
        use_temporal_rope: bool = True,
        audio_fps: float = 48000 / 960,
        vae_temporal_stride: int = 4,
        zero_init_cross_attn: bool = False,
        zero_init_video_cross_attn: bool | None = None,
        zero_init_audio_cross_attn: bool | None = None,
        a2v_cross_attn_layers: LayerSelection = None,
        v2a_cross_attn_layers: LayerSelection = None,
        use_gating: bool = True,
        use_a2v_gating: bool | None = None,
        use_v2a_gating: bool | None = None,
        zero_init_gating: bool = False,
        zero_init_a2v_gating: bool | None = None,
        zero_init_v2a_gating: bool | None = None,
        gate_init_value: Optional[float] = None,
        a2v_gate_init_value: Optional[float] = None,
        v2a_gate_init_value: Optional[float] = None,
        a2v_gate_alphas: LayerAlphas = None,
        v2a_gate_alphas: LayerAlphas = None,
    ):
        nn.Module.__init__(self)
        self.video_model = video_model
        self.audio_model = audio_model

        self.video_dim = int(video_model.dim)
        self.audio_dim = int(audio_model.dim)
        self.video_num_heads = int(video_model.num_heads)
        self.audio_num_heads = int(audio_model.num_heads)
        self.num_layers = int(video_model.num_layers)
        self.video_patch_size = tuple(int(value) for value in video_model.patch_size)
        self.audio_patch_size = tuple(int(value) for value in audio_model.patch_size)

        assert int(audio_model.num_layers) == self.num_layers, (
            f"Video ({self.num_layers}) and audio ({audio_model.num_layers}) must have same number of layers"
        )

        a2v_enabled = _expand_layer_selection(
            a2v_cross_attn_layers, self.num_layers, "a2v_cross_attn_layers"
        )
        v2a_enabled = _expand_layer_selection(
            v2a_cross_attn_layers, self.num_layers, "v2a_cross_attn_layers"
        )
        self.a2v_cross_attn_layers = a2v_enabled
        self.v2a_cross_attn_layers = v2a_enabled
        resolved_a2v_gate_alphas = _expand_layer_alphas(
            a2v_gate_alphas, self.num_layers, "a2v_gate_alphas"
        )
        resolved_v2a_gate_alphas = _expand_layer_alphas(
            v2a_gate_alphas, self.num_layers, "v2a_gate_alphas"
        )
        self.a2v_gate_alphas = resolved_a2v_gate_alphas
        self.v2a_gate_alphas = resolved_v2a_gate_alphas

        resolved_use_a2v_gating = use_gating if use_a2v_gating is None else use_a2v_gating
        resolved_use_v2a_gating = use_gating if use_v2a_gating is None else use_v2a_gating
        resolved_zero_init_a2v_gating = (
            zero_init_gating if zero_init_a2v_gating is None else zero_init_a2v_gating
        )
        resolved_zero_init_v2a_gating = (
            zero_init_gating if zero_init_v2a_gating is None else zero_init_v2a_gating
        )
        resolved_a2v_gate_init_value = (
            gate_init_value if a2v_gate_init_value is None else a2v_gate_init_value
        )
        resolved_v2a_gate_init_value = (
            gate_init_value if v2a_gate_init_value is None else v2a_gate_init_value
        )
        video_blocks = list(video_model.blocks)
        audio_blocks = list(audio_model.blocks)
        self.joint_blocks = nn.ModuleList([
            GatedJointBlock(
                video_block=video_block,
                audio_block=audio_block,
                video_dim=self.video_dim,
                audio_dim=self.audio_dim,
                video_num_heads=self.video_num_heads,
                audio_num_heads=self.audio_num_heads,
                enable_a2v_cross_attn=a2v_enabled[layer_idx],
                enable_v2a_cross_attn=v2a_enabled[layer_idx],
                zero_init_output=zero_init_cross_attn,
                zero_init_video_cross_attn=zero_init_video_cross_attn,
                zero_init_audio_cross_attn=zero_init_audio_cross_attn,
                use_a2v_gating=resolved_use_a2v_gating,
                use_v2a_gating=resolved_use_v2a_gating,
                zero_init_a2v_gating=resolved_zero_init_a2v_gating,
                zero_init_v2a_gating=resolved_zero_init_v2a_gating,
                a2v_gate_init_value=resolved_a2v_gate_init_value,
                v2a_gate_init_value=resolved_v2a_gate_init_value,
                a2v_gate_alpha=resolved_a2v_gate_alphas[layer_idx],
                v2a_gate_alpha=resolved_v2a_gate_alphas[layer_idx],
            )
            for layer_idx, (video_block, audio_block) in enumerate(zip(video_blocks, audio_blocks))
        ])

        video_model.blocks = nn.ModuleList()
        audio_model.blocks = nn.ModuleList()

        self.use_temporal_rope = use_temporal_rope
        self.audio_fps = audio_fps
        self.vae_temporal_stride = vae_temporal_stride
        if use_temporal_rope:
            head_dim = self.video_dim // self.video_num_heads
            self.temporal_rope_inv_freq = 1.0 / (
                10000.0 ** (torch.arange(0, head_dim, 2, dtype=torch.float64) / head_dim)
            )
        else:
            self.temporal_rope_inv_freq = None

        # Ulysses-style sequence-parallel inference state. All ranks keep a
        # complete copy of the weights and exchange token/head dimensions in
        # attention, matching the Wan2.2 inference design.
        self.sp_world_size = 1
        self.sp_world_rank = 0
        self.sp_group = None

    def enable_multi_gpus_inference(self, group=None) -> None:
        """Enable raw-process-group sequence parallelism for inference."""
        import torch.distributed as dist

        if not dist.is_initialized():
            raise RuntimeError("Sequence-parallel inference requires an initialized process group")
        self.sp_world_size = dist.get_world_size(group)
        self.sp_world_rank = dist.get_rank(group)
        self.sp_group = group
        if self.video_num_heads % self.sp_world_size != 0:
            raise ValueError(
                f"Video attention heads ({self.video_num_heads}) must be divisible by "
                f"SP size ({self.sp_world_size})"
            )
        if self.audio_num_heads % self.sp_world_size != 0:
            raise ValueError(
                f"Audio attention heads ({self.audio_num_heads}) must be divisible by "
                f"SP size ({self.sp_world_size})"
            )

    def _apply_sequence_parallel(self, video_state: Dict[str, Any], audio_state: Dict[str, Any]):
        """Shard prepared video/audio token states along their sequence axes."""
        if self.sp_world_size <= 1:
            video_state["sp_rank"] = 0
            video_state["sp_world_size"] = 1
            audio_state["sp_rank"] = 0
            audio_state["sp_world_size"] = 1
            return video_state, audio_state

        def shard(state: Dict[str, Any], *, audio: bool):
            local_len = state["x"].size(1) // self.sp_world_size
            rank = self.sp_world_rank
            state["x"] = torch.chunk(state["x"], self.sp_world_size, dim=1)[rank]
            if not audio:
                if state["e"].dim() >= 3:
                    state["e"] = torch.chunk(state["e"], self.sp_world_size, dim=1)[rank]
                if state["e0"].dim() >= 4:
                    state["e0"] = torch.chunk(state["e0"], self.sp_world_size, dim=1)[rank]
            state["freqs"] = (
                torch.chunk(state["freqs"], self.sp_world_size, dim=0)[rank]
                if audio else state["freqs"]
            )
            if state.get("temporal_positions") is not None:
                state["temporal_positions"] = torch.chunk(
                    state["temporal_positions"], self.sp_world_size, dim=1
                )[rank]
            state["local_seq_lens"] = (
                state["seq_lens"] - rank * local_len
            ).clamp(min=0, max=local_len)
            state["sp_rank"] = rank
            state["sp_world_size"] = self.sp_world_size
            return state

        return shard(video_state, audio=False), shard(audio_state, audio=True)

    def _prepare_video(self, video_inputs: Dict[str, Any], dtype: torch.dtype) -> Dict[str, Any]:
        video_model = self.video_model
        device = video_model.patch_embedding.weight.device

        if video_model.freqs.device != device:
            video_model.freqs = video_model.freqs.to(device)

        x_list = _to_list(video_inputs["x"])
        y = video_inputs.get("y")
        if y is not None:
            y_list = _to_list(y)
            x_list = [torch.cat([sample, condition], dim=0) for sample, condition in zip(x_list, y_list)]

        x_list = [video_model.patch_embedding(sample.unsqueeze(0)) for sample in x_list]
        grid_sizes = torch.stack([
            torch.tensor(sample.shape[2:], dtype=torch.long, device=device) for sample in x_list
        ])
        x_list = [sample.flatten(2).transpose(1, 2) for sample in x_list]
        seq_lens = torch.tensor([sample.size(1) for sample in x_list], dtype=torch.long, device=device)

        seq_len = self._round_seq_len(int(video_inputs["seq_len"]))
        assert int(seq_lens.max().item()) <= seq_len
        x = torch.cat([
            torch.cat([sample, sample.new_zeros(1, seq_len - sample.size(1), sample.size(2))], dim=1)
            for sample in x_list
        ])

        timesteps = video_inputs["t"].to(device)
        embedding, modulation = _build_time_embeddings(
            video_model.time_embedding,
            video_model.time_projection,
            int(video_model.freq_dim),
            int(video_model.dim),
            timesteps,
            seq_len,
        )

        context = _embed_context(video_model.text_embedding, video_inputs["context"], int(video_model.text_len))

        return {
            "x": x,
            "e": embedding,
            "e0": modulation,
            "seq_lens": seq_lens,
            "grid_sizes": grid_sizes,
            "freqs": video_model.freqs,
            "context": context,
            "context_lens": None,
            "seq_len": seq_len,
        }

    def _prepare_audio(self, audio_inputs: Dict[str, Any], dtype: torch.dtype) -> Dict[str, Any]:
        audio_model = self.audio_model
        device = audio_model.patch_embedding.weight.device

        x_list = _to_list(audio_inputs["x"])
        y = audio_inputs.get("y")
        if y is not None:
            y_list = _to_list(y)
            x_list = [torch.cat([sample, condition], dim=0) for sample, condition in zip(x_list, y_list)]

        original_audio_shapes = [tuple(sample.shape) for sample in x_list]

        patchified = []
        grid_sizes_list = []
        for sample in x_list:
            tokens = audio_model.patch_embedding(sample.unsqueeze(0).to(device))
            tokens = rearrange(tokens, "1 c f -> f c").contiguous()
            patchified.append(tokens)
            grid_sizes_list.append(tokens.shape[0])

        grid_sizes = torch.tensor([[grid_size] for grid_size in grid_sizes_list], dtype=torch.long, device=device)
        seq_lens = grid_sizes[:, 0]
        seq_len = self._round_seq_len(int(audio_inputs["seq_len"]))
        assert int(seq_lens.max().item()) <= seq_len

        x = torch.stack([
            torch.cat([sample, sample.new_zeros(seq_len - sample.size(0), sample.size(1))], dim=0)
            for sample in patchified
        ])

        timesteps = audio_inputs["t"].to(device)
        embedding, modulation = _build_time_embeddings(
            audio_model.time_embedding,
            audio_model.time_projection,
            int(audio_model.freq_dim),
            int(audio_model.dim),
            timesteps,
            seq_len,
        )

        context = _embed_context(audio_model.text_embedding, audio_inputs["context"], int(audio_model.text_len))
        freqs = audio_model._build_freqs(seq_len, device)

        clip_fea = audio_inputs.get("clip_fea")
        if audio_model.has_image_input and clip_fea is not None:
            clip_embedding = audio_model.img_emb(clip_fea)
            context = torch.cat([clip_embedding, context], dim=1)

        return {
            "x": x,
            "e": embedding,
            "e0": modulation,
            "seq_lens": seq_lens,
            "grid_sizes": grid_sizes,
            "freqs": freqs,
            "context": context,
            "context_lens": None,
            "seq_len": seq_len,
            "original_audio_shapes": original_audio_shapes,
        }

    def forward(
        self,
        video: Dict[str, Any],
        audio: Dict[str, Any],
        dtype: torch.dtype = torch.bfloat16,
        return_dict: bool = True,
        enable_a2v: Optional[Union[bool, torch.Tensor]] = None,
        enable_v2a: Optional[Union[bool, torch.Tensor]] = None,
    ):
        """Forward pass of the joint audio-video model.

        Args:
            video: video input dict with keys 'x', 't', 'context', 'seq_len', etc.
            audio: audio input dict with keys 'x', 't', 'context', 'seq_len', etc.
            dtype: computation dtype for attention ops (default bfloat16).
            return_dict: if True, return dict with 'video'/'audio' keys; else tuple.
            enable_a2v: gate A2V cross-attention (video attending to audio).
                A bool tensor can select the enabled CFG branches per sample.
            enable_v2a: gate V2A cross-attention (audio attending to video), same
                semantics as enable_a2v.

        Returns:
            dict or tuple of (video_output, audio_output) tensors.
        """
        video_state = self._prepare_video(video, dtype)
        audio_state = self._prepare_audio(audio, dtype)

        device = video_state["x"].device

        video_temporal_pos = None
        audio_temporal_pos = None
        temporal_rope_inv_freq = None
        if self.use_temporal_rope and self.temporal_rope_inv_freq is not None:
            video_fps = float(video.get("video_fps", 16.0))
            temporal_rope_inv_freq = self.temporal_rope_inv_freq
            video_temporal_pos = compute_video_temporal_positions(
                video_state["grid_sizes"],
                video_state["x"].size(1),
                device,
                audio_fps=self.audio_fps,
                video_fps=video_fps,
                vae_temporal_stride=self.vae_temporal_stride,
            )
            audio_temporal_pos = compute_audio_temporal_positions(
                audio_state["seq_lens"],
                audio_state["x"].size(1),
                device,
            )

        video_state["temporal_positions"] = video_temporal_pos
        audio_state["temporal_positions"] = audio_temporal_pos
        video_state, audio_state = self._apply_sequence_parallel(video_state, audio_state)
        video_x = video_state["x"]
        audio_x = audio_state["x"]

        video_kwargs = {
            "e0": video_state["e0"],
            "seq_lens": video_state["seq_lens"],
            "grid_sizes": video_state["grid_sizes"],
            "freqs": video_state["freqs"],
            "context": video_state["context"],
            "context_lens": video_state["context_lens"],
            "temporal_positions": video_state["temporal_positions"],
            "temporal_rope_inv_freq": temporal_rope_inv_freq,
            "sp_rank": video_state["sp_rank"],
            "sp_world_size": video_state["sp_world_size"],
        }
        audio_kwargs = {
            "e0": audio_state["e0"],
            "seq_lens": audio_state["seq_lens"],
            "freqs": audio_state["freqs"],
            "context": audio_state["context"],
            "temporal_positions": audio_state["temporal_positions"],
            "sp_rank": audio_state["sp_rank"],
            "sp_world_size": audio_state["sp_world_size"],
        }

        runtime_cross_attn = {
            "enable_a2v": enable_a2v,
            "enable_v2a": enable_v2a,
            "sequence_parallel": self.sp_world_size > 1,
            "sp_group": self.sp_group,
        }

        for joint_block in self.joint_blocks:
            video_x, audio_x = joint_block(
                video_x, audio_x, video_kwargs, audio_kwargs, dtype,
                **runtime_cross_attn,
            )

        video_output = self.video_model.head(video_x, video_state["e"])
        audio_output = self.audio_model.head(audio_x, audio_state["e"])

        if self.sp_world_size > 1:
            video_output = all_gather_sequence(video_output, group=self.sp_group)
            audio_output = all_gather_sequence(audio_output, group=self.sp_group)

        video_output = torch.stack(
            self.video_model.unpatchify(video_output, video_state["grid_sizes"])
        )
        audio_output = torch.stack(
            self.audio_model.unpatchify(
                audio_output,
                audio_state["grid_sizes"],
                audio_state["original_audio_shapes"],
            )
        )
        result = {"video": video_output, "audio": audio_output}
        return result if return_dict else (video_output, audio_output)

    def _round_seq_len(self, seq_len: int) -> int:
        if self.sp_world_size > 1:
            return int(math.ceil(seq_len / self.sp_world_size) * self.sp_world_size)
        return int(seq_len)

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_path: Optional[str] = None,
        video_pretrained_model_path: Optional[str] = None,
        audio_pretrained_model_path: Optional[str] = None,
        video_subfolder: Optional[str] = None,
        audio_subfolder: Optional[str] = None,
        video_kwargs: Optional[Dict] = None,
        audio_kwargs: Optional[Dict] = None,
        video_model_cls=Wan2_2Transformer3DModel,
        audio_model_cls=CreatorAudioModel,
        low_cpu_mem_usage: bool = False,
        torch_dtype: torch.dtype = torch.bfloat16,
        use_temporal_rope: bool = True,
        audio_fps: float = 48000.0 / 960.0,
        vae_temporal_stride: int = 4,
        zero_init_cross_attn: bool = False,
        zero_init_video_cross_attn: bool | None = None,
        zero_init_audio_cross_attn: bool | None = None,
        a2v_cross_attn_layers: LayerSelection = None,
        v2a_cross_attn_layers: LayerSelection = None,
        use_gating: bool = True,
        use_a2v_gating: bool | None = None,
        use_v2a_gating: bool | None = None,
        zero_init_gating: bool = False,
        zero_init_a2v_gating: bool | None = None,
        zero_init_v2a_gating: bool | None = None,
        gate_init_value: Optional[float] = None,
        a2v_gate_init_value: Optional[float] = None,
        v2a_gate_init_value: Optional[float] = None,
        a2v_gate_alphas: LayerAlphas = None,
        v2a_gate_alphas: LayerAlphas = None,
    ):
        video_kwargs = dict(video_kwargs or {})
        audio_kwargs = dict(audio_kwargs or {})

        if video_pretrained_model_path is not None and audio_pretrained_model_path is not None:
            video_path = video_pretrained_model_path
            audio_path = audio_pretrained_model_path
        elif pretrained_model_path is not None:
            video_path = os.path.join(pretrained_model_path, "video_model")
            audio_path = os.path.join(pretrained_model_path, "audio_model")
        else:
            raise ValueError(
                "Must provide either pretrained_model_path or both video_pretrained_model_path "
                "and audio_pretrained_model_path"
            )

        logging.info("Loading video model from: %s", video_path)
        video_model = video_model_cls.from_pretrained(
            video_path,
            subfolder=video_subfolder,
            transformer_additional_kwargs=video_kwargs,
            low_cpu_mem_usage=low_cpu_mem_usage,
            torch_dtype=torch_dtype,
        )

        logging.info("Loading audio model from: %s", audio_path)
        audio_model = audio_model_cls.from_pretrained(
            audio_path,
            subfolder=audio_subfolder,
            transformer_additional_kwargs=audio_kwargs,
            low_cpu_mem_usage=low_cpu_mem_usage,
            torch_dtype=torch_dtype,
        )

        model = cls(
            video_model=video_model,
            audio_model=audio_model,
            use_temporal_rope=use_temporal_rope,
            audio_fps=audio_fps,
            vae_temporal_stride=vae_temporal_stride,
            zero_init_cross_attn=zero_init_cross_attn,
            zero_init_video_cross_attn=zero_init_video_cross_attn,
            zero_init_audio_cross_attn=zero_init_audio_cross_attn,
            a2v_cross_attn_layers=a2v_cross_attn_layers,
            v2a_cross_attn_layers=v2a_cross_attn_layers,
            use_gating=use_gating,
            use_a2v_gating=use_a2v_gating,
            use_v2a_gating=use_v2a_gating,
            zero_init_gating=zero_init_gating,
            zero_init_a2v_gating=zero_init_a2v_gating,
            zero_init_v2a_gating=zero_init_v2a_gating,
            gate_init_value=gate_init_value,
            a2v_gate_init_value=a2v_gate_init_value,
            v2a_gate_init_value=v2a_gate_init_value,
            a2v_gate_alphas=a2v_gate_alphas,
            v2a_gate_alphas=v2a_gate_alphas,
        ).to(torch_dtype)
        if pretrained_model_path is not None:
            cross_attn_file = os.path.join(pretrained_model_path, "cross_attn_weights.safetensors")
            cross_attn_file_bin = os.path.join(pretrained_model_path, "cross_attn_weights.bin")

            if os.path.exists(cross_attn_file):
                from safetensors.torch import load_file

                cross_attn_state = load_file(cross_attn_file)
                logging.info("Loading cross-attn weights from: %s (%d keys)", cross_attn_file, len(cross_attn_state))
            elif os.path.exists(cross_attn_file_bin):
                cross_attn_state = torch.load(cross_attn_file_bin, map_location="cpu")
                logging.info("Loading cross-attn weights from: %s (%d keys)", cross_attn_file_bin, len(cross_attn_state))
            else:
                cross_attn_state = None
                logging.warning("No cross_attn_weights found in %s, skipping.", pretrained_model_path)

            if cross_attn_state is not None:
                missing, unexpected = model.load_state_dict(cross_attn_state, strict=False)
                logging.info("Cross-attn load: %d missing, %d unexpected keys", len(missing), len(unexpected))
                if unexpected:
                    logging.warning("Unexpected keys in cross_attn_weights: %s", unexpected[:10])

        return model


WanCreatorCrossAttnGatingAVModel = WanCreatorGatingAVModel
WanCreatorGatedCrossAttnAVModel = WanCreatorGatingAVModel

__all__ = [
    "GatedCrossModalAttention",
    "GatedJointBlock",
    "WanCreatorGatingAVModel",
    "WanCreatorCrossAttnGatingAVModel",
    "WanCreatorGatedCrossAttnAVModel",
    "compute_audio_temporal_positions",
    "compute_video_temporal_positions",
]