| 58 | |
| 59 | |
| 60 | class StylizationBlock(nn.Module): |
| 61 | |
| 62 | def __init__(self, latent_dim, time_embed_dim, dropout): |
| 63 | super().__init__() |
| 64 | self.emb_layers = nn.Sequential( |
| 65 | nn.SiLU(), |
| 66 | nn.Linear(time_embed_dim, 2 * latent_dim), |
| 67 | ) |
| 68 | self.norm = nn.LayerNorm(latent_dim) |
| 69 | self.out_layers = nn.Sequential( |
| 70 | nn.SiLU(), |
| 71 | nn.Dropout(p=dropout), |
| 72 | zero_module(nn.Linear(latent_dim, latent_dim)), |
| 73 | ) |
| 74 | |
| 75 | def forward(self, h, emb): |
| 76 | """ |
| 77 | h: B, T, D |
| 78 | emb: B, D |
| 79 | """ |
| 80 | # B, 1, 2D |
| 81 | emb_out = self.emb_layers(emb).unsqueeze(1) |
| 82 | # scale: B, 1, D / shift: B, 1, D |
| 83 | scale, shift = torch.chunk(emb_out, 2, dim=2) |
| 84 | h = self.norm(h) * (1 + scale) + shift |
| 85 | h = self.out_layers(h) |
| 86 | return h |
| 87 | |
| 88 | |
| 89 | class LinearTemporalSelfAttention(nn.Module): |