| 194 | return x |
| 195 | |
| 196 | class TemporalSelfAttention(nn.Module): |
| 197 | |
| 198 | def __init__(self, seq_len, latent_dim, num_head, dropout, time_embed_dim): |
| 199 | super().__init__() |
| 200 | self.num_head = num_head |
| 201 | self.norm = nn.LayerNorm(latent_dim) |
| 202 | self.query = nn.Linear(latent_dim, latent_dim) |
| 203 | self.key = nn.Linear(latent_dim, latent_dim) |
| 204 | self.value = nn.Linear(latent_dim, latent_dim) |
| 205 | self.dropout = nn.Dropout(dropout) |
| 206 | self.proj_out = StylizationBlock(latent_dim, time_embed_dim, dropout) |
| 207 | |
| 208 | def forward(self, x, emb, src_mask): |
| 209 | """ |
| 210 | x: B, T, D |
| 211 | """ |
| 212 | B, T, D = x.shape |
| 213 | H = self.num_head |
| 214 | # B, T, 1, D |
| 215 | query = self.query(self.norm(x)).unsqueeze(2) |
| 216 | # B, 1, T, D |
| 217 | key = self.key(self.norm(x)).unsqueeze(1) |
| 218 | query = query.view(B, T, H, -1) |
| 219 | key = key.view(B, T, H, -1) |
| 220 | # B, T, T, H |
| 221 | attention = torch.einsum('bnhd,bmhd->bnmh', query, key) / math.sqrt(D // H) |
| 222 | attention = attention + (1 - src_mask.unsqueeze(-1)) * -100000 |
| 223 | weight = self.dropout(F.softmax(attention, dim=2)) |
| 224 | value = self.value(self.norm(x)).view(B, T, H, -1) |
| 225 | y = torch.einsum('bnmh,bmhd->bnhd', weight, value).reshape(B, T, D) |
| 226 | y = x + self.proj_out(y, emb) |
| 227 | return y |
| 228 | |
| 229 | class TemporalCrossAttention(nn.Module): |
| 230 | |