ApoorvBrooklyn/stable-diffusion-implementation
0
1import torch2from torch import nn3from torch.nn import functional as F4from attention import SelfAttention, CrossAttention5 6class TimeEmbedding(nn.Module):7 def __init__(self, n_embd):8 super().__init__()9 self.linear_1 = nn.Linear(n_embd, 4 * n_embd)10 self.linear_2 = nn.Linear(4 * n_embd, 4 * n_embd)11 12 def forward(self, x):13 # x: (1, 320)14 15 # (1, 320) -> (1, 1280)16 x = self.linear_1(x)17 18 # (1, 1280) -> (1, 1280)19 x = F.silu(x) 20 21 # (1, 1280) -> (1, 1280)22 x = self.linear_2(x)23 24 return x25 26class UNET_ResidualBlock(nn.Module):27 def __init__(self, in_channels, out_channels, n_time=1280):28 super().__init__()29 self.groupnorm_feature = nn.GroupNorm(32, in_channels)30 self.conv_feature = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)31 self.linear_time = nn.Linear(n_time, out_channels)32 33 self.groupnorm_merged = nn.GroupNorm(32, out_channels)34 self.conv_merged = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)35 36 if in_channels == out_channels:37 self.residual_layer = nn.Identity()38 else:39 self.residual_layer = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)40 41 def forward(self, feature, time):42 # feature: (Batch_Size, In_Channels, Height, Width)43 # time: (1, 1280)44 45 residue = feature46 47 # (Batch_Size, In_Channels, Height, Width) -> (Batch_Size, In_Channels, Height, Width)48 feature = self.groupnorm_feature(feature)49 50 # (Batch_Size, In_Channels, Height, Width) -> (Batch_Size, In_Channels, Height, Width)51 feature = F.silu(feature)52 53 # (Batch_Size, In_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)54 feature = self.conv_feature(feature)55 56 # (1, 1280) -> (1, 1280)57 time = F.silu(time)58 59 # (1, 1280) -> (1, Out_Channels)60 time = self.linear_time(time)61 62 # Add width and height dimension to time. 63 # (Batch_Size, Out_Channels, Height, Width) + (1, Out_Channels, 1, 1) -> (Batch_Size, Out_Channels, Height, Width)64 merged = feature + time.unsqueeze(-1).unsqueeze(-1)65 66 # (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)67 merged = self.groupnorm_merged(merged)68 69 # (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)70 merged = F.silu(merged)71 72 # (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)73 merged = self.conv_merged(merged)74 75 # (Batch_Size, Out_Channels, Height, Width) + (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)76 return merged + self.residual_layer(residue)77 78class UNET_AttentionBlock(nn.Module):79 def __init__(self, n_head: int, n_embd: int, d_context=768):80 super().__init__()81 channels = n_head * n_embd82 83 self.groupnorm = nn.GroupNorm(32, channels, eps=1e-6)84 self.conv_input = nn.Conv2d(channels, channels, kernel_size=1, padding=0)85 86 self.layernorm_1 = nn.LayerNorm(channels)87 self.attention_1 = SelfAttention(n_head, channels, in_proj_bias=False)88 self.layernorm_2 = nn.LayerNorm(channels)89 self.attention_2 = CrossAttention(n_head, channels, d_context, in_proj_bias=False)90 self.layernorm_3 = nn.LayerNorm(channels)91 self.linear_geglu_1 = nn.Linear(channels, 4 * channels * 2)92 self.linear_geglu_2 = nn.Linear(4 * channels, channels)93 94 self.conv_output = nn.Conv2d(channels, channels, kernel_size=1, padding=0)95 96 def forward(self, x, context):97 # x: (Batch_Size, Features, Height, Width)98 # context: (Batch_Size, Seq_Len, Dim)99 100 residue_long = x101 102 # (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height, Width)103 x = self.groupnorm(x)104 105 # (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height, Width)106 x = self.conv_input(x)107 108 n, c, h, w = x.shape109 110 # (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height * Width)111 x = x.view((n, c, h * w))112 113 # (Batch_Size, Features, Height * Width) -> (Batch_Size, Height * Width, Features)114 x = x.transpose(-1, -2)115 116 # Normalization + Self-Attention with skip connection117 118 # (Batch_Size, Height * Width, Features)119 residue_short = x120 121 # (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)122 x = self.layernorm_1(x)123 124 # (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)125 x = self.attention_1(x)126 127 # (Batch_Size, Height * Width, Features) + (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)128 x += residue_short129 130 # (Batch_Size, Height * Width, Features)131 residue_short = x132 133 # Normalization + Cross-Attention with skip connection134 135 # (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)136 x = self.layernorm_2(x)137 138 # (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)139 x = self.attention_2(x, context)140 141 # (Batch_Size, Height * Width, Features) + (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)142 x += residue_short143 144 # (Batch_Size, Height * Width, Features)145 residue_short = x146 147 # Normalization + FFN with GeGLU and skip connection148 149 # (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)150 x = self.layernorm_3(x)151 152 # GeGLU as implemented in the original code: https://github.com/CompVis/stable-diffusion/blob/21f890f9da3cfbeaba8e2ac3c425ee9e998d5229/ldm/modules/attention.py#L37C10-L37C10153 # (Batch_Size, Height * Width, Features) -> two tensors of shape (Batch_Size, Height * Width, Features * 4)154 x, gate = self.linear_geglu_1(x).chunk(2, dim=-1) 155 156 # Element-wise product: (Batch_Size, Height * Width, Features * 4) * (Batch_Size, Height * Width, Features * 4) -> (Batch_Size, Height * Width, Features * 4)157 x = x * F.gelu(gate)158 159 # (Batch_Size, Height * Width, Features * 4) -> (Batch_Size, Height * Width, Features)160 x = self.linear_geglu_2(x)161 162 # (Batch_Size, Height * Width, Features) + (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)163 x += residue_short164 165 # (Batch_Size, Height * Width, Features) -> (Batch_Size, Features, Height * Width)166 x = x.transpose(-1, -2)167 168 # (Batch_Size, Features, Height * Width) -> (Batch_Size, Features, Height, Width)169 x = x.view((n, c, h, w))170 171 # Final skip connection between initial input and output of the block172 # (Batch_Size, Features, Height, Width) + (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height, Width)173 return self.conv_output(x) + residue_long174 175class Upsample(nn.Module):176 def __init__(self, channels):177 super().__init__()178 self.conv = nn.Conv2d(channels, channels, kernel_size=3, padding=1)179 180 def forward(self, x):181 # (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height * 2, Width * 2)182 x = F.interpolate(x, scale_factor=2, mode='nearest') 183 return self.conv(x)184 185class SwitchSequential(nn.Sequential):186 def forward(self, x, context, time):187 for layer in self:188 if isinstance(layer, UNET_AttentionBlock):189 x = layer(x, context)190 elif isinstance(layer, UNET_ResidualBlock):191 x = layer(x, time)192 else:193 x = layer(x)194 return x195 196class UNET(nn.Module):197 def __init__(self):198 super().__init__()199 self.encoders = nn.ModuleList([200 # (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)201 SwitchSequential(nn.Conv2d(4, 320, kernel_size=3, padding=1)),202 203 # (Batch_Size, 320, Height / 8, Width / 8) -> # (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)204 SwitchSequential(UNET_ResidualBlock(320, 320), UNET_AttentionBlock(8, 40)),205 206 # (Batch_Size, 320, Height / 8, Width / 8) -> # (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)207 SwitchSequential(UNET_ResidualBlock(320, 320), UNET_AttentionBlock(8, 40)),208 209 # (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 16, Width / 16)210 SwitchSequential(nn.Conv2d(320, 320, kernel_size=3, stride=2, padding=1)),211 212 # (Batch_Size, 320, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16)213 SwitchSequential(UNET_ResidualBlock(320, 640), UNET_AttentionBlock(8, 80)),214 215 # (Batch_Size, 640, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16)216 SwitchSequential(UNET_ResidualBlock(640, 640), UNET_AttentionBlock(8, 80)),217 218 # (Batch_Size, 640, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 32, Width / 32)219 SwitchSequential(nn.Conv2d(640, 640, kernel_size=3, stride=2, padding=1)),220 221 # (Batch_Size, 640, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32)222 SwitchSequential(UNET_ResidualBlock(640, 1280), UNET_AttentionBlock(8, 160)),223 224 # (Batch_Size, 1280, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32)225 SwitchSequential(UNET_ResidualBlock(1280, 1280), UNET_AttentionBlock(8, 160)),226 227 # (Batch_Size, 1280, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 64, Width / 64)228 SwitchSequential(nn.Conv2d(1280, 1280, kernel_size=3, stride=2, padding=1)),229 230 # (Batch_Size, 1280, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64)231 SwitchSequential(UNET_ResidualBlock(1280, 1280)),232 233 # (Batch_Size, 1280, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64)234 SwitchSequential(UNET_ResidualBlock(1280, 1280)),235 ])236 237 self.bottleneck = SwitchSequential(238 # (Batch_Size, 1280, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64)239 UNET_ResidualBlock(1280, 1280), 240 241 # (Batch_Size, 1280, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64)242 UNET_AttentionBlock(8, 160), 243 244 # (Batch_Size, 1280, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64)245 UNET_ResidualBlock(1280, 1280), 246 )247 248 self.decoders = nn.ModuleList([249 # (Batch_Size, 2560, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64)250 SwitchSequential(UNET_ResidualBlock(2560, 1280)),251 252 # (Batch_Size, 2560, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64)253 SwitchSequential(UNET_ResidualBlock(2560, 1280)),254 255 # (Batch_Size, 2560, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 32, Width / 32) 256 SwitchSequential(UNET_ResidualBlock(2560, 1280), Upsample(1280)),257 258 # (Batch_Size, 2560, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32)259 SwitchSequential(UNET_ResidualBlock(2560, 1280), UNET_AttentionBlock(8, 160)),260 261 # (Batch_Size, 2560, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32)262 SwitchSequential(UNET_ResidualBlock(2560, 1280), UNET_AttentionBlock(8, 160)),263 264 # (Batch_Size, 1920, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 16, Width / 16)265 SwitchSequential(UNET_ResidualBlock(1920, 1280), UNET_AttentionBlock(8, 160), Upsample(1280)),266 267 # (Batch_Size, 1920, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16)268 SwitchSequential(UNET_ResidualBlock(1920, 640), UNET_AttentionBlock(8, 80)),269 270 # (Batch_Size, 1280, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16)271 SwitchSequential(UNET_ResidualBlock(1280, 640), UNET_AttentionBlock(8, 80)),272 273 # (Batch_Size, 960, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 8, Width / 8)274 SwitchSequential(UNET_ResidualBlock(960, 640), UNET_AttentionBlock(8, 80), Upsample(640)),275 276 # (Batch_Size, 960, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)277 SwitchSequential(UNET_ResidualBlock(960, 320), UNET_AttentionBlock(8, 40)),278 279 # (Batch_Size, 640, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)280 SwitchSequential(UNET_ResidualBlock(640, 320), UNET_AttentionBlock(8, 40)),281 282 # (Batch_Size, 640, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)283 SwitchSequential(UNET_ResidualBlock(640, 320), UNET_AttentionBlock(8, 40)),284 ])285 286 def forward(self, x, context, time):287 # x: (Batch_Size, 4, Height / 8, Width / 8)288 # context: (Batch_Size, Seq_Len, Dim) 289 # time: (1, 1280)290 291 skip_connections = []292 for layers in self.encoders:293 x = layers(x, context, time)294 skip_connections.append(x)295 296 x = self.bottleneck(x, context, time)297 298 for layers in self.decoders:299 # Since we always concat with the skip connection of the encoder, the number of features increases before being sent to the decoder's layer300 x = torch.cat((x, skip_connections.pop()), dim=1) 301 x = layers(x, context, time)302 303 return x304 305 306class UNET_OutputLayer(nn.Module):307 def __init__(self, in_channels, out_channels):308 super().__init__()309 self.groupnorm = nn.GroupNorm(32, in_channels)310 self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)311 312 def forward(self, x):313 # x: (Batch_Size, 320, Height / 8, Width / 8)314 315 # (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)316 x = self.groupnorm(x)317 318 # (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)319 x = F.silu(x)320 321 # (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 4, Height / 8, Width / 8)322 x = self.conv(x)323 324 # (Batch_Size, 4, Height / 8, Width / 8) 325 return x326 327class Diffusion(nn.Module):328 def __init__(self):329 super().__init__()330 self.time_embedding = TimeEmbedding(320)331 self.unet = UNET()332 self.final = UNET_OutputLayer(320, 4)333 334 def forward(self, latent, context, time):335 # latent: (Batch_Size, 4, Height / 8, Width / 8)336 # context: (Batch_Size, Seq_Len, Dim)337 # time: (1, 320)338 339 # (1, 320) -> (1, 1280)340 time = self.time_embedding(time)341 342 # (Batch, 4, Height / 8, Width / 8) -> (Batch, 320, Height / 8, Width / 8)343 output = self.unet(latent, context, time)344 345 # (Batch, 320, Height / 8, Width / 8) -> (Batch, 4, Height / 8, Width / 8)346 output = self.final(output)347 348 # (Batch, 4, Height / 8, Width / 8)349 return output