{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "411c59b3-f177-4a10-8925-d931ce572eaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from types import MethodType\n",
    "\n",
    "import torch\n",
    "from diffusers import StableDiffusionControlNetPipeline, DDIMScheduler, AutoencoderKL, ControlNetModel\n",
    "from PIL import Image\n",
    "\n",
    "from ip_adapter import IPAdapter\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6b6dc69c-192d-4d74-8b1e-f0d9ccfbdb49",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/modelzoo/IP-Adapter\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "current_dir = os.getcwd()\n",
    "print(current_dir)\n",
    "\n",
    "base_model_path = f\"{current_dir}/pretrained_models/sd1.5/Realistic_Vision_v4.0_noVAE\"\n",
    "vae_model_path = f\"{current_dir}/pretrained_models/sd1.5/sd-vae-ft-mse\"\n",
    "image_encoder_path = f\"{current_dir}/pretrained_models/models/image_encoder/\"\n",
    "ip_ckpt = f\"{current_dir}/pretrained_models/models/ip-adapter_sd15.safetensors\"\n",
    "device = \"cuda\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "63ec542f-8474-4f38-9457-073425578073",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "An error occurred while trying to fetch /home/modelzoo/IP-Adapter/pretrained_models/sd1.5/sd-vae-ft-mse: Error no file named diffusion_pytorch_model.safetensors found in directory /home/modelzoo/IP-Adapter/pretrained_models/sd1.5/sd-vae-ft-mse.\n",
      "Defaulting to unsafe serialization. Pass `allow_pickle=False` to raise an error instead.\n"
     ]
    }
   ],
   "source": [
    "def image_grid(imgs, rows, cols):\n",
    "    assert len(imgs) == rows*cols\n",
    "\n",
    "    w, h = imgs[0].size\n",
    "    grid = Image.new('RGB', size=(cols*w, rows*h))\n",
    "    grid_w, grid_h = grid.size\n",
    "    \n",
    "    for i, img in enumerate(imgs):\n",
    "        grid.paste(img, box=(i%cols*w, i//cols*h))\n",
    "    return grid\n",
    "\n",
    "noise_scheduler = DDIMScheduler(\n",
    "    num_train_timesteps=1000,\n",
    "    beta_start=0.00085,\n",
    "    beta_end=0.012,\n",
    "    beta_schedule=\"scaled_linear\",\n",
    "    clip_sample=False,\n",
    "    set_alpha_to_one=False,\n",
    "    steps_offset=1,\n",
    ")\n",
    "vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fed4baf0-eb72-47c6-8b5b-4046c7c7c72e",
   "metadata": {},
   "source": [
    "## ControlNet Depth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "3849f9d0-5f68-4a49-9190-69dd50720cae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "82bc276a642b4c4286171ddcdf40ab05",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# load controlnet\n",
    "controlnet_model_path = f\"{current_dir}/pretrained_models/sd1.5/lllyasviel/control_v11f1p_sd15_depth\"\n",
    "controlnet = ControlNetModel.from_pretrained(controlnet_model_path, torch_dtype=torch.float16, use_safetensors=True, variant=\"fp16\")\n",
    "# load SD pipeline\n",
    "pipe = StableDiffusionControlNetPipeline.from_pretrained(\n",
    "    base_model_path,\n",
    "    controlnet=controlnet,\n",
    "    torch_dtype=torch.float16,\n",
    "    scheduler=noise_scheduler,\n",
    "    vae=vae,\n",
    "    feature_extractor=None,\n",
    "    safety_checker=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ec09e937-3904-4d8e-a559-9066502ded36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 15388 characters]",
      "image/png": "[encoded data, 118704 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=512x256>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# read image prompt\n",
    "image = Image.open(\"assets/images/statue.png\")\n",
    "depth_map = Image.open(\"assets/structure_controls/depth.png\")\n",
    "image_grid([image.resize((256, 256)), depth_map.resize((256, 256))], 1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "81b1ab06-d3ed-4a7e-a356-9ddf1a2eecd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load ip-adapter\n",
    "ip_model = IPAdapter(pipe, image_encoder_path, ip_ckpt, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "42cd9556-988c-4a16-9408-17c8919d839c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "77fed9b113d0469cb6a8914337541e85",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 144592 characters]",
      "image/png": "[encoded data, 2060880 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=2048x512>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# generate image variations\n",
    "images = ip_model.generate(pil_image=image, image=depth_map, num_samples=4, num_inference_steps=50, seed=42)\n",
    "grid = image_grid(images, 1, 4)\n",
    "grid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf199405-7cb5-4f78-9973-5fe51c632a41",
   "metadata": {},
   "source": [
    "## ControlNet OpenPose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "6f089ad0-4683-46d7-ab58-9e5fe8f34c67",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "75b526a428084db587a11ab2d3773b47",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# load SD pipe\n",
    "del pipe, ip_model\n",
    "torch.cuda.empty_cache()\n",
    "# load controlnet\n",
    "controlnet_model_path = f\"{current_dir}/pretrained_models/sd1.5/lllyasviel/control_v11p_sd15_openpose\"\n",
    "controlnet = ControlNetModel.from_pretrained(controlnet_model_path, torch_dtype=torch.float16, variant=\"fp16\", use_safetensors=True)\n",
    "# load SD pipeline\n",
    "pipe = StableDiffusionControlNetPipeline.from_pretrained(\n",
    "    base_model_path,\n",
    "    controlnet=controlnet,\n",
    "    torch_dtype=torch.float16,\n",
    "    scheduler=noise_scheduler,\n",
    "    vae=vae,\n",
    "    feature_extractor=None,\n",
    "    safety_checker=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b8db2b55-2f56-4eef-b2ca-c5126b14feb7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 182684 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGBA size=256x256>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# read image prompt\n",
    "image = Image.open(\"assets/images/girl.png\")\n",
    "image.resize((256, 256))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "346d1e25-5b50-4d2e-851d-4ba620a55ce3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 9392 characters]",
      "image/png": "[encoded data, 19196 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=256x384>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "openpose_image = Image.open(\"assets/structure_controls/openpose.png\")\n",
    "openpose_image.resize((256, 384))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "a501f284-f295-4673-96ab-e34378da62ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load ip-adapter\n",
    "ip_model = IPAdapter(pipe, image_encoder_path, ip_ckpt, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "f58fff74-9ff2-46e6-bc8a-2ad4ae1fbe0f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ec492d35de794b6788b6093570b76328",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 239964 characters]",
      "image/png": "[encoded data, 3355212 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=2048x768>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# generate\n",
    "images = ip_model.generate(pil_image=image, image=openpose_image, width=512, height=768, num_samples=4, num_inference_steps=50, seed=42)\n",
    "grid = image_grid(images, 1, 4)\n",
    "grid"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
