| 737 | |
| 738 | |
| 739 | class SpatialUpsample2x(Module): |
| 740 | def __init__(self, dim, dim_out=None): |
| 741 | super().__init__() |
| 742 | dim_out = default(dim_out, dim) |
| 743 | conv = nn.Conv2d(dim, dim_out * 4, 1) |
| 744 | |
| 745 | self.net = nn.Sequential(conv, nn.SiLU(), Rearrange("b (c p1 p2) h w -> b c (h p1) (w p2)", p1=2, p2=2)) |
| 746 | |
| 747 | self.init_conv_(conv) |
| 748 | |
| 749 | def init_conv_(self, conv): |
| 750 | o, i, h, w = conv.weight.shape |
| 751 | conv_weight = torch.empty(o // 4, i, h, w) |
| 752 | nn.init.kaiming_uniform_(conv_weight) |
| 753 | conv_weight = repeat(conv_weight, "o ... -> (o 4) ...") |
| 754 | |
| 755 | conv.weight.data.copy_(conv_weight) |
| 756 | nn.init.zeros_(conv.bias.data) |
| 757 | |
| 758 | def forward(self, x): |
| 759 | x = rearrange(x, "b c t h w -> b t c h w") |
| 760 | x, ps = pack_one(x, "* c h w") |
| 761 | |
| 762 | out = self.net(x) |
| 763 | |
| 764 | out = unpack_one(out, ps, "* c h w") |
| 765 | out = rearrange(out, "b t c h w -> b c t h w") |
| 766 | return out |
| 767 | |
| 768 | |
| 769 | class TimeUpsample2x(Module): |