diff --git a/images/dog.jpg b/images/dog.jpg new file mode 100644 index 0000000..4455579 Binary files /dev/null and b/images/dog.jpg differ diff --git a/sd/demo.ipynb b/sd/demo.ipynb new file mode 100644 index 0000000..51b53db --- /dev/null +++ b/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 +} diff --git a/sd/diffusion.py b/sd/diffusion.py index b12970b..5a887fa 100644 --- a/sd/diffusion.py +++ b/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) diff --git a/sd/model_converter.py b/sd/model_converter.py index 92dc25e..1a5f0d4 100644 --- a/sd/model_converter.py +++ b/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'] = {} diff --git a/sd/model_loader.py b/sd/model_loader.py index b574844..13c91d5 100644 --- a/sd/model_loader.py +++ b/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) diff --git a/sd/pipeline.py b/sd/pipeline.py index 0867360..2ba5d0a 100644 --- a/sd/pipeline.py +++ b/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,7 +124,7 @@ 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 @@ -132,20 +132,16 @@ def generate(prompt: str, uncond_prompt: str, input_image: None | torch.Tensor, 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) diff --git a/sd/vae_decoder.py b/sd/vae_decoder.py index 0817b7a..55001c4 100644 --- a/sd/vae_decoder.py +++ b/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: