| 227 | return y |
| 228 | |
| 229 | class TemporalCrossAttention(nn.Module): |
| 230 | |
| 231 | def __init__(self, seq_len, latent_dim, text_latent_dim, num_head, dropout, time_embed_dim): |
| 232 | super().__init__() |
| 233 | self.num_head = num_head |
| 234 | self.norm = nn.LayerNorm(latent_dim) |
| 235 | self.text_norm = nn.LayerNorm(text_latent_dim) |
| 236 | self.query = nn.Linear(latent_dim, latent_dim) |
| 237 | self.key = nn.Linear(text_latent_dim, latent_dim) |
| 238 | self.value = nn.Linear(text_latent_dim, latent_dim) |
| 239 | self.dropout = nn.Dropout(dropout) |
| 240 | self.proj_out = StylizationBlock(latent_dim, time_embed_dim, dropout) |
| 241 | |
| 242 | def forward(self, x, xf, emb): |
| 243 | """ |
| 244 | x: B, T, D |
| 245 | xf: B, N, L |
| 246 | """ |
| 247 | B, T, D = x.shape |
| 248 | N = xf.shape[1] |
| 249 | H = self.num_head |
| 250 | # B, T, 1, D |
| 251 | query = self.query(self.norm(x)).unsqueeze(2) |
| 252 | # B, 1, N, D |
| 253 | key = self.key(self.text_norm(xf)).unsqueeze(1) |
| 254 | query = query.view(B, T, H, -1) |
| 255 | key = key.view(B, N, H, -1) |
| 256 | # B, T, N, H |
| 257 | attention = torch.einsum('bnhd,bmhd->bnmh', query, key) / math.sqrt(D // H) |
| 258 | weight = self.dropout(F.softmax(attention, dim=2)) |
| 259 | value = self.value(self.text_norm(xf)).view(B, N, H, -1) |
| 260 | y = torch.einsum('bnmh,bmhd->bnhd', weight, value).reshape(B, T, D) |
| 261 | y = x + self.proj_out(y, emb) |
| 262 | return y |
| 263 | |
| 264 | class TemporalDiffusionTransformerDecoderLayer(nn.Module): |
| 265 | |