| 87 | |
| 88 | |
| 89 | class LinearTemporalSelfAttention(nn.Module): |
| 90 | |
| 91 | def __init__(self, seq_len, latent_dim, num_head, dropout, time_embed_dim): |
| 92 | super().__init__() |
| 93 | self.num_head = num_head |
| 94 | self.norm = nn.LayerNorm(latent_dim) |
| 95 | self.query = nn.Linear(latent_dim, latent_dim) |
| 96 | self.key = nn.Linear(latent_dim, latent_dim) |
| 97 | self.value = nn.Linear(latent_dim, latent_dim) |
| 98 | self.dropout = nn.Dropout(dropout) |
| 99 | self.proj_out = StylizationBlock(latent_dim, time_embed_dim, dropout) |
| 100 | |
| 101 | def forward(self, x, emb, src_mask): |
| 102 | """ |
| 103 | x: B, T, D |
| 104 | """ |
| 105 | B, T, D = x.shape |
| 106 | H = self.num_head |
| 107 | # B, T, D |
| 108 | query = self.query(self.norm(x)) |
| 109 | # B, T, D |
| 110 | key = (self.key(self.norm(x)) + (1 - src_mask) * -1000000) |
| 111 | query = F.softmax(query.view(B, T, H, -1), dim=-1) |
| 112 | key = F.softmax(key.view(B, T, H, -1), dim=1) |
| 113 | # B, T, H, HD |
| 114 | value = (self.value(self.norm(x)) * src_mask).view(B, T, H, -1) |
| 115 | # B, H, HD, HD |
| 116 | attention = torch.einsum('bnhd,bnhl->bhdl', key, value) |
| 117 | y = torch.einsum('bnhd,bhdl->bnhl', query, attention).reshape(B, T, D) |
| 118 | y = x + self.proj_out(y, emb) |
| 119 | return y |
| 120 | |
| 121 | |
| 122 | class LinearTemporalCrossAttention(nn.Module): |