{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/lib/python3.9/site-packages/diffusers/utils/outputs.py:63: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "  torch.utils._pytree._register_pytree_node(\n"
     ]
    }
   ],
   "source": [
    "#!/usr/bin/env python\n",
    "# coding: utf-8\n",
    "import torch\n",
    "from diffusers import StableDiffusionXLPipeline, DDIMScheduler, AutoencoderKL\n",
    "from PIL import Image\n",
    "\n",
    "from ip_adapter import IPAdapterPlusXL\n",
    "from ip_adapter.custom_pipelines import StableDiffusionXLCustomPipeline\n",
    "import wandb\n",
    "from ip_adapter.utils import register_cross_attention_hook, get_net_attn_map, attnmaps2images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_model_path = \"stabilityai/stable-diffusion-xl-base-1.0\"\n",
    "image_encoder_path = \"models/image_encoder/\"\n",
    "vae_model_path = \"stabilityai/sd-vae-ft-mse\"\n",
    "ip_ckpt = \"sdxl_models/ip-adapter-plus-face_sdxl_vit-h.bin\" # a experimental version\n",
    "device = \"cuda\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install `accelerate` for faster and less memory-intense model loading. You can do so with: \n",
      "```\n",
      "pip install accelerate\n",
      "```\n",
      ".\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9863bfcfcc2f4a448b9a917cd9c23041",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/7 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pipe = StableDiffusionXLCustomPipeline.from_pretrained(\n",
    "    base_model_path,\n",
    "    torch_dtype=torch.float16,\n",
    "    add_watermarker=False,\n",
    ")\n",
    "pipe.unet = register_cross_attention_hook(pipe.unet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe.unet = register_cross_attention_hook(pipe.unet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "ip_model = IPAdapterPlusXL(pipe, image_encoder_path, ip_ckpt, device, num_tokens=16)\n",
    "# load ip-adapter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 133532 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGBA size=256x256>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "face_image = Image.open(\"assets/images/ai_face.png\")\n",
    "# generate image\n",
    "prompt = \"photo of a beautiful girl wearing casual shirt in a garden\"\n",
    "negative_prompt = \"monochrome, lowres, bad anatomy, worst quality, low quality, blurry\"\n",
    "face_image.resize((256, 256))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "20eb7dabcecf4b50874d58b755dd2867",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 649964 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=512x768>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "images = ip_model.generate(\n",
    "        pil_image=face_image,\n",
    "        prompt=prompt,\n",
    "        num_samples=1,\n",
    "        num_inference_steps=30,\n",
    "        width=512, height=768,\n",
    "        seed=420)\n",
    "images[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 768, 512])\n"
     ]
    }
   ],
   "source": [
    "attn_maps = get_net_attn_map((768, 512))\n",
    "print(attn_maps.shape)\n",
    "attn_hot = attnmaps2images(attn_maps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 69960 characters]",
      "text/plain": [
       "<Figure size 1200x400 with 18 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "#axes[0].imshow(attn_hot[0], cmap='gray')\n",
    "display_images = [face_image] + attn_hot + [images[0]]\n",
    "fig, axes = plt.subplots(1, len(display_images), figsize=(12, 4))\n",
    "for axe, image in zip(axes, display_images):\n",
    "    axe.imshow(image, cmap='gray')\n",
    "    axe.axis('off')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
