(self, dim, num_attention_heads)
| 205 | |
| 206 | class FluxSingleTransformerBlock(torch.nn.Module): |
| 207 | def __init__(self, dim, num_attention_heads): |
| 208 | super().__init__() |
| 209 | self.num_heads = num_attention_heads |
| 210 | self.head_dim = dim // num_attention_heads |
| 211 | self.dim = dim |
| 212 | |
| 213 | self.norm = AdaLayerNormSingle(dim) |
| 214 | self.to_qkv_mlp = torch.nn.Linear(dim, dim * (3 + 4)) |
| 215 | self.norm_q_a = RMSNorm(self.head_dim, eps=1e-6) |
| 216 | self.norm_k_a = RMSNorm(self.head_dim, eps=1e-6) |
| 217 | |
| 218 | self.proj_out = torch.nn.Linear(dim * 5, dim) |
| 219 | |
| 220 | |
| 221 | def apply_rope(self, xq, xk, freqs_cis): |
nothing calls this directly
no test coverage detected