{
 "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 IPAdapterFull"
   ]
  },
  {
   "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-full-face_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": "326a774751fd48408789f8092ec2e520",
       "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": "a23de3d2-169e-470b-8012-960e3d07b04b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load ip-adapter\n",
    "# we use 256 patch tokens and 1 cls token\n",
    "ip_model = IPAdapterFull(pipe, image_encoder_path, ip_ckpt, device, num_tokens=257)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fcbdd2d-e69b-469a-8c04-42ee45fb5c9e",
   "metadata": {},
   "source": [
    "## Best practices\n",
    "- You should use a cropped face image, segment the face, and fill the background with pure white (255, 255, 255)\n",
    "- Select a proper weight (scale), 0.7 can achieve good results in most casese"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e1e0016a-3be8-408b-906e-73375dad9f9a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 172368 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGBA size=256x256>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# read image prompt\n",
    "image = Image.open(\"assets/images/ai_face2.png\")\n",
    "image.resize((256, 256))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "70e77d65-262f-415f-9cbd-057d57c4222d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2ca151f387284f02ba9082152291045d",
       "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, 192112 characters]",
      "image/png": "[encoded data, 2817856 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=2048x704>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# use face as image prompt\n",
    "images = ip_model.generate(\n",
    "    pil_image=image, num_samples=4, prompt=\"A photo of a girl wearing a black dress, holding red roses in hand, upper body, behind is the Eiffel Tower\",\n",
    "    scale=0.7, width=512, height=704, num_inference_steps=50, seed=42)\n",
    "grid = image_grid(images, 1, 4)\n",
    "grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d83df45f-717d-4bb3-a5fd-0ea30930a431",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "848b25afba9044b8a08c417d33144bfd",
       "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, 273592 characters]",
      "image/png": "[encoded data, 3272560 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=2048x704>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Use a lower scale to mix faces\n",
    "images = ip_model.generate(\n",
    "    pil_image=image, num_samples=4, prompt=\"photo of Einstein wearing colorful casual shirt in a garden\",\n",
    "    scale=0.4, width=512, height=704, num_inference_steps=50, seed=42)\n",
    "grid = image_grid(images, 1, 4)\n",
    "grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "10d6359e-6eb3-432a-a890-b814c505d005",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b046daa2c06e4600a973f607f58e7bf4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Change to a anime model\n",
    "del pipe, ip_model\n",
    "torch.cuda.empty_cache()\n",
    "base_model_path = f\"{current_dir}/pretrained_models/sd1.5/dreamlike-anime-1.0/\"\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": 11,
   "id": "c5059a81-0f1f-41b7-bdcc-9e266be3c0a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load ip-adapter\n",
    "ip_model = IPAdapterFull(pipe, image_encoder_path, ip_ckpt, device, num_tokens=257)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f915046e-1fcf-44ad-8bf8-4fe427d19e0c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "11cc9ec14c6f4c2ab603a8db73a94c02",
       "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, 248904 characters]",
      "image/png": "[encoded data, 3040316 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=2048x704>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# use face as image prompt\n",
    "images = ip_model.generate(\n",
    "    pil_image=image, num_samples=4, prompt=\"A photo of a girl wearing a black dress, holding red roses in hand, upper body, behind is the Eiffel Tower\",\n",
    "    scale=0.6, width=512, height=704, 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"
  },
  "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
}
