(self, dim, num_attention_heads)
| 108 | |
| 109 | class FluxJointTransformerBlock(torch.nn.Module): |
| 110 | def __init__(self, dim, num_attention_heads): |
| 111 | super().__init__() |
| 112 | self.norm1_a = AdaLayerNorm(dim) |
| 113 | self.norm1_b = AdaLayerNorm(dim) |
| 114 | |
| 115 | self.attn = FluxJointAttention(dim, dim, num_attention_heads, dim // num_attention_heads) |
| 116 | |
| 117 | self.norm2_a = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 118 | self.ff_a = torch.nn.Sequential( |
| 119 | torch.nn.Linear(dim, dim*4), |
| 120 | torch.nn.GELU(approximate="tanh"), |
| 121 | torch.nn.Linear(dim*4, dim) |
| 122 | ) |
| 123 | |
| 124 | self.norm2_b = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| 125 | self.ff_b = torch.nn.Sequential( |
| 126 | torch.nn.Linear(dim, dim*4), |
| 127 | torch.nn.GELU(approximate="tanh"), |
| 128 | torch.nn.Linear(dim*4, dim) |
| 129 | ) |
| 130 | |
| 131 | |
| 132 | def forward(self, hidden_states_a, hidden_states_b, temb, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None): |
nothing calls this directly
no test coverage detected