| 120 | |
| 121 | |
| 122 | class LinearTemporalCrossAttention(nn.Module): |
| 123 | |
| 124 | def __init__(self, seq_len, latent_dim, text_latent_dim, num_head, dropout, time_embed_dim): |
| 125 | super().__init__() |
| 126 | self.num_head = num_head |
| 127 | self.norm = nn.LayerNorm(latent_dim) |
| 128 | self.text_norm = nn.LayerNorm(text_latent_dim) |
| 129 | self.query = nn.Linear(latent_dim, latent_dim) |
| 130 | self.key = nn.Linear(text_latent_dim, latent_dim) |
| 131 | self.value = nn.Linear(text_latent_dim, latent_dim) |
| 132 | self.dropout = nn.Dropout(dropout) |
| 133 | self.proj_out = StylizationBlock(latent_dim, time_embed_dim, dropout) |
| 134 | |
| 135 | def forward(self, x, xf, emb): |
| 136 | """ |
| 137 | x: B, T, D |
| 138 | xf: B, N, L |
| 139 | """ |
| 140 | B, T, D = x.shape |
| 141 | N = xf.shape[1] |
| 142 | H = self.num_head |
| 143 | # B, T, D |
| 144 | query = self.query(self.norm(x)) |
| 145 | # B, N, D |
| 146 | key = self.key(self.text_norm(xf)) |
| 147 | query = F.softmax(query.view(B, T, H, -1), dim=-1) |
| 148 | key = F.softmax(key.view(B, N, H, -1), dim=1) |
| 149 | # B, N, H, HD |
| 150 | value = self.value(self.text_norm(xf)).view(B, N, H, -1) |
| 151 | # B, H, HD, HD |
| 152 | attention = torch.einsum('bnhd,bnhl->bhdl', key, value) |
| 153 | y = torch.einsum('bnhd,bhdl->bnhl', query, attention).reshape(B, T, D) |
| 154 | y = x + self.proj_out(y, emb) |
| 155 | return y |
| 156 | |
| 157 | class FFN(nn.Module): |
| 158 | |