notebooks/272-paint-by-example/272-paint-by-example.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Paint By Example: Exemplar-based Image Editing with Diffusion Models\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stable Diffusion in Diffusers library\n",
    "To work with Stable Diffusion, we will use the Hugging Face [Diffusers](https://github.com/huggingface/diffusers) library. To experiment with in-painting we can use Diffusers which exposes the [StableDiffusionInpaintPipeline](https://huggingface.co/docs/diffusers/using-diffusers/conditional_image_generation) similar to the [other Diffusers pipelines](https://huggingface.co/docs/diffusers/api/pipelines/overview). The code below demonstrates how to create `StableDiffusionInpaintPipeline` using `stable-diffusion-2-inpainting`.\n",
    "To create the drawing tool we will install Gradio for handling user interaction.\n",
    "\n",
    "This is the overall flow of the application:\n",
    "![Flow Diagram](https://user-images.githubusercontent.com/103226580/236954918-f364b227-293c-4f78-a9bf-9dcebcb1034a.png)\n",
    "\n",
    "This is the detailed flowchart for the pipeline:\n",
    "![pipeline-flowchart](https://github.com/openvinotoolkit/openvino_notebooks/assets/103226580/cde2d5c4-2540-4a45-ad9c-339f7a69459d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "[notice] A new release of pip is available: 23.1 -> 23.3.1\n",
      "[notice] To update, run: python.exe -m pip install --upgrade pip\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "[notice] A new release of pip is available: 23.1 -> 23.3.1\n",
      "[notice] To update, run: python.exe -m pip install --upgrade pip\n"
     ]
    }
   ],
   "source": [
    "%pip install -q \"gradio>=4.10.0\"\n",
    "%pip install -q torch torchvision --extra-index-url \"https://download.pytorch.org/whl/cpu\"\n",
    "%pip install -q \"diffusers>=0.25.0\" \"peft<=0.6.2\" \"openvino>=2023.2.0\" \"transformers>=4.25.1\" ipywidgets opencv_python"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download the model from [HuggingFace Paint-by-Example](https://huggingface.co/Fantasy-Studio/Paint-by-Example). This might take several minutes because it is over 5GB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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",
      "You are using a model of type clip_vision_model to instantiate a model of type clip. This is not supported for all configurations of models and can yield errors.\n"
     ]
    }
   ],
   "source": [
    "from diffusers import DPMSolverMultistepScheduler, DiffusionPipeline\n",
    "\n",
    "pipeline = DiffusionPipeline.from_pretrained(\"Fantasy-Studio/Paint-By-Example\")\n",
    "\n",
    "scheduler_inpaint = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "\n",
    "extractor = pipeline.feature_extractor\n",
    "image_encoder = pipeline.image_encoder\n",
    "image_encoder.eval()\n",
    "unet_inpaint = pipeline.unet\n",
    "unet_inpaint.eval()\n",
    "vae_inpaint = pipeline.vae\n",
    "vae_inpaint.eval()\n",
    "\n",
    "del pipeline\n",
    "gc.collect();"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download default images\n",
    "\n",
    "Download default images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6eeaae0ba9794b8880d6e36402fdc7af",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "data\\image\\0.png:   0%|          | 0.00/453k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8cfddc9286674e72a5ef2ad1865384e7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "data\\image\\1.png:   0%|          | 0.00/545k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d6e922cd79094f628162d820db7316aa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "data\\image\\2.png:   0%|          | 0.00/431k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7c7d7e04215449b0827160ded02f5e70",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "data\\reference\\bird.jpg:   0%|          | 0.00/835k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d5c60718d1d34fbc87ef86f8fd0589a4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "data\\reference\\car.jpg:   0%|          | 0.00/414k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1f04be4501a449d290ee3533d72640b9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "data\\reference\\dog.jpg:   0%|          | 0.00/543k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "WindowsPath('C:/hackathon/openvino_notebooks/notebooks/272-paint-by-example/data/reference/dog.jpg')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Fetch `notebook_utils` module\n",
    "import urllib.request\n",
    "urllib.request.urlretrieve(\n",
    "    url='https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/main/notebooks/utils/notebook_utils.py',\n",
    "    filename='notebook_utils.py'\n",
    ")\n",
    "\n",
    "from notebook_utils import download_file\n",
    "\n",
    "download_file(\"https://github-production-user-asset-6210df.s3.amazonaws.com/103226580/286377210-edc98e97-0e43-4796-b771-dacd074c39ea.png\", \"0.png\", \"data/image\")\n",
    "\n",
    "download_file(\"https://github-production-user-asset-6210df.s3.amazonaws.com/103226580/286377233-b2c2d902-d379-415a-8183-5bdd37c52429.png\", \"1.png\", \"data/image\")\n",
    "\n",
    "download_file(\"https://github-production-user-asset-6210df.s3.amazonaws.com/103226580/286377248-da1db61e-3521-4cdb-85c8-1386d360ce22.png\", \"2.png\", \"data/image\")\n",
    "\n",
    "download_file(\"https://github-production-user-asset-6210df.s3.amazonaws.com/103226580/286377279-fa496f17-e850-4351-87c5-2552dfbc4633.jpg\", \"bird.jpg\", \"data/reference\")\n",
    "\n",
    "download_file(\"https://github-production-user-asset-6210df.s3.amazonaws.com/103226580/286377298-06a25ff2-84d8-4d46-95cd-8c25efa690d8.jpg\", \"car.jpg\", \"data/reference\")\n",
    "\n",
    "download_file(\"https://github-production-user-asset-6210df.s3.amazonaws.com/103226580/286377318-8841a801-1933-4523-a433-7d2fb64c47e6.jpg\", \"dog.jpg\", \"data/reference\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert models to OpenVINO Intermediate representation (IR) format\n",
    "\n",
    "Adapted from [236 Stable Diffusion v2 Infinite Zoom notebook](../236-stable-diffusion-v2/236-stable-diffusion-v2-infinite-zoom.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import torch\n",
    "import numpy as np\n",
    "import openvino as ov\n",
    "\n",
    "model_dir = Path(\"model\")\n",
    "model_dir.mkdir(exist_ok=True)\n",
    "sd2_inpainting_model_dir = Path(\"model/paint_by_example\")\n",
    "sd2_inpainting_model_dir.mkdir(exist_ok=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Functions to convert to OpenVINO IR format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cleanup_torchscript_cache():\n",
    "    \"\"\"\n",
    "    Helper for removing cached model representation\n",
    "    \"\"\"\n",
    "    torch._C._jit_clear_class_registry()\n",
    "    torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore()\n",
    "    torch.jit._state._clear_class_state()\n",
    "\n",
    "\n",
    "def convert_image_encoder(image_encoder: torch.nn.Module, ir_path:Path):\n",
    "    \"\"\"\n",
    "    Convert Image Encoder model to IR. \n",
    "    Function accepts pipeline, prepares example inputs for conversion\n",
    "    Parameters: \n",
    "        image_encoder (torch.nn.Module): image encoder PyTorch model\n",
    "        ir_path (Path): File for storing model\n",
    "    Returns:\n",
    "        None\n",
    "    \"\"\"\n",
    "    class ImageEncoderWrapper(torch.nn.Module):\n",
    "        def __init__(self, image_encoder):\n",
    "            super().__init__()\n",
    "            self.image_encoder = image_encoder\n",
    "\n",
    "        def forward(self, image):\n",
    "            image_embeddings, negative_prompt_embeds = self.image_encoder(image, return_uncond_vector=True)\n",
    "            return image_embeddings, negative_prompt_embeds\n",
    "\n",
    "    if not ir_path.exists():\n",
    "        image_encoder = ImageEncoderWrapper(image_encoder)\n",
    "        image_encoder.eval()\n",
    "        input_ids = torch.randn((1,3,224,224))\n",
    "        # switch model to inference mode\n",
    "\n",
    "        # disable gradients calculation for reducing memory consumption\n",
    "        with torch.no_grad():\n",
    "            ov_model = ov.convert_model(\n",
    "                image_encoder,\n",
    "                example_input=input_ids,\n",
    "                input=([1,3,224,224],)\n",
    "            )\n",
    "            ov.save_model(ov_model, ir_path)\n",
    "            del ov_model\n",
    "            cleanup_torchscript_cache()\n",
    "        print('Image Encoder successfully converted to IR')\n",
    "\n",
    "        \n",
    "def convert_unet(unet:torch.nn.Module, ir_path:Path, num_channels:int = 4, width:int = 64, height:int = 64):\n",
    "    \"\"\"\n",
    "    Convert Unet model to IR format. \n",
    "    Function accepts pipeline, prepares example inputs for conversion \n",
    "    Parameters: \n",
    "        unet (torch.nn.Module): UNet PyTorch model\n",
    "        ir_path (Path): File for storing model\n",
    "        num_channels (int, optional, 4): number of input channels\n",
    "        width (int, optional, 64): input width\n",
    "        height (int, optional, 64): input height\n",
    "    Returns:\n",
    "        None\n",
    "    \"\"\"\n",
    "    dtype_mapping = {\n",
    "        torch.float32: ov.Type.f32,\n",
    "        torch.float64: ov.Type.f64\n",
    "    }\n",
    "    if not ir_path.exists():\n",
    "        # prepare inputs\n",
    "        encoder_hidden_state = torch.ones((2, 1, 768))\n",
    "        latents_shape = (2, num_channels, width, height)\n",
    "        latents = torch.randn(latents_shape)\n",
    "        t = torch.from_numpy(np.array(1, dtype=np.float32))\n",
    "        unet.eval()\n",
    "        dummy_inputs = (latents, t, encoder_hidden_state)\n",
    "        input_info = []\n",
    "        for input_tensor in dummy_inputs:\n",
    "            shape = ov.PartialShape(tuple(input_tensor.shape))\n",
    "            element_type = dtype_mapping[input_tensor.dtype]\n",
    "            input_info.append((shape, element_type))\n",
    "\n",
    "        with torch.no_grad():\n",
    "            ov_model = ov.convert_model(\n",
    "                unet, \n",
    "                example_input=dummy_inputs,\n",
    "                input=input_info\n",
    "            )\n",
    "            ov.save_model(ov_model, ir_path)\n",
    "            del ov_model\n",
    "            cleanup_torchscript_cache()\n",
    "        print('U-Net successfully converted to IR')\n",
    "\n",
    "\n",
    "def convert_vae_encoder(vae: torch.nn.Module, ir_path: Path, width:int = 512, height:int = 512):\n",
    "    \"\"\"\n",
    "    Convert VAE model to IR format. \n",
    "    Function accepts VAE model, creates wrapper class for export only necessary for inference part, \n",
    "    prepares example inputs for conversion, \n",
    "    Parameters: \n",
    "        vae (torch.nn.Module): VAE PyTorch model\n",
    "        ir_path (Path): File for storing model\n",
    "        width (int, optional, 512): input width\n",
    "        height (int, optional, 512): input height\n",
    "    Returns:\n",
    "        None\n",
    "    \"\"\"\n",
    "    class VAEEncoderWrapper(torch.nn.Module):\n",
    "        def __init__(self, vae):\n",
    "            super().__init__()\n",
    "            self.vae = vae\n",
    "\n",
    "        def forward(self, image):\n",
    "            latents = self.vae.encode(image).latent_dist.sample()\n",
    "            return latents\n",
    "\n",
    "    if not ir_path.exists():\n",
    "        vae_encoder = VAEEncoderWrapper(vae)\n",
    "        vae_encoder.eval()\n",
    "        image = torch.zeros((1, 3, width, height))\n",
    "        with torch.no_grad():\n",
    "            ov_model = ov.convert_model(vae_encoder, example_input=image, input=([1,3, width, height],))\n",
    "        ov.save_model(ov_model, ir_path)\n",
    "        del ov_model\n",
    "        cleanup_torchscript_cache()\n",
    "        print('VAE encoder successfully converted to IR')\n",
    "\n",
    "\n",
    "def convert_vae_decoder(vae: torch.nn.Module, ir_path: Path, width:int = 64, height:int = 64):\n",
    "    \"\"\"\n",
    "    Convert VAE decoder model to IR format. \n",
    "    Function accepts VAE model, creates wrapper class for export only necessary for inference part, \n",
    "    prepares example inputs for conversion, \n",
    "    Parameters: \n",
    "        vae (torch.nn.Module): VAE model \n",
    "        ir_path (Path): File for storing model\n",
    "        width (int, optional, 64): input width\n",
    "        height (int, optional, 64): input height\n",
    "    Returns:\n",
    "        None\n",
    "    \"\"\"\n",
    "    class VAEDecoderWrapper(torch.nn.Module):\n",
    "        def __init__(self, vae):\n",
    "            super().__init__()\n",
    "            self.vae = vae\n",
    "\n",
    "        def forward(self, latents):\n",
    "            latents = 1 / 0.18215 * latents\n",
    "            return self.vae.decode(latents)\n",
    "\n",
    "    if not ir_path.exists():\n",
    "        vae_decoder = VAEDecoderWrapper(vae)\n",
    "        latents = torch.zeros((1, 4, width, height))\n",
    "\n",
    "        vae_decoder.eval()\n",
    "        with torch.no_grad():\n",
    "            ov_model = ov.convert_model(vae_decoder, example_input=latents, input=([1, 4, width, height],))\n",
    "        ov.save_model(ov_model, ir_path)\n",
    "        del ov_model\n",
    "        cleanup_torchscript_cache()\n",
    "        print('VAE decoder successfully converted to ')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do the conversion of the in-painting model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image encoder will be loaded from model\\paint_by_example\\image_encoder.xml\n"
     ]
    }
   ],
   "source": [
    "IMAGE_ENCODER_OV_PATH_INPAINT = sd2_inpainting_model_dir / \"image_encoder.xml\"\n",
    "\n",
    "if not IMAGE_ENCODER_OV_PATH_INPAINT.exists():\n",
    "    convert_image_encoder(image_encoder, IMAGE_ENCODER_OV_PATH_INPAINT)\n",
    "else:\n",
    "    print(f\"Image encoder will be loaded from {IMAGE_ENCODER_OV_PATH_INPAINT}\")\n",
    "\n",
    "del image_encoder\n",
    "gc.collect();"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do the conversion of the Unet model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "U-Net will be loaded from model\\paint_by_example\\unet.xml\n"
     ]
    }
   ],
   "source": [
    "UNET_OV_PATH_INPAINT = sd2_inpainting_model_dir / 'unet.xml'\n",
    "if not UNET_OV_PATH_INPAINT.exists():\n",
    "    convert_unet(unet_inpaint, UNET_OV_PATH_INPAINT, num_channels=9, width=64, height=64)\n",
    "    del unet_inpaint\n",
    "    gc.collect()\n",
    "else:\n",
    "    del unet_inpaint\n",
    "    print(f\"U-Net will be loaded from {UNET_OV_PATH_INPAINT}\")\n",
    "gc.collect();"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do the conversion of the VAE Encoder model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VAE encoder will be loaded from model\\paint_by_example\\vae_encoder.xml\n",
      "VAE decoder will be loaded from model\\paint_by_example\\vae_decoder.xml\n"
     ]
    }
   ],
   "source": [
    "VAE_ENCODER_OV_PATH_INPAINT = sd2_inpainting_model_dir / 'vae_encoder.xml'\n",
    "\n",
    "if not VAE_ENCODER_OV_PATH_INPAINT.exists():\n",
    "    convert_vae_encoder(vae_inpaint, VAE_ENCODER_OV_PATH_INPAINT, 512, 512)\n",
    "else:\n",
    "    print(f\"VAE encoder will be loaded from {VAE_ENCODER_OV_PATH_INPAINT}\")\n",
    "\n",
    "VAE_DECODER_OV_PATH_INPAINT = sd2_inpainting_model_dir / 'vae_decoder.xml'\n",
    "if not VAE_DECODER_OV_PATH_INPAINT.exists():\n",
    "    convert_vae_decoder(vae_inpaint, VAE_DECODER_OV_PATH_INPAINT, 64, 64)\n",
    "else:\n",
    "    print(f\"VAE decoder will be loaded from {VAE_DECODER_OV_PATH_INPAINT}\")\n",
    "\n",
    "del vae_inpaint\n",
    "gc.collect();"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare Inference pipeline\n",
    "\n",
    "Function to prepare the mask and masked image.\n",
    "\n",
    "Adapted from [236 Stable Diffusion v2 Infinite Zoom notebook](../236-stable-diffusion-v2/236-stable-diffusion-v2-infinite-zoom.ipynb)\n",
    "\n",
    "The main difference is that instead of encoding a text prompt it will now encode an image as the prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect\n",
    "from typing import Optional, Union, Dict\n",
    "\n",
    "import PIL\n",
    "import cv2\n",
    "\n",
    "from transformers import CLIPImageProcessor\n",
    "from diffusers.pipelines.pipeline_utils import DiffusionPipeline\n",
    "from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler\n",
    "from openvino.runtime import Model\n",
    "\n",
    "\n",
    "def prepare_mask_and_masked_image(image:PIL.Image.Image, mask:PIL.Image.Image):\n",
    "    \"\"\"\n",
    "    Prepares a pair (image, mask) to be consumed by the Stable Diffusion pipeline. This means that those inputs will be\n",
    "    converted to ``np.array`` with shapes ``batch x channels x height x width`` where ``channels`` is ``3`` for the\n",
    "    ``image`` and ``1`` for the ``mask``.\n",
    "\n",
    "    The ``image`` will be converted to ``np.float32`` and normalized to be in ``[-1, 1]``. The ``mask`` will be\n",
    "    binarized (``mask > 0.5``) and cast to ``np.float32`` too.\n",
    "\n",
    "    Args:\n",
    "        image (Union[np.array, PIL.Image]): The image to inpaint.\n",
    "            It can be a ``PIL.Image``, or a ``height x width x 3`` ``np.array``\n",
    "        mask (_type_): The mask to apply to the image, i.e. regions to inpaint.\n",
    "            It can be a ``PIL.Image``, or a ``height x width`` ``np.array``.\n",
    "\n",
    "    Returns:\n",
    "        tuple[np.array]: The pair (mask, masked_image) as ``torch.Tensor`` with 4\n",
    "            dimensions: ``batch x channels x height x width``.\n",
    "    \"\"\"\n",
    "    if isinstance(image, (PIL.Image.Image, np.ndarray)):\n",
    "        image = [image]\n",
    "\n",
    "    if isinstance(image, list) and isinstance(image[0], PIL.Image.Image):\n",
    "        image = [np.array(i.convert(\"RGB\"))[None, :] for i in image]\n",
    "        image = np.concatenate(image, axis=0)\n",
    "    elif isinstance(image, list) and isinstance(image[0], np.ndarray):\n",
    "        image = np.concatenate([i[None, :] for i in image], axis=0)\n",
    "\n",
    "    image = image.transpose(0, 3, 1, 2)\n",
    "    image = image.astype(np.float32) / 127.5 - 1.0\n",
    "\n",
    "    # preprocess mask\n",
    "    if isinstance(mask, (PIL.Image.Image, np.ndarray)):\n",
    "        mask = [mask]\n",
    "\n",
    "    if isinstance(mask, list) and isinstance(mask[0], PIL.Image.Image):\n",
    "        mask = np.concatenate([np.array(m.convert(\"L\"))[None, None, :] for m in mask], axis=0)\n",
    "        mask = mask.astype(np.float32) / 255.0\n",
    "    elif isinstance(mask, list) and isinstance(mask[0], np.ndarray):\n",
    "        mask = np.concatenate([m[None, None, :] for m in mask], axis=0)\n",
    "\n",
    "    mask = 1 - mask\n",
    "\n",
    "    mask[mask < 0.5] = 0\n",
    "    mask[mask >= 0.5] = 1\n",
    "\n",
    "    masked_image = image * mask\n",
    "\n",
    "    return mask, masked_image"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Class for the pipeline which will connect all the models together: VAE decode --> image encode --> tokenizer --> Unet --> VAE model --> scheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class OVStableDiffusionInpaintingPipeline(DiffusionPipeline):\n",
    "    def __init__(\n",
    "        self,\n",
    "        vae_decoder: Model,\n",
    "        image_encoder: Model,\n",
    "        image_processor: CLIPImageProcessor,\n",
    "        unet: Model,\n",
    "        scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],\n",
    "        vae_encoder: Model = None,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Pipeline for text-to-image generation using Stable Diffusion.\n",
    "        Parameters:\n",
    "            vae_decoder (Model):\n",
    "                Variational Auto-Encoder (VAE) Model to decode images to and from latent representations.\n",
    "            image_encoder (Model):\n",
    "                https://huggingface.co/Fantasy-Studio/Paint-by-Example/blob/main/image_encoder/config.json\n",
    "            tokenizer (CLIPTokenizer):\n",
    "                Tokenizer of class CLIPTokenizer(https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).\n",
    "            unet (Model): Conditional U-Net architecture to denoise the encoded image latents.\n",
    "            vae_encoder (Model):\n",
    "                Variational Auto-Encoder (VAE) Model to encode images to latent representation.\n",
    "            scheduler (SchedulerMixin):\n",
    "                A scheduler to be used in combination with unet to denoise the encoded image latents. Can be one of\n",
    "                DDIMScheduler, LMSDiscreteScheduler, or PNDMScheduler.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.scheduler = scheduler\n",
    "        self.vae_decoder = vae_decoder\n",
    "        self.vae_encoder = vae_encoder\n",
    "        self.image_encoder = image_encoder\n",
    "        self.unet = unet\n",
    "        self._unet_output = unet.output(0)\n",
    "        self._vae_d_output = vae_decoder.output(0)\n",
    "        self._vae_e_output = vae_encoder.output(0) if vae_encoder is not None else None\n",
    "        self.height = self.unet.input(0).shape[2] * 8\n",
    "        self.width = self.unet.input(0).shape[3] * 8\n",
    "        self.image_processor = image_processor\n",
    "\n",
    "    def prepare_mask_latents(\n",
    "        self,\n",
    "        mask,\n",
    "        masked_image,\n",
    "        height=512,\n",
    "        width=512,\n",
    "        do_classifier_free_guidance=True,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Prepare mask as Unet nput and encode input masked image to latent space using vae encoder\n",
    "\n",
    "        Parameters:\n",
    "          mask (np.array): input mask array\n",
    "          masked_image (np.array): masked input image tensor\n",
    "          heigh (int, *optional*, 512): generated image height\n",
    "          width (int, *optional*, 512): generated image width\n",
    "          do_classifier_free_guidance (bool, *optional*, True): whether to use classifier free guidance or not\n",
    "        Returns:\n",
    "          mask (np.array): resized mask tensor\n",
    "          masked_image_latents (np.array): masked image encoded into latent space using VAE\n",
    "        \"\"\"\n",
    "        mask = torch.nn.functional.interpolate(torch.from_numpy(mask), size=(height // 8, width // 8))\n",
    "        mask = mask.numpy()\n",
    "\n",
    "        # encode the mask image into latents space so we can concatenate it to the latents\n",
    "        masked_image_latents = self.vae_encoder(masked_image)[self._vae_e_output]\n",
    "        masked_image_latents = 0.18215 * masked_image_latents\n",
    "\n",
    "        mask = np.concatenate([mask] * 2) if do_classifier_free_guidance else mask\n",
    "        masked_image_latents = (\n",
    "            np.concatenate([masked_image_latents] * 2)\n",
    "            if do_classifier_free_guidance\n",
    "            else masked_image_latents\n",
    "        )\n",
    "        return mask, masked_image_latents\n",
    "\n",
    "    def __call__(\n",
    "        self,\n",
    "        image: PIL.Image.Image,\n",
    "        mask_image: PIL.Image.Image,\n",
    "        reference_image: PIL.Image.Image,\n",
    "        num_inference_steps: Optional[int] = 50,\n",
    "        guidance_scale: Optional[float] = 7.5,\n",
    "        eta: Optional[float] = 0,\n",
    "        output_type: Optional[str] = \"pil\",\n",
    "        seed: Optional[int] = None,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Function invoked when calling the pipeline for generation.\n",
    "        Parameters:\n",
    "            image (PIL.Image.Image):\n",
    "                 Source image for inpainting.\n",
    "            mask_image (PIL.Image.Image):\n",
    "                 Mask area for inpainting\n",
    "            reference_image (PIL.Image.Image):\n",
    "                 Reference image to inpaint in mask area\n",
    "            num_inference_steps (int, *optional*, defaults to 50):\n",
    "                The number of denoising steps. More denoising steps usually lead to a higher quality image at the\n",
    "                expense of slower inference.\n",
    "            guidance_scale (float, *optional*, defaults to 7.5):\n",
    "                Guidance scale as defined in Classifier-Free Diffusion Guidance(https://arxiv.org/abs/2207.12598).\n",
    "                guidance_scale is defined as `w` of equation 2.\n",
    "                Higher guidance scale encourages to generate images that are closely linked to the text prompt,\n",
    "                usually at the expense of lower image quality.\n",
    "            eta (float, *optional*, defaults to 0.0):\n",
    "                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to\n",
    "                [DDIMScheduler], will be ignored for others.\n",
    "            output_type (`str`, *optional*, defaults to \"pil\"):\n",
    "                The output format of the generate image. Choose between\n",
    "                [PIL](https://pillow.readthedocs.io/en/stable/): PIL.Image.Image or np.array.\n",
    "            seed (int, *optional*, None):\n",
    "                Seed for random generator state initialization.\n",
    "        Returns:\n",
    "            Dictionary with keys:\n",
    "                sample - the last generated image PIL.Image.Image or np.array\n",
    "        \"\"\"\n",
    "        if seed is not None:\n",
    "            np.random.seed(seed)\n",
    "        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)\n",
    "        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`\n",
    "        # corresponds to doing no classifier free guidance.\n",
    "        do_classifier_free_guidance = guidance_scale > 1.0\n",
    "\n",
    "        # get reference image embeddings\n",
    "        image_embeddings = self._encode_image(reference_image, do_classifier_free_guidance=do_classifier_free_guidance)\n",
    "\n",
    "        # prepare mask\n",
    "        mask, masked_image = prepare_mask_and_masked_image(image, mask_image)\n",
    "        # set timesteps\n",
    "        accepts_offset = \"offset\" in set(\n",
    "            inspect.signature(self.scheduler.set_timesteps).parameters.keys()\n",
    "        )\n",
    "        extra_set_kwargs = {}\n",
    "        if accepts_offset:\n",
    "            extra_set_kwargs[\"offset\"] = 1\n",
    "\n",
    "        self.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)\n",
    "        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, 1)\n",
    "        latent_timestep = timesteps[:1]\n",
    "\n",
    "        # get the initial random noise unless the user supplied it\n",
    "        latents, meta = self.prepare_latents(latent_timestep)\n",
    "        mask, masked_image_latents = self.prepare_mask_latents(\n",
    "            mask,\n",
    "            masked_image,\n",
    "            do_classifier_free_guidance=do_classifier_free_guidance,\n",
    "        )\n",
    "\n",
    "        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature\n",
    "        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.\n",
    "        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502\n",
    "        # and should be between [0, 1]\n",
    "        accepts_eta = \"eta\" in set(\n",
    "            inspect.signature(self.scheduler.step).parameters.keys()\n",
    "        )\n",
    "        extra_step_kwargs = {}\n",
    "        if accepts_eta:\n",
    "            extra_step_kwargs[\"eta\"] = eta\n",
    "\n",
    "        for t in self.progress_bar(timesteps):\n",
    "            # expand the latents if we are doing classifier free guidance\n",
    "            latent_model_input = (\n",
    "                np.concatenate([latents] * 2)\n",
    "                if do_classifier_free_guidance\n",
    "                else latents\n",
    "            )\n",
    "            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)\n",
    "            latent_model_input = np.concatenate(\n",
    "                [latent_model_input, masked_image_latents, mask], axis=1\n",
    "            )\n",
    "            # predict the noise residual\n",
    "            noise_pred = self.unet(\n",
    "                [latent_model_input, np.array(t, dtype=np.float32), image_embeddings]\n",
    "            )[self._unet_output]\n",
    "            # perform guidance\n",
    "            if do_classifier_free_guidance:\n",
    "                noise_pred_uncond, noise_pred_text = noise_pred[0], noise_pred[1]\n",
    "                noise_pred = noise_pred_uncond + guidance_scale * (\n",
    "                    noise_pred_text - noise_pred_uncond\n",
    "                )\n",
    "\n",
    "            # compute the previous noisy sample x_t -> x_t-1\n",
    "            latents = self.scheduler.step(\n",
    "                torch.from_numpy(noise_pred),\n",
    "                t,\n",
    "                torch.from_numpy(latents),\n",
    "                **extra_step_kwargs,\n",
    "            )[\"prev_sample\"].numpy()\n",
    "        # scale and decode the image latents with vae\n",
    "        image = self.vae_decoder(latents)[self._vae_d_output]\n",
    "\n",
    "        image = self.postprocess_image(image, meta, output_type)\n",
    "        return {\"sample\": image}\n",
    "\n",
    "    def _encode_image(self, image:PIL.Image.Image, do_classifier_free_guidance:bool = True):\n",
    "        \"\"\"\n",
    "        Encodes the image into image encoder hidden states.\n",
    "\n",
    "        Parameters:\n",
    "            image (PIL.Image.Image): base image to encode\n",
    "            do_classifier_free_guidance (bool): whether to use classifier free guidance or not\n",
    "        Returns:\n",
    "            image_embeddings (np.ndarray): image encoder hidden states\n",
    "        \"\"\"\n",
    "        processed_image = self.image_processor(image)\n",
    "        processed_image = processed_image['pixel_values'][0]\n",
    "        processed_image = np.expand_dims(processed_image, axis=0)\n",
    "\n",
    "        output = self.image_encoder(processed_image)\n",
    "        image_embeddings = output[self.image_encoder.output(0)]\n",
    "        negative_embeddings = output[self.image_encoder.output(1)]\n",
    "\n",
    "        image_embeddings = np.concatenate([negative_embeddings, image_embeddings])\n",
    "\n",
    "        return image_embeddings\n",
    "\n",
    "    def prepare_latents(self, latent_timestep:torch.Tensor = None):\n",
    "        \"\"\"\n",
    "        Function for getting initial latents for starting generation\n",
    "        \n",
    "        Parameters:\n",
    "            latent_timestep (torch.Tensor, *optional*, None):\n",
    "                Predicted by scheduler initial step for image generation, required for latent image mixing with nosie\n",
    "        Returns:\n",
    "            latents (np.ndarray):\n",
    "                Image encoded in latent space\n",
    "        \"\"\"\n",
    "        latents_shape = (1, 4, self.height // 8, self.width // 8)\n",
    "        noise = np.random.randn(*latents_shape).astype(np.float32)\n",
    "        # if we use LMSDiscreteScheduler, let's make sure latents are mulitplied by sigmas\n",
    "        if isinstance(self.scheduler, LMSDiscreteScheduler):\n",
    "            noise = noise * self.scheduler.sigmas[0].numpy()\n",
    "        return noise, {}\n",
    "\n",
    "    def postprocess_image(self, image:np.ndarray, meta:Dict, output_type:str = \"pil\"):\n",
    "        \"\"\"\n",
    "        Postprocessing for decoded image. Takes generated image decoded by VAE decoder, unpad it to initila image size (if required), \n",
    "        normalize and convert to [0, 255] pixels range. Optionally, convertes it from np.ndarray to PIL.Image format\n",
    "        \n",
    "        Parameters:\n",
    "            image (np.ndarray):\n",
    "                Generated image\n",
    "            meta (Dict):\n",
    "                Metadata obtained on latents preparing step, can be empty\n",
    "            output_type (str, *optional*, pil):\n",
    "                Output format for result, can be pil or numpy\n",
    "        Returns:\n",
    "            image (List of np.ndarray or PIL.Image.Image):\n",
    "                Postprocessed images\n",
    "        \"\"\"\n",
    "        if \"padding\" in meta:\n",
    "            pad = meta[\"padding\"]\n",
    "            (_, end_h), (_, end_w) = pad[1:3]\n",
    "            h, w = image.shape[2:]\n",
    "            unpad_h = h - end_h\n",
    "            unpad_w = w - end_w\n",
    "            image = image[:, :, :unpad_h, :unpad_w]\n",
    "        image = np.clip(image / 2 + 0.5, 0, 1)\n",
    "        image = np.transpose(image, (0, 2, 3, 1))\n",
    "        # 9. Convert to PIL\n",
    "        if output_type == \"pil\":\n",
    "            image = self.numpy_to_pil(image)\n",
    "            if \"src_height\" in meta:\n",
    "                orig_height, orig_width = meta[\"src_height\"], meta[\"src_width\"]\n",
    "                image = [img.resize((orig_width, orig_height),\n",
    "                                    PIL.Image.Resampling.LANCZOS) for img in image]\n",
    "        else:\n",
    "            if \"src_height\" in meta:\n",
    "                orig_height, orig_width = meta[\"src_height\"], meta[\"src_width\"]\n",
    "                image = [cv2.resize(img, (orig_width, orig_width))\n",
    "                         for img in image]\n",
    "        return image\n",
    "\n",
    "    def get_timesteps(self, num_inference_steps:int, strength:float):\n",
    "        \"\"\"\n",
    "        Helper function for getting scheduler timesteps for generation\n",
    "        In case of image-to-image generation, it updates number of steps according to strength\n",
    "        \n",
    "        Parameters:\n",
    "           num_inference_steps (int):\n",
    "              number of inference steps for generation\n",
    "           strength (float):\n",
    "               value between 0.0 and 1.0, that controls the amount of noise that is added to the input image. \n",
    "               Values that approach 1.0 allow for lots of variations but will also produce images that are not semantically consistent with the input.\n",
    "        \"\"\"\n",
    "        # get the original timestep using init_timestep\n",
    "        init_timestep = min(int(num_inference_steps * strength), num_inference_steps)\n",
    "\n",
    "        t_start = max(num_inference_steps - init_timestep, 0)\n",
    "        timesteps = self.scheduler.timesteps[t_start:]\n",
    "\n",
    "        return timesteps, num_inference_steps - t_start "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Select inference device\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "select device from dropdown list for running inference using OpenVINO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bec54240bdda4dec9eb7341f048767c9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', index=2, options=('CPU', 'GPU', 'AUTO'), value='AUTO')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from openvino.runtime import Core\n",
    "import ipywidgets as widgets\n",
    "\n",
    "core = Core()\n",
    "\n",
    "device = widgets.Dropdown(\n",
    "    options=core.available_devices + [\"AUTO\"],\n",
    "    value='AUTO',\n",
    "    description='Device:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "device"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configure Inference Pipeline\n",
    "\n",
    "Configuration steps:\n",
    "1. Load models on device\n",
    "2. Configure tokenizer and scheduler\n",
    "3. Create instance of OvStableDiffusionInpaintingPipeline class\n",
    "\n",
    "This can take a while to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "ov_config = {\"INFERENCE_PRECISION_HINT\": \"f32\"} if device.value != \"CPU\" else {}\n",
    "\n",
    "image_encoder_inpaint = core.compile_model(IMAGE_ENCODER_OV_PATH_INPAINT, device.value)\n",
    "unet_model_inpaint = core.compile_model(UNET_OV_PATH_INPAINT, device.value)\n",
    "vae_decoder_inpaint = core.compile_model(VAE_DECODER_OV_PATH_INPAINT, device.value, ov_config)\n",
    "vae_encoder_inpaint = core.compile_model(VAE_ENCODER_OV_PATH_INPAINT, device.value, ov_config)\n",
    "\n",
    "ov_pipe_inpaint = OVStableDiffusionInpaintingPipeline(\n",
    "    image_processor=extractor,\n",
    "    image_encoder=image_encoder_inpaint,\n",
    "    unet=unet_model_inpaint,\n",
    "    vae_encoder=vae_encoder_inpaint,\n",
    "    vae_decoder=vae_decoder_inpaint,\n",
    "    scheduler=scheduler_inpaint,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on local URL:  http://127.0.0.1:7860\n",
      "\n",
      "To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7860/\" width=\"100%\" height=\"680\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Code adapated from https://huggingface.co/spaces/Fantasy-Studio/Paint-by-Example/blob/main/app.py\n",
    "\n",
    "import os\n",
    "import gradio as gr\n",
    "\n",
    "def predict(dict:gr.components.Image, reference:PIL.Image.Image, seed:int, step:int):\n",
    "    \"\"\"\n",
    "        This function runs when the 'paint' button is pressed. It takes 3 input images. Takes generated image decoded by VAE decoder, unpad it to initila image size (if required), \n",
    "        normalize and convert to [0, 255] pixels range. Optionally, convertes it from np.ndarray to PIL.Image format\n",
    "        \n",
    "        Parameters:\n",
    "            dict (Dict):\n",
    "                Contains two images in a dictionary\n",
    "                    'image' is the image that will be painted on\n",
    "                    'mask' is the black/white image specifying where to paint (white) and not to paint (black)\n",
    "            image (PIL.Image.Image):\n",
    "                Reference image that will be used by the model to know what to paint in the specified area\n",
    "            seed (int):\n",
    "                Used to initialize the random number generator state\n",
    "            step (int):\n",
    "                The number of denoising steps to run during inference. Low = fast/low quality, High = slow/higher quality\n",
    "        Returns:\n",
    "            image (PIL.Image.Image):\n",
    "                Postprocessed images\n",
    "    \"\"\"\n",
    "    width,height = dict[\"image\"].size\n",
    "\n",
    "    # If the image is not 512x512 then resize\n",
    "    if width < height:\n",
    "        factor = width / 512.0\n",
    "        width = 512\n",
    "        height = int((height / factor) / 8.0) * 8\n",
    "    else:\n",
    "        factor = height / 512.0\n",
    "        height = 512\n",
    "        width = int((width / factor) / 8.0) * 8\n",
    "\n",
    "    init_image = dict[\"image\"].convert(\"RGB\").resize((width,height))\n",
    "    mask = dict[\"mask\"].convert(\"RGB\").resize((width,height))\n",
    "\n",
    "    # If the image is not a 512x512 square then crop\n",
    "    if width > height:\n",
    "        buffer = (width - height) / 2\n",
    "        input_image = init_image.crop((buffer, 0, width - buffer, 512))\n",
    "        mask = mask.crop((buffer, 0, width - buffer, 512))\n",
    "    elif width < height:\n",
    "        buffer = (height - width) / 2\n",
    "        input_image = init_image.crop((0, buffer, 512, height - buffer))\n",
    "        mask = mask.crop((0, buffer, 512, height - buffer))\n",
    "    else:\n",
    "        input_image = init_image\n",
    "\n",
    "    if not os.path.exists('output'):\n",
    "        os.mkdir('output')\n",
    "    input_image.save('output/init.png')\n",
    "    mask.save('output/mask.png')\n",
    "    reference.save('output/ref.png')\n",
    "\n",
    "    mask = [mask]\n",
    "\n",
    "    result = ov_pipe_inpaint(\n",
    "        image=input_image,\n",
    "        mask_image=mask,\n",
    "        reference_image=reference,\n",
    "        seed=seed,\n",
    "        num_inference_steps=step,\n",
    "    )[\"sample\"][0]\n",
    "\n",
    "    out_dir = Path(\"output\")\n",
    "    out_dir.mkdir(exist_ok=True)\n",
    "    result.save('output/result.png')\n",
    "\n",
    "    return result\n",
    "\n",
    "\n",
    "example = {}\n",
    "ref_dir = 'data/reference'\n",
    "image_dir = 'data/image'\n",
    "ref_list = [os.path.join(ref_dir,file) for file in os.listdir(ref_dir) if file.endswith(\".jpg\")]\n",
    "ref_list.sort()\n",
    "image_list = [os.path.join(image_dir,file) for file in os.listdir(image_dir) if file.endswith(\".png\")]\n",
    "image_list.sort()\n",
    "\n",
    "\n",
    "image_blocks = gr.Blocks()\n",
    "with image_blocks as demo:\n",
    "    with gr.Group():\n",
    "        with gr.Row():\n",
    "            with gr.Column():\n",
    "                image = gr.ImageEditor(sources=['upload'], type=\"pil\", label=\"Source Image\")\n",
    "                reference = gr.Image(sources=['upload'], type=\"pil\", label=\"Reference Image\")\n",
    "\n",
    "            with gr.Column():\n",
    "                image_out = gr.Image(label=\"Output\", elem_id=\"output-img\")\n",
    "                steps = gr.Slider(label=\"Steps\", value=15, minimum=2, maximum=75, step=1,interactive=True)\n",
    "\n",
    "                seed = gr.Slider(0, 10000, label='Seed (0 = random)', value=0, step=1)\n",
    "\n",
    "                with gr.Row(elem_id=\"prompt-container\"):\n",
    "                    btn = gr.Button(\"Paint!\")\n",
    "                       \n",
    "        with gr.Row():\n",
    "            with gr.Column():\n",
    "                gr.Examples(image_list, inputs=[image],label=\"Examples - Source Image\",examples_per_page=12)\n",
    "            with gr.Column():\n",
    "                gr.Examples(ref_list, inputs=[reference],label=\"Examples - Reference Image\",examples_per_page=12)\n",
    "        \n",
    "        btn.click(fn=predict, inputs=[image, reference, seed, steps], outputs=[image_out])\n",
    "\n",
    "# Launching the Gradio app\n",
    "try:\n",
    "    image_blocks.launch(debug=True, height=680)\n",
    "except Exception:\n",
    "    image_blocks.queue().launch(share=True, debug=True, height=680)\n",
    "# if you are launching remotely, specify server_name and server_port\n",
    "# image_blocks.launch(server_name='your server name', server_port='server port in int')\n",
    "# Read more in the docs: https://gradio.app/docs/"
   ]
  }
 ],
 "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.8.10"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}