wan_s2v_audio ¶

def __init__(self, dim: int, eps: float = 1e-5) -> None:
    super().__init__()
    self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
    self.act = nn.SiLU()
    self.linear = nn.Linear(dim, 2 * dim)

def __init__(self, dim: int, num_heads: int, qk_norm: bool = True, eps: float = 1e-6) -> None:
    super().__init__()
    assert dim % num_heads == 0
    self.num_heads = num_heads
    self.head_dim = dim // num_heads
    self.to_q = ReplicatedLinear(dim, dim)
    self.to_k = ReplicatedLinear(dim, dim)
    self.to_v = ReplicatedLinear(dim, dim)
    self.to_out = ReplicatedLinear(dim, dim)
    self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
    self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
    self.attn = LocalAttention(num_heads=num_heads, head_size=self.head_dim, dropout_rate=0,
                               softmax_scale=None, causal=False,
                               supported_attention_backends=S2V_ATTENTION_BACKENDS)

def __init__(self, dim: int, num_heads: int, inject_layers: tuple[int, ...],
             enable_adain: bool = True, adain_dim: int | None = None, eps: float = 1e-6) -> None:
    super().__init__()
    self.injected_block_id = {block_id: i for i, block_id in enumerate(sorted(inject_layers))}
    n_inject = len(self.injected_block_id)
    assert n_inject == len(inject_layers), "duplicate entries in audio_inject_layers"

    self.injector = nn.ModuleList(
        [AudioCrossAttention(dim=dim, num_heads=num_heads, qk_norm=True, eps=eps) for _ in range(n_inject)])
    self.injector_pre_norm_feat = nn.ModuleList(
        [nn.LayerNorm(dim, elementwise_affine=False, eps=eps) for _ in range(n_inject)])
    self.injector_pre_norm_vec = nn.ModuleList(
        [nn.LayerNorm(dim, elementwise_affine=False, eps=eps) for _ in range(n_inject)])
    self.enable_adain = enable_adain
    if enable_adain:
        self.injector_adain_layers = nn.ModuleList(
            # eps is AdaLayerNorm's own default (1e-5, from diffusers), not the
            # model-wide 1e-6 -- upstream constructs these without passing eps.
            [AdaLayerNorm(adain_dim or dim) for _ in range(n_inject)])

def __init__(self, dim: int = 1024, num_layers: int = 25, out_dim: int = 5120,
             num_token: int = 4, need_global: bool = False) -> None:
    super().__init__()
    self.encoder = MotionEncoder(in_dim=dim, hidden_dim=out_dim, num_heads=num_token,
                                 need_global=need_global)
    # num_layers is the wav2vec2 hidden-state count (25 for wav2vec2-large),
    # NOT a function of num_token. Deriving it is a silent-failure bug.
    self.weights = nn.Parameter(torch.ones((1, num_layers, 1, 1)) * 0.01)
    self.act = nn.SiLU()

def __init__(self, chan_in: int, chan_out: int, kernel_size: int = 3, stride: int = 1,
             pad_mode: str = "replicate") -> None:
    super().__init__()
    self.pad_mode = pad_mode
    self.time_causal_padding = (kernel_size - 1, 0)
    self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, stride=stride)

def __init__(self, in_dim: int, hidden_dim: int, num_heads: int = 4, need_global: bool = True) -> None:
    super().__init__()
    self.num_heads = num_heads
    self.need_global = need_global

    self.conv1_local = CausalConv1d(in_dim, hidden_dim // 4 * num_heads, 3, stride=1)
    if need_global:
        self.conv1_global = CausalConv1d(in_dim, hidden_dim // 4, 3, stride=1)
    self.conv2 = CausalConv1d(hidden_dim // 4, hidden_dim // 2, 3, stride=2)
    self.conv3 = CausalConv1d(hidden_dim // 2, hidden_dim, 3, stride=2)
    if need_global:
        self.final_linear = nn.Linear(hidden_dim, hidden_dim)

    self.norm1 = nn.LayerNorm(hidden_dim // 4, elementwise_affine=False, eps=1e-6)
    self.norm2 = nn.LayerNorm(hidden_dim // 2, elementwise_affine=False, eps=1e-6)
    self.norm3 = nn.LayerNorm(hidden_dim, elementwise_affine=False, eps=1e-6)
    self.act = nn.SiLU()
    self.padding_tokens = nn.Parameter(torch.zeros(1, 1, 1, hidden_dim))