{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bcafe28c-9f59-4d96-b52f-3e3f45a16c3a",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Copyright 2022 Google LLC. Double-click for license information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a952f9ab-ed6f-4e99-8304-99a3716734b5",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Copyright 2022 Google LLC\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#      http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ed9ffa91",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d2b343b-1c90-4747-a753-71eb7071a289",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Null-text inversion + Editing with Prompt-to-Prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7b1b6199-9dfe-4055-8a84-66ff4bfa8901",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.24.2\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"
     ]
    }
   ],
   "source": [
    "from typing import Optional, Union, Tuple, List, Callable, Dict\n",
    "from tqdm.notebook import tqdm\n",
    "import torch\n",
    "from diffusers import StableDiffusionPipeline, DDIMScheduler\n",
    "import torch.nn.functional as nnf\n",
    "import numpy as np\n",
    "import abc\n",
    "import ptp_utils\n",
    "import seq_aligner\n",
    "import shutil\n",
    "from torch.optim.adam import Adam\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c6a4bd1-3130-408b-ae2d-a166b9f19cb7",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "For loading the Stable Diffusion using Diffusers, follow the instuctions https://huggingface.co/blog/stable_diffusion and update MY_TOKEN with your token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7558a4b4-fec6-4bd2-9c8f-139809b1a1a1",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Fetching 13 files: 100%|██████████| 13/13 [00:00<00:00, 76581.39it/s]\n",
      "/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/transformers/models/clip/feature_extraction_clip.py:28: FutureWarning: The class CLIPFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please use CLIPImageProcessor instead.\n",
      "  warnings.warn(\n",
      "/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py:106: FutureWarning: The configuration file of this scheduler: DDIMScheduler {\n",
      "  \"_class_name\": \"DDIMScheduler\",\n",
      "  \"_diffusers_version\": \"0.14.0\",\n",
      "  \"beta_end\": 0.012,\n",
      "  \"beta_schedule\": \"scaled_linear\",\n",
      "  \"beta_start\": 0.00085,\n",
      "  \"clip_sample\": false,\n",
      "  \"num_train_timesteps\": 1000,\n",
      "  \"prediction_type\": \"epsilon\",\n",
      "  \"set_alpha_to_one\": false,\n",
      "  \"steps_offset\": 0,\n",
      "  \"trained_betas\": null\n",
      "}\n",
      " is outdated. `steps_offset` should be set to 1 instead of 0. Please make sure to update the config accordingly as leaving `steps_offset` might led to incorrect results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json` file\n",
      "  deprecate(\"steps_offset!=1\", \"1.0.0\", deprecation_message, standard_warn=False)\n"
     ]
    }
   ],
   "source": [
    "scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule=\"scaled_linear\", clip_sample=False, set_alpha_to_one=False)\n",
    "MY_TOKEN = ''\n",
    "LOW_RESOURCE = False \n",
    "NUM_DDIM_STEPS = 50\n",
    "GUIDANCE_SCALE = 7.5\n",
    "MAX_NUM_WORDS = 77\n",
    "device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')\n",
    "ldm_stable = StableDiffusionPipeline.from_pretrained(\"stabilityai/stable-diffusion-2-1-base\", use_auth_token=MY_TOKEN, scheduler=scheduler).to(device)\n",
    "try:\n",
    "    ldm_stable.disable_xformers_memory_efficient_attention()\n",
    "except AttributeError:\n",
    "    print(\"Attribute disable_xformers_memory_efficient_attention() is missing\")\n",
    "tokenizer = ldm_stable.tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01422991-cafe-4cf0-8406-66f052a75d9b",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Prompt-to-Prompt code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "dc083590-15de-4216-8d8d-50336f9f1d34",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "class LocalBlend:\n",
    "    \n",
    "    def get_mask(self, maps, alpha, use_pool):\n",
    "        k = 1\n",
    "        maps = (maps * alpha).sum(-1).mean(1)\n",
    "        if use_pool:\n",
    "            maps = nnf.max_pool2d(maps, (k * 2 + 1, k * 2 +1), (1, 1), padding=(k, k))\n",
    "        mask = nnf.interpolate(maps, size=(x_t.shape[2:]))\n",
    "        mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0]\n",
    "        mask = mask.gt(self.th[1-int(use_pool)])\n",
    "        mask = mask[:1] + mask\n",
    "        return mask\n",
    "    \n",
    "    def __call__(self, x_t, attention_store):\n",
    "        self.counter += 1\n",
    "        if self.counter > self.start_blend:\n",
    "           \n",
    "            maps = attention_store[\"down_cross\"][2:4] + attention_store[\"up_cross\"][:3]\n",
    "            maps = [item.reshape(self.alpha_layers.shape[0], -1, 1, 16, 16, MAX_NUM_WORDS) for item in maps]\n",
    "            maps = torch.cat(maps, dim=1)\n",
    "            mask = self.get_mask(maps, self.alpha_layers, True)\n",
    "            if self.substruct_layers is not None:\n",
    "                maps_sub = ~self.get_mask(maps, self.substruct_layers, False)\n",
    "                mask = mask * maps_sub\n",
    "            mask = mask.float()\n",
    "            x_t = x_t[:1] + mask * (x_t - x_t[:1])\n",
    "        return x_t\n",
    "       \n",
    "    def __init__(self, prompts: List[str], words: [List[List[str]]], substruct_words=None, start_blend=0.2, th=(.3, .3)):\n",
    "        alpha_layers = torch.zeros(len(prompts),  1, 1, 1, 1, MAX_NUM_WORDS)\n",
    "        for i, (prompt, words_) in enumerate(zip(prompts, words)):\n",
    "            if type(words_) is str:\n",
    "                words_ = [words_]\n",
    "            for word in words_:\n",
    "                ind = ptp_utils.get_word_inds(prompt, word, tokenizer)\n",
    "                alpha_layers[i, :, :, :, :, ind] = 1\n",
    "        \n",
    "        if substruct_words is not None:\n",
    "            substruct_layers = torch.zeros(len(prompts),  1, 1, 1, 1, MAX_NUM_WORDS)\n",
    "            for i, (prompt, words_) in enumerate(zip(prompts, substruct_words)):\n",
    "                if type(words_) is str:\n",
    "                    words_ = [words_]\n",
    "                for word in words_:\n",
    "                    ind = ptp_utils.get_word_inds(prompt, word, tokenizer)\n",
    "                    substruct_layers[i, :, :, :, :, ind] = 1\n",
    "            self.substruct_layers = substruct_layers.to(device)\n",
    "        else:\n",
    "            self.substruct_layers = None\n",
    "        self.alpha_layers = alpha_layers.to(device)\n",
    "        self.start_blend = int(start_blend * NUM_DDIM_STEPS)\n",
    "        self.counter = 0 \n",
    "        self.th=th\n",
    "\n",
    "\n",
    "        \n",
    "        \n",
    "class EmptyControl:\n",
    "    \n",
    "    \n",
    "    def step_callback(self, x_t):\n",
    "        return x_t\n",
    "    \n",
    "    def between_steps(self):\n",
    "        return\n",
    "    \n",
    "    def __call__(self, attn, is_cross: bool, place_in_unet: str):\n",
    "        return attn\n",
    "\n",
    "    \n",
    "class AttentionControl(abc.ABC):\n",
    "    \n",
    "    def step_callback(self, x_t):\n",
    "        return x_t\n",
    "    \n",
    "    def between_steps(self):\n",
    "        return\n",
    "    \n",
    "    @property\n",
    "    def num_uncond_att_layers(self):\n",
    "        return self.num_att_layers if LOW_RESOURCE else 0\n",
    "    \n",
    "    @abc.abstractmethod\n",
    "    def forward (self, attn, is_cross: bool, place_in_unet: str):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def __call__(self, attn, is_cross: bool, place_in_unet: str):\n",
    "        if self.cur_att_layer >= self.num_uncond_att_layers:\n",
    "            if LOW_RESOURCE:\n",
    "                attn = self.forward(attn, is_cross, place_in_unet)\n",
    "            else:\n",
    "                h = attn.shape[0]\n",
    "                attn[h // 2:] = self.forward(attn[h // 2:], is_cross, place_in_unet)\n",
    "        self.cur_att_layer += 1\n",
    "        if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers:\n",
    "            self.cur_att_layer = 0\n",
    "            self.cur_step += 1\n",
    "            self.between_steps()\n",
    "        return attn\n",
    "    \n",
    "    def reset(self):\n",
    "        self.cur_step = 0\n",
    "        self.cur_att_layer = 0\n",
    "\n",
    "    def __init__(self):\n",
    "        self.cur_step = 0\n",
    "        self.num_att_layers = -1\n",
    "        self.cur_att_layer = 0\n",
    "\n",
    "class SpatialReplace(EmptyControl):\n",
    "    \n",
    "    def step_callback(self, x_t):\n",
    "        if self.cur_step < self.stop_inject:\n",
    "            b = x_t.shape[0]\n",
    "            x_t = x_t[:1].expand(b, *x_t.shape[1:])\n",
    "        return x_t\n",
    "\n",
    "    def __init__(self, stop_inject: float):\n",
    "        super(SpatialReplace, self).__init__()\n",
    "        self.stop_inject = int((1 - stop_inject) * NUM_DDIM_STEPS)\n",
    "        \n",
    "\n",
    "class AttentionStore(AttentionControl):\n",
    "\n",
    "    @staticmethod\n",
    "    def get_empty_store():\n",
    "        return {\"down_cross\": [], \"mid_cross\": [], \"up_cross\": [],\n",
    "                \"down_self\": [],  \"mid_self\": [],  \"up_self\": []}\n",
    "\n",
    "    def forward(self, attn, is_cross: bool, place_in_unet: str):\n",
    "        key = f\"{place_in_unet}_{'cross' if is_cross else 'self'}\"\n",
    "        if attn.shape[1] <= 32 ** 2:  # avoid memory overhead\n",
    "            self.step_store[key].append(attn)\n",
    "        return attn\n",
    "\n",
    "    def between_steps(self):\n",
    "        if len(self.attention_store) == 0:\n",
    "            self.attention_store = self.step_store\n",
    "        else:\n",
    "            for key in self.attention_store:\n",
    "                for i in range(len(self.attention_store[key])):\n",
    "                    self.attention_store[key][i] += self.step_store[key][i]\n",
    "        self.step_store = self.get_empty_store()\n",
    "\n",
    "    def get_average_attention(self):\n",
    "        average_attention = {key: [item / self.cur_step for item in self.attention_store[key]] for key in self.attention_store}\n",
    "        return average_attention\n",
    "\n",
    "\n",
    "    def reset(self):\n",
    "        super(AttentionStore, self).reset()\n",
    "        self.step_store = self.get_empty_store()\n",
    "        self.attention_store = {}\n",
    "\n",
    "    def __init__(self):\n",
    "        super(AttentionStore, self).__init__()\n",
    "        self.step_store = self.get_empty_store()\n",
    "        self.attention_store = {}\n",
    "\n",
    "        \n",
    "class AttentionControlEdit(AttentionStore, abc.ABC):\n",
    "    \n",
    "    def step_callback(self, x_t):\n",
    "        if self.local_blend is not None:\n",
    "            x_t = self.local_blend(x_t, self.attention_store)\n",
    "        return x_t\n",
    "        \n",
    "    def replace_self_attention(self, attn_base, att_replace, place_in_unet):\n",
    "        if att_replace.shape[2] <= 32 ** 2:\n",
    "            attn_base = attn_base.unsqueeze(0).expand(att_replace.shape[0], *attn_base.shape)\n",
    "            return attn_base\n",
    "        else:\n",
    "            return att_replace\n",
    "    \n",
    "    @abc.abstractmethod\n",
    "    def replace_cross_attention(self, attn_base, att_replace):\n",
    "        raise NotImplementedError\n",
    "    \n",
    "    def forward(self, attn, is_cross: bool, place_in_unet: str):\n",
    "        super(AttentionControlEdit, self).forward(attn, is_cross, place_in_unet)\n",
    "        if is_cross or (self.num_self_replace[0] <= self.cur_step < self.num_self_replace[1]):\n",
    "            h = attn.shape[0] // (self.batch_size)\n",
    "            attn = attn.reshape(self.batch_size, h, *attn.shape[1:])\n",
    "            attn_base, attn_repalce = attn[0], attn[1:]\n",
    "            if is_cross:\n",
    "                alpha_words = self.cross_replace_alpha[self.cur_step]\n",
    "                attn_repalce_new = self.replace_cross_attention(attn_base, attn_repalce) * alpha_words + (1 - alpha_words) * attn_repalce\n",
    "                attn[1:] = attn_repalce_new\n",
    "            else:\n",
    "                attn[1:] = self.replace_self_attention(attn_base, attn_repalce, place_in_unet)\n",
    "            attn = attn.reshape(self.batch_size * h, *attn.shape[2:])\n",
    "        return attn\n",
    "    \n",
    "    def __init__(self, prompts, num_steps: int,\n",
    "                 cross_replace_steps: Union[float, Tuple[float, float], Dict[str, Tuple[float, float]]],\n",
    "                 self_replace_steps: Union[float, Tuple[float, float]],\n",
    "                 local_blend: Optional[LocalBlend]):\n",
    "        super(AttentionControlEdit, self).__init__()\n",
    "        self.batch_size = len(prompts)\n",
    "        self.cross_replace_alpha = ptp_utils.get_time_words_attention_alpha(prompts, num_steps, cross_replace_steps, tokenizer).to(device)\n",
    "        if type(self_replace_steps) is float:\n",
    "            self_replace_steps = 0, self_replace_steps\n",
    "        self.num_self_replace = int(num_steps * self_replace_steps[0]), int(num_steps * self_replace_steps[1])\n",
    "        self.local_blend = local_blend\n",
    "\n",
    "class AttentionReplace(AttentionControlEdit):\n",
    "\n",
    "    def replace_cross_attention(self, attn_base, att_replace):\n",
    "        return torch.einsum('hpw,bwn->bhpn', attn_base, self.mapper)\n",
    "      \n",
    "    def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float,\n",
    "                 local_blend: Optional[LocalBlend] = None):\n",
    "        super(AttentionReplace, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)\n",
    "        self.mapper = seq_aligner.get_replacement_mapper(prompts, tokenizer).to(device)\n",
    "        \n",
    "\n",
    "class AttentionRefine(AttentionControlEdit):\n",
    "\n",
    "    def replace_cross_attention(self, attn_base, att_replace):\n",
    "        attn_base_replace = attn_base[:, :, self.mapper].permute(2, 0, 1, 3)\n",
    "        attn_replace = attn_base_replace * self.alphas + att_replace * (1 - self.alphas)\n",
    "        # attn_replace = attn_replace / attn_replace.sum(-1, keepdims=True)\n",
    "        return attn_replace\n",
    "\n",
    "    def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float,\n",
    "                 local_blend: Optional[LocalBlend] = None):\n",
    "        super(AttentionRefine, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)\n",
    "        self.mapper, alphas = seq_aligner.get_refinement_mapper(prompts, tokenizer)\n",
    "        self.mapper, alphas = self.mapper.to(device), alphas.to(device)\n",
    "        self.alphas = alphas.reshape(alphas.shape[0], 1, 1, alphas.shape[1])\n",
    "\n",
    "\n",
    "class AttentionReweight(AttentionControlEdit):\n",
    "\n",
    "    def replace_cross_attention(self, attn_base, att_replace):\n",
    "        if self.prev_controller is not None:\n",
    "            attn_base = self.prev_controller.replace_cross_attention(attn_base, att_replace)\n",
    "        attn_replace = attn_base[None, :, :, :] * self.equalizer[:, None, None, :]\n",
    "        # attn_replace = attn_replace / attn_replace.sum(-1, keepdims=True)\n",
    "        return attn_replace\n",
    "\n",
    "    def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float, equalizer,\n",
    "                local_blend: Optional[LocalBlend] = None, controller: Optional[AttentionControlEdit] = None):\n",
    "        super(AttentionReweight, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)\n",
    "        self.equalizer = equalizer.to(device)\n",
    "        self.prev_controller = controller\n",
    "\n",
    "\n",
    "def get_equalizer(text: str, word_select: Union[int, Tuple[int, ...]], values: Union[List[float],\n",
    "                  Tuple[float, ...]]):\n",
    "    if type(word_select) is int or type(word_select) is str:\n",
    "        word_select = (word_select,)\n",
    "    equalizer = torch.ones(1, 77)\n",
    "    \n",
    "    for word, val in zip(word_select, values):\n",
    "        inds = ptp_utils.get_word_inds(text, word, tokenizer)\n",
    "        equalizer[:, inds] = val\n",
    "    return equalizer\n",
    "\n",
    "def aggregate_attention(attention_store: AttentionStore, res: int, from_where: List[str], is_cross: bool, select: int):\n",
    "    out = []\n",
    "    attention_maps = attention_store.get_average_attention()\n",
    "    num_pixels = res ** 2\n",
    "    for location in from_where:\n",
    "        for item in attention_maps[f\"{location}_{'cross' if is_cross else 'self'}\"]:\n",
    "            if item.shape[1] == num_pixels:\n",
    "                cross_maps = item.reshape(len(prompts), -1, res, res, item.shape[-1])[select]\n",
    "                out.append(cross_maps)\n",
    "    out = torch.cat(out, dim=0)\n",
    "    out = out.sum(0) / out.shape[0]\n",
    "    return out.cpu()\n",
    "\n",
    "\n",
    "def make_controller(prompts: List[str], is_replace_controller: bool, cross_replace_steps: Dict[str, float], self_replace_steps: float, blend_words=None, equilizer_params=None) -> AttentionControlEdit:\n",
    "    if blend_words is None:\n",
    "        lb = None\n",
    "    else:\n",
    "        lb = LocalBlend(prompts, blend_word)\n",
    "    if is_replace_controller:\n",
    "        controller = AttentionReplace(prompts, NUM_DDIM_STEPS, cross_replace_steps=cross_replace_steps, self_replace_steps=self_replace_steps, local_blend=lb)\n",
    "    else:\n",
    "        controller = AttentionRefine(prompts, NUM_DDIM_STEPS, cross_replace_steps=cross_replace_steps, self_replace_steps=self_replace_steps, local_blend=lb)\n",
    "    if equilizer_params is not None:\n",
    "        eq = get_equalizer(prompts[1], equilizer_params[\"words\"], equilizer_params[\"values\"])\n",
    "        controller = AttentionReweight(prompts, NUM_DDIM_STEPS, cross_replace_steps=cross_replace_steps,\n",
    "                                       self_replace_steps=self_replace_steps, equalizer=eq, local_blend=lb, controller=controller)\n",
    "    return controller\n",
    "\n",
    "\n",
    "def show_cross_attention(attention_store: AttentionStore, res: int, from_where: List[str], select: int = 0):\n",
    "    tokens = tokenizer.encode(prompts[select])\n",
    "    decoder = tokenizer.decode\n",
    "    attention_maps = aggregate_attention(attention_store, res, from_where, True, select)\n",
    "    images = []\n",
    "    for i in range(len(tokens)):\n",
    "        image = attention_maps[:, :, i]\n",
    "        image = 255 * image / image.max()\n",
    "        image = image.unsqueeze(-1).expand(*image.shape, 3)\n",
    "        image = image.numpy().astype(np.uint8)\n",
    "        image = np.array(Image.fromarray(image).resize((256, 256)))\n",
    "        image = ptp_utils.text_under_image(image, decoder(int(tokens[i])))\n",
    "        images.append(image)\n",
    "    ptp_utils.view_images(np.stack(images, axis=0))\n",
    "    \n",
    "\n",
    "def show_self_attention_comp(attention_store: AttentionStore, res: int, from_where: List[str],\n",
    "                        max_com=10, select: int = 0):\n",
    "    attention_maps = aggregate_attention(attention_store, res, from_where, False, select).numpy().reshape((res ** 2, res ** 2))\n",
    "    u, s, vh = np.linalg.svd(attention_maps - np.mean(attention_maps, axis=1, keepdims=True))\n",
    "    images = []\n",
    "    for i in range(max_com):\n",
    "        image = vh[i].reshape(res, res)\n",
    "        image = image - image.min()\n",
    "        image = 255 * image / image.max()\n",
    "        image = np.repeat(np.expand_dims(image, axis=2), 3, axis=2).astype(np.uint8)\n",
    "        image = Image.fromarray(image).resize((256, 256))\n",
    "        image = np.array(image)\n",
    "        images.append(image)\n",
    "    ptp_utils.view_images(np.concatenate(images, axis=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a914bda0-c191-4db6-b891-101cde74ddaf",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Null Text Inversion code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c442992d-8156-4dfc-a2a5-1fbf8bedb4b2",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def load_512(image_path, left=0, right=0, top=0, bottom=0):\n",
    "    if type(image_path) is str:\n",
    "        image = np.array(Image.open(image_path))[:, :, :3]\n",
    "    else:\n",
    "        image = image_path\n",
    "    h, w, c = image.shape\n",
    "    left = min(left, w-1)\n",
    "    right = min(right, w - left - 1)\n",
    "    top = min(top, h - left - 1)\n",
    "    bottom = min(bottom, h - top - 1)\n",
    "    image = image[top:h-bottom, left:w-right]\n",
    "    h, w, c = image.shape\n",
    "    if h < w:\n",
    "        offset = (w - h) // 2\n",
    "        image = image[:, offset:offset + h]\n",
    "    elif w < h:\n",
    "        offset = (h - w) // 2\n",
    "        image = image[offset:offset + w]\n",
    "    image = np.array(Image.fromarray(image).resize((512, 512)))\n",
    "    return image\n",
    "\n",
    "\n",
    "class NullInversion:\n",
    "    \n",
    "    def prev_step(self, model_output: Union[torch.FloatTensor, np.ndarray], timestep: int, sample: Union[torch.FloatTensor, np.ndarray]):\n",
    "        prev_timestep = timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps\n",
    "        alpha_prod_t = self.scheduler.alphas_cumprod[timestep]\n",
    "        alpha_prod_t_prev = self.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.scheduler.final_alpha_cumprod\n",
    "        beta_prod_t = 1 - alpha_prod_t\n",
    "        pred_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5\n",
    "        pred_sample_direction = (1 - alpha_prod_t_prev) ** 0.5 * model_output\n",
    "        prev_sample = alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction\n",
    "        return prev_sample\n",
    "    \n",
    "    def next_step(self, model_output: Union[torch.FloatTensor, np.ndarray], timestep: int, sample: Union[torch.FloatTensor, np.ndarray]):\n",
    "        timestep, next_timestep = min(timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps, 999), timestep\n",
    "        alpha_prod_t = self.scheduler.alphas_cumprod[timestep] if timestep >= 0 else self.scheduler.final_alpha_cumprod\n",
    "        alpha_prod_t_next = self.scheduler.alphas_cumprod[next_timestep]\n",
    "        beta_prod_t = 1 - alpha_prod_t\n",
    "        next_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5\n",
    "        next_sample_direction = (1 - alpha_prod_t_next) ** 0.5 * model_output\n",
    "        next_sample = alpha_prod_t_next ** 0.5 * next_original_sample + next_sample_direction\n",
    "        return next_sample\n",
    "    \n",
    "    def get_noise_pred_single(self, latents, t, context):\n",
    "        noise_pred = self.model.unet(latents, t, encoder_hidden_states=context)[\"sample\"]\n",
    "        return noise_pred\n",
    "\n",
    "    def get_noise_pred(self, latents, t, is_forward=True, context=None):\n",
    "        latents_input = torch.cat([latents] * 2)\n",
    "        if context is None:\n",
    "            context = self.context\n",
    "        guidance_scale = 1 if is_forward else GUIDANCE_SCALE\n",
    "        noise_pred = self.model.unet(latents_input, t, encoder_hidden_states=context)[\"sample\"]\n",
    "        noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2)\n",
    "        noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond)\n",
    "        if is_forward:\n",
    "            latents = self.next_step(noise_pred, t, latents)\n",
    "        else:\n",
    "            latents = self.prev_step(noise_pred, t, latents)\n",
    "        return latents\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def latent2image(self, latents, return_type='np'):\n",
    "        latents = 1 / 0.18215 * latents.detach()\n",
    "        image = self.model.vae.decode(latents)['sample']\n",
    "        if return_type == 'np':\n",
    "            image = (image / 2 + 0.5).clamp(0, 1)\n",
    "            image = image.cpu().permute(0, 2, 3, 1).numpy()[0]\n",
    "            image = (image * 255).astype(np.uint8)\n",
    "        return image\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def image2latent(self, image):\n",
    "        with torch.no_grad():\n",
    "            if type(image) is Image:\n",
    "                image = np.array(image)\n",
    "            if type(image) is torch.Tensor and image.dim() == 4:\n",
    "                latents = image\n",
    "            else:\n",
    "                image = torch.from_numpy(image).float() / 127.5 - 1\n",
    "                image = image.permute(2, 0, 1).unsqueeze(0).to(device)\n",
    "                latents = self.model.vae.encode(image)['latent_dist'].mean\n",
    "                latents = latents * 0.18215\n",
    "        return latents\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def init_prompt(self, prompt: str):\n",
    "        uncond_input = self.model.tokenizer(\n",
    "            [\"\"], padding=\"max_length\", max_length=self.model.tokenizer.model_max_length,\n",
    "            return_tensors=\"pt\"\n",
    "        )\n",
    "        uncond_embeddings = self.model.text_encoder(uncond_input.input_ids.to(self.model.device))[0]\n",
    "        text_input = self.model.tokenizer(\n",
    "            [prompt],\n",
    "            padding=\"max_length\",\n",
    "            max_length=self.model.tokenizer.model_max_length,\n",
    "            truncation=True,\n",
    "            return_tensors=\"pt\",\n",
    "        )\n",
    "        text_embeddings = self.model.text_encoder(text_input.input_ids.to(self.model.device))[0]\n",
    "        self.context = torch.cat([uncond_embeddings, text_embeddings])\n",
    "        self.prompt = prompt\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def ddim_loop(self, latent):\n",
    "        uncond_embeddings, cond_embeddings = self.context.chunk(2)\n",
    "        all_latent = [latent]\n",
    "        latent = latent.clone().detach()\n",
    "        for i in range(NUM_DDIM_STEPS):\n",
    "            t = self.model.scheduler.timesteps[len(self.model.scheduler.timesteps) - i - 1]\n",
    "            noise_pred = self.get_noise_pred_single(latent, t, cond_embeddings)\n",
    "            latent = self.next_step(noise_pred, t, latent)\n",
    "            all_latent.append(latent)\n",
    "        return all_latent\n",
    "\n",
    "    @property\n",
    "    def scheduler(self):\n",
    "        return self.model.scheduler\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def ddim_inversion(self, image):\n",
    "        latent = self.image2latent(image)\n",
    "        image_rec = self.latent2image(latent)\n",
    "        ddim_latents = self.ddim_loop(latent)\n",
    "        return image_rec, ddim_latents\n",
    "\n",
    "    def null_optimization(self, latents, num_inner_steps, epsilon):\n",
    "        uncond_embeddings, cond_embeddings = self.context.chunk(2)\n",
    "        uncond_embeddings_list = []\n",
    "        latent_cur = latents[-1]\n",
    "        bar = tqdm(total=num_inner_steps * NUM_DDIM_STEPS)\n",
    "        for i in range(NUM_DDIM_STEPS):\n",
    "            uncond_embeddings = uncond_embeddings.clone().detach()\n",
    "            uncond_embeddings.requires_grad = True\n",
    "            optimizer = Adam([uncond_embeddings], lr=1e-2 * (1. - i / 100.))\n",
    "            latent_prev = latents[len(latents) - i - 2]\n",
    "            t = self.model.scheduler.timesteps[i]\n",
    "            with torch.no_grad():\n",
    "                noise_pred_cond = self.get_noise_pred_single(latent_cur, t, cond_embeddings)\n",
    "            for j in range(num_inner_steps):\n",
    "                noise_pred_uncond = self.get_noise_pred_single(latent_cur, t, uncond_embeddings)\n",
    "                noise_pred = noise_pred_uncond + GUIDANCE_SCALE * (noise_pred_cond - noise_pred_uncond)\n",
    "                latents_prev_rec = self.prev_step(noise_pred, t, latent_cur)\n",
    "                loss = nnf.mse_loss(latents_prev_rec, latent_prev)\n",
    "                optimizer.zero_grad()\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                loss_item = loss.item()\n",
    "                bar.update()\n",
    "                if loss_item < epsilon + i * 2e-5:\n",
    "                    break\n",
    "            for j in range(j + 1, num_inner_steps):\n",
    "                bar.update()\n",
    "            uncond_embeddings_list.append(uncond_embeddings[:1].detach())\n",
    "            with torch.no_grad():\n",
    "                context = torch.cat([uncond_embeddings, cond_embeddings])\n",
    "                latent_cur = self.get_noise_pred(latent_cur, t, False, context)\n",
    "        bar.close()\n",
    "        return uncond_embeddings_list\n",
    "    \n",
    "    def invert(self, image_path: str, prompt: str, offsets=(0,0,0,0), num_inner_steps=10, early_stop_epsilon=1e-5, verbose=False):\n",
    "        self.init_prompt(prompt)\n",
    "        ptp_utils.register_attention_control(self.model, None)\n",
    "        image_gt = load_512(image_path, *offsets)\n",
    "        if verbose:\n",
    "            print(\"DDIM inversion...\")\n",
    "        image_rec, ddim_latents = self.ddim_inversion(image_gt)\n",
    "        if verbose:\n",
    "            print(\"Null-text optimization...\")\n",
    "        uncond_embeddings = self.null_optimization(ddim_latents, num_inner_steps, early_stop_epsilon)\n",
    "        return (image_gt, image_rec), ddim_latents[-1], uncond_embeddings\n",
    "        \n",
    "    \n",
    "    def __init__(self, model):\n",
    "        scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule=\"scaled_linear\", clip_sample=False,\n",
    "                                  set_alpha_to_one=False)\n",
    "        self.model = model\n",
    "        self.tokenizer = self.model.tokenizer\n",
    "        self.model.scheduler.set_timesteps(NUM_DDIM_STEPS)\n",
    "        self.prompt = None\n",
    "        self.context = None\n",
    "\n",
    "null_inversion = NullInversion(ldm_stable)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c919e093-998c-4e4c-92a2-dc9517ef8ea4",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Infernce Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f9499145-1a2b-4c91-900e-093c0c08043c",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def text2image_ldm_stable(\n",
    "    model,\n",
    "    prompt:  List[str],\n",
    "    controller,\n",
    "    num_inference_steps: int = 50,\n",
    "    guidance_scale: Optional[float] = 7.5,\n",
    "    generator: Optional[torch.Generator] = None,\n",
    "    latent: Optional[torch.FloatTensor] = None,\n",
    "    uncond_embeddings=None,\n",
    "    start_time=50,\n",
    "    return_type='image'\n",
    "):\n",
    "    batch_size = len(prompt)\n",
    "    ptp_utils.register_attention_control(model, controller)\n",
    "    height = width = 512\n",
    "    \n",
    "    text_input = model.tokenizer(\n",
    "        prompt,\n",
    "        padding=\"max_length\",\n",
    "        max_length=model.tokenizer.model_max_length,\n",
    "        truncation=True,\n",
    "        return_tensors=\"pt\",\n",
    "    )\n",
    "    text_embeddings = model.text_encoder(text_input.input_ids.to(model.device))[0]\n",
    "    max_length = text_input.input_ids.shape[-1]\n",
    "    if uncond_embeddings is None:\n",
    "        uncond_input = model.tokenizer(\n",
    "            [\"\"] * batch_size, padding=\"max_length\", max_length=max_length, return_tensors=\"pt\"\n",
    "        )\n",
    "        uncond_embeddings_ = model.text_encoder(uncond_input.input_ids.to(model.device))[0]\n",
    "    else:\n",
    "        uncond_embeddings_ = None\n",
    "\n",
    "    latent, latents = ptp_utils.init_latent(latent, model, height, width, generator, batch_size)\n",
    "    model.scheduler.set_timesteps(num_inference_steps)\n",
    "    for i, t in enumerate(tqdm(model.scheduler.timesteps[-start_time:])):\n",
    "        if uncond_embeddings_ is None:\n",
    "            context = torch.cat([uncond_embeddings[i].expand(*text_embeddings.shape), text_embeddings])\n",
    "        else:\n",
    "            context = torch.cat([uncond_embeddings_, text_embeddings])\n",
    "        latents = ptp_utils.diffusion_step(model, controller, latents, context, t, guidance_scale, low_resource=False)\n",
    "        \n",
    "    if return_type == 'image':\n",
    "        image = ptp_utils.latent2image(model.vae, latents)\n",
    "    else:\n",
    "        image = latents\n",
    "    return image, latent\n",
    "\n",
    "\n",
    "\n",
    "# def run_and_display(prompts, controller, latent=None, run_baseline=False, generator=None, uncond_embeddings=None, verbose=True):\n",
    "#     if run_baseline:\n",
    "#         print(\"w.o. prompt-to-prompt\")\n",
    "#         images, latent = run_and_display(prompts, EmptyControl(), latent=latent, run_baseline=False, generator=generator)\n",
    "#         print(\"with prompt-to-prompt\")\n",
    "#     images, x_t = text2image_ldm_stable(ldm_stable, prompts, controller, latent=latent, num_inference_steps=NUM_DDIM_STEPS, guidance_scale=GUIDANCE_SCALE, generator=generator, uncond_embeddings=uncond_embeddings)\n",
    "#     if verbose:\n",
    "#         ptp_utils.view_images(images)\n",
    "#     return images, x_t\n",
    "def run_and_display(prompts, latent=None, run_baseline=False, generator=None, uncond_embeddings=None, verbose=True):\n",
    "    images, latent = run_and_display(prompts, EmptyControl(), latent=latent, run_baseline=False, generator=generator)\n",
    "    if verbose:\n",
    "        ptp_utils.view_images(images)\n",
    "    return images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6ef5cbc6-2581-415f-b608-67f2e87c32f5",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 718792 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=512x512 at 0x7FB68BDAFD00>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "image_path = \"../data/dragon_statue_1/image.png\"\n",
    "prompt = \"A high-resolution DSLR image of a grey dragon statue\"\n",
    "offsets = (0, 0, 0, 0)\n",
    "img = load_512(image_path, *offsets)\n",
    "ptp_utils.view_images(img)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "27fd08bb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DDIM inversion...\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800000; text-decoration-color: #800000\">╭─────────────────────────────── </span><span style=\"color: #800000; text-decoration-color: #800000; font-weight: bold\">Traceback </span><span style=\"color: #bf7f7f; text-decoration-color: #bf7f7f; font-weight: bold\">(most recent call last)</span><span style=\"color: #800000; text-decoration-color: #800000\"> ────────────────────────────────╮</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">&lt;module&gt;</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">1</span>                                                                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>1 (image_gt, image_enc), x_t, uncond_embeddings = null_inversion.invert(image_path, prompt     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">2 </span>                                                                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">3 </span><span style=\"color: #00ffff; text-decoration-color: #00ffff\">print</span>(<span style=\"color: #808000; text-decoration-color: #808000\">\"Modify or remove offsets according to your image!\"</span>)                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">4 </span>                                                                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">invert</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">168</span>                                                                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">165 │   │   </span>image_gt = load_512(image_path, *offsets)                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">166 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> verbose:                                                                        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">167 │   │   │   </span><span style=\"color: #00ffff; text-decoration-color: #00ffff\">print</span>(<span style=\"color: #808000; text-decoration-color: #808000\">\"DDIM inversion...\"</span>)                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>168 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   </span>image_rec, ddim_latents = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.ddim_inversion(image_gt)                            <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">169 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> verbose:                                                                        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">170 │   │   │   </span><span style=\"color: #00ffff; text-decoration-color: #00ffff\">print</span>(<span style=\"color: #808000; text-decoration-color: #808000\">\"Null-text optimization...\"</span>)                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">171 │   │   </span>uncond_embeddings = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.null_optimization(ddim_latents, num_inner_steps, early_   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/torch/autograd/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">grad_mode.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">27</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">decorate_context</span>                                                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 24 │   │   </span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff; font-weight: bold\">@functools</span>.wraps(func)                                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 25 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">def</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">decorate_context</span>(*args, **kwargs):                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 26 │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">with</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.clone():                                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span> 27 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> func(*args, **kwargs)                                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 28 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> cast(F, decorate_context)                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 29 │   </span>                                                                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 30 │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">def</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_wrap_generator</span>(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>, func):                                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">ddim_inversion</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">125</span>                                                                            <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">122 │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">def</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">ddim_inversion</span>(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>, image):                                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">123 │   │   </span>latent = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.image2latent(image)                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">124 │   │   </span>image_rec = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.latent2image(latent)                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>125 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   </span>ddim_latents = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.ddim_loop(latent)                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">126 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> image_rec, ddim_latents                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">127 │   </span>                                                                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">128 │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">def</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">null_optimization</span>(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>, latents, num_inner_steps, epsilon):                        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/torch/autograd/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">grad_mode.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">27</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">decorate_context</span>                                                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 24 │   │   </span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff; font-weight: bold\">@functools</span>.wraps(func)                                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 25 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">def</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">decorate_context</span>(*args, **kwargs):                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 26 │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">with</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.clone():                                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span> 27 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> func(*args, **kwargs)                                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 28 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> cast(F, decorate_context)                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 29 │   </span>                                                                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 30 │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">def</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_wrap_generator</span>(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>, func):                                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">ddim_loop</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">112</span>                                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">109 │   │   </span>latent = latent.clone().detach()                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">110 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">for</span> i <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">in</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">range</span>(NUM_DDIM_STEPS):                                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">111 │   │   │   </span>t = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.model.scheduler.timesteps[<span style=\"color: #00ffff; text-decoration-color: #00ffff\">len</span>(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.model.scheduler.timesteps) - i -   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>112 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span>noise_pred = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.get_noise_pred_single(latent, t, cond_embeddings)            <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">113 │   │   │   </span>latent = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.next_step(noise_pred, t, latent)                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">114 │   │   │   </span>all_latent.append(latent)                                                      <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">115 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> all_latent                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">get_noise_pred_single</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">46</span>                                                                      <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 43 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> next_sample                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 44 │   </span>                                                                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 45 │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">def</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">get_noise_pred_single</span>(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>, latents, t, context):                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span> 46 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   </span>noise_pred = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.model.unet(latents, t, encoder_hidden_states=context)[<span style=\"color: #808000; text-decoration-color: #808000\">\"sample\"</span>   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 47 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> noise_pred                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 48 │   </span>                                                                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 49 │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">def</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">get_noise_pred</span>(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>, latents, t, is_forward=<span style=\"color: #0000ff; text-decoration-color: #0000ff\">True</span>, context=<span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>):                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/torch/nn/modules/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">module.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">111</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_call_impl</span>                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1107 │   │   # this function, and just call forward.</span>                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1108 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span> (<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._forward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._forward_pre_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">o</span>  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1109 │   │   │   │   </span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_pre_hooks):                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>1110 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> forward_call(*<span style=\"color: #00ffff; text-decoration-color: #00ffff\">input</span>, **kwargs)                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1111 │   │   # Do not call functions when jit is used</span>                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1112 │   │   </span>full_backward_hooks, non_full_backward_hooks = [], []                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1113 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_backward_hooks:                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/diffusers/models/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">unet_2d_condi</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">tion.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">582</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">forward</span>                                                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">579 │   │   </span>down_block_res_samples = (sample,)                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">580 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">for</span> downsample_block <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">in</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.down_blocks:                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">581 │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">hasattr</span>(downsample_block, <span style=\"color: #808000; text-decoration-color: #808000\">\"has_cross_attention\"</span>) <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">and</span> downsample_block.has   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>582 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   │   </span>sample, res_samples = downsample_block(                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">583 │   │   │   │   │   </span>hidden_states=sample,                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">584 │   │   │   │   │   </span>temb=emb,                                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">585 │   │   │   │   │   </span>encoder_hidden_states=encoder_hidden_states,                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/torch/nn/modules/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">module.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">111</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_call_impl</span>                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1107 │   │   # this function, and just call forward.</span>                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1108 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span> (<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._forward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._forward_pre_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">o</span>  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1109 │   │   │   │   </span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_pre_hooks):                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>1110 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> forward_call(*<span style=\"color: #00ffff; text-decoration-color: #00ffff\">input</span>, **kwargs)                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1111 │   │   # Do not call functions when jit is used</span>                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1112 │   │   </span>full_backward_hooks, non_full_backward_hooks = [], []                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1113 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_backward_hooks:                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/diffusers/models/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">unet_2d_block</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">s.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">837</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">forward</span>                                                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 834 │   │   │   │   </span>)[<span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span>]                                                                      <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 835 │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">else</span>:                                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 836 │   │   │   │   </span>hidden_states = resnet(hidden_states, temb)                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span> 837 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   │   </span>hidden_states = attn(                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 838 │   │   │   │   │   </span>hidden_states,                                                        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 839 │   │   │   │   │   </span>encoder_hidden_states=encoder_hidden_states,                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 840 │   │   │   │   │   </span>cross_attention_kwargs=cross_attention_kwargs,                        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/torch/nn/modules/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">module.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">111</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_call_impl</span>                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1107 │   │   # this function, and just call forward.</span>                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1108 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span> (<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._forward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._forward_pre_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">o</span>  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1109 │   │   │   │   </span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_pre_hooks):                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>1110 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> forward_call(*<span style=\"color: #00ffff; text-decoration-color: #00ffff\">input</span>, **kwargs)                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1111 │   │   # Do not call functions when jit is used</span>                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1112 │   │   </span>full_backward_hooks, non_full_backward_hooks = [], []                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1113 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_backward_hooks:                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/diffusers/models/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">transformer_2</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">d.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">265</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">forward</span>                                                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">262 │   │   </span>                                                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">263 │   │   # 2. Blocks</span>                                                                        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">264 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">for</span> block <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">in</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.transformer_blocks:                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>265 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span>hidden_states = block(                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">266 │   │   │   │   </span>hidden_states,                                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">267 │   │   │   │   </span>encoder_hidden_states=encoder_hidden_states,                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">268 │   │   │   │   </span>timestep=timestep,                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/torch/nn/modules/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">module.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">111</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_call_impl</span>                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1107 │   │   # this function, and just call forward.</span>                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1108 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span> (<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._forward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._forward_pre_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">o</span>  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1109 │   │   │   │   </span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_pre_hooks):                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>1110 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> forward_call(*<span style=\"color: #00ffff; text-decoration-color: #00ffff\">input</span>, **kwargs)                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1111 │   │   # Do not call functions when jit is used</span>                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1112 │   │   </span>full_backward_hooks, non_full_backward_hooks = [], []                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1113 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_backward_hooks:                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/diffusers/models/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">attention.py</span>: <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">291</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">forward</span>                                                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">288 │   │   </span>                                                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">289 │   │   # 1. Self-Attention</span>                                                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">290 │   │   </span>cross_attention_kwargs = cross_attention_kwargs <span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> cross_attention_kwargs <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">is</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span>   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>291 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   </span>attn_output = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.attn1(                                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">292 │   │   │   </span>norm_hidden_states,                                                            <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">293 │   │   │   </span>encoder_hidden_states=encoder_hidden_states <span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.only_cross_attention <span style=\"color: #0000ff; text-decoration-color: #0000ff\">els</span>   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">294 │   │   │   </span>attention_mask=attention_mask,                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/torch/nn/modules/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">module.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">111</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_call_impl</span>                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1107 │   │   # this function, and just call forward.</span>                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1108 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span> (<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._forward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._forward_pre_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">o</span>  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1109 │   │   │   │   </span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_pre_hooks):                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>1110 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> forward_call(*<span style=\"color: #00ffff; text-decoration-color: #00ffff\">input</span>, **kwargs)                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1111 │   │   # Do not call functions when jit is used</span>                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1112 │   │   </span>full_backward_hooks, non_full_backward_hooks = [], []                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1113 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_backward_hooks:                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">╰──────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
       "<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-weight: bold\">TypeError: </span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">forward</span><span style=\"font-weight: bold\">()</span> got an unexpected keyword argument <span style=\"color: #008000; text-decoration-color: #008000\">'encoder_hidden_states'</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[31m╭─\u001b[0m\u001b[31m──────────────────────────────\u001b[0m\u001b[31m \u001b[0m\u001b[1;31mTraceback \u001b[0m\u001b[1;2;31m(most recent call last)\u001b[0m\u001b[31m \u001b[0m\u001b[31m───────────────────────────────\u001b[0m\u001b[31m─╮\u001b[0m\n",
       "\u001b[31m│\u001b[0m in \u001b[92m<module>\u001b[0m:\u001b[94m1\u001b[0m                                                                                    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1 (image_gt, image_enc), x_t, uncond_embeddings = null_inversion.invert(image_path, prompt     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m2 \u001b[0m                                                                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m3 \u001b[0m\u001b[96mprint\u001b[0m(\u001b[33m\"\u001b[0m\u001b[33mModify or remove offsets according to your image!\u001b[0m\u001b[33m\"\u001b[0m)                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m4 \u001b[0m                                                                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m in \u001b[92minvert\u001b[0m:\u001b[94m168\u001b[0m                                                                                    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m165 \u001b[0m\u001b[2m│   │   \u001b[0mimage_gt = load_512(image_path, *offsets)                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m166 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m verbose:                                                                        \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m167 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[96mprint\u001b[0m(\u001b[33m\"\u001b[0m\u001b[33mDDIM inversion...\u001b[0m\u001b[33m\"\u001b[0m)                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m168 \u001b[2m│   │   \u001b[0mimage_rec, ddim_latents = \u001b[96mself\u001b[0m.ddim_inversion(image_gt)                            \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m169 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m verbose:                                                                        \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m170 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[96mprint\u001b[0m(\u001b[33m\"\u001b[0m\u001b[33mNull-text optimization...\u001b[0m\u001b[33m\"\u001b[0m)                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m171 \u001b[0m\u001b[2m│   │   \u001b[0muncond_embeddings = \u001b[96mself\u001b[0m.null_optimization(ddim_latents, num_inner_steps, early_   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/torch/autograd/\u001b[0m\u001b[1;33mgrad_mode.py\u001b[0m:\u001b[94m27\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m in \u001b[92mdecorate_context\u001b[0m                                                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 24 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[1;95m@functools\u001b[0m.wraps(func)                                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 25 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mdef\u001b[0m \u001b[92mdecorate_context\u001b[0m(*args, **kwargs):                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 26 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[94mwith\u001b[0m \u001b[96mself\u001b[0m.clone():                                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m 27 \u001b[2m│   │   │   │   \u001b[0m\u001b[94mreturn\u001b[0m func(*args, **kwargs)                                               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 28 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mreturn\u001b[0m cast(F, decorate_context)                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 29 \u001b[0m\u001b[2m│   \u001b[0m                                                                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 30 \u001b[0m\u001b[2m│   \u001b[0m\u001b[94mdef\u001b[0m \u001b[92m_wrap_generator\u001b[0m(\u001b[96mself\u001b[0m, func):                                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m in \u001b[92mddim_inversion\u001b[0m:\u001b[94m125\u001b[0m                                                                            \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m122 \u001b[0m\u001b[2m│   \u001b[0m\u001b[94mdef\u001b[0m \u001b[92mddim_inversion\u001b[0m(\u001b[96mself\u001b[0m, image):                                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m123 \u001b[0m\u001b[2m│   │   \u001b[0mlatent = \u001b[96mself\u001b[0m.image2latent(image)                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m124 \u001b[0m\u001b[2m│   │   \u001b[0mimage_rec = \u001b[96mself\u001b[0m.latent2image(latent)                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m125 \u001b[2m│   │   \u001b[0mddim_latents = \u001b[96mself\u001b[0m.ddim_loop(latent)                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m126 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mreturn\u001b[0m image_rec, ddim_latents                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m127 \u001b[0m\u001b[2m│   \u001b[0m                                                                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m128 \u001b[0m\u001b[2m│   \u001b[0m\u001b[94mdef\u001b[0m \u001b[92mnull_optimization\u001b[0m(\u001b[96mself\u001b[0m, latents, num_inner_steps, epsilon):                        \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/torch/autograd/\u001b[0m\u001b[1;33mgrad_mode.py\u001b[0m:\u001b[94m27\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m in \u001b[92mdecorate_context\u001b[0m                                                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 24 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[1;95m@functools\u001b[0m.wraps(func)                                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 25 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mdef\u001b[0m \u001b[92mdecorate_context\u001b[0m(*args, **kwargs):                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 26 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[94mwith\u001b[0m \u001b[96mself\u001b[0m.clone():                                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m 27 \u001b[2m│   │   │   │   \u001b[0m\u001b[94mreturn\u001b[0m func(*args, **kwargs)                                               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 28 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mreturn\u001b[0m cast(F, decorate_context)                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 29 \u001b[0m\u001b[2m│   \u001b[0m                                                                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 30 \u001b[0m\u001b[2m│   \u001b[0m\u001b[94mdef\u001b[0m \u001b[92m_wrap_generator\u001b[0m(\u001b[96mself\u001b[0m, func):                                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m in \u001b[92mddim_loop\u001b[0m:\u001b[94m112\u001b[0m                                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m109 \u001b[0m\u001b[2m│   │   \u001b[0mlatent = latent.clone().detach()                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m110 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mfor\u001b[0m i \u001b[95min\u001b[0m \u001b[96mrange\u001b[0m(NUM_DDIM_STEPS):                                                    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m111 \u001b[0m\u001b[2m│   │   │   \u001b[0mt = \u001b[96mself\u001b[0m.model.scheduler.timesteps[\u001b[96mlen\u001b[0m(\u001b[96mself\u001b[0m.model.scheduler.timesteps) - i -   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m112 \u001b[2m│   │   │   \u001b[0mnoise_pred = \u001b[96mself\u001b[0m.get_noise_pred_single(latent, t, cond_embeddings)            \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m113 \u001b[0m\u001b[2m│   │   │   \u001b[0mlatent = \u001b[96mself\u001b[0m.next_step(noise_pred, t, latent)                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m114 \u001b[0m\u001b[2m│   │   │   \u001b[0mall_latent.append(latent)                                                      \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m115 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mreturn\u001b[0m all_latent                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m in \u001b[92mget_noise_pred_single\u001b[0m:\u001b[94m46\u001b[0m                                                                      \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 43 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mreturn\u001b[0m next_sample                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 44 \u001b[0m\u001b[2m│   \u001b[0m                                                                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 45 \u001b[0m\u001b[2m│   \u001b[0m\u001b[94mdef\u001b[0m \u001b[92mget_noise_pred_single\u001b[0m(\u001b[96mself\u001b[0m, latents, t, context):                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m 46 \u001b[2m│   │   \u001b[0mnoise_pred = \u001b[96mself\u001b[0m.model.unet(latents, t, encoder_hidden_states=context)[\u001b[33m\"\u001b[0m\u001b[33msample\u001b[0m\u001b[33m\"\u001b[0m   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 47 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mreturn\u001b[0m noise_pred                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 48 \u001b[0m\u001b[2m│   \u001b[0m                                                                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 49 \u001b[0m\u001b[2m│   \u001b[0m\u001b[94mdef\u001b[0m \u001b[92mget_noise_pred\u001b[0m(\u001b[96mself\u001b[0m, latents, t, is_forward=\u001b[94mTrue\u001b[0m, context=\u001b[94mNone\u001b[0m):                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/torch/nn/modules/\u001b[0m\u001b[1;33mmodule.py\u001b[0m:\u001b[94m111\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[94m0\u001b[0m in \u001b[92m_call_impl\u001b[0m                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1107 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# this function, and just call forward.\u001b[0m                                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1108 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[95mnot\u001b[0m (\u001b[96mself\u001b[0m._backward_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._forward_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._forward_pre_hooks \u001b[95mo\u001b[0m  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1109 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[95mor\u001b[0m _global_forward_hooks \u001b[95mor\u001b[0m _global_forward_pre_hooks):                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1110 \u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m forward_call(*\u001b[96minput\u001b[0m, **kwargs)                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1111 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# Do not call functions when jit is used\u001b[0m                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1112 \u001b[0m\u001b[2m│   │   \u001b[0mfull_backward_hooks, non_full_backward_hooks = [], []                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1113 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[96mself\u001b[0m._backward_hooks \u001b[95mor\u001b[0m _global_backward_hooks:                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/diffusers/models/\u001b[0m\u001b[1;33munet_2d_condi\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[1;33mtion.py\u001b[0m:\u001b[94m582\u001b[0m in \u001b[92mforward\u001b[0m                                                                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m579 \u001b[0m\u001b[2m│   │   \u001b[0mdown_block_res_samples = (sample,)                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m580 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mfor\u001b[0m downsample_block \u001b[95min\u001b[0m \u001b[96mself\u001b[0m.down_blocks:                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m581 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[96mhasattr\u001b[0m(downsample_block, \u001b[33m\"\u001b[0m\u001b[33mhas_cross_attention\u001b[0m\u001b[33m\"\u001b[0m) \u001b[95mand\u001b[0m downsample_block.has   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m582 \u001b[2m│   │   │   │   \u001b[0msample, res_samples = downsample_block(                                    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m583 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0mhidden_states=sample,                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m584 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0mtemb=emb,                                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m585 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0mencoder_hidden_states=encoder_hidden_states,                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/torch/nn/modules/\u001b[0m\u001b[1;33mmodule.py\u001b[0m:\u001b[94m111\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[94m0\u001b[0m in \u001b[92m_call_impl\u001b[0m                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1107 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# this function, and just call forward.\u001b[0m                                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1108 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[95mnot\u001b[0m (\u001b[96mself\u001b[0m._backward_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._forward_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._forward_pre_hooks \u001b[95mo\u001b[0m  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1109 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[95mor\u001b[0m _global_forward_hooks \u001b[95mor\u001b[0m _global_forward_pre_hooks):                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1110 \u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m forward_call(*\u001b[96minput\u001b[0m, **kwargs)                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1111 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# Do not call functions when jit is used\u001b[0m                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1112 \u001b[0m\u001b[2m│   │   \u001b[0mfull_backward_hooks, non_full_backward_hooks = [], []                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1113 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[96mself\u001b[0m._backward_hooks \u001b[95mor\u001b[0m _global_backward_hooks:                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/diffusers/models/\u001b[0m\u001b[1;33munet_2d_block\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[1;33ms.py\u001b[0m:\u001b[94m837\u001b[0m in \u001b[92mforward\u001b[0m                                                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 834 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m)[\u001b[94m0\u001b[0m]                                                                      \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 835 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[94melse\u001b[0m:                                                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 836 \u001b[0m\u001b[2m│   │   │   │   \u001b[0mhidden_states = resnet(hidden_states, temb)                               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m 837 \u001b[2m│   │   │   │   \u001b[0mhidden_states = attn(                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 838 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0mhidden_states,                                                        \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 839 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0mencoder_hidden_states=encoder_hidden_states,                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 840 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0mcross_attention_kwargs=cross_attention_kwargs,                        \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/torch/nn/modules/\u001b[0m\u001b[1;33mmodule.py\u001b[0m:\u001b[94m111\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[94m0\u001b[0m in \u001b[92m_call_impl\u001b[0m                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1107 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# this function, and just call forward.\u001b[0m                                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1108 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[95mnot\u001b[0m (\u001b[96mself\u001b[0m._backward_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._forward_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._forward_pre_hooks \u001b[95mo\u001b[0m  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1109 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[95mor\u001b[0m _global_forward_hooks \u001b[95mor\u001b[0m _global_forward_pre_hooks):                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1110 \u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m forward_call(*\u001b[96minput\u001b[0m, **kwargs)                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1111 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# Do not call functions when jit is used\u001b[0m                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1112 \u001b[0m\u001b[2m│   │   \u001b[0mfull_backward_hooks, non_full_backward_hooks = [], []                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1113 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[96mself\u001b[0m._backward_hooks \u001b[95mor\u001b[0m _global_backward_hooks:                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/diffusers/models/\u001b[0m\u001b[1;33mtransformer_2\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[1;33md.py\u001b[0m:\u001b[94m265\u001b[0m in \u001b[92mforward\u001b[0m                                                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m262 \u001b[0m\u001b[2m│   │   \u001b[0m                                                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m263 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# 2. Blocks\u001b[0m                                                                        \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m264 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mfor\u001b[0m block \u001b[95min\u001b[0m \u001b[96mself\u001b[0m.transformer_blocks:                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m265 \u001b[2m│   │   │   \u001b[0mhidden_states = block(                                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m266 \u001b[0m\u001b[2m│   │   │   │   \u001b[0mhidden_states,                                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m267 \u001b[0m\u001b[2m│   │   │   │   \u001b[0mencoder_hidden_states=encoder_hidden_states,                               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m268 \u001b[0m\u001b[2m│   │   │   │   \u001b[0mtimestep=timestep,                                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/torch/nn/modules/\u001b[0m\u001b[1;33mmodule.py\u001b[0m:\u001b[94m111\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[94m0\u001b[0m in \u001b[92m_call_impl\u001b[0m                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1107 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# this function, and just call forward.\u001b[0m                                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1108 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[95mnot\u001b[0m (\u001b[96mself\u001b[0m._backward_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._forward_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._forward_pre_hooks \u001b[95mo\u001b[0m  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1109 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[95mor\u001b[0m _global_forward_hooks \u001b[95mor\u001b[0m _global_forward_pre_hooks):                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1110 \u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m forward_call(*\u001b[96minput\u001b[0m, **kwargs)                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1111 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# Do not call functions when jit is used\u001b[0m                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1112 \u001b[0m\u001b[2m│   │   \u001b[0mfull_backward_hooks, non_full_backward_hooks = [], []                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1113 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[96mself\u001b[0m._backward_hooks \u001b[95mor\u001b[0m _global_backward_hooks:                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/diffusers/models/\u001b[0m\u001b[1;33mattention.py\u001b[0m: \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[94m291\u001b[0m in \u001b[92mforward\u001b[0m                                                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m288 \u001b[0m\u001b[2m│   │   \u001b[0m                                                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m289 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# 1. Self-Attention\u001b[0m                                                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m290 \u001b[0m\u001b[2m│   │   \u001b[0mcross_attention_kwargs = cross_attention_kwargs \u001b[94mif\u001b[0m cross_attention_kwargs \u001b[95mis\u001b[0m \u001b[95mnot\u001b[0m   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m291 \u001b[2m│   │   \u001b[0mattn_output = \u001b[96mself\u001b[0m.attn1(                                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m292 \u001b[0m\u001b[2m│   │   │   \u001b[0mnorm_hidden_states,                                                            \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m293 \u001b[0m\u001b[2m│   │   │   \u001b[0mencoder_hidden_states=encoder_hidden_states \u001b[94mif\u001b[0m \u001b[96mself\u001b[0m.only_cross_attention \u001b[94mels\u001b[0m   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m294 \u001b[0m\u001b[2m│   │   │   \u001b[0mattention_mask=attention_mask,                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/home/qiang/anaconda3/envs/realfusion/lib/python3.8/site-packages/torch/nn/modules/\u001b[0m\u001b[1;33mmodule.py\u001b[0m:\u001b[94m111\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[94m0\u001b[0m in \u001b[92m_call_impl\u001b[0m                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1107 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# this function, and just call forward.\u001b[0m                                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1108 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[95mnot\u001b[0m (\u001b[96mself\u001b[0m._backward_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._forward_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._forward_pre_hooks \u001b[95mo\u001b[0m  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1109 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[95mor\u001b[0m _global_forward_hooks \u001b[95mor\u001b[0m _global_forward_pre_hooks):                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1110 \u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m forward_call(*\u001b[96minput\u001b[0m, **kwargs)                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1111 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# Do not call functions when jit is used\u001b[0m                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1112 \u001b[0m\u001b[2m│   │   \u001b[0mfull_backward_hooks, non_full_backward_hooks = [], []                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1113 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[96mself\u001b[0m._backward_hooks \u001b[95mor\u001b[0m _global_backward_hooks:                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n",
       "\u001b[1;91mTypeError: \u001b[0m\u001b[1;35mforward\u001b[0m\u001b[1m(\u001b[0m\u001b[1m)\u001b[0m got an unexpected keyword argument \u001b[32m'encoder_hidden_states'\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the the current cell or a previous cell. Please review the code in the cell(s) to identify a possible cause of the failure. Click <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. View Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "\n",
    "(image_gt, image_enc), x_t, uncond_embeddings = null_inversion.invert(image_path, prompt, offsets=offsets, verbose=True)\n",
    "\n",
    "print(\"Modify or remove offsets according to your image!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a36d450d-6764-4195-9ff1-842b2f60249e",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "426b5cc834ce4f4cb86d0b6266b9210d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "showing from left to right: the ground truth image, the vq-autoencoder reconstruction, the null-text inverted image\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 1584700 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=1556x512>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 225024 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=2352x307>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "prompts = [prompt]\n",
    "controller = AttentionStore()\n",
    "image_inv, x_t = run_and_display(prompts, controller, run_baseline=False, latent=x_t, uncond_embeddings=uncond_embeddings, verbose=False)\n",
    "print(\"showing from left to right: the ground truth image, the vq-autoencoder reconstruction, the null-text inverted image\")\n",
    "ptp_utils.view_images([image_gt, image_enc, image_inv[0]])\n",
    "show_cross_attention(controller, 16, [\"up\", \"down\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "da742d87-b48e-40e1-8b8e-c0f9b7528cc9",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a43bd01468e4b978e3ffa11b89699ee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 1157460 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=1034x512>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image is highly affected by the self_replace_steps, usually 0.4 is a good default value, but you may want to try the range 0.3,0.4,0.5,0.7 \n"
     ]
    }
   ],
   "source": [
    "prompts = [\"a cat sitting next to a mirror\",\n",
    "           \"a tiger sitting next to a mirror\"\n",
    "        ]\n",
    "\n",
    "cross_replace_steps = {'default_': .8,}\n",
    "self_replace_steps = .5\n",
    "blend_word = ((('cat',), (\"tiger\",))) # for local edit. If it is not local yet - use only the source object: blend_word = ((('cat',), (\"cat\",))).\n",
    "eq_params = {\"words\": (\"tiger\",), \"values\": (2,)} # amplify attention to the word \"tiger\" by *2 \n",
    "\n",
    "controller = make_controller(prompts, True, cross_replace_steps, self_replace_steps, blend_word, eq_params)\n",
    "images, _ = run_and_display(prompts, controller, run_baseline=False, latent=x_t, uncond_embeddings=uncond_embeddings)\n",
    "\n",
    "print(\"Image is highly affected by the self_replace_steps, usually 0.4 is a good default value, but you may want to try the range 0.3,0.4,0.5,0.7 \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "861db4be-72cc-4d8f-969b-bba1bf45bb50",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1171205e989f4c48898dbb50c5803de6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 1172196 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=1034x512>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "prompts = [\"a cat sitting next to a mirror\",\n",
    "           \"a silver cat sculpture sitting next to a mirror\"\n",
    "        ]\n",
    "\n",
    "cross_replace_steps = {'default_': .8, }\n",
    "self_replace_steps = .6\n",
    "blend_word = ((('cat',), (\"cat\",))) # for local edit\n",
    "eq_params = {\"words\": (\"silver\", 'sculpture', ), \"values\": (2,2,)}  # amplify attention to the words \"silver\" and \"sculpture\" by *2 \n",
    " \n",
    "controller = make_controller(prompts, False, cross_replace_steps, self_replace_steps, blend_word, eq_params)\n",
    "images, _ = run_and_display(prompts, controller, run_baseline=False, latent=x_t, uncond_embeddings=uncond_embeddings)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4b768dd2-a139-4163-823e-15318441ea49",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8d200aaed0624a14ba14b5389c9eb64c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 1149052 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=1034x512>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "prompts = [\"a cat sitting next to a mirror\",\n",
    "           \"watercolor painting of a cat sitting next to a mirror\"\n",
    "        ]\n",
    "\n",
    "cross_replace_steps = {'default_': .8, }\n",
    "self_replace_steps = .7\n",
    "blend_word = None\n",
    "eq_params = {\"words\": (\"watercolor\",  ), \"values\": (5, 2,)}  # amplify attention to the word \"watercolor\" by 5\n",
    " \n",
    "controller = make_controller(prompts, False, cross_replace_steps, self_replace_steps, blend_word, eq_params)\n",
    "images, _ = run_and_display(prompts, controller, run_baseline=False, latent=x_t, uncond_embeddings=uncond_embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6adfb554-148b-45af-ae8b-17b213f9070f",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "pytorch-gpu.1-12.m97",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/pytorch-gpu.1-12:m97"
  },
  "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
