| 4 | |
| 5 | |
| 6 | class TemporalTransformerBlock(torch.nn.Module): |
| 7 | |
| 8 | def __init__(self, dim, num_attention_heads, attention_head_dim, max_position_embeddings=32): |
| 9 | super().__init__() |
| 10 | |
| 11 | # 1. Self-Attn |
| 12 | self.pe1 = torch.nn.Parameter(torch.zeros(1, max_position_embeddings, dim)) |
| 13 | self.norm1 = torch.nn.LayerNorm(dim, elementwise_affine=True) |
| 14 | self.attn1 = Attention(q_dim=dim, num_heads=num_attention_heads, head_dim=attention_head_dim, bias_out=True) |
| 15 | |
| 16 | # 2. Cross-Attn |
| 17 | self.pe2 = torch.nn.Parameter(torch.zeros(1, max_position_embeddings, dim)) |
| 18 | self.norm2 = torch.nn.LayerNorm(dim, elementwise_affine=True) |
| 19 | self.attn2 = Attention(q_dim=dim, num_heads=num_attention_heads, head_dim=attention_head_dim, bias_out=True) |
| 20 | |
| 21 | # 3. Feed-forward |
| 22 | self.norm3 = torch.nn.LayerNorm(dim, elementwise_affine=True) |
| 23 | self.act_fn = GEGLU(dim, dim * 4) |
| 24 | self.ff = torch.nn.Linear(dim * 4, dim) |
| 25 | |
| 26 | |
| 27 | def forward(self, hidden_states, batch_size=1): |
| 28 | |
| 29 | # 1. Self-Attention |
| 30 | norm_hidden_states = self.norm1(hidden_states) |
| 31 | norm_hidden_states = rearrange(norm_hidden_states, "(b f) h c -> (b h) f c", b=batch_size) |
| 32 | attn_output = self.attn1(norm_hidden_states + self.pe1[:, :norm_hidden_states.shape[1]]) |
| 33 | attn_output = rearrange(attn_output, "(b h) f c -> (b f) h c", b=batch_size) |
| 34 | hidden_states = attn_output + hidden_states |
| 35 | |
| 36 | # 2. Cross-Attention |
| 37 | norm_hidden_states = self.norm2(hidden_states) |
| 38 | norm_hidden_states = rearrange(norm_hidden_states, "(b f) h c -> (b h) f c", b=batch_size) |
| 39 | attn_output = self.attn2(norm_hidden_states + self.pe2[:, :norm_hidden_states.shape[1]]) |
| 40 | attn_output = rearrange(attn_output, "(b h) f c -> (b f) h c", b=batch_size) |
| 41 | hidden_states = attn_output + hidden_states |
| 42 | |
| 43 | # 3. Feed-forward |
| 44 | norm_hidden_states = self.norm3(hidden_states) |
| 45 | ff_output = self.act_fn(norm_hidden_states) |
| 46 | ff_output = self.ff(ff_output) |
| 47 | hidden_states = ff_output + hidden_states |
| 48 | |
| 49 | return hidden_states |
| 50 | |
| 51 | |
| 52 | class TemporalBlock(torch.nn.Module): |