| 26 | |
| 27 | |
| 28 | class LatentTransparencyOffsetEncoder(torch.nn.Module): |
| 29 | def __init__(self, latent_c=4, *args, **kwargs): |
| 30 | super().__init__(*args, **kwargs) |
| 31 | self.blocks = torch.nn.Sequential( |
| 32 | torch.nn.Conv2d(4, 32, kernel_size=3, padding=1, stride=1), |
| 33 | nn.SiLU(), |
| 34 | torch.nn.Conv2d(32, 32, kernel_size=3, padding=1, stride=1), |
| 35 | nn.SiLU(), |
| 36 | torch.nn.Conv2d(32, 64, kernel_size=3, padding=1, stride=2), |
| 37 | nn.SiLU(), |
| 38 | torch.nn.Conv2d(64, 64, kernel_size=3, padding=1, stride=1), |
| 39 | nn.SiLU(), |
| 40 | torch.nn.Conv2d(64, 128, kernel_size=3, padding=1, stride=2), |
| 41 | nn.SiLU(), |
| 42 | torch.nn.Conv2d(128, 128, kernel_size=3, padding=1, stride=1), |
| 43 | nn.SiLU(), |
| 44 | torch.nn.Conv2d(128, 256, kernel_size=3, padding=1, stride=2), |
| 45 | nn.SiLU(), |
| 46 | torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1), |
| 47 | nn.SiLU(), |
| 48 | zero_module(torch.nn.Conv2d(256, latent_c, kernel_size=3, padding=1, stride=1)), |
| 49 | ) |
| 50 | |
| 51 | def __call__(self, x): |
| 52 | return self.blocks(x) |
| 53 | |
| 54 | |
| 55 | # 1024 * 1024 * 3 -> 16 * 16 * 512 -> 1024 * 1024 * 3 |