ApoorvBrooklyn/stable-diffusion-implementation

1import torch2from torch import nn3from torch.nn import functional as F4from decoder import VAE_AttentionBlock, VAE_ResidualBlock5 6class VAE_Encoder(nn.Sequential):7    def __init__(self):8        super().__init__(9            # (Batch_Size, Channel, Height, Width) -> (Batch_Size, 128, Height, Width)10            nn.Conv2d(3, 128, kernel_size=3, padding=1),11            12             # (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height, Width)13            VAE_ResidualBlock(128, 128),14            15            # (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height, Width)16            VAE_ResidualBlock(128, 128),17            18            # (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height / 2, Width / 2)19            nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=0),20            21            # (Batch_Size, 128, Height / 2, Width / 2) -> (Batch_Size, 256, Height / 2, Width / 2)22            VAE_ResidualBlock(128, 256), 23            24            # (Batch_Size, 256, Height / 2, Width / 2) -> (Batch_Size, 256, Height / 2, Width / 2)25            VAE_ResidualBlock(256, 256), 26            27            # (Batch_Size, 256, Height / 2, Width / 2) -> (Batch_Size, 256, Height / 4, Width / 4)28            nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=0), 29            30            # (Batch_Size, 256, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 4, Width / 4)31            VAE_ResidualBlock(256, 512), 32            33            # (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 4, Width / 4)34            VAE_ResidualBlock(512, 512), 35            36            # (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 8, Width / 8)37            nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=0), 38            39            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)40            VAE_ResidualBlock(512, 512), 41            42            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)43            VAE_ResidualBlock(512, 512), 44            45            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)46            VAE_ResidualBlock(512, 512), 47            48            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)49            VAE_AttentionBlock(512), 50            51            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)52            VAE_ResidualBlock(512, 512), 53            54            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)55            nn.GroupNorm(32, 512), 56            57            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)58            nn.SiLU(), 59 60            # Because the padding=1, it means the width and height will increase by 261            # Out_Height = In_Height + Padding_Top + Padding_Bottom62            # Out_Width = In_Width + Padding_Left + Padding_Right63            # Since padding = 1 means Padding_Top = Padding_Bottom = Padding_Left = Padding_Right = 1,64            # Since the Out_Width = In_Width + 2 (same for Out_Height), it will compensate for the Kernel size of 365            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 8, Height / 8, Width / 8). 66            nn.Conv2d(512, 8, kernel_size=3, padding=1), 67 68            # (Batch_Size, 8, Height / 8, Width / 8) -> (Batch_Size, 8, Height / 8, Width / 8)69            nn.Conv2d(8, 8, kernel_size=1, padding=0), 70        )71 72    def forward(self, x, noise):73        # x: (Batch_Size, Channel, Height, Width)74        # noise: (Batch_Size, 4, Height / 8, Width / 8)75 76        for module in self:77 78            if getattr(module, 'stride', None) == (2, 2):  # Padding at downsampling should be asymmetric (see #8)79                # Pad: (Padding_Left, Padding_Right, Padding_Top, Padding_Bottom).80                # Pad with zeros on the right and bottom.81                # (Batch_Size, Channel, Height, Width) -> (Batch_Size, Channel, Height + Padding_Top + Padding_Bottom, Width + Padding_Left + Padding_Right) = (Batch_Size, Channel, Height + 1, Width + 1)82                x = F.pad(x, (0, 1, 0, 1))83            84            x = module(x)85        # (Batch_Size, 8, Height / 8, Width / 8) -> two tensors of shape (Batch_Size, 4, Height / 8, Width / 8)86        mean, log_variance = torch.chunk(x, 2, dim=1)87        # Clamp the log variance between -30 and 20, so that the variance is between (circa) 1e-14 and 1e8. 88        # (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 4, Height / 8, Width / 8)89        log_variance = torch.clamp(log_variance, -30, 20)90        # (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 4, Height / 8, Width / 8)91        variance = log_variance.exp()92        # (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 4, Height / 8, Width / 8)93        stdev = variance.sqrt()94        95        # Transform N(0, 1) -> N(mean, stdev) 96        # (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 4, Height / 8, Width / 8)97        x = mean + stdev * noise98        99        # Scale by a constant100        # Constant taken from: https://github.com/CompVis/stable-diffusion/blob/21f890f9da3cfbeaba8e2ac3c425ee9e998d5229/configs/stable-diffusion/v1-inference.yaml#L17C1-L17C1101        x *= 0.18215102        103        return x