{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/varadh/Project/DreamBooth_finetuning\n"
     ]
    }
   ],
   "source": [
    "!cd /home/varadh/Project/DreamBooth_finetuning\n",
    "!pwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "NVIDIA GeForce RTX 4070, 12282 MiB, 11626 MiB\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "print(torch.cuda.is_available())\n",
    "!nvidia-smi --query-gpu=name,memory.total,memory.free --format=csv,noheader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/varadh/Project/.venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Token will not been saved to git credential helper. Pass `add_to_git_credential=True` if you want to set the git credential as well.\n",
      "Token is valid (permission: write).\n",
      "Your token has been saved to /home/varadh/.cache/huggingface/token\n",
      "Login successful\n"
     ]
    }
   ],
   "source": [
    "import huggingface_hub\n",
    "huggingface_hub.login(\"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\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Weights will be saved at ./output_model/model_weights/run1\n"
     ]
    }
   ],
   "source": [
    "MODEL_NAME = \"runwayml/stable-diffusion-v1-5\"\n",
    "OUTPUT_DIR = \"./output_model/model_weights/run1\"\n",
    "\n",
    "print(f\"Weights will be saved at {OUTPUT_DIR}\")\n",
    "\n",
    "!mkdir -p $OUTPUT_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GPU Memory requirements\n",
    "\n",
    "Use the table below to choose the best flags based on your memory and speed requirements. Tested on Tesla T4 GPU.\n",
    "\n",
    "\n",
    "| `fp16` | `train_batch_size` | `gradient_accumulation_steps` | `gradient_checkpointing` | `use_8bit_adam` | GB VRAM usage | Speed (it/s) |\n",
    "| ---- | ------------------ | ----------------------------- | ----------------------- | --------------- | ---------- | ------------ |\n",
    "| fp16 | 1                  | 1                             | TRUE                    | TRUE            | 9.92       | 0.93         |\n",
    "| no   | 1                  | 1                             | TRUE                    | TRUE            | 10.08      | 0.42         |\n",
    "| fp16 | 2                  | 1                             | TRUE                    | TRUE            | 10.4       | 0.66         |\n",
    "| fp16 | 1                  | 1                             | FALSE                   | TRUE            | 11.17      | 1.14         |\n",
    "| no   | 1                  | 1                             | FALSE                   | TRUE            | 11.17      | 0.49         |\n",
    "| fp16 | 1                  | 2                             | TRUE                    | TRUE            | 11.56      | 1            |\n",
    "| fp16 | 2                  | 1                             | FALSE                   | TRUE            | 13.67      | 0.82         |\n",
    "| fp16 | 1                  | 2                             | FALSE                   | TRUE            | 13.7       | 0.83          |\n",
    "| fp16 | 1                  | 1                             | TRUE                    | FALSE           | 15.79      | 0.77         |\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define each class, the corresponding prompt and point to the directory they are stored in.\n",
    "concepts_list = [\n",
    "    {\n",
    "        \"instance_prompt\":      \"1girl, depth of field, holding weapon, glowstick, akali, cosplay photo\",\n",
    "        \"class_prompt\":         \"photo of a video game character\",\n",
    "        \"instance_data_dir\":    \"./data/AKALI\",\n",
    "        \"class_data_dir\":       \"./data/AKALI_instance\"\n",
    "    },\n",
    "    {\n",
    "        \"instance_prompt\":      \"1man, 3d, depth of field, weapon, gun, spy, secret, intel, traps\",\n",
    "        \"class_prompt\":         \"photo of a video game character\",\n",
    "        \"instance_data_dir\":    \"./data/Cypher\",\n",
    "        \"class_data_dir\":       \"./data/Cypher_instance\"\n",
    "    }\n",
    "]\n",
    "\n",
    "# `class_data_dir` contains regularization images\n",
    "import json\n",
    "import os\n",
    "for c in concepts_list:\n",
    "    os.makedirs(c[\"instance_data_dir\"], exist_ok=True)\n",
    "\n",
    "with open(\"concepts_list.json\", \"w\") as f:\n",
    "    json.dump(concepts_list, f, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/varadh/Project/.venv/lib/python3.11/site-packages/accelerate/accelerator.py:384: UserWarning: `log_with=tensorboard` was passed but no supported trackers are currently installed.\n",
      "  warnings.warn(f\"`log_with={log_with}` was passed but no supported trackers are currently installed.\")\n",
      "unet/diffusion_pytorch_model.safetensors not found\n",
      "/home/varadh/Project/.venv/lib/python3.11/site-packages/transformers/models/clip/feature_extraction_clip.py:28: FutureWarning: The class CLIPFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please use CLIPImageProcessor instead.\n",
      "  warnings.warn(\n",
      "You have disabled the safety checker for <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'> by passing `safety_checker=None`. Ensure that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered results in services or applications open to the public. Both the diffusers team and Hugging Face strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling it only for use-cases that involve analyzing network behavior or auditing its results. For more information, please have a look at https://github.com/huggingface/diffusers/pull/254 .\n",
      "12/06/2023 04:16:43 - INFO - __main__ - Number of class images to sample: 50.\n",
      "Generating class images: 100%|██████████████████| 50/50 [01:00<00:00,  1.22s/it]\n",
      "/home/varadh/Project/.venv/lib/python3.11/site-packages/diffusers/configuration_utils.py:203: FutureWarning: It is deprecated to pass a pretrained model name or path to `from_config`.If you were trying to load a scheduler, please use <class 'diffusers.schedulers.scheduling_ddpm.DDPMScheduler'>.from_pretrained(...) instead. Otherwise, please make sure to pass a configuration dictionary instead. This functionality will be removed in v1.0.0.\n",
      "  deprecate(\"config-passed-as-path\", \"1.0.0\", deprecation_message, standard_warn=False)\n",
      "Caching latents: 100%|████████████████████████| 100/100 [00:08<00:00, 11.98it/s]\n",
      "12/06/2023 04:17:56 - INFO - __main__ - ***** Running training *****\n",
      "12/06/2023 04:17:56 - INFO - __main__ -   Num examples = 100\n",
      "12/06/2023 04:17:56 - INFO - __main__ -   Num batches each epoch = 100\n",
      "12/06/2023 04:17:56 - INFO - __main__ -   Num Epochs = 8\n",
      "12/06/2023 04:17:56 - INFO - __main__ -   Instantaneous batch size per device = 1\n",
      "12/06/2023 04:17:56 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 1\n",
      "12/06/2023 04:17:56 - INFO - __main__ -   Gradient Accumulation steps = 1\n",
      "12/06/2023 04:17:56 - INFO - __main__ -   Total optimization steps = 800\n",
      "Steps: 100%|█████████████| 800/800 [04:16<00:00,  3.14it/s, loss=0.316, lr=1e-6]unet/diffusion_pytorch_model.safetensors not found\n",
      "You have disabled the safety checker for <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'> by passing `safety_checker=None`. Ensure that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered results in services or applications open to the public. Both the diffusers team and Hugging Face strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling it only for use-cases that involve analyzing network behavior or auditing its results. For more information, please have a look at https://github.com/huggingface/diffusers/pull/254 .\n",
      "\n",
      "Generating samples:   0%|                                 | 0/4 [00:00<?, ?it/s]\u001b[A\n",
      "Generating samples:  25%|██████▎                  | 1/4 [00:01<00:04,  1.48s/it]\u001b[A\n",
      "Generating samples:  50%|████████████▌            | 2/4 [00:02<00:02,  1.33s/it]\u001b[A\n",
      "Generating samples:  75%|██████████████████▊      | 3/4 [00:03<00:01,  1.28s/it]\u001b[A\n",
      "Generating samples: 100%|█████████████████████████| 4/4 [00:05<00:00,  1.29s/it]\u001b[A\n",
      "[*] Weights saved at ./output_model/model_weights/run1/800\n",
      "Steps: 100%|█████████████| 800/800 [04:26<00:00,  3.01it/s, loss=0.316, lr=1e-6]\n"
     ]
    }
   ],
   "source": [
    "# Start fine-tuning the model\n",
    "import bitsandbytes\n",
    "import scipy\n",
    "\n",
    "\n",
    "\n",
    "!python3 train_dreambooth.py \\\n",
    "  --pretrained_model_name_or_path=$MODEL_NAME \\\n",
    "  --pretrained_vae_name_or_path=\"stabilityai/sd-vae-ft-mse\" \\\n",
    "  \\\n",
    "  --output_dir=$OUTPUT_DIR \\\n",
    "  \\\n",
    "  --revision=\"fp16\" \\\n",
    "  --with_prior_preservation --prior_loss_weight=1.0 \\\n",
    "  \\\n",
    "  --seed=14 \\\n",
    "  \\\n",
    "  --resolution=512 \\\n",
    "  \\\n",
    "  --train_batch_size=1 \\\n",
    "  --train_text_encoder \\\n",
    "  --mixed_precision=\"fp16\" \\\n",
    "  --use_8bit_adam\\\n",
    "  --gradient_accumulation_steps=1 \\\n",
    "  \\\n",
    "  --learning_rate=1e-6 \\\n",
    "  --lr_scheduler=\"constant\" \\\n",
    "  --lr_warmup_steps=0 \\\n",
    "  \\\n",
    "  --num_class_images=50 \\\n",
    "  --sample_batch_size=1 \\\n",
    "  --max_train_steps=800 \\\n",
    "  --save_interval=400 \\\n",
    "  --save_sample_prompt=\"3d, depth of field, weapon, gun, spy, secret, intel, traps\" \\\n",
    "  --concepts_list=\"concepts_list.json\"\\\n",
    "\n",
    "\n",
    "# Reduce the `--save_interval` to lower than `--max_train_steps` to save weights from intermediate steps.\n",
    "# `--save_sample_prompt` can be same as `--instance_prompt` to generate intermediate samples (saved along with weights in samples directory)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1082292 characters]",
      "text/plain": [
       "<Figure size 1600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Run to generate a grid of preview images from the last saved weights.\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n",
    "\n",
    "weights_folder = OUTPUT_DIR\n",
    "folders = sorted([f for f in os.listdir(weights_folder) if f != \"0\"], key=lambda x: int(x))\n",
    "\n",
    "row = len(folders)\n",
    "col = len(os.listdir(os.path.join(weights_folder, folders[0], \"samples\")))\n",
    "scale = 4\n",
    "fig, axes = plt.subplots(row, col, figsize=(col*scale, row*scale), gridspec_kw={'hspace': 0, 'wspace': 0})\n",
    "\n",
    "for i, folder in enumerate(folders):\n",
    "    folder_path = os.path.join(weights_folder, folder)\n",
    "    image_folder = os.path.join(folder_path, \"samples\")\n",
    "    images = [f for f in os.listdir(image_folder)]\n",
    "    for j, image in enumerate(images):\n",
    "        if row == 1:\n",
    "            currAxes = axes[j]\n",
    "        else:\n",
    "            currAxes = axes[i, j]\n",
    "        if i == 0:\n",
    "            currAxes.set_title(f\"Image {j}\")\n",
    "        if j == 0:\n",
    "            currAxes.text(-0.1, 0.5, folder, rotation=0, va='center', ha='center', transform=currAxes.transAxes)\n",
    "        image_path = os.path.join(image_folder, image)\n",
    "        img = mpimg.imread(image_path)\n",
    "        currAxes.imshow(img, cmap='gray')\n",
    "        currAxes.axis('off')\n",
    "        \n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/varadh/Project/.venv/lib/python3.11/site-packages/transformers/models/clip/feature_extraction_clip.py:28: FutureWarning: The class CLIPFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please use CLIPImageProcessor instead.\n",
      "  warnings.warn(\n",
      "You have disabled the safety checker for <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'> by passing `safety_checker=None`. Ensure that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered results in services or applications open to the public. Both the diffusers team and Hugging Face strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling it only for use-cases that involve analyzing network behavior or auditing its results. For more information, please have a look at https://github.com/huggingface/diffusers/pull/254 .\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch import autocast\n",
    "from diffusers import StableDiffusionPipeline, DDIMScheduler\n",
    "from IPython.display import display\n",
    "\n",
    "model_path = \"/home/varadh/Project/DreamBooth_finetuning/output_model/model_weights/run1/800\"\n",
    "pipe = StableDiffusionPipeline.from_pretrained(model_path, safety_checker=None, torch_dtype=torch.float16).to(\"cuda\")\n",
    "pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)\n",
    "pipe.enable_xformers_memory_efficient_attention()\n",
    "g_cuda = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7f413d4542f0>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g_cuda = torch.Generator(device='cuda')\n",
    "seed = 14\n",
    "g_cuda.manual_seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50/50 [00:03<00:00, 16.28it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 63132 characters]",
      "image/png": "[encoded data, 602056 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=512x512>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Specify the prompt\n",
    "prompt = \"3d, depth of field, weapon, gun, spy, secret, intel, traps\"  # replace this with your desired prompt\n",
    "\n",
    "# Set parameters (if needed)\n",
    "num_inference_steps = 50\n",
    "guidance_scale = 7.5      \n",
    "\n",
    "# Generate the image\n",
    "with autocast(\"cuda\"):\n",
    "    generated_image = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=g_cuda).images[0]\n",
    "\n",
    "# Display the image\n",
    "display(generated_image)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Baseline SD model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`text_config_dict` is provided which will be used to initialize `CLIPTextConfig`. The value `text_config[\"id2label\"]` will be overriden.\n",
      "`text_config_dict` is provided which will be used to initialize `CLIPTextConfig`. The value `text_config[\"bos_token_id\"]` will be overriden.\n",
      "`text_config_dict` is provided which will be used to initialize `CLIPTextConfig`. The value `text_config[\"eos_token_id\"]` will be overriden.\n",
      "100%|██████████| 50/50 [00:02<00:00, 18.04it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 40756 characters]",
      "image/png": "[encoded data, 547100 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=512x512>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from diffusers import StableDiffusionPipeline\n",
    "import torch\n",
    "\n",
    "prompt = \"3d, depth of field, weapon, gun, spy, secret, intel, traps\"\n",
    "\n",
    "model_id = \"runwayml/stable-diffusion-v1-5\"\n",
    "pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)\n",
    "pipe = pipe.to(\"cuda\")\n",
    "\n",
    "image = pipe(prompt).images[0]  \n",
    "display(image)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
