(self, latent_dim, text_latent_dim, num_heads, num_text_heads,
num_experts, topk, gate_type, gate_noise, ffn_dim,
time_embed_dim, max_seq_len, max_text_seq_len, temporal_comb,
dropout)
| 63 | class SAMI(nn.Module): |
| 64 | |
| 65 | def __init__(self, latent_dim, text_latent_dim, num_heads, num_text_heads, |
| 66 | num_experts, topk, gate_type, gate_noise, ffn_dim, |
| 67 | time_embed_dim, max_seq_len, max_text_seq_len, temporal_comb, |
| 68 | dropout): |
| 69 | super().__init__() |
| 70 | self.latent_dim = latent_dim |
| 71 | self.num_heads = num_heads |
| 72 | self.num_text_heads = num_text_heads |
| 73 | self.max_seq_len = max_seq_len |
| 74 | |
| 75 | self.norm = nn.LayerNorm(latent_dim) |
| 76 | self.text_norm = nn.LayerNorm(text_latent_dim) |
| 77 | |
| 78 | self.sigma = nn.Parameter(torch.Tensor([100])) |
| 79 | self.time = torch.arange(max_seq_len) / max_seq_len |
| 80 | self.text_moe = MOE(num_experts, topk, text_latent_dim, |
| 81 | text_latent_dim * 4, 2 * latent_dim, |
| 82 | num_text_heads, max_text_seq_len, gate_type, |
| 83 | gate_noise) |
| 84 | self.motion_moe = MOE(num_experts, topk, latent_dim, latent_dim * 4, |
| 85 | 3 * latent_dim, num_heads, max_seq_len, |
| 86 | gate_type, gate_noise) |
| 87 | self.key_motion = nn.Parameter(torch.randn(max_seq_len, latent_dim)) |
| 88 | self.body_weight = nn.Parameter(torch.randn(num_heads, num_heads)) |
| 89 | |
| 90 | self.template_s = get_ffn(latent_dim, ffn_dim) |
| 91 | self.template_v = get_ffn(latent_dim, ffn_dim) |
| 92 | self.template_a = get_ffn(latent_dim, ffn_dim) |
| 93 | self.template_j = get_ffn(latent_dim, ffn_dim) |
| 94 | self.template_t = nn.Sequential(nn.Linear(latent_dim, ffn_dim), |
| 95 | nn.GELU(), nn.Linear(ffn_dim, 1)) |
| 96 | self.t_sigma = nn.Parameter(torch.Tensor([1])) |
| 97 | self.proj_out = StylizationBlock(latent_dim * num_heads, |
| 98 | time_embed_dim, dropout) |
| 99 | self.temporal_comb = temporal_comb |
| 100 | |
| 101 | def forward(self, x, xf, emb, src_mask, cond_type, motion_length, |
| 102 | num_intervals, **kwargs): |
nothing calls this directly
no test coverage detected