Upload 31 files

Implemented and trained a stable diffusion model from scratch (CLIP + VAE + UNet + Cross-Attention), optimizing schedulers (DDPM, DDIM, Euler Ancestral, DPM-Solver++) and attention mechanisms (Flash Attention, xFormers) to reduce generation time by 22% and improve image clarity using MPS on RunPod GPUs.

.gitattributes

@@ -33,3 +33,17 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+pytorch-stablediffusion/images/ddim_images/1_82.png filter=lfs diff=lfs merge=lfs -text
+pytorch-stablediffusion/images/ddim_images/3_112.png filter=lfs diff=lfs merge=lfs -text
+pytorch-stablediffusion/images/ddim_images/4_80.png filter=lfs diff=lfs merge=lfs -text
+pytorch-stablediffusion/images/ddpm_images/1_1022.png filter=lfs diff=lfs merge=lfs -text
+pytorch-stablediffusion/images/ddpm_images/3_114.png filter=lfs diff=lfs merge=lfs -text
+pytorch-stablediffusion/images/ddpm_images/4_333.png filter=lfs diff=lfs merge=lfs -text
+pytorch-stablediffusion/sd/1_114.png filter=lfs diff=lfs merge=lfs -text
+pytorch-stablediffusion/sd/100_97.png filter=lfs diff=lfs merge=lfs -text
+pytorch-stablediffusion/sd/101_225.png filter=lfs diff=lfs merge=lfs -text
+pytorch-stablediffusion/sd/201_328.png filter=lfs diff=lfs merge=lfs -text
+pytorch-stablediffusion/sd/201_948.png filter=lfs diff=lfs merge=lfs -text
+pytorch-stablediffusion/sd/3_{time_process}.png filter=lfs diff=lfs merge=lfs -text
+pytorch-stablediffusion/sd/4_114.png filter=lfs diff=lfs merge=lfs -text
+pytorch-stablediffusion/sd/8_{time_process}.png filter=lfs diff=lfs merge=lfs -text

pytorch-stablediffusion/data/v1-5-pruned-emaonly.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc6cb27103417325ff94f52b7a5d2dde45a7515b25c255d8e396c90014281516
+size 4265380512

pytorch-stablediffusion/sd/attention.py

@@ -0,0 +1,113 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+import math
+
+class SelfAttention(nn.Module):
+    def __init__(self, n_heads: int, d_embed: int, in_proj_bias=True, out_proj_bias=True):
+        super().__init__()
+
+        self.in_proj = nn.Linear(d_embed, 3 * d_embed, bias=in_proj_bias)
+        self.out_proj = nn.Linear(d_embed, d_embed, bias=out_proj_bias)
+        self.n_heads = n_heads
+        self.d_head = d_embed // n_heads #embedding divided by number of heads
+
+    def forward(self, x: torch.Tensor, causal_mask=False):
+        # x: (Batch_Size, Seq_Len, Dim)
+
+        input_shape = x.shape
+        batch_size, sequence_length, d_embed = input_shape
+
+        intermim_shape = (batch_size, sequence_length, self.n_heads, self.d_head)
+
+        # (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, Dim * 3) -> 3 tensors of shape (Batch_Size, Seq_Len, Dim)
+        q, k, v = self.in_proj(x).chunk(3, dim=-1)
+
+        # (Batch_Size, Seq_len, Dim) -> (Batch_Size, Seq_Len, H, Dim / H) -> (Batch_Size, H, Seq_Len, Dim / H)
+        # Each head will watch the full sequence but only a part of the embedding of each pizel
+        q = q.view(intermim_shape).transpose(1, 2)
+        k = k.view(intermim_shape).transpose(1, 2)
+        v = v.view(intermim_shape).transpose(1, 2)
+
+        # (Batch_Size, H, Seq_Len, Seq_Len)
+        weight = q @ k.transpose(-1, -2)
+
+        if causal_mask:
+            # Mask where the upper triangle (above the principle diagonal)  is made up of 1
+            mask = torch.ones_like(weight, dtype=torch.bool).triu(1)
+            weight.masked_fill_(mask, -torch.inf)
+        
+        weight /= math.sqrt(self.d_head)
+
+        weight = F.softmax(weight, dim=-1)
+
+        # (Batch_Size, H, Seq_Len, Seq_Len) @ (Batch_Size, H, Seq_Len, Dim / H) -> (Batch_Size, H, Seq_Len, Dim / H)
+        output = weight @ v
+
+        #(Batch_Size, H, Seq_Len, Dim / H) -> (Batch_Size, Seq_Len, H, Dim / H)
+        output = output.transpose(1, 2) 
+
+        output = output.reshape(input_shape)
+
+        output = self.out_proj(output)
+
+        # (Batch_Size, Seq_Len, Dim)
+        return output
+
+class CrossAttention(nn.Module):
+    def __init__(self, n_heads, d_embed, d_cross, in_proj_bias=True, out_proj_bias=True):
+        super().__init__()
+        self.q_proj   = nn.Linear(d_embed, d_embed, bias=in_proj_bias)
+        self.k_proj   = nn.Linear(d_cross, d_embed, bias=in_proj_bias)
+        self.v_proj   = nn.Linear(d_cross, d_embed, bias=in_proj_bias)
+        self.out_proj = nn.Linear(d_embed, d_embed, bias=out_proj_bias)
+        self.n_heads = n_heads
+        self.d_head = d_embed // n_heads
+    
+    def forward(self, x, y):
+        # x (latent): # (Batch_Size, Seq_Len_Q, Dim_Q)
+        # y (context): # (Batch_Size, Seq_Len_KV, Dim_KV) = (Batch_Size, 77, 768)
+
+        input_shape = x.shape
+        batch_size, sequence_length, d_embed = input_shape
+        # Divide each embedding of Q into multiple heads such that d_heads * n_heads = Dim_Q
+        interim_shape = (batch_size, -1, self.n_heads, self.d_head)
+        
+        # (Batch_Size, Seq_Len_Q, Dim_Q) -> (Batch_Size, Seq_Len_Q, Dim_Q)
+        q = self.q_proj(x)
+        # (Batch_Size, Seq_Len_KV, Dim_KV) -> (Batch_Size, Seq_Len_KV, Dim_Q)
+        k = self.k_proj(y)
+        # (Batch_Size, Seq_Len_KV, Dim_KV) -> (Batch_Size, Seq_Len_KV, Dim_Q)
+        v = self.v_proj(y)
+
+        # (Batch_Size, Seq_Len_Q, Dim_Q) -> (Batch_Size, Seq_Len_Q, H, Dim_Q / H) -> (Batch_Size, H, Seq_Len_Q, Dim_Q / H)
+        q = q.view(interim_shape).transpose(1, 2) 
+        # (Batch_Size, Seq_Len_KV, Dim_Q) -> (Batch_Size, Seq_Len_KV, H, Dim_Q / H) -> (Batch_Size, H, Seq_Len_KV, Dim_Q / H)
+        k = k.view(interim_shape).transpose(1, 2) 
+        # (Batch_Size, Seq_Len_KV, Dim_Q) -> (Batch_Size, Seq_Len_KV, H, Dim_Q / H) -> (Batch_Size, H, Seq_Len_KV, Dim_Q / H)
+        v = v.view(interim_shape).transpose(1, 2) 
+        
+        # (Batch_Size, H, Seq_Len_Q, Dim_Q / H) @ (Batch_Size, H, Dim_Q / H, Seq_Len_KV) -> (Batch_Size, H, Seq_Len_Q, Seq_Len_KV)
+        weight = q @ k.transpose(-1, -2)
+        
+        # (Batch_Size, H, Seq_Len_Q, Seq_Len_KV)
+        weight /= math.sqrt(self.d_head)
+        
+        # (Batch_Size, H, Seq_Len_Q, Seq_Len_KV)
+        weight = F.softmax(weight, dim=-1)
+        
+        # (Batch_Size, H, Seq_Len_Q, Seq_Len_KV) @ (Batch_Size, H, Seq_Len_KV, Dim_Q / H) -> (Batch_Size, H, Seq_Len_Q, Dim_Q / H)
+        output = weight @ v
+        
+        # (Batch_Size, H, Seq_Len_Q, Dim_Q / H) -> (Batch_Size, Seq_Len_Q, H, Dim_Q / H)
+        output = output.transpose(1, 2).contiguous()
+        
+        # (Batch_Size, Seq_Len_Q, H, Dim_Q / H) -> (Batch_Size, Seq_Len_Q, Dim_Q)
+        output = output.view(input_shape)
+        
+        # (Batch_Size, Seq_Len_Q, Dim_Q) -> (Batch_Size, Seq_Len_Q, Dim_Q)
+        output = self.out_proj(output)
+
+        # (Batch_Size, Seq_Len_Q, Dim_Q)
+        return output
+

pytorch-stablediffusion/sd/clip.py

@@ -0,0 +1,88 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+from attention import SelfAttention
+
+#CLIP is similar to encoded layer of the Transformer, input embedding, multi head attentino, feed forward, add and norm in between
+
+class CLIPEmbedding(nn.Module):
+
+    def __init__(self, n_vocab: int, n_embd: int, n_tokens: int):
+        super().__init__()
+
+        self.token_embedding = nn.Embedding(n_vocab, n_embd)
+        self.position_embedding = nn.Parameter(torch.zeros(n_tokens, n_embd))
+    
+    def forward(self, tokens):
+        # (Batch_Size, Seq_Len) -> (Batch_Size, Seq_Len, Dim)
+        x = self.token_embedding(tokens)
+
+        x += self.position_embedding
+
+        return x
+
+class CLIPLayer(nn.Module):
+
+    def __init__(self, n_head: int, n_embd: int):
+        super().__init__()
+
+        self.layernorm_1 = nn.LayerNorm(n_embd)
+        self.attention = SelfAttention(n_head, n_embd)
+        self.layernorm_2 = nn.LayerNorm(n_embd)
+        self.linear_1 = nn.Linear(n_embd, 4 * n_embd)
+        self.linear_2 = nn.Linear(4 * n_embd, n_embd)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # (Batch_Size, Seq_Len, Dim)
+        residue = x
+
+        ##SELF ATTENTION
+
+        x = self.layernorm_1(x)
+
+        x = self.attention(x, causal_mask=True)
+
+        x += residue
+
+        ## FEEDFORWARD LAYER
+
+        residue = x
+
+        x = self.layernorm_2(x)
+
+        x = self.linear_1(x)
+
+        x = x * torch.sigmoid(1.702 * x) #QuickGELU activation function
+
+        x = self.linear_2(x)
+
+        x += residue
+
+        return x
+class CLIP(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.embedding = CLIPEmbedding(49408, 768, 77) #vocab size, embedding size,max sequence length (padding)
+
+        self.layers = nn.ModuleList([
+            CLIPLayer(12, 768) for i in range(12) #number of heads in multi head attention, embedding size, layers
+        ])
+
+        self.layernorm = nn.LayerNorm(768)  #number of features
+    
+    def forward(self, tokens: torch.LongTensor) -> torch.FloatTensor:
+        tokens = tokens.type(torch.long)
+
+        # (Batch_Size, Seq_Len) -> (Batch_Size, Seq_Len, Dim)
+        state = self.embedding(tokens)
+
+        for layer in self.layers:
+            state = layer(state)
+        
+        # (Batch_Size, Seq_Len, Dim)
+        output = self.layernorm(state)
+
+        return output
+
+

pytorch-stablediffusion/sd/decoder.py

@@ -0,0 +1,177 @@
+import torch
+from torch import nn
+from torch.nn import 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

pytorch-stablediffusion/sd/diffusion.py

@@ -0,0 +1,349 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+from attention import SelfAttention, CrossAttention
+
+class TimeEmbedding(nn.Module):
+    def __init__(self, n_embd):
+        super().__init__()
+        self.linear_1 = nn.Linear(n_embd, 4 * n_embd)
+        self.linear_2 = nn.Linear(4 * n_embd, 4 * n_embd)
+
+    def forward(self, x):
+        # x: (1, 320)
+
+        # (1, 320) -> (1, 1280)
+        x = self.linear_1(x)
+        
+        # (1, 1280) -> (1, 1280)
+        x = F.silu(x) 
+        
+        # (1, 1280) -> (1, 1280)
+        x = self.linear_2(x)
+
+        return x
+
+class UNET_ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, n_time=1280):
+        super().__init__()
+        self.groupnorm_feature = nn.GroupNorm(32, in_channels)
+        self.conv_feature = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
+        self.linear_time = nn.Linear(n_time, out_channels)
+
+        self.groupnorm_merged = nn.GroupNorm(32, out_channels)
+        self.conv_merged = 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, feature, time):
+        # feature: (Batch_Size, In_Channels, Height, Width)
+        # time: (1, 1280)
+
+        residue = feature
+        
+        # (Batch_Size, In_Channels, Height, Width) -> (Batch_Size, In_Channels, Height, Width)
+        feature = self.groupnorm_feature(feature)
+        
+        # (Batch_Size, In_Channels, Height, Width) -> (Batch_Size, In_Channels, Height, Width)
+        feature = F.silu(feature)
+        
+        # (Batch_Size, In_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)
+        feature = self.conv_feature(feature)
+        
+        # (1, 1280) -> (1, 1280)
+        time = F.silu(time)
+
+        # (1, 1280) -> (1, Out_Channels)
+        time = self.linear_time(time)
+        
+        # Add width and height dimension to time. 
+        # (Batch_Size, Out_Channels, Height, Width) + (1, Out_Channels, 1, 1) -> (Batch_Size, Out_Channels, Height, Width)
+        merged = feature + time.unsqueeze(-1).unsqueeze(-1)
+        
+        # (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)
+        merged = self.groupnorm_merged(merged)
+        
+        # (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)
+        merged = F.silu(merged)
+        
+        # (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)
+        merged = self.conv_merged(merged)
+        
+        # (Batch_Size, Out_Channels, Height, Width) + (Batch_Size, Out_Channels, Height, Width) -> (Batch_Size, Out_Channels, Height, Width)
+        return merged + self.residual_layer(residue)
+
+class UNET_AttentionBlock(nn.Module):
+    def __init__(self, n_head: int, n_embd: int, d_context=768):
+        super().__init__()
+        channels = n_head * n_embd
+        
+        self.groupnorm = nn.GroupNorm(32, channels, eps=1e-6)
+        self.conv_input = nn.Conv2d(channels, channels, kernel_size=1, padding=0)
+
+        self.layernorm_1 = nn.LayerNorm(channels)
+        self.attention_1 = SelfAttention(n_head, channels, in_proj_bias=False)
+        self.layernorm_2 = nn.LayerNorm(channels)
+        self.attention_2 = CrossAttention(n_head, channels, d_context, in_proj_bias=False)
+        self.layernorm_3 = nn.LayerNorm(channels)
+        self.linear_geglu_1  = nn.Linear(channels, 4 * channels * 2)
+        self.linear_geglu_2 = nn.Linear(4 * channels, channels)
+
+        self.conv_output = nn.Conv2d(channels, channels, kernel_size=1, padding=0)
+    
+    def forward(self, x, context):
+        # x: (Batch_Size, Features, Height, Width)
+        # context: (Batch_Size, Seq_Len, Dim)
+
+        residue_long = x
+
+        # (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height, Width)
+        x = self.groupnorm(x)
+        
+        # (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height, Width)
+        x = self.conv_input(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)
+        x = x.transpose(-1, -2)
+        
+        # Normalization + Self-Attention with skip connection
+
+        # (Batch_Size, Height * Width, Features)
+        residue_short = x
+        
+        # (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)
+        x = self.layernorm_1(x)
+        
+        # (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)
+        x = self.attention_1(x)
+        
+        # (Batch_Size, Height * Width, Features) + (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)
+        x += residue_short
+        
+        # (Batch_Size, Height * Width, Features)
+        residue_short = x
+
+        # Normalization + Cross-Attention with skip connection
+        
+        # (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)
+        x = self.layernorm_2(x)
+        
+        # (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)
+        x = self.attention_2(x, context)
+        
+        # (Batch_Size, Height * Width, Features) + (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)
+        x += residue_short
+        
+        # (Batch_Size, Height * Width, Features)
+        residue_short = x
+
+        # Normalization + FFN with GeGLU and skip connection
+        
+        # (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)
+        x = self.layernorm_3(x)
+        
+        # GeGLU as implemented in the original code: https://github.com/CompVis/stable-diffusion/blob/21f890f9da3cfbeaba8e2ac3c425ee9e998d5229/ldm/modules/attention.py#L37C10-L37C10
+        # (Batch_Size, Height * Width, Features) -> two tensors of shape (Batch_Size, Height * Width, Features * 4)
+        x, gate = self.linear_geglu_1(x).chunk(2, dim=-1) 
+        
+        # Element-wise product: (Batch_Size, Height * Width, Features * 4) * (Batch_Size, Height * Width, Features * 4) -> (Batch_Size, Height * Width, Features * 4)
+        x = x * F.gelu(gate)
+        
+        # (Batch_Size, Height * Width, Features * 4) -> (Batch_Size, Height * Width, Features)
+        x = self.linear_geglu_2(x)
+        
+        # (Batch_Size, Height * Width, Features) + (Batch_Size, Height * Width, Features) -> (Batch_Size, Height * Width, Features)
+        x += residue_short
+        
+        # (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))
+
+        # Final skip connection between initial input and output of the block
+        # (Batch_Size, Features, Height, Width) + (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height, Width)
+        return self.conv_output(x) + residue_long
+
+class Upsample(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+    
+    def forward(self, x):
+        # (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height * 2, Width * 2)
+        x = F.interpolate(x, scale_factor=2, mode='nearest') 
+        return self.conv(x)
+
+class SwitchSequential(nn.Sequential):
+    def forward(self, x, context, time):
+        for layer in self:
+            if isinstance(layer, UNET_AttentionBlock):
+                x = layer(x, context)
+            elif isinstance(layer, UNET_ResidualBlock):
+                x = layer(x, time)
+            else:
+                x = layer(x)
+        return x
+
+class UNET(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoders = nn.ModuleList([
+            # (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)
+            SwitchSequential(nn.Conv2d(4, 320, kernel_size=3, padding=1)),
+            
+            # (Batch_Size, 320, Height / 8, Width / 8) -> # (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)
+            SwitchSequential(UNET_ResidualBlock(320, 320), UNET_AttentionBlock(8, 40)),
+            
+            # (Batch_Size, 320, Height / 8, Width / 8) -> # (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)
+            SwitchSequential(UNET_ResidualBlock(320, 320), UNET_AttentionBlock(8, 40)),
+            
+            # (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 16, Width / 16)
+            SwitchSequential(nn.Conv2d(320, 320, kernel_size=3, stride=2, padding=1)),
+            
+            # (Batch_Size, 320, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16)
+            SwitchSequential(UNET_ResidualBlock(320, 640), UNET_AttentionBlock(8, 80)),
+            
+            # (Batch_Size, 640, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16)
+            SwitchSequential(UNET_ResidualBlock(640, 640), UNET_AttentionBlock(8, 80)),
+            
+            # (Batch_Size, 640, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 32, Width / 32)
+            SwitchSequential(nn.Conv2d(640, 640, kernel_size=3, stride=2, padding=1)),
+            
+            # (Batch_Size, 640, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32)
+            SwitchSequential(UNET_ResidualBlock(640, 1280), UNET_AttentionBlock(8, 160)),
+            
+            # (Batch_Size, 1280, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32)
+            SwitchSequential(UNET_ResidualBlock(1280, 1280), UNET_AttentionBlock(8, 160)),
+            
+            # (Batch_Size, 1280, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 64, Width / 64)
+            SwitchSequential(nn.Conv2d(1280, 1280, kernel_size=3, stride=2, padding=1)),
+            
+            # (Batch_Size, 1280, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64)
+            SwitchSequential(UNET_ResidualBlock(1280, 1280)),
+            
+            # (Batch_Size, 1280, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64)
+            SwitchSequential(UNET_ResidualBlock(1280, 1280)),
+        ])
+
+        self.bottleneck = SwitchSequential(
+            # (Batch_Size, 1280, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64)
+            UNET_ResidualBlock(1280, 1280), 
+            
+            # (Batch_Size, 1280, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64)
+            UNET_AttentionBlock(8, 160), 
+            
+            # (Batch_Size, 1280, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64)
+            UNET_ResidualBlock(1280, 1280), 
+        )
+        
+        self.decoders = nn.ModuleList([
+            # (Batch_Size, 2560, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64)
+            SwitchSequential(UNET_ResidualBlock(2560, 1280)),
+            
+            # (Batch_Size, 2560, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64)
+            SwitchSequential(UNET_ResidualBlock(2560, 1280)),
+            
+            # (Batch_Size, 2560, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 64, Width / 64) -> (Batch_Size, 1280, Height / 32, Width / 32) 
+            SwitchSequential(UNET_ResidualBlock(2560, 1280), Upsample(1280)),
+            
+            # (Batch_Size, 2560, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32)
+            SwitchSequential(UNET_ResidualBlock(2560, 1280), UNET_AttentionBlock(8, 160)),
+            
+            # (Batch_Size, 2560, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32) -> (Batch_Size, 1280, Height / 32, Width / 32)
+            SwitchSequential(UNET_ResidualBlock(2560, 1280), UNET_AttentionBlock(8, 160)),
+            
+            # (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)
+            SwitchSequential(UNET_ResidualBlock(1920, 1280), UNET_AttentionBlock(8, 160), Upsample(1280)),
+            
+            # (Batch_Size, 1920, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16)
+            SwitchSequential(UNET_ResidualBlock(1920, 640), UNET_AttentionBlock(8, 80)),
+            
+            # (Batch_Size, 1280, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16) -> (Batch_Size, 640, Height / 16, Width / 16)
+            SwitchSequential(UNET_ResidualBlock(1280, 640), UNET_AttentionBlock(8, 80)),
+            
+            # (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)
+            SwitchSequential(UNET_ResidualBlock(960, 640), UNET_AttentionBlock(8, 80), Upsample(640)),
+            
+            # (Batch_Size, 960, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)
+            SwitchSequential(UNET_ResidualBlock(960, 320), UNET_AttentionBlock(8, 40)),
+            
+            # (Batch_Size, 640, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)
+            SwitchSequential(UNET_ResidualBlock(640, 320), UNET_AttentionBlock(8, 40)),
+            
+            # (Batch_Size, 640, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)
+            SwitchSequential(UNET_ResidualBlock(640, 320), UNET_AttentionBlock(8, 40)),
+        ])
+
+    def forward(self, x, context, time):
+        # x: (Batch_Size, 4, Height / 8, Width / 8)
+        # context: (Batch_Size, Seq_Len, Dim) 
+        # time: (1, 1280)
+
+        skip_connections = []
+        for layers in self.encoders:
+            x = layers(x, context, time)
+            skip_connections.append(x)
+
+        x = self.bottleneck(x, context, time)
+
+        for layers in self.decoders:
+            # Since we always concat with the skip connection of the encoder, the number of features increases before being sent to the decoder's layer
+            x = torch.cat((x, skip_connections.pop()), dim=1) 
+            x = layers(x, context, time)
+        
+        return x
+
+
+class UNET_OutputLayer(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.groupnorm = nn.GroupNorm(32, in_channels)
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
+    
+    def forward(self, x):
+        # x: (Batch_Size, 320, Height / 8, Width / 8)
+
+        # (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)
+        x = self.groupnorm(x)
+        
+        # (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 320, Height / 8, Width / 8)
+        x = F.silu(x)
+        
+        # (Batch_Size, 320, Height / 8, Width / 8) -> (Batch_Size, 4, Height / 8, Width / 8)
+        x = self.conv(x)
+        
+        # (Batch_Size, 4, Height / 8, Width / 8) 
+        return x
+
+class Diffusion(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.time_embedding = TimeEmbedding(320)
+        self.unet = UNET()
+        self.final = UNET_OutputLayer(320, 4)
+    
+    def forward(self, latent, context, time):
+        # latent: (Batch_Size, 4, Height / 8, Width / 8)
+        # context: (Batch_Size, Seq_Len, Dim)
+        # time: (1, 320)
+
+        # (1, 320) -> (1, 1280)
+        time = self.time_embedding(time)
+        
+        # (Batch, 4, Height / 8, Width / 8) -> (Batch, 320, Height / 8, Width / 8)
+        output = self.unet(latent, context, time)
+        
+        # (Batch, 320, Height / 8, Width / 8) -> (Batch, 4, Height / 8, Width / 8)
+        output = self.final(output)
+        
+        # (Batch, 4, Height / 8, Width / 8)
+        return output

pytorch-stablediffusion/sd/encoder.py

@@ -0,0 +1,96 @@
+#Variational Encoder
+import torch
+from torch import nn
+from torch.nn import functional as F
+from decoder import VAE_AttentionBlock, VAE_ResidualBlock
+
+class VAE_Encoder(nn.Sequential): #encoder is a sequence of submodels
+    def __init__(self):
+        super().__init__( #each model reduces the dimension of data and increases the number of features
+            #convert from 3 to 128 channels
+            #(Batch_Size, Channel, Height, Width) -> (Batch_Size, 128, Height, Width)
+            nn.Conv2d(3, 128, kernel_size=3, padding=1),
+
+            #Residual block, combination of convolutions and normalization
+            # (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/2, Width/2)
+            nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=0),
+            
+            # (Batch_Size, 128, Height/2, Width/2) -> (Batch_Size, 256, Height/2, Width/2)
+            VAE_ResidualBlock(128, 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/4, Width/4)
+            nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=0),
+
+            # (Batch_Size, 256, Height/4, Width/4) -> (Batch_Size, 512, Height/4, Width/4)
+            VAE_ResidualBlock(256, 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/8, Width/8)
+            nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=0),
+
+            # (Batch_Size, 512, Height/4, Width/4) -> (Batch_Size, 512, Height/8, Width/8)
+            VAE_ResidualBlock(512, 512),
+
+            VAE_ResidualBlock(512, 512),
+
+            # (Batch_Size, 512, Height/8, Width/8) -> (Batch_Size, 512, Height/8, Width/8)
+            VAE_ResidualBlock(512, 512),
+
+            #self attention over each pixel, relate pixels to each other
+            # (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),
+
+            #Group normalization (Batch_Size, 512, Height/8, Width/8) -> (Batch_Size, 512, Height/8, Width/8)
+            nn.GroupNorm(32, 512), #32 groups, 512 channels
+
+            nn.SiLU(), #activation function, sigmoid linear unit, similar to ReLU
+
+            # (Batch_Size, 512, Height/8, Width/8) -> (Batch_Size, 512, Height/8, Width/8)           
+            nn.Conv2d(512, 8, kernel_size=3, padding=1),
+
+            # (Batch_Size, 8, Height/8, Width/8) -> (Batch_Size, 8, Height/8, Width/8)           
+            nn.Conv2d(8, 8, kernel_size=1, padding=0)
+        )
+    def forward(self, x: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
+        # x: (Batch_Size, Channel, Height, Width)
+        # noise: (Batch_Size, Out_Channels, Height / 8, Width / 8)
+
+        #VAE learns the mu and the sigma which is the mean and variance of the distribution
+        for module in self:
+            if getattr(module, 'stride', None) == (2, 2): #apply to convolutions that only have stride 2
+                # (Padding_Left, Padding_Right, Padding-Top, Padding_Bottom)
+                x = F.pad(x, (0, 1, 0, 1))
+            x = module(x)
+        # (Batch_Size, 8, Height, Height / 8, Width / 8) -> two tensors of shape (Batch_Size, 4, Height / 8, Width / 8)
+        mean, log_variance = torch.chunk(x, 2, dim=1)
+
+        log_variance = torch.clamp(log_variance, -30, 20)
+
+        variance = log_variance.exp()
+
+        stdev = variance.sqrt()
+
+        # Z -> N(0,1) -> N(mean, variance)?
+        #X = mean + stdev*Z
+        x = mean + stdev * noise
+        
+        #Scale the output by a constant
+        x *= 0.18215
+
+        return x
+        
+

pytorch-stablediffusion/sd/model_loader.py

@@ -0,0 +1,28 @@
+from clip import CLIP
+from encoder import VAE_Encoder
+from decoder import VAE_Decoder
+from diffusion import Diffusion
+
+import model_converter
+#the ckpt file is a dictionary that contains many keys where each key corresponds to one matrix of the model
+def preload_models_from_standard_weights(ckpt_path, device):
+    state_dict = model_converter.load_from_standard_weights(ckpt_path, device)
+
+    encoder = VAE_Encoder().to(device)
+    encoder.load_state_dict(state_dict['encoder'], strict=True)
+
+    decoder = VAE_Decoder().to(device)
+    decoder.load_state_dict(state_dict['decoder'], strict=True)
+
+    diffusion = Diffusion().to(device)
+    diffusion.load_state_dict(state_dict['diffusion'], strict=True)
+
+    clip = CLIP().to(device)
+    clip.load_state_dict(state_dict['clip'], strict=True)
+
+    return {
+        'clip': clip,
+        'encoder': encoder,
+        'decoder': decoder,
+        'diffusion': diffusion,
+    }

pytorch-stablediffusion/sd/pipeline.py

@@ -0,0 +1,181 @@
+import torch
+import numpy as np
+from tqdm import tqdm
+from sampler import Sampler
+
+WIDTH = 512 #stable diffusion only takes in this dimension
+HEIGHT = 512
+LATENTS_WIDTH = WIDTH // 8
+LATENTS_HEIGHT = HEIGHT // 8
+
+#strength is how much attention we want to put into the starting image
+def generate(prompt: str, 
+            uncond_prompt: str, # Negative prompt or empty string
+            input_image=None, 
+            strength=0.8, 
+            do_cfg=True, 
+            cfg_scale=7.5, 
+            sampler_name="ddpm", 
+            n_inference_steps=50, 
+            models={}, 
+            seed=None, 
+            device=None,
+            idle_device=None,
+            tokenizer=None,
+            eta = 0.0):
+    with torch.no_grad():
+        if not (0 < strength <= 1):
+            raise ValueError("strength must be between 0 and 1")
+        if idle_device:
+            to_idle = lambda x: x.to(idle_device)
+        else: 
+            to_idle = lambda x: x
+        generator = torch.Generator(device=device)
+        if seed is None:
+            generator.seed()
+        else:
+            generator.manual_seed(seed)
+        
+        clip = models["clip"]
+        clip.to(device)
+
+        if do_cfg:
+            # Convert the prompt into tokens using tokenizer
+            cond_tokens = tokenizer.batch_encode_plus(
+                [prompt], padding="max_length", max_length=77
+            ).input_ids
+            # (Batch_Size, Seq_Len)
+            cond_tokens = torch.tensor(cond_tokens, dtype=torch.long, device=device)
+            # (Batch_Size, Seq_Len) -> (Batch_Size, Seq_Len, Dim)
+            cond_context = clip(cond_tokens)
+
+            #with no conditioins now
+            uncond_tokens = tokenizer.batch_encode_plus(
+                [uncond_prompt], padding="max_length", max_length=77
+            ).input_ids
+            uncond_tokens = torch.tensor(uncond_tokens, dtype=torch.long, device=device)
+            uncond_context = clip(uncond_tokens)
+            
+            # (Batch_2, Seq_Len, Dim) = (2, 77, 768)
+            context = torch.cat([cond_context, uncond_context])
+        else:
+            # Convert it into a list of tokens
+            tokens = tokenizer.batch_encode_plus(
+                [prompt], padding="max_length", max_length=77
+            ).input_ids
+            tokens = torch.tensor(tokens, dtype=torch.long, device=device)
+            # (1, 77, 768)
+            context = clip(tokens)
+        to_idle(clip) #very useful if you have a limited gpu and want to offload to cpu
+
+        if sampler_name in ["ddpm", "ddim", "euler", "dpm_solver"]:
+            sampler = Sampler(generator)
+            sampler.set_inference_timesteps(n_inference_steps)
+        else:
+            raise ValueError(f"Unknown sampler name")
+        
+        latents_shape = (1, 4, LATENTS_HEIGHT, LATENTS_WIDTH)
+
+        if input_image:
+            encoder = models["encoder"]
+            encoder.to(device)
+
+            input_image_tensor = input_image.resize((WIDTH, HEIGHT))
+            input_image_tensor = np.array(input_image_tensor)
+            # (Height, Width, Channel)
+            input_image_tensor = torch.tensor(input_image_tensor, dtype=torch.float32, device=device)
+            input_image_tensor = rescale(input_image_tensor, (0,255), (-1,1))
+            # (Height, Width, Channel) -> (Batch_Size, Height, Width, Channel)
+            input_image_tensor = input_image_tensor.unsqueeze(0)
+            # (Batch_Size, Height, Width, Channel) -> (Batch_Size, Channel, Height, Width)
+            input_image_tensor = input_image_tensor.permute(0,3,1,2)
+
+            encoder_noise = torch.randn(latents_shape, generator=generator, device=device)
+            #run the image through the encoder of the VAE
+            latents = encoder(input_image_tensor, encoder_noise)
+
+            #Add noise to the latent, the more the strength, the stronger the noise, making the model more creative
+            sampler.set_strength(strength=strength)
+            latents = sampler.add_noise(latents, sampler.timesteps[0])
+
+            to_idle(encoder)
+        else:
+            # If we are doing text to image, start with random noise N(0,1)
+            latents = torch.randn(latents_shape, generator=generator, device=device)
+        
+        #999...0
+        #1000 980 940 920 900 880....0, each of these time steps indicates a nosie level
+        #can tell the scheduler to reduce noise according to particular time steps, defined by n_inference_steps 
+        diffusion = models["diffusion"]
+        diffusion.to(device)
+
+        timesteps = tqdm(sampler.timesteps)
+        for i, timestep in enumerate(timesteps):
+            # (1, 320)
+            time_embedding = get_time_embedding(timestep).to(device)
+
+            # (Batch_Size, 4, Latents_Height, Latents_Width)
+            model_input = latents
+
+            if do_cfg:
+                # (Batch_Size, 4, Latents_Height, Latents_Width) -> (4 * Batch_Size, 4, Latents_Height, Latents_Width)
+                model_input = model_input.repeat(2, 1, 1, 1)
+
+            # model_output is the predicted noise by the UNET
+            model_output = diffusion(model_input, context, time_embedding)
+
+            if do_cfg:
+                output_cond, output_uncond = model_output.chunk(2)
+                model_output = cfg_scale * (output_cond - output_uncond) + output_uncond
+
+            # how to remove the noise from image? using the scheduler
+            #Remove noise predicted by the UNET
+            if sampler_name == "ddpm":
+                latents = sampler.ddpm_step(timestep, latents, model_output)
+            elif sampler_name == "ddim":
+                latents = sampler.ddim_step(timestep, latents, model_output, eta=eta)
+            elif sampler_name == "euler":
+                latents = sampler.euler_ancestral_step(timestep, latents, model_output, eta=eta)
+            elif sampler_name == "dpm_solver":
+                latents = sampler.dpm_solver_pp_2m_step(timestep, latents, model_output)
+            else:
+                raise ValueError(f"Unknown sampler name {sampler_name}")
+        
+        to_idle(diffusion)
+
+        decoder = models["decoder"]
+        decoder.to(device)
+
+        images = decoder(latents)
+        to_idle(decoder)
+
+        images = rescale(images, (-1, 1), (0, 255), clamp=True)
+        # (Batch_Size, Channel, Height, Width) -> (Batch_Size, Height, Width, Channel)
+        images = images.permute(0, 2, 3, 1)
+        images = images.to("cpu", torch.uint8).numpy()
+        return images[0]
+
+def rescale(x, old_range, new_range, clamp=False):
+    old_min, old_max = old_range
+    new_min, new_max = new_range
+    x -= old_min
+    x *= (new_max - new_min) / (old_max - old_min)
+    x += new_min
+    if clamp:
+        x = x.clamp(new_min, new_max)
+    return x
+
+def get_time_embedding(timestep):
+    freqs = torch.pow(10000, -torch.arange(start=0, end=160, dtype=torch.float32) / 160)
+    # (1, 160)
+    x = torch.tensor([timestep], dtype=torch.float32)[:, None] * freqs[None]
+    # (1, 320)
+    return torch.cat([torch.cos(x), torch.sin(x)], dim=-1)
+
+
+
+
+            
+
+
+

pytorch-stablediffusion/sd/sampler.py

@@ -0,0 +1,198 @@
+import torch
+import numpy as np
+
+class Sampler:
+
+    def __init__(self, generator: torch.Generator, num_training_steps=1000, beta_start: float = 0.00085, beta_end: float=0.0120):
+        #beta is a series of numbers that indicates the variance of the noise that we add with each of these steps
+        # the start and end values were a choice made by the authors
+        # will be using a linear scheduler, 1000 numbers between start and end
+
+        self.betas = torch.linspace(beta_start ** 0.5, beta_end ** 0.5, num_training_steps, dtype=torch.float32) ** 2
+
+        # alpha bar is the product of alpha going from 1 to T
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, 0) 
+        self.one = torch.tensor(1.0)
+
+        self.generator = generator
+        self.num_training_steps = num_training_steps
+        self.timesteps = torch.from_numpy(np.arange(0, num_training_steps)[::-1].copy())
+    
+    def set_inference_timesteps(self, num_inference_steps=50):
+        self.num_inference_steps = num_inference_steps
+        # 999, 998, 997, ... 0 = 1000 steps
+        # 999, 999-20, 999-40, ... 0 = 50 steps
+        step_ratio = self.num_training_steps // num_inference_steps
+        timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
+        self.timesteps = torch.from_numpy(timesteps)
+    
+    def _get_previous_timestep(self, timestep:int) -> int:
+        prev_t = timestep - (self.num_training_steps // self.num_inference_steps)
+        return prev_t
+    
+    def _get_variance(self, timestep: int) -> torch.Tensor:
+        prev_t = self._get_previous_timestep(timestep)
+
+        alpha_prod_t = self.alphas_cumprod[timestep]
+        alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.one
+        current_beta_t = 1 - alpha_prod_t / alpha_prod_t_prev
+
+        # Computed using formula (7) of the DDPM paper
+        variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * current_beta_t
+        variance = torch.clamp(variance, min=1e-20)
+
+        return variance
+    
+    def set_strength(self, strength=1):
+        start_step = self.num_inference_steps - int(self.num_inference_steps * strength)
+        self.timesteps = self.timesteps[start_step:]
+        self.start_step = start_step
+
+    def ddpm_step(self, timestep: int, latents: torch.Tensor, model_output: torch.Tensor):
+        t = timestep
+        prev_t = self._get_previous_timestep(t)
+
+        alpha_prod_t = self.alphas_cumprod[t]
+        alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.one
+        beta_prod_t = 1 - alpha_prod_t
+        beta_prod_t_prev = 1 - alpha_prod_t_prev
+        current_alpha_t = alpha_prod_t / alpha_prod_t_prev
+        current_beta_t = 1 - current_alpha_t
+
+        # Compute the predicted original sample using formula (15) of the DDPM paper
+        pred_original_sample = (latents - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
+
+        # Compute the coefficient for pred_original_sample and current sample x_t
+        pred_original_sample_coeff = (alpha_prod_t_prev ** 0.5 * current_beta_t) / beta_prod_t
+        current_sample_coeff = current_alpha_t ** 0.5 * beta_prod_t_prev / beta_prod_t
+
+        # Compute the predicted previous sample mean
+        pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * latents
+
+        variance = 0
+        if t > 0:
+            device = model_output.device
+            noise = torch.randn(model_output.shape, generator=self.generator, device=device, dtype=model_output.dtype)
+            variance = (self._get_variance(t) ** 0.5) * noise
+        
+        # N(0,1) --> N(mu, sigma)
+        # X = mu + sigma * Z where Z ~ N(0, 1)
+        pred_prev_sample = pred_prev_sample + variance
+        return pred_prev_sample
+    
+    def ddim_step(self, timestep: int, latents: torch.Tensor, model_output: torch.Tensor, eta=0.0):
+        t = timestep
+        prev_t = self._get_previous_timestep(t)
+        
+        alpha_t = self.alphas_cumprod[t]
+        alpha_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else torch.tensor(1.0, device=latents.device, dtype=latents.dtype)
+        
+        # Predicted original clean sample x_0
+        pred_original_sample = (latents - torch.sqrt(1 - alpha_t) * model_output) / torch.sqrt(alpha_t)
+        
+        # Direction pointing to x_t
+        #dir_xt = torch.sqrt(1 - alpha_prev - (eta ** 2) * ((1 - alpha_prev) / (1 - alpha_t)) * (1 - alpha_t / alpha_prev)) * model_output
+        
+        # Noise term
+        noise = torch.randn_like(latents) if eta > 0 else torch.zeros_like(latents)
+        
+        sigma_t = eta * torch.sqrt((1 - alpha_prev) / (1 - alpha_t)) * torch.sqrt(1 - alpha_t / alpha_prev)
+        
+        # Compute previous latent x_{t-1}
+        #prev_latent = torch.sqrt(alpha_prev) * pred_original_sample + dir_xt + sigma_t * noise
+        prev_latent = torch.sqrt(alpha_prev) * pred_original_sample + torch.sqrt(1 - alpha_prev - sigma_t ** 2) * model_output + sigma_t * noise
+
+        
+        return prev_latent
+    
+    def euler_ancestral_step(self, timestep: int, latents: torch.Tensor, model_output: torch.Tensor, eta=1.0):
+        t = timestep
+        prev_t = self._get_previous_timestep(t)
+        
+        # Convert alphas to sigmas (standard deviation of noise at each timestep)
+        alpha_t = self.alphas_cumprod[t]
+        alpha_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else torch.tensor(1.0, device=latents.device, dtype=latents.dtype)
+
+        sigma_t = torch.sqrt(1 - alpha_t)
+        sigma_prev = torch.sqrt(1 - alpha_prev)
+        
+        # Predict x_0
+        x0_pred = (latents - sigma_t * model_output) / torch.sqrt(alpha_t)
+        
+        # Euler drift step (toward next timestep)
+        dt = sigma_prev - sigma_t
+        x_drift = latents + dt * model_output
+
+        # Stochastic noise addition
+        if eta > 0.0:
+            noise = torch.randn_like(latents)
+            sigma = torch.sqrt(torch.clamp(eta * (sigma_prev**2 - sigma_t**2), min=1e-20))
+            x_drift += sigma * noise
+
+        return x_drift
+    
+    def dpm_solver_pp_2m_step(self, timestep: int, latents: torch.Tensor, model_output: torch.Tensor):
+        """
+        One DPM-Solver++(2M) step with DDIM-style signature.
+
+        Args:
+            timestep:       Current timestep index t.
+            latents:        Latents at current timestep x_t.
+            model_output:   Model prediction ε_θ(x_t, t).
+
+        Returns:
+            x_{t-1}:        Estimated latent at previous timestep.
+        """
+        t = self.timesteps[timestep]
+        prev_t = self.timesteps[timestep + 1] if timestep + 1 < len(self.timesteps) else 0.0  # t_{prev}
+
+        h = prev_t - t  # Note: time goes backward
+
+        # Extract alpha and sigma for current and previous timesteps
+        alpha_t = self.alphas_cumprod[timestep] ** 0.5
+        alpha_prev = self.alphas_cumprod[timestep + 1] ** 0.5 if timestep + 1 < len(self.alphas_cumprod) else self.one
+        sigma_t = (1 - self.alphas_cumprod[timestep]) ** 0.5
+        sigma_prev = (1 - self.alphas_cumprod[timestep + 1]) ** 0.5 if timestep + 1 < len(self.alphas_cumprod) else self.zero
+
+        # Store previous model output if not already done
+        if not hasattr(self, "_prev_model_output"):
+            self._prev_model_output = model_output  # Just initialize on first call
+
+        model_output_t = model_output
+        model_output_prev = self._prev_model_output
+
+        # Compute x0_t and x0_prev estimates
+        x0_t = (latents - sigma_t * model_output_t) / alpha_t
+        x0_prev = (latents - sigma_t * model_output_prev) / alpha_t
+
+        # 2nd-order multistep estimate
+        x0_hat = x0_t + 0.5 * h * (model_output_t - model_output_prev)
+
+        # Estimate x_{t-1}
+        x_prev = alpha_prev * x0_hat + sigma_prev * model_output_prev
+
+        # Update previous model output for next step
+        self._prev_model_output = model_output
+
+        return x_prev
+
+
+    def add_noise(self, original_samples: torch.FloatTensor, timesteps: torch.IntTensor) -> torch.FloatTensor:
+        #at what time we want to add the timestep
+        alpha_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
+        timesteps = timesteps.to(original_samples.device)
+    
+        sqrt_alpha_prod = alpha_cumprod[timesteps] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) #adds a new dimension with length one at a specific pos within tensors shape
+        sqrt_one_minus_alpha_prod = (1 - alpha_cumprod[timesteps]) ** 0.5 #standard deviation
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod) < len(original_samples.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+        
+        # According to the euation (4) of the DDM paper
+        noise = torch.randn(original_samples.shape, generator=self.generator, device=original_samples.device, dtype=original_samples.dtype)
+        noisy_samples = (sqrt_alpha_prod * original_samples) + (sqrt_one_minus_alpha_prod) * noise
+