{ "cells": [ { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "2024-09-10 11:23:42.358210: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n", "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n", "2024-09-10 11:23:43.031780: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n" ] } ], "source": [ "from base64 import b64encode\n", "\n", "import numpy\n", "import torch\n", "from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel\n", "from huggingface_hub import notebook_login\n", "\n", "# For video display:\n", "from IPython.display import HTML\n", "from matplotlib import pyplot as plt\n", "from pathlib import Path\n", "from PIL import Image\n", "from torch import autocast\n", "from torchvision import transforms as tfms\n", "from tqdm.auto import tqdm\n", "from transformers import CLIPTextModel, CLIPTokenizer, logging\n", "import os\n", "import os\n", "os.environ['HF_HOME'] = '/raid/users/mohammadibrahim-st/ModelCache'\n", "torch.manual_seed(1)\n", "if not (Path.home()/'.cache/huggingface'/'token').exists(): notebook_login()\n", "\n", "# Supress some unnecessary warnings when loading the CLIPTextModel\n", "logging.set_verbosity_error()\n", "\n", "# Set device\n", "torch_device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n", "if \"mps\" == torch_device: os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = \"1\"" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/home/mohammadibrahim-st/.local/lib/python3.8/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n", " warnings.warn(\n" ] } ], "source": [ "# Load the autoencoder model which will be used to decode the latents into image space.\n", "import os\n", "os.environ[\"https_proxy\"] = \"http://185.46.212.90:80\"\n", "os.environ[\"http_proxy\"] = \"http://185.46.212.90:80\"\n", "vae = AutoencoderKL.from_pretrained(\"CompVis/stable-diffusion-v1-4\", subfolder=\"vae\")\n", "\n", "# Load the tokenizer and text encoder to tokenize and encode the text.\n", "tokenizer = CLIPTokenizer.from_pretrained(\"openai/clip-vit-large-patch14\")\n", "text_encoder = CLIPTextModel.from_pretrained(\"openai/clip-vit-large-patch14\")\n", "\n", "# The UNet model for generating the latents.\n", "unet = UNet2DConditionModel.from_pretrained(\"CompVis/stable-diffusion-v1-4\", subfolder=\"unet\")\n", "\n", "# The noise scheduler\n", "scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule=\"scaled_linear\", num_train_timesteps=1000)\n", "\n", "# To the GPU we go!\n", "vae = vae.to(torch_device)\n", "text_encoder = text_encoder.to(torch_device)\n", "unet = unet.to(torch_device);" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "def pil_to_latent(input_im):\n", " # Single image -> single latent in a batch (so size 1, 4, 64, 64)\n", " with torch.no_grad():\n", " latent = vae.encode(tfms.ToTensor()(input_im).unsqueeze(0).to(torch_device)*2-1) # Note scaling\n", " return 0.18215 * latent.latent_dist.sample()\n", "\n", "def latents_to_pil(latents):\n", " # bath of latents -> list of images\n", " latents = (1 / 0.18215) * latents\n", " with torch.no_grad():\n", " image = vae.decode(latents).sample\n", " image = (image / 2 + 0.5).clamp(0, 1)\n", " image = image.detach().cpu().permute(0, 2, 3, 1).numpy()\n", " images = (image * 255).round().astype(\"uint8\")\n", " pil_images = [Image.fromarray(image) for image in images]\n", " # pil_images[0].save('/raid/users/mohammadibrahim-st/TSAI/Assignment24/Depth/mouseseed64bright.png')\n", " return pil_images" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "def set_timesteps(scheduler, num_inference_steps):\n", " scheduler.set_timesteps(num_inference_steps)\n", " scheduler.timesteps = scheduler.timesteps.to(torch.float32)" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "def brightness_loss(images, target_brightness=0.9):\n", " # Convert images to grayscale to calculate brightness\n", " grayscale_images = images.mean(dim=1, keepdim=True)\n", " error = torch.abs(grayscale_images - target_brightness).mean()\n", " return error\n" ] }, { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [], "source": [ "def generate_with_embs(text_embeddings):\n", " height = 512 # default height of Stable Diffusion\n", " width = 512 # default width of Stable Diffusion\n", " num_inference_steps = 30 # Number of denoising steps\n", " guidance_scale = 7.5 # Scale for classifier-free guidance\n", " generator = torch.manual_seed(164) # Seed generator to create the inital latent noise\n", " batch_size = 1\n", " blue_loss_scale=200\n", "\n", " max_length = text_input.input_ids.shape[-1]\n", " uncond_input = tokenizer(\n", " [\"\"] * batch_size, padding=\"max_length\", max_length=max_length, return_tensors=\"pt\"\n", " )\n", " with torch.no_grad():\n", " uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]\n", " text_embeddings = torch.cat([uncond_embeddings, text_embeddings])\n", "\n", " # Prep Scheduler\n", " set_timesteps(scheduler, num_inference_steps)\n", "\n", " # Prep latents\n", " latents = torch.randn(\n", " (batch_size, unet.in_channels, height // 8, width // 8),\n", " generator=generator,\n", " )\n", " latents = latents.to(torch_device)\n", " latents = latents * scheduler.init_noise_sigma\n", "\n", " # Loop\n", " for i, t in tqdm(enumerate(scheduler.timesteps), total=len(scheduler.timesteps)):\n", " # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.\n", " latent_model_input = torch.cat([latents] * 2)\n", " sigma = scheduler.sigmas[i]\n", " latent_model_input = scheduler.scale_model_input(latent_model_input, t)\n", "\n", " # predict the noise residual\n", " with torch.no_grad():\n", " noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)[\"sample\"]\n", "\n", " # perform guidance\n", " noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n", " noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)\n", "\n", " # compute the previous noisy sample x_t -> x_t-1\n", " if i%5 == 0:\n", " # Requires grad on the latents\n", " latents = latents.detach().requires_grad_()\n", "\n", " # Get the predicted x0:\n", " latents_x0 = latents - sigma * noise_pred\n", " # latents_x0 = scheduler.step(noise_pred, t, latents).pred_original_sample\n", "\n", " # Decode to image space\n", " denoised_images = vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5 # range (0, 1)\n", "\n", " # Calculate loss\n", " loss = brightness_loss(denoised_images) * blue_loss_scale\n", "\n", " # Occasionally print it out\n", " if i%10==0:\n", " print(i, 'loss:', loss.item())\n", "\n", " # Get gradient\n", " cond_grad = torch.autograd.grad(loss, latents)[0]\n", "\n", " # Modify the latents based on this gradient\n", " latents = latents.detach() - cond_grad * sigma**2\n", "\n", " # Now step with scheduler\n", " latents = scheduler.step(noise_pred, t, latents).prev_sample\n", " return latents_to_pil(latents)[0]" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [ "def build_causal_attention_mask(bsz, seq_len, dtype):\n", " # lazily create causal attention mask, with full attention between the vision tokens\n", " # pytorch uses additive attention mask; fill with -inf\n", " mask = torch.empty(bsz, seq_len, seq_len, dtype=dtype)\n", " mask.fill_(torch.tensor(torch.finfo(dtype).min))\n", " mask.triu_(1) # zero out the lower diagonal\n", " mask = mask.unsqueeze(1) # expand mask\n", " return mask\n", "def get_output_embeds(input_embeddings):\n", " # CLIP's text model uses causal mask, so we prepare it here:\n", " bsz, seq_len = input_embeddings.shape[:2]\n", " causal_attention_mask = build_causal_attention_mask(bsz, seq_len, dtype=input_embeddings.dtype)\n", "\n", " # Getting the output embeddings involves calling the model with passing output_hidden_states=True\n", " # so that it doesn't just return the pooled final predictions:\n", " encoder_outputs = text_encoder.text_model.encoder(\n", " inputs_embeds=input_embeddings,\n", " attention_mask=None, # We aren't using an attention mask so that can be None\n", " causal_attention_mask=causal_attention_mask.to(torch_device),\n", " output_attentions=None,\n", " output_hidden_states=True, # We want the output embs not the final output\n", " return_dict=None,\n", " )\n", "\n", " # We're interested in the output hidden state only\n", " output = encoder_outputs[0]\n", "\n", " # There is a final layer norm we need to pass these through\n", " output = text_encoder.text_model.final_layer_norm(output)\n", "\n", " # And now they're ready!\n", " return output\n", "\n", "# out_embs_test = get_output_embeds(input_embeddings) # Feed through the model with our new function\n", "# print(out_embs_test.shape) # Check the output shape\n", "# out_embs_test # Inspect the output" ] }, { "cell_type": "code", "execution_count": 19, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "dict_keys([''])\n", "dict_keys([''])\n", "dict_keys([''])\n", "dict_keys(['oil_style'])\n", "dict_keys([''])\n" ] } ], "source": [ "current_directory = os.path.dirname(__file__)\n", "\n", "# Construct the paths dynamically\n", "birb_embed = torch.load(os.path.join(current_directory, 'Depth', 'learned_embeds.bin'))\n", "\n", "birb_embedjerry = torch.load(os.path.join(current_directory, 'Jerry mouse', 'learned_embeds.bin'))\n", "\n", "birb_embedmobius = torch.load(os.path.join(current_directory, 'Mobius', 'learned_embeds.bin'))\n", "\n", "birb_embedoilpaint = torch.load(os.path.join(current_directory, 'Oil paint', 'learned_embeds.bin'))\n", "\n", "birb_embedpolygon = torch.load(os.path.join(current_directory, 'Polygon', 'learned_embeds.bin'))\n", "\n", "print(birb_embed.keys())\n", "print(birb_embedjerry.keys())\n", "print(birb_embedmobius.keys())\n", "print(birb_embedoilpaint.keys())\n", "print(birb_embedpolygon.keys())\n", "\n" ] }, { "cell_type": "code", "execution_count": 12, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/tmp/ipykernel_1959221/3185483617.py:23: FutureWarning: Accessing config attribute `in_channels` directly via 'UNet2DConditionModel' object attribute is deprecated. Please access 'in_channels' over 'UNet2DConditionModel's config object instead, e.g. 'unet.config.in_channels'.\n", " (batch_size, unet.in_channels, height // 8, width // 8),\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "7a30a6f43b3442668ab8569e9e7a8fda", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/30 [00:00" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "prompt = 'A mouse in the style of puppy'\n", "token_emb_layer = text_encoder.text_model.embeddings.token_embedding\n", "pos_emb_layer = text_encoder.text_model.embeddings.position_embedding\n", "position_ids = text_encoder.text_model.embeddings.position_ids[:, :77]\n", "position_embeddings = pos_emb_layer(position_ids)\n", "# Tokenize\n", "text_input = tokenizer(prompt, padding=\"max_length\", max_length=tokenizer.model_max_length, truncation=True, return_tensors=\"pt\")\n", "input_ids = text_input.input_ids.to(torch_device)\n", "\n", "# Get token embeddings\n", "token_embeddings = token_emb_layer(input_ids)\n", "\n", "# The new embedding - our special birb word\n", "replacement_token_embedding = birb_embed[''].to(torch_device)\n", "\n", "# Insert this into the token embeddings\n", "token_embeddings[0, torch.where(input_ids[0]==6829)] = replacement_token_embedding.to(torch_device)\n", "\n", "# Combine with pos embs\n", "input_embeddings = token_embeddings + position_embeddings\n", "\n", "# Feed through to get final output embs\n", "modified_output_embeddings = get_output_embeds(input_embeddings)\n", "\n", "# And generate an image with this:\n", "a = generate_with_embs(modified_output_embeddings)\n", "display(a)" ] }, { "cell_type": "code", "execution_count": 21, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Running on local URL: http://127.0.0.1:7865\n", "IMPORTANT: You are using gradio version 4.24.0, however version 4.29.0 is available, please upgrade.\n", "--------\n", "Running on public URL: https://25baf7d4fc18ef497d.gradio.live\n", "\n", "This share link expires in 72 hours. For free permanent hosting and GPU upgrades, run `gradio deploy` from Terminal to deploy to Spaces (https://huggingface.co/spaces)\n" ] }, { "data": { "text/html": [ "
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [] }, "execution_count": 21, "metadata": {}, "output_type": "execute_result" }, { "name": "stderr", "output_type": "stream", "text": [ "/tmp/ipykernel_1959221/3185483617.py:23: FutureWarning: Accessing config attribute `in_channels` directly via 'UNet2DConditionModel' object attribute is deprecated. Please access 'in_channels' over 'UNet2DConditionModel's config object instead, e.g. 'unet.config.in_channels'.\n", " (batch_size, unet.in_channels, height // 8, width // 8),\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "9762c8b63e5b49bba277e7dc2ca9d99e", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/30 [00:00'),\n", " \"Jerry mouse\": (birb_embedjerry, ''),\n", " \"Mobius\": (birb_embedmobius, ''),\n", " \"Oil paint\": (birb_embedoilpaint, 'oil_style'),\n", " \"Polygon\": (birb_embedpolygon, '')\n", " }\n", " \n", " token_emb_layer = text_encoder.text_model.embeddings.token_embedding\n", " pos_emb_layer = text_encoder.text_model.embeddings.position_embedding\n", " position_ids = text_encoder.text_model.embeddings.position_ids[:, :77]\n", " position_embeddings = pos_emb_layer(position_ids)\n", "\n", " # Tokenize\n", " text_input = tokenizer(prompt, padding=\"max_length\", max_length=tokenizer.model_max_length, truncation=True, return_tensors=\"pt\")\n", " input_ids = text_input.input_ids.to(torch_device)\n", "\n", " # Get token embeddings\n", " token_embeddings = token_emb_layer(input_ids)\n", "\n", " # Select the appropriate birb embedding and key based on user input\n", " selected_embedding_file, embedding_key = embedding_dict[selected_embedding]\n", " replacement_token_embedding = selected_embedding_file[embedding_key].to(torch_device)\n", "\n", " # Insert this into the token embeddings\n", " token_embeddings[0, torch.where(input_ids[0] == 6829)] = replacement_token_embedding.to(torch_device)\n", "\n", " # Combine with pos embs\n", " input_embeddings = token_embeddings + position_embeddings\n", "\n", " # Feed through to get final output embs\n", " modified_output_embeddings = get_output_embeds(input_embeddings)\n", "\n", " # Generate an image with this and return it\n", " generated_image = generate_with_embs(modified_output_embeddings)\n", " return generated_image\n", "\n", "# Define options for the dropdown\n", "embedding_options = [\"Depth\", \"Jerry mouse\", \"Mobius\", \"Oil paint\", \"Polygon\"]\n", "\n", "# Create Gradio interface\n", "iface = gr.Interface(\n", " fn=generate_image, \n", " inputs=[\n", " \"text\", \n", " gr.Dropdown(choices=embedding_options, label=\"Select Style\")\n", " ],\n", " outputs=\"image\",\n", " title=\"Image Generation App (Please use the word 'puppy' in the prompt)\"\n", ")\n", "\n", "iface.launch(share=True)\n" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "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.8.10" } }, "nbformat": 4, "nbformat_minor": 2 }