| 210 | |
| 211 | |
| 212 | class GateModule(nn.Module): |
| 213 | def __init__(self,): |
| 214 | super().__init__() |
| 215 | |
| 216 | def forward(self, x, gate, residual): |
| 217 | return x + gate * residual |
| 218 | |
| 219 | class DiTBlock(nn.Module): |
| 220 | def __init__(self, has_image_input: bool, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6): |