(self, dim, dim_cond, num_heads)
| 44 | |
| 45 | class CogDiTBlock(torch.nn.Module): |
| 46 | def __init__(self, dim, dim_cond, num_heads): |
| 47 | super().__init__() |
| 48 | self.norm1 = CogAdaLayerNorm(dim, dim_cond) |
| 49 | self.attn1 = Attention(q_dim=dim, num_heads=48, head_dim=dim//num_heads, bias_q=True, bias_kv=True, bias_out=True) |
| 50 | self.norm_q = torch.nn.LayerNorm((dim//num_heads,), eps=1e-06, elementwise_affine=True) |
| 51 | self.norm_k = torch.nn.LayerNorm((dim//num_heads,), eps=1e-06, elementwise_affine=True) |
| 52 | |
| 53 | self.norm2 = CogAdaLayerNorm(dim, dim_cond) |
| 54 | self.ff = torch.nn.Sequential( |
| 55 | torch.nn.Linear(dim, dim*4), |
| 56 | torch.nn.GELU(approximate="tanh"), |
| 57 | torch.nn.Linear(dim*4, dim) |
| 58 | ) |
| 59 | |
| 60 | |
| 61 | def apply_rotary_emb(self, x, freqs_cis): |
nothing calls this directly
no test coverage detected