| 38 | return x * weights |
| 39 | |
| 40 | class ChannelEmbeddingCompression(nn.Module): |
| 41 | def __init__(self, embed_dim, embed_dim_next): |
| 42 | super().__init__() |
| 43 | self.patch_unembed = PatchUnEmbedIR(embed_dim=embed_dim) |
| 44 | self.conv = nn.Conv2d(embed_dim, embed_dim_next, 1, 1, 0) |
| 45 | self.patch_embed = PatchEmbedIR(embed_dim=embed_dim_next) |
| 46 | |
| 47 | def forward(self, x, x_size): |
| 48 | x = self.patch_unembed(x, x_size) |
| 49 | x = self.conv(x) |
| 50 | x = self.patch_embed(x) |
| 51 | return x |
| 52 | |
| 53 | class InvertedConvolution(nn.Module): |
| 54 | def __init__(self, in_channels, out_channels, kernel_size, padding='same', bias=True): |