360Diffusion_Public.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sadnow/360Diffusion/blob/main/360Diffusion_Public.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tkav03SvGq7g"
      },
      "source": [
        "#360Diffusion - ESRGAN-Embedded Fast CLIP Guided Diffusion\n",
        "\n",
        "This is `sadnow`'s stable fork of `Daniel Russ`’s fork of `Katherine Crowson`’s CLIP Guided Diffusion notebook! Daniel also credits `Dango233` and `nsheppard` for Quick Diffusion developments.\n",
        "\n",
        "###This fork’s main feature is Real-ESRGAN integration! Automatically upscales from 256px or 512px all the way up to 4096px!\n",
        "---\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "###Huge thanks to:\n",
        "\n",
        "- `Katherine Crowson` (https://github.com/crowsonkb, https://twitter.com/RiversHaveWings)\n",
        "- `Daniel Russell` (https://github.com/russelldc, https://twitter.com/danielrussruss)\n",
        "- VQLIPSE Dev & his Discord (https://www.patreon.com/sportsracer48)\n",
        "\n",
        "---\n",
        "**Latest Update:** Alpha 1.61 [Main Branch] - 01/11/22\n",
        "- Layout adjustments\n",
        "- Added seed handling for users\n",
        "- Misc: dprint function, changed default prompts\n",
        "\n",
        "\n",
        "Updates can be found at https://github.com/sadnow/ESRGAN-UltraFast-CLIP-Guided-Diffusion-Colab if available.\n",
        "\n",
        "Full documentation can be found at https://thisnameismine.com\n",
        "\n",
        "Sadnow also occasionally posts developments on https://twitter.com/sadly_existent\n",
        "\n",
        "---\n",
        "##***See the bottom of the notebook for documentation and examples-***\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "ztPt3l_I5CLU"
      },
      "source": [
        "# @title Licensed under the MIT License\n",
        "\n",
        "# Copyright (c) 2021 Katherine Crowson\n",
        "\n",
        "# Permission is hereby granted, free of charge, to any person obtaining a copy\n",
        "# of this software and associated documentation files (the \"Software\"), to deal\n",
        "# in the Software without restriction, including without limitation the rights\n",
        "# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
        "# copies of the Software, and to permit persons to whom the Software is\n",
        "# furnished to do so, subject to the following conditions:\n",
        "\n",
        "# The above copyright notice and this permission notice shall be included in\n",
        "# all copies or substantial portions of the Software.\n",
        "\n",
        "# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
        "# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
        "# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
        "# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
        "# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
        "# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\n",
        "# THE SOFTWARE."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "qZ3rNuAWAewx"
      },
      "outputs": [],
      "source": [
        "#@markdown ##Check GPU\n",
        "#@markdown ***This Notebook supports \"Run All. Once everything is set to your liking, you can do `Ctrl+F9`!***\n",
        "import torch\n",
        "# Check the GPU status\n",
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "print('Using device:', device)\n",
        "\n",
        "# a100_support = False #@param {type:\"boolean\"}\n",
        "\n",
        "enable_error_checking = False #saves ram\n",
        "if enable_error_checking:\n",
        "  !nvidia-smi\n",
        "else:\n",
        "  !nvidia-smi\n",
        "  !nvidia-smi -i 0 -e 0\n",
        "#@markdown *Special thanks to sportracer48 and his Discord channel. If you want to make AI animations, he has the meats in closed beta.* https://www.patreon.com/sportsracer48\n",
        "# if a100_support:\n",
        "#   print(\"If you have an A100: Turn your cutn to 64-96 and keep your batch_n low (1 or 2), and it will be fast!\")\n",
        "#   !pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html\n",
        "\n",
        "# function runAllCells() {\n",
        "#   const F9Event = {key: \"F9\", code: \"F9\", metaKey: true, keyCode: 120};\n",
        "#   document.dispatchEvent(new KeyboardEvent(\"keydown\", F9Event));\n",
        "# } #https://stackoverflow.com/questions/65984431/run-all-cells-command-in-google-colab-programmatically"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "yZsjzwS0YGo6"
      },
      "outputs": [],
      "source": [
        "def dir_make(a):\n",
        "  import os.path\n",
        "  from os import path\n",
        "  if not path.exists(a):\n",
        "    print(\"Creating\"+a+\"...\")\n",
        "    !mkdir -p $a\n",
        "\n",
        "def dir_check(a):\n",
        "  import os.path\n",
        "  from os import path\n",
        "  if not path.exists(a):\n",
        "    return False\n",
        "  else:\n",
        "    return True\n",
        "\n",
        "def file_check(a):\n",
        "  import os.path\n",
        "  from os import path\n",
        "  if not path.isfile(a):\n",
        "    return False\n",
        "  else:\n",
        "    return True\n",
        "\n",
        "def image_upscale(model_path,scale,input,output):\n",
        "    %cd /content/Real-ESRGAN/\n",
        "    !python inference_realesrgan.py --model_name $_upscale_model --netscale $scale --face_enhance --input $input --output $output --ext jpg\n",
        "    %cd /content/\n",
        "\n",
        "from google.colab import drive\n",
        "\n",
        "#@title Choose model here:\n",
        "diffusion_model = \"512x512_diffusion_uncond_finetune_008100\" #@param [\"256x256_diffusion_uncond\", \"512x512_diffusion_uncond_finetune_008100\"]\n",
        "\n",
        "#@markdown `512 diffusion model` is new and has very little testing behind it. Play with `cutn` and `cutn_batches` to fix VRAM errors.\n",
        "#diffusion_model = \"256x256_diffusion_uncond\"\n",
        "\n",
        "#@markdown If you connect your Google Drive, you can save the final image of each run on your drive.\n",
        "\n",
        "google_drive = True #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown You can use your mounted Google Drive to load the model checkpoint file if you've already got a copy downloaded there. This will save time (and resources!) when you re-visit this notebook in the future.\n",
        "\n",
        "#@markdown Click here if you'd like to save the diffusion model checkpoint file to (and/or load from) your Google Drive:\n",
        "yes_please = True #@param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "\n",
        "_drive_location = '/content/drive/MyDrive/AI/360Diffusion/' #@param{type:\"string\"}\n",
        "contains_slash = (_drive_location.endswith('/'))\n",
        "if not contains_slash:\n",
        "  _drive_location = _drive_location + '/'\n",
        "\n",
        "\n",
        "#####################################################################\n",
        "##@title Google Drive & Download diffusion model\n",
        "\n",
        "model_path = '/content/'\n",
        "if google_drive:\n",
        "    from google.colab import drive\n",
        "    drive.mount('/content/drive')\n",
        "    if yes_please:\n",
        "        dir_make(_drive_location)\n",
        "        model_path = _drive_location\n",
        "\n",
        "\n",
        "if diffusion_model == '256x256_diffusion_uncond':\n",
        "    !wget --continue 'https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_diffusion_uncond.pt' -P {model_path}\n",
        "elif diffusion_model == '512x512_diffusion_uncond_finetune_008100':\n",
        "    !wget --continue 'https://huggingface.co/lowlevelware/512x512_diffusion_unconditional_ImageNet/resolve/main/512x512_diffusion_uncond_finetune_008100.pt' -P {model_path}\n",
        "if google_drive and not yes_please:\n",
        "    model_path = _drive_location\n",
        "\n",
        "#https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WnN6j5WzvFHg"
      },
      "source": [
        "# Install and import dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-_UVMZCIAq_r"
      },
      "outputs": [],
      "source": [
        "def install_ESRGAN():\n",
        "    %cd /content/\n",
        "    print(\"Installing libraries for Real-ESRGAN upscaling.\")\n",
        "    !git clone https://github.com/xinntao/Real-ESRGAN.git\n",
        "    %cd /content/Real-ESRGAN\n",
        "    !pip install basicsr -q\n",
        "    !pip install facexlib -q\n",
        "    !pip install gfpgan -q\n",
        "    !pip install -r requirements.txt -q\n",
        "    %cd /content/Real-ESRGAN\n",
        "    !python setup.py develop -q\n",
        "    !wget -nc https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P experiments/pretrained_models\n",
        "    !wget -nc https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth -P experiments/pretrained_models\n",
        "    !wget -nc https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth -P experiments/pretrained_models\n",
        "    !wget -nc https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth -P experiments/pretrained_models\n",
        "    !wget -nc https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth -P experiments/pretrained_models  #regular esrgan\n",
        "\n",
        "\n",
        "    print(\"Finished Installing libraries for Real-ESRGAN upscaling.\")\n",
        "    %cd /content/\n",
        "    #\n",
        "\n",
        "#@markdown ##Git and pip install\n",
        "!git clone https://github.com/openai/CLIP\n",
        "!git clone https://github.com/crowsonkb/guided-diffusion\n",
        "!pip install -e ./CLIP\n",
        "!pip install -e ./guided-diffusion\n",
        "!pip install lpips datetime\n",
        "install_ESRGAN()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "JmbrcrhpBPC6"
      },
      "outputs": [],
      "source": [
        "#@markdown ##imports\n",
        "import time\n",
        "import gc\n",
        "import io\n",
        "import math\n",
        "import sys\n",
        "from IPython import display\n",
        "import lpips\n",
        "from PIL import Image, ImageOps\n",
        "import requests\n",
        "import torch\n",
        "from torch import nn\n",
        "from torch.nn import functional as F\n",
        "import torchvision.transforms as T\n",
        "import torchvision.transforms.functional as TF\n",
        "from tqdm.notebook import tqdm\n",
        "sys.path.append('./CLIP')\n",
        "sys.path.append('./guided-diffusion')\n",
        "import clip\n",
        "from guided_diffusion.script_util import create_model_and_diffusion, model_and_diffusion_defaults\n",
        "from datetime import datetime #filename\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import random\n",
        "\n",
        "#import subprocess #future implementation\n",
        "import os"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FpZczxnOnPIU",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title  { form-width: \"100px\" }\n",
        "\n",
        "# https://gist.github.com/adefossez/0646dbe9ed4005480a2407c62aac8869\n",
        "\n",
        "def dprint(input_str): #debug print\n",
        "    global msg_runtime\n",
        "    print(\"\\033[1;31;40m\",input_str,\"\\033[0m\")\n",
        "    msg_runtime = msg_runtime + input_str + '\\n'\n",
        "    return msg_runtime\n",
        "\n",
        "def add_command(var,string):\n",
        "  var = (var + string + ' ')\n",
        "  return var\n",
        "\n",
        "def image_resize(filepath,width):\n",
        "  from PIL import Image\n",
        "  basewidth = width\n",
        "  img = Image.open(filepath)\n",
        "  wpercent = (basewidth/float(img.size[0]))\n",
        "  hsize = int((float(img.size[1])*float(wpercent)))\n",
        "  #img = img.resize((basewidth,hsize), Image.ANTIALIAS)\n",
        "  if width == 1024: img = img.resize((basewidth,hsize), Image.LANCZOS)\n",
        "  else: img = img.resize((basewidth,hsize), Image.BICUBIC)\n",
        "  img.save(filepath)\n",
        "\n",
        "#####################################################################################################################\n",
        "# https://gist.github.com/adefossez/0646dbe9ed4005480a2407c62aac8869\n",
        "\n",
        "def interp(t):\n",
        "    return 3 * t**2 - 2 * t ** 3\n",
        "\n",
        "def perlin(width, height, scale=10, device=None):\n",
        "    gx, gy = torch.randn(2, width + 1, height + 1, 1, 1, device=device)\n",
        "    xs = torch.linspace(0, 1, scale + 1)[:-1, None].to(device)\n",
        "    ys = torch.linspace(0, 1, scale + 1)[None, :-1].to(device)\n",
        "    wx = 1 - interp(xs)\n",
        "    wy = 1 - interp(ys)\n",
        "    dots = 0\n",
        "    dots += wx * wy * (gx[:-1, :-1] * xs + gy[:-1, :-1] * ys)\n",
        "    dots += (1 - wx) * wy * (-gx[1:, :-1] * (1 - xs) + gy[1:, :-1] * ys)\n",
        "    dots += wx * (1 - wy) * (gx[:-1, 1:] * xs - gy[:-1, 1:] * (1 - ys))\n",
        "    dots += (1 - wx) * (1 - wy) * (-gx[1:, 1:] * (1 - xs) - gy[1:, 1:] * (1 - ys))\n",
        "    return dots.permute(0, 2, 1, 3).contiguous().view(width * scale, height * scale)\n",
        "\n",
        "def perlin_ms(octaves, width, height, grayscale, device=device):\n",
        "    out_array = [0.5] if grayscale else [0.5, 0.5, 0.5]\n",
        "    # out_array = [0.0] if grayscale else [0.0, 0.0, 0.0]\n",
        "    for i in range(1 if grayscale else 3):\n",
        "        scale = 2 ** len(octaves)\n",
        "        oct_width = width\n",
        "        oct_height = height\n",
        "        for oct in octaves:\n",
        "            p = perlin(oct_width, oct_height, scale, device)\n",
        "            out_array[i] += p * oct\n",
        "            scale //= 2\n",
        "            oct_width *= 2\n",
        "            oct_height *= 2\n",
        "    return torch.cat(out_array)\n",
        "\n",
        "def create_perlin_noise(octaves=[1, 1, 1, 1], width=2, height=2, grayscale=True):\n",
        "    out = perlin_ms(octaves, width, height, grayscale)\n",
        "    if grayscale:\n",
        "        out = TF.resize(size=(side_x, side_y), img=out.unsqueeze(0))\n",
        "        out = TF.to_pil_image(out.clamp(0, 1)).convert('RGB')\n",
        "    else:\n",
        "        out = out.reshape(-1, 3, out.shape[0]//3, out.shape[1])\n",
        "        out = TF.resize(size=(side_x, side_y), img=out)\n",
        "        out = TF.to_pil_image(out.clamp(0, 1).squeeze())\n",
        "\n",
        "    out = ImageOps.autocontrast(out)\n",
        "    return out\n",
        "\n",
        "#################################################################################################################\n",
        "def fetch(url_or_path):\n",
        "    if str(url_or_path).startswith('http://') or str(url_or_path).startswith('https://'):\n",
        "        r = requests.get(url_or_path)\n",
        "        r.raise_for_status()\n",
        "        fd = io.BytesIO()\n",
        "        fd.write(r.content)\n",
        "        fd.seek(0)\n",
        "        return fd\n",
        "    return open(url_or_path, 'rb')\n",
        "\n",
        "\n",
        "def parse_prompt(prompt):\n",
        "    if prompt.startswith('http://') or prompt.startswith('https://'):\n",
        "        vals = prompt.rsplit(':', 2)\n",
        "        vals = [vals[0] + ':' + vals[1], *vals[2:]]\n",
        "    else:\n",
        "        vals = prompt.rsplit(':', 1)\n",
        "    vals = vals + ['', '1'][len(vals):]\n",
        "    return vals[0], float(vals[1])\n",
        "\n",
        "def sinc(x):\n",
        "    return torch.where(x != 0, torch.sin(math.pi * x) / (math.pi * x), x.new_ones([]))\n",
        "\n",
        "def lanczos(x, a):\n",
        "    cond = torch.logical_and(-a < x, x < a)\n",
        "    out = torch.where(cond, sinc(x) * sinc(x/a), x.new_zeros([]))\n",
        "    return out / out.sum()\n",
        "\n",
        "def ramp(ratio, width):\n",
        "    n = math.ceil(width / ratio + 1)\n",
        "    out = torch.empty([n])\n",
        "    cur = 0\n",
        "    for i in range(out.shape[0]):\n",
        "        out[i] = cur\n",
        "        cur += ratio\n",
        "    return torch.cat([-out[1:].flip([0]), out])[1:-1]\n",
        "\n",
        "def resample(input, size, align_corners=True):\n",
        "    n, c, h, w = input.shape\n",
        "    dh, dw = size\n",
        "\n",
        "    input = input.reshape([n * c, 1, h, w])\n",
        "\n",
        "    if dh < h:\n",
        "        kernel_h = lanczos(ramp(dh / h, 2), 2).to(input.device, input.dtype)\n",
        "        pad_h = (kernel_h.shape[0] - 1) // 2\n",
        "        input = F.pad(input, (0, 0, pad_h, pad_h), 'reflect')\n",
        "        input = F.conv2d(input, kernel_h[None, None, :, None])\n",
        "\n",
        "    if dw < w:\n",
        "        kernel_w = lanczos(ramp(dw / w, 2), 2).to(input.device, input.dtype)\n",
        "        pad_w = (kernel_w.shape[0] - 1) // 2\n",
        "        input = F.pad(input, (pad_w, pad_w, 0, 0), 'reflect')\n",
        "        input = F.conv2d(input, kernel_w[None, None, None, :])\n",
        "\n",
        "    input = input.reshape([n, c, h, w])\n",
        "    return F.interpolate(input, size, mode='bicubic', align_corners=align_corners)\n",
        "\n",
        "# class MakeCutouts(nn.Module):\n",
        "#     def __init__(self, cut_size, cutn, skip_augs=False):\n",
        "#         super().__init__()\n",
        "#         self.cut_size = cut_size\n",
        "#         self.cutn = cutn\n",
        "#         self.skip_augs = skip_augs\n",
        "#         self.augs = T.Compose([\n",
        "#             T.RandomHorizontalFlip(p=0.5),\n",
        "#             T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
        "#             T.RandomAffine(degrees=15, translate=(0.1, 0.1)),\n",
        "#             T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
        "#             T.RandomPerspective(distortion_scale=0.4, p=0.7),\n",
        "#             T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
        "#             T.RandomGrayscale(p=0.15),\n",
        "#             T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
        "#             # T.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),\n",
        "#         ])\n",
        "\n",
        "class MakeCutouts(nn.Module):\n",
        "    def __init__(self, cut_size, cutn, skip_augs=False):\n",
        "        super().__init__()\n",
        "        self.cut_size = cut_size\n",
        "        self.cutn = cutn\n",
        "        self.skip_augs = skip_augs\n",
        "        self.augs = T.Compose([\n",
        "            T.RandomHorizontalFlip(p=0.5),\n",
        "            # T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
        "            T.RandomAffine(degrees=15, translate=(0.1, 0.1)),\n",
        "            # T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
        "            T.RandomPerspective(distortion_scale=0.4, p=0.7),\n",
        "            # T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
        "            T.RandomGrayscale(p=0.15),\n",
        "            # T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
        "            # T.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),\n",
        "        ])\n",
        "\n",
        "\n",
        "    def forward(self, input):\n",
        "        input = T.Pad(input.shape[2]//4, fill=0)(input)\n",
        "        sideY, sideX = input.shape[2:4]\n",
        "        max_size = min(sideX, sideY)\n",
        "\n",
        "        cutouts = []\n",
        "        for ch in range(cutn):\n",
        "            if ch > cutn - cutn//4:\n",
        "                cutout = input.clone()\n",
        "            else:\n",
        "                size = int(max_size * torch.zeros(1,).normal_(mean=.8, std=.3).clip(float(self.cut_size/max_size), 1.))\n",
        "                offsetx = torch.randint(0, abs(sideX - size + 1), ())\n",
        "                offsety = torch.randint(0, abs(sideY - size + 1), ())\n",
        "                cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size]\n",
        "\n",
        "            if not self.skip_augs:\n",
        "                cutout = self.augs(cutout)\n",
        "            cutouts.append(resample(cutout, (self.cut_size, self.cut_size)))\n",
        "            del cutout\n",
        "\n",
        "        cutouts = torch.cat(cutouts, dim=0)\n",
        "        return cutouts\n",
        "\n",
        "\n",
        "def spherical_dist_loss(x, y):\n",
        "    x = F.normalize(x, dim=-1)\n",
        "    y = F.normalize(y, dim=-1)\n",
        "    return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2)\n",
        "\n",
        "\n",
        "def tv_loss(input):\n",
        "    \"\"\"L2 total variation loss, as in Mahendran et al.\"\"\"\n",
        "    input = F.pad(input, (0, 1, 0, 1), 'replicate')\n",
        "    x_diff = input[..., :-1, 1:] - input[..., :-1, :-1]\n",
        "    y_diff = input[..., 1:, :-1] - input[..., :-1, :-1]\n",
        "    return (x_diff**2 + y_diff**2).mean([1, 2, 3])\n",
        "\n",
        "\n",
        "def range_loss(input):\n",
        "    return (input - input.clamp(-1, 1)).pow(2).mean([1, 2, 3])\n",
        "\n",
        "###\n",
        "# NEWEST PERLIN NOISE EDITS\n",
        "def unitwise_norm(x, norm_type=2.0):\n",
        "    if x.ndim <= 1:\n",
        "        return x.norm(norm_type)\n",
        "    else:\n",
        "        # works for nn.ConvNd and nn,Linear where output dim is first in the kernel/weight tensor\n",
        "        # might need special cases for other weights (possibly MHA) where this may not be true\n",
        "        return x.norm(norm_type, dim=tuple(range(1, x.ndim)), keepdim=True)\n",
        "def adaptive_clip_grad(parameters, clip_factor=0.01, eps=1e-3, norm_type=2.0):\n",
        "    if isinstance(parameters, torch.Tensor):\n",
        "        parameters = [parameters]\n",
        "    for p in parameters:\n",
        "        if p.grad is None:\n",
        "            continue\n",
        "        p_data = p.detach()\n",
        "        g_data = p.grad.detach()\n",
        "        max_norm = unitwise_norm(p_data, norm_type=norm_type).clamp_(min=eps).mul_(clip_factor)\n",
        "        grad_norm = unitwise_norm(g_data, norm_type=norm_type)\n",
        "        clipped_grad = g_data * (max_norm / grad_norm.clamp(min=1e-6))\n",
        "        new_grads = torch.where(grad_norm < max_norm, g_data, clipped_grad)\n",
        "        p.grad.detach().copy_(new_grads)\n",
        "\n",
        "def regen_perlin(): #NEWEST PERLIN UPDATE\n",
        "    if perlin_mode == 'color':\n",
        "        init = create_perlin_noise([1.5**-i*0.5 for i in range(12)], 1, 1, False)\n",
        "        init2 = create_perlin_noise([1.5**-i*0.5 for i in range(8)], 4, 4, False)\n",
        "    elif perlin_mode == 'gray':\n",
        "        init = create_perlin_noise([1.5**-i*0.5 for i in range(12)], 1, 1, True)\n",
        "        init2 = create_perlin_noise([1.5**-i*0.5 for i in range(8)], 4, 4, True)\n",
        "    else:\n",
        "        init = create_perlin_noise([1.5**-i*0.5 for i in range(12)], 1, 1, False)\n",
        "        init2 = create_perlin_noise([1.5**-i*0.5 for i in range(8)], 4, 4, True)\n",
        "\n",
        "    init = TF.to_tensor(init).add(TF.to_tensor(init2)).div(2).to(device).unsqueeze(0).mul(2).sub(1)\n",
        "    del init2\n",
        "    return init\n",
        "\n",
        "##########################################################################################################\n",
        "\n",
        "#@markdown ##def do_run()\n",
        "def do_run():\n",
        "    global firstRun,_scale_multiplier\n",
        "    loss_values = []\n",
        "\n",
        "    if seed is not None:\n",
        "        np.random.seed(seed)\n",
        "        random.seed(seed)\n",
        "        torch.manual_seed(seed)\n",
        "        torch.cuda.manual_seed_all(seed)\n",
        "        torch.backends.cudnn.deterministic = True\n",
        "\n",
        "    make_cutouts = MakeCutouts(clip_size, cutn, skip_augs=skip_augs)\n",
        "    target_embeds, weights = [], []\n",
        "\n",
        "    for prompt in text_prompts:\n",
        "        txt, weight = parse_prompt(prompt)\n",
        "        txt = clip_model.encode_text(clip.tokenize(prompt).to(device)).float()\n",
        "        #----\n",
        "        print('This world could be signficantly less terrible\\n')\n",
        "        time.sleep(0.2)\n",
        "        print('if we cared for one another just a little bit more.\\n')\n",
        "        time.sleep(0.4)\n",
        "        print('Fork by sadnow aka @sadly_existent\\n')\n",
        "        time.sleep(0.2)\n",
        "        print('Huge thanks to everyone involved! Especially the core founders of these tools.\\n')\n",
        "        time.sleep(0.2)\n",
        "        print('Bugs can be reported through Github, Twitter, and VQLIPSE Discord...')\n",
        "        time.sleep(0.5)\n",
        "        if not _debug_mode:\n",
        "              display.clear_output(wait=True)\n",
        "        #----\n",
        "        if fuzzy_prompt:\n",
        "            for i in range(25):\n",
        "                # target_embeds.append((txt + torch.randn(txt.shape).cuda() * rand_mag).clamp(0,1))\n",
        "                target_embeds.append(txt + torch.randn(txt.shape).cuda() * rand_mag)\n",
        "                weights.append(weight)\n",
        "        else:\n",
        "            target_embeds.append(txt)\n",
        "            weights.append(weight)\n",
        "\n",
        "    for prompt in image_prompts:\n",
        "        path, weight = parse_prompt(prompt)\n",
        "        img = Image.open(fetch(path)).convert('RGB')\n",
        "        img = TF.resize(img, min(side_x, side_y, *img.size), T.InterpolationMode.LANCZOS)\n",
        "        batch = make_cutouts(TF.to_tensor(img).to(device).unsqueeze(0).mul(2).sub(1))\n",
        "        embed = clip_model.encode_image(normalize(batch)).float()\n",
        "        if fuzzy_prompt:\n",
        "            for i in range(25):\n",
        "                # target_embeds.append((embed + torch.randn(embed.shape).cuda() * rand_mag).clamp(0,1))\n",
        "                target_embeds.append(embed + torch.randn(embed.shape).cuda() * rand_mag)\n",
        "                weights.extend([weight / cutn] * cutn)\n",
        "        else:\n",
        "            target_embeds.append(embed)\n",
        "            weights.extend([weight / cutn] * cutn)\n",
        "\n",
        "    target_embeds = torch.cat(target_embeds)\n",
        "    weights = torch.tensor(weights, device=device)\n",
        "    if weights.sum().abs() < 1e-3:\n",
        "        raise RuntimeError('The weights must not sum to 0.')\n",
        "    weights /= weights.sum().abs()\n",
        "\n",
        "    init = None\n",
        "    if init_image is not None:\n",
        "        init = Image.open(fetch(init_image)).convert('RGB')\n",
        "        init = init.resize((side_x, side_y), Image.LANCZOS)\n",
        "        init = TF.to_tensor(init).to(device).unsqueeze(0).mul(2).sub(1)\n",
        "    \n",
        "    #DISABLED WITH PERLIN UPDATE\n",
        "    # if perlin_init:\n",
        "    #     if perlin_mode == 'color':\n",
        "    #         init = create_perlin_noise([1.5**-i*0.5 for i in range(12)], 1, 1, False)\n",
        "    #         init2 = create_perlin_noise([1.5**-i*0.5 for i in range(8)], 4, 4, False)\n",
        "    #     elif perlin_mode == 'gray':\n",
        "    #         init = create_perlin_noise([1.5**-i*0.5 for i in range(12)], 1, 1, True)\n",
        "    #         init2 = create_perlin_noise([1.5**-i*0.5 for i in range(8)], 4, 4, True)\n",
        "    #     else:\n",
        "    #         init = create_perlin_noise([1.5**-i*0.5 for i in range(12)], 1, 1, False)\n",
        "    #         init2 = create_perlin_noise([1.5**-i*0.5 for i in range(8)], 4, 4, True)\n",
        "        \n",
        "        # init = TF.to_tensor(init).add(TF.to_tensor(init2)).div(2).to(device)\n",
        "        # init = TF.to_tensor(init).add(TF.to_tensor(init2)).div(2).to(device).unsqueeze(0).mul(2).sub(1)\n",
        "        # del init2\n",
        "\n",
        "    cur_t = None\n",
        "\n",
        "    def cond_fn(x, t, y=None):\n",
        "        with torch.enable_grad():\n",
        "            x = x.detach().requires_grad_()\n",
        "            n = x.shape[0]\n",
        "            my_t = torch.ones([n], device=device, dtype=torch.long) * cur_t\n",
        "            out = diffusion.p_mean_variance(model, x, my_t, clip_denoised=False, model_kwargs={'y': y})\n",
        "            fac = diffusion.sqrt_one_minus_alphas_cumprod[cur_t]\n",
        "            x_in = out['pred_xstart'] * fac + x * (1 - fac)\n",
        "            x_in_grad = torch.zeros_like(x_in)\n",
        "            for i in range(cutn_batches):\n",
        "                clip_in = normalize(make_cutouts(x_in.add(1).div(2)))\n",
        "                image_embeds = clip_model.encode_image(clip_in).float()\n",
        "                dists = spherical_dist_loss(image_embeds.unsqueeze(1), target_embeds.unsqueeze(0))\n",
        "                dists = dists.view([cutn, n, -1])\n",
        "                losses = dists.mul(weights).sum(2).mean(0)\n",
        "                loss_values.append(losses.sum().item()) # log loss, probably shouldn't do per cutn_batch\n",
        "                x_in_grad += torch.autograd.grad(losses.sum() * clip_guidance_scale, x_in)[0] / cutn_batches\n",
        "            tv_losses = tv_loss(x_in)\n",
        "            range_losses = range_loss(out['pred_xstart'])\n",
        "            sat_losses = torch.abs(x_in - x_in.clamp(min=-1,max=1)).mean()\n",
        "            loss = tv_losses.sum() * tv_scale + range_losses.sum() * range_scale + sat_losses.sum() * sat_scale\n",
        "            if init is not None and init_scale:\n",
        "                init_losses = lpips_model(x_in, init)\n",
        "                loss = loss + init_losses.sum() * init_scale\n",
        "            x_in_grad += torch.autograd.grad(loss, x_in)[0]\n",
        "            grad = -torch.autograd.grad(x_in, x, x_in_grad)[0]\n",
        "        if clamp_grad:\n",
        "            adaptive_clip_grad([x]) #ADDED WITH PERLIN UPDATE\n",
        "            magnitude = grad.square().mean().sqrt()\n",
        "            return grad * magnitude.clamp(max=0.05) / magnitude\n",
        "        return grad\n",
        "\n",
        "    if model_config['timestep_respacing'].startswith('ddim'):\n",
        "        sample_fn = diffusion.ddim_sample_loop_progressive\n",
        "    else:\n",
        "        sample_fn = diffusion.p_sample_loop_progressive\n",
        "\n",
        "    for i in range(n_batches):\n",
        "        cur_t = diffusion.num_timesteps - skip_timesteps - 1\n",
        "        \n",
        "        if perlin_init: #ADDED WITH PERLIN UPDATE \n",
        "            init = regen_perlin()\n",
        "\n",
        "        if model_config['timestep_respacing'].startswith('ddim'):\n",
        "            samples = sample_fn(\n",
        "                model,\n",
        "                (batch_size, 3, side_y, side_x),\n",
        "                clip_denoised=clip_denoised,\n",
        "                model_kwargs={},\n",
        "                cond_fn=cond_fn,\n",
        "                progress=True,\n",
        "                skip_timesteps=skip_timesteps,\n",
        "                init_image=init,\n",
        "                randomize_class=randomize_class,\n",
        "                eta=eta,\n",
        "            )\n",
        "        else:\n",
        "            samples = sample_fn(\n",
        "                model,\n",
        "                (batch_size, 3, side_y, side_x),\n",
        "                clip_denoised=clip_denoised,\n",
        "                model_kwargs={},\n",
        "                cond_fn=cond_fn,\n",
        "                progress=True,\n",
        "                skip_timesteps=skip_timesteps,\n",
        "                init_image=init,\n",
        "                randomize_class=randomize_class,\n",
        "            )\n",
        "\n",
        "        for j, sample in enumerate(samples):\n",
        "            if not _debug_mode:\n",
        "              display.clear_output(wait=True)\n",
        "            cur_t -= 1\n",
        "            if j % display_rate == 0 or cur_t == -1:  #Only single iteration has finished\n",
        "                for k, image in enumerate(sample['pred_xstart']):\n",
        "                    #################################################################\n",
        "                    tqdm.write(f'Batch {i}, step {j}, output {k}:')\n",
        "                    if firstRun: max_mutation_amount = 0\n",
        "                    tqdm.write(f'Max mutation {max_mutation_amount}, Guidance scale {clip_guidance_scale}, TV scale {tv_scale}, Range scale {range_scale}')\n",
        "                    tqdm.write(msg_runtime)\n",
        "                    #################################################################\n",
        "                    current_time = datetime.now().strftime('%y%m%d-%H%M%S_%f')\n",
        "                    filename = f'batch{i:04}_iteration{j:04}_output{k:04}_{current_time}.png'  #reduced padding\n",
        "                    image = TF.to_pil_image(image.add(1).div(2).clamp(0, 1))\n",
        "                    image.save('/content/image_storage/' + _project_name + filename) \n",
        "                    display.display(display.Image('/content/image_storage/' + _project_name + filename))\n",
        "                    \n",
        "                    if _batch_genetics or _init_genetics:\n",
        "                      if firstRun:\n",
        "                        print(firstRun)  #debug\n",
        "                        max_mutation_amount = (_max_genetic_variance) / (n_batches)\n",
        "                        #print(\"Max mutation per output is \",max_mutation_amount)  #debug\n",
        "                        firstRun = False #we don't want this calculating again until next run\n",
        "                      #import random\n",
        "                      actual_mutation_amount = random.uniform(0,max_mutation_amount)\n",
        "                      _scale_multiplier = _scale_multiplier + actual_mutation_amount\n",
        "                      calculate_scale_multiplier()\n",
        "                      #cond_fn()\n",
        "                      cond_fn=cond_fn #unsure whether thse will work\n",
        "\n",
        "                    if google_drive and cur_t == -1:  #IMAGE HAS FINISHED & DRIVE IS ENABLED\n",
        "                      gc.collect()\n",
        "                      torch.cuda.empty_cache()\n",
        "                      #Image_upscale(filename)  #unsued resizer function\n",
        "                      #INSERT CODE ON THIS LINE\n",
        "                      if not _skip_upscaling: Image_upscale(filename,image)  #brought to you by world hunger  \n",
        "                      else:\n",
        "                        print(\"\\n Saving to \",output_folder_images,filename)\n",
        "                        os.path.join(_drive_location,_project_name)\n",
        "                        image.save(os.path.join(output_folder_images,filename))\n",
        "                        #image.save(output_folder_images + filename)\n",
        "                      #else:                   \n",
        "                      #anything placed here will run before the upscale has completed\n",
        "                      #will be ran every single step\n",
        "                      #image.save(output_folder_images + filename)\n",
        "\n",
        "                        \n",
        "\n",
        "        plt.plot(np.array(loss_values), 'r')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9zY-8I90LkC6"
      },
      "source": [
        "# Settings & Generation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Fpbody2NCR7w"
      },
      "outputs": [],
      "source": [
        "timestep_respacing = 'ddim25' #@param ['25','50','100','150','250','500','1000','ddim25','ddim50','ddim100','ddim150','ddim250','ddim500','ddim1000']  \n",
        "#@markdown *Modify this value to decrease the number of iterations/prompt.\n",
        "# timestep_respacing = '25'\n",
        "diffusion_steps = 1000\n",
        "\n",
        "model_config = model_and_diffusion_defaults()\n",
        "if diffusion_model == '512x512_diffusion_uncond_finetune_008100':\n",
        "    model_config.update({\n",
        "        'attention_resolutions': '32, 16, 8',\n",
        "        'class_cond': False,\n",
        "        'diffusion_steps': diffusion_steps,\n",
        "        'rescale_timesteps': True,\n",
        "        'timestep_respacing': timestep_respacing,\n",
        "        'image_size': 512,\n",
        "        'learn_sigma': True,\n",
        "        'noise_schedule': 'linear',\n",
        "        'num_channels': 256,\n",
        "        'num_head_channels': 64,\n",
        "        'num_res_blocks': 2,\n",
        "        'resblock_updown': True,\n",
        "        'use_fp16': True,\n",
        "        'use_scale_shift_norm': True,\n",
        "    })\n",
        "elif diffusion_model == '256x256_diffusion_uncond':\n",
        "    model_config.update({\n",
        "        'attention_resolutions': '32, 16, 8',\n",
        "        'class_cond': False,\n",
        "        'diffusion_steps': diffusion_steps,\n",
        "        'rescale_timesteps': True,\n",
        "        'timestep_respacing': timestep_respacing,\n",
        "        'image_size': 256,\n",
        "        'learn_sigma': True,\n",
        "        'noise_schedule': 'linear',\n",
        "        'num_channels': 256,\n",
        "        'num_head_channels': 64,\n",
        "        'num_res_blocks': 2,\n",
        "        'resblock_updown': True,\n",
        "        'use_fp16': True,\n",
        "        'use_scale_shift_norm': True,\n",
        "    })\n",
        "side_x = side_y = model_config['image_size']\n",
        "\n",
        "model, diffusion = create_model_and_diffusion(**model_config)\n",
        "model.load_state_dict(torch.load(f'{model_path}{diffusion_model}.pt', map_location='cpu'))\n",
        "model.requires_grad_(False).eval().to(device)\n",
        "for name, param in model.named_parameters():\n",
        "    if 'qkv' in name or 'norm' in name or 'proj' in name:\n",
        "        param.requires_grad_()\n",
        "if model_config['use_fp16']:\n",
        "    model.convert_to_fp16()\n",
        "\n",
        "################################################\n",
        "\n",
        "clip_model = clip.load('ViT-B/16', jit=False)[0].eval().requires_grad_(False).to(device)\n",
        "clip_size = clip_model.visual.input_resolution\n",
        "normalize = T.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711])\n",
        "lpips_model = lpips.LPIPS(net='vgg').to(device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U0PwzFZbLfcy",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#INITIAL VALUES\n",
        "#@title  { form-width: \"300px\" }\n",
        "firstRun = True\n",
        "msg_runtime = ''\n",
        "_keep_first_upscale = False\n",
        "_run_upscaler = True\n",
        "batch_size =  1\n",
        "clamp_grad = True # True - Experimental: Using adaptive clip grad in the cond_fn\n",
        "skip_augs = False # False - Controls whether to skip torchvision augmentations\n",
        "randomize_class = True # True - Controls whether the imagenet class is randomly changed each iteration\n",
        "#############################################################################################\n",
        "#@markdown ## Basic Settings\n",
        "_text_prompt =  \"Gloomy Alaskan Cabin,' surrealist painting inspired by Salvador Dali\" #@param {type:\"string\"} \n",
        "n_batches =  200#@param{type:\"raw\"}\n",
        "  #Controls the starting point along the diffusion timesteps\n",
        "_init_image = '' #@param{type:\"string\"}\n",
        "#@markdown >`init_image` is optional. See scaling settings in docs- \n",
        "_seed = \"None\" #@param{type:\"string\"}\n",
        "#@markdown >Set `seed` to *None* to use a random seed.\n",
        "_batch_genetics = False #future implementation\n",
        "_init_genetics = False  #not currently implemented\n",
        "_max_genetic_variance =  0.1\n",
        "_saturation_scale =  0  #mysterious\n",
        "skip_timesteps =  5#@param{type:\"raw\"}\n",
        "#@markdown ---\n",
        "#@markdown ##Advanced Settings\n",
        "#@markdown `Performance Settings` (Optimized for P100; increase `cutn_batches` if you get memory errors.)\n",
        "cutn =  16#@param{type:\"raw\"}\n",
        "  #Controls how many crops to take from the image. Increase for higher quality.\n",
        "cutn_batches = 8 #@param [1,2,4,8,16] {type:\"raw\"} \n",
        "  #Accumulate CLIP gradient from multiple batches of cuts [Can help with OOM errors / Low VRAM]\n",
        "_esrgan_tilesize = \"256\" #@param[16,32,64,128,256,512,1024]\n",
        "#_upscale_performance_mode = False\n",
        "\n",
        "#@markdown `Visual Settings`\n",
        "_scale_multiplier = 1 #@param {type:\"slider\", min:0.1, max:5, step:0.1}\n",
        "clip_denoised = False #@param{type:\"boolean\"}\n",
        "fuzzy_prompt = False #@param{type:\"boolean\"}\n",
        "  # False - Controls whether to add multiple noisy prompts to the prompt losses\n",
        "_noise_mode = 'gray' #@param ['mixed','gray','color']\n",
        "_noise_amount = 0.07 #@param {type:\"slider\", min:0.00, max:1, step:0.01}\n",
        "_image_prompts = \"\" #@param{type:\"string\"}\n",
        "eta =  0.5\n",
        "#@markdown `Scaling`\n",
        "_clip_guidance_scale =  5000#@param {type:\"raw\"}\n",
        "_tv_scale =  100#@param {type:\"raw\"}\n",
        "_range_scale =  150#@param {type:\"raw\"}\n",
        "  #This enhances the effect of the init image, a good value is 1000\n",
        "_init_scale =  1000#@param{type:\"raw\"}\n",
        "\n",
        "# seed = random.randint(0, 2**32) # Choose a random seed and print it at end of run for reproduction\n",
        "\n",
        "_skip_upscaling = False\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown ##Session Settings\n",
        "display_rate =  1#@param{type:\"raw\"}\n",
        "#@markdown >Raise `display_rate` if your image previews aren't rendering fast enough.\n",
        "_upscale_model='RealESRGAN_x4plus' #@param ['RealESRGAN_x4plus_anime_6B','RealESRGAN_x4plus','RealESRGAN_x2plus'] {type:\"string\"}\n",
        "_target_resolution = \"2048\" #@param[256,512,1024,2048,4096]\n",
        "_project_name = 'gloomy-snowdays' #@param{type:\"string\"}\n",
        "_debug_mode = False #@param{type:\"boolean\"}\n",
        "global _debug_mode\n",
        "##@markdown Original Defaults: `clip_guidance_scale 5000`,`tv_scale 150`,`range scale 150`\n",
        "##@markdown Recommended defaults for init_images (ddim50): `clip_guidance_scale 2000`,`tv_scale 150`,`range scale 50`, `init_scale 1000`, `skip_timesteps 16 (7-9 for ddim25)`\n",
        "##@markdown There is a possibility that `tv_scale` can be set between `0` to `10000`\n",
        "##@markdown `skip_timesteps` does a lot for the similarity in `init_settings`\n",
        "##@markdown Special thanks to many people on the VQLIPSE Discord\n",
        "##---\n",
        "\n",
        "#--------------------------------------------------------------------------------------------------------\n",
        "text_prompts = [\n",
        "    # \"an abstract painting of 'ravioli on a plate'\",\n",
        "    # 'cyberpunk wizard on top of a skyscraper, trending on artstation, photorealistic depiction of a cyberpunk wizard',\n",
        "    _text_prompt]\n",
        "    # 'cyberpunk wizard',\n",
        "if _debug_mode: dprint('Note: _debug_mode is on.')\n",
        "if diffusion_model == \"512x512_diffusion_uncond_finetune_008100\": model_size = 512\n",
        "else: model_size = 256\n",
        "if int(_target_resolution) == model_size:\n",
        "  print(\"\\n Due to your _target_resolution, _skip_upsscaling will be enabled. \\n\")\n",
        "  _skip_upscaling = True\n",
        "if int(model_size) > int(_target_resolution):\n",
        "  print(\"\\n NOTICE: Your _upscale_reoslution is higher than your model size! Setting to match and disabling upscaling... \\n\")\n",
        "  _target_resolution = int(model_size)\n",
        "  _skip_upscaling = True\n",
        "#---------------------------------------------------------------------------------------\n",
        "#-init & perlin handling\n",
        "if _init_image == '':\n",
        "  init_image = None\n",
        "else:\n",
        "  init_image = _init_image  # None - URL or local path\n",
        "  dprint('Init image: ' + init_image)\n",
        "if _noise_amount > 0: add_random_noise = True\n",
        "else: add_random_noise = False\n",
        "perlin_init = add_random_noise\n",
        "if init_image is not None: # Can't combine init_image and perlin options\n",
        "  perlin_init = False\n",
        "  dprint('NOTICE: You may want to disable _noise_amount when using _init_images')\n",
        "rand_mag = _noise_amount # 0.1 - Controls the magnitude of the random noise\n",
        "if _init_image == '':\n",
        "  print(\"No init image detected. Setting init_scale to 0...\")\n",
        "  init_scale = 0\n",
        "else:\n",
        "  init_scale = _init_scale\n",
        "  if init_scale == 0:\n",
        "    dprint('Notice: init_image is set but there is no init_scale! Try 1000 as a default.')\n",
        "if not _image_prompts == \"\":\n",
        "  dprint('Image prompt: ' + _image_prompts)\n",
        "  image_prompts = [\n",
        "      _image_prompts,\n",
        "  ]\n",
        "else: image_prompts = []\n",
        "perlin_mode = _noise_mode # 'mixed' ('gray', 'color')\n",
        "sat_scale = _saturation_scale #mysterious # 0 - Controls how much saturation is allowed. From nshepperd's JAX notebook.\n",
        "#---------------------------------------------------------------------------------------\n",
        "#-seed handling\n",
        "if _seed == \"None\" or _seed == \"None: \":\n",
        "  dprint('Seed: Undefined (Random)')\n",
        "  seed = None\n",
        "else:\n",
        "  dprint('Seed: ' + _seed)\n",
        "  seed = int(_seed)\n",
        "#---------------------------------------------------------------------------------------\n",
        "#Output image directory handling\n",
        "import os\n",
        "output_folder_images = os.path.join(_drive_location,_project_name)\n",
        "print (\"\\n output_folder_images is \",output_folder_images)\n",
        "dir_make(output_folder_images)\n",
        "temp_image_storage = '/content/image_storage/'  #for non-upscaled images\n",
        "dir_make(temp_image_storage)\n",
        "#---------------------------------------------------------------------------------------\n",
        "def calculate_scale_multiplier():\n",
        "  global clip_guidance_scale,tv_scale,range_scale\n",
        "  clip_guidance_scale = _scale_multiplier * _clip_guidance_scale\n",
        "  tv_scale = _scale_multiplier * _tv_scale\n",
        "  range_scale = _scale_multiplier * _range_scale\n",
        "  return clip_guidance_scale,tv_scale,range_scale\n",
        "calculate_scale_multiplier()\n",
        "#---------------------------------\n",
        "# UPSCALING\n",
        "dprint('Upscaler model: ' + _upscale_model + '; Target resolution: ' + str(_target_resolution))\n",
        "\n",
        "def Image_upscale(filename,image):   \n",
        "  \n",
        "  def round_up_to_even(f):\n",
        "    return math.ceil(f / 2.) * 2\n",
        "\n",
        "  def File_addSuffix(filename,needed_passes,completed_passes):  #returns the filename of the previously upscaled file\n",
        "    stripped_filename = os.path.splitext(filename)[0] #get non-upscaled image filename\n",
        "    stripped_filename = _project_name + stripped_filename\n",
        "    z = 1\n",
        "    while z < completed_passes + 1:\n",
        "      stripped_filename = stripped_filename + '_out'\n",
        "      z = z + 1\n",
        "    if completed_passes > 1 or needed_passes == 1:\n",
        "      target_extension = 'jpg'  #dont forget the dot\n",
        "    else: target_extension = 'png'\n",
        "    stripped_filename = stripped_filename + '.' + target_extension  #this identifies the fullpath output from the 1st upscale\n",
        "    current_upscale_target = os.path.join(output_folder_images,stripped_filename)\n",
        "    print('Current upscale target is ',current_upscale_target)\n",
        "    return current_upscale_target\n",
        "  \n",
        "  %cd /content/Real-ESRGAN/\n",
        "  temp_image_filepath = temp_image_storage + _project_name + filename\n",
        "  #image.save(temp_image_filepath) #here we save the un-upscaled so we can load it IN the upscaler\n",
        "  \n",
        "  retry_attempt = 0\n",
        "\n",
        "  \n",
        "  def Load_config(retry_attempt,_upscale_performance_mode):\n",
        "    needed_passes = 1\n",
        "    global _skip_upscaling\n",
        "    if retry_attempt > 1:\n",
        "      print(\"\\n UPSCALING FAILED! Try raising your cutn_batch settings.\")\n",
        "      print(\"Attempting to save un-upscaled file... \\n\")\n",
        "      image.save((os.path.join(output_folder_images,filename)))\n",
        "      _skip_upscaling = True\n",
        "      import sys\n",
        "      sys.exit()\n",
        "\n",
        "    global model_size\n",
        "    \n",
        "    global _target_resolution\n",
        "    outscale_resolution_int = int(_target_resolution)\n",
        "    if int(model_size) == 256 and int(_target_resolution) > 2048:\n",
        "      print(\"\\n NOTE: 256 model not currently compatible with 4096 upscaling. \\n You can use the 512 model by picking it at the top of the Notebook. \\n\")\n",
        "      outscale_resolution_int = 2048\n",
        "    _outscale = outscale_resolution_int / model_size\n",
        "      #ex: 2048 / 512 = _outscale of 4\n",
        "    if not (_outscale % 2) == 0: #number is not even\n",
        "      _outscale = round_up_to_even(_outscale)\n",
        "    if _outscale > 4 :\n",
        "      needed_passes = _outscale / 4\n",
        "      _outscale = _outscale / 2    \n",
        "    if _upscale_performance_mode:\n",
        "      _outscale = _outscale / 2\n",
        "      needed_passes = needed_passes * 2\n",
        "      half_precision = True\n",
        "    else:\n",
        "      _skip_upscaling = False\n",
        "      half_precision = False\n",
        "    needed_passes = round(needed_passes)\n",
        "    print(\"\\n NOTICE: Performing \",needed_passes,\"needed_passes...\")\n",
        "    print(\"Outscale is\", _outscale)\n",
        "    print(\"Needed passes is\",needed_passes,\"\\n\")\n",
        "    return _outscale,needed_passes,_skip_upscaling,half_precision,model_size\n",
        "\n",
        "  run_config = True\n",
        "  _upscale_performance_mode = False\n",
        "\n",
        "  completed_passes = 0\n",
        "\n",
        "  retry_attempt = 0\n",
        "  upscale_complete = False\n",
        "  next_outscale = 0\n",
        "  \n",
        "  while not upscale_complete:\n",
        "    if run_config:\n",
        "      _outscale,needed_passes,_skip_upscaling,half_precision,model_size = Load_config(retry_attempt,_upscale_performance_mode)\n",
        "      completed_passes = 0\n",
        "    end_flag = '--face_enhance '\n",
        "    _outscale = str(_outscale)\n",
        "    if half_precision: end_flag = end_flag + '--half '\n",
        "\n",
        "    if completed_passes < 1:\n",
        "      end_flag = end_flag + '--input ' + temp_image_filepath + ' '\n",
        "    else:\n",
        "      ##################################################################################\n",
        "      current_upscale_target = File_addSuffix(filename,needed_passes,completed_passes)    \n",
        "      #################################################################################\n",
        "      end_flag = end_flag + '--input ' + current_upscale_target + ' '\n",
        "\n",
        "    if completed_passes > 0 or needed_passes == 1:\n",
        "      output_extension = 'jpg'\n",
        "    else: output_extension = 'png'\n",
        "    \n",
        "    if (int(model_size) * int(float(_outscale)) * (int(needed_passes) ** 2)) > int(_target_resolution):\n",
        "      next_outscale = 2\n",
        "    if not _upscale_model == 'RealESRGAN_x2plus' and next_outscale == 0:\n",
        "      end_flag = end_flag + '--outscale ' + _outscale + ' '\n",
        "    else:\n",
        "      if completed_passes + 1 == needed_passes:\n",
        "        end_flag = end_flag + '--outscale ' + str(next_outscale) + ' '\n",
        "        print(\"next_outscale is\",next_outscale)\n",
        "    if end_flag.endswith(' '): end_flag = end_flag[:-1] #removes extra space\n",
        "\n",
        "    if not upscale_complete:\n",
        "      !python inference_realesrgan.py --model_name $_upscale_model --output $output_folder_images --tile $_esrgan_tilesize --ext $output_extension $end_flag\n",
        "    \n",
        "\n",
        "\n",
        "    if completed_passes > 0: _upscaled_path = file_check(current_upscale_target)\n",
        "    else: _upscaled_path = file_check(temp_image_filepath)  #attempting to load in performance_mode\n",
        "    \n",
        "    if _upscaled_path:\n",
        "      print(\"\\n NOTICE: Upscale pass \",completed_passes,\" COMPLETE! \\n\")\n",
        "      run_config = False  #because we are done\n",
        "      completed_passes = completed_passes + 1\n",
        "      if completed_passes >= needed_passes and needed_passes > 1:\n",
        "        print(\"\\n Removing previous file: \",current_upscale_target,\"\\n\")\n",
        "        !rm $current_upscale_target\n",
        "    else: \n",
        "      if completed_passes < 1: print(\"\\n ERROR: CAN'T FIND UPSCALED FILE \", temp_image_filepath)\n",
        "      else: print(\"\\n ERROR: CAN'T FIND UPSCALED FILE \", current_upscale_target)\n",
        "      if not _upscale_performance_mode:\n",
        "        print(\"\\n RECOVERY STATUS: ATTEMPTING TO RUN IN upscale_performance mode... \\n\")\n",
        "        _upscale_performance_mode = True\n",
        "      completed_passes = 0\n",
        "      retry_attempt = retry_attempt + 1\n",
        "      run_config = True\n",
        "    if completed_passes >= needed_passes:\n",
        "      upscale_complete = True\n",
        "    if upscale_complete:\n",
        "      return\n",
        "\n",
        "if _debug_mode:\n",
        "  print(\"\\n Setting skip to large number... \\n\")\n",
        "  skip_timesteps = 20\n",
        "\n",
        "\n",
        "#-------------------------------------------------------\n",
        "\n",
        "# 1 - Controls how many consecutive batches of images are generated\n",
        "\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()\n",
        "try:\n",
        "    do_run()\n",
        "except KeyboardInterrupt:\n",
        "    gc.collect()\n",
        "    torch.cuda.empty_cache()\n",
        "    pass\n",
        "finally:\n",
        "    print('\\n seed', seed)\n",
        "    print('\\n Output(s) saved to ',output_folder_images)\n",
        "    gc.collect()\n",
        "    torch.cuda.empty_cache()\n",
        "\n",
        "###############################################"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xc1W6_SAAvJU"
      },
      "source": [
        "# Prompt Engineering (Documentation)\n",
        "\n",
        "This part may be outdated outdated; will update eventually. Includes setting documentation-"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u4Ut78Lol7EN"
      },
      "source": [
        "**This documentation may be out of date. See full documentation and settings research at https://thisnameismine.com\n",
        "\n",
        "##Performance Settings (VRAM Troubleshooting)\n",
        "\n",
        "- `cutn`: Controls image quality\n",
        "- `cutn_batches` : Lessens the amount of cutn required by doing bursts of cutn. If you want to use a cutn of 60, but don’t have enough VRAM, set cutn to 30 and cutn_n to 2.\n",
        "- `esrgan_tilesize`: If you need to squeeze a little more VRAM to complete the upscale, this is the setting.\n",
        "---\n",
        "- **Defaults (P100)**: `cutn=30`, `cutn_batches = 2 or 4`, `esrgan_tilesize = 512 or 1024`\n",
        "\n",
        "---\n",
        "\n",
        "##Image Quality Settings\n",
        "- `clip_guidance_scale` controls how much the image should look like the prompt\n",
        "- `tv_scale` controls smoothing\n",
        " - set to 0 for crispy, otherwise 50-150 works good. Can go up to 10,000.\n",
        "- `range_scale` Controls how far out of range RGB values are allowed to be.\n",
        "- `scale_multiplier` multiplies clip_guidance_scale,tv_scale,range_scale\n",
        " - On init images, 50 might be a good value; otherwise 150.\n",
        "- `skip_timesteps` Controls the starting point along the diffusion timesteps.\n",
        " - If you set `timestep_respacing` to ddim50 and `skip_timesteps` to `10`, your image will do 40 iterations.\n",
        "- `clip_denoised` tries to filter out noise from generation.\n",
        "- `fuzzy_prompt` attempts to include noise in generation\n",
        "---\n",
        "- **Defaults for non-init images**: `clip_guidance_scale=5000`,`tv_scale=100`,`range_scale=150`\n",
        "- **Defaults for init images**: `clip_guidance_scale=2000`,`tv_scale=100`,`range_scale=50`\n",
        "---\n",
        "##Init Settings\n",
        "- Leave `_init_image` blank if you want to generate from scratch.\n",
        "- Set `_noise_amount` to 0 (Perlin noise not currently compatible)\n",
        "- Set `skip_timesteps` to `30-40%` of your `timestep_respacing` value when using an init image. This setting also controls blur to a certain extent. \n",
        "A possible good value for `_init_scale` is 1000. If set to `0` it will default back to 1000 when an init image is present.\n",
        "- `_noise_amount` settings should be deactivated if an `init_image` is detected.\n",
        "- `init_scale` controls how strictly the image adheres to the init image. `1000` is a good value, but for lower `skip_timesteps` you can set `init scale` to `5000`. \n",
        "\n",
        "---\n",
        "##A100 settings\n",
        "The a100 does much better with `cutn_batch` set to 1. You can set `cutn` up to `128` replicate results.\n",
        "\n",
        "---\n",
        "**Examples of diffusion prompts:**\n",
        "\n",
        "- `Representation of chronic anxiety:2, illustrated by Luigi Serafini:1, inspired by the Codex Seraphinianus:1`\n",
        "\n",
        "- `Sad panda vector art made in Blender 3d:4, nature background:2, 4k parallax vibrant colorful panda bear trending on artstation:1` with init scale of 5, init image of https://www.pngitem.com/pimgs/m/16-161339_giant-panda-bear-silhouette-drawing-clip-art-silhouette.png\n",
        "\n",
        "- `Representation of depression | isometric vector art titled 'Depression'`\n",
        "\n",
        "- `Liminal space, liminal hotel hallway rendered in unreal engine, top post on r/liminalspaces`\n",
        "\n",
        "- `art piece titled 'The meaning of life is to find meaning IN life', existential clipart featuring your friend Dave, vector clipart anime faces`\n",
        "\n",
        "- `Vector art named 'Chronic procrastination equals suicidal ideation', vector HD clipart featuring suicide`\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "*If you use `:` (weights), you must weigh out everything or you'll receive an error. For example, if you have three things and one has a weight of `:2`, the other two things need to have a weight of `:1`.\n",
        "\n",
        "*I am currently unsure whether `|` (pipes) do the same thing as commas. Also unsure whether having spaces between commas makes a different outcome. Weights work for both.*\n",
        "\n",
        "---\n",
        "\n",
        "Diffusion prompting might seem trickier to master than VQGAN models, but it still allows for some level of control. Sometimes less it better.\n",
        "\n",
        "It doesn't always need the `trending on artstation` type of lingo; sometimes it will benefit more from something like `photorealistic 4k nature replication`.\n",
        "\n",
        "Diffusion is kind of unique in how it often doesn't duplicate the desired subject. For example, `cat photo` will usually only give you 1 cat.\n",
        "\n",
        "------------------------------------------------\n",
        "\n",
        "### A Giant List of Terms to Try\n",
        "\n",
        "Credit for this list goes to @Atman on Discord and/or whoever else contributed to the Pastebin.\n",
        "\n",
        "*'8k resolution'\n",
        ",'pencil sketch'\n",
        ",'8K 3D'\n",
        ",'creative commons attribution'\n",
        ",'deviantart'\n",
        ",'CryEngine'\n",
        ",'Unreal Engine'\n",
        ",'concept art'\n",
        ",'photoillustration'\n",
        ",'pixiv'\n",
        ",'Flickr'\n",
        ",'ArtStation HD'\n",
        ",'Behance HD'\n",
        ",'HDR'\n",
        ",'anime'\n",
        ",'filmic'\n",
        ",'Stock photo'\n",
        ",'Ambient occlusion'\n",
        ",'Global illumination'\n",
        ",'Chalk art'\n",
        ",'Low poly'\n",
        ",'Booru'\n",
        ",'Polycount'\n",
        ",'Acrylic art'\n",
        ",'Hyperrealism'\n",
        ",'Zbrush Central'\n",
        ",'Rendered in Cinema4D'\n",
        ",'Rendered in Maya'\n",
        ",'Photo taken with Nikon D750'\n",
        ",'Tilt shift'\n",
        ",'Mixed media'\n",
        ",'Depth of field'\n",
        ",'DSLR'\n",
        ",'Detailed painting'\n",
        ",'Volumetric lighting'\n",
        ",'Storybook illustration'\n",
        ",'Unsplash contest winner'\n",
        ",'#vfxfriday'\n",
        ",'Ultrafine detail'\n",
        ",'20 megapixels'\n",
        ",'Photo taken with Fujifilm Superia'\n",
        ",'Photo taken with Ektachrome'\n",
        ",'matte painting'\n",
        ",'reimagined by industrial light and magic'\n",
        ",'Watercolor'\n",
        ",'CGSociety'\n",
        ",'childs drawing'\n",
        ",'marble sculpture'\n",
        ",'airbrush art'\n",
        ",'renaissance painting'\n",
        ",'Velvia'\n",
        ",'Provia'\n",
        ",'photo taken with Provia'\n",
        ",'prerendered graphics'\n",
        ",'criterion collection'\n",
        ",'dye-transfer'\n",
        ",'stipple'\n",
        ",'Parallax'\n",
        ",'Bryce 3D'\n",
        ",'Terragen'\n",
        ",'(2013) directed by cinematography by'\n",
        ",'Bokeh'\n",
        ",'1990s 1995'\n",
        ",'1970s 1975'\n",
        ",'1920s 1925'\n",
        ",'charcoal drawing'\n",
        ",'commission for'\n",
        ",'furaffinity'\n",
        ",'flat shading'\n",
        ",'ink drawing'\n",
        ",'artwork'\n",
        ",'oil on canvas'\n",
        ",'macro photography'\n",
        ",'hall of mirrors'\n",
        ",'polished'\n",
        ",'sunrays shine upon it'\n",
        ",'aftereffects'\n",
        ",'iridescent'\n",
        ",'#film'\n",
        ",'datamosh'\n",
        ",'(1962) directed by cinematography'\n",
        ",'holographic'\n",
        ",'dutch golden age'\n",
        ",'digitally enhanced'\n",
        ",'National Geographic photo'\n",
        ",'Associated Press photo'\n",
        ",'matte background'\n",
        ",'Art on Instagram'\n",
        ",'#myportfolio'\n",
        ",'digital illustration'\n",
        ",'stock photo'\n",
        ",'aftereffects'\n",
        ",'speedpainting'\n",
        ",'colorized'\n",
        ",'detailed'\n",
        ",'psychedelic'\n",
        ",'wavy'\n",
        ",'groovy'\n",
        ",'movie poster'\n",
        ",'pop art'\n",
        ",'made of beads and yarn'\n",
        ",'made of feathers'\n",
        ",'made of crystals'\n",
        ",'made of liquid metal'\n",
        ",'made of glass'\n",
        ",'made of cardboard'\n",
        ",'made of vines'\n",
        ",'made of cheese'\n",
        ",'made of flowers'\n",
        ",'made of insects'\n",
        ",'made of mist'\n",
        ",'made of paperclips'\n",
        ",'made of rubber'\n",
        ",'made of plastic'\n",
        ",'made of wire'\n",
        ",'made of trash'\n",
        ",'made of wrought iron'\n",
        ",'made of all of the above'\n",
        ",'tattoo'\n",
        ",'woodcut'\n",
        ",'American propaganda'\n",
        ",'Soviet propaganda'\n",
        ",'PS1 graphics'\n",
        ",'Fine art'\n",
        ",'HD mod'\n",
        ",'Photorealistic'\n",
        ",'Poster art'\n",
        ",'Constructivism'\n",
        ",'pre-Raphaelite'\n",
        ",'Impressionism'\n",
        ",'Lowbrow'\n",
        ",'RTX on'\n",
        ",'chiaroscuro'\n",
        ",'Egyptian art'\n",
        ",'Fauvism'\n",
        ",'shot on 70mm'\n",
        ",'Art Deco'\n",
        ",'Picasso'\n",
        ",'Da Vinci'\n",
        ",'Academic art'\n",
        ",'3840x2160'\n",
        ",'Photocollage'\n",
        ",'Cubism'\n",
        ",'Surrealist'\n",
        ",'THX Sound'\n",
        ",'ZBrush'\n",
        ",'Panorama'\n",
        ",'smooth'\n",
        ",'DC Comics'\n",
        ",'Marvel Comics'\n",
        ",'Ukiyo-e'\n",
        ",'Flemish Baroque'\n",
        ",'vray tracing'\n",
        ",'pixel perfect'\n",
        ",'quantum wavetracing'\n",
        ",'Zbrush central contest winner'\n",
        ",'ISO 200'\n",
        ",'Bob Ross'\n",
        ",'32k HUHD'\n",
        ",'Photocopy'\n",
        ",'DeviantArt HD'\n",
        ",'infrared'\n",
        ",'Angelic photograph'\n",
        ",'Demonic photograph'\n",
        ",'Biomorphic'\n",
        ",'Windows Vista'\n",
        ",'Skeuomorphic'\n",
        ",'Physically based rendering'\n",
        ",'Trance compilation CD'\n",
        ",'Concert poster'\n",
        ",'Steampunk'\n",
        ",'Sketchfab'\n",
        ",'Goth'\n",
        ",'Wiccan'\n",
        ",'trending on artstation'\n",
        ",'featured on artstation'\n",
        ",'artstation HQ'\n",
        ",'artstation contest winner'\n",
        ",'ultra HD'\n",
        ",'high quality photo'\n",
        ",'instax'\n",
        ",'ilford HP5'\n",
        ",'infrared'\n",
        ",'Lomo'\n",
        ",'Matte drawing'\n",
        ",'matte photo'\n",
        ",'glowing neon'\n",
        ",'Xbox 360 graphics'\n",
        ",'flickering light'\n",
        ",'Playstation 5 screenshot'\n",
        ",'Kodak Gold 200'\n",
        ",'by Edward Hopper'\n",
        ",'rough'\n",
        ",'maximalist'\n",
        ",'minimalist'\n",
        ",'Kodak Ektar'\n",
        ",'Kodak Portra'\n",
        ",'geometric'\n",
        ",'cluttered'\n",
        ",'Rococo'\n",
        ",'destructive'\n",
        ",'by James Gurney'\n",
        ",'by Thomas Kinkade'\n",
        ",'by Vincent Di Fate'\n",
        ",'by Jim Burns'\n",
        ",'androgynous'\n",
        ",'masculine'\n",
        ",'genderless'\n",
        ",'feminine'\n",
        ",'extremely gendered, masculine and feminine'\n",
        ",'4k result'\n",
        ",'#pixelart'\n",
        ",'voxel art'\n",
        ",'wimmelbilder'\n",
        ",'dystopian art'\n",
        ",'apocalypse art'\n",
        ",'apocalypse landscape'\n",
        ",'2D game art'\n",
        ",'Windows XP'\n",
        ",'y2k aesthetic'\n",
        ",'#screenshotsaturday'\n",
        ",'seapunk'\n",
        ",'vaporwave'\n",
        ",'Ilya Kuvshinov'\n",
        ",'Paul Cezanne'\n",
        ",'Henry Moore'\n",
        ",'phallic'\n",
        ",'creepypasta'\n",
        ",'retrowave'\n",
        ",'synthwave'\n",
        ",'outrun'*"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Other Projects by Sadnow\n",
        "\n",
        "---\n",
        "\n",
        "###[AnimationKit](https://github.com/sadnow/AnimationKit-AI) (Github)\n",
        "\n",
        "AnKit is a Colab-based video enhancement tool that uses deep learning models (currently Real-ESRGAN and Practical-RIFE) to synthetically smooth and upscale videos.\n",
        "\n",
        "Pick your mp4 file or frame directory, set your desired length and framerate, and watch the magic happen!\n",
        "\n"
      ],
      "metadata": {
        "id": "Hza14whDV4G3"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "1YwMUyt9LHG1",
        "I6Vt-98xGZVu",
        "WnN6j5WzvFHg",
        "xc1W6_SAAvJU",
        "u4Ut78Lol7EN"
      ],
      "machine_shape": "hm",
      "name": "360Diffusion_Public.ipynb",
      "private_outputs": true,
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}