(self, num_attention_heads, attention_head_dim, in_channels, num_layers=1, norm_num_groups=32, eps=1e-5)
| 52 | class TemporalBlock(torch.nn.Module): |
| 53 | |
| 54 | def __init__(self, num_attention_heads, attention_head_dim, in_channels, num_layers=1, norm_num_groups=32, eps=1e-5): |
| 55 | super().__init__() |
| 56 | inner_dim = num_attention_heads * attention_head_dim |
| 57 | |
| 58 | self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=eps, affine=True) |
| 59 | self.proj_in = torch.nn.Linear(in_channels, inner_dim) |
| 60 | |
| 61 | self.transformer_blocks = torch.nn.ModuleList([ |
| 62 | TemporalTransformerBlock( |
| 63 | inner_dim, |
| 64 | num_attention_heads, |
| 65 | attention_head_dim |
| 66 | ) |
| 67 | for d in range(num_layers) |
| 68 | ]) |
| 69 | |
| 70 | self.proj_out = torch.nn.Linear(inner_dim, in_channels) |
| 71 | |
| 72 | def forward(self, hidden_states, time_emb, text_emb, res_stack, batch_size=1): |
| 73 | batch, _, height, width = hidden_states.shape |
no test coverage detected