notebooks/252-fastcomposer-image-generation/252-fastcomposer-image-generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "#  [FastComposer: Tuning-Free Multi-Subject Image Generation with Localized Attention](https://fastcomposer.mit.edu/)\n",
    "\n",
    "FastComposer uses subject embeddings extracted by an image encoder to augment the generic text conditioning in diffusion models, enabling personalized image generation based on subject images and textual instructions with only forward passes.\n",
    "Moreover it addresses two problems:\n",
    "\n",
    " - **The identity blending problem.**\n",
    "To address the problem in the multi-subject generation it proposes cross-attention localization supervision during training, enforcing the attention of reference subjects localized to the correct regions in the target images.\n",
    "\n",
    " - **Subject overfitting.**\n",
    "Naively conditioning on subject embeddings results in subject overfitting. FastComposer proposes delayed subject conditioning in the denoising step to maintain both identity and editability in subject-driven image generation.\n",
    "\n",
    "\n",
    "FastComposer generates images of multiple unseen individuals with different styles, actions, and contexts.\n",
    "\n",
    "\n",
    "<img src=\"https://github.com/mit-han-lab/fastcomposer/blob/main/figures/multi-subject.png?raw=True\" width=\"969\">\n",
    "\n",
    "\n",
    "> **NOTE**: `model.py` is slightly changed `model.py` from fastcomposer repository. There are two main changes:\n",
    ">  - some unused lines of code are removed to avoid errors if there are no CUDA drivers in the system\n",
    ">  - changes to have compatibility with transformers >= 4.30.1 (due to security vulnerability)\n",
    "\n",
    "\n",
    "#### Table of content:",
    "- [Install Prerequisites](#Install-Prerequisites-Uparrow)\n",
    "- [Convert models to OpenVINO Intermediate representation (IR) format](#Convert-models-to-OpenVINO-Intermediate-representation-(IR)-format-Uparrow)\n",
    "    - [Convert text_encoder](#Convert-text_encoder-Uparrow)\n",
    "    - [The Object Transform](#The-Object-Transform-Uparrow)\n",
    "    - [The Image Encoder](#The-Image-Encoder-Uparrow)\n",
    "    - [Postfuse module](#Postfuse-module-Uparrow)\n",
    "    - [Convert Unet](#Convert-Unet-Uparrow)\n",
    "- [Rebuild pipeline](#Rebuild-pipeline-Uparrow)\n",
    "- [Inference](#Inference-Uparrow)\n",
    "- [Run Gradio](#Run-Gradio-Uparrow)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    This tutorial requires about 25-28GB of free memory to generate one image. Each extra image requires ~11GB of free memory.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Install Prerequisites [$\\Uparrow$](#Table-of-content:)\n",
    "Install required packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "is_executing": true
   },
   "outputs": [],
   "source": [
    "!pip install -q --upgrade pip\n",
    "!pip install -q torch torchvision huggingface-hub\n",
    "!pip install -q transformers accelerate \"diffusers==0.16.1\" gradio\n",
    "!pip install -q \"openvino>=2023.1.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Clone FastComposer project from GitHub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "\n",
    "# clone FastComposer repo\n",
    "if not Path(\"fastcomposer\").exists():\n",
    "    !git clone https://github.com/mit-han-lab/fastcomposer.git\n",
    "else:\n",
    "    print(\"FastComposer repo already cloned\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Download pretrained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "\n",
    "model_path = hf_hub_download(repo_id='mit-han-lab/fastcomposer', filename='pytorch_model.bin')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Convert models to OpenVINO Intermediate representation (IR) format [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "Define a configuration and make instance of `FastComposerModel`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "\n",
    "import torch\n",
    "\n",
    "from model import FastComposerModel\n",
    "\n",
    "\n",
    "@dataclass()\n",
    "class Config:\n",
    "    finetuned_model_path = str(model_path)\n",
    "    image_encoder_name_or_path = 'openai/clip-vit-large-patch14'\n",
    "    localization_layers = 5\n",
    "    mask_loss = False\n",
    "    mask_loss_prob = 0.5\n",
    "    non_ema_revision = None\n",
    "    object_localization = False\n",
    "    object_localization_weight = 0.01\n",
    "    object_resolution = 256\n",
    "    pretrained_model_name_or_path = 'runwayml/stable-diffusion-v1-5'\n",
    "    revision = None\n",
    "\n",
    "\n",
    "config = Config()\n",
    "model = FastComposerModel.from_pretrained(config)\n",
    "model.load_state_dict(torch.load(config.finetuned_model_path, map_location=\"cpu\"), strict=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Pipeline consist of next models: `Unet`, `TextEncoder`, `ImageEncoder` and `PostfuseModule` (MLP), `object_transforms` . \n",
    "\n",
    "\n",
    "![inference-pipeline](https://github.com/openvinotoolkit/openvino_notebooks/assets/29454499/1d858a65-e7c7-43f8-83df-1e896d745725)\n",
    "\n",
    "\n",
    "So, convert the models into OpenVINO IR format."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Convert text_encoder [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "Model components are PyTorch modules, that can be converted with openvino.convert_model function directly. We also use openvino.save_model function to serialize the result of conversion. Let's create a helper function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import gc\n",
    "import openvino\n",
    "\n",
    "\n",
    "def convert(model: torch.nn.Module, xml_path: str, example_input):\n",
    "    xml_path = Path(xml_path)\n",
    "    if not xml_path.exists():\n",
    "        xml_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "        with torch.no_grad():\n",
    "            converted_model = openvino.convert_model(model, example_input=example_input)\n",
    "        openvino.save_model(converted_model, xml_path)\n",
    "        \n",
    "        # cleanup memory\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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "The text encoder is responsible for converting the input prompt into an embedding space that can be fed to the next stage's U-Net. Typically, it is a transformer-based encoder that maps a sequence of input tokens to a sequence of text embeddings.\n",
    "\n",
    "The input for the text encoder consists of a tensor `input_ids`, which contains token indices from the text processed by the tokenizer and padded to the maximum length accepted by the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "text_encoder_ir_xml_path = Path('models/text_encoder_ir.xml')\n",
    "example_input = torch.zeros((1, 77), dtype=torch.int64)\n",
    "\n",
    "model.text_encoder.eval()\n",
    "convert(model.text_encoder, text_encoder_ir_xml_path, example_input)\n",
    "\n",
    "del model.text_encoder\n",
    "gc.collect();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### The Object Transform [$\\Uparrow$](#Table-of-content:)\n",
    "It pads an incoming user image to square and resize it. An input is a tensor of size [3, height, width]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "from torchvision import transforms as T\n",
    "from fastcomposer.fastcomposer.transforms import PadToSquare\n",
    "\n",
    "\n",
    "object_transforms = torch.nn.Sequential(\n",
    "    OrderedDict(\n",
    "        [\n",
    "            (\"pad_to_square\", PadToSquare(fill=0, padding_mode=\"constant\")),\n",
    "            (\n",
    "                \"resize\",\n",
    "                T.Resize(\n",
    "                    (config.object_resolution, config.object_resolution),\n",
    "                    interpolation=T.InterpolationMode.BILINEAR,\n",
    "                    antialias=True,\n",
    "                ),\n",
    "            ),\n",
    "            (\"convert_to_float\", T.ConvertImageDtype(torch.float32)),\n",
    "        ]\n",
    "    )\n",
    ")\n",
    "\n",
    "object_transforms_ir_xml_path = Path('models/object_transforms_ir.xml')\n",
    "example_input = torch.zeros([3, 1500, 1453], dtype=torch.uint8)\n",
    "\n",
    "object_transforms.eval()\n",
    "convert(object_transforms, object_transforms_ir_xml_path, example_input)\n",
    "\n",
    "del object_transforms\n",
    "gc.collect();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### The Image Encoder [$\\Uparrow$](#Table-of-content:)\n",
    "The image encoder is a CLIP (Contrastive Language-Image Pretraining) Image Encoder. It takes a transformed image from the previous step as input and transforms it into a high-dimensional vector or embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "image_encoder_ir_xml_path = Path('models/image_encoder_ir.xml')\n",
    "example_input = torch.zeros((1, 2, 3, 256, 256), dtype=torch.float32)\n",
    "\n",
    "model.image_encoder.eval()\n",
    "convert(model.image_encoder, image_encoder_ir_xml_path, example_input)\n",
    "\n",
    "del model.image_encoder\n",
    "gc.collect();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Postfuse module [$\\Uparrow$](#Table-of-content:)\n",
    "On this step it is employed a multilayer perceptron (MLP) to augment the text embeddings with visual features extracted from the reference subjects. The Postfuse module concatenates the word embeddings with the visual features and feeds the resulting augmented embeddings into the MLP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "postfuse_module_ir_xml_path = Path('models/postfuse_module_ir.xml')\n",
    "\n",
    "example_input = [\n",
    "    torch.zeros((1, 77, 768), dtype=torch.float32),\n",
    "    torch.zeros((1, 2, 1, 768), dtype=torch.float32),\n",
    "    torch.zeros((1, 77), dtype=torch.bool),\n",
    "    torch.zeros((1,), dtype=torch.int64)\n",
    "]\n",
    "\n",
    "model.postfuse_module.eval()\n",
    "convert(model.postfuse_module, postfuse_module_ir_xml_path, example_input)\n",
    "\n",
    "del model.postfuse_module\n",
    "gc.collect();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Convert Unet [$\\Uparrow$](#Table-of-content:)\n",
    "U-Net model gradually denoises latent image representation guided by text encoder hidden state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "unet_ir_xml_path = Path('models/unet_ir.xml')\n",
    "\n",
    "example_input = [\n",
    "    torch.zeros((8, 4, 64, 64), dtype=torch.float32),\n",
    "    torch.zeros((), dtype=torch.int64),\n",
    "    torch.zeros((8, 77, 768), dtype=torch.float32)\n",
    "]\n",
    "model.unet.eval()\n",
    "convert(model.unet, unet_ir_xml_path, example_input)\n",
    "\n",
    "\n",
    "del model\n",
    "del example_input\n",
    "\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Rebuild pipeline [$\\Uparrow$](#Table-of-content:)\n",
    "Also, it needs to modify some internal FastComposer entities, to use OpenVINO models. First of all, how to get results. For example, to convert outputs from numpy to torch types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker\n",
    "from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput\n",
    "from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline\n",
    "from diffusers.loaders import TextualInversionLoaderMixin\n",
    "from diffusers.models import AutoencoderKL, UNet2DConditionModel\n",
    "from typing import Any, Callable, Dict, List, Optional, Union\n",
    "from diffusers.schedulers import KarrasDiffusionSchedulers\n",
    "from transformers import CLIPImageProcessor, CLIPTokenizer\n",
    "from PIL import Image\n",
    "\n",
    "from model import FastComposerTextEncoder\n",
    "\n",
    "\n",
    "class StableDiffusionFastCompposerPipeline(StableDiffusionPipeline):\n",
    "    r\"\"\"\n",
    "    Pipeline for text-to-image generation using FastComposer (https://arxiv.org/abs/2305.10431).\n",
    "\n",
    "    This model inherits from [`StableDiffusionPipeline`]. Check the superclass documentation for the generic methods the\n",
    "    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        vae: AutoencoderKL,\n",
    "        text_encoder: FastComposerTextEncoder,\n",
    "        tokenizer: CLIPTokenizer,\n",
    "        unet: UNet2DConditionModel,\n",
    "        scheduler: KarrasDiffusionSchedulers,\n",
    "        safety_checker: StableDiffusionSafetyChecker,\n",
    "        feature_extractor: CLIPImageProcessor,\n",
    "        requires_safety_checker: bool = True,\n",
    "    ):\n",
    "        super().__init__(\n",
    "            vae,\n",
    "            text_encoder,\n",
    "            tokenizer,\n",
    "            unet,\n",
    "            scheduler,\n",
    "            safety_checker,\n",
    "            feature_extractor,\n",
    "            requires_safety_checker,\n",
    "        )\n",
    "\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def _tokenize_and_mask_noun_phrases_ends(self, caption):\n",
    "        input_ids = self.special_tokenizer.encode(caption)\n",
    "        noun_phrase_end_mask = [False for _ in input_ids]\n",
    "        clean_input_ids = []\n",
    "        clean_index = 0\n",
    "\n",
    "        for i, id in enumerate(input_ids):\n",
    "            if id == self.image_token_id:\n",
    "                noun_phrase_end_mask[clean_index - 1] = True\n",
    "            else:\n",
    "                clean_input_ids.append(id)\n",
    "                clean_index += 1\n",
    "\n",
    "        max_len = self.special_tokenizer.model_max_length\n",
    "\n",
    "        if len(clean_input_ids) > max_len:\n",
    "            clean_input_ids = clean_input_ids[:max_len]\n",
    "        else:\n",
    "            clean_input_ids = clean_input_ids + [self.tokenizer.pad_token_id] * (\n",
    "                max_len - len(clean_input_ids)\n",
    "            )\n",
    "\n",
    "        if len(noun_phrase_end_mask) > max_len:\n",
    "            noun_phrase_end_mask = noun_phrase_end_mask[:max_len]\n",
    "        else:\n",
    "            noun_phrase_end_mask = noun_phrase_end_mask + [False] * (\n",
    "                max_len - len(noun_phrase_end_mask)\n",
    "            )\n",
    "\n",
    "        clean_input_ids = torch.tensor(clean_input_ids, dtype=torch.long)\n",
    "        noun_phrase_end_mask = torch.tensor(noun_phrase_end_mask, dtype=torch.bool)\n",
    "        return clean_input_ids.unsqueeze(0), noun_phrase_end_mask.unsqueeze(0)\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def _encode_augmented_prompt(self, prompt: str, reference_images: List[Image.Image], device: torch.device, weight_dtype: torch.dtype):\n",
    "        # TODO: check this\n",
    "        # encode reference images\n",
    "        object_pixel_values = []\n",
    "        for image in reference_images:\n",
    "            image_tensor = torch.from_numpy(np.array(image.convert(\"RGB\"))).permute(2, 0, 1)\n",
    "            image = torch.from_numpy((self.object_transforms(image_tensor)[0]))\n",
    "            object_pixel_values.append(image)\n",
    "\n",
    "        object_pixel_values = torch.stack(object_pixel_values, dim=0).to(memory_format=torch.contiguous_format).float()\n",
    "        object_pixel_values = object_pixel_values.unsqueeze(0).to(dtype=torch.float32, device=device)\n",
    "        object_embeds = self.image_encoder(object_pixel_values)[0]\n",
    "        object_embeds = torch.from_numpy(object_embeds)\n",
    "\n",
    "        # augment the text embedding\n",
    "        input_ids, image_token_mask = self._tokenize_and_mask_noun_phrases_ends(prompt)\n",
    "        input_ids, image_token_mask = input_ids.to(device), image_token_mask.to(device)\n",
    "\n",
    "        num_objects = image_token_mask.sum(dim=1)\n",
    "\n",
    "        text_embeds = torch.from_numpy(self.text_encoder(input_ids)[0])\n",
    "        augmented_prompt_embeds = self.postfuse_module([\n",
    "            text_embeds,\n",
    "            object_embeds,\n",
    "            image_token_mask,\n",
    "            num_objects\n",
    "        ])[0]\n",
    "        return torch.from_numpy(augmented_prompt_embeds)\n",
    "\n",
    "    def _encode_prompt(\n",
    "        self,\n",
    "        prompt,\n",
    "        device,\n",
    "        num_images_per_prompt,\n",
    "        do_classifier_free_guidance,\n",
    "        negative_prompt=None\n",
    "    ):\n",
    "        r\"\"\"\n",
    "        Encodes the prompt into text encoder hidden states.\n",
    "\n",
    "        Args:\n",
    "             prompt (`str` or `List[str]`, *optional*):\n",
    "                prompt to be encoded\n",
    "            device: (`torch.device`):\n",
    "                torch device\n",
    "            num_images_per_prompt (`int`):\n",
    "                number of images that should be generated per prompt\n",
    "            do_classifier_free_guidance (`bool`):\n",
    "                whether to use classifier free guidance or not\n",
    "            negative_prompt (`str` or `List[str]`, *optional*):\n",
    "                The prompt or prompts not to guide the image generation. If not defined, one has to pass\n",
    "                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is\n",
    "                less than `1`).\n",
    "        \"\"\"\n",
    "        if isinstance(prompt, str):\n",
    "            batch_size = 1\n",
    "        elif isinstance(prompt, list):\n",
    "            batch_size = len(prompt)\n",
    "\n",
    "        # textual inversion: procecss multi-vector tokens if necessary\n",
    "        if isinstance(self, TextualInversionLoaderMixin):\n",
    "            prompt = self.maybe_convert_prompt(prompt, self.tokenizer)\n",
    "\n",
    "        text_inputs = self.tokenizer(\n",
    "            prompt,\n",
    "            padding=\"max_length\",\n",
    "            max_length=self.tokenizer.model_max_length,\n",
    "            truncation=True,\n",
    "            return_tensors=\"pt\",\n",
    "        )\n",
    "        text_input_ids = text_inputs.input_ids\n",
    "        untruncated_ids = self.tokenizer(prompt, padding=\"longest\", return_tensors=\"pt\").input_ids\n",
    "\n",
    "        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(\n",
    "            text_input_ids, untruncated_ids\n",
    "        ):\n",
    "            removed_text = self.tokenizer.batch_decode(\n",
    "                untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]\n",
    "            )\n",
    "            print(\n",
    "                \"The following part of your input was truncated because CLIP can only handle sequences up to\"\n",
    "                f\" {self.tokenizer.model_max_length} tokens: {removed_text}\"\n",
    "            )\n",
    "\n",
    "        prompt_embeds = self.text_encoder(text_input_ids.to(device))[0]\n",
    "        prompt_embeds = torch.from_numpy(prompt_embeds)\n",
    "\n",
    "        bs_embed, seq_len, _ = prompt_embeds.shape\n",
    "        # duplicate text embeddings for each generation per prompt, using mps friendly method\n",
    "        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)\n",
    "        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)\n",
    "\n",
    "        # get unconditional embeddings for classifier free guidance\n",
    "        if do_classifier_free_guidance:\n",
    "            uncond_tokens: List[str]\n",
    "            if negative_prompt is None:\n",
    "                uncond_tokens = [\"\"] * batch_size\n",
    "            elif type(prompt) is not type(negative_prompt):\n",
    "                raise TypeError(\n",
    "                    f\"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=\"\n",
    "                    f\" {type(prompt)}.\"\n",
    "                )\n",
    "            elif isinstance(negative_prompt, str):\n",
    "                uncond_tokens = [negative_prompt]\n",
    "            elif batch_size != len(negative_prompt):\n",
    "                raise ValueError(\n",
    "                    f\"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:\"\n",
    "                    f\" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches\"\n",
    "                    \" the batch size of `prompt`.\"\n",
    "                )\n",
    "            else:\n",
    "                uncond_tokens = negative_prompt\n",
    "\n",
    "            # textual inversion: procecss multi-vector tokens if necessary\n",
    "            if isinstance(self, TextualInversionLoaderMixin):\n",
    "                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)\n",
    "\n",
    "            max_length = prompt_embeds.shape[1]\n",
    "            uncond_input = self.tokenizer(\n",
    "                uncond_tokens,\n",
    "                padding=\"max_length\",\n",
    "                max_length=max_length,\n",
    "                truncation=True,\n",
    "                return_tensors=\"pt\",\n",
    "            )\n",
    "\n",
    "            negative_prompt_embeds = self.text_encoder(uncond_input.input_ids.to(device))[0]\n",
    "            negative_prompt_embeds = torch.from_numpy(negative_prompt_embeds)\n",
    "\n",
    "        if do_classifier_free_guidance:\n",
    "            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method\n",
    "            seq_len = negative_prompt_embeds.shape[1]\n",
    "\n",
    "            negative_prompt_embeds = negative_prompt_embeds.to(dtype=torch.float32, device=device)\n",
    "\n",
    "            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)\n",
    "            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)\n",
    "\n",
    "            # For classifier free guidance, we need to do two forward passes.\n",
    "            # Here we concatenate the unconditional and text embeddings into a single batch\n",
    "            # to avoid doing two forward passes\n",
    "            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])\n",
    "\n",
    "        return prompt_embeds\n",
    "\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def __call__(\n",
    "        self,\n",
    "        prompt: Union[str, List[str]] = None,\n",
    "        height: Optional[int] = None,\n",
    "        width: Optional[int] = None,\n",
    "        num_inference_steps: int = 50,\n",
    "        guidance_scale: float = 7.5,\n",
    "        negative_prompt: Optional[Union[str, List[str]]] = None,\n",
    "        num_images_per_prompt: Optional[int] = 1,\n",
    "        eta: float = 0.0,\n",
    "        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,\n",
    "        latents: Optional[torch.FloatTensor] = None,\n",
    "        prompt_embeds: Optional[torch.FloatTensor] = None,\n",
    "        negative_prompt_embeds: Optional[torch.FloatTensor] = None,\n",
    "        output_type: Optional[str] = \"pil\",\n",
    "        return_dict: bool = True,\n",
    "        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,\n",
    "        callback_steps: int = 1,\n",
    "        cross_attention_kwargs: Optional[Dict[str, Any]] = None,\n",
    "        alpha_: float = 0.7,\n",
    "        reference_subject_images: List[Image.Image] = None,\n",
    "        augmented_prompt_embeds: Optional[torch.FloatTensor] = None\n",
    "    ):\n",
    "        r\"\"\"\n",
    "        Function invoked when calling the pipeline for generation.\n",
    "\n",
    "        Args:\n",
    "            prompt (`str` or `List[str]`, *optional*):\n",
    "                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.\n",
    "                instead.\n",
    "            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):\n",
    "                The height in pixels of the generated image.\n",
    "            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):\n",
    "                The width in pixels of the generated image.\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. of [Imagen\n",
    "                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >\n",
    "                1`. 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",
    "            negative_prompt (`str` or `List[str]`, *optional*):\n",
    "                The prompt or prompts not to guide the image generation. If not defined, one has to pass\n",
    "                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is\n",
    "                less than `1`).\n",
    "            num_images_per_prompt (`int`, *optional*, defaults to 1):\n",
    "                The number of images to generate per prompt.\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",
    "                [`schedulers.DDIMScheduler`], will be ignored for others.\n",
    "            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):\n",
    "                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)\n",
    "                to make generation deterministic.\n",
    "            latents (`torch.FloatTensor`, *optional*):\n",
    "                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image\n",
    "                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents\n",
    "                tensor will ge generated by sampling using the supplied random `generator`.\n",
    "            prompt_embeds (`torch.FloatTensor`, *optional*):\n",
    "                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not\n",
    "                provided, text embeddings will be generated from `prompt` input argument.\n",
    "            negative_prompt_embeds (`torch.FloatTensor`, *optional*):\n",
    "                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt\n",
    "                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input\n",
    "                argument.\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",
    "            return_dict (`bool`, *optional*, defaults to `True`):\n",
    "                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a\n",
    "                plain tuple.\n",
    "            callback (`Callable`, *optional*):\n",
    "                A function that will be called every `callback_steps` steps during inference. The function will be\n",
    "                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.\n",
    "            callback_steps (`int`, *optional*, defaults to 1):\n",
    "                The frequency at which the `callback` function will be called. If not specified, the callback will be\n",
    "                called at every step.\n",
    "            cross_attention_kwargs (`dict`, *optional*):\n",
    "                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under\n",
    "                `self.processor` in\n",
    "                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).\n",
    "            alpha_ (`float`, defaults to 0.7):\n",
    "                The ratio of subject conditioning. If `alpha_` is 0.7, the beginning 30% of denoising steps use text prompts, while the\n",
    "                last 70% utilize image-augmented prompts. Increase alpha for identity preservation, decrease it for prompt consistency.\n",
    "            reference_subject_images (`List[PIL.Image.Image]`):\n",
    "                a list of PIL images that are used as reference subjects. The number of images should be equal to the number of augmented\n",
    "                tokens in the prompts.\n",
    "            augmented_prompt_embeds: (`torch.FloatTensor`, *optional*):\n",
    "                Pre-generated image augmented text embeddings. If not provided, embeddings will be generated from `prompt` and\n",
    "                `reference_subject_images`.\n",
    "        Examples:\n",
    "\n",
    "        Returns:\n",
    "            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:\n",
    "            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.\n",
    "            When returning a tuple, the first element is a list with the generated images, and the second element is a\n",
    "            list of `bool`s denoting whether the corresponding generated image likely represents \"not-safe-for-work\"\n",
    "            (nsfw) content, according to the `safety_checker`.\n",
    "        \"\"\"\n",
    "        # 0. Default height and width to unet\n",
    "        height = height or self.unet.config.sample_size * self.vae_scale_factor\n",
    "        width = width or self.unet.config.sample_size * self.vae_scale_factor\n",
    "\n",
    "        # 1. Check inputs. Raise error if not correct\n",
    "        self.check_inputs(\n",
    "            prompt,\n",
    "            height,\n",
    "            width,\n",
    "            callback_steps,\n",
    "            negative_prompt,\n",
    "            prompt_embeds,\n",
    "            negative_prompt_embeds,\n",
    "        )\n",
    "\n",
    "        assert (prompt is not None and reference_subject_images is not None) or (prompt_embeds is not None and augmented_prompt_embeds is not None),  \\\n",
    "            \"Prompt and reference subject images or prompt_embeds and augmented_prompt_embeds must be provided.\"\n",
    "\n",
    "        # 2. Define call parameters\n",
    "        if prompt is not None and isinstance(prompt, str):\n",
    "            batch_size = 1\n",
    "        elif prompt is not None and isinstance(prompt, list):\n",
    "            batch_size = len(prompt)\n",
    "        else:\n",
    "            batch_size = prompt_embeds.shape[0]\n",
    "\n",
    "        device = self._execution_device\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",
    "        assert do_classifier_free_guidance\n",
    "\n",
    "        # 3. Encode input prompt\n",
    "        prompt_text_only = prompt.replace(\"<image>\", \"\")\n",
    "\n",
    "        prompt_embeds = self._encode_prompt(\n",
    "            prompt_text_only,\n",
    "            device,\n",
    "            num_images_per_prompt,\n",
    "            do_classifier_free_guidance,\n",
    "            negative_prompt,\n",
    "        )\n",
    "\n",
    "        if augmented_prompt_embeds is None:\n",
    "            augmented_prompt_embeds = self._encode_augmented_prompt(prompt, reference_subject_images, device, prompt_embeds.dtype)\n",
    "            augmented_prompt_embeds = augmented_prompt_embeds.repeat(num_images_per_prompt, 1, 1)\n",
    "\n",
    "        prompt_embeds = torch.cat([prompt_embeds, augmented_prompt_embeds], dim=0)\n",
    "\n",
    "        # 4. Prepare timesteps\n",
    "        self.scheduler.set_timesteps(num_inference_steps, device=device)\n",
    "        timesteps = self.scheduler.timesteps\n",
    "\n",
    "        # 5. Prepare latent variables\n",
    "        # num_channels_latents = self.unet.in_channels\n",
    "        num_channels_latents = 4\n",
    "        latents = self.prepare_latents(\n",
    "            batch_size * num_images_per_prompt,\n",
    "            num_channels_latents,\n",
    "            height,\n",
    "            width,\n",
    "            prompt_embeds.dtype,\n",
    "            device,\n",
    "            generator,\n",
    "            latents,\n",
    "        )\n",
    "\n",
    "        start_subject_conditioning_step = (1 - alpha_) * num_inference_steps\n",
    "\n",
    "        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)\n",
    "        (\n",
    "            null_prompt_embeds,\n",
    "            text_prompt_embeds,\n",
    "            augmented_prompt_embeds\n",
    "        ) = prompt_embeds.chunk(3)\n",
    "\n",
    "        # 7. Denoising loop\n",
    "        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order\n",
    "        with self.progress_bar(total=num_inference_steps) as progress_bar:\n",
    "            for i, t in enumerate(timesteps):\n",
    "                latent_model_input = (\n",
    "                    torch.cat([latents] * 2) if do_classifier_free_guidance else latents\n",
    "                )\n",
    "                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)\n",
    "\n",
    "                if i <= start_subject_conditioning_step:\n",
    "                    current_prompt_embeds = torch.cat(\n",
    "                        [null_prompt_embeds, text_prompt_embeds], dim=0\n",
    "                    )\n",
    "                else:\n",
    "                    current_prompt_embeds = torch.cat(\n",
    "                        [null_prompt_embeds, augmented_prompt_embeds], dim=0\n",
    "                    )\n",
    "\n",
    "                # predict the noise residual\n",
    "                noise_pred = self.unet([\n",
    "                    latent_model_input,\n",
    "                    t,\n",
    "                    current_prompt_embeds,\n",
    "                    # cross_attention_kwargs\n",
    "                ],\n",
    "                )[0]\n",
    "                noise_pred = torch.from_numpy(noise_pred)\n",
    "\n",
    "\n",
    "                # perform guidance\n",
    "                if do_classifier_free_guidance:\n",
    "                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n",
    "                    noise_pred = noise_pred_uncond + guidance_scale * (\n",
    "                        noise_pred_text - noise_pred_uncond\n",
    "                    )\n",
    "                else:\n",
    "                    assert 0, \"Not Implemented\"\n",
    "\n",
    "                # compute the previous noisy sample x_t -> x_t-1\n",
    "                latents = self.scheduler.step(\n",
    "                    noise_pred, t, latents, **extra_step_kwargs\n",
    "                ).prev_sample\n",
    "\n",
    "                # call the callback, if provided\n",
    "                if i == len(timesteps) - 1 or (\n",
    "                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0\n",
    "                ):\n",
    "                    progress_bar.update()\n",
    "                    if callback is not None and i % callback_steps == 0:\n",
    "                        callback(i, t, latents)\n",
    "\n",
    "        if output_type == \"latent\":\n",
    "            image = latents\n",
    "            has_nsfw_concept = None\n",
    "        elif output_type == \"pil\":\n",
    "            # 8. Post-processing\n",
    "            image = self.decode_latents(latents)\n",
    "\n",
    "            # 9. Run safety checker\n",
    "            image, has_nsfw_concept = self.run_safety_checker(\n",
    "                image, device, prompt_embeds.dtype\n",
    "            )\n",
    "\n",
    "            # 10. Convert to PIL\n",
    "            image = self.numpy_to_pil(image)\n",
    "        else:\n",
    "            # 8. Post-processing\n",
    "            image = self.decode_latents(latents)\n",
    "\n",
    "            # 9. Run safety checker\n",
    "            image, has_nsfw_concept = self.run_safety_checker(\n",
    "                image, device, prompt_embeds.dtype\n",
    "            )\n",
    "\n",
    "        # Offload last model to CPU\n",
    "        if hasattr(self, \"final_offload_hook\") and self.final_offload_hook is not None:\n",
    "            self.final_offload_hook.offload()\n",
    "\n",
    "        if not return_dict:\n",
    "            return (image, has_nsfw_concept)\n",
    "\n",
    "        return StableDiffusionPipelineOutput(\n",
    "            images=image, nsfw_content_detected=has_nsfw_concept\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "And replace all model in the pipeline by converted models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import PIL\n",
    "from transformers import CLIPTokenizer\n",
    "\n",
    "\n",
    "def create_pipeline(\n",
    "        args,\n",
    "        *,\n",
    "        text_encoder,\n",
    "        image_encoder,\n",
    "        unet,\n",
    "        object_transforms,\n",
    "        postfuse_module,\n",
    "        device\n",
    "):\n",
    "    weight_dtype = torch.float32\n",
    "\n",
    "    tokenizer = CLIPTokenizer.from_pretrained(\n",
    "        args.pretrained_model_name_or_path,\n",
    "        subfolder=\"tokenizer\",\n",
    "        revision=args.revision,\n",
    "    )\n",
    "    tokenizer.add_tokens([\"img\"], special_tokens=True)\n",
    "    image_token_id = tokenizer.convert_tokens_to_ids(\"img\")\n",
    "\n",
    "    pipe = StableDiffusionFastCompposerPipeline.from_pretrained(\n",
    "        args.pretrained_model_name_or_path, torch_dtype=weight_dtype\n",
    "    ).to(device)\n",
    "\n",
    "    pipe.object_transforms = object_transforms\n",
    "    pipe.unet = unet\n",
    "    pipe.text_encoder = text_encoder\n",
    "    pipe.postfuse_module = postfuse_module\n",
    "    pipe.image_encoder = image_encoder\n",
    "    pipe.image_token_id = image_token_id\n",
    "    pipe.special_tokenizer = tokenizer\n",
    "\n",
    "    return pipe\n",
    "\n",
    "\n",
    "class ModelWrapper:\n",
    "    def __init__(self, model):\n",
    "        super().__init__()\n",
    "        self.model = model\n",
    "\n",
    "    def inference(\n",
    "        self,\n",
    "        image1: PIL.Image.Image,\n",
    "        image2: PIL.Image.Image,\n",
    "        prompt: str,\n",
    "        negative_prompt: str,\n",
    "        seed: int,\n",
    "        guidance_scale: float,\n",
    "        alpha_: float,\n",
    "        num_steps: int,\n",
    "        num_images: int,\n",
    "    ):\n",
    "        print(\"Running model inference...\")\n",
    "        image = []\n",
    "        if image1 is not None:\n",
    "            image.append(image1)\n",
    "\n",
    "        if image2 is not None:\n",
    "            image.append(image2)\n",
    "\n",
    "        if len(image) == 0:\n",
    "            return [], \"You need to upload at least one image.\"\n",
    "\n",
    "        num_subject_in_text = (\n",
    "            np.array(self.model.special_tokenizer.encode(prompt))\n",
    "            == self.model.image_token_id\n",
    "        ).sum()\n",
    "        if num_subject_in_text != len(image):\n",
    "            return (\n",
    "                [],\n",
    "                f\"Number of subjects in the text description doesn't match the number of reference images, #text subjects: {num_subject_in_text} #reference image: {len(image)}\",\n",
    "            )\n",
    "\n",
    "        if seed == -1:\n",
    "            seed = np.random.randint(0, 1000000)\n",
    "\n",
    "        generator = torch.manual_seed(seed)\n",
    "\n",
    "        return (\n",
    "            self.model(\n",
    "                prompt=prompt,\n",
    "                negative_prompt=negative_prompt,\n",
    "                height=512,\n",
    "                width=512,\n",
    "                num_inference_steps=num_steps,\n",
    "                guidance_scale=guidance_scale,\n",
    "                num_images_per_prompt=num_images,\n",
    "                generator=generator,\n",
    "                alpha_=alpha_,\n",
    "                reference_subject_images=image,\n",
    "            ).images,\n",
    "            \"run successfully\",\n",
    "        )\n",
    "\n",
    "\n",
    "core = openvino.Core()\n",
    "compiled_unet = core.compile_model(unet_ir_xml_path)\n",
    "compiled_text_encoder = core.compile_model(text_encoder_ir_xml_path)\n",
    "compiled_image_encoder = core.compile_model(image_encoder_ir_xml_path)\n",
    "compiled_postfuse_module = core.compile_model(postfuse_module_ir_xml_path)\n",
    "compiled_object_transforms = core.compile_model(object_transforms_ir_xml_path)\n",
    "\n",
    "wrapped_model = ModelWrapper(\n",
    "    create_pipeline(\n",
    "        config,\n",
    "        text_encoder=compiled_text_encoder,\n",
    "        image_encoder=compiled_image_encoder,\n",
    "        unet=compiled_unet,\n",
    "        object_transforms=compiled_object_transforms,\n",
    "        postfuse_module=compiled_postfuse_module,\n",
    "        device='cpu'\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Inference [$\\Uparrow$](#Table-of-content:)\n",
    "And now it is possible to make inference. You can provide 1 or 2 images (`image1` and `image2`). If you want to provide only one image pass in inference `None` instead image. \n",
    "`prompt` describes context in what objects from user images will be generated. Word `img` is a token that correlates with input images.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "image1 = Image.open('fastcomposer/data/newton_einstein/einstein/0.png')\n",
    "image2 = Image.open('fastcomposer/data/newton_einstein/newton/0.png')\n",
    "prompt = 'A man img and a man img sitting in a park'\n",
    "negative_prompt = '((((ugly)))), (((duplicate))), ((morbid)), ((mutilated)), [out of frame], extra fingers, mutated hands, ((poorly drawn hands)), ((poorly drawn face)), (((mutation))), (((deformed))), ((ugly)), blurry, ((bad anatomy)), (((bad proportions))), ((extra limbs)), cloned face, (((disfigured))). out of frame, ugly, extra limbs, (bad anatomy), gross proportions, (malformed limbs), ((missing arms)), ((missing legs)), (((extra arms))), (((extra legs))), mutated hands, (fused fingers), (too many fingers), (((long neck)))'\n",
    "alpha_ = 0.7\n",
    "num_images = 1  # each extra image requires ~11GB of free memory\n",
    "num_steps = 50\n",
    "guidance_scale = 5\n",
    "seed = -1\n",
    "\n",
    "\n",
    "result = wrapped_model.inference(\n",
    "    image1,\n",
    "    image2,\n",
    "    prompt,\n",
    "    negative_prompt,\n",
    "    seed,\n",
    "    guidance_scale,\n",
    "    alpha_,\n",
    "    num_steps,\n",
    "    num_images\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Result consists of several (`num_images`) images and now it possible to display them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "display(result[0][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Run Gradio [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "Also, it is possible to run with Gradio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "\n",
    "\n",
    "def create_demo():\n",
    "    TITLE = \"# [FastComposer Demo](https://github.com/mit-han-lab/fastcomposer) with OpenVINO\"\n",
    "\n",
    "    DESCRIPTION = \"\"\"To run the demo, you should:\n",
    "    1. Upload your images. The order of image1 and image2 needs to match the order of the subects in the prompt. You only need 1 image for single subject generation.\n",
    "    2. Input proper text prompts, such as \"A woman img and a man img in the snow\" or \"A painting of a man img in the style of Van Gogh\", where \"img\" specifies the token you want to augment and comes after the word.\n",
    "    3. Click the Run button. You can also adjust the hyperparameters to improve the results. Look at the job status to see if there are any errors with your input.\n",
    "    As a result, pictures with person or persons from input images will be generated in accordance with the description in the prompt.\n",
    "    \"\"\"\n",
    "\n",
    "    with gr.Blocks() as demo:\n",
    "        gr.Markdown(TITLE)\n",
    "        gr.Markdown(DESCRIPTION)\n",
    "        with gr.Row():\n",
    "            with gr.Column():\n",
    "                with gr.Box():\n",
    "                    image1 = gr.Image(label=\"Image 1\", type=\"pil\")\n",
    "                    gr.Examples(\n",
    "                        examples=[\"fastcomposer/data/newton.jpeg\"],\n",
    "                        inputs=image1,\n",
    "                    )\n",
    "                    image2 = gr.Image(label=\"Image 2\", type=\"pil\")\n",
    "                    gr.Examples(\n",
    "                        examples=[\"fastcomposer/data/einstein.jpeg\"],\n",
    "                        inputs=image2,\n",
    "                    )\n",
    "                    gr.Markdown(\"Upload the image for your subject\")\n",
    "\n",
    "                prompt = gr.Text(\n",
    "                    value=\"A man img and a man img sitting in a park\",\n",
    "                    label=\"Prompt\",\n",
    "                    placeholder='e.g. \"A woman img and a man img in the snow\", \"A painting of a man img in the style of Van Gogh\"',\n",
    "                    info='Use \"img\" to specify the word you want to augment.',\n",
    "                )\n",
    "                negative_prompt = gr.Text(\n",
    "                    value=\"((((ugly)))), (((duplicate))), ((morbid)), ((mutilated)), [out of frame], extra fingers, mutated hands, ((poorly drawn hands)), ((poorly drawn face)), (((mutation))), (((deformed))), ((ugly)), blurry, ((bad anatomy)), (((bad proportions))), ((extra limbs)), cloned face, (((disfigured))). out of frame, ugly, extra limbs, (bad anatomy), gross proportions, (malformed limbs), ((missing arms)), ((missing legs)), (((extra arms))), (((extra legs))), mutated hands, (fused fingers), (too many fingers), (((long neck)))\",\n",
    "                    label=\"Negative Prompt\",\n",
    "                    info='Features that you want to avoid.',\n",
    "                )\n",
    "                alpha_ = gr.Slider(\n",
    "                    label=\"alpha\",\n",
    "                    minimum=0,\n",
    "                    maximum=1,\n",
    "                    step=0.05,\n",
    "                    value=0.75,\n",
    "                    info=\"A smaller alpha aligns images with text better, but may deviate from the subject image. Increase alpha to improve identity preservation, decrease it for prompt consistency.\",\n",
    "                )\n",
    "                num_images = gr.Slider(\n",
    "                    label=\"Number of generated images\",\n",
    "                    minimum=1,\n",
    "                    maximum=8,\n",
    "                    step=1,\n",
    "                    value=1,\n",
    "                    info=\"Each extra image requires ~11GB of free memory.\",\n",
    "                )\n",
    "                run_button = gr.Button(\"Run\")\n",
    "                with gr.Accordion(label=\"Advanced options\", open=False):\n",
    "                    seed = gr.Slider(\n",
    "                        label=\"Seed\",\n",
    "                        minimum=-1,\n",
    "                        maximum=1000000,\n",
    "                        step=1,\n",
    "                        value=-1,\n",
    "                        info=\"If set to -1, a different seed will be used each time.\",\n",
    "                    )\n",
    "                    guidance_scale = gr.Slider(\n",
    "                        label=\"Guidance scale\",\n",
    "                        minimum=1,\n",
    "                        maximum=10,\n",
    "                        step=1,\n",
    "                        value=5,\n",
    "                    )\n",
    "                    num_steps = gr.Slider(\n",
    "                        label=\"Steps\",\n",
    "                        minimum=1,\n",
    "                        maximum=300,\n",
    "                        step=1,\n",
    "                        value=50,\n",
    "                    )\n",
    "            with gr.Column():\n",
    "                result = gr.Gallery(label=\"Generated Images\").style(\n",
    "                    grid=[2], height=\"auto\"\n",
    "                )\n",
    "                error_message = gr.Text(label=\"Job Status\")\n",
    "\n",
    "        inputs = [\n",
    "            image1,\n",
    "            image2,\n",
    "            prompt,\n",
    "            negative_prompt,\n",
    "            seed,\n",
    "            guidance_scale,\n",
    "            alpha_,\n",
    "            num_steps,\n",
    "            num_images,\n",
    "        ]\n",
    "        run_button.click(\n",
    "            fn=wrapped_model.inference, inputs=inputs, outputs=[result, error_message]\n",
    "        )\n",
    "    return demo\n",
    "\n",
    "\n",
    "demo = create_demo()\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    try:\n",
    "        demo.launch(debug=True)\n",
    "    except Exception:\n",
    "        demo.launch(share=True, debug=True)\n",
    "# if you are launching remotely, specify server_name and server_port\n",
    "# demo.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",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}