| 497 | |
| 498 | class Head(nn.Module): |
| 499 | def __init__(self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float): |
| 500 | super().__init__() |
| 501 | self.dim = dim |
| 502 | self.patch_size = patch_size |
| 503 | self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) |
| 504 | self.head = nn.Linear(dim, out_dim * math.prod(patch_size)) |
| 505 | self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) |
| 506 | |
| 507 | def forward(self, x, t_mod): |
| 508 | shift, scale = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(2, dim=1) |