360Diffusion_Simplified-Demo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "360Diffusion_Simplified-Demo.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [
        "1YwMUyt9LHG1",
        "I6Vt-98xGZVu",
        "WnN6j5WzvFHg",
        "xc1W6_SAAvJU",
        "u4Ut78Lol7EN"
      ],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "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_Simplified-Demo.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 [Demo] - Text-to-image Generation for Beginners\n",
        "\n",
        "##Are you new to AI image generation? \n",
        "\n",
        "This demo has been developed with beginners in mind. Many of the advanced settings have been hidden, so as not to overwhelm those unfamiliar with Google Colab. Depending on the GPU Google assigns you, you can generate 1 image every 20-60 seconds. This demo uses the 256 model, as that one requires less resources and is therefore less prone to errors. Images are upscaled to 1024x1024px after each generation.\n",
        "\n",
        "---\n",
        "\n",
        "###Major components in this notebook were originally developed and/or revised by:\n",
        "\n",
        "- `Katherine Crowson` (https://github.com/crowsonkb, https://twitter.com/RiversHaveWings); Original CLIP Guided Diffusion Notebook Founder; Plays a big role in the AI gen scene\n",
        "- `Daniel Russell` (https://github.com/russelldc, https://twitter.com/danielrussruss); Forked Katherin Crowson's Diffusion notebook into Quick / Fast Diffusion\n",
        "- `sadnow` (https://github.com/sadnow); Forked Daniel's notebook into 360Diffusion with integration of Real-ESRGAN upscaling\n",
        "- Parameter research assisted in-part by the following community (https://www.patreon.com/sportsracer48); pyTTI / VQLIPSE animation\n",
        "---\n",
        "To begin, click the play buttons on the left hand sides. There are 2 steps you will want to run. After you have ran the 1st step, you do not need to run it again."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FpZczxnOnPIU",
        "cellView": "form"
      },
      "source": [
        "#@title  { form-width: \"100px\" }\n",
        "\n",
        "#@markdown #Step 1 (Only needs ran once per session)\n",
        "\n",
        "\n",
        "##@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",
        "enable_error_checking = True #saves ram\n",
        "if enable_error_checking:\n",
        "  !nvidia-smi\n",
        "else:\n",
        "  !nvidia-smi\n",
        "  !nvidia-smi -i 0 -e 0\n",
        "\n",
        "##########\n",
        "\n",
        "#universal functions\n",
        "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_path $model_path --netscale $scale --face_enhance --input $input --output $output --ext jpg\n",
        "    %cd /content/\n",
        "\n",
        "##@title Choose model here:\n",
        "# diffusion_model = \"256x256_diffusion_uncond\" #@param [\"256x256_diffusion_uncond\", \"512x512_diffusion_uncond_finetune_008100\"]\n",
        "diffusion_model = \"256x256_diffusion_uncond\" #demo-specific\n",
        "\n",
        "\n",
        "google_drive = True #@param {type:\"boolean\"}\n",
        "if google_drive:\n",
        "  from google.colab import drive\n",
        "\n",
        "#@markdown Using Google Drive is optional but recommended. If enabled, your images will be saved in \"AI/Diffusion-Demo.\" You can find newly-generated images in your \"Recent\" window when opening Google Drive.\n",
        "\n",
        "_drive_location = '/content/drive/MyDrive/AI/Diffusion-Demo/'\n",
        "\n",
        "\n",
        "#####################################################################\n",
        "##@title Google Drive & Download diffusion model\n",
        "\n",
        "model_path = '/content/'\n",
        "if google_drive:\n",
        "    yes_please = True #demo-specifc\n",
        "    print(\"NOTE: Because you are using Google Drive, the model file will be saved in \",_drive_location,\" for faster reloading\")\n",
        "    #^ demo specific\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://the-eye.eu/public/AI/models/512x512_diffusion_unconditional_ImageNet/512x512_diffusion_uncond_finetune_008100.pt' -P {model_path}\n",
        "    #_diffusion_int = 512\n",
        "if google_drive and not yes_please:\n",
        "    model_path = _drive_location\n",
        "#################################################\n",
        "\n",
        "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",
        "\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()\n",
        "\n",
        "#########################################################\n",
        "\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\n",
        "\n",
        "###########################################################\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",
        "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",
        "\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",
        "\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('Perception is quite a ride... \\n')\n",
        "        time.sleep(0.2)\n",
        "        print('Considering this is just the demo. \\n')\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",
        "\n",
        "##############################################################################################################\n",
        "#Settings and generation\n",
        "timestep_respacing = 'ddim25' \n",
        "  ##@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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U0PwzFZbLfcy",
        "cellView": "form"
      },
      "source": [
        "########################################################################\n",
        "#INITIAL VALUES\n",
        "#@title  { form-width: \"300px\" }\n",
        "#@markdown #Step 2\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",
        "_text_prompt =  \"'College Campus' art piece inspired by Edward Hopper\" #@param {type:\"string\"} \n",
        "\n",
        "_batch_genetics = False #future implementation\n",
        "_init_genetics = False  #not currently implemented\n",
        "_max_genetic_variance =  0.1\n",
        "_saturation_scale =  0\n",
        "\n",
        "#_enhance_upscale = True #@param{type:\"boolean\"}\n",
        "\n",
        "cutn = 16\n",
        "  #Controls how many crops to take from the image. Increase for higher quality.\n",
        "quality_preset = \"4\" #@param [2,4,8,16,32]\n",
        "cutn_batches = int(quality_preset) #demo-specific\n",
        "\n",
        "  #Accumulate CLIP gradient from multiple batches of cuts [Can help with OOM errors / Low VRAM]\n",
        "_esrgan_tilesize = \"128\"\n",
        "#_upscale_performance_mode = False\n",
        "##@markdown `Performance Settings`\n",
        "\n",
        "##@markdown ---\n",
        "clip_denoised = False\n",
        "fuzzy_prompt = False\n",
        "  # False - Controls whether to add multiple noisy prompts to the prompt losses\n",
        "eta =  0.5\n",
        "_clip_guidance_scale =  5000\n",
        "_tv_scale = 150\n",
        "_range_scale =  100\n",
        "_scale_multiplier = 1\n",
        "##@markdown `Visual Settings`\n",
        "seed = None #@param {type:\"raw\"}\n",
        "# seed = random.randint(0, 2**32) # Choose a random seed and print it at end of run for reproduction\n",
        "\n",
        "##@markdown ---\n",
        "_noise_mode = 'gray'\n",
        "_noise_amount = 0.1\n",
        "_init_image = ''\n",
        "_image_prompts = \"\"\n",
        "_init_scale = 1000\n",
        "skip_timesteps = 5\n",
        "n_batches =  50#@param{type:\"raw\"}\n",
        "  #Controls the starting point along the diffusion timesteps\n",
        "\n",
        "##@markdown `Generation Settings`\n",
        "\n",
        "##@markdown ---\n",
        "_project_name = 'images' #@param{type:\"string\"}\n",
        "global _debug_mode\n",
        "_debug_mode = False\n",
        "_upscale_model='RealESRGAN 4x'\n",
        "_target_resolution = \"1024\"\n",
        "_skip_upscaling = False\n",
        "display_rate =  2\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",
        "\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",
        "\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",
        "\n",
        "if _init_image == '':\n",
        "  init_image = None\n",
        "else:\n",
        "  init_image = _init_image  # None - URL or local path\n",
        "  msg_runtime = msg_runtime + 'Init image: ' + init_image + ' \\n'\n",
        "\n",
        "if _noise_amount > 0: \n",
        "  add_random_noise = True\n",
        "else:\n",
        "  add_random_noise = False\n",
        "\n",
        "# if not _init_image == '':  #to prevent noise from messing with init\n",
        "#   if add_random_noise == True:\n",
        "#     msg_runtime = msg_runtime + 'Notice: init_images have mixed results when _noise_amount > 0 \\n'\n",
        "#   # _noise_amount = 0\n",
        "#   # 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",
        "  msg_runtime = msg_runtime + 'NOTICE: You may want to disable _noise_amount when using _init_images \\n'\n",
        "\n",
        "\n",
        "rand_mag = _noise_amount # 0.1 - Controls the magnitude of the random noise\n",
        "\n",
        "#default_init_scale = 1000 #in case user forgets to set\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",
        "    msg_runtime = msg_runtime + 'Notice: init_image is set but there is no init_scale! Try 1000 as a default. \\n'\n",
        "    #init_scale = default_init_scale\n",
        "if not _image_prompts == \"\":\n",
        "  msg_runtime = msg_runtime + 'Image prompt: ' + _image_prompts + ' \\n'\n",
        "  image_prompts = [\n",
        "      _image_prompts,\n",
        "  ]\n",
        "else: image_prompts = []\n",
        "perlin_mode = _noise_mode # 'mixed' ('gray', 'color')\n",
        "\n",
        "sat_scale = _saturation_scale\n",
        "  # 0 - Controls how much saturation is allowed. From nshepperd's JAX notebook.\n",
        "\n",
        "##@markdown `skip_timesteps` best 5 (thx steven), 10 for dd50\n",
        "\n",
        "###@markdown False - Determines whether CLIP discriminates a noisy or denoised image\n",
        "# if _diffusion_int == 512: display_rate = 2\n",
        "# else: display_rate = 1\n",
        "###@markdown False - Controls whether to add multiple noisy prompts to the prompt losses\n",
        "###@markdown 0.0 - DDIM hyperparameter\n",
        "\n",
        "\n",
        "################################################\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",
        "# contains_slash = (output_folder_images.endswith('/'))\n",
        "# if not contains_slash:\n",
        "#   output_folder_images = 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",
        "#---------------------------------\n",
        "# UPSCALING\n",
        "msg_runtime = msg_runtime + 'Upscaler model: ' + _upscale_model + '; Target resolution: ' + str(_target_resolution) + ' \\n'\n",
        "if _upscale_model == '(Regular) ESRGAN 4x': _upscale_model = 'ESRGAN_SRx4_DF2KOST_official-ff704c30.pth'\n",
        "if _upscale_model == 'RealESR_NET 4x': _upscale_model = 'RealESRNet_x4plus.pth'\n",
        "if _upscale_model == 'RealESRGAN 2x': _upscale_model = 'RealESRGAN_x2plus.pth'\n",
        "if _upscale_model == 'RealESRGAN 4x': _upscale_model = 'RealESRGAN_x4plus.pth'\n",
        "if _upscale_model == 'RealESRGAN x4 Anime_6B': _upscale_model = 'RealESRGAN_x4plus_anime_6B.pth'\n",
        "upscale_model_fullpath='/content/Real-ESRGAN/experiments/pretrained_models/' + _upscale_model\n",
        "if upscale_model_fullpath == '/content/Real-ESRGAN/experiments/pretrained_models/RealESRGAN_x2plus.pth': upscale_value=\"2\"\n",
        "#else: upscale_value=\"4\"\n",
        " \n",
        "\n",
        "\n",
        "\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):\n",
        "    #returns the filename of the previously upscaled file\n",
        "    #if not completed_passes == needed_passes:\n",
        "    stripped_filename = os.path.splitext(filename)[0] #get non-upscaled image filename\n",
        "    #stripped_filename = os.path.join(_project_name,stripped_filename)\n",
        "    stripped_filename = _project_name + stripped_filename\n",
        "\n",
        "    #stripped_filename = os.path.join(output_folder_images,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",
        "    # if completed_passes == needed_passes:\n",
        "    #   return\n",
        "    # else:\n",
        "    current_upscale_target = os.path.join(output_folder_images,stripped_filename)\n",
        "    #prior_upscale_target = current_upscale_target #newest addition\n",
        "    #print('Prior upscale target is ',prior_upscale_target)\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: _outscale = 4\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",
        "      # _outscale = 4\n",
        "      # needed_passes = 2\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",
        "  #double_pass = False\n",
        "  #second_pass = False\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",
        "    #--input $current_upscale_target\n",
        "    _outscale = str(_outscale)\n",
        "    if half_precision: end_flag = end_flag + '--half '\n",
        "    # if not _upscale_model == 'RealESRGAN 2x':\n",
        "    #   if next_outscale == 0:\n",
        "    #     end_flag = end_flag + '--outscale ' + _outscale + ' '\n",
        "    #   else: end_flag = end_flag + '--outscale ' + next_outscale + ' '\n",
        "\n",
        "\n",
        "    #current_upscale_target = File_addSuffix(filename,needed_passes,completed_passes)\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 2x' 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_path $upscale_model_fullpath --output $output_folder_images --tile $_esrgan_tilesize --ext $output_extension $end_flag\n",
        "    \n",
        "\n",
        "    # if double_pass: completed_passes = completed_passes + 1\n",
        "    # else: completed_passes = completed_passes + 2\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",
        "      #if double_pass: second_pass = True\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(\"the fuckn thing is\", current_upscale_target)\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 - 1:\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",
        "###############################################"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u4Ut78Lol7EN"
      },
      "source": [
        "#Example prompts\n",
        "\n",
        "- 'College Campus' art piece inspired by Edward Hopper\n",
        "- Liminal space, liminal hotel hallway rendered in unreal engine, top post on r/liminalspaces\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'*"
      ]
    }
  ]
}