ApoorvBrooklyn/stable-diffusion-implementation

1import torch2from torch import nn3from torch.nn import functional as F4from attention import SelfAttention5 6class VAE_AttentionBlock(nn.Module):7    def __init__(self, channels):8        super().__init__()9        self.groupnorm = nn.GroupNorm(32, channels)10        self.attention = SelfAttention(1, channels)11    12    def forward(self, x):13        # x: (Batch_Size, Features, Height, Width)14 15        residue = x 16 17        # (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height, Width)18        x = self.groupnorm(x)19 20        n, c, h, w = x.shape21        22        # (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height * Width)23        x = x.view((n, c, h * w))24        25        # (Batch_Size, Features, Height * Width) -> (Batch_Size, Height * Width, Features). Each pixel becomes a feature of size "Features", the sequence length is "Height * Width".26        x = x.transpose(-1, -2)27        28        # Perform self-attention WITHOUT mask29        # (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)30        x = self.attention(x)31        32        # (Batch_Size, Height * Width, Features) -> (Batch_Size, Features, Height * Width)33        x = x.transpose(-1, -2)34        35        # (Batch_Size, Features, Height * Width) -> (Batch_Size, Features, Height, Width)36        x = x.view((n, c, h, w))37        38        # (Batch_Size, Features, Height, Width) + (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height, Width) 39        x += residue40 41        # (Batch_Size, Features, Height, Width)42        return x 43 44class VAE_ResidualBlock(nn.Module):45    def __init__(self, in_channels, out_channels):46        super().__init__()47        self.groupnorm_1 = nn.GroupNorm(32, in_channels)48        self.conv_1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)49 50        self.groupnorm_2 = nn.GroupNorm(32, out_channels)51        self.conv_2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)52 53        if in_channels == out_channels:54            self.residual_layer = nn.Identity()55        else:56            self.residual_layer = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)57    58    def forward(self, x):59        # x: (Batch_Size, In_Channels, Height, Width)60 61        residue = x62 63        # (Batch_Size, In_Channels, Height, Width) -> (Batch_Size, In_Channels, Height, Width)64        x = self.groupnorm_1(x)65        66        # (Batch_Size, In_Channels, Height, Width) -> (Batch_Size, In_Channels, Height, Width)67        x = F.silu(x)68        69        # (Batch_Size, In_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)70        x = self.conv_1(x)71        72        # (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)73        x = self.groupnorm_2(x)74        75        # (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)76        x = F.silu(x)77        78        # (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)79        x = self.conv_2(x)80        81        # (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)82        return x + self.residual_layer(residue)83 84class VAE_Decoder(nn.Sequential):85    def __init__(self):86        super().__init__(87            # (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 4, Height / 8, Width / 8)88            nn.Conv2d(4, 4, kernel_size=1, padding=0),89 90            # (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)91            nn.Conv2d(4, 512, kernel_size=3, padding=1),92            93            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)94            VAE_ResidualBlock(512, 512), 95            96            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)97            VAE_AttentionBlock(512), 98            99            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)100            VAE_ResidualBlock(512, 512), 101            102            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)103            VAE_ResidualBlock(512, 512), 104            105            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)106            VAE_ResidualBlock(512, 512), 107            108            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)109            VAE_ResidualBlock(512, 512), 110            111            # Repeats the rows and columns of the data by scale_factor (like when you resize an image by doubling its size).112            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 4, Width / 4)113            nn.Upsample(scale_factor=2),114            115            # (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 4, Width / 4)116            nn.Conv2d(512, 512, kernel_size=3, padding=1), 117            118            # (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 4, Width / 4)119            VAE_ResidualBlock(512, 512), 120            121            # (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 4, Width / 4)122            VAE_ResidualBlock(512, 512), 123            124            # (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 4, Width / 4)125            VAE_ResidualBlock(512, 512), 126            127            # (Batch_Size, 512, Height / 4, Width / 4) -> (Batch_Size, 512, Height / 2, Width / 2)128            nn.Upsample(scale_factor=2), 129            130            # (Batch_Size, 512, Height / 2, Width / 2) -> (Batch_Size, 512, Height / 2, Width / 2)131            nn.Conv2d(512, 512, kernel_size=3, padding=1), 132            133            # (Batch_Size, 512, Height / 2, Width / 2) -> (Batch_Size, 256, Height / 2, Width / 2)134            VAE_ResidualBlock(512, 256), 135            136            # (Batch_Size, 256, Height / 2, Width / 2) -> (Batch_Size, 256, Height / 2, Width / 2)137            VAE_ResidualBlock(256, 256), 138            139            # (Batch_Size, 256, Height / 2, Width / 2) -> (Batch_Size, 256, Height / 2, Width / 2)140            VAE_ResidualBlock(256, 256), 141            142            # (Batch_Size, 256, Height / 2, Width / 2) -> (Batch_Size, 256, Height, Width)143            nn.Upsample(scale_factor=2), 144            145            # (Batch_Size, 256, Height, Width) -> (Batch_Size, 256, Height, Width)146            nn.Conv2d(256, 256, kernel_size=3, padding=1), 147            148            # (Batch_Size, 256, Height, Width) -> (Batch_Size, 128, Height, Width)149            VAE_ResidualBlock(256, 128), 150            151            # (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height, Width)152            VAE_ResidualBlock(128, 128), 153            154            # (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height, Width)155            VAE_ResidualBlock(128, 128), 156            157            # (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height, Width)158            nn.GroupNorm(32, 128), 159            160            # (Batch_Size, 128, Height, Width) -> (Batch_Size, 128, Height, Width)161            nn.SiLU(), 162            163            # (Batch_Size, 128, Height, Width) -> (Batch_Size, 3, Height, Width)164            nn.Conv2d(128, 3, kernel_size=3, padding=1), 165        )166 167    def forward(self, x):168        # x: (Batch_Size, 4, Height / 8, Width / 8)169        170        # Remove the scaling added by the Encoder.171        x /= 0.18215172 173        for module in self:174            x = module(x)175 176        # (Batch_Size, 3, Height, Width)177        return x178