{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "acab479f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "===================================BUG REPORT===================================\n",
      "Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
      "================================================================================\n"
     ]
    }
   ],
   "source": [
    "import argparse\n",
    "import gc\n",
    "import hashlib\n",
    "import itertools\n",
    "import logging\n",
    "import math\n",
    "import os\n",
    "import threading\n",
    "import warnings\n",
    "from pathlib import Path\n",
    "from typing import Optional\n",
    "import psutil\n",
    "import json\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.checkpoint\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "import datasets\n",
    "import diffusers\n",
    "import transformers\n",
    "from accelerate import Accelerator\n",
    "from accelerate.logging import get_logger\n",
    "from accelerate.utils import set_seed\n",
    "from diffusers import AutoencoderKL, DDPMScheduler, DiffusionPipeline, UNet2DConditionModel\n",
    "from diffusers import DDPMScheduler, PNDMScheduler, StableDiffusionPipeline\n",
    "from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker\n",
    "from diffusers.optimization import get_scheduler\n",
    "from diffusers.utils import check_min_version\n",
    "from diffusers.utils.import_utils import is_xformers_available\n",
    "from huggingface_hub import HfFolder, Repository, whoami\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "from tqdm.auto import tqdm\n",
    "from transformers import AutoTokenizer, PretrainedConfig, CLIPFeatureExtractor\n",
    "from peft import PeftModel, LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict\n",
    "\n",
    "# Will error if the minimal version of diffusers is not installed. Remove at your own risks.\n",
    "check_min_version(\"0.10.0.dev0\")\n",
    "\n",
    "logger = get_logger(__name__)\n",
    "\n",
    "\n",
    "MODEL_NAME = \"CompVis/stable-diffusion-v1-4\"  # \"stabilityai/stable-diffusion-2-1-base\"\n",
    "INSTANCE_PROMPT = \"a photo of sks dog\"\n",
    "base_path = \"/home/sourab/temp/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "06cfd506",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_lora_sd_pipeline(\n",
    "    ckpt_dir, base_model_name_or_path=None, dtype=torch.float16, device=\"cuda\", adapter_name=\"default\"\n",
    "):\n",
    "    unet_sub_dir = os.path.join(ckpt_dir, \"unet\")\n",
    "    text_encoder_sub_dir = os.path.join(ckpt_dir, \"text_encoder\")\n",
    "    if os.path.exists(text_encoder_sub_dir) and base_model_name_or_path is None:\n",
    "        config = LoraConfig.from_pretrained(text_encoder_sub_dir)\n",
    "        base_model_name_or_path = config.base_model_name_or_path\n",
    "\n",
    "    if base_model_name_or_path is None:\n",
    "        raise ValueError(\"Please specify the base model name or path\")\n",
    "\n",
    "    pipe = StableDiffusionPipeline.from_pretrained(\n",
    "        base_model_name_or_path, torch_dtype=dtype, requires_safety_checker=False\n",
    "    ).to(device)\n",
    "    pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)\n",
    "\n",
    "    if os.path.exists(text_encoder_sub_dir):\n",
    "        pipe.text_encoder = PeftModel.from_pretrained(\n",
    "            pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name\n",
    "        )\n",
    "\n",
    "    if dtype in (torch.float16, torch.bfloat16):\n",
    "        pipe.unet.half()\n",
    "        pipe.text_encoder.half()\n",
    "\n",
    "    pipe.to(device)\n",
    "    return pipe\n",
    "\n",
    "\n",
    "def load_adapter(pipe, ckpt_dir, adapter_name):\n",
    "    unet_sub_dir = os.path.join(ckpt_dir, \"unet\")\n",
    "    text_encoder_sub_dir = os.path.join(ckpt_dir, \"text_encoder\")\n",
    "    pipe.unet.load_adapter(unet_sub_dir, adapter_name=adapter_name)\n",
    "    if os.path.exists(text_encoder_sub_dir):\n",
    "        pipe.text_encoder.load_adapter(text_encoder_sub_dir, adapter_name=adapter_name)\n",
    "\n",
    "\n",
    "def set_adapter(pipe, adapter_name):\n",
    "    pipe.unet.set_adapter(adapter_name)\n",
    "    if isinstance(pipe.text_encoder, PeftModel):\n",
    "        pipe.text_encoder.set_adapter(adapter_name)\n",
    "\n",
    "\n",
    "def merging_lora_with_base(pipe, ckpt_dir, adapter_name=\"default\"):\n",
    "    unet_sub_dir = os.path.join(ckpt_dir, \"unet\")\n",
    "    text_encoder_sub_dir = os.path.join(ckpt_dir, \"text_encoder\")\n",
    "    if isinstance(pipe.unet, PeftModel):\n",
    "        pipe.unet.set_adapter(adapter_name)\n",
    "    else:\n",
    "        pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)\n",
    "    pipe.unet = pipe.unet.merge_and_unload()\n",
    "\n",
    "    if os.path.exists(text_encoder_sub_dir):\n",
    "        if isinstance(pipe.text_encoder, PeftModel):\n",
    "            pipe.text_encoder.set_adapter(adapter_name)\n",
    "        else:\n",
    "            pipe.text_encoder = PeftModel.from_pretrained(\n",
    "                pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name\n",
    "            )\n",
    "        pipe.text_encoder = pipe.text_encoder.merge_and_unload()\n",
    "\n",
    "    return pipe\n",
    "\n",
    "\n",
    "def create_weighted_lora_adapter(pipe, adapters, weights, adapter_name=\"default\"):\n",
    "    pipe.unet.add_weighted_adapter(adapters, weights, adapter_name)\n",
    "    if isinstance(pipe.text_encoder, PeftModel):\n",
    "        pipe.text_encoder.add_weighted_adapter(adapters, weights, adapter_name)\n",
    "\n",
    "    return pipe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d4e888d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9f12b2cca0784cba9dc14ec48de929d5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 19 files:   0%|          | 0/19 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 13.2 s, sys: 7.55 s, total: 20.8 s\n",
      "Wall time: 5.89 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "pipe = get_lora_sd_pipeline(os.path.join(base_path, \"dog_dreambooth_updated\"), adapter_name=\"dog\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "22012468",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 90 ms, sys: 890 µs, total: 90.9 ms\n",
      "Wall time: 90.2 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "load_adapter(pipe, os.path.join(base_path, \"toy_dreambooth\"), adapter_name=\"toy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4ae4a74e",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe = create_weighted_lora_adapter(pipe, [\"toy\", \"dog\"], [1.0, 1.05], adapter_name=\"toy_dog\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9b8e2c69",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.74 ms, sys: 373 µs, total: 6.11 ms\n",
      "Wall time: 6.04 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "set_adapter(pipe, adapter_name=\"dog\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6d8b5e5d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2234aaa88bd34c898e175d8280ee6d4d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 545144 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=512x512>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt = \"sks dog playing fetch in the park\"\n",
    "negative_prompt = \"low quality, blurry, unfinished\"\n",
    "image = pipe(prompt, num_inference_steps=50, guidance_scale=7, negative_prompt=negative_prompt).images[0]\n",
    "image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1e1d1f30",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7.74 ms, sys: 0 ns, total: 7.74 ms\n",
      "Wall time: 7.31 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "set_adapter(pipe, adapter_name=\"toy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0c50c03d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f978437709b44b391744cf972415027",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 648432 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=512x512>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt = \"narendra modi rendered in the style of <1>\"\n",
    "negative_prompt = \"low quality, blurry, unfinished\"\n",
    "image = pipe(prompt, num_inference_steps=50, guidance_scale=7, negative_prompt=negative_prompt).images[0]\n",
    "image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "e3b9a681",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8871731127904802ba6123128ba22ecf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 453688 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=512x512>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "set_adapter(pipe, adapter_name=\"dog\")\n",
    "prompt = \"sks dog in a big red bucket\"\n",
    "negative_prompt = \"low quality, blurry, unfinished\"\n",
    "image = pipe(prompt, num_inference_steps=50, guidance_scale=7, negative_prompt=negative_prompt).images[0]\n",
    "image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a659ca6e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b701336557e4417d8d39d500699c298b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 437912 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=512x512>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "set_adapter(pipe, adapter_name=\"toy\")\n",
    "prompt = \"superman rendered in the style of <1>, close up potrait\"\n",
    "negative_prompt = \"low quality, blurry, unfinished\"\n",
    "image = pipe(prompt, num_inference_steps=50, guidance_scale=7, negative_prompt=negative_prompt).images[0]\n",
    "image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "1f0ecb40",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b47385e953b14c768d82bea776129904",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 297748 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=512x512>"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "set_adapter(pipe, adapter_name=\"toy_dog\")\n",
    "prompt = \"sks dog rendered in the style of <1>, close up potrait, 4K HD\"\n",
    "negative_prompt = \"low quality, blurry, unfinished\"\n",
    "image = pipe(prompt, num_inference_steps=50, guidance_scale=7, negative_prompt=negative_prompt).images[0]\n",
    "image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29720cdb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.4"
  },
  "vscode": {
   "interpreter": {
    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
