(self, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6)
| 439 | |
| 440 | class DiTBlock(nn.Module): |
| 441 | def __init__(self, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6): |
| 442 | super().__init__() |
| 443 | self.dim = dim |
| 444 | self.num_heads = num_heads |
| 445 | self.ffn_dim = ffn_dim |
| 446 | |
| 447 | self.self_attn = SelfAttention(dim, num_heads, eps) |
| 448 | self.cross_attn = CrossAttention(dim, num_heads, eps) |
| 449 | |
| 450 | self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) |
| 451 | self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) |
| 452 | self.norm3 = nn.LayerNorm(dim, eps=eps) |
| 453 | self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU( |
| 454 | approximate='tanh'), nn.Linear(ffn_dim, dim)) |
| 455 | self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) |
| 456 | self.gate = GateModule() |
| 457 | |
| 458 | def forward(self, x, context, t_mod, freqs, f, h, w, local_num=None, topk=None, |
| 459 | train_img=False, block_id=None, kv_len=None, is_full_block=False, |
nothing calls this directly
no test coverage detected