| 494 | |
| 495 | |
| 496 | class Head(nn.Module): |
| 497 | |
| 498 | def __init__(self, dim, out_dim, patch_size, eps=1e-6): |
| 499 | super().__init__() |
| 500 | self.dim = dim |
| 501 | self.out_dim = out_dim |
| 502 | self.patch_size = patch_size |
| 503 | self.eps = eps |
| 504 | |
| 505 | # layers |
| 506 | out_dim = math.prod(patch_size) * out_dim |
| 507 | self.norm = WanLayerNorm(dim, eps) |
| 508 | self.head = nn.Linear(dim, out_dim) |
| 509 | |
| 510 | # modulation |
| 511 | self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) |
| 512 | |
| 513 | def forward(self, x, e): |
| 514 | assert e.dtype == torch.float32 |
| 515 | with amp.autocast(dtype=torch.float32, device_type="cuda"): |
| 516 | e = (self.modulation.to(dtype=e.dtype, device=e.device) + e.unsqueeze(1)).chunk(2, dim=1) |
| 517 | x = (self.head(self.norm(x) * (1 + e[1]) + e[0])) |
| 518 | return x |
| 519 | |
| 520 | |
| 521 | class MLPProj(torch.nn.Module): |