(self, num_attention_heads, attention_head_dim, in_channels, cross_attention_dim=None, add_positional_conv=None)
| 138 | class TemporalAttentionBlock(torch.nn.Module): |
| 139 | |
| 140 | def __init__(self, num_attention_heads, attention_head_dim, in_channels, cross_attention_dim=None, add_positional_conv=None): |
| 141 | super().__init__() |
| 142 | |
| 143 | self.positional_embedding_proj = torch.nn.Sequential( |
| 144 | torch.nn.Linear(in_channels, in_channels * 4), |
| 145 | torch.nn.SiLU(), |
| 146 | torch.nn.Linear(in_channels * 4, in_channels) |
| 147 | ) |
| 148 | if add_positional_conv is not None: |
| 149 | self.positional_embedding = TrainableTemporalTimesteps(in_channels, True, 0, add_positional_conv) |
| 150 | self.positional_conv = torch.nn.Conv3d(in_channels, in_channels, kernel_size=3, padding=1, padding_mode="reflect") |
| 151 | else: |
| 152 | self.positional_embedding = TemporalTimesteps(in_channels, True, 0) |
| 153 | self.positional_conv = None |
| 154 | |
| 155 | self.norm_in = torch.nn.LayerNorm(in_channels) |
| 156 | self.act_fn_in = GEGLU(in_channels, in_channels * 4) |
| 157 | self.ff_in = torch.nn.Linear(in_channels * 4, in_channels) |
| 158 | |
| 159 | self.norm1 = torch.nn.LayerNorm(in_channels) |
| 160 | self.attn1 = Attention( |
| 161 | q_dim=in_channels, |
| 162 | num_heads=num_attention_heads, |
| 163 | head_dim=attention_head_dim, |
| 164 | bias_out=True |
| 165 | ) |
| 166 | |
| 167 | self.norm2 = torch.nn.LayerNorm(in_channels) |
| 168 | self.attn2 = Attention( |
| 169 | q_dim=in_channels, |
| 170 | kv_dim=cross_attention_dim, |
| 171 | num_heads=num_attention_heads, |
| 172 | head_dim=attention_head_dim, |
| 173 | bias_out=True |
| 174 | ) |
| 175 | |
| 176 | self.norm_out = torch.nn.LayerNorm(in_channels) |
| 177 | self.act_fn_out = GEGLU(in_channels, in_channels * 4) |
| 178 | self.ff_out = torch.nn.Linear(in_channels * 4, in_channels) |
| 179 | |
| 180 | def forward(self, hidden_states, time_emb, text_emb, res_stack, **kwargs): |
| 181 |
nothing calls this directly
no test coverage detected