(self, dim, out_dim, patch_size, eps=1e-6)
| 454 | class Head(nn.Module): |
| 455 | |
| 456 | def __init__(self, dim, out_dim, patch_size, eps=1e-6): |
| 457 | super().__init__() |
| 458 | self.dim = dim |
| 459 | self.out_dim = out_dim |
| 460 | self.patch_size = patch_size |
| 461 | self.eps = eps |
| 462 | |
| 463 | # layers |
| 464 | out_dim = math.prod(patch_size) * out_dim |
| 465 | self.norm = WanLayerNorm(dim, eps) |
| 466 | self.head = nn.Linear(dim, out_dim) |
| 467 | |
| 468 | # modulation |
| 469 | self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) |
| 470 | |
| 471 | def forward(self, x, e): |
| 472 | assert e.dtype == torch.float32 |