files upload

app.py

@@ -0,0 +1,53 @@
+import torch
+import spaces
+from PIL import Image
+from generationPipeline import generate
+from transformers import CLIPTokenizer
+from loadModel import preload_models_from_standard_weights
+import gradio as gr
+
+
+Device = 'cuda'
+
+print(f"Using device: {Device}")
+
+
+tokenizer = CLIPTokenizer("vocab.json", merges_file="merges.txt")
+model_file = "weights2.ckpt"
+models = preload_models_from_standard_weights(model_file, Device)
+
+
+@spaces.GPU(duration = 242)
+def generate_image(prompt, strength, seed):
+    # Your generate function adapted to accept parameters
+    output_image = generate(
+        prompt=prompt,
+        uncond_prompt="",
+        input_image=None,
+        strength=strength,
+        do_cfg=True,
+        cfg_scale=8,
+        sampler_name="ddpm",
+        n_inference_steps=50,
+        seed=seed,
+        models=models,
+        device=Device,
+        idle_device="cpu",
+        tokenizer=tokenizer,
+    )
+    return Image.fromarray(output_image)
+
+
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=[
+        gr.inputs.Textbox(label="Prompt"),
+        gr.inputs.Slider(0, 1, step=0.01, label="Strength (For Image-2-Image): Strength = 1 (Output further from input image), Strength = 0 (Output similar as Input image)"),
+        gr.inputs.Number(default=42, label="Seed"),
+    ],
+    outputs=gr.outputs.Image(label="Generated Image"),
+    title="Stable Diffusion Image Generator",
+    description="Generate images from text prompts using Stable Diffusion.",
+)
+
+iface.launch(debug = True)

attention.py

@@ -0,0 +1,131 @@
+import math
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+#Attention: softmax(q @ k.transpose / sqrt(dk)) @ w
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, n_heads, d_embed, in_proj_bias=True, out_proj_bias=True):
+        super().__init__()
+        # This combines the Wq, Wk and Wv matrices into one matrix
+        self.in_proj = nn.Linear(d_embed, 3 * d_embed, bias=in_proj_bias)
+        # This one represents the Wo matrix
+        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, causal_mask=False):
+        # (Batch_Size, Seq_Len, Dim)
+        input_shape = x.shape
+
+        # (Batch_Size, Seq_Len, Dim)
+        batch_size, sequence_length, d_embed = input_shape
+
+        # (Batch_Size, Seq_Len, H, Dim / H)
+        interim_shape = (batch_size, sequence_length, self.n_heads, self.d_head)
+
+        # (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, Dim * 3) -> 3 tensor 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)
+        q = q.view(interim_shape).transpose(1, 2)
+        k = k.view(interim_shape).transpose(1, 2)
+        v = v.view(interim_shape).transpose(1, 2)
+
+        # (Batch_Size, H, Seq_Len, Dim / H) @ (Batch_Size, H, Dim / H, Seq_Len) -> (Batch_Size, H, Seq_Len, Seq_Len)
+        weight = q @ k.transpose(-1, -2)
+
+        if causal_mask:
+            # It masks the token after the current tokens so that the future tokens are not accessible
+            # Mask where the upper triangle (above the principal diagonal) is 1
+            mask = torch.ones_like(weight, dtype=torch.bool).triu(1)
+            # Fill the upper triangle with -inf
+            weight.masked_fill_(mask, -torch.inf)
+
+        # Divide by d_k (Dim / H).
+        # (Batch_Size, H, Seq_Len, Seq_Len) -> (Batch_Size, H, Seq_Len, Seq_Len)
+        weight /= math.sqrt(self.d_head)
+
+        # (Batch_Size, H, Seq_Len, Seq_Len) -> (Batch_Size, H, Seq_Len, Seq_Len)
+        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)
+
+        # (Batch_Size, Seq_Len, H, Dim / H) -> (Batch_Size, Seq_Len, Dim)
+        output = output.reshape(input_shape)
+
+        # (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, Dim)
+        output = self.out_proj(output)
+
+        # (Batch_Size, Seq_Len, Dim)
+        return output
+    
+
+# Calculate Attention between latent and prompt(context)
+
+
+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: (b, h*w, c) -> (b, seq_legth, d_model) = (b, h/8*w/8, 512)
+        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)
+
+        # In cross attention query is taken from one element (latent here) and key, values are taken from another element (context)
+        # (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, Dim) -> (b, h/8*w/8, 512) = (b, h*w, d_model)
+        return output

clip.py

@@ -0,0 +1,97 @@
+import torch
+import torch.nn as nn
+from attention import SelfAttention
+
+class CLIPEmbedding(nn.Module):
+    def __init__(self, n_vocab: int, n_embd: int, n_token: int):
+        super().__init__()
+
+        self.token_embedding = nn.Embedding(n_vocab, n_embd)  #(vocab_Size, embedding_dim)
+        # A learnable weight matrix encodes the position information for each token
+        self.position_embedding = nn.Parameter(torch.zeros((n_token, n_embd)))  #(seq_legth, embedding_dim)
+
+    def forward(self, tokens):
+        # (Batch_Size, Seq_Len) -> (Batch_Size, Seq_Len, Dim)
+        x = self.token_embedding(tokens)
+        # (Batch_Size, Seq_Len) -> (Batch_Size, Seq_Len, Dim)
+        x += self.position_embedding
+
+        return x
+
+
+class CLIPLayer(nn.Module):
+    def __init__(self, n_head: int, n_embd: int):
+        super().__init__()
+
+        # Pre-attention norm
+        self.layernorm_1 = nn.LayerNorm(n_embd)
+        # Self attention
+        self.attention = SelfAttention(n_head, n_embd)
+        # Pre-FNN norm
+        self.layernorm_2 = nn.LayerNorm(n_embd)
+        # Feedforward layer
+        self.linear_1 = nn.Linear(n_embd, 4 * n_embd)
+        self.linear_2 = nn.Linear(4 * n_embd, n_embd)
+
+    def forward(self, x):
+        # (Batch_Size, Seq_Len, Dim)
+        residue = x
+
+        ### SELF ATTENTION ###
+
+        # (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, Dim)
+        x = self.layernorm_1(x)
+
+        # (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, Dim)
+        x = self.attention(x, causal_mask=True)
+
+        # (Batch_Size, Seq_Len, Dim) + (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, Dim)
+        x += residue
+
+        ### FEEDFORWARD LAYER ###
+        # Apply a feedforward layer where the hidden dimension is 4 times the embedding dimension.
+
+        residue = x
+        # (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, Dim)
+        x = self.layernorm_2(x)
+
+        # (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, 4 * Dim)
+        x = self.linear_1(x)
+
+        # (Batch_Size, Seq_Len, 4 * Dim) -> (Batch_Size, Seq_Len, 4 * Dim)
+        x = x * torch.sigmoid(1.702 * x)   # QuickGELU activation function
+
+        # (Batch_Size, Seq_Len, 4 * Dim) -> (Batch_Size, Seq_Len, Dim)
+        x = self.linear_2(x)
+
+        # (Batch_Size, Seq_Len, Dim) + (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, Dim)
+        x += residue
+
+        return x
+    
+
+class CLIP(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embedding = CLIPEmbedding(49408, 768, 77)
+
+        self.layers = nn.ModuleList([
+            CLIPLayer(12, 768) for i in range(12)
+        ])
+
+        self.layernorm = nn.LayerNorm(768)
+
+    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)
+
+        # Apply encoder layers similar to the Transformer's encoder.
+        for layer in self.layers:
+            # (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, Dim)
+            state = layer(state)
+        # (Batch_Size, Seq_Len, Dim) -> (Batch_Size, Seq_Len, Dim)
+        output = self.layernorm(state)
+
+        return output

decoder.py

@@ -0,0 +1,99 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from helperVAE import VAE_AttentionBlock, VAE_ResidualBlock
+
+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

encoder.py

@@ -0,0 +1,102 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from helperVAE import VAE_ResidualBlock, VAE_AttentionBlock
+
+class VAE_Encoder(nn.Sequential):
+    def __init__(self):
+        super().__init__(
+            # (Batch_Size, Channel, Height, Width) -> (Batch_Size, 128, Height, Width)
+            nn.Conv2d(3, 128, kernel_size=3, padding=1),
+
+             # (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 / 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_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)
+            nn.GroupNorm(32, 512),
+
+            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 512, Height / 8, Width / 8)
+            nn.SiLU(),
+
+            # Because the padding=1, it means the width and height will increase by 2
+            # Out_Height = In_Height + Padding_Top + Padding_Bottom
+            # Out_Width = In_Width + Padding_Left + Padding_Right
+            # Since padding = 1 means Padding_Top = Padding_Bottom = Padding_Left = Padding_Right = 1,
+            # Since the Out_Width = In_Width + 2 (same for Out_Height), it will compensate for the Kernel size of 3
+            # (Batch_Size, 512, Height / 8, Width / 8) -> (Batch_Size, 8, 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, noise):
+        # x: (Batch_Size, Channel, Height, Width)
+        # noise: (Batch_Size, 4, Height / 8, Width / 8)
+
+        for module in self:
+
+            if getattr(module, 'stride', None) == (2, 2):  # Padding at downsampling should be asymmetric (see #8)
+                # Pad: (Padding_Left, Padding_Right, Padding_Top, Padding_Bottom).
+                # Pad with zeros on the right and bottom.
+                # (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)
+                x = F.pad(x, (0, 1, 0, 1))
+
+            x = module(x)
+        # (Batch_Size, 8, Height / 8, Width / 8) -> two tensors of shape (Batch_Size, 4, Height / 8, Width / 8)
+        mean, log_variance = torch.chunk(x, 2, dim=1)
+        # Clamp the log variance between -30 and 20, so that the variance is between (circa) 1e-14 and 1e8.
+        # (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 4, Height / 8, Width / 8)
+        log_variance = torch.clamp(log_variance, -30, 20)
+        # (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 4, Height / 8, Width / 8)
+        variance = log_variance.exp()
+        # (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 4, Height / 8, Width / 8)
+        stdev = variance.sqrt()
+
+        # Transform N(0, 1) -> N(mean, stdev)
+        # (Batch_Size, 4, Height / 8, Width / 8) -> (Batch_Size, 4, Height / 8, Width / 8)
+        x = mean + stdev * noise
+
+        # Scale by a constant
+        x *= 0.18215
+
+        return x

helperUNET.py

@@ -0,0 +1,179 @@
+import torch.nn as nn
+from attention import SelfAttention, CrossAttention
+from torch.nn import functional as F
+
+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')  #upsampling using nearest neighbor interpolation
+        return self.conv(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 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
+#switch between attention and residual layer

helperVAE.py

@@ -0,0 +1,83 @@
+import torch.nn as nn
+from attention import SelfAttention
+from torch.nn import functional as F
+
+
+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)

sampler.py

@@ -0,0 +1,127 @@
+import torch
+import numpy as np
+
+class DDPMSampler:
+
+    def __init__(self, generator: torch.Generator, num_training_steps=1000, beta_start: float = 0.00085, beta_end: float = 0.0120):
+        # Params "beta_start" and "beta_end" taken from: https://github.com/CompVis/stable-diffusion/blob/21f890f9da3cfbeaba8e2ac3c425ee9e998d5229/configs/stable-diffusion/v1-inference.yaml#L5C8-L5C8
+        # For the naming conventions, refer to the DDPM paper (https://arxiv.org/pdf/2006.11239.pdf)
+        self.betas = torch.linspace(beta_start ** 0.5, beta_end ** 0.5, num_training_steps, dtype=torch.float32) ** 2  #beta
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)  # alpha bar
+        self.one = torch.tensor(1.0)
+
+        self.generator = generator
+
+        self.num_train_timesteps = num_training_steps
+        self.timesteps = torch.from_numpy(np.arange(0, num_training_steps)[::-1].copy()) ##[999, 998, ...0]
+
+    def set_inference_timesteps(self, num_inference_steps=50):
+      # num_inference_steps = 50
+      # step ratio = num_training_steps // inference_steps = 20
+        self.num_inference_steps = num_inference_steps
+        step_ratio = self.num_train_timesteps // self.num_inference_steps # 1000/50 = 20
+        timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)  #[980, 960, ..0]
+        self.timesteps = torch.from_numpy(timesteps)
+
+    def _get_previous_timestep(self, timestep: int) -> int:
+        prev_t = timestep - self.num_train_timesteps // self.num_inference_steps  #eg: t = 960, t-1 = 960-20 = 940
+        return prev_t
+
+    def _get_variance(self, timestep: int) -> torch.Tensor:
+        prev_t = self._get_previous_timestep(timestep)  #t-1
+
+        alpha_prod_t = self.alphas_cumprod[timestep] #alpha bar t
+        alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.one  #alpha bar t minus 1
+        current_beta_t = 1 - alpha_prod_t / alpha_prod_t_prev  #beta t
+
+        # For t > 0, compute predicted variance βt (see formula (6) and (7) from https://arxiv.org/pdf/2006.11239.pdf)
+        # and sample from it to get previous sample
+        # x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample
+        variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * current_beta_t  #variance#
+
+        # we always take the log of variance, so clamp it to ensure it's not 0
+        variance = torch.clamp(variance, min=1e-20)
+
+        return variance
+
+    def set_strength(self, strength=1):
+        """
+            Set how much noise to add to the input image.
+            More noise (strength ~ 1) means that the output will be further from the input image.
+            Less noise (strength ~ 0) means that the output will be closer to the input image.
+        """
+        # more strength -> start step is approximately 0 that is model starts from pure noise and generates the image from it, strength = 1,  start step = 50 - (50 * 1) = 0
+        # less strenght -> start step is skipped till 50 so model has the less noisified image a time step 50, model reconstructs the image from the less noisified image, strength = 0, start_step = 50
+
+        # start_step is the number of noise levels to skip
+        #eg inf_steps = 50, strength = 1,  start step = 50 - (50 * 1) = 0, strength = 0, start_step = 50
+        start_step = self.num_inference_steps - int(self.num_inference_steps * strength)
+        self.timesteps = self.timesteps[start_step:]  #skip time_steps, if start_step = 50 8#
+        self.start_step = start_step  #50, in this case
+
+    def step(self, timestep: int, latents: torch.Tensor, model_output: torch.Tensor):
+        t = timestep #t
+        prev_t = self._get_previous_timestep(t)  #t-1
+
+        # 1. compute alphas, betas
+        alpha_prod_t = self.alphas_cumprod[t] #alpha_bar_t
+        alpha_prod_t_prev = self.alphas_cumprod[prev_t] if prev_t >= 0 else self.one  #alpha_bar_t-1
+        beta_prod_t = 1 - alpha_prod_t  #beta_bar_t
+        beta_prod_t_prev = 1 - alpha_prod_t_prev #beta_bar_t-1
+        current_alpha_t = alpha_prod_t / alpha_prod_t_prev  #alpha_t
+        current_beta_t = 1 - current_alpha_t  #beta_t
+
+        # 2. compute predicted original sample from predicted noise also called
+        # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
+        pred_original_sample = (latents - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)  #x_0 - gaussian noise
+
+        # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
+        # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
+        pred_original_sample_coeff = (alpha_prod_t_prev ** (0.5) * current_beta_t) / beta_prod_t  #coeff_x_0
+        current_sample_coeff = current_alpha_t ** (0.5) * beta_prod_t_prev / beta_prod_t  #coff_x_t
+
+        # 5. Compute predicted previous sample µ_t
+        # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
+        pred_prev_sample = pred_original_sample_coeff * pred_original_sample + current_sample_coeff * latents  #
+
+        # 6. Add noise
+        variance = 0
+        if t > 0:
+            device = model_output.device
+            noise = torch.randn(model_output.shape, generator=self.generator, device=device, dtype=model_output.dtype)
+            # Compute the variance as per formula (7) from https://arxiv.org/pdf/2006.11239.pdf
+            variance = (self._get_variance(t) ** 0.5) * noise
+
+        # sample from N(mu, sigma) = X can be obtained by X = mu + sigma * N(0, 1)
+        # the variable "variance" is already multiplied by the noise N(0, 1)
+        pred_prev_sample = pred_prev_sample + variance  #predicted xt-1
+
+        return pred_prev_sample
+
+    def add_noise(
+        self,
+        original_samples: torch.FloatTensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.FloatTensor:
+    #forward noisification
+    #q(xt | x_not) = N(xt; sqrt(alpha_cumprod); (1 - alpha_cumprod)I)
+        alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)  #alpha_bar
+        timesteps = timesteps.to(original_samples.device)
+
+        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5  #sqrt(alpha_bar_t)
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()  #flatten
+        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):  #for boardcasting
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5  #sqrt(1 - alpha_bar_t)
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        # Sample from q(x_t | x_0) as in equation (4) of https://arxiv.org/pdf/2006.11239.pdf
+        # Because N(mu, sigma) = X can be obtained by X = mu + sigma * N(0, 1)
+        # here mu = sqrt_alpha_prod * original_samples and sigma = sqrt_one_minus_alpha_prod
+        noise = torch.randn(original_samples.shape, generator=self.generator, device=original_samples.device, dtype=original_samples.dtype)  #noise
+        noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise  #noisy samples
+        return noisy_samples

unet.py

@@ -0,0 +1,183 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from helperUNET import SwitchSequential, UNET_AttentionBlock, UNET_ResidualBlock, Upsample
+
+
+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 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 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