Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| from torch.nn import functional as F | |
| from helperVAE import VAE_AttentionBlock, VAE_ResidualBlock | |
| class VAE_Decoder(nn.Sequential): | |
| def __init__(self): | |
| super().__init__( | |
| # (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 4, Height / 8, Width / 8) | |
| nn.Conv2d(4, 4, kernel_size=1, padding=0), | |
| # (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8) | |
| nn.Conv2d(4, 512, kernel_size=3, padding=1), | |
| # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8) | |
| VAE_ResidualBlock(512, 512), | |
| # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8) | |
| VAE_AttentionBlock(512), | |
| # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8) | |
| VAE_ResidualBlock(512, 512), | |
| # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8) | |
| VAE_ResidualBlock(512, 512), | |
| # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8) | |
| VAE_ResidualBlock(512, 512), | |
| # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8) | |
| VAE_ResidualBlock(512, 512), | |
| # Repeats the rows and columns of the data by scale_factor (like when you resize an image by doubling its size). | |
| # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 4, Width / 4) | |
| nn.Upsample(scale_factor=2), | |
| # (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 4, Width / 4) | |
| nn.Conv2d(512, 512, kernel_size=3, padding=1), | |
| # (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 4, Width / 4) | |
| VAE_ResidualBlock(512, 512), | |
| # (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 4, Width / 4) | |
| VAE_ResidualBlock(512, 512), | |
| # (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 4, Width / 4) | |
| VAE_ResidualBlock(512, 512), | |
| # (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 2, Width / 2) | |
| nn.Upsample(scale_factor=2), | |
| # (Batch_Size, 512, Height / 2, Width / 2) -> (Batch_Size, 512, Height / 2, Width / 2) | |
| nn.Conv2d(512, 512, kernel_size=3, padding=1), | |
| # (Batch_Size, 512, Height / 2, Width / 2) -> (Batch_Size, 256, Height / 2, Width / 2) | |
| VAE_ResidualBlock(512, 256), | |
| # (Batch_Size, 256, Height / 2, Width / 2) -> (Batch_Size, 256, Height / 2, Width / 2) | |
| VAE_ResidualBlock(256, 256), | |
| # (Batch_Size, 256, Height / 2, Width / 2) -> (Batch_Size, 256, Height / 2, Width / 2) | |
| VAE_ResidualBlock(256, 256), | |
| # (Batch_Size, 256, Height / 2, Width / 2) -> (Batch_Size, 256, Height, Width) | |
| nn.Upsample(scale_factor=2), | |
| # (Batch_Size, 256, Height, Width) -> (Batch_Size, 256, Height, Width) | |
| nn.Conv2d(256, 256, kernel_size=3, padding=1), | |
| # (Batch_Size, 256, Height, Width) -> (Batch_Size, 128, Height, Width) | |
| VAE_ResidualBlock(256, 128), | |
| # (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height, Width) | |
| VAE_ResidualBlock(128, 128), | |
| # (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height, Width) | |
| VAE_ResidualBlock(128, 128), | |
| # (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height, Width) | |
| nn.GroupNorm(32, 128), | |
| # (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height, Width) | |
| nn.SiLU(), | |
| # (Batch_Size, 128, Height, Width) -> (Batch_Size, 3, Height, Width) | |
| nn.Conv2d(128, 3, kernel_size=3, padding=1), | |
| ) | |
| def forward(self, x): | |
| # x: (Batch_Size, 4, Height / 8, Width / 8) | |
| # Remove the scaling added by the Encoder. | |
| x /= 0.18215 | |
| for module in self: | |
| x = module(x) | |
| # (Batch_Size, 3, Height, Width) | |
| return x |