Neo_Hidamari_Diffusion_wip.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/FuouM/HidamariDiffusionColab/blob/main/Neo_Hidamari_Diffusion_wip.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NtJ1grxqfG8u"
      },
      "source": [
        "Code taken from https://colab.research.google.com/drive/1yf3-bUhTcfxRmAphJHVQQAD2ArYO1CRZ\n",
        "\n",
        "Guides:\n",
        "\n",
        "https://rentry.org/retardsguide - Main txt2img guide\n",
        "https://rentry.org/kretard - K-Diffusion guide\n",
        "https://rentry.org/img2img - img2img guide\n",
        "\n",
        "Updated to use the new model with old code\n",
        "\n",
        "Recommend using Optimized SD + K-diffusion & hlky's fork (most up to date)\n",
        "\n",
        "Updated as of 8/27\n",
        "\n",
        "Archive: https://github.com/FuouM/HidamariDiffusionColab\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!nvidia-smi"
      ],
      "metadata": {
        "id": "dTlNpVz7d_v9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jwZ0GT0eObBW"
      },
      "source": [
        "### Download and set up the environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hTjD9Ij7Nuh4"
      },
      "outputs": [],
      "source": [
        "# !git clone https://github.com/FuouM/stable-diffusion-hidamari stable-diffusion\\\n",
        "!git clone https://github.com/basujindal/stable-diffusion \n",
        "%cd stable-diffusion"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9pxSsF1HOi4V"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install albumentations==0.4.3\n",
        "!pip install opencv-python==4.1.2.30\n",
        "!pip install pudb==2019.2\n",
        "!pip install imageio==2.9.0\n",
        "!pip install imageio-ffmpeg==0.4.2\n",
        "!pip install pytorch-lightning \n",
        "!pip install omegaconf==2.1.1\n",
        "!pip install test-tube>=0.7.5\n",
        "!pip install streamlit>=0.73.1\n",
        "!pip install einops==0.3.0\n",
        "!pip install torch-fidelity==0.3.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wQILJaTuZCLU"
      },
      "outputs": [],
      "source": [
        "!pip install transformers==4.19.2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "52gTy8xUZD_5"
      },
      "outputs": [],
      "source": [
        "!yes w | pip install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers\n",
        "!yes w | pip install -e git+https://github.com/openai/CLIP.git@main#egg=clip\n",
        "!git clone https://github.com/crowsonkb/k-diffusion.git src/k-diffusion\n",
        "!pip install src/k-diffusion\n",
        "!pip install kornia"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dPIqZIAu_SNj"
      },
      "outputs": [],
      "source": [
        "!mkdir -p '/content/stable-diffusion/Source'\n",
        "!mkdir -p '/content/stable-diffusion/Output'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rnULtFtROkqr"
      },
      "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": "XvHTXI7KOnu4"
      },
      "source": [
        "### Download the model (Choose only 1 option)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OAw5HFJRO4b-"
      },
      "outputs": [],
      "source": [
        "!gdown 1athuOO6kKtvLm3x1zqxNN8fqBIzV15ss -O /content/stable-diffusion/model.ckpt\n",
        "## Link seems dead"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b1CqjnK1Ik0J"
      },
      "outputs": [],
      "source": [
        "%cd stable-diffusion/\n",
        "!wget -O model.ckpt \"https://drive.yerf.org/wl/?id=EBfTrmcCCUAGaQBXVIj5lJmEhjoP1tgl&mode=grid&download=1\""
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### For Gdrive"
      ],
      "metadata": {
        "id": "yaHfjOHyJgdC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ],
      "metadata": {
        "id": "s1CMRG1FJo2R"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "user_path = \"\" #@param {type: 'string'} \n",
        "## exclude the model name ^ (Gdrive/MyDrive/some_folder)\n",
        "file_name = \"sd-v1-4.ckpt\" #@param {type: 'string'}\n",
        "%cd {user_path}\n",
        "\n",
        "!cp {file_name} \"/content/stable-diffusion/\"\n",
        "\n",
        "%cd /content/stable-diffusion\n"
      ],
      "metadata": {
        "id": "rb-O9TBEJt2j"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Rename your model file (if it isn't model.ckpt)"
      ],
      "metadata": {
        "id": "Gbo4OrWjcAsc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!mv {file_name} model.ckpt"
      ],
      "metadata": {
        "id": "4ZJ_5_8zb8EX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Huggingface (Fastest)"
      ],
      "metadata": {
        "id": "N-ObhCcRHhFv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "user_token = \"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\" #@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/model.ckpt"
      ],
      "metadata": {
        "id": "3BQoPx_8Hj08"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Optimized SD - Updated as of 8/27\n",
        "\n",
        "https://github.com/basujindal/stable-diffusion"
      ],
      "metadata": {
        "id": "Oudkt8AEKw05"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OmQfJWm8QGoA"
      },
      "source": [
        "## Text 2 Image (Optimized SD PLMS)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-7syo80FPa4T"
      },
      "source": [
        "### Set up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UY4lXPAhbU88"
      },
      "outputs": [],
      "source": [
        "%cd stable-diffusion"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fgwY7MOYXfja"
      },
      "outputs": [],
      "source": [
        "import argparse, os, sys, glob, random\n",
        "import torch\n",
        "import numpy as np\n",
        "from random import randint\n",
        "import math\n",
        "\n",
        "from omegaconf import OmegaConf\n",
        "from PIL import Image\n",
        "from tqdm import tqdm, trange\n",
        "from itertools import islice\n",
        "\n",
        "from einops import rearrange\n",
        "import time\n",
        "from pytorch_lightning import seed_everything\n",
        "from torch import autocast\n",
        "from contextlib import contextmanager, nullcontext\n",
        "from ldm.util import instantiate_from_config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KvnfjJvQXjiY"
      },
      "outputs": [],
      "source": [
        "def chunk(it, size):\n",
        "    it = iter(it)\n",
        "    return iter(lambda: tuple(islice(it, size)), ())\n",
        "\n",
        "def load_model_from_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",
        "    return sd\n",
        "\n",
        "def torch_gc():\n",
        "    torch.cuda.empty_cache()\n",
        "    torch.cuda.ipc_collect()"
      ]
    },
    {
      "cell_type": "code",
      "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/NotoSansJP-Bold.otf\", 20, encoding='utf-8')\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"
      ],
      "metadata": {
        "id": "H0LAycy3dbul"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class User_OSD1:\n",
        "  def __init__(self, prompt: str, seed: int, samples: int, steps: int, scale: float, height:int, width: int,\n",
        "               rows: int, iter: int, skip_grid: bool, skip_save: bool):\n",
        "    self.prompt = prompt\n",
        "    self.seed = seed\n",
        "    self.n_samples = samples\n",
        "\n",
        "    self.ddim_steps = steps\n",
        "    self.cfg_scale = scale\n",
        "    \n",
        "    self.height = height\n",
        "    self.width = width\n",
        "\n",
        "    self.n_rows = rows\n",
        "\n",
        "    self.n_iter = iter\n",
        "\n",
        "    self.skip_grid = skip_grid\n",
        "    self.skip_save = skip_save\n",
        "    "
      ],
      "metadata": {
        "id": "ueE6FqE28XYj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ihQD4BD1P1Vr"
      },
      "source": [
        "### Load model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UcSFxhq4P8kl"
      },
      "outputs": [],
      "source": [
        "config = \"optimizedSD/v1-inference.yaml\"\n",
        "ckpt = \"model.ckpt\"\n",
        "sd = load_model_from_config(f\"{ckpt}\")\n",
        "li, lo = [], []\n",
        "\n",
        "for key, value in sd.items():\n",
        "    sp = key.split('.')\n",
        "    if(sp[0]) == 'model':\n",
        "        if('input_blocks' in sp):\n",
        "            li.append(key)\n",
        "        elif('middle_block' in sp):\n",
        "            li.append(key)\n",
        "        elif('time_embed' in sp):\n",
        "            li.append(key)\n",
        "        else:\n",
        "            lo.append(key)\n",
        "            \n",
        "for key in li:\n",
        "    sd['model1.' + key[6:]] = sd.pop(key)\n",
        "for key in lo:\n",
        "    sd['model2.' + key[6:]] = sd.pop(key)\n",
        "\n",
        "config = OmegaConf.load(f\"{config}\")\n",
        "device = \"cuda\"\n",
        "\n",
        "model = instantiate_from_config(config.modelUNet)\n",
        "_, _ = model.load_state_dict(sd, strict=False)\n",
        "model.eval()\n",
        "\n",
        "model.unet_bs = True\n",
        "model.cdevice = device\n",
        "model.turbo = True\n",
        "\n",
        "modelCS = instantiate_from_config(config.modelCondStage)\n",
        "_, _ = modelCS.load_state_dict(sd, strict=False)\n",
        "modelCS.cond_stage_model.device = device\n",
        "modelCS.eval()\n",
        "    \n",
        "modelFS = instantiate_from_config(config.modelFirstStage)\n",
        "_, _ = modelFS.load_state_dict(sd, strict=False)\n",
        "modelFS.eval()\n",
        "del sd"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Set generator"
      ],
      "metadata": {
        "id": "ie9xTKWi8b1e"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def generate(user: User_OSD1, out_name: str):\n",
        "    torch_gc()\n",
        "    \n",
        "    device = \"cuda\"\n",
        "    C = 4\n",
        "    f = 8\n",
        "    ddim_eta = 0.0\n",
        "    start_code = None\n",
        "\n",
        "    model.half()\n",
        "    modelCS.half()\n",
        "\n",
        "    batch_size = user.n_samples\n",
        "      \n",
        "    if user.seed == -1:\n",
        "      user.seed = randint(0, 1000000)\n",
        "\n",
        "    init_seed = user.seed\n",
        "\n",
        "    seed_everything(user.seed)\n",
        "\n",
        "    assert prompt is not None\n",
        "    data = [batch_size * [prompt]]\n",
        "\n",
        "    precision_scope = autocast\n",
        "\n",
        "    with torch.no_grad():\n",
        "\n",
        "        all_samples = list()\n",
        "        for _ in trange(user.n_iter, desc=\"Sampling\"):\n",
        "            for prompts in tqdm(data, desc=\"data\"):\n",
        "                with precision_scope(\"cuda\"):\n",
        "                    modelCS.to(device)\n",
        "                    uc = None\n",
        "                    if user.cfg_scale != 1.0:\n",
        "                        uc = modelCS.get_learned_conditioning(batch_size * [\"\"])\n",
        "                    if isinstance(prompts, tuple):\n",
        "                        prompts = list(prompts)\n",
        "                    \n",
        "                    c = modelCS.get_learned_conditioning(prompts)                    \n",
        "\n",
        "                    shape = [C, Height // f, Width // f]\n",
        "                    modelCS.to(\"cpu\")                    \n",
        "\n",
        "                    samples_ddim = model.sample(S=user.ddim_steps,\n",
        "                                   conditioning=c,\n",
        "                                   batch_size=batch_size,\n",
        "                                   seed = user.seed,\n",
        "                                   shape=shape,\n",
        "                                   verbose=False,\n",
        "                                   unconditional_guidance_scale=user.cfg_scale,\n",
        "                                   unconditional_conditioning=uc,\n",
        "                                   eta=ddim_eta,\n",
        "                                   x_T=start_code)\n",
        "\n",
        "                    modelFS.to(device)\n",
        "\n",
        "                    for i in range(batch_size):\n",
        "                        \n",
        "                        x_samples_ddim = modelFS.decode_first_stage(samples_ddim[i].unsqueeze(0))\n",
        "                        x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)\n",
        "                        x_sample = 255. * rearrange(x_sample[0].cpu().numpy(), 'c h w -> h w c')\n",
        "\n",
        "                        out = Image.fromarray(x_sample.astype(np.uint8))\n",
        "                        if not user.skip_save:\n",
        "                          out.save(f\"/content/stable-diffusion/Output/{out_name}_{init_seed}[{i}].png\")\n",
        "\n",
        "                        all_samples.append(out)\n",
        "                        user.seed+=1\n",
        "\n",
        "                    modelFS.to(\"cpu\")\n",
        "\n",
        "                    del samples_ddim\n",
        "                    del x_sample\n",
        "                    del x_samples_ddim\n",
        "\n",
        "    if not user.skip_grid:\n",
        "              grid = image_grid(all_samples, batch_size, user.n_rows)\n",
        "              all_samples.insert(0, grid)\n",
        "\n",
        "    torch_gc()\n",
        "    return all_samples, init_seed"
      ],
      "metadata": {
        "id": "l6UxQ6dGcS6s"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def txt2img(prompt, seed, samples, steps, scale, height, width, rows, iter, skip_grid, skip_save, out_name: str):\n",
        "  if(rows > samples):\n",
        "    rows = samples\n",
        "  user = User_OSD1(prompt, seed, samples, steps, scale, height, width, rows, iter, skip_grid, skip_save)\n",
        "  return generate(user, out_name)"
      ],
      "metadata": {
        "id": "w9OcVFfL8giT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P95iTXcZQbnk"
      },
      "source": [
        "### Run Text 2 Image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CmNXWU20QesI"
      },
      "outputs": [],
      "source": [
        "prompt = \"A blue corgi. At the table. Drinking tea. Next to the window. Dusk. Sun shines through the window. Light reflection on table.\" #@param {type:\"string\"}\n",
        "\n",
        "steps = 50 #@param {type:\"slider\", min:1, max:150, step:1}\n",
        "scale = 7 #@param {type:\"slider\", min:1, max:30, step:0.5}\n",
        "\n",
        "batch_size = 4 #@param {type:'integer'}\n",
        "n_rows = 2 #@param {type:'integer'}\n",
        "\n",
        "Height = 512 #@param {type:'integer'}\n",
        "Width = 512 #@param {type:'integer'}\n",
        "\n",
        "skip_grid = False #@param {type: 'boolean'}\n",
        "skip_save = True #@param {type: 'boolean'}\n",
        "out_name = \"test\" #@param {type:\"string\"}\n",
        "seed = -1 #@param {type:'integer'}\n",
        "## Seed = -1 ==> Random\n",
        "\n",
        "images, seed_new = txt2img(prompt, seed, batch_size,\n",
        "                            steps, scale, Height, Width,\n",
        "                            n_rows, 1, skip_grid, skip_save,\n",
        "                            out_name)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PeaR0VMfr_im"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.kill(os.getpid(), 9) # Crash colab if runs out of gpu memory / Funny errors (Run from Set up again)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Caption\n",
        "index = 0 #@param {type:\"integer\"}\n",
        "info = f\"seed = {seed_new}, cfg = {scale}, steps = {steps}, sampler = PLMS\"\n",
        "captioned_image = caption(images[index], prompt, info)\n",
        "captioned_image"
      ],
      "metadata": {
        "id": "JCrLSAUHddq6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3IAzahTMZVIE"
      },
      "outputs": [],
      "source": [
        "images[0] # First one is grid"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IDQmNv_9ZTGc"
      },
      "source": [
        "#### Save images individually "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y8AwaB3XZhfb"
      },
      "outputs": [],
      "source": [
        "images[1]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fWys2AGRZhVg"
      },
      "outputs": [],
      "source": [
        "images[2]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e72kpFjNZhL8"
      },
      "outputs": [],
      "source": [
        "images[3]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h2D4kZNNJ_Uv"
      },
      "outputs": [],
      "source": [
        "images[4]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c6jepvv8Q-1a"
      },
      "source": [
        "## Image 2 Image (Optimized SD DDIM)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IFdeC2LZRBm-"
      },
      "source": [
        "### Set up"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%cd stable-diffusion"
      ],
      "metadata": {
        "id": "Exl7IBQKc8bC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bw-VyYfuRBPn"
      },
      "outputs": [],
      "source": [
        "import argparse, os, re\n",
        "import numpy as np\n",
        "import torch\n",
        "from random import randint\n",
        "\n",
        "from omegaconf import OmegaConf\n",
        "\n",
        "from PIL import Image\n",
        "from tqdm import tqdm, trange\n",
        "from itertools import islice\n",
        "from einops import rearrange\n",
        "import time\n",
        "import math\n",
        "\n",
        "from pytorch_lightning import seed_everything\n",
        "from torch import autocast\n",
        "from contextlib import contextmanager, nullcontext\n",
        "from einops import rearrange, repeat\n",
        "from ldm.util import instantiate_from_config"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def chunk(it, size):\n",
        "    it = iter(it)\n",
        "    return iter(lambda: tuple(islice(it, size)), ())\n",
        "\n",
        "def load_model_from_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",
        "    return sd\n",
        "\n",
        "def torch_gc():\n",
        "    torch.cuda.empty_cache()\n",
        "    torch.cuda.ipc_collect()"
      ],
      "metadata": {
        "id": "nRoKSRTZb0mb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def load_img(init_image, h0, w0):\n",
        "   \n",
        "    image = init_image.convert(\"RGB\")\n",
        "    w, h = image.size\n",
        "\n",
        "    # print(f\"loaded input image of size ({w}, {h}) from {path}\")   \n",
        "    if(h0 is not None and w0 is not None):\n",
        "        h, w = h0, w0\n",
        "    \n",
        "    w, h = map(lambda x: x - x % 64, (w, h))  # resize to integer multiple of 32\n",
        "\n",
        "    print(f\"New image size ({w}, {h})\")\n",
        "    image = image.resize((w, h), resample = Image.LANCZOS)\n",
        "    image = np.array(image).astype(np.float32) / 255.0\n",
        "    image = image[None].transpose(0, 3, 1, 2)\n",
        "    image = torch.from_numpy(image)\n",
        "    return 2.*image - 1.\n",
        "\n",
        "LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)\n",
        "invalid_filename_chars = '<>:\"/\\|?*\\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",
        "      if im.width != width or im.height != height:\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",
        "        else:\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",
        "      else:\n",
        "        return im\n",
        "\n",
        "    return res"
      ],
      "metadata": {
        "id": "mPp1wgnoc_Vg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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/NotoSansJP-Bold.otf\", 20, encoding='utf-8')\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"
      ],
      "metadata": {
        "id": "2YY_5NWndZ_M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class User_OSD2:\n",
        "    def __init__(self, prompt: str, seed: int, samples: int, steps: int, scale: float, strength: float,\n",
        "                 height:int, width: int, rows: int, iter: int, skip_grid: bool, skip_save: bool):\n",
        "        self.prompt = prompt\n",
        "        self.seed = seed\n",
        "\n",
        "        self.n_samples = samples\n",
        "\n",
        "        self.ddim_steps = steps\n",
        "        self.cfg_scale = scale\n",
        "        self.strength = strength\n",
        "\n",
        "        self.height = height\n",
        "        self.width = width\n",
        "\n",
        "        self.n_rows = rows\n",
        "        self.n_iter = iter\n",
        "\n",
        "        self.skip_grid = skip_grid\n",
        "        self.skip_save = skip_save"
      ],
      "metadata": {
        "id": "IIeZdyOd834L"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7oPo2rboRWfH"
      },
      "source": [
        "### Load model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o2pAF_LzRZSr"
      },
      "outputs": [],
      "source": [
        "config = \"optimizedSD/v1-inference.yaml\"\n",
        "ckpt = \"model.ckpt\"\n",
        "device = \"cuda\"\n",
        "\n",
        "sd = load_model_from_config(f\"{ckpt}\")\n",
        "li, lo = [], []\n",
        "for key, v_ in sd.items():\n",
        "    sp = key.split('.')\n",
        "    if(sp[0]) == 'model':\n",
        "        if('input_blocks' in sp):\n",
        "            li.append(key)\n",
        "        elif('middle_block' in sp):\n",
        "            li.append(key)\n",
        "        elif('time_embed' in sp):\n",
        "            li.append(key)\n",
        "        else:\n",
        "            lo.append(key)\n",
        "for key in li:\n",
        "    sd['model1.' + key[6:]] = sd.pop(key)\n",
        "for key in lo:\n",
        "    sd['model2.' + key[6:]] = sd.pop(key)\n",
        "\n",
        "config = OmegaConf.load(f\"{config}\")\n",
        "\n",
        "model = instantiate_from_config(config.modelUNet)\n",
        "_, _ = model.load_state_dict(sd, strict=False)\n",
        "model.eval()    \n",
        "\n",
        "modelCS = instantiate_from_config(config.modelCondStage)\n",
        "_, _ = modelCS.load_state_dict(sd, strict=False)\n",
        "modelCS.cond_stage_model.device = device\n",
        "modelCS.eval()\n",
        "\n",
        "modelFS = instantiate_from_config(config.modelFirstStage)\n",
        "_, _ = modelFS.load_state_dict(sd, strict=False)\n",
        "modelFS.eval()\n",
        "\n",
        "del sd"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Set generator"
      ],
      "metadata": {
        "id": "zUGoyvzD86IL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def generate(user: User_OSD2, input_image, out_name: str):\n",
        "    torch_gc()\n",
        "    device = \"cuda\"\n",
        "    batch_size = user.n_samples\n",
        "    model.small_batch = False\n",
        "\n",
        "    model.cdevice = device\n",
        "    model.unet_bs = True\n",
        "    model.turbo = True\n",
        "\n",
        "    model.half()\n",
        "    modelCS.half()\n",
        "    modelFS.half()\n",
        "    \n",
        "    init_image = load_img(input_image, user.height, user.width).to(device).half()\n",
        "    \n",
        "    if user.seed == -1:\n",
        "      user.seed = randint(0, 1000000)\n",
        "\n",
        "    init_seed = user.seed\n",
        "\n",
        "    seed_everything(user.seed)\n",
        "\n",
        "    assert prompt is not None\n",
        "    data = [batch_size * [prompt]]\n",
        "\n",
        "    modelFS.to(device)\n",
        "\n",
        "    init_image = repeat(init_image, '1 ... -> b ...', b=batch_size)\n",
        "    init_latent = modelFS.get_first_stage_encoding(modelFS.encode_first_stage(init_image))  # move to latent space\n",
        "\n",
        "    modelFS.to(\"cpu\")\n",
        "\n",
        "    assert 0. <= user.strength <= 1., 'can only work with strength in [0.0, 1.0]'\n",
        "    t_enc = int(user.strength * user.ddim_steps)\n",
        "    print(f\"target t_enc is {t_enc} steps\")\n",
        "\n",
        "    precision_scope = autocast\n",
        "\n",
        "    with torch.no_grad():\n",
        "        all_samples = list()\n",
        "        for _ in trange(user.n_iter, desc=\"Sampling\"):\n",
        "            for prompts in tqdm(data, desc=\"data\"):\n",
        "                with precision_scope(\"cuda\"):\n",
        "                    modelCS.to(device)\n",
        "                    uc = None\n",
        "                    if user.cfg_scale != 1.0:\n",
        "                        uc = modelCS.get_learned_conditioning(batch_size * [\"\"])\n",
        "                    if isinstance(prompts, tuple):\n",
        "                        prompts = list(prompts)\n",
        "                    \n",
        "                    c = modelCS.get_learned_conditioning(prompts)\n",
        "\n",
        "                    modelCS.to(\"cpu\")\n",
        "\n",
        "                    # encode (scaled latent)\n",
        "                    z_enc = model.stochastic_encode(init_latent, torch.tensor([t_enc]*batch_size).to(device), user.seed,ddim_steps=user.ddim_steps, ddim_eta=0.0)\n",
        "                    # decode it\n",
        "                    samples_ddim = model.decode(z_enc, c, t_enc, unconditional_guidance_scale=user.cfg_scale,\n",
        "                                                    unconditional_conditioning=uc,)\n",
        "\n",
        "                    modelFS.to(device)\n",
        "                    # print(\"saving images\")\n",
        "                    for i in range(batch_size):\n",
        "                        \n",
        "                        x_samples_ddim = modelFS.decode_first_stage(samples_ddim[i].unsqueeze(0))\n",
        "                        x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)\n",
        "                        x_sample = 255. * rearrange(x_sample[0].cpu().numpy(), 'c h w -> h w c')\n",
        "\n",
        "                        out = Image.fromarray(x_sample.astype(np.uint8))\n",
        "                        if not user.skip_save:\n",
        "                          out.save(f\"/content/stable-diffusion/Output/{out_name}_{init_seed}[{i}].png\")\n",
        "                        all_samples.append(out)\n",
        "\n",
        "                        user.seed+=1\n",
        "\n",
        "                    modelFS.to(\"cpu\")\n",
        "\n",
        "                    del samples_ddim\n",
        "                    del x_sample\n",
        "                    del x_samples_ddim\n",
        "\n",
        "    if not user.skip_grid:\n",
        "              grid = image_grid(all_samples, batch_size, user.n_rows)\n",
        "              all_samples.insert(0, grid)\n",
        "    torch_gc()\n",
        "    return all_samples, init_seed"
      ],
      "metadata": {
        "id": "vdfIbOnkdDMv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def img2img(prompt, seed, samples, steps, scale, strength, height, width, rows,\n",
        "            iter, skip_grid, skip_save, mode, init_image, out_name):\n",
        "  if mode == \"Just resize\":\n",
        "    resize_mode = 0\n",
        "  elif mode == \"Crop and resize\":\n",
        "    resize_mode = 1\n",
        "  else:\n",
        "    resize_mode = 2\n",
        "  if(rows > samples):\n",
        "      rows = samples\n",
        "  user = User_OSD2(prompt, seed, samples, steps, scale, strength, height, width, rows, iter, skip_grid, skip_save)\n",
        "  init_image = resize_image(resize_mode, init_image, width, height)\n",
        "\n",
        "  return generate(user, init_image, out_name) + (resize_mode,)"
      ],
      "metadata": {
        "id": "lVoVkuXQ9JSe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NMj_JbQ9Rrxq"
      },
      "source": [
        "### Run Image 2 Image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xdR7s1yQR3m4"
      },
      "outputs": [],
      "source": [
        "%cd '/content/stable-diffusion/Source'\n",
        "link = \"https://cdn.discordapp.com/attachments/692612547345514549/1010734092687519744/take_00474.png\" #@param {type:\"string\"}\n",
        "!wget {link}\n",
        "%cd  '/content/stable-diffusion/'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5q1Qeg3_RwKU"
      },
      "outputs": [],
      "source": [
        "prompt = \"a male hand holding a gun shooting a bullet with gun powder and explosion, hyperreal, artstation, masterpiece, oil painting, close shot\" #@param {type:\"string\"}\n",
        "image = \"/content/stable-diffusion/Source/take_00474.png\" #@param {type:\"string\"}\n",
        "resize_mode = \"Resize and fill\" #@param [\"Just resize\", \"Crop and resize\", \"Resize and fill\"] {allow-input: false}\n",
        "\n",
        "batch_size = 4 #@param {type:\"integer\"}\n",
        "\n",
        "scale = 7 #@param {type:\"slider\", min:1, max:30, step:0.5}\n",
        "steps = 50 #@param {type:\"slider\", min:30, max:250, step:5}\n",
        "strength = 0.7 #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "\n",
        "height = 512 #@param {type:\"integer\"}\n",
        "width = 512 #@param {type:\"integer\"}\n",
        "\n",
        "n_rows = 2 #@param {type:\"integer\"}\n",
        "skip_grid = False #@param {type: 'boolean'}\n",
        "skip_save = True #@param {type: 'boolean'}\n",
        "out_name = \"test\" #@param {type:\"string\"}\n",
        "seed = -1 #@param {type:\"integer\"};\n",
        "## Seed = -1 ==> Random\n",
        "init_image = Image.open(image)\n",
        "images, seed_new, mode = img2img(prompt, seed, batch_size, steps, scale, strength, height, width, n_rows, 1, skip_grid, skip_save, resize_mode, init_image, out_name)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.kill(os.getpid(), 9) # Crash colab if runs out of gpu memory / Funny errors (Run from Set up again)"
      ],
      "metadata": {
        "id": "L-NXHyjcdmRm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Caption\n",
        "index = 0 #@param {type:\"integer\"}\n",
        "original = resize_image(mode, init_image, width, height)\n",
        "\n",
        "concatted = get_concat_h_blank(images[index], original)\n",
        "info = f\"seed = {seed_new}, cfg = {scale}, steps = {steps}, strength = {strength}, sampler = DDIM\"\n",
        "captioned_image = caption(concatted, prompt, info)\n",
        "captioned_image"
      ],
      "metadata": {
        "id": "bB674hn6dQ7Z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u3zuCpuYNJ6_"
      },
      "outputs": [],
      "source": [
        "images[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Individual outputs"
      ],
      "metadata": {
        "id": "iIBkjkG_dHAu"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UgSuZn4DOKAs"
      },
      "outputs": [],
      "source": [
        "images[1]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-rHaSK3VP6ja"
      },
      "outputs": [],
      "source": [
        "images[2]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ldZcKISoP8E5"
      },
      "outputs": [],
      "source": [
        "images[3]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cQ1-eVNhP8-k"
      },
      "outputs": [],
      "source": [
        "images[4]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5kQ2bDBrPe6Q"
      },
      "source": [
        "### Running image sequences"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1DtBv6OVPiR_"
      },
      "outputs": [],
      "source": [
        "from PIL import Image\n",
        "from IPython.display import clear_output \n",
        "import numpy as np\n",
        "import timeit\n",
        "#@title Image sequence\n",
        "#@markdown ##### Enter text prompt:\n",
        "prompt = \"man standing in front of a painting of a mountain in the distance, wide angle, Van Gogh, futuristic, surreal, holy\" #@param {type:\"string\"}\n",
        "resize_mode = \"Resize and fill\" #@param [\"Just resize\", \"Crop and resize\", \"Resize and fill\"] {allow-input: false}\n",
        "#@markdown ---\n",
        "#@markdown ##### Set sequence duration (/Source/prefix_XXXXX.png):\n",
        "start_frame = \"00024\" #@param {type:\"string\"}\n",
        "end_frame = \"00028\" #@param {type:\"string\"}\n",
        "prefix = \"take\" #@param {type:\"string\"}\n",
        "current_frame = start_frame\n",
        "#@markdown ---\n",
        "#@markdown ##### SD settings:\n",
        "height = 512 #@param {type:\"integer\"}\n",
        "width = 512 #@param {type:\"integer\"}\n",
        "scale = 7 #@param {type:\"slider\", min:1, max:30, step:0.5}\n",
        "steps = 100 #@param {type:\"slider\", min:30, max:250, step:5}\n",
        "strength = 0.7 #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "batch_size = 1 #@param {type:\"integer\"}\n",
        "n_rows = 1 #@param {type:\"integer\"}\n",
        "\n",
        "seed = -1 #@param {type:\"integer\"};\n",
        "timeStart = timeit.default_timer()\n",
        "for i in range(int(start_frame), int(end_frame) + 1):\n",
        "  current_frame = \"{:05d}\".format(i)\n",
        "  print(\"PROCESSING \" + current_frame)\n",
        "  img = \"/content/stable-diffusion/Source/{}_{}.png\".format(prefix, current_frame)\n",
        "  init_image = Image.open(image)\n",
        "  images, seed_new, mode = img2img(prompt, seed, 1, steps, scale, strength, height, width, 1, 1, True, True, resize_mode, init_image, \"\")\n",
        "  images[0].save(\"/content/stable-diffusion/Output/{}_{}.png\".format(prefix, current_frame))\n",
        "  images[0]\n",
        "\n",
        "  clear_output()\n",
        "timeStop = timeit.default_timer()\n",
        "print(\"{} frames done in {} seconds.\\nPrompt: {}\\nStrength: {}\".format(int(end_frame) - int(start_frame) + 1, timeStop - timeStart, prompt, strength))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m1cRK744Ch4c"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.kill(os.getpid(), 9) # Crash colab if runs out of gpu memory / Funny errors (Run from Set up again)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZFnTkna0dP1K"
      },
      "outputs": [],
      "source": [
        "from google.colab import files\n",
        "!zip -r \"/content/stable-diffusion/result.zip\" \"/content/stable-diffusion/Output\"\n",
        "files.download('/content/stable-diffusion/result.zip')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NBmrTroGjzxt"
      },
      "outputs": [],
      "source": [
        "!rm \"/content/stable-diffusion/result.zip\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kwtcgOzRhll5"
      },
      "outputs": [],
      "source": [
        "#@title Export video\n",
        "from google.colab import files\n",
        "file_name = \"out.mp4\" #@param {type: \"string\"}\n",
        "always_overwrite = True #@param {type:\"boolean\"}\n",
        "framerate = 30 #@param (type:number)\n",
        "if always_overwrite:\n",
        "  replace = 'y'\n",
        "else:\n",
        "  replace = 'n'\n",
        "!yes {replace} | ffmpeg -framerate {framerate} -start_number {int(start_frame)} -i /content/stable-diffusion/Output/{prefix}_%05d.png -c:v libx264 -crf 0 {file_name}\n",
        "files.download('/content/stable-diffusion/{}'.format(file_name))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fViEXufPjam6"
      },
      "outputs": [],
      "source": [
        "!rm -rf /content/stable-diffusion/Source/*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Np75eoywjbn8"
      },
      "outputs": [],
      "source": [
        "!rm -rf /content/stable-diffusion/Output/*"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Optimized SD + K-diffusion (Updated as of 8/27)"
      ],
      "metadata": {
        "id": "CszrKJDe-66T"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Set up"
      ],
      "metadata": {
        "id": "5deJM1EQ_BIC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/stable-diffusion"
      ],
      "metadata": {
        "id": "L70j4gz6_Aq1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import argparse, os\n",
        "import torch\n",
        "import numpy as np\n",
        "from random import randint\n",
        "from omegaconf import OmegaConf\n",
        "from PIL import Image\n",
        "from tqdm import tqdm, trange\n",
        "from itertools import islice\n",
        "from einops import rearrange, repeat\n",
        "import time\n",
        "from pytorch_lightning import seed_everything\n",
        "from torch import autocast\n",
        "from contextlib import nullcontext\n",
        "from ldm.util import instantiate_from_config\n",
        "import k_diffusion as K\n",
        "import math"
      ],
      "metadata": {
        "id": "DkCi-pIf_I-Q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def chunk(it, size):\n",
        "    it = iter(it)\n",
        "    return iter(lambda: tuple(islice(it, size)), ())\n",
        "\n",
        "\n",
        "def load_model_from_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",
        "    return sd\n",
        "\n",
        "\n",
        "class CFGDenoiser(torch.nn.Module):\n",
        "    def __init__(self, model):\n",
        "        super().__init__()\n",
        "        self.inner_model = model\n",
        "\n",
        "    def forward(self, x, sigma, uncond, cond, cond_scale):\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",
        "        return uncond + (cond - uncond) * cond_scale\n",
        "\n",
        "def torch_gc():\n",
        "    torch.cuda.empty_cache()\n",
        "    torch.cuda.ipc_collect()\n",
        "\n",
        "def create_random_tensors(shape, seeds):\n",
        "    xs = []\n",
        "    for seed in seeds:\n",
        "        torch.manual_seed(seed)\n",
        "\n",
        "        xs.append(torch.randn(shape, device=device))\n",
        "    x = torch.stack(xs)\n",
        "    return x\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"
      ],
      "metadata": {
        "id": "Ph6x3jUF_KJ0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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/NotoSansJP-Bold.otf\", 20, encoding='utf-8')\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"
      ],
      "metadata": {
        "id": "gqoTA8Pd_M-X"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "## UPDATE\n",
        "\n",
        "LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)\n",
        "invalid_filename_chars = '<>:\"/\\|?*\\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",
        "      if im.width != width or im.height != height:\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",
        "        else:\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",
        "      else:\n",
        "        return im\n",
        "\n",
        "    return res"
      ],
      "metadata": {
        "id": "6QvTJ5xY_QYJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def load_img(path, h0, w0):\n",
        "    image = Image.open(path).convert(\"RGB\")\n",
        "    w, h = image.size\n",
        "\n",
        "    print(f\"loaded input image of size ({w}, {h}) from {path}\")\n",
        "    if (h0 is not None and w0 is not None):\n",
        "        h, w = h0, w0\n",
        "\n",
        "    w, h = map(lambda x: x - x % 32, (w0, h0))  # resize to integer multiple of 32\n",
        "\n",
        "    print(f\"New image size ({w}, {h})\")\n",
        "    image = image.resize((w, h), resample=Image.LANCZOS)\n",
        "    image = np.array(image).astype(np.float32) / 255.0\n",
        "    image = image[None].transpose(0, 3, 1, 2)\n",
        "    image = torch.from_numpy(image)\n",
        "    return 2. * image - 1."
      ],
      "metadata": {
        "id": "23-jsUl1_RoV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Info_OSD:\n",
        "    def __init__(self, prompt, seed, sampler, ddim_steps, cfg_scale, strength, n_samples, n_rows, height, width, skip_grid, n_iter, resize_mode):\n",
        "        self.prompt = prompt\n",
        "        self.seed = seed\n",
        "\n",
        "        self.sampler = sampler\n",
        "        self.ddim_steps = ddim_steps\n",
        "        self.cfg_scale = cfg_scale\n",
        "        self.strength = strength\n",
        "\n",
        "        self.n_samples = n_samples\n",
        "        self.n_rows = n_rows\n",
        "\n",
        "\n",
        "        self.height = height\n",
        "        self.width = width\n",
        "\n",
        "        self.skip_grid = skip_grid\n",
        "        self.n_iter = n_iter\n",
        "        self.resize_mode = resize_mode\n"
      ],
      "metadata": {
        "id": "1Dz3IYRl_3Un"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load model"
      ],
      "metadata": {
        "id": "YDcdDSiX_HAo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "config = \"/content/stable-diffusion/optimizedSD/v1-inference.yaml\"\n",
        "ckpt = \"model.ckpt\"\n",
        "device = \"cuda\"\n",
        "sd = load_model_from_config(f\"{ckpt}\")\n",
        "li = []\n",
        "lo = []\n",
        "for key, value in sd.items():\n",
        "    sp = key.split('.')\n",
        "    if(sp[0]) == 'model':\n",
        "        if('input_blocks' in sp):\n",
        "            li.append(key)\n",
        "        elif('middle_block' in sp):\n",
        "            li.append(key)\n",
        "        elif('time_embed' in sp):\n",
        "            li.append(key)\n",
        "        else:\n",
        "            lo.append(key)\n",
        "for key in li:\n",
        "    sd['model1.' + key[6:]] = sd.pop(key)\n",
        "for key in lo:\n",
        "    sd['model2.' + key[6:]] = sd.pop(key)\n",
        "\n",
        "config = OmegaConf.load(f\"{config}\")\n",
        "\n",
        "\n",
        "model = instantiate_from_config(config.modelUNet)\n",
        "_, _ = model.load_state_dict(sd, strict=False)\n",
        "model.eval()\n",
        "    \n",
        "modelCS = instantiate_from_config(config.modelCondStage)\n",
        "_, _ = modelCS.load_state_dict(sd, strict=False)\n",
        "modelCS.cond_stage_model.device = device\n",
        "modelCS.eval()\n",
        "    \n",
        "modelFS = instantiate_from_config(config.modelFirstStage)\n",
        "_, _ = modelFS.load_state_dict(sd, strict=False)\n",
        "modelFS.eval()\n",
        "\n",
        "model.unet_bs = True\n",
        "model.turbo = True\n",
        "model.cdevice = device\n",
        "modelCS.cond_stage_model.device = device\n",
        "model.to(device)"
      ],
      "metadata": {
        "id": "cyt5iSO5--Vg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Set generators\n",
        "\n",
        "txt2img https://files.catbox.moe/mmtgd5.py\n",
        "\n",
        "img2img https://files.catbox.moe/n7jqeu.py"
      ],
      "metadata": {
        "id": "4y4c6LWDAAVB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def txt2img_generate(user: Info_OSD):\n",
        "  torch_gc()\n",
        "  opt_C = 4\n",
        "  opt_f = 8\n",
        "\n",
        "  device = \"cuda\"\n",
        "  # model.unet_bs = True\n",
        "  # model.turbo = True\n",
        "  # model.cdevice = device  \n",
        "\n",
        "  model.half()\n",
        "  modelCS.half()  \n",
        "  modelFS.half()\n",
        "\n",
        "  batch_size = user.n_samples\n",
        "  prompt = user.prompt\n",
        "  assert prompt is not None\n",
        "  data = [batch_size * [prompt]]\n",
        "\n",
        "  if user.seed == -1:\n",
        "      user.seed = randint(0, 1000000)\n",
        "  init_seed = user.seed\n",
        "  batch_size = user.n_samples\n",
        "  all_prompts = batch_size * user.n_iter * [prompt]\n",
        "  all_seeds = [user.seed + x for x in range(len(all_prompts))]\n",
        "  precision_scope = autocast\n",
        "  \n",
        "  model_wrap = K.external.CompVisDenoiser(model).half()\n",
        "  # sigma_min, sigma_max = model_wrap.sigmas[0].item(), model_wrap.sigmas[-1].item()  \n",
        "  with torch.no_grad():\n",
        "      all_samples = list()\n",
        "      for n in trange(user.n_iter, desc=\"Sampling\"):\n",
        "          for prompts in tqdm(data, desc=\"data\"):\n",
        "              with precision_scope(\"cuda\"):\n",
        "                  modelCS.to(device)\n",
        "                  modelFS.to(device)\n",
        "                  uc = None\n",
        "                  if user.cfg_scale != 1.0:\n",
        "                      uc = modelCS.get_learned_conditioning(batch_size * [\"\"])\n",
        "                  if isinstance(prompts, tuple):\n",
        "                      prompts = list(prompts)\n",
        "                  \n",
        "                  c = modelCS.get_learned_conditioning(prompts)\n",
        "                  shape = [opt_C, user.height // opt_f, user.width // opt_f]\n",
        "                  modelCS.to(\"cpu\")\n",
        "                  \n",
        "                  model_wrap_cfg = CFGDenoiser(model_wrap)\n",
        "                  sigmas = model_wrap.get_sigmas(user.ddim_steps)\n",
        "                  extra_args = {'cond': c, 'uncond': uc, 'cond_scale': user.cfg_scale}\n",
        "                  x = create_random_tensors(shape, seeds=all_seeds)\n",
        "                  x = x * sigmas[0]\n",
        "                  samples_k = user.sampler(model_wrap_cfg, x, sigmas, extra_args=extra_args)\n",
        "                  \n",
        "                  for i in range(batch_size):\n",
        "                      x_samples_k = modelFS.decode_first_stage(samples_k[i].unsqueeze(0))\n",
        "                      x_sample = torch.clamp((x_samples_k + 1.0) / 2.0, min=0.0, max=1.0)\n",
        "                      x_sample = 255. * rearrange(x_sample[0].cpu().numpy(), 'c h w -> h w c')\n",
        "                      \n",
        "                      all_samples.append(Image.fromarray(x_sample.astype(np.uint8)))\n",
        "\n",
        "                  modelFS.to(\"cpu\")\n",
        "                  torch_gc()\n",
        "                  del samples_k\n",
        "                  del x_sample\n",
        "\n",
        "          if not user.skip_grid:\n",
        "              grid = image_grid(all_samples, batch_size, user.n_rows)\n",
        "              all_samples.insert(0, grid)\n",
        "  return all_samples, init_seed"
      ],
      "metadata": {
        "id": "9Jf1a_YV_-29"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def img2img_generate(user: Info_OSD, init_image):\n",
        "  torch_gc()\n",
        "  device = \"cuda\"\n",
        "  # model.to(device)\n",
        "\n",
        "  model.half()\n",
        "  modelCS.half()\n",
        "  modelFS.half()\n",
        "\n",
        "  # model.unet_bs = True\n",
        "  # model.turbo = True\n",
        "  # model.cdevice = device\n",
        "  # modelCS.cond_stage_model.device = device\n",
        "  \n",
        "  model_wrap = K.external.CompVisDenoiser(model).half()\n",
        "  sigma_min, sigma_max = model_wrap.sigmas[0].item(), model_wrap.sigmas[-1].item()\n",
        "\n",
        "  opt_C = 4\n",
        "  opt_f = 8\n",
        "  batch_size = user.n_samples\n",
        "\n",
        "  load_image = load_img(init_image, height, width).to(device)\n",
        "  load_image = load_image.half()\n",
        "  \n",
        "  prompt = user.prompt\n",
        "  assert prompt is not None\n",
        "  data = [batch_size * [prompt]]\n",
        "\n",
        "  if user.seed == -1:\n",
        "      user.seed = randint(0, 1000000)\n",
        "  init_seed = user.seed\n",
        "  batch_size = user.n_samples\n",
        "\n",
        "  all_prompts = batch_size * user.n_iter * [prompt]\n",
        "  all_seeds = [user.seed + x for x in range(len(all_prompts))]\n",
        "\n",
        "  modelFS.to(device)\n",
        "\n",
        "  load_image = repeat(load_image, '1 ... -> b ...', b=batch_size)\n",
        "  init_latent = modelFS.get_first_stage_encoding(modelFS.encode_first_stage(load_image))  # move to latent space\n",
        "\n",
        "  modelFS.to(\"cpu\")\n",
        "  \n",
        "  assert 0. <= user.strength <= 1., 'can only work with strength in [0.0, 1.0]'\n",
        "  t_enc = int(user.strength * user.ddim_steps)\n",
        "  print(f\"target t_enc is {t_enc} steps\")\n",
        "\n",
        "  precision_scope = autocast\n",
        "  with torch.no_grad():\n",
        "\n",
        "      all_samples = list()\n",
        "      for n in trange(user.n_iter, desc=\"Sampling\"):\n",
        "          for prompts in tqdm(data, desc=\"data\"):\n",
        "              with precision_scope(\"cuda\"):\n",
        "                  modelCS.to(device)\n",
        "                  uc = None\n",
        "                  if user.cfg_scale != 1.0:\n",
        "                      uc = modelCS.get_learned_conditioning(batch_size * [\"\"])\n",
        "                  if isinstance(prompts, tuple):\n",
        "                      prompts = list(prompts)\n",
        "                  \n",
        "                  c = modelCS.get_learned_conditioning(prompts)\n",
        "                  modelCS.to(\"cpu\")\n",
        "                  sigmas = model_wrap.get_sigmas(user.ddim_steps)\n",
        "                  shape = [opt_C, user.height // opt_f, user.width // opt_f]\n",
        "                  x = create_random_tensors(shape, seeds=all_seeds)\n",
        "                  x0 = init_latent\n",
        "                  sigmas = model_wrap.get_sigmas(user.ddim_steps)\n",
        "                  noise = x * sigmas[user.ddim_steps - t_enc - 1]\n",
        "\n",
        "                  xi = x0 + noise\n",
        "                  sigma_sched = sigmas[user.ddim_steps - t_enc - 1:]\n",
        "                  model_wrap_cfg = CFGDenoiser(model_wrap)\n",
        "                  extra_args = {'cond': c, 'uncond': uc, 'cond_scale': user.cfg_scale}\n",
        "                  samples_k = user.sampler(model_wrap_cfg, xi, sigma_sched, extra_args=extra_args)\n",
        "\n",
        "                  # modelCS.to(\"cpu\")\n",
        "                  modelFS.to(device)\n",
        "                  torch_gc()\n",
        "                  for i in range(batch_size):\n",
        "                      x_samples_k = modelFS.decode_first_stage(samples_k[i].unsqueeze(0))\n",
        "                      x_sample = torch.clamp((x_samples_k + 1.0) / 2.0, min=0.0, max=1.0)\n",
        "                      x_sample = 255. * rearrange(x_sample[0].cpu().numpy(), 'c h w -> h w c')\n",
        "                      \n",
        "                      all_samples.append(Image.fromarray(x_sample.astype(np.uint8)))\n",
        "                      \n",
        "                  modelFS.to(\"cpu\")\n",
        "                  torch_gc()\n",
        "\n",
        "                  del samples_k\n",
        "\n",
        "          if not user.skip_grid:\n",
        "              grid = image_grid(all_samples, batch_size, user.n_rows)\n",
        "              all_samples.insert(0, grid)\n",
        "  return all_samples, init_seed"
      ],
      "metadata": {
        "id": "F-2Kahe7AEnZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def txt2img(prompt: str, seed: int, sampler_name: str, steps: int, scale: int, strength: int, samples: int, rows: int, height: int, width: int, skip_grid: bool):\n",
        "  \n",
        "  if sampler_name == 'k-diffusion':\n",
        "    sampler = K.sampling.sample_lms\n",
        "  elif sampler_name == 'k_euler_a':\n",
        "    sampler = K.sampling.sample_euler_ancestral\n",
        "\n",
        "  user = Info_OSD(prompt, seed, sampler, steps, scale, strength, samples, rows, height, width, skip_grid, 1, 0)\n",
        "\n",
        "  return txt2img_generate(user)"
      ],
      "metadata": {
        "id": "ufNVs9AbAH4q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def img2img(prompt: str, init_image, seed: int, sampler_name: str, steps: int, scale: int, strength: int, samples: int, rows: int, height: int, width: int, skip_grid: bool, mode):\n",
        "  if sampler_name == 'k-diffusion':\n",
        "    sampler = K.sampling.sample_lms\n",
        "  elif sampler_name == 'k_euler_a':\n",
        "    sampler = K.sampling.sample_euler_ancestral\n",
        "\n",
        "  if mode == \"Just resize\":\n",
        "    resize_mode = 0\n",
        "  elif mode == \"Crop and resize\":\n",
        "    resize_mode = 1\n",
        "  else:\n",
        "    resize_mode = 2\n",
        "  user = Info_OSD(prompt, seed, sampler, steps, scale, strength, samples, rows, height, width, skip_grid, 1, resize_mode)\n",
        "  \n",
        "  return img2img_generate(user, init_image) + (resize_mode, )"
      ],
      "metadata": {
        "id": "vjOidPrgAIg9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Text 2 Image\n",
        "\n"
      ],
      "metadata": {
        "id": "9lnxMQBZ_asY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = \"aaa, bbb, ccc, octane 3D render\" #@param {type:\"string\"}\n",
        "sampler = 'k-diffusion' #@param [\"k_euler_a\",\"k-diffusion\"] {allow-input: false}\n",
        "\n",
        "width = 512 #@param {type:\"integer\"}\n",
        "height = 512 #@param {type:\"integer\"}\n",
        "\n",
        "scale = 7 #@param {type:'integer'}\n",
        "steps = 50 #@param {type:\"slider\", min:1, max:150, step:1}\n",
        "\n",
        "samples = 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",
        "images, seed_new = txt2img(prompt, seed, sampler, steps, scale, 0, samples, rows, height, width, \n",
        "                           skip_grid)\n"
      ],
      "metadata": {
        "id": "ImGOidHaAOZn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.kill(os.getpid(), 9) # Crash colab if runs out of gpu memory / Funny errors (Run from Set up again)"
      ],
      "metadata": {
        "id": "WrqxFhj4hbAC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Caption\n",
        "index = 0 #@param {type:\"integer\"}\n",
        "info = f\"seed = {seed_new}, cfg = {scale}, steps = {steps}, sampler = {sampler}\"\n",
        "captioned_image = caption(images[index], prompt, info)\n",
        "captioned_image"
      ],
      "metadata": {
        "id": "lw1Mt9SfDV51"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[0]"
      ],
      "metadata": {
        "id": "dv1QmzMxAbe6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Individual results"
      ],
      "metadata": {
        "id": "NK8ZJXL3DOkY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "images[0]"
      ],
      "metadata": {
        "id": "F4LXUYzHDSGJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[1]"
      ],
      "metadata": {
        "id": "k_lNh25JDRz_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[2]"
      ],
      "metadata": {
        "id": "lmaM9FnfDRry"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[3]"
      ],
      "metadata": {
        "id": "1MWVsZIJDRYU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Image 2 Image\n",
        "\n"
      ],
      "metadata": {
        "id": "L7jnzh7O_vQT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = \"kek, lmao, octane 3d render\" #@param {type:\"string\"}\n",
        "sampler = 'k_euler_a' #@param [\"k_euler_a\",\"k-diffusion\"] {allow-input: false}\n",
        "init_image_path = \"/content/stable-diffusion/Source/take_00474.png\" #@param {type: 'string'}\n",
        "\n",
        "resize_mode = \"Resize and fill\" #@param [\"Just resize\", \"Crop and resize\", \"Resize and fill\"] {allow-input: false}\n",
        "\n",
        "\n",
        "width = 512 #@param {type:\"integer\"}\n",
        "height = 512 #@param {type:\"integer\"}\n",
        "\n",
        "scale = 7 #@param {type:'integer'}\n",
        "steps = 50 #@param {type:\"slider\", min:1, max:150, step:1}\n",
        "strength = 0.7 #@param {type: \"slider\", min:0.00, max:1.00, step:0.01}\n",
        "\n",
        "samples = 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",
        "init_image = Image.open(init_image_path)\n",
        "\n",
        "images, seed_new, mode = img2img(prompt, init_image_path, seed, sampler, steps, scale, strength, \n",
        "                                 samples, rows, height, width, skip_grid, resize_mode)\n"
      ],
      "metadata": {
        "id": "xOcjG58T_wzp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.kill(os.getpid(), 9) # Crash colab if runs out of gpu memory / Funny errors (Run from Set up again)"
      ],
      "metadata": {
        "id": "cunFP4J0hZkt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Caption\n",
        "index = 0 #@param {type:\"integer\"}\n",
        "original = resize_image(mode, init_image, width, height)\n",
        "# if init_mask is not None:\n",
        "#   original = get_concat_v_blank(original, resize_image(mode, init_mask, width, height))\n",
        "concatted = get_concat_h_blank(images[index], original)\n",
        "info = f\"seed = {seed_new}, cfg = {scale}, steps = {steps}, strength = {strength}, sampler = {sampler}\"\n",
        "captioned_image = caption(concatted, prompt, info)\n",
        "captioned_image"
      ],
      "metadata": {
        "id": "EhR5DwxyDdgr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[0]"
      ],
      "metadata": {
        "id": "D_XNFzLXAZr2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Individual results"
      ],
      "metadata": {
        "id": "NBtpXi1JDXCv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "images[1]"
      ],
      "metadata": {
        "id": "o1XLood5DW4J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[2]"
      ],
      "metadata": {
        "id": "B-N5e_29DZvC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[3]"
      ],
      "metadata": {
        "id": "l7NI8r4kDa5q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[4]"
      ],
      "metadata": {
        "id": "5d1Bu8ZQDbvC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# hlky fork (no GFPGAN) - Gradio-less (Tested as of 8/27)\n",
        "\n",
        "Has K-diffusion and more features but uses more VRAM + Mask support\n",
        "\n",
        "https://github.com/hlky/stable-diffusion"
      ],
      "metadata": {
        "id": "YbgZNNTxLOEo"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Set up"
      ],
      "metadata": {
        "id": "fVPP3LONLmdN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/stable-diffusion"
      ],
      "metadata": {
        "id": "SqqSCkAJ-Uwq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import k_diffusion as K\n",
        "import math\n",
        "import numpy as np\n",
        "import random\n",
        "\n",
        "import time\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from typing import List, Union\n",
        "\n",
        "from contextlib import contextmanager, nullcontext\n",
        "from einops import rearrange, repeat\n",
        "from itertools import islice\n",
        "from omegaconf import OmegaConf\n",
        "from PIL import Image, ImageFont, ImageDraw, ImageFilter, ImageOps\n",
        "from PIL.PngImagePlugin import PngInfo\n",
        "\n",
        "from torch import autocast\n",
        "# from ldm.models.diffusion.ddim import DDIMSampler\n",
        "# from ldm.models.diffusion.plms import PLMSSampler\n",
        "from ldm.util import instantiate_from_config"
      ],
      "metadata": {
        "id": "y0Ss3f7ViDJJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def chunk(it, size):\n",
        "    it = iter(it)\n",
        "    return iter(lambda: tuple(islice(it, size)), ())\n",
        "\n",
        "def torch_gc():\n",
        "    torch.cuda.empty_cache()\n",
        "    torch.cuda.ipc_collect()\n",
        "\n",
        "# def load_img_pil(img_pil):\n",
        "#     image = img_pil.convert(\"RGB\")\n",
        "#     w, h = image.size\n",
        "#     print(f\"loaded input image of size ({w}, {h})\")\n",
        "#     w, h = map(lambda x: x - x % 64, (w, h))  # resize to integer multiple of 64\n",
        "#     image = image.resize((w, h), resample=PIL.Image.LANCZOS)\n",
        "#     print(f\"cropped image to size ({w}, {h})\")\n",
        "#     image = np.array(image).astype(np.float32) / 255.0\n",
        "#     image = image[None].transpose(0, 3, 1, 2)\n",
        "#     image = torch.from_numpy(image)\n",
        "#     return 2. * image - 1.\n",
        "\n",
        "\n",
        "# def load_img(path):\n",
        "#     return load_img_pil(Image.open(path))\n",
        "\n",
        "class CFGDenoiser(nn.Module):\n",
        "    def __init__(self, model):\n",
        "        super().__init__()\n",
        "        self.inner_model = model\n",
        "\n",
        "    def forward(self, x, sigma, uncond, cond, cond_scale):\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",
        "        \n",
        "        return uncond + (cond - uncond) * cond_scale\n",
        "\n",
        "class KDiffusionSampler:\n",
        "    def __init__(self, m, sampler):\n",
        "        self.model = m\n",
        "        self.model_wrap = K.external.CompVisDenoiser(m)\n",
        "        self.schedule = sampler\n",
        "\n",
        "    def sample(self, S, conditioning, batch_size, shape, verbose, unconditional_guidance_scale, unconditional_conditioning, eta, x_T):\n",
        "        sigmas = self.model_wrap.get_sigmas(S)\n",
        "        x = x_T * sigmas[0]\n",
        "        model_wrap_cfg = CFGDenoiser(self.model_wrap)\n",
        "        samples_ddim = K.sampling.sample_lms(model_wrap_cfg, x, sigmas, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': unconditional_guidance_scale}, disable=False)\n",
        "\n",
        "        return samples_ddim, None\n",
        "\n",
        "def create_random_tensors(shape, seeds):\n",
        "    xs = []\n",
        "    for seed in seeds:\n",
        "        torch.manual_seed(seed)\n",
        "\n",
        "        xs.append(torch.randn(shape, device=device))\n",
        "    x = torch.stack(xs)\n",
        "    return x\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"
      ],
      "metadata": {
        "id": "3I78LM7iLX4I"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Updates"
      ],
      "metadata": {
        "id": "gkLTqrwyRcAy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## UPDATE\n",
        "\n",
        "LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)\n",
        "invalid_filename_chars = '<>:\"/\\|?*\\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",
        "      if im.width != width or im.height != height:\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",
        "        else:\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",
        "      else:\n",
        "        return im\n",
        "\n",
        "    return res"
      ],
      "metadata": {
        "id": "e3PYezWfOUYm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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/NotoSansJP-Bold.otf\", 20, encoding='utf-8')\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"
      ],
      "metadata": {
        "id": "vqRs04yoRett"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Load model"
      ],
      "metadata": {
        "id": "na17SiCBMY9Y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def load_model_from_config(config, ckpt, verbose=False):\n",
        "    print(f\"Loading model from {ckpt}\")\n",
        "    ckpt = '/content/stable-diffusion/model.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",
        "    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.cuda()\n",
        "    model.eval()\n",
        "    return model"
      ],
      "metadata": {
        "id": "miGskb1_l511"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "7CebqbB-l-9t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def load_model_from_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",
        "    return sd"
      ],
      "metadata": {
        "id": "_5kHmX0nnXZ7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "config = \"/content/stable-diffusion/optimizedSD/v1-inference.yaml\"\n",
        "ckpt = \"model.ckpt\"\n",
        "device = \"cuda\"\n",
        "sd = load_model_from_config(f\"{ckpt}\")\n",
        "li = []\n",
        "lo = []\n",
        "for key, value in sd.items():\n",
        "    sp = key.split('.')\n",
        "    if(sp[0]) == 'model':\n",
        "        if('input_blocks' in sp):\n",
        "            li.append(key)\n",
        "        elif('middle_block' in sp):\n",
        "            li.append(key)\n",
        "        elif('time_embed' in sp):\n",
        "            li.append(key)\n",
        "        else:\n",
        "            lo.append(key)\n",
        "for key in li:\n",
        "    sd['model1.' + key[6:]] = sd.pop(key)\n",
        "for key in lo:\n",
        "    sd['model2.' + key[6:]] = sd.pop(key)\n",
        "\n",
        "config = OmegaConf.load(f\"{config}\")\n",
        "\n",
        "\n",
        "model = instantiate_from_config(config.modelUNet)\n",
        "_, _ = model.load_state_dict(sd, strict=False)\n",
        "model.eval()\n",
        "    \n",
        "modelCS = instantiate_from_config(config.modelCondStage)\n",
        "_, _ = modelCS.load_state_dict(sd, strict=False)\n",
        "modelCS.cond_stage_model.device = device\n",
        "modelCS.eval()\n",
        "    \n",
        "modelFS = instantiate_from_config(config.modelFirstStage)\n",
        "_, _ = modelFS.load_state_dict(sd, strict=False)\n",
        "modelFS.eval()\n",
        "\n",
        "model.unet_bs = True\n",
        "model.turbo = True\n",
        "model.cdevice = device\n",
        "modelCS.cond_stage_model.device = device\n",
        "model.to(device)"
      ],
      "metadata": {
        "id": "S28R0KC8l3G2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "modelCS.to(device)\n",
        "modelFS.to(device)"
      ],
      "metadata": {
        "id": "CVX4v83todfl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def load_model_from_config(config, ckpt, verbose=False):\n",
        "    print(f\"Loading model from {ckpt}\")\n",
        "    ckpt = '/content/stable-diffusion/model.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",
        "    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.cuda()\n",
        "    model.eval()\n",
        "    return model\n",
        "\n",
        "config = OmegaConf.load(\"configs/stable-diffusion/v1-inference.yaml\")\n",
        "model = load_model_from_config(config, \"models/ldm/stable-diffusion-v1/model.ckpt\")"
      ],
      "metadata": {
        "id": "C8J52y2xMcti"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Set generators"
      ],
      "metadata": {
        "id": "occevEG_kUXk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "opt_C = 4\n",
        "opt_f = 8\n",
        "# device = torch.device(\"cuda\")\n",
        "# model = model.half().to(device)"
      ],
      "metadata": {
        "id": "wkAr1hKRMLJ8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Info:\n",
        "    def __init__(self, prompt, seed, sampler, ddim_steps, cfg_scale, strength, n_samples, n_rows, height, width, skip_grid, n_iter, resize_mode):\n",
        "        self.prompt = prompt\n",
        "        self.seed = seed\n",
        "\n",
        "        self.sampler = sampler\n",
        "        self.ddim_steps = ddim_steps\n",
        "        self.cfg_scale = cfg_scale\n",
        "        self.strength = strength\n",
        "\n",
        "        self.n_samples = n_samples\n",
        "        self.n_rows = n_rows\n",
        "\n",
        "\n",
        "        self.height = height\n",
        "        self.width = width\n",
        "\n",
        "        self.skip_grid = skip_grid\n",
        "        self.n_iter = n_iter\n",
        "        self.resize_mode = resize_mode\n"
      ],
      "metadata": {
        "id": "nsnYNQyE-krD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def generate(user: Info, mode: str, init_image, init_mask, func_init, func_sample, sampler_name):\n",
        "  torch_gc()\n",
        "\n",
        "  device = \"cuda\"\n",
        "  model.half()\n",
        "  modelCS.half()\n",
        "  modelFS.half()\n",
        "  assert user.prompt is not None\n",
        " \n",
        "  if user.seed == -1:\n",
        "    user.seed = random.randrange(4294967294)\n",
        "\n",
        "  init_seed = user.seed\n",
        "\n",
        "  batch_size = user.n_samples\n",
        "  all_prompts = batch_size * user.n_iter * [prompt]\n",
        "  all_seeds = [user.seed + x for x in range(len(all_prompts))]\n",
        "\n",
        "  # modelFS.to(device)\n",
        "\n",
        "  output_images = []\n",
        "  precision_scope = autocast\n",
        "  flag_img2img = False\n",
        "  if mode == \"img2img\":\n",
        "    flag_img2img = True\n",
        "    \n",
        "  with torch.no_grad(), precision_scope(\"cuda\"):#, model.ema_scope():\n",
        "\n",
        "    modelCS.to(device)\n",
        "    modelFS.to(device)\n",
        "    init_data = func_init()\n",
        "    \n",
        "    for n in range(user.n_iter):\n",
        "      prompts = all_prompts[n * batch_size:(n + 1) * batch_size]\n",
        "      seeds = all_seeds[n * batch_size:(n + 1) * batch_size]\n",
        "\n",
        "      uc = None\n",
        "\n",
        "      if user.cfg_scale != 1.0:\n",
        "        uc = modelCS.get_learned_conditioning(batch_size * [\"\"])\n",
        "\n",
        "      if isinstance(prompts, tuple):\n",
        "        prompts = list(prompts)\n",
        "\n",
        "      c = modelCS.get_learned_conditioning(prompts)\n",
        "      modelCS.to(\"cpu\")\n",
        "\n",
        "      shape = [opt_C, height // opt_f, width // opt_f]\n",
        "      x = create_random_tensors([opt_C, height // opt_f, width // opt_f], seeds=seeds)\n",
        "      samples_ddim = func_sample(init_data=init_data, x=x, conditioning=c, unconditional_conditioning=uc, sampler_name=sampler_name)\n",
        "\n",
        "      x_samples_ddim = modelFS.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",
        "\n",
        "      for i, x_sample in enumerate(x_samples_ddim):\n",
        "        x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')\n",
        "        x_sample = x_sample.astype(np.uint8)\n",
        "\n",
        "        image = Image.fromarray(x_sample)\n",
        "        \n",
        "        if flag_img2img and init_mask is not None:\n",
        "          init_mask = init_mask.filter(ImageFilter.GaussianBlur(3))\n",
        "          init_mask = init_mask.convert('L')\n",
        "          init_image = init_image.convert('RGB')\n",
        "          image = image.convert('RGB')\n",
        "          image = Image.composite(init_image, image, init_mask)\n",
        "\n",
        "        output_images.append(image)\n",
        "        modelFS.to(\"cpu\")\n",
        "        torch_gc()\n",
        "        \n",
        "    if not user.skip_grid:\n",
        "      # additionally, save as grid\n",
        "      grid = image_grid(output_images, batch_size, user.n_rows)\n",
        "      output_images.insert(0, grid)\n",
        "\n",
        "  return output_images, init_seed\n",
        "  "
      ],
      "metadata": {
        "id": "DsB_jx_QMM9I"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch_gc()"
      ],
      "metadata": {
        "id": "77KTJBr1pdXW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def txt2img(prompt: str, seed: int, sampler_name: str, steps: int, scale: int, strength: int, samples: int, rows: int, height: int, width: int, skip_grid: bool):\n",
        "\n",
        "  if sampler_name == 'k-diffusion':\n",
        "    sampler = KDiffusionSampler(model, 'lms')\n",
        "  elif sampler_name == 'k_euler_a':\n",
        "    sampler = KDiffusionSampler(model,'euler_ancestral')\n",
        "  elif sampler_name == 'k_euler':\n",
        "      sampler = KDiffusionSampler(model,'euler')\n",
        "  elif sampler_name == 'k_heun':\n",
        "      sampler = KDiffusionSampler(model,'heun')\n",
        "  elif sampler_name == 'k_dpm_2_a':\n",
        "      sampler = KDiffusionSampler(model,'dpm_2_ancestral')\n",
        "  elif sampler_name == 'k_dpm_2':\n",
        "      sampler = KDiffusionSampler(model,'dpm_2')\n",
        "  else:\n",
        "      raise Exception(\"Unknown sampler: \" + sampler_name)\n",
        "    \n",
        "\n",
        "  user = Info(prompt, seed, sampler, steps, scale, strength, samples, rows, height, width, skip_grid, 1, 0)\n",
        "\n",
        "  def init():\n",
        "          pass\n",
        "\n",
        "  def sample(init_data, x, conditioning, unconditional_conditioning, sampler_name):\n",
        "    samples_ddim, _ = sampler.sample(S=steps, conditioning=conditioning, batch_size=int(x.shape[0]), shape=x[0].shape, verbose=False, unconditional_guidance_scale=scale, unconditional_conditioning=unconditional_conditioning, eta=0.0, x_T=x)\n",
        "    return samples_ddim\n",
        "  return generate(user, \"txt2img\", None, None, init, sample, sampler_name)"
      ],
      "metadata": {
        "id": "0LVk8_jY-4Dq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def img2img(prompt: str, seed: int, sampler_name: str, steps: int, scale: int, strength: int, samples: int, rows: int, height: int, width: int, skip_grid: bool, init_image, mode: str, mask, invert):\n",
        "\n",
        "  if sampler_name == 'k-diffusion':\n",
        "    sampler = KDiffusionSampler(model, 'lms')\n",
        "  elif sampler_name == 'k_euler_a':\n",
        "    sampler = KDiffusionSampler(model,'euler_ancestral')\n",
        "  elif sampler_name == 'k_euler':\n",
        "      sampler = KDiffusionSampler(model,'euler')\n",
        "  elif sampler_name == 'k_heun':\n",
        "      sampler = KDiffusionSampler(model,'heun')\n",
        "  elif sampler_name == 'k_dpm_2_a':\n",
        "      sampler = KDiffusionSampler(model,'dpm_2_ancestral')\n",
        "  elif sampler_name == 'k_dpm_2':\n",
        "      sampler = KDiffusionSampler(model,'dpm_2')\n",
        "  else:\n",
        "      raise Exception(\"Unknown sampler: \" + sampler_name)\n",
        "\n",
        "  user = Info(prompt, seed, sampler, steps, scale, strength, samples, rows, height, width, skip_grid, 1, mode)\n",
        "\n",
        "  if mode == \"Just resize\":\n",
        "    user.resize_mode = 0\n",
        "  elif mode == \"Crop and resize\":\n",
        "    user.resize_mode = 1\n",
        "  else:\n",
        "    user.resize_mode = 2\n",
        "\n",
        "  init_image = init_image.convert(\"RGB\")\n",
        "  init_image = resize_image(user.resize_mode, init_image, 512, 512)\n",
        "  \n",
        "  try:\n",
        "    init_mask = mask\n",
        "    if invert:\n",
        "      init_mask = ImageOps.invert(init_mask)\n",
        "    init_mask = init_mask.convert(\"RGB\")\n",
        "    init_mask = resize_image(user.resize_mode, init_mask, user.width, user.height)\n",
        "  except:\n",
        "    init_mask = None\n",
        "\n",
        "  assert 0. <= strength <= 1., 'can only work with strength in [0.0, 1.0]'\n",
        "  t_enc = int(strength * user.ddim_steps)\n",
        "\n",
        "  def init():\n",
        "    image = init_image.convert(\"RGB\")\n",
        "    image = resize_image(user.resize_mode, image, user.width, user.height)\n",
        "    image = np.array(image).astype(np.float32) / 255.0\n",
        "    image = image[None].transpose(0, 3, 1, 2)\n",
        "    image = torch.from_numpy(image)\n",
        "\n",
        "    original = 2. * image - 1.\n",
        "    original = original.to(device)\n",
        "    original = repeat(original, '1 ... -> b ...', b=samples)\n",
        "    modelFS.to(device)\n",
        "    init_latent = modelFS.get_first_stage_encoding(modelFS.encode_first_stage(original))  # move to latent space\n",
        "    # modelFS.to(\"cpu\")\n",
        "    return init_latent,\n",
        "\n",
        "\n",
        "  def sample(init_data, x, conditioning, unconditional_conditioning, sampler_name):\n",
        "        x0, = init_data\n",
        "\n",
        "        sigmas = sampler.model_wrap.get_sigmas(steps)\n",
        "        noise = x * sigmas[user.ddim_steps - t_enc - 1]\n",
        "\n",
        "        xi = x0 + noise\n",
        "        sigma_sched = sigmas[user.ddim_steps - t_enc - 1:]\n",
        "        model_wrap_cfg = CFGDenoiser(sampler.model_wrap)\n",
        "        samples_ddim = K.sampling.sample_lms(model_wrap_cfg, xi, sigma_sched, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': scale}, disable=False)\n",
        "       \n",
        "        return samples_ddim\n",
        "\n",
        "  return generate(user, \"img2img\", init_image, init_mask, init, sample, sampler_name) + (user.resize_mode, )"
      ],
      "metadata": {
        "id": "n5pvbz9a-49U"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "modelFS.device"
      ],
      "metadata": {
        "id": "Fiei4DMptTS9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Text 2 Image\n",
        "\n",
        "512x512 works well, 640x640 might crash colab"
      ],
      "metadata": {
        "id": "thSWMltDLYNy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = \"www, octane 3D render\" #@param {type:\"string\"}\n",
        "sampler = 'k_dpm_2' #@param [\"k_euler_a\",\"k-diffusion\", \"k_dpm_2\", \"k_dpm_2_a\", \"k_euler\", \"k_heun\"] {allow-input: false}\n",
        "\n",
        "width = 512 #@param {type:\"integer\"}\n",
        "height = 512 #@param {type:\"integer\"}\n",
        "\n",
        "scale = 7 #@param {type:'integer'}\n",
        "steps = 9 #@param {type:\"slider\", min:1, max:150, step:1}\n",
        "\n",
        "samples = 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",
        "## Seed = -1 ==> Random\n",
        "images, seed_new = txt2img(prompt, seed, sampler, steps, scale, 0, samples, rows, height, width, skip_grid)\n"
      ],
      "metadata": {
        "id": "g0Sy0dk0Mkkg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Caption\n",
        "index = 0 #@param {type:\"integer\"}\n",
        "info = f\"seed = {seed_new}, cfg = {scale}, steps = {steps}, sampler = {sampler}\"\n",
        "captioned_image = caption(images[index], prompt, info)\n",
        "captioned_image"
      ],
      "metadata": {
        "id": "e5i1NKmCU8z8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.kill(os.getpid(), 9) # Clear GPU if OOM (Run from set up again)"
      ],
      "metadata": {
        "id": "F73YlrzHQxNr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[0]"
      ],
      "metadata": {
        "id": "akuObYB5Mv8I"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Individual outputs"
      ],
      "metadata": {
        "id": "UeAcuHx-Em0U"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "images[1]"
      ],
      "metadata": {
        "id": "VwiRZr9xMxd0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[2]"
      ],
      "metadata": {
        "id": "gaR2KDdoMy7w"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[3]"
      ],
      "metadata": {
        "id": "7-YR0dsdM0JB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[4]"
      ],
      "metadata": {
        "id": "cusiIi8SM1Jw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Image 2 Image"
      ],
      "metadata": {
        "id": "NwsSUAbMLe04"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = \"kek, lmao, octane 3d render\" #@param {type:\"string\"}\n",
        "ssampler = 'k_euler' #@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/Source/take_00474.png\" #@param {type: 'string'}\n",
        "\n",
        "#@markdown White = Original, Black = SD, Leave empty for no masking\n",
        "init_mask_path = \"/content/stable-diffusion/Source/take_00474.png\" #@param {type: 'string'}\n",
        "invert = True #@param {type: 'boolean'}\n",
        "resize_mode = \"Resize and fill\" #@param [\"Just resize\", \"Crop and resize\", \"Resize and fill\"] {allow-input: false}\n",
        "\n",
        "\n",
        "width = 512 #@param {type:\"integer\"}\n",
        "height = 512 #@param {type:\"integer\"}\n",
        "\n",
        "scale = 7 #@param {type:'integer'}\n",
        "steps = 11 #@param {type:\"slider\", min:1, max:150, step:1}\n",
        "strength = 0.7 #@param {type: \"slider\", min:0.00, max:1.00, step:0.01}\n",
        "\n",
        "samples = 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",
        "## Seed = -1 ==> Random\n",
        "init_image = Image.open(init_image_path)\n",
        "try:\n",
        "  init_mask = Image.open(init_mask_path)\n",
        "  init_mask = resize_image(resize_mode, init_mask, init_image.width, init_image.height)\n",
        "except:\n",
        "  init_mask = None\n",
        "\n",
        "\n",
        "images, seed_new, mode = img2img(prompt, seed, sampler, steps, scale, strength, samples, rows, height, width, skip_grid, init_image, resize_mode, init_mask, invert)\n"
      ],
      "metadata": {
        "id": "SHOh80nzT4jH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Caption\n",
        "index = 0 #@param {type:\"integer\"}\n",
        "original = resize_image(mode, init_image, width, height)\n",
        "if init_mask is not None:\n",
        "  original = get_concat_v_blank(original, resize_image(mode, init_mask, width, height))\n",
        "concatted = get_concat_h_blank(images[index], original)\n",
        "info = f\"seed = {seed_new}, cfg = {scale}, steps = {steps}, strength = {strength}, sampler = {sampler}\"\n",
        "captioned_image = caption(concatted, prompt, info)\n",
        "captioned_image"
      ],
      "metadata": {
        "id": "AfXX5n1TmBEg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "os.kill(os.getpid(), 9) # Clear GPU if OOM or funny errors (Run from set up again)"
      ],
      "metadata": {
        "id": "mcDhMNZGVoyf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[0]"
      ],
      "metadata": {
        "id": "LLKFosZiVpqm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Individual results"
      ],
      "metadata": {
        "id": "7fplwcZ_EZaZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "images[1]"
      ],
      "metadata": {
        "id": "sEG1FyrGVpgz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[2]"
      ],
      "metadata": {
        "id": "2e1NLC_BVpZA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[3]"
      ],
      "metadata": {
        "id": "1b6Qdz8BVpRz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images[4]"
      ],
      "metadata": {
        "id": "xytXKAcmVpJe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Image 2 Image (Sequence)"
      ],
      "metadata": {
        "id": "4D_g45M-F9e2"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Running I2I Sequence"
      ],
      "metadata": {
        "id": "8lQVGJJFGGcK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from PIL import Image\n",
        "from IPython.display import clear_output \n",
        "import numpy as np\n",
        "import timeit\n",
        "#@title Image sequence\n",
        "#@markdown ##### Enter text prompt:\n",
        "prompt = \"man standing in front of a painting of a mountain in the distance, wide angle, Van Gogh, futuristic, surreal, holy\" #@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",
        "#@markdown ---\n",
        "#@markdown ##### Set sequence duration (/Source/prefix_XXXXX.png):\n",
        "start_frame = \"00024\" #@param {type:\"string\"}\n",
        "end_frame = \"00028\" #@param {type:\"string\"}\n",
        "prefix = \"take\" #@param {type:\"string\"}\n",
        "current_frame = start_frame\n",
        "#@markdown ---\n",
        "#@markdown ##### SD settings:\n",
        "\n",
        "resize_mode = \"Resize and fill\" #@param [\"Just resize\", \"Crop and resize\", \"Resize and fill\"] {allow-input: false}\n",
        "\n",
        "\n",
        "width = 512 #@param {type:\"integer\"}\n",
        "height = 512 #@param {type:\"integer\"}\n",
        "\n",
        "scale = 15 #@param {type:'integer'}\n",
        "steps = 50 #@param {type:\"slider\", min:1, max:150, step:1}\n",
        "strength = 0.7 #@param {type: \"slider\", min:0.00, max:1.00, step:0.01}\n",
        "\n",
        "seed = -1 #@param {type:'integer'}\n",
        "## Seed = -1 ==> Random\n",
        "timeStart = timeit.default_timer()\n",
        "for i in range(int(start_frame), int(end_frame) + 1):\n",
        "  current_frame = \"{:05d}\".format(i)\n",
        "  print(\"PROCESSING \" + current_frame)\n",
        "\n",
        "  init_image_path = f\"/content/stable-diffusion/Source/{prefix}_{current_frame}.png\"\n",
        "  init_image = Image.open(init_image_path)\n",
        "\n",
        "  images, seed_new, mode = img2img(prompt, seed, sampler, steps, scale, strength, 1, 1, height, width, True, init_image, resize_mode, None, False)\n",
        "  images[0].save(f\"/content/stable-diffusion/Output/{prefix}_{current_frame}.png\")\n",
        "  images[0]\n",
        "  \n",
        "  clear_output()\n",
        "timeStop = timeit.default_timer()\n",
        "\n",
        "print(\"{} frames done in {} seconds.\\nPrompt: {}\\nStrength: {}\".format(int(end_frame) - int(start_frame), timeStop - timeStart, prompt, strength))"
      ],
      "metadata": {
        "id": "-LUoUqavGM2Y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Export video"
      ],
      "metadata": {
        "id": "mtsnExYRGKBO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Export video\n",
        "from google.colab import files\n",
        "file_name = \"out.mp4\" #@param {type: \"string\"}\n",
        "always_overwrite = True #@param {type:\"boolean\"}\n",
        "framerate = 30 #@param (type:number)\n",
        "if always_overwrite:\n",
        "  replace = 'y'\n",
        "else:\n",
        "  replace = 'n'\n",
        "!yes {replace} | ffmpeg -framerate {framerate} -start_number {int(start_frame)} -i /content/stable-diffusion/Output/{prefix}_%05d.png -c:v libx264 -crf 0 {file_name}\n",
        "files.download('/content/stable-diffusion/{}'.format(file_name))"
      ],
      "metadata": {
        "id": "7rK2voJIHOt4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Saving outputs"
      ],
      "metadata": {
        "id": "uZ13WzesoXdq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "path = \"/content/stable-diffusion/Output/\"\n",
        "name = \"out\" #@param {type:\"string\"}\n",
        "\n",
        "save_all = True #@param {type:\"boolean\"}\n",
        "\n",
        "if save_all:\n",
        "  k = 0\n",
        "  for i in images:\n",
        "    i.save(f'{path}{name}_{k}.png')\n",
        "    k += 1\n",
        "else:\n",
        "  index = 1 #@param {type:\"integer\"}\n",
        "  images[index].save(f'{path}{name}_{index}.png')"
      ],
      "metadata": {
        "id": "W_xIkXoMB9iQ"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "XvHTXI7KOnu4",
        "yaHfjOHyJgdC",
        "Oudkt8AEKw05",
        "OmQfJWm8QGoA",
        "-7syo80FPa4T",
        "ihQD4BD1P1Vr",
        "ie9xTKWi8b1e",
        "P95iTXcZQbnk",
        "IDQmNv_9ZTGc",
        "c6jepvv8Q-1a",
        "IFdeC2LZRBm-",
        "7oPo2rboRWfH",
        "zUGoyvzD86IL",
        "NMj_JbQ9Rrxq",
        "iIBkjkG_dHAu",
        "5kQ2bDBrPe6Q",
        "CszrKJDe-66T",
        "5deJM1EQ_BIC",
        "YDcdDSiX_HAo",
        "4y4c6LWDAAVB",
        "9lnxMQBZ_asY",
        "NK8ZJXL3DOkY",
        "L7jnzh7O_vQT",
        "NBtpXi1JDXCv",
        "YbgZNNTxLOEo",
        "fVPP3LONLmdN",
        "na17SiCBMY9Y",
        "thSWMltDLYNy",
        "UeAcuHx-Em0U",
        "NwsSUAbMLe04",
        "7fplwcZ_EZaZ",
        "uZ13WzesoXdq"
      ],
      "name": "Neo Hidamari Diffusion.ipynb",
      "provenance": [],
      "private_outputs": true,
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}