import torch import torch.nn as nn import torch.nn.functional as F from attention import SelfAttention class VAE_AttentionBlock(nn.Module): def __init__(self, channels): super().__init__() self.groupnorm = nn.GroupNorm(32, channels) self.attention = SelfAttention(1, channels) def forward(self, x): # x: (Batch_Size, Features, Height, Width) residue = x # (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height, Width) x = self.groupnorm(x) n, c, h, w = x.shape # (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height * Width) x = x.view((n, c, h * w)) # (Batch_Size, Features, Height * Width) -> (Batch_Size, Height * Width, Features). Each pixel becomes a feature of size "Features", the sequence length is "Height * Width". x = x.transpose(-1, -2) # Perform self-attention WITHOUT mask # (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features) x = self.attention(x) # (Batch_Size, Height * Width, Features) -> (Batch_Size, Features, Height * Width) x = x.transpose(-1, -2) # (Batch_Size, Features, Height * Width) -> (Batch_Size, Features, Height, Width) x = x.view((n, c, h, w)) # (Batch_Size, Features, Height, Width) + (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height, Width) x += residue # (Batch_Size, Features, Height, Width) return x class VAE_ResidualBlock(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.groupnorm_1 = nn.GroupNorm(32, in_channels) self.conv_1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) self.groupnorm_2 = nn.GroupNorm(32, out_channels) self.conv_2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1) if in_channels == out_channels: self.residual_layer = nn.Identity() else: self.residual_layer = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0) def forward(self, x): # x: (Batch_Size, In_Channels, Height, Width) residue = x # (Batch_Size, In_Channels, Height, Width) -> (Batch_Size, In_Channels, Height, Width) x = self.groupnorm_1(x) # (Batch_Size, In_Channels, Height, Width) -> (Batch_Size, In_Channels, Height, Width) x = F.silu(x) # (Batch_Size, In_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width) x = self.conv_1(x) # (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width) x = self.groupnorm_2(x) # (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width) x = F.silu(x) # (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width) x = self.conv_2(x) # (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width) return x + self.residual_layer(residue) 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