| 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): |