(self, dim, num_attention_heads, use_rms_norm=False)
| 196 | |
| 197 | class DualTransformerBlock(torch.nn.Module): |
| 198 | def __init__(self, dim, num_attention_heads, use_rms_norm=False): |
| 199 | super().__init__() |
| 200 | self.norm1_a = AdaLayerNorm(dim, dual=True) |
| 201 | self.norm1_b = AdaLayerNorm(dim) |
| 202 | |
| 203 | self.attn = JointAttention(dim, dim, num_attention_heads, dim // num_attention_heads, use_rms_norm=use_rms_norm) |
| 204 | self.attn2 = JointAttention(dim, dim, num_attention_heads, dim // num_attention_heads, use_rms_norm=use_rms_norm) |
| 205 | |
| 206 | self.norm2_a = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 207 | self.ff_a = torch.nn.Sequential( |
| 208 | torch.nn.Linear(dim, dim*4), |
| 209 | torch.nn.GELU(approximate="tanh"), |
| 210 | torch.nn.Linear(dim*4, dim) |
| 211 | ) |
| 212 | |
| 213 | self.norm2_b = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 214 | self.ff_b = torch.nn.Sequential( |
| 215 | torch.nn.Linear(dim, dim*4), |
| 216 | torch.nn.GELU(approximate="tanh"), |
| 217 | torch.nn.Linear(dim*4, dim) |
| 218 | ) |
| 219 | |
| 220 | |
| 221 | def forward(self, hidden_states_a, hidden_states_b, temb): |
nothing calls this directly
no test coverage detected