Yuno_Diffusion.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/FuouM/Yuno-Diffusion/blob/main/Yuno_Diffusion.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xuhEr-E2Y6K1"
      },
      "source": [
        "Adapted from https://colab.research.google.com/drive/1Iy-xW9t1-OQWhb0hNxueGij8phCyluOh\n",
        "\n",
        "Archive: https://github.com/FuouM/Yuno-Diffusion\n",
        "\n",
        "No upscale and gradio\n",
        "\n",
        "Colab de Voldemort"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n58jZbYMjnY_"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6raQLlYdjoHn"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.kill(os.getpid(), 9) # This will crash Colab (required, everything will still be intact so dont worry)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DHfrnDNPY5FO"
      },
      "source": [
        "### Download and set up environment (Just run)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AZN5mpwbj2Bq"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui\n",
        "%cd stable-diffusion-webui"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir -p /content/stable-diffusion-webui/Source\n",
        "!mkdir -p /content/stable-diffusion-webui/Output"
      ],
      "metadata": {
        "id": "X4LCqt7b1oZ-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!wget  https://github.com/raw/FuouM/stable-diffusion-hidamari/main/NotoSansJP-Bold.otf"
      ],
      "metadata": {
        "id": "bc3ad1UP60x1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uuymjGaNkSDc"
      },
      "outputs": [],
      "source": [
        "# The requirements.txt has \"invisible-watermark\" which is pretty sus\n",
        "!pip install numpy\n",
        "!pip install Pillow\n",
        "!pip install realesrgan\n",
        "!pip install torch\n",
        "!pip install transformers\n",
        "!pip install omegaconf\n",
        "!pip install pytorch_lightning\n",
        "!pip install diffusers\n",
        "!pip install scikit-image\n",
        "!pip install git+https://github.com/crowsonkb/k-diffusion.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zMlAM4yhknyW"
      },
      "outputs": [],
      "source": [
        "!mkdir repositories\n",
        "!git clone https://github.com/CompVis/stable-diffusion.git repositories/stable-diffusion\n",
        "!git clone https://github.com/CompVis/taming-transformers.git repositories/taming-transformers"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "current = \"\""
      ],
      "metadata": {
        "id": "5Y20WOCkQqzM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G_wA3uADlO6n"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.kill(os.getpid(), 9) # This will crash Colab (required, everything will still be intact so dont worry)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XZHZTAgdjymb"
      },
      "source": [
        "### Model Download (Expand)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c753aosFpukU"
      },
      "source": [
        "#### Huggingface"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3RVNpLfdpjvl"
      },
      "outputs": [],
      "source": [
        "#@title Normal 1.4 model\n",
        "%cd /content/stable-diffusion-webui\n",
        "user_token = \"\" #@param {type:\"string\"}\n",
        "user_header = f\"\\\"Authorization: Bearer {user_token}\\\"\"\n",
        "!wget --header={user_header} https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt -O /content/stable-diffusion-webui/model.ckpt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "OSP7c8hKpmvv"
      },
      "outputs": [],
      "source": [
        "%cd /content/stable-diffusion-webui\n",
        "#@title Anime 1.2 model\n",
        "user_token = \"\" #@param {type:\"string\"}\n",
        "user_header = f\"\\\"Authorization: Bearer {user_token}\\\"\"\n",
        "!wget --header={user_header} https://huggingface.co/naclbit/trinart_stable_diffusion/resolve/main/trinart_stable_diffusion_epoch3.ckpt \n",
        "!mv trinart_stable_diffusion_epoch3.ckpt model_anime.ckpt\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LJPrA0GepxQ7"
      },
      "source": [
        "#### (Danbooru) - DOESN'T WORK\n",
        "\n",
        "Takes a very long time ~ 16 minutes"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/stable-diffusion-webui\n",
        "!wget https://thisanimedoesnotexist.ai/downloads/wd-v1-2-full-ema.ckpt\n",
        "!mv wd-v1-2-full-ema.ckpt model_danbooru.ckpt"
      ],
      "metadata": {
        "id": "abLpVvnbQ1tu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bGkPdIqVp0St"
      },
      "source": [
        "#### Merged Danbooru and SD (Mega)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown merged danbooru sd\n",
        "import sys, os, urllib.request\n",
        "import time\n",
        "import subprocess\n",
        "import contextlib\n",
        "from IPython.display import clear_output\n",
        "\n",
        "HOME = os.path.expanduser(\"~\")\n",
        "if not os.path.exists(f\"{HOME}/.ipython/ocr.py\"):\n",
        "    hCode = \"https://github.com/raw/biplobsd/\" \\\n",
        "                \"OneClickRun/master/res/ocr.py\"\n",
        "    urllib.request.urlretrieve(hCode, f\"{HOME}/.ipython/ocr.py\")\n",
        "\n",
        "from ocr import (\n",
        "    runSh,\n",
        "    loadingAn,\n",
        ")\n",
        "\n",
        "URL = \"https://mega.nz/file/IwxQTQCb#tsfYZeHrd6U3fg5QCmBazEL-OjOmy1gmhvKAtp1YWFs\" #@param {type:'string'}\n",
        "OUTPUT_PATH = \"/content/stable-diffusion-webui\"\n",
        "if not OUTPUT_PATH:\n",
        "  os.makedirs(\"downloads\", exist_ok=True)\n",
        "  OUTPUT_PATH = \"downloads\"\n",
        "# MEGAcmd installing\n",
        "if not os.path.exists(\"/usr/bin/mega-cmd\"):\n",
        "    loadingAn()\n",
        "    print(\"Installing MEGA ...\")\n",
        "    runSh('sudo apt-get -y update')\n",
        "    runSh('sudo apt-get -y install libmms0 libc-ares2 libc6 libcrypto++6 libgcc1 libmediainfo0v5 libpcre3 libpcrecpp0v5 libssl1.1 libstdc++6 libzen0v5 zlib1g apt-transport-https')\n",
        "    runSh('sudo curl -sL -o /var/cache/apt/archives/MEGAcmd.deb https://mega.nz/linux/MEGAsync/Debian_9.0/amd64/megacmd-Debian_9.0_amd64.deb', output=True)\n",
        "    runSh('sudo dpkg -i /var/cache/apt/archives/MEGAcmd.deb', output=True)\n",
        "    print(\"MEGA is installed.\")\n",
        "    clear_output()\n",
        "\n",
        "# Unix, Windows and old Macintosh end-of-line\n",
        "newlines = ['\\n', '\\r\\n', '\\r']\n",
        "\n",
        "def unbuffered(proc, stream='stdout'):\n",
        "    stream = getattr(proc, stream)\n",
        "    with contextlib.closing(stream):\n",
        "        while True:\n",
        "            out = []\n",
        "            last = stream.read(1)\n",
        "            # Don't loop forever\n",
        "            if last == '' and proc.poll() is not None:\n",
        "                break\n",
        "            while last not in newlines:\n",
        "                # Don't loop forever\n",
        "                if last == '' and proc.poll() is not None:\n",
        "                    break\n",
        "                out.append(last)\n",
        "                last = stream.read(1)\n",
        "            out = ''.join(out)\n",
        "            yield out\n",
        "\n",
        "\n",
        "def transfare():\n",
        "    import codecs\n",
        "    decoder = codecs.getincrementaldecoder(\"UTF-8\")()\n",
        "    cmd = [\"mega-get\", URL, OUTPUT_PATH]\n",
        "    proc = subprocess.Popen(\n",
        "        cmd,\n",
        "        stdout=subprocess.PIPE,\n",
        "        stderr=subprocess.STDOUT,\n",
        "        # Make all end-of-lines '\\n'\n",
        "        universal_newlines=True,\n",
        "    )\n",
        "    for line in unbuffered(proc):\n",
        "        print(line, end=\"\\r\")\n",
        "        \n",
        "\n",
        "\n",
        "transfare()\n",
        "\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "7oXlhgBZaT7j"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown change model name to model_merged.ckpt\n",
        "file_name = 'model (1).ckpt' #@param {type: 'string'}\n",
        "file_name = f\"\\\"{file_name}\\\"\"\n",
        "!mv {file_name} model_merged.ckpt"
      ],
      "metadata": {
        "cellView": "form",
        "id": "d4yLMA1Oc_kh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RE_wPYZLqA1S"
      },
      "source": [
        "### Change model \n",
        "\n",
        "Default is \"normal\" (no need to run if that is the case). \n",
        "\n",
        "Otherwise, run every runtime (Restart every change)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.kill(os.getpid(), 9)"
      ],
      "metadata": {
        "id": "ub_ytuGITu49"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yA-4bCb3p885",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "curr_model = \"merged\" #@param [\"normal\",\"danbooru\", \"anime\", \"merged\"] {allow-input: false}\n",
        "if curr_model == \"anime\":\n",
        "  current = \"_anime\"\n",
        "elif curr_model == \"danbooru\":\n",
        "  current = \"_danbooru\"\n",
        "elif curr_model == \"merged\": \n",
        "  current = \"_merged\"\n",
        "else:\n",
        "  current = \"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vf5WJ4b3lSOD"
      },
      "source": [
        "# Prepare (Just run)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vlWxMWnTtsIW"
      },
      "source": [
        "##### First"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aKdUb9oQlYBP"
      },
      "outputs": [],
      "source": [
        "%cd stable-diffusion-webui\n",
        "!git pull"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/stable-diffusion-webui"
      ],
      "metadata": {
        "id": "Bb9T3Aas1eoi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u-9AzVB_mPgU"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "from omegaconf import OmegaConf\n",
        "from PIL import Image, ImageFilter, ImageOps, ImageChops, ImageDraw, ImageFont\n",
        "import random\n",
        "import tqdm\n",
        "from collections import namedtuple\n",
        "import contextlib\n",
        "import json\n",
        "import math\n",
        "import cv2\n",
        "\n",
        "import k_diffusion.sampling\n",
        "%cd /content/stable-diffusion-webui/repositories/stable-diffusion\n",
        "from ldm.util import instantiate_from_config\n",
        "import ldm.models.diffusion.ddim\n",
        "import ldm.models.diffusion.plms\n",
        "\n",
        "%cd /content/stable-diffusion-webui/repositories/taming-transformers\n",
        "import taming\n",
        "\n",
        "%cd /content/stable-diffusion-webui"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IEdOfzjXtONk"
      },
      "outputs": [],
      "source": [
        "config_filename = \"config.json\"\n",
        "sd_model_file = f\"/content/stable-diffusion-webui/model{current}.ckpt\"\n",
        "device = torch.device(\"cuda\")\n",
        "batch_cond_uncond = True\n",
        "parallel_processing_allowed = False"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gzEgNxZ6xn1G"
      },
      "source": [
        "##### Caption maker"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Awg98YBqxq5h"
      },
      "outputs": [],
      "source": [
        "import PIL\n",
        "from PIL import Image, ImageFont, ImageDraw \n",
        "\n",
        "def add_margin(pil_img, top, right, bottom, left, color):\n",
        "    width, height = pil_img.size\n",
        "    new_width = width + right + left\n",
        "    new_height = height + top + bottom\n",
        "    result = Image.new(pil_img.mode, (new_width, new_height), color)\n",
        "    result.paste(pil_img, (left, top))\n",
        "    return result\n",
        "\n",
        "def text_wrap(text, font, max_width):\n",
        "    lines = []\n",
        "    if font.getsize(text)[0]  <= max_width:\n",
        "        lines.append(text)\n",
        "    else:\n",
        "        words = text.split(' ')\n",
        "        i = 0\n",
        "        while i < len(words):\n",
        "            line = ''\n",
        "            while i < len(words) and font.getsize(line + words[i])[0] <= max_width:\n",
        "                line = line + words[i]+ \" \"\n",
        "                i += 1\n",
        "            if not line:\n",
        "                line = words[i]\n",
        "                i += 1\n",
        "            lines.append(line)\n",
        "    return lines\n",
        "\n",
        "def caption(image, prompt, info):\n",
        "    width, height = image.size\n",
        "\n",
        "    font = ImageFont.truetype(\"/content/stable-diffusion-webui/NotoSansJP-Bold.otf\", 20)\n",
        "    print(font)\n",
        "    lines = text_wrap(prompt, font, image.size[0])\n",
        "    lines.append(f\"{info}\")\n",
        "    line_height = font.getsize('hg')[1]\n",
        "    cap_img = add_margin(image, 0, 0, line_height * (len(lines) + 1), 0, (255, 255, 255))\n",
        "    draw = ImageDraw.Draw(cap_img)\n",
        "    pad = 2\n",
        "    x = pad * 2\n",
        "    y = height + pad\n",
        "    for line in lines:\n",
        "        draw.text((x,y), line, fill=(0, 0, 0), font=font)\n",
        "        y = y + line_height\n",
        "    return cap_img\n",
        "\n",
        "def get_concat_h_blank(im1, im2, color=(255, 255, 255)):\n",
        "    dst = Image.new('RGB', (im1.width + im2.width, max(im1.height, im2.height)), color)\n",
        "    dst.paste(im1, (0, 0))\n",
        "    dst.paste(im2, (im1.width, 0))\n",
        "    return dst\n",
        "\n",
        "def get_concat_v_blank(im1, im2, color=(255, 255, 255)):\n",
        "    dst = Image.new('RGB', (max(im1.width, im2.width), im1.height + im2.height), color)\n",
        "    dst.paste(im1, (0, 0))\n",
        "    dst.paste(im2, (0, im1.height))\n",
        "    return dst\n",
        "\n",
        "def image_grid(imgs, batch_size, n_rows:int):\n",
        "    if n_rows > 0:\n",
        "        rows = n_rows\n",
        "    elif n_rows == 0:\n",
        "        rows = batch_size\n",
        "    else:\n",
        "        rows = math.sqrt(len(imgs))\n",
        "        rows = round(rows)\n",
        "\n",
        "    cols = math.ceil(len(imgs) / rows)\n",
        "\n",
        "    w, h = imgs[0].size\n",
        "    grid = Image.new('RGB', size=(cols * w, rows * h), color='black')\n",
        "\n",
        "    for i, img in enumerate(imgs):\n",
        "        grid.paste(img, box=(i % cols * w, i // cols * h))\n",
        "\n",
        "    return grid"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6pfVOIy0pL4e"
      },
      "source": [
        "##### shared.py"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2BlxdAgnpjiN"
      },
      "outputs": [],
      "source": [
        "def load_model_from_config(config, ckpt, verbose=False):\n",
        "    print(f\"Loading model from {ckpt}\")\n",
        "    pl_sd = torch.load(ckpt, map_location=\"cpu\")\n",
        "    if \"global_step\" in pl_sd:\n",
        "        print(f\"Global Step: {pl_sd['global_step']}\")\n",
        "    sd = pl_sd[\"state_dict\"]\n",
        "\n",
        "    model = instantiate_from_config(config.model)\n",
        "    m, u = model.load_state_dict(sd, strict=False)\n",
        "    if len(m) > 0 and verbose:\n",
        "        print(\"missing keys:\")\n",
        "        print(m)\n",
        "    if len(u) > 0 and verbose:\n",
        "        print(\"unexpected keys:\")\n",
        "        print(u)\n",
        "\n",
        "    model.eval()\n",
        "    return model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ecjbHSOmsw--"
      },
      "outputs": [],
      "source": [
        "class State:\n",
        "    interrupted = False\n",
        "    job = \"\"\n",
        "    job_no = 0\n",
        "    job_count = 0\n",
        "    sampling_step = 0\n",
        "    sampling_steps = 0\n",
        "    current_latent = None\n",
        "    current_image = None\n",
        "    current_image_sampling_step = 0\n",
        "\n",
        "\n",
        "    def interrupt(self):\n",
        "        self.interrupted = True\n",
        "\n",
        "    def nextjob(self):\n",
        "        self.job_no += 1\n",
        "        self.sampling_step = 0\n",
        "        self.current_image_sampling_step = 0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M7EOopVRsuRE"
      },
      "outputs": [],
      "source": [
        "state = State()\n",
        "sd_upscalers = []"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eDakK97gsPHY"
      },
      "source": [
        "##### sd_samplers.py (Done)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def name_index(user_sampler: str):\n",
        "  if user_sampler == 'k_euler_a':\n",
        "    return 0\n",
        "  elif user_sampler == 'k_euler':\n",
        "    return 1\n",
        "  elif user_sampler == 'k_lms':\n",
        "    return 2\n",
        "  elif user_sampler == 'k_heun':\n",
        "    return 3\n",
        "  elif user_sampler == 'k_dpm_2':\n",
        "    return 4\n",
        "  elif user_sampler == 'k_dpm_2_a':\n",
        "    return 5  "
      ],
      "metadata": {
        "id": "Xjo5IyWCzkYt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_1e5CzXFsTaf"
      },
      "outputs": [],
      "source": [
        "SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases'])\n",
        "\n",
        "samplers_k_diffusion = [\n",
        "    ('Euler a', 'sample_euler_ancestral', ['k_euler_a']),\n",
        "    ('Euler', 'sample_euler', ['k_euler']),\n",
        "    ('LMS', 'sample_lms', ['k_lms']),\n",
        "    ('Heun', 'sample_heun', ['k_heun']),\n",
        "    ('DPM2', 'sample_dpm_2', ['k_dpm_2']),\n",
        "    ('DPM2 a', 'sample_dpm_2_ancestral', ['k_dpm_2_a']),\n",
        "]\n",
        "\n",
        "samplers_data_k_diffusion = [\n",
        "    SamplerData(label, lambda model, funcname=funcname: KDiffusionSampler(funcname, model), aliases)\n",
        "    for label, funcname, aliases in samplers_k_diffusion\n",
        "    if hasattr(k_diffusion.sampling, funcname)\n",
        "]\n",
        "\n",
        "samplers = [\n",
        "    *samplers_data_k_diffusion,\n",
        "    SamplerData('DDIM', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.ddim.DDIMSampler, model), []),\n",
        "    SamplerData('PLMS', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.plms.PLMSSampler, model), []),\n",
        "]\n",
        "samplers_for_img2img = [x for x in samplers if x.name != 'PLMS']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T6wPEopHsbxy"
      },
      "outputs": [],
      "source": [
        "def sample_to_image(samples):\n",
        "    x_sample = sd_model.decode_first_stage(samples[0:1].type(sd_model.dtype))[0]\n",
        "    x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)\n",
        "    x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)\n",
        "    x_sample = x_sample.astype(np.uint8)\n",
        "    return Image.fromarray(x_sample)\n",
        "\n",
        "\n",
        "def store_latent(decoded):\n",
        "    state.current_latent = decoded"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8lo3Bh5Ds100"
      },
      "outputs": [],
      "source": [
        "def p_sample_ddim_hook(sampler_wrapper, x_dec, cond, ts, *args, **kwargs):\n",
        "    if sampler_wrapper.mask is not None:\n",
        "        img_orig = sampler_wrapper.sampler.model.q_sample(sampler_wrapper.init_latent, ts)\n",
        "        x_dec = img_orig * sampler_wrapper.mask + sampler_wrapper.nmask * x_dec\n",
        "\n",
        "        store_latent(sampler_wrapper.init_latent * sampler_wrapper.mask + sampler_wrapper.nmask * x_dec)\n",
        "\n",
        "    else:\n",
        "        store_latent(x_dec)\n",
        "\n",
        "    return sampler_wrapper.orig_p_sample_ddim(x_dec, cond, ts, *args, **kwargs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2O-4Zn4Es4B9"
      },
      "outputs": [],
      "source": [
        "def extended_tdqm(sequence, *args, desc=None, **kwargs):\n",
        "    state.sampling_steps = len(sequence)\n",
        "    state.sampling_step = 0\n",
        "\n",
        "    for x in tqdm.tqdm(sequence, *args, desc=state.job, **kwargs):\n",
        "        if state.interrupted:\n",
        "            break\n",
        "\n",
        "        yield x\n",
        "\n",
        "        state.sampling_step += 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DnQTKURqs9h1"
      },
      "outputs": [],
      "source": [
        "def extended_trange(count, *args, **kwargs):\n",
        "    state.sampling_steps = count\n",
        "    state.sampling_step = 0\n",
        "\n",
        "    for x in tqdm.trange(count, *args, desc=state.job, **kwargs):\n",
        "        if state.interrupted:\n",
        "            break\n",
        "\n",
        "        yield x\n",
        "\n",
        "        state.sampling_step += 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "71jbs6TOs4oz"
      },
      "outputs": [],
      "source": [
        "ldm.models.diffusion.ddim.tqdm = lambda *args, desc=None, **kwargs: extended_tdqm(*args, desc=desc, **kwargs)\n",
        "ldm.models.diffusion.plms.tqdm = lambda *args, desc=None, **kwargs: extended_tdqm(*args, desc=desc, **kwargs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yVkxzBafs6kb"
      },
      "outputs": [],
      "source": [
        "class VanillaStableDiffusionSampler:\n",
        "    def __init__(self, constructor, sd_model):\n",
        "        self.sampler = constructor(sd_model)\n",
        "        self.orig_p_sample_ddim = self.sampler.p_sample_ddim if hasattr(self.sampler, 'p_sample_ddim') else None\n",
        "        self.mask = None\n",
        "        self.nmask = None\n",
        "        self.init_latent = None\n",
        "\n",
        "    def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning):\n",
        "        t_enc = int(min(p.denoising_strength, 0.999) * p.steps)\n",
        "\n",
        "        # existing code fails with cetin step counts, like 9\n",
        "        try:\n",
        "            self.sampler.make_schedule(ddim_num_steps=p.steps, verbose=False)\n",
        "        except Exception:\n",
        "            self.sampler.make_schedule(ddim_num_steps=p.steps+1, verbose=False)\n",
        "\n",
        "        x1 = self.sampler.stochastic_encode(x, torch.tensor([t_enc] * int(x.shape[0])).to(shared.device), noise=noise)\n",
        "\n",
        "        self.sampler.p_sample_ddim = lambda x_dec, cond, ts, *args, **kwargs: p_sample_ddim_hook(self, x_dec, cond, ts, *args, **kwargs)\n",
        "        self.mask = p.mask\n",
        "        self.nmask = p.nmask\n",
        "        self.init_latent = p.init_latent\n",
        "\n",
        "        samples = self.sampler.decode(x1, conditioning, t_enc, unconditional_guidance_scale=p.cfg_scale, unconditional_conditioning=unconditional_conditioning)\n",
        "\n",
        "        return samples\n",
        "\n",
        "    def sample(self, p, x, conditioning, unconditional_conditioning):\n",
        "        samples_ddim, _ = self.sampler.sample(S=p.steps, conditioning=conditioning, batch_size=int(x.shape[0]), shape=x[0].shape, verbose=False, unconditional_guidance_scale=p.cfg_scale, unconditional_conditioning=unconditional_conditioning, x_T=x)\n",
        "        return samples_ddim\n",
        "\n",
        "\n",
        "class CFGDenoiser(torch.nn.Module):\n",
        "    def __init__(self, model):\n",
        "        super().__init__()\n",
        "        self.inner_model = model\n",
        "        self.mask = None\n",
        "        self.nmask = None\n",
        "        self.init_latent = None\n",
        "\n",
        "    def forward(self, x, sigma, uncond, cond, cond_scale):\n",
        "        if batch_cond_uncond:\n",
        "            x_in = torch.cat([x] * 2)\n",
        "            sigma_in = torch.cat([sigma] * 2)\n",
        "            cond_in = torch.cat([uncond, cond])\n",
        "            uncond, cond = self.inner_model(x_in, sigma_in, cond=cond_in).chunk(2)\n",
        "            denoised = uncond + (cond - uncond) * cond_scale\n",
        "        else:\n",
        "            uncond = self.inner_model(x, sigma, cond=uncond)\n",
        "            cond = self.inner_model(x, sigma, cond=cond)\n",
        "            denoised = uncond + (cond - uncond) * cond_scale\n",
        "\n",
        "        if self.mask is not None:\n",
        "            denoised = self.init_latent * self.mask + self.nmask * denoised\n",
        "\n",
        "        return denoised\n",
        "\n",
        "class KDiffusionSampler:\n",
        "    def __init__(self, funcname, sd_model):\n",
        "        self.model_wrap = k_diffusion.external.CompVisDenoiser(sd_model)\n",
        "        self.funcname = funcname\n",
        "        self.func = getattr(k_diffusion.sampling, self.funcname)\n",
        "        self.model_wrap_cfg = CFGDenoiser(self.model_wrap)\n",
        "\n",
        "    def callback_state(self, d):\n",
        "        store_latent(d[\"denoised\"])\n",
        "\n",
        "    def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning):\n",
        "        t_enc = int(min(p.denoising_strength, 0.999) * p.steps)\n",
        "        sigmas = self.model_wrap.get_sigmas(p.steps)\n",
        "        noise = noise * sigmas[p.steps - t_enc - 1]\n",
        "\n",
        "        xi = x + noise\n",
        "\n",
        "        sigma_sched = sigmas[p.steps - t_enc - 1:]\n",
        "\n",
        "        self.model_wrap_cfg.mask = p.mask\n",
        "        self.model_wrap_cfg.nmask = p.nmask\n",
        "        self.model_wrap_cfg.init_latent = p.init_latent\n",
        "\n",
        "        if hasattr(k_diffusion.sampling, 'trange'):\n",
        "            k_diffusion.sampling.trange = lambda *args, **kwargs: extended_trange(*args, **kwargs)\n",
        "\n",
        "        return self.func(self.model_wrap_cfg, xi, sigma_sched, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False, callback=self.callback_state)\n",
        "\n",
        "    def sample(self, p, x, conditioning, unconditional_conditioning):\n",
        "        sigmas = self.model_wrap.get_sigmas(p.steps)\n",
        "        x = x * sigmas[0]\n",
        "\n",
        "        if hasattr(k_diffusion.sampling, 'trange'):\n",
        "            k_diffusion.sampling.trange = lambda *args, **kwargs: extended_trange(*args, **kwargs)\n",
        "\n",
        "        samples_ddim = self.func(self.model_wrap_cfg, x, sigmas, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False, callback=self.callback_state)\n",
        "        return samples_ddim"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s-Yr1u3so4X6"
      },
      "source": [
        "##### processing.py (Done)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NTp7EeNyr4ER"
      },
      "outputs": [],
      "source": [
        "def torch_gc():\n",
        "    if torch.cuda.is_available():\n",
        "        torch.cuda.empty_cache()\n",
        "        torch.cuda.ipc_collect()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wnI5ghn7r4oV"
      },
      "outputs": [],
      "source": [
        "class StableDiffusionProcessing:\n",
        "    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt=\"\", seed=-1, sampler_index=0, batch_size=1, n_iter=1, steps=50, cfg_scale=7.0, width=512, height=512, restore_faces=False, tiling=False, do_not_save_samples=False, do_not_save_grid=False, extra_generation_params=None, overlay_images=None, negative_prompt=None):\n",
        "        self.sd_model = sd_model\n",
        "        self.outpath_samples: str = outpath_samples\n",
        "        self.outpath_grids: str = outpath_grids\n",
        "        self.prompt: str = prompt\n",
        "        self.prompt_for_display: str = None\n",
        "        self.negative_prompt: str = (negative_prompt or \"\")\n",
        "        self.seed: int = seed\n",
        "        self.sampler_index: int = sampler_index\n",
        "        self.batch_size: int = batch_size\n",
        "        self.n_iter: int = n_iter\n",
        "        self.steps: int = steps\n",
        "        self.cfg_scale: float = cfg_scale\n",
        "        self.width: int = width\n",
        "        self.height: int = height\n",
        "        self.restore_faces: bool = restore_faces\n",
        "        self.tiling: bool = tiling\n",
        "        self.do_not_save_samples: bool = do_not_save_samples\n",
        "        self.do_not_save_grid: bool = do_not_save_grid\n",
        "        self.extra_generation_params: dict = extra_generation_params\n",
        "        self.overlay_images = overlay_images\n",
        "        self.paste_to = None\n",
        "\n",
        "    def init(self, seed):\n",
        "        pass\n",
        "\n",
        "    def sample(self, x, conditioning, unconditional_conditioning):\n",
        "        raise NotImplementedError()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iGPitgorsC2U"
      },
      "outputs": [],
      "source": [
        "class Processed:\n",
        "    def __init__(self, p: StableDiffusionProcessing, images_list, seed, info):\n",
        "        self.images = images_list\n",
        "        self.prompt = p.prompt\n",
        "        self.seed = seed\n",
        "        self.info = info\n",
        "        self.width = p.width\n",
        "        self.height = p.height\n",
        "        self.sampler = samplers[p.sampler_index].name\n",
        "        self.cfg_scale = p.cfg_scale\n",
        "        self.steps = p.steps\n",
        "\n",
        "    def js(self):\n",
        "        obj = {\n",
        "            \"prompt\": self.prompt if type(self.prompt) != list else self.prompt[0],\n",
        "            \"seed\": int(self.seed if type(self.seed) != list else self.seed[0]),\n",
        "            \"width\": self.width,\n",
        "            \"height\": self.height,\n",
        "            \"sampler\": self.sampler,\n",
        "            \"cfg_scale\": self.cfg_scale,\n",
        "            \"steps\": self.steps,\n",
        "        }\n",
        "\n",
        "        return json.dumps(obj)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9ncyJ1AlsDft"
      },
      "outputs": [],
      "source": [
        "def create_random_tensors(shape, seeds):\n",
        "    xs = []\n",
        "    for seed in seeds:\n",
        "        torch.manual_seed(seed)\n",
        "\n",
        "        # randn results depend on device; gpu and cpu get different results for same seed;\n",
        "        # the way I see it, it's better to do this on CPU, so that everyone gets same result;\n",
        "        # but the original script had it like this so I do not dare change it for now because\n",
        "        # it will break everyone's seeds.\n",
        "        xs.append(torch.randn(shape, device=device))\n",
        "    x = torch.stack(xs)\n",
        "    return x\n",
        "\n",
        "\n",
        "def set_seed(seed):\n",
        "    return int(random.randrange(4294967294)) if seed is None or seed == -1 else seed"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QrNLVfu6usPq"
      },
      "outputs": [],
      "source": [
        "opt_C = 4\n",
        "opt_f = 8"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TwZ7dCNHuNBc"
      },
      "outputs": [],
      "source": [
        "def process_images(p: StableDiffusionProcessing) -> Processed:\n",
        "    \"\"\"this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch\"\"\"\n",
        "\n",
        "    prompt = p.prompt\n",
        "\n",
        "    assert p.prompt is not None\n",
        "    torch_gc()\n",
        "\n",
        "    seed = set_seed(p.seed)\n",
        "\n",
        "    # os.makedirs(p.outpath_samples, exist_ok=True)\n",
        "    # os.makedirs(p.outpath_grids, exist_ok=True)\n",
        "\n",
        "    # modules.sd_hijack.model_hijack.apply_circular(p.tiling)\n",
        "\n",
        "    comments = []\n",
        "\n",
        "    if type(prompt) == list:\n",
        "        all_prompts = prompt\n",
        "    else:\n",
        "        all_prompts = p.batch_size * p.n_iter * [prompt]\n",
        "\n",
        "    if type(seed) == list:\n",
        "        all_seeds = seed\n",
        "    else:\n",
        "        all_seeds = [int(seed + x) for x in range(len(all_prompts))]\n",
        "\n",
        "    def infotext(iteration=0, position_in_batch=0):\n",
        "        generation_params = {\n",
        "            \"Steps\": p.steps,\n",
        "            \"Sampler\": samplers[p.sampler_index].name,\n",
        "            \"CFG scale\": p.cfg_scale,\n",
        "            \"Seed\": all_seeds[position_in_batch + iteration * p.batch_size],\n",
        "            # \"Face restoration\": (opts.face_restoration_model if p.restore_faces else None),\n",
        "            \"Batch size\": (None if p.batch_size < 2 else p.batch_size),\n",
        "            \"Batch pos\": (None if p.batch_size < 2 else position_in_batch),\n",
        "        }\n",
        "\n",
        "        if p.extra_generation_params is not None:\n",
        "            generation_params.update(p.extra_generation_params)\n",
        "\n",
        "        generation_params_text = \", \".join([k if k == v else f'{k}: {v}' for k, v in generation_params.items() if v is not None])\n",
        "\n",
        "        return f\"{p.prompt_for_display or prompt}\\n{generation_params_text}\".strip() + \"\".join([\"\\n\\n\" + x for x in comments])\n",
        "\n",
        "    # if os.path.exists(cmd_opts.embeddings_dir):\n",
        "    #     model_hijack.load_textual_inversion_embeddings(cmd_opts.embeddings_dir, p.sd_model)\n",
        "\n",
        "    output_images = []\n",
        "    precision_scope = torch.autocast\n",
        "    ema_scope = (contextlib.nullcontext)\n",
        "    with torch.no_grad(), precision_scope(\"cuda\"), ema_scope():\n",
        "        p.init(seed=all_seeds[0])\n",
        "\n",
        "        if state.job_count == -1:\n",
        "            state.job_count = p.n_iter\n",
        "\n",
        "        for n in range(p.n_iter):\n",
        "            if state.interrupted:\n",
        "                break\n",
        "\n",
        "            prompts = all_prompts[n * p.batch_size:(n + 1) * p.batch_size]\n",
        "            seeds = all_seeds[n * p.batch_size:(n + 1) * p.batch_size]\n",
        "\n",
        "            uc = p.sd_model.get_learned_conditioning(len(prompts) * [p.negative_prompt])\n",
        "            c = p.sd_model.get_learned_conditioning(prompts)\n",
        "\n",
        "            # we manually generate all input noises because each one should have a specific seed\n",
        "            x = create_random_tensors([opt_C, p.height // opt_f, p.width // opt_f], seeds=seeds)\n",
        "\n",
        "            if p.n_iter > 1:\n",
        "                state.job = f\"Batch {n+1} out of {p.n_iter}\"\n",
        "\n",
        "            samples_ddim = p.sample(x=x, conditioning=c, unconditional_conditioning=uc)\n",
        "            if state.interrupted:\n",
        "\n",
        "                # if we are interruped, sample returns just noise\n",
        "                # use the image collected previously in sampler loop\n",
        "                samples_ddim = state.current_latent\n",
        "\n",
        "            x_samples_ddim = p.sd_model.decode_first_stage(samples_ddim)\n",
        "            x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)\n",
        "\n",
        "            for i, x_sample in enumerate(x_samples_ddim):\n",
        "                x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)\n",
        "                x_sample = x_sample.astype(np.uint8)\n",
        "\n",
        "                if p.restore_faces:\n",
        "                    torch_gc()\n",
        "\n",
        "                    # x_sample = modules.face_restoration.restore_faces(x_sample)\n",
        "\n",
        "                image = Image.fromarray(x_sample)\n",
        "\n",
        "                if p.overlay_images is not None and i < len(p.overlay_images):\n",
        "                    overlay = p.overlay_images[i]\n",
        "\n",
        "                    if p.paste_to is not None:\n",
        "                        x, y, w, h = p.paste_to\n",
        "                        base_image = Image.new('RGBA', (overlay.width, overlay.height))\n",
        "                        image = resize_image(1, image, w, h)\n",
        "                        base_image.paste(image, (x, y))\n",
        "                        image = base_image\n",
        "\n",
        "                    image = image.convert('RGBA')\n",
        "                    image.alpha_composite(overlay)\n",
        "                    image = image.convert('RGB')\n",
        "\n",
        "                # if opts.samples_save and not p.do_not_save_samples:\n",
        "                #     images.save_image(image, p.outpath_samples, \"\", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i))\n",
        "\n",
        "                output_images.append(image)\n",
        "\n",
        "            state.nextjob()\n",
        "\n",
        "        # unwanted_grid_because_of_img_count = len(output_images) < 2 and opts.grid_only_if_multiple\n",
        "        if not p.do_not_save_grid: # and not unwanted_grid_because_of_img_count:\n",
        "            # return_grid = opts.return_grid\n",
        "\n",
        "            grid = image_grid(output_images, p.batch_size, 2)\n",
        "\n",
        "            # if return_grid:\n",
        "            output_images.insert(0, grid)\n",
        "\n",
        "            # if opts.grid_save:\n",
        "            #     images.save_image(grid, p.outpath_grids, \"grid\", seed, all_prompts[0], opts.grid_format, info=infotext(), short_filename=not opts.grid_extended_filename)\n",
        "\n",
        "    torch_gc()\n",
        "    return Processed(p, output_images, seed, infotext())\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0fFLFqfbvJzM"
      },
      "outputs": [],
      "source": [
        "class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):\n",
        "    sampler = None\n",
        "\n",
        "    def init(self, seed):\n",
        "        self.sampler = samplers[self.sampler_index].constructor(self.sd_model)\n",
        "\n",
        "    def sample(self, x, conditioning, unconditional_conditioning):\n",
        "        samples_ddim = self.sampler.sample(self, x, conditioning, unconditional_conditioning)\n",
        "        return samples_ddim\n",
        "\n",
        "\n",
        "def get_crop_region(mask, pad=0):\n",
        "    h, w = mask.shape\n",
        "\n",
        "    crop_left = 0\n",
        "    for i in range(w):\n",
        "        if not (mask[:, i] == 0).all():\n",
        "            break\n",
        "        crop_left += 1\n",
        "\n",
        "    crop_right = 0\n",
        "    for i in reversed(range(w)):\n",
        "        if not (mask[:, i] == 0).all():\n",
        "            break\n",
        "        crop_right += 1\n",
        "\n",
        "    crop_top = 0\n",
        "    for i in range(h):\n",
        "        if not (mask[i] == 0).all():\n",
        "            break\n",
        "        crop_top += 1\n",
        "\n",
        "    crop_bottom = 0\n",
        "    for i in reversed(range(h)):\n",
        "        if not (mask[i] == 0).all():\n",
        "            break\n",
        "        crop_bottom += 1\n",
        "\n",
        "    return (\n",
        "        int(max(crop_left-pad, 0)),\n",
        "        int(max(crop_top-pad, 0)),\n",
        "        int(min(w - crop_right + pad, w)),\n",
        "        int(min(h - crop_bottom + pad, h))\n",
        "    )\n",
        "\n",
        "\n",
        "def fill(image, mask):\n",
        "    image_mod = Image.new('RGBA', (image.width, image.height))\n",
        "\n",
        "    image_masked = Image.new('RGBa', (image.width, image.height))\n",
        "    image_masked.paste(image.convert(\"RGBA\").convert(\"RGBa\"), mask=ImageOps.invert(mask.convert('L')))\n",
        "\n",
        "    image_masked = image_masked.convert('RGBa')\n",
        "\n",
        "    for radius, repeats in [(256, 1), (64, 1), (16, 2), (4, 4), (2, 2), (0, 1)]:\n",
        "        blurred = image_masked.filter(ImageFilter.GaussianBlur(radius)).convert('RGBA')\n",
        "        for _ in range(repeats):\n",
        "            image_mod.alpha_composite(blurred)\n",
        "\n",
        "    return image_mod.convert(\"RGB\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G6DxNrJSvMgi"
      },
      "outputs": [],
      "source": [
        "class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):\n",
        "    sampler = None\n",
        "\n",
        "    def __init__(self, init_images=None, resize_mode=0, denoising_strength=0.75, mask=None, mask_blur=4, inpainting_fill=0, inpaint_full_res=True, inpainting_mask_invert=0, **kwargs):\n",
        "        super().__init__(**kwargs)\n",
        "\n",
        "        self.init_images = init_images\n",
        "        self.resize_mode: int = resize_mode\n",
        "        self.denoising_strength: float = denoising_strength\n",
        "        self.init_latent = None\n",
        "        self.image_mask = mask\n",
        "        #self.image_unblurred_mask = None\n",
        "        self.latent_mask = None\n",
        "        self.mask_for_overlay = None\n",
        "        self.mask_blur = mask_blur\n",
        "        self.inpainting_fill = inpainting_fill\n",
        "        self.inpaint_full_res = inpaint_full_res\n",
        "        self.inpainting_mask_invert = inpainting_mask_invert\n",
        "        self.mask = None\n",
        "        self.nmask = None\n",
        "\n",
        "    def init(self, seed):\n",
        "        self.sampler = samplers_for_img2img[self.sampler_index].constructor(self.sd_model)\n",
        "        crop_region = None\n",
        "\n",
        "        if self.image_mask is not None:\n",
        "            self.image_mask = self.image_mask.convert('L')\n",
        "\n",
        "            if self.inpainting_mask_invert:\n",
        "                self.image_mask = ImageOps.invert(self.image_mask)\n",
        "\n",
        "            #self.image_unblurred_mask = self.image_mask\n",
        "\n",
        "            if self.mask_blur > 0:\n",
        "                self.image_mask = self.image_mask.filter(ImageFilter.GaussianBlur(self.mask_blur))\n",
        "\n",
        "            if self.inpaint_full_res:\n",
        "                self.mask_for_overlay = self.image_mask\n",
        "                mask = self.image_mask.convert('L')\n",
        "                crop_region = get_crop_region(np.array(mask), False)\n",
        "                x1, y1, x2, y2 = crop_region\n",
        "\n",
        "                mask = mask.crop(crop_region)\n",
        "                self.image_mask = resize_image(2, mask, self.width, self.height)\n",
        "                self.paste_to = (x1, y1, x2-x1, y2-y1)\n",
        "            else:\n",
        "                self.image_mask = resize_image(self.resize_mode, self.image_mask, self.width, self.height)\n",
        "                np_mask = np.array(self.image_mask)\n",
        "                np_mask = np.clip((np_mask.astype(np.float)) * 2, 0, 255).astype(np.uint8)\n",
        "                self.mask_for_overlay = Image.fromarray(np_mask)\n",
        "\n",
        "            self.overlay_images = []\n",
        "\n",
        "        latent_mask = self.latent_mask if self.latent_mask is not None else self.image_mask\n",
        "\n",
        "        imgs = []\n",
        "        for img in self.init_images:\n",
        "            image = img.convert(\"RGB\")\n",
        "\n",
        "            if crop_region is None:\n",
        "                image = resize_image(self.resize_mode, image, self.width, self.height)\n",
        "\n",
        "            if self.image_mask is not None:\n",
        "                image_masked = Image.new('RGBa', (image.width, image.height))\n",
        "                image_masked.paste(image.convert(\"RGBA\").convert(\"RGBa\"), mask=ImageOps.invert(self.mask_for_overlay.convert('L')))\n",
        "\n",
        "                self.overlay_images.append(image_masked.convert('RGBA'))\n",
        "\n",
        "            if crop_region is not None:\n",
        "                image = image.crop(crop_region)\n",
        "                image = resize_image(2, image, self.width, self.height)\n",
        "\n",
        "            if self.image_mask is not None:\n",
        "                if self.inpainting_fill != 1:\n",
        "                    image = fill(image, latent_mask)\n",
        "\n",
        "            image = np.array(image).astype(np.float32) / 255.0\n",
        "            image = np.moveaxis(image, 2, 0)\n",
        "\n",
        "            imgs.append(image)\n",
        "\n",
        "        if len(imgs) == 1:\n",
        "            batch_images = np.expand_dims(imgs[0], axis=0).repeat(self.batch_size, axis=0)\n",
        "            if self.overlay_images is not None:\n",
        "                self.overlay_images = self.overlay_images * self.batch_size\n",
        "        elif len(imgs) <= self.batch_size:\n",
        "            self.batch_size = len(imgs)\n",
        "            batch_images = np.array(imgs)\n",
        "        else:\n",
        "            raise RuntimeError(f\"bad number of images passed: {len(imgs)}; expecting {self.batch_size} or less\")\n",
        "\n",
        "        image = torch.from_numpy(batch_images)\n",
        "        image = 2. * image - 1.\n",
        "        image = image.to(device)\n",
        "\n",
        "        self.init_latent = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image))\n",
        "\n",
        "        if self.image_mask is not None:\n",
        "            init_mask = latent_mask\n",
        "            latmask = init_mask.convert('RGB').resize((self.init_latent.shape[3], self.init_latent.shape[2]))\n",
        "            latmask = np.moveaxis(np.array(latmask, dtype=np.float64), 2, 0) / 255\n",
        "            latmask = latmask[0]\n",
        "            latmask = np.around(latmask)\n",
        "            latmask = np.tile(latmask[None], (4, 1, 1))\n",
        "\n",
        "            self.mask = torch.asarray(1.0 - latmask).to(device).type(self.sd_model.dtype)\n",
        "            self.nmask = torch.asarray(latmask).to(device).type(self.sd_model.dtype)\n",
        "\n",
        "            if self.inpainting_fill == 2:\n",
        "                self.init_latent = self.init_latent * self.mask + create_random_tensors(self.init_latent.shape[1:], [seed + x + 1 for x in range(self.init_latent.shape[0])]) * self.nmask\n",
        "            elif self.inpainting_fill == 3:\n",
        "                self.init_latent = self.init_latent * self.mask\n",
        "\n",
        "    def sample(self, x, conditioning, unconditional_conditioning):\n",
        "        samples = self.sampler.sample_img2img(self, self.init_latent, x, conditioning, unconditional_conditioning)\n",
        "\n",
        "        if self.mask is not None:\n",
        "            samples = samples * self.nmask + self.init_latent * self.mask\n",
        "\n",
        "        return samples\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "##### scripts.py"
      ],
      "metadata": {
        "id": "U0vz6D412oSY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Script:\n",
        "    filename = None\n",
        "    args_from = None\n",
        "    args_to = None\n",
        "\n",
        "    def title(self):\n",
        "        raise NotImplementedError()\n",
        "\n",
        "    def ui(self, is_img2img):\n",
        "        pass\n",
        "\n",
        "    def show(self, is_img2img):\n",
        "        return True\n",
        "\n",
        "    def run(self, *args):\n",
        "        raise NotImplementedError()\n",
        "\n",
        "    def describe(self):\n",
        "        return \"\"\n",
        "\n",
        "\n",
        "scripts_data = []\n",
        "\n",
        "\n",
        "def load_scripts(basedir):\n",
        "    if not os.path.exists(basedir):\n",
        "        return\n",
        "\n",
        "    for filename in os.listdir(basedir):\n",
        "        path = os.path.join(basedir, filename)\n",
        "\n",
        "        if not os.path.isfile(path):\n",
        "            continue\n",
        "\n",
        "        with open(path, \"r\", encoding=\"utf8\") as file:\n",
        "            text = file.read()\n",
        "\n",
        "        try:\n",
        "            from types import ModuleType\n",
        "            compiled = compile(text, path, 'exec')\n",
        "            module = ModuleType(filename)\n",
        "            exec(compiled, module.__dict__)\n",
        "\n",
        "            for key, script_class in module.__dict__.items():\n",
        "                if type(script_class) == type and issubclass(script_class, Script):\n",
        "                    scripts_data.append((script_class, path))\n",
        "\n",
        "        except Exception:\n",
        "            print(f\"Error loading script: {filename}\", file=sys.stderr)\n",
        "            print(traceback.format_exc(), file=sys.stderr)\n",
        "\n",
        "\n",
        "def wrap_call(func, filename, funcname, *args, default=None, **kwargs):\n",
        "    try:\n",
        "        res = func(*args, **kwargs)\n",
        "        return res\n",
        "    except Exception:\n",
        "        print(f\"Error calling: {filename}/{funcname}\", file=sys.stderr)\n",
        "        print(traceback.format_exc(), file=sys.stderr)\n",
        "\n",
        "    return default\n",
        "\n",
        "\n",
        "class ScriptRunner:\n",
        "    def __init__(self):\n",
        "        self.scripts = []\n",
        "\n",
        "    def setup_ui(self, is_img2img):\n",
        "        for script_class, path in scripts_data:\n",
        "            script = script_class()\n",
        "            script.filename = path\n",
        "\n",
        "            if not script.show(is_img2img):\n",
        "                continue\n",
        "\n",
        "            self.scripts.append(script)\n",
        "\n",
        "        titles = [wrap_call(script.title, script.filename, \"title\") or f\"{script.filename} [error]\" for script in self.scripts]\n",
        "\n",
        "        dropdown = gr.Dropdown(label=\"Script\", choices=[\"None\"] + titles, value=\"None\", type=\"index\")\n",
        "        inputs = [dropdown]\n",
        "\n",
        "        for script in self.scripts:\n",
        "            script.args_from = len(inputs)\n",
        "\n",
        "            controls = wrap_call(script.ui, script.filename, \"ui\", is_img2img)\n",
        "\n",
        "            if controls is None:\n",
        "                continue\n",
        "\n",
        "            for control in controls:\n",
        "                control.visible = False\n",
        "\n",
        "            inputs += controls\n",
        "            script.args_to = len(inputs)\n",
        "\n",
        "        def select_script(script_index):\n",
        "            if 0 < script_index <= len(self.scripts):\n",
        "                script = self.scripts[script_index-1]\n",
        "                args_from = script.args_from\n",
        "                args_to = script.args_to\n",
        "            else:\n",
        "                args_from = 0\n",
        "                args_to = 0\n",
        "\n",
        "            return [ui.gr_show(True if i == 0 else args_from <= i < args_to) for i in range(len(inputs))]\n",
        "\n",
        "        dropdown.change(\n",
        "            fn=select_script,\n",
        "            inputs=[dropdown],\n",
        "            outputs=inputs\n",
        "        )\n",
        "\n",
        "        return inputs\n",
        "\n",
        "\n",
        "    def run(self, p: StableDiffusionProcessing, *args):\n",
        "        script_index = args[0]\n",
        "\n",
        "        if script_index == 0:\n",
        "            return None\n",
        "\n",
        "        script = self.scripts[script_index-1]\n",
        "\n",
        "        if script is None:\n",
        "            return None\n",
        "\n",
        "        script_args = args[script.args_from:script.args_to]\n",
        "        processed = script.run(p, *script_args)\n",
        "\n",
        "        shared.total_tqdm.clear()\n",
        "\n",
        "        return processed\n",
        "\n",
        "\n",
        "scripts_txt2img = ScriptRunner()\n",
        "scripts_img2img = ScriptRunner()"
      ],
      "metadata": {
        "id": "ebevDdEe2n1c"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "36c5Zztlo8vh"
      },
      "source": [
        "##### images.py"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yfIxxdfQq2-S"
      },
      "outputs": [],
      "source": [
        "LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vRLqNhD4q3hg"
      },
      "outputs": [],
      "source": [
        "def image_grid(imgs, batch_size=1, rows=2):\n",
        "    if rows is None:\n",
        "      rows = math.sqrt(len(imgs))\n",
        "      rows = round(rows)\n",
        "\n",
        "    cols = math.ceil(len(imgs) / rows)\n",
        "\n",
        "    w, h = imgs[0].size\n",
        "    grid = Image.new('RGB', size=(cols * w, rows * h), color='black')\n",
        "\n",
        "    for i, img in enumerate(imgs):\n",
        "        grid.paste(img, box=(i % cols * w, i // cols * h))\n",
        "\n",
        "    return grid\n",
        "\n",
        "Grid = namedtuple(\"Grid\", [\"tiles\", \"tile_w\", \"tile_h\", \"image_w\", \"image_h\", \"overlap\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gE9NPfjpq8eH"
      },
      "outputs": [],
      "source": [
        "def split_grid(image, tile_w=512, tile_h=512, overlap=64):\n",
        "    w = image.width\n",
        "    h = image.height\n",
        "\n",
        "    non_overlap_width = tile_w - overlap\n",
        "    non_overlap_height = tile_h - overlap\n",
        "\n",
        "    cols = math.ceil((w - overlap) / non_overlap_width)\n",
        "    rows = math.ceil((h - overlap) / non_overlap_height)\n",
        "\n",
        "    dx = (w - tile_w) / (cols-1) if cols > 1 else 0\n",
        "    dy = (h - tile_h) / (rows-1) if rows > 1 else 0\n",
        "\n",
        "    grid = Grid([], tile_w, tile_h, w, h, overlap)\n",
        "    for row in range(rows):\n",
        "        row_images = []\n",
        "\n",
        "        y = int(row * dy)\n",
        "\n",
        "        if y + tile_h >= h:\n",
        "            y = h - tile_h\n",
        "\n",
        "        for col in range(cols):\n",
        "            x = int(col * dx)\n",
        "\n",
        "            if x+tile_w >= w:\n",
        "                x = w - tile_w\n",
        "\n",
        "            tile = image.crop((x, y, x + tile_w, y + tile_h))\n",
        "\n",
        "            row_images.append([x, tile_w, tile])\n",
        "\n",
        "        grid.tiles.append([y, tile_h, row_images])\n",
        "\n",
        "    return grid\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BHh5RVktrA4D"
      },
      "outputs": [],
      "source": [
        "def combine_grid(grid):\n",
        "    def make_mask_image(r):\n",
        "        r = r * 255 / grid.overlap\n",
        "        r = r.astype(np.uint8)\n",
        "        return Image.fromarray(r, 'L')\n",
        "\n",
        "    mask_w = make_mask_image(np.arange(grid.overlap, dtype=np.float32).reshape((1, grid.overlap)).repeat(grid.tile_h, axis=0))\n",
        "    mask_h = make_mask_image(np.arange(grid.overlap, dtype=np.float32).reshape((grid.overlap, 1)).repeat(grid.image_w, axis=1))\n",
        "\n",
        "    combined_image = Image.new(\"RGB\", (grid.image_w, grid.image_h))\n",
        "    for y, h, row in grid.tiles:\n",
        "        combined_row = Image.new(\"RGB\", (grid.image_w, h))\n",
        "        for x, w, tile in row:\n",
        "            if x == 0:\n",
        "                combined_row.paste(tile, (0, 0))\n",
        "                continue\n",
        "\n",
        "            combined_row.paste(tile.crop((0, 0, grid.overlap, h)), (x, 0), mask=mask_w)\n",
        "            combined_row.paste(tile.crop((grid.overlap, 0, w, h)), (x + grid.overlap, 0))\n",
        "\n",
        "        if y == 0:\n",
        "            combined_image.paste(combined_row, (0, 0))\n",
        "            continue\n",
        "\n",
        "        combined_image.paste(combined_row.crop((0, 0, combined_row.width, grid.overlap)), (0, y), mask=mask_h)\n",
        "        combined_image.paste(combined_row.crop((0, grid.overlap, combined_row.width, h)), (0, y + grid.overlap))\n",
        "\n",
        "    return combined_image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rzLJ48hdrE95"
      },
      "outputs": [],
      "source": [
        "class GridAnnotation:\n",
        "    def __init__(self, text='', is_active=True):\n",
        "        self.text = text\n",
        "        self.is_active = is_active\n",
        "        self.size = None\n",
        "\n",
        "\n",
        "def draw_grid_annotations(im, width, height, hor_texts, ver_texts):\n",
        "    def wrap(drawing, text, font, line_length):\n",
        "        lines = ['']\n",
        "        for word in text.split():\n",
        "            line = f'{lines[-1]} {word}'.strip()\n",
        "            if drawing.textlength(line, font=font) <= line_length:\n",
        "                lines[-1] = line\n",
        "            else:\n",
        "                lines.append(word)\n",
        "        return lines\n",
        "\n",
        "    def draw_texts(drawing, draw_x, draw_y, lines):\n",
        "        for i, line in enumerate(lines):\n",
        "            drawing.multiline_text((draw_x, draw_y + line.size[1] / 2), line.text, font=fnt, fill=color_active if line.is_active else color_inactive, anchor=\"mm\", align=\"center\")\n",
        "\n",
        "            if not line.is_active:\n",
        "                drawing.line((draw_x - line.size[0]//2, draw_y + line.size[1]//2, draw_x + line.size[0]//2, draw_y + line.size[1]//2), fill=color_inactive, width=4)\n",
        "\n",
        "            draw_y += line.size[1] + line_spacing\n",
        "\n",
        "    fontsize = (width + height) // 25\n",
        "    line_spacing = fontsize // 2\n",
        "    fnt = ImageFont.truetype(\"/content/stable-diffusion-webui/NotoSansJP-Bold.otf\", fontsize)\n",
        "    color_active = (0, 0, 0)\n",
        "    color_inactive = (153, 153, 153)\n",
        "\n",
        "    pad_left = width * 3 // 4 if len(ver_texts) > 0 else 0\n",
        "\n",
        "    cols = im.width // width\n",
        "    rows = im.height // height\n",
        "\n",
        "    assert cols == len(hor_texts), f'bad number of horizontal texts: {len(hor_texts)}; must be {cols}'\n",
        "    assert rows == len(ver_texts), f'bad number of vertical texts: {len(ver_texts)}; must be {rows}'\n",
        "\n",
        "    calc_img = Image.new(\"RGB\", (1, 1), \"white\")\n",
        "    calc_d = ImageDraw.Draw(calc_img)\n",
        "\n",
        "    for texts, allowed_width in zip(hor_texts + ver_texts, [width] * len(hor_texts) + [pad_left] * len(ver_texts)):\n",
        "        items = [] + texts\n",
        "        texts.clear()\n",
        "\n",
        "        for line in items:\n",
        "            wrapped = wrap(calc_d, line.text, fnt, allowed_width)\n",
        "            texts += [GridAnnotation(x, line.is_active) for x in wrapped]\n",
        "\n",
        "        for line in texts:\n",
        "            bbox = calc_d.multiline_textbbox((0, 0), line.text, font=fnt)\n",
        "            line.size = (bbox[2] - bbox[0], bbox[3] - bbox[1])\n",
        "\n",
        "    hor_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing for lines in hor_texts]\n",
        "    ver_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing * len(lines) for lines in ver_texts]\n",
        "\n",
        "    pad_top = max(hor_text_heights) + line_spacing * 2\n",
        "\n",
        "    result = Image.new(\"RGB\", (im.width + pad_left, im.height + pad_top), \"white\")\n",
        "    result.paste(im, (pad_left, pad_top))\n",
        "\n",
        "    d = ImageDraw.Draw(result)\n",
        "\n",
        "    for col in range(cols):\n",
        "        x = pad_left + width * col + width / 2\n",
        "        y = pad_top / 2 - hor_text_heights[col] / 2\n",
        "\n",
        "        draw_texts(d, x, y, hor_texts[col])\n",
        "\n",
        "    for row in range(rows):\n",
        "        x = pad_left / 2\n",
        "        y = pad_top + height * row + height / 2 - ver_text_heights[row] / 2\n",
        "\n",
        "        draw_texts(d, x, y, ver_texts[row])\n",
        "\n",
        "    return result"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vNqU11iGrFl1"
      },
      "outputs": [],
      "source": [
        "def draw_prompt_matrix(im, width, height, all_prompts):\n",
        "    prompts = all_prompts[1:]\n",
        "    boundary = math.ceil(len(prompts) / 2)\n",
        "\n",
        "    prompts_horiz = prompts[:boundary]\n",
        "    prompts_vert = prompts[boundary:]\n",
        "\n",
        "    hor_texts = [[GridAnnotation(x, is_active=pos & (1 << i) != 0) for i, x in enumerate(prompts_horiz)] for pos in range(1 << len(prompts_horiz))]\n",
        "    ver_texts = [[GridAnnotation(x, is_active=pos & (1 << i) != 0) for i, x in enumerate(prompts_vert)] for pos in range(1 << len(prompts_vert))]\n",
        "\n",
        "    return draw_grid_annotations(im, width, height, hor_texts, ver_texts)\n",
        "\n",
        "\n",
        "def resize_image(resize_mode, im, width, height):\n",
        "    if resize_mode == 0:\n",
        "        res = im.resize((width, height), resample=LANCZOS)\n",
        "    elif resize_mode == 1:\n",
        "        ratio = width / height\n",
        "        src_ratio = im.width / im.height\n",
        "\n",
        "        src_w = width if ratio > src_ratio else im.width * height // im.height\n",
        "        src_h = height if ratio <= src_ratio else im.height * width // im.width\n",
        "\n",
        "        resized = im.resize((src_w, src_h), resample=LANCZOS)\n",
        "        res = Image.new(\"RGB\", (width, height))\n",
        "        res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))\n",
        "    else:\n",
        "        ratio = width / height\n",
        "        src_ratio = im.width / im.height\n",
        "\n",
        "        src_w = width if ratio < src_ratio else im.width * height // im.height\n",
        "        src_h = height if ratio >= src_ratio else im.height * width // im.width\n",
        "\n",
        "        resized = im.resize((src_w, src_h), resample=LANCZOS)\n",
        "        res = Image.new(\"RGB\", (width, height))\n",
        "        res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))\n",
        "\n",
        "        if ratio < src_ratio:\n",
        "            fill_height = height // 2 - src_h // 2\n",
        "            res.paste(resized.resize((width, fill_height), box=(0, 0, width, 0)), box=(0, 0))\n",
        "            res.paste(resized.resize((width, fill_height), box=(0, resized.height, width, resized.height)), box=(0, fill_height + src_h))\n",
        "        elif ratio > src_ratio:\n",
        "            fill_width = width // 2 - src_w // 2\n",
        "            res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0))\n",
        "            res.paste(resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)), box=(fill_width + src_w, 0))\n",
        "\n",
        "    return res"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oacrnMXqrLVy"
      },
      "outputs": [],
      "source": [
        "class Upscaler:\n",
        "    name = \"Lanczos\"\n",
        "\n",
        "    def do_upscale(self, img):\n",
        "        return img\n",
        "\n",
        "    def upscale(self, img, w, h):\n",
        "        for i in range(3):\n",
        "            if img.width >= w and img.height >= h:\n",
        "                break\n",
        "\n",
        "            img = self.do_upscale(img)\n",
        "\n",
        "        if img.width != w or img.height != h:\n",
        "            img = img.resize((w, h), resample=LANCZOS)\n",
        "\n",
        "        return img\n",
        "\n",
        "\n",
        "class UpscalerNone(Upscaler):\n",
        "    name = \"None\"\n",
        "\n",
        "    def upscale(self, img, w, h):\n",
        "        return img\n",
        "\n",
        "sd_upscalers.append(UpscalerNone())\n",
        "sd_upscalers.append(Upscaler())"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def make_init_masked(init_image: str, matte: str, blur: int, invert: bool):\n",
        "  transparency = Image.new('RGBA', (width, height))\n",
        "  mask_image = Image.open(matte)\n",
        "  if invert:\n",
        "    mask_image = ImageOps.invert(mask_image)\n",
        "  mask_blurred = mask_image.filter(ImageFilter.GaussianBlur(blur)).convert('L')\n",
        "  masked_image = Image.composite(transparency, Image.open(init_image), mask_blurred)\n",
        "  return masked_image"
      ],
      "metadata": {
        "id": "bnBvsGSYJY-b"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qywis6QspZwv"
      },
      "source": [
        "##### lowvram.py (Done)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v1dCECY0qukX"
      },
      "outputs": [],
      "source": [
        "module_in_gpu = None\n",
        "\n",
        "cpu = torch.device(\"cpu\")\n",
        "\n",
        "device = gpu = torch.device(\"cuda\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W3EQzeTsq9Ut"
      },
      "outputs": [],
      "source": [
        "def setup_for_low_vram(sd_model, use_medvram):\n",
        "    parents = {}\n",
        "\n",
        "    def send_me_to_gpu(module, _):\n",
        "        \"\"\"send this module to GPU; send whatever tracked module was previous in GPU to CPU;\n",
        "        we add this as forward_pre_hook to a lot of modules and this way all but one of them will\n",
        "        be in CPU\n",
        "        \"\"\"\n",
        "        global module_in_gpu\n",
        "\n",
        "        module = parents.get(module, module)\n",
        "\n",
        "        if module_in_gpu == module:\n",
        "            return\n",
        "\n",
        "        if module_in_gpu is not None:\n",
        "            module_in_gpu.to(cpu)\n",
        "\n",
        "        module.to(gpu)\n",
        "        module_in_gpu = module\n",
        "\n",
        "    # see below for register_forward_pre_hook;\n",
        "    # first_stage_model does not use forward(), it uses encode/decode, so register_forward_pre_hook is\n",
        "    # useless here, and we just replace those methods\n",
        "    def first_stage_model_encode_wrap(self, encoder, x):\n",
        "        send_me_to_gpu(self, None)\n",
        "        return encoder(x)\n",
        "\n",
        "    def first_stage_model_decode_wrap(self, decoder, z):\n",
        "        send_me_to_gpu(self, None)\n",
        "        return decoder(z)\n",
        "\n",
        "    # remove three big modules, cond, first_stage, and unet from the model and then\n",
        "    # send the model to GPU. Then put modules back. the modules will be in CPU.\n",
        "    stored = sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.model\n",
        "    sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.model = None, None, None\n",
        "    sd_model.to(device)\n",
        "    sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.model = stored\n",
        "\n",
        "    # register hooks for those the first two models\n",
        "    sd_model.cond_stage_model.transformer.register_forward_pre_hook(send_me_to_gpu)\n",
        "    sd_model.first_stage_model.register_forward_pre_hook(send_me_to_gpu)\n",
        "    sd_model.first_stage_model.encode = lambda x, en=sd_model.first_stage_model.encode: first_stage_model_encode_wrap(sd_model.first_stage_model, en, x)\n",
        "    sd_model.first_stage_model.decode = lambda z, de=sd_model.first_stage_model.decode: first_stage_model_decode_wrap(sd_model.first_stage_model, de, z)\n",
        "    parents[sd_model.cond_stage_model.transformer] = sd_model.cond_stage_model\n",
        "\n",
        "    if use_medvram:\n",
        "        sd_model.model.register_forward_pre_hook(send_me_to_gpu)\n",
        "    else:\n",
        "        diff_model = sd_model.model.diffusion_model\n",
        "\n",
        "        # the third remaining model is still too big for 4 GB, so we also do the same for its submodules\n",
        "        # so that only one of them is in GPU at a time\n",
        "        stored = diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed\n",
        "        diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed = None, None, None, None\n",
        "        sd_model.model.to(device)\n",
        "        diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed = stored\n",
        "\n",
        "        # install hooks for bits of third model\n",
        "        diff_model.time_embed.register_forward_pre_hook(send_me_to_gpu)\n",
        "        for block in diff_model.input_blocks:\n",
        "            block.register_forward_pre_hook(send_me_to_gpu)\n",
        "        diff_model.middle_block.register_forward_pre_hook(send_me_to_gpu)\n",
        "        for block in diff_model.output_blocks:\n",
        "            block.register_forward_pre_hook(send_me_to_gpu)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iB5fIbQ7pbZD"
      },
      "source": [
        "##### txt2img.py (Done)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VY8xjBOWrgHb"
      },
      "outputs": [],
      "source": [
        "def txt2img(prompt: str, negative_prompt: str, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, height: int, width: int, *args):\n",
        "    p = StableDiffusionProcessingTxt2Img(\n",
        "        sd_model=sd_model,\n",
        "        # outpath_samples=opts.outdir_samples or opts.outdir_txt2img_samples,\n",
        "        # outpath_grids=opts.outdir_grids or opts.outdir_txt2img_grids,\n",
        "        prompt=prompt,\n",
        "        negative_prompt=negative_prompt,\n",
        "        seed=seed,\n",
        "        sampler_index=sampler_index,\n",
        "        batch_size=batch_size,\n",
        "        n_iter=n_iter,\n",
        "        steps=steps,\n",
        "        cfg_scale=cfg_scale,\n",
        "        width=width,\n",
        "        height=height,\n",
        "        restore_faces=restore_faces,\n",
        "        tiling=tiling,\n",
        "    )\n",
        "\n",
        "    processed = scripts_txt2img.run(p, *args)\n",
        "\n",
        "    if processed is not None:\n",
        "        pass\n",
        "    else:\n",
        "        processed = process_images(p)\n",
        "\n",
        "    return processed.images, processed.seed, None #plaintext_to_html(processed.info)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l_kIaCEEpcmX"
      },
      "source": [
        "##### img2img"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZGwlntkClnZ9"
      },
      "outputs": [],
      "source": [
        "def img2img(prompt: str, init_img, init_img_with_mask, mask, steps: int, sampler_index: int, mask_blur: int, inpainting_fill: int, restore_faces: bool, tiling: bool, mode: int, n_iter: int, batch_size: int, cfg_scale: float, denoising_strength: float, denoising_strength_change_factor: float, seed: int, height: int, width: int, resize_mode: int, upscaler_index: str, upscale_overlap: int, inpaint_full_res: bool, inpainting_mask_invert: int, *args):\n",
        "    is_inpaint = mode == 1\n",
        "    is_loopback = mode == 2\n",
        "    is_upscale = mode == 3\n",
        "    print(denoising_strength)\n",
        "    if is_inpaint:\n",
        "        image = init_img_with_mask\n",
        "        alpha_mask = ImageOps.invert(image.split()[-1]).convert('L').point(lambda x: 255 if x > 0 else 0, mode='1')\n",
        "        mask = ImageChops.lighter(alpha_mask, mask.convert('L')).convert('RGBA')\n",
        "        image = image.convert('RGB')\n",
        "    else:\n",
        "        image = init_img\n",
        "        mask = None\n",
        "\n",
        "    assert 0. <= denoising_strength <= 1., 'can only work with strength in [0.0, 1.0]'\n",
        "\n",
        "    p = StableDiffusionProcessingImg2Img(\n",
        "        sd_model=sd_model,\n",
        "        # outpath_samples=opts.outdir_samples or opts.outdir_img2img_samples,\n",
        "        # outpath_grids=opts.outdir_grids or opts.outdir_img2img_grids,\n",
        "        prompt=prompt,\n",
        "        seed=seed,\n",
        "        sampler_index=sampler_index,\n",
        "        batch_size=batch_size,\n",
        "        n_iter=n_iter,\n",
        "        steps=steps,\n",
        "        cfg_scale=cfg_scale,\n",
        "        width=width,\n",
        "        height=height,\n",
        "        restore_faces=restore_faces,\n",
        "        tiling=tiling,\n",
        "        init_images=[image],\n",
        "        mask=mask,\n",
        "        mask_blur=mask_blur,\n",
        "        inpainting_fill=inpainting_fill,\n",
        "        resize_mode=resize_mode,\n",
        "        denoising_strength=denoising_strength,\n",
        "        inpaint_full_res=inpaint_full_res,\n",
        "        inpainting_mask_invert=inpainting_mask_invert,\n",
        "        extra_generation_params={\n",
        "            \"Denoising strength\": denoising_strength,\n",
        "            \"Denoising strength change factor\": denoising_strength_change_factor\n",
        "        }\n",
        "    )\n",
        "\n",
        "    if is_loopback:\n",
        "        output_images, info = None, None\n",
        "        history = []\n",
        "        initial_seed = None\n",
        "        initial_info = None\n",
        "\n",
        "        state.job_count = n_iter\n",
        "\n",
        "        do_color_correction = False\n",
        "        try:\n",
        "            from skimage import exposure\n",
        "            do_color_correction = True\n",
        "        except:\n",
        "            print(\"Install scikit-image to perform color correction on loopback\")\n",
        "\n",
        "\n",
        "        for i in range(n_iter):\n",
        "\n",
        "            if do_color_correction and i == 0:\n",
        "                correction_target = cv2.cvtColor(np.asarray(init_img.copy()), cv2.COLOR_RGB2LAB)\n",
        "\n",
        "            p.n_iter = 1\n",
        "            p.batch_size = 1\n",
        "            p.do_not_save_grid = True\n",
        "\n",
        "            state.job = f\"Batch {i + 1} out of {n_iter}\"\n",
        "            processed = process_images(p)\n",
        "\n",
        "            if initial_seed is None:\n",
        "                initial_seed = processed.seed\n",
        "                initial_info = processed.info\n",
        "            \n",
        "            init_img = processed.images[0]\n",
        "\n",
        "            if do_color_correction and correction_target is not None:\n",
        "                init_img = Image.fromarray(cv2.cvtColor(exposure.match_histograms(\n",
        "                    cv2.cvtColor(\n",
        "                        np.asarray(init_img),\n",
        "                        cv2.COLOR_RGB2LAB\n",
        "                    ),\n",
        "                    correction_target,\n",
        "                    channel_axis=2\n",
        "                ), cv2.COLOR_LAB2RGB).astype(\"uint8\"))\n",
        "\n",
        "            p.init_images = [init_img]\n",
        "            p.seed = processed.seed + 1\n",
        "            p.denoising_strength = min(max(p.denoising_strength * denoising_strength_change_factor, 0.1), 1)\n",
        "            history.append(processed.images[0])\n",
        "\n",
        "        grid = images.image_grid(history, batch_size, rows=1)\n",
        "\n",
        "        # images.save_image(grid, p.outpath_grids, \"grid\", initial_seed, prompt, 'png', info=info, short_filename= None)\n",
        "\n",
        "        processed = Processed(p, history, initial_seed, initial_info)\n",
        "\n",
        "    elif is_upscale:\n",
        "        initial_seed = None\n",
        "        initial_info = None\n",
        "\n",
        "        upscaler = sd_upscalers[upscaler_index]\n",
        "        img = upscaler.upscale(init_img, init_img.width * 2, init_img.height * 2)\n",
        "\n",
        "        processing.torch_gc()\n",
        "\n",
        "        grid = images.split_grid(img, tile_w=width, tile_h=height, overlap=upscale_overlap)\n",
        "\n",
        "        p.n_iter = 1\n",
        "        p.do_not_save_grid = True\n",
        "        p.do_not_save_samples = True\n",
        "\n",
        "        work = []\n",
        "        work_results = []\n",
        "\n",
        "        for y, h, row in grid.tiles:\n",
        "            for tiledata in row:\n",
        "                work.append(tiledata[2])\n",
        "\n",
        "        batch_count = math.ceil(len(work) / p.batch_size)\n",
        "        print(f\"SD upscaling will process a total of {len(work)} images tiled as {len(grid.tiles[0][2])}x{len(grid.tiles)} in a total of {batch_count} batches.\")\n",
        "\n",
        "        state.job_count = batch_count\n",
        "\n",
        "        for i in range(batch_count):\n",
        "            p.init_images = work[i*p.batch_size:(i+1)*p.batch_size]\n",
        "\n",
        "            state.job = f\"Batch {i + 1} out of {batch_count}\"\n",
        "            processed = process_images(p)\n",
        "\n",
        "            if initial_seed is None:\n",
        "                initial_seed = processed.seed\n",
        "                initial_info = processed.info\n",
        "\n",
        "            p.seed = processed.seed + 1\n",
        "            work_results += processed.images\n",
        "\n",
        "        image_index = 0\n",
        "        for y, h, row in grid.tiles:\n",
        "            for tiledata in row:\n",
        "                tiledata[2] = work_results[image_index] if image_index < len(work_results) else Image.new(\"RGB\", (p.width, p.height))\n",
        "                image_index += 1\n",
        "\n",
        "        combined_image = images.combine_grid(grid)\n",
        "\n",
        "        # if opts.samples_save:\n",
        "        #     images.save_image(combined_image, p.outpath_samples, \"\", initial_seed, prompt, opts.grid_format, info=initial_info)\n",
        "\n",
        "        processed = Processed(p, [combined_image], initial_seed, initial_info)\n",
        "\n",
        "    else:\n",
        "\n",
        "        processed = scripts_img2img.run(p, *args)\n",
        "\n",
        "        if processed is None:\n",
        "            processed = process_images(p)\n",
        "\n",
        "\n",
        "    return processed.images, processed.seed, None # plaintext_to_html(processed.info)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JGRHg3EfrWcp"
      },
      "source": [
        "#### Load model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FlP6AflBrZHn"
      },
      "outputs": [],
      "source": [
        "sd_config = OmegaConf.load(\"/content/stable-diffusion-webui/repositories/stable-diffusion/configs/stable-diffusion/v1-inference.yaml\")\n",
        "sd_model = load_model_from_config(sd_config, sd_model_file)\n",
        "sd_model = sd_model.half()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CMPZ0xM4xhto"
      },
      "outputs": [],
      "source": [
        "medvram = True #NOT ENOUGH VRAM\n",
        "setup_for_low_vram(sd_model, medvram)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DmwYnTpwlWfI"
      },
      "source": [
        "# Inference"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qSFTFJGzw2qs"
      },
      "source": [
        "## Txt2Img"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jup3Ow7LsM_o",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "prompt = \"1girl blush cowboy_shot detached_sleeves green_eyes green_hair hatsune_miku headset long_hair looking_at_viewer matching_hair/eyes necktie nekono_rin simple_background skirt solo thighhighs twintails v very_long_hair vocaloid white_background\" #@param {type:\"string\"}\n",
        "negative_prompt = \"\" #@param {type:\"string\"}\n",
        "sampler = 'k_euler_a' #@param [\"k_euler_a\",\"k_euler\", \"k_lms\", \"k_heun\", \"k_dpm_2\", \"k_dpm_2_a\"] {allow-input: false}\n",
        "width = 512 #@param {type:\"integer\"}\n",
        "height = 512 #@param {type:\"integer\"}\n",
        "\n",
        "scale = 7 #@param {type:\"slider\", min:1, max:30, step:0.5}\n",
        "steps = 10 #@param {type:\"slider\", min:1, max:150, step:1}\n",
        "\n",
        "batch_size = 4 #@param {type:'integer'}\n",
        "skip_grid = False #@param {type:\"boolean\"}\n",
        "rows = 2 #@param {type:'integer'}\n",
        "\n",
        "seed = -1 #@param {type:'integer'}\n",
        "n_iter = 1 #@param {type:'integer'}\n",
        "tiling = False # This doesn't work properly yet \n",
        "## Seed = -1 ==> Random\n",
        "\n",
        "images, seed_new, c = txt2img(prompt, negative_prompt, steps, name_index(sampler), False, \n",
        "        tiling, n_iter, batch_size, scale, seed, height, width, 0)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sEp7hzVwyBx_",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Caption\n",
        "index = 0 #@param {type:\"integer\"}\n",
        "info = f\"seed = {seed_new}, cfg = {scale}, steps = {steps}, sampler = {sampler}\"\n",
        "if negative_prompt == \"\":\n",
        "  new_prompt = prompt\n",
        "else:\n",
        "  new_prompt = f\"{prompt} - ({negative_prompt})\"\n",
        "captioned_image = caption(images[index], new_prompt, info)\n",
        "captioned_image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1EC7FLu4v1o4",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Show Image\n",
        "index = 0 #@param {type: 'integer'}\n",
        "images[index]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NaJntT7Bw4ZP"
      },
      "source": [
        "## Img2Img"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Download image from link\n",
        "%cd '/content/stable-diffusion-webui/Source'\n",
        "link = \"https://cdn.discordapp.com/attachments/692612547345514549/1010734092687519744/take_00474.png\" #@param {type:\"string\"}\n",
        "!wget {link}\n",
        "%cd  '/content/stable-diffusion-webui'"
      ],
      "metadata": {
        "cellView": "form",
        "id": "YYf2bEfK1YZw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OXqkgbEuseRf",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "prompt = \"a photo of a cat sitting on a chair\" #@param {type:\"string\"}\n",
        "sampler = 'k_euler_a' #@param [\"k_euler_a\",\"k-diffusion\", \"k_dpm_2\", \"k_dpm_2_a\", \"k_euler\", \"k_heun\"] {allow-input: false}\n",
        "init_image_path = \"/content/stable-diffusion-webui/Source/init.png\" #@param {type: 'string'}\n",
        "sampler = 'k_euler_a' #@param [\"k_euler_a\",\"k_euler\", \"k_lms\", \"k_heun\", \"k_dpm_2\", \"k_dpm_2_a\"] {allow-input: false}\n",
        "gen_mode = 'normal' #@param [\"normal\", \"inpaint\"]\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown INPAINT (Set at gen_mode)\n",
        "\n",
        "#@markdown White = SD, Black = Original\n",
        "\n",
        "init_mask_path = \"/content/stable-diffusion-webui/Source/mask.png\" #@param {type: 'string'}\n",
        "invert = False #@param {type: 'boolean'}\n",
        "inpainting_fill = 3 #@param {type: \"slider\", min:1, max:3, step:1}\n",
        "#@markdown ---\n",
        "resize_mode = \"Resize and fill\" #@param [\"Just resize\", \"Crop and resize\", \"Resize and fill\"] {allow-input: false}\n",
        "\n",
        "width = 512 #@param {type:\"integer\"}\n",
        "height = 512 #@param {type:\"integer\"}\n",
        "\n",
        "scale = 7.5 #@param {type:\"slider\", min:1, max:30, step:0.5}\n",
        "steps = 10 #@param {type:\"slider\", min:1, max:150, step:1}\n",
        "strength = 0.5 #@param {type: \"slider\", min:0.00, max:1.00, step:0.01}\n",
        "\n",
        "batch_size = 4 #@param {type:'integer'}\n",
        "skip_grid = False #@param {type:\"boolean\"}\n",
        "rows = 2 #@param {type:'integer'}\n",
        "\n",
        "n_iter = 1 #@param {type:'integer'}\n",
        "seed = -1 #@param {type:'integer'}\n",
        "## Seed = -1 ==> Random\n",
        "#@markdown Seed = -1 --> Random\n",
        "if resize_mode == 'Just resize':\n",
        "  resize_mode_user = 0\n",
        "elif resize_mode == 'Crop and resize':\n",
        "  resize_mode_user = 1\n",
        "else:\n",
        "  resize_mode_user = 2\n",
        "  \n",
        "init_image = Image.open(init_image_path)\n",
        "\n",
        "if gen_mode == 'normal':\n",
        "  mode = 0\n",
        "  init_mask = None\n",
        "else: # gen_mode == 'inpaint':\n",
        "  mode = 1\n",
        "  init_mask = Image.open(init_mask_path)\n",
        "\n",
        "init_img_with_mask = make_init_masked(init_image_path, init_mask_path, 20, invert)\n",
        "\n",
        "mask_blur = 10\n",
        "tiling = False\n",
        "denoising_strength_change_factor = 0.1\n",
        "inpaint_full_res = True\n",
        "inpainting_mask_invert = 0\n",
        "\n",
        "a, b, c = img2img(prompt, init_image, init_img_with_mask, init_mask, steps, name_index(sampler), mask_blur, \n",
        "        inpainting_fill, False, tiling, mode, n_iter, \n",
        "        batch_size, scale, strength, denoising_strength_change_factor, seed, \n",
        "        height, width, resize_mode_user, None, None, \n",
        "        inpaint_full_res, inpainting_mask_invert, 0)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Caption\n",
        "index = 0 #@param {type:\"integer\"}\n",
        "original = resize_image(resize_mode_user, init_image, width, height)\n",
        "if init_mask is not None:\n",
        "  masked = resize_image(resize_mode_user, init_mask, width, height)\n",
        "  original = get_concat_v_blank(original, masked)\n",
        "\n",
        "concatted = get_concat_h_blank(a[index], original)\n",
        "info = f\"seed = {b}, cfg = {scale}, steps = {steps}, strength = {strength}, sampler = {sampler}\"\n",
        "captioned_image = caption(concatted, prompt, info)\n",
        "captioned_image"
      ],
      "metadata": {
        "id": "PK60-67R7Q17",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Show Image\n",
        "index = 0 #@param {type: 'integer'}\n",
        "images[index]"
      ],
      "metadata": {
        "id": "3I4_g5xZ48bg",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "DHfrnDNPY5FO",
        "LJPrA0GepxQ7",
        "Vf5WJ4b3lSOD",
        "vlWxMWnTtsIW",
        "gzEgNxZ6xn1G",
        "6pfVOIy0pL4e",
        "eDakK97gsPHY",
        "s-Yr1u3so4X6",
        "U0vz6D412oSY",
        "36c5Zztlo8vh",
        "Qywis6QspZwv",
        "iB5fIbQ7pbZD",
        "l_kIaCEEpcmX",
        "JGRHg3EfrWcp"
      ],
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyPPWGpcGEfJk3ghjypz5XtC",
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}