| 767 | |
| 768 | |
| 769 | class TimeUpsample2x(Module): |
| 770 | def __init__(self, dim, dim_out=None): |
| 771 | super().__init__() |
| 772 | dim_out = default(dim_out, dim) |
| 773 | conv = nn.Conv1d(dim, dim_out * 2, 1) |
| 774 | |
| 775 | self.net = nn.Sequential(conv, nn.SiLU(), Rearrange("b (c p) t -> b c (t p)", p=2)) |
| 776 | |
| 777 | self.init_conv_(conv) |
| 778 | |
| 779 | def init_conv_(self, conv): |
| 780 | o, i, t = conv.weight.shape |
| 781 | conv_weight = torch.empty(o // 2, i, t) |
| 782 | nn.init.kaiming_uniform_(conv_weight) |
| 783 | conv_weight = repeat(conv_weight, "o ... -> (o 2) ...") |
| 784 | |
| 785 | conv.weight.data.copy_(conv_weight) |
| 786 | nn.init.zeros_(conv.bias.data) |
| 787 | |
| 788 | def forward(self, x): |
| 789 | x = rearrange(x, "b c t h w -> b h w c t") |
| 790 | x, ps = pack_one(x, "* c t") |
| 791 | |
| 792 | out = self.net(x) |
| 793 | |
| 794 | out = unpack_one(out, ps, "* c t") |
| 795 | out = rearrange(out, "b h w c t -> b c t h w") |
| 796 | return out |
| 797 | |
| 798 | |
| 799 | # autoencoder - only best variant here offered, with causal conv 3d |