Dev

sd/demo.ipynb

@@ -0,0 +1,125 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "83e7e681",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import model_loader\n",
+    "import pipeline\n",
+    "from PIL import Image\n",
+    "from transformers import CLIPTokenizer\n",
+    "import torch\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "8e730da6",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'cuda'"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ALLOW_CUDA = True\n",
+    "DEVICE = 'cuda' if torch.cuda.is_available() and ALLOW_CUDA else 'cpu'\n",
+    "DEVICE"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cea60e6f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizer = CLIPTokenizer(vocab_file='../data/vocab.json', merges_file='../data/merges.txt')\n",
+    "model_file = '../data/v1-5-pruned-emaonly.ckpt'\n",
+    "models = model_loader.preload_models_from_standard_weights(model_file, DEVICE)\n",
+    "\n",
+    "## TEXT TO IMAGE\n",
+    "\n",
+    "prompt = \"A dog with sunglasses, wearing comfy hat, looking at camera, highly detailed, ultra sharp, cinematic, 100mm lens, 8k resolution.\"\n",
+    "# prompt = \"A cat stretching on the floor, highly detailed, ultra sharp, cinematic, 100mm lens, 8k resolution.\"\n",
+    "uncond_prompt = \"\"  # Also known as negative prompt\n",
+    "do_cfg = True\n",
+    "cfg_scale = 8  # min: 1, max: 14\n",
+    "\n",
+    "## IMAGE TO IMAGE\n",
+    "\n",
+    "input_image = None\n",
+    "# Comment to disable image to image\n",
+    "image_path = \"../images/dog.jpg\"\n",
+    "# input_image = Image.open(image_path)\n",
+    "# Higher values means more noise will be added to the input image, so the result will further from the input image.\n",
+    "# Lower values means less noise is added to the input image, so output will be closer to the input image.\n",
+    "strength = 0.9\n",
+    "\n",
+    "## SAMPLER\n",
+    "\n",
+    "sampler = \"ddpm\"\n",
+    "num_inference_steps = 50\n",
+    "seed = 42\n",
+    "\n",
+    "output_image = pipeline.generate(\n",
+    "    prompt=prompt,\n",
+    "    uncond_prompt=uncond_prompt,\n",
+    "    input_image=input_image,\n",
+    "    strength=strength,\n",
+    "    do_cfg=do_cfg,\n",
+    "    cfg_scale=cfg_scale,\n",
+    "    sampler_name=sampler,\n",
+    "    n_inference_steps=num_inference_steps,\n",
+    "    seed=seed,\n",
+    "    models=models,\n",
+    "    device=DEVICE,\n",
+    "    idle_device=\"cpu\",\n",
+    "    tokenizer=tokenizer,\n",
+    ")\n",
+    "\n",
+    "# Combine the input image and the output image into a single image.\n",
+    "Image.fromarray(output_image)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4bfedbd6",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "rl_venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

sd/diffusion.py

@@ -82,7 +82,7 @@ def __init__(self, n_head: int, n_embd: int, d_context=768):
         self.layernorm_1 = nn.LayerNorm(normalized_shape=channels)
         self.attention_1 = SelfAttention(n_heads=n_head, d_embed=channels, in_proj_bias=False)
         self.layernorm_2 = nn.LayerNorm(normalized_shape=channels)
-        self.attention_2 = CrossAttention(n_heads=n_head, d_cross=channels, d_embed=d_context, in_proj_bias=False)
+        self.attention_2 = CrossAttention(n_heads=n_head, d_embed=channels, d_cross=d_context, in_proj_bias=False)
         self.layernorm_3 = nn.LayerNorm(normalized_shape=channels)
         self.linear_geglu_1  = nn.Linear(in_features=channels, out_features=4 * channels * 2)
         self.linear_geglu_2 = nn.Linear(in_features=4 * channels, out_features=channels)
@@ -168,12 +168,12 @@ def forward(self, x:torch.Tensor, context:torch.Tensor):
         # (Batch_Size, Features, Height, Width) + (Batch_Size, Features, Height, Width) -> (Batch_Size, Features, Height, Width)
         return self.conv_output(x) + residue_long
 
-class SwitchSequential(nn.Module):
+class SwitchSequential(nn.Sequential):
     def forward(self, x: torch.Tensor, context: torch.Tensor, time: torch.Tensor) -> torch.Tensor:
         for layer in self:
             if isinstance(layer, UNET_AttentionBlock):
                 x = layer(x, context)
-            elif isinstance(x, UNET_ResidualBlock):
+            elif isinstance(layer, UNET_ResidualBlock):
                 x = layer(x, time)
             else:
                 x = layer(x)
@@ -192,7 +192,7 @@ def forward(self, x):
 
 
 class UNET(nn.Module):
-    def __init__(self,):
+    def __init__(self):
         super().__init__()
 
         self.encoders = nn.ModuleList([
@@ -307,6 +307,7 @@ def forward(self, x):
 class Diffusion(nn.Module):
 
     def __init__(self):
+        super().__init__()
         self.time_embedding = TimeEmbedding(n_embed=320)
         self.unet = UNET()
         self.final = UNET_OutputLayer(in_channels=320, out_channels=4)

sd/model_converter.py

@@ -2,7 +2,7 @@
 
 def load_from_standard_weights(input_file: str, device: str) -> dict[str, torch.Tensor]:
     # Taken from: https://github.com/kjsman/stable-diffusion-pytorch/issues/7#issuecomment-1426839447
-    original_model = torch.load(input_file, map_location=device, weights_only = False)["state_dict"]
+    original_model = torch.load(input_file, map_location='cpu', weights_only = False)["state_dict"]
 
     converted = {}
     converted['diffusion'] = {}

sd/model_loader.py

@@ -7,7 +7,7 @@
 
 def preload_models_from_standard_weights(ckpt_path: str, device: str):
     state_dict = model_converter.load_from_standard_weights(ckpt_path, device)
-
+    device = 'cpu'
     encoder = VAE_Encoder().to(device)
     encoder.load_state_dict(state_dict['encoder'], strict=True)
 

sd/pipeline.py

@@ -40,15 +40,15 @@ def generate(prompt: str, uncond_prompt: str, input_image: None | torch.Tensor,
             raise ValueError('Strength must be between 0 and 1')
 
         if idle_device:
-            to_idle: lambda x: x.to(idle_device)
+            to_idle = lambda x: x.to(idle_device)
         else:
-            to_idle: lambda x: x
+            to_idle = lambda x: x
 
         generator = torch.Generator(device=device)
         if seed is None:
             generator.seed()
         else:
-            generator.manual_seed(seed=seed)
+            generator.manual_seed(seed)
 
         clip = models['clip']
         clip.to(device) 
@@ -57,7 +57,7 @@ def generate(prompt: str, uncond_prompt: str, input_image: None | torch.Tensor,
             # convert the prompt into tokens
             cond_tokens = tokenizer.batch_encode_plus([prompt], padding='max_length', max_length=77).input_ids
             # convert input ids into tensor (batch_size, seq len)
-            cond_tokens = torch.Tensor(cond_tokens, dtype=torch.long, device=device)
+            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)
@@ -124,28 +124,24 @@ def generate(prompt: str, uncond_prompt: str, input_image: None | torch.Tensor,
         timesteps = tqdm(sampler.timesteps)
         for i, timestep in enumerate(timesteps):
             # (1, 320)
-            time_embedding = get_time_embedding(timestep=timestep).to_device
+            time_embedding = get_time_embedding(timestep=timestep).to(device)
 
             # (batch_szie, 4, latent_height, latent_width)
             model_input = latents
 
             if do_cfg:
                 # (Batch_Size, 4, Latents_Height, Latents_Width) -> (2 * Batch_Size, 4, Latents_Height, Latents_Width)
                 model_input = model_input.repeat(2, 1, 1, 1)
-
-            else:
-                # model_output is the predicted noise
-                # (Batch_Size, 4, Latents_Height, Latents_Width) -> (Batch_Size, 4, Latents_Height, Latents_Width)
-                model_output = diffusion(model_input, context, time_embedding)
+            # model_output is the predicted noise
+            # (Batch_Size, 4, Latents_Height, Latents_Width) -> (Batch_Size, 4, Latents_Height, Latents_Width)
+            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
-
-            else:
-                # (Batch_Size, 4, Latents_Height, Latents_Width) -> (Batch_Size, 4, Latents_Height, Latents_Width)
-                latents = sampler.step(timestep, latents, model_output)
+            # (Batch_Size, 4, Latents_Height, Latents_Width) -> (Batch_Size, 4, Latents_Height, Latents_Width)
+            latents = sampler.step(timestep, latents, model_output)
 
         to_idle(diffusion)
 

sd/vae_decoder.py

@@ -48,7 +48,7 @@ def __init__(self, in_channels, out_channels):
         self.groupnorm_1 = nn.GroupNorm(num_groups=32, num_channels=in_channels)
         self.conv_1 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, padding=1)
 
-        self.groupnorm_2 = nn.GroupNorm(num_groups=32, out_channels=out_channels)
+        self.groupnorm_2 = nn.GroupNorm(num_groups=32, num_channels=out_channels)
         self.conv_2 = nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, padding=1)
 
         if in_channels == out_channels: