| 619 | return rearrange(x, "... h w (nh nw e) -> ... (h nh) (w nw) e", nh=self.h, nw=self.w) |
| 620 | |
| 621 | class TokenSplit(nn.Module): |
| 622 | def __init__(self, in_features, out_features, patch_size=(2, 2)): |
| 623 | super().__init__() |
| 624 | self.h = patch_size[0] |
| 625 | self.w = patch_size[1] |
| 626 | self.proj = apply_wd(Linear(in_features, out_features * self.h * self.w, bias=False)) |
| 627 | self.fac = nn.Parameter(torch.ones(1) * 0.5) |
| 628 | |
| 629 | def forward(self, x, skip): |
| 630 | # print("x before proj in Tokensplit is ", x.mean(), x.std()) |
| 631 | x = self.proj(x) |
| 632 | # print("x after proj in Tokensplit is ", x.mean(), x.std()) |
| 633 | x = rearrange(x, "... h w (nh nw e) -> ... (h nh) (w nw) e", nh=self.h, nw=self.w) |
| 634 | # print("x after rearange in Tokensplit is ", x.mean(), x.std()) |
| 635 | x = torch.lerp(skip, x, self.fac.to(x.dtype)) |
| 636 | # x = mp_sum_t(x, skip, self.fac.to(x.dtype)) |
| 637 | # print("x after lerp in Tokensplit is ", x.mean(), x.std()) |
| 638 | return x |
| 639 | |
| 640 | |
| 641 | class LocalCondProj(nn.Module): |