| from torch import nn |
| import torch |
| from einops import rearrange |
| import torch.nn.functional as F |
| from ..attention import attention |
|
|
| class CausalConv1d(nn.Module): |
| def __init__(self, chan_in, chan_out, kernel_size=3, stride=1, dilation=1, pad_mode="replicate", **kwargs): |
| super().__init__() |
|
|
| self.pad_mode = pad_mode |
| padding = (kernel_size - 1, 0) |
| self.time_causal_padding = padding |
|
|
| self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) |
|
|
| def forward(self, x): |
| x = F.pad(x, self.time_causal_padding, mode=self.pad_mode) |
| return self.conv(x) |
|
|
|
|
| class FaceEncoder(nn.Module): |
| def __init__(self, in_dim: int, out_dim: int, num_heads: int, dtype=None, device=None): |
| super().__init__() |
| self.dtype = dtype |
| self.device = device |
|
|
| self.num_heads = num_heads |
| self.conv1_local = CausalConv1d(in_dim, 1024 * num_heads, 3, stride=1) |
| self.norm1 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, device=device, dtype=dtype) |
| self.act = nn.SiLU() |
| self.conv2 = CausalConv1d(1024, 1024, 3, stride=2) |
| self.conv3 = CausalConv1d(1024, 1024, 3, stride=2) |
|
|
| self.out_proj = nn.Linear(1024, out_dim) |
|
|
| self.norm2 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, device=device, dtype=dtype) |
| self.norm3 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, device=device, dtype=dtype) |
|
|
| self.padding_tokens = nn.Parameter(torch.zeros(1, 1, 1, out_dim)) |
|
|
| def forward(self, x): |
| x = rearrange(x, "b t c -> b c t") |
| b = x.shape[0] |
|
|
| x = self.conv1_local(x) |
| x = rearrange(x, "b (n c) t -> (b n) t c", n=self.num_heads) |
|
|
| x = self.norm1(x) |
| x = self.act(x) |
| x = rearrange(x, "b t c -> b c t") |
| x = self.conv2(x) |
| x = rearrange(x, "b c t -> b t c") |
| x = self.norm2(x) |
| x = self.act(x) |
| x = rearrange(x, "b t c -> b c t") |
| x = self.conv3(x) |
| x = rearrange(x, "b c t -> b t c") |
| x = self.norm3(x) |
| x = self.act(x) |
| x = self.out_proj(x) |
| x = rearrange(x, "(b n) t c -> b t n c", b=b) |
| padding = self.padding_tokens.repeat(b, x.shape[1], 1, 1) |
|
|
| return torch.cat([x, padding], dim=-2) |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim, elementwise_affine=True, eps=1e-6, device=None, dtype=None): |
| super().__init__() |
| self.eps = eps |
| if elementwise_affine: |
| self.weight = nn.Parameter(torch.ones(dim, device=device, dtype=dtype)) |
|
|
| def _norm(self, x): |
| return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) |
|
|
| def forward(self, x): |
| output = self._norm(x.float()).type_as(x) |
| if hasattr(self, "weight"): |
| output = output * self.weight |
| return output |
|
|
|
|
| class FaceBlock(nn.Module): |
| def __init__(self, feature_dim, num_heads, dtype=None, device=None): |
| super().__init__() |
|
|
| self.feature_dim = feature_dim |
| self.num_heads = num_heads |
| head_dim = feature_dim // num_heads |
|
|
| self.linear1_kv = nn.Linear(feature_dim, feature_dim * 2, device=device, dtype=dtype) |
| self.linear1_q = nn.Linear(feature_dim, feature_dim, device=device, dtype=dtype) |
| self.linear2 = nn.Linear(feature_dim, feature_dim, device=device, dtype=dtype) |
|
|
| self.q_norm = (RMSNorm(head_dim, elementwise_affine=True, eps=1e-6, device=device, dtype=dtype)) |
| self.k_norm = (RMSNorm(head_dim, elementwise_affine=True, eps=1e-6, device=device, dtype=dtype)) |
|
|
| self.pre_norm_feat = nn.LayerNorm(feature_dim, elementwise_affine=False, eps=1e-6, device=device, dtype=dtype) |
| self.pre_norm_motion = nn.LayerNorm(feature_dim, elementwise_affine=False, eps=1e-6, device=device, dtype=dtype) |
|
|
|
|
| def forward(self, x, motion_vec, motion_mask=None): |
| B, T, N, C = motion_vec.shape |
|
|
| x_motion = self.pre_norm_motion(motion_vec) |
| x_feat = self.pre_norm_feat(x) |
|
|
| kv = self.linear1_kv(x_motion) |
| q = self.linear1_q(x_feat) |
|
|
| k, v = rearrange(kv, "B L N (K H D) -> K B L N H D", K=2, H=self.num_heads) |
| q = rearrange(q, "B S (H D) -> B S H D", H=self.num_heads) |
|
|
| q = self.q_norm(q).to(v) |
| k = self.k_norm(k).to(v) |
|
|
| k = rearrange(k, "B L N H D -> (B L) N H D") |
| v = rearrange(v, "B L N H D -> (B L) N H D") |
| q = rearrange(q, "B (L S) H D -> (B L) S H D", L=T) |
|
|
| attn = attention(q, k, v) |
| attn = attn.reshape(attn.shape[0], attn.shape[1], -1) |
| attn = rearrange(attn, "(B L) S C -> B (L S) C", L=T) |
| output = self.linear2(attn) |
|
|
| if motion_mask is not None: |
| output = output * rearrange(motion_mask, "B T H W -> B (T H W)").unsqueeze(-1) |
|
|
| return output |
