{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "411c59b3-f177-4a10-8925-d931ce572eaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, StableDiffusionInpaintPipelineLegacy, DDIMScheduler, AutoencoderKL\n",
    "from PIL import Image\n",
    "\n",
    "from ip_adapter import IPAdapter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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": 3,
   "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"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/site-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return self.fget.__get__(instance, owner)()\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": "code",
   "execution_count": 4,
   "id": "3849f9d0-5f68-4a49-9190-69dd50720cae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "024f4fbef59a4408b4af0f2d910f1ca4",
       "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 pipeline\n",
    "pipe = StableDiffusionPipeline.from_pretrained(\n",
    "    base_model_path,\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": 5,
   "id": "ec09e937-3904-4d8e-a559-9066502ded36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 11420 characters]",
      "image/png": "[encoded data, 136424 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=256x256>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# read image prompt\n",
    "image = Image.open(\"assets/images/woman.png\")\n",
    "image.resize((256, 256))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a23de3d2-169e-470b-8012-960e3d07b04b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load ip-adapter\n",
    "ip_model = IPAdapter(pipe, image_encoder_path, ip_ckpt, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d83df45f-717d-4bb3-a5fd-0ea30930a431",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b98a708c2007494f9108509c04a66cbb",
       "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, 112540 characters]",
      "image/png": "[encoded data, 1944560 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=2048x512>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# only image prompt\n",
    "images = ip_model.generate(pil_image=image, num_samples=4, num_inference_steps=50, seed=42)\n",
    "grid = image_grid(images, 1, 4)\n",
    "grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b77f52de-a9e4-44e1-aeec-8165414f1273",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "147d5c0f7b1c4b76b75b8561bd87ec85",
       "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, 139200 characters]",
      "image/png": "[encoded data, 2091852 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=2048x512>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# multimodal prompts\n",
    "images = ip_model.generate(pil_image=image, num_samples=4, num_inference_steps=50, seed=42,\n",
    "        prompt=\"best quality, high quality, wearing a hat on the beach\", scale=0.6)\n",
    "grid = image_grid(images, 1, 4)\n",
    "grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1e411dbc-423a-4c14-bf93-6423c2444341",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 10624 characters]",
      "image/png": "[encoded data, 72888 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=256x256>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# read image prompt\n",
    "image = Image.open(\"assets/images/statue.png\")\n",
    "image.resize((256, 256))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "36081858-b354-430e-9ff1-015576c0a339",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4849769ef1244c3b81551c0db17be9bd",
       "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, 165564 characters]",
      "image/png": "[encoded data, 1806700 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=2048x512>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# only image prompt\n",
    "images = ip_model.generate(pil_image=image, num_samples=4, num_inference_steps=50, seed=42)\n",
    "grid = image_grid(images, 1, 4)\n",
    "grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5d3d874a-49b2-4c7e-ad58-b0ecc085c1fd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a211e932ce204d32812b9c366c082ff9",
       "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, 286448 characters]",
      "image/png": "[encoded data, 2818908 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=2048x512>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# multimodal prompts\n",
    "images = ip_model.generate(pil_image=image, num_samples=4, num_inference_steps=50, seed=42,\n",
    "        prompt=\"best quality, high quality, in a garden with flowers\", scale=0.5)\n",
    "grid = image_grid(images, 1, 4)\n",
    "grid"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
