{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "acab479f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from PIL import Image\n",
    "\n",
    "import torch\n",
    "from accelerate.logging import get_logger\n",
    "from diffusers import StableDiffusionPipeline\n",
    "from diffusers.utils import check_min_version\n",
    "\n",
    "from peft import PeftModel\n",
    "\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",
    "MODEL_NAME = \"stabilityai/stable-diffusion-2-1\"\n",
    "\n",
    "PEFT_TYPE=\"hra\"\n",
    "HRA_R=8\n",
    "SELECTED_SUBJECT=\"backpack\"\n",
    "EPOCH_IDX = 1000\n",
    "\n",
    "PROJECT_NAME=f\"dreambooth_{PEFT_TYPE}\"\n",
    "RUN_NAME=f\"{SELECTED_SUBJECT}_{PEFT_TYPE}_{HRA_R}\"\n",
    "OUTPUT_DIR=f\"./data/output/{PEFT_TYPE}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "06cfd506",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_hra_sd_pipeline(\n",
    "    ckpt_dir, base_model_name_or_path=None, epoch=int, dtype=torch.float32, device=\"cuda\", adapter_name=\"default\"\n",
    "):\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",
    "    \n",
    "    load_adapter(pipe, ckpt_dir, epoch, adapter_name)\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, epoch, adapter_name=\"default\"):\n",
    "    \n",
    "    unet_sub_dir = os.path.join(ckpt_dir, f\"unet/{epoch}\", adapter_name)\n",
    "    text_encoder_sub_dir = os.path.join(ckpt_dir, f\"text_encoder/{epoch}\", adapter_name)\n",
    "    \n",
    "    if isinstance(pipe.unet, PeftModel):\n",
    "        pipe.unet.load_adapter(unet_sub_dir, adapter_name=adapter_name)\n",
    "    else:\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",
    "        if isinstance(pipe.text_encoder, PeftModel):\n",
    "            pipe.text_encoder.load_adapter(text_encoder_sub_dir, adapter_name=adapter_name)\n",
    "        else:\n",
    "            pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, 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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98a0d8ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"a purple qwe backpack.\"\n",
    "negative_prompt = \"low quality, blurry, unfinished\"\n",
    "device = torch.accelerator.current_accelerator().type if hasattr(torch, \"accelerator\") else \"cuda\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4e888d2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading pipeline components...: 100%|██████████| 6/6 [00:00<00:00, 14.47it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.72 s, sys: 495 ms, total: 2.22 s\n",
      "Wall time: 2.28 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "pipe = get_hra_sd_pipeline(OUTPUT_DIR, MODEL_NAME, EPOCH_IDX, adapter_name=RUN_NAME, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f1c1a1c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/50 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50/50 [00:09<00:00,  5.28it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 8.77 s, sys: 861 ms, total: 9.63 s\n",
      "Wall time: 9.6 s\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 36932 characters]",
      "image/png": "[encoded data, 480020 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=512x512>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5, negative_prompt=negative_prompt).images[0]\n",
    "image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "60fa38d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 72508 characters]",
      "image/png": "[encoded data, 681100 characters]",
      "text/plain": [
       "<PIL.PngImagePlugin.PngImageFile image mode=RGB size=512x512>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# This is an example.\n",
    "example_image = Image.open(\"./a_purple_qwe_backpack.png\")\n",
    "example_image"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
