wan_animate_face ¶

def __init__(self, bias_channels: int) -> None:
    super().__init__()
    self.bias = nn.Parameter(torch.zeros(bias_channels))

def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1,
             padding: int = 0, bias: bool = True, blur_kernel: tuple[int, ...] | None = None,
             use_activation: bool = True) -> None:
    super().__init__()
    self.in_channels = in_channels
    self.stride = stride
    self.padding = padding

    self.blur = blur_kernel is not None
    self._blur_kernel_taps = tuple(blur_kernel) if blur_kernel is not None else None
    if self.blur:
        p = (len(blur_kernel) - stride) + (kernel_size - 1)
        self.blur_padding = ((p + 1) // 2, p // 2)
        self.register_buffer("blur_kernel", self._build_blur_kernel(), persistent=False)

    self.weight = nn.Parameter(torch.randn(out_channels, in_channels, kernel_size, kernel_size))
    # The load-bearing constant: applied to the *stored* weight every forward.
    self.scale = 1 / math.sqrt(in_channels * kernel_size**2)

    # With an activation the bias is fused into it (act_fn.bias); without,
    # it lives on the conv. Mutually exclusive, exactly like the checkpoint.
    self.bias = nn.Parameter(torch.zeros(out_channels)) if (bias and not use_activation) else None
    self.act_fn = FusedLeakyReLU(bias_channels=out_channels) if use_activation else None

def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3,
             kernel_size_skip: int = 1, blur_kernel: tuple[int, ...] = (1, 3, 3, 1),
             downsample_factor: int = 2) -> None:
    super().__init__()
    self.conv1 = MotionConv2d(in_channels, in_channels, kernel_size, stride=1,
                              padding=kernel_size // 2, use_activation=True)
    self.conv2 = MotionConv2d(in_channels, out_channels, kernel_size, stride=downsample_factor,
                              padding=0, blur_kernel=blur_kernel, use_activation=True)
    self.conv_skip = MotionConv2d(in_channels, out_channels, kernel_size_skip, stride=downsample_factor,
                                  padding=0, bias=False, blur_kernel=blur_kernel, use_activation=False)

def __init__(self, in_dim: int, out_dim: int) -> None:
    super().__init__()
    self.weight = nn.Parameter(torch.randn(out_dim, in_dim))
    self.scale = 1 / math.sqrt(in_dim)
    self.bias = nn.Parameter(torch.zeros(out_dim))

def __init__(self, dim: int, num_heads: int, eps: float = 1e-6) -> None:
    super().__init__()
    assert dim % num_heads == 0
    self.num_heads = num_heads
    self.head_dim = dim // num_heads
    self.pre_norm_q = nn.LayerNorm(dim, eps, elementwise_affine=False)
    self.pre_norm_kv = nn.LayerNorm(dim, eps, elementwise_affine=False)
    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(self.head_dim, eps=eps)
    self.norm_k = RMSNorm(self.head_dim, eps=eps)
    self.attn = LocalAttention(num_heads=num_heads, head_size=self.head_dim, dropout_rate=0,
                               softmax_scale=None, causal=False,
                               supported_attention_backends=ANIMATE_ATTENTION_BACKENDS)

def __init__(self, in_dim: int, out_dim: int, hidden_dim: int = 1024, num_heads: int = 4,
             kernel_size: int = 3, eps: float = 1e-6, pad_mode: str = "replicate") -> None:
    super().__init__()
    self.num_heads = num_heads
    self.time_causal_padding = (kernel_size - 1, 0)
    self.pad_mode = pad_mode

    self.act = nn.SiLU()
    self.conv1_local = nn.Conv1d(in_dim, hidden_dim * num_heads, kernel_size, stride=1)
    self.conv2 = nn.Conv1d(hidden_dim, hidden_dim, kernel_size, stride=2)
    self.conv3 = nn.Conv1d(hidden_dim, hidden_dim, kernel_size, stride=2)
    self.norm1 = nn.LayerNorm(hidden_dim, eps, elementwise_affine=False)
    self.norm2 = nn.LayerNorm(hidden_dim, eps, elementwise_affine=False)
    self.norm3 = nn.LayerNorm(hidden_dim, eps, elementwise_affine=False)
    self.out_proj = nn.Linear(hidden_dim, out_dim)
    self.padding_tokens = nn.Parameter(torch.zeros(1, 1, 1, out_dim))

def __init__(self, size: int = 512, style_dim: int = 512, motion_dim: int = 20,
             out_dim: int = 512, motion_blocks: int = 5,
             channels: dict[str, int] | None = None) -> None:
    super().__init__()
    self.size = size
    if channels is None:
        channels = WAN_ANIMATE_MOTION_ENCODER_CHANNEL_SIZES

    self.conv_in = MotionConv2d(3, channels[str(size)], 1, use_activation=True)
    self.res_blocks = nn.ModuleList()
    in_channels = channels[str(size)]
    # bit_length is exact where float log2 can land just below the integer;
    # a size missing from the channel table fails loudly on the lookup.
    for i in range(size.bit_length() - 1, 2, -1):
        out_channels = channels[str(2**(i - 1))]
        self.res_blocks.append(MotionEncoderResBlock(in_channels, out_channels))
        in_channels = out_channels
    self.conv_out = MotionConv2d(in_channels, style_dim, 4, padding=0, bias=False, use_activation=False)

    # No activations between these linears -- LIA's design, kept verbatim.
    linears = [MotionLinear(style_dim, style_dim) for _ in range(motion_blocks - 1)]
    linears.append(MotionLinear(style_dim, motion_dim))
    self.motion_network = nn.ModuleList(linears)

    self.motion_synthesis_weight = nn.Parameter(torch.randn(out_dim, motion_dim))