diff --git "a/Learned.ipynb" "b/Learned.ipynb" new file mode 100644--- /dev/null +++ "b/Learned.ipynb" @@ -0,0 +1,537 @@ +{ + "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 +}