Multi_Perceptor_VQGAN_+_CLIP_Local.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Multi Perceptor VQGAN + CLIP Local",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "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/MSFTserver/AI-Colab-Notebooks-Local/blob/main/Multi_Perceptor_VQGAN_%2B_CLIP_Local.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TL3qnIS2qC9K"
      },
      "source": [
        "# Multi-Perceptor VQGAN + CLIP (v.3.2021.11.29)\n",
        "by [@remi_durant](https://twitter.com/remi_durant)\n",
        "\n",
        "edited by [HostsServer](https://github.com/MSFTserver) for local runtime\n",
        "\n",
        "Lots drawn from or inspired by other colabs, chief among them is [@jbusted1](https://twitter.com/jbusted1)'s MSE regularized VQGAN + Clip, and [@RiversHaveWings](https://twitter.com/RiversHaveWings) VQGAN + Clip with Z+Quantize. Standing on the shoulders of giants. \n",
        "\n",
        "\n",
        "- Multi-clip mode sends the same cutouts to whatever clip models you want. If they have different perceptor resolutions, the cuts are generated at each required size, replaying the same augments across both scales\n",
        "- Alternate random noise generation options to use as start point (perlin, pyramid, or vqgan random z tokens)\n",
        "- MSE Loss doesn't apply if you have no init_image until after reaching the first epoch.\n",
        "- MSE epoch targets z.tensor, not z.average, to allow for more creativity\n",
        "- Grayscale augment added for better structure\n",
        "- Padding fix for perspective and affine augments to not always be black barred\n",
        "- Automatic disable of cudnn for A100\n",
        "\n",
        "Also be sure to check out my [Artist Studies](https://remidurant.com/artists/) project, to help you find good artist names to use in your prompts. \n",
        "\n",
        "![visitors](https://visitor-badge.glitch.me/badge?page_id=remi_multiclipvqgan3)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "93QxgHDB6jkr",
        "cellView": "form"
      },
      "source": [
        "#@markdown #**Check GPU type**\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "#@markdown 4Gb = Poor\n",
        "\n",
        "#@markdown 6GB = Average\n",
        "\n",
        "#@markdown 8GB = Good\n",
        "\n",
        "#@markdown 10 GB = Great\n",
        "\n",
        "#@markdown 12+GB = MUZZZ BE NICEEE.....IMMMM TRYYNNA BE LIKE YOUU BOOZZZ\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "!nvidia-smi --query-gpu=name --format=csv\n",
        "vram = !nvidia-smi --query-gpu=memory.total --format=csv\n",
        "vram = str(int(vram[1][:len(vram[1])-3])/1000)\n",
        "print(\"Total Vram: \"+vram+\" GB\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NEAdjUn3RcNm"
      },
      "source": [
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GdC-w8aCQPAQ",
        "cellView": "form"
      },
      "source": [
        "#@title memory footprint support libraries/code\n",
        "\n",
        "!pip install gputil\n",
        "!pip install psutil\n",
        "!pip install humanize\n",
        "\n",
        "import psutil\n",
        "import humanize\n",
        "import os\n",
        "import GPUtil as GPU\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PZnX4Z6ScpKh",
        "cellView": "form"
      },
      "source": [
        "#@title Print GPU details\n",
        "\n",
        "!nvidia-smi\n",
        "\n",
        "GPUs = GPU.getGPUs()\n",
        "# XXX: only one GPU on Colab and isn’t guaranteed\n",
        "gpu = GPUs[0]\n",
        "def printm():\n",
        "    process = psutil.Process(os.getpid())\n",
        "    print(\"Gen RAM Free: \" + humanize.naturalsize(psutil.virtual_memory().available), \" |     Proc size: \" + humanize.naturalsize(process.memory_info().rss))\n",
        "    print(\"GPU RAM Free: {0:.0f}MB | Used: {1:.0f}MB | Util {2:3.0f}% | Total     {3:.0f}MB\".format(gpu.memoryFree, gpu.memoryUsed, gpu.memoryUtil*100, gpu.memoryTotal))\n",
        "printm()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zNIJyIYLQRsH",
        "cellView": "form"
      },
      "source": [
        "#@title Install Dependencies\n",
        "\n",
        "# Fix for A100 issues\n",
        "#!pip install tensorflow==1.15.2\n",
        "\n",
        "# Install normal dependencies\n",
        "!git clone https://github.com/openai/CLIP\n",
        "!git clone https://github.com/CompVis/taming-transformers\n",
        "!pip install ftfy regex tqdm omegaconf pytorch-lightning\n",
        "!pip install kornia\n",
        "!pip install einops\n",
        "!pip install transformers\n",
        "!pip install fastprogress\n",
        "!pip install torchvision"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VtLgtENkRDk0",
        "cellView": "form"
      },
      "source": [
        "#@title Load libraries and variables\n",
        "import argparse\n",
        "import math\n",
        "from pathlib import Path\n",
        "import sys\n",
        "import fastprogress\n",
        "\n",
        "sys.path.append('./taming-transformers')\n",
        "\n",
        "from IPython import display\n",
        "from omegaconf import OmegaConf\n",
        "from PIL import Image\n",
        "from taming.models import cond_transformer, vqgan\n",
        "import torch\n",
        "from torch import nn, optim\n",
        "from torch.nn import functional as F\n",
        "from torchvision import transforms\n",
        "from torchvision.transforms import functional as TF\n",
        "from tqdm.notebook import tqdm\n",
        "import numpy as np\n",
        "import os.path\n",
        "from os import path\n",
        "from urllib.request import Request, urlopen\n",
        " \n",
        "from CLIP import clip\n",
        "import kornia\n",
        "import kornia.augmentation as K\n",
        "from torch.utils.checkpoint import checkpoint\n",
        "\n",
        "from matplotlib import pyplot as plt\n",
        "from fastprogress.fastprogress import master_bar, progress_bar\n",
        "import random\n",
        "import gc\n",
        "\n",
        "import re\n",
        "from datetime import datetime\n",
        "\n",
        "from base64 import b64encode\n",
        "\n",
        "import warnings\n",
        "\n",
        "warnings.filterwarnings('ignore')\n",
        "torch.set_printoptions( sci_mode=False )\n",
        "\n",
        "def noise_gen(shape, octaves=5):\n",
        "    n, c, h, w = shape\n",
        "    noise = torch.zeros([n, c, 1, 1])\n",
        "    max_octaves = min(octaves, math.log(h)/math.log(2), math.log(w)/math.log(2))\n",
        "    for i in reversed(range(max_octaves)):\n",
        "        h_cur, w_cur = h // 2**i, w // 2**i\n",
        "        noise = F.interpolate(noise, (h_cur, w_cur), mode='bicubic', align_corners=False)\n",
        "        noise += torch.randn([n, c, h_cur, w_cur]) / 5\n",
        "    return noise\n",
        "\n",
        "\n",
        "def sinc(x):\n",
        "    return torch.where(x != 0, torch.sin(math.pi * x) / (math.pi * x), x.new_ones([]))\n",
        "\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",
        "\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",
        "\n",
        "def resample(input, size, align_corners=True):\n",
        "    n, c, h, w = input.shape\n",
        "    dh, dw = size\n",
        "\n",
        "    input = input.view([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.view([n, c, h, w])\n",
        "    return F.interpolate(input, size, mode='bicubic', align_corners=align_corners)\n",
        "    \n",
        "\n",
        "# def replace_grad(fake, real):\n",
        "#     return fake.detach() - real.detach() + real\n",
        "\n",
        "\n",
        "class ReplaceGrad(torch.autograd.Function):\n",
        "    @staticmethod\n",
        "    def forward(ctx, x_forward, x_backward):\n",
        "        ctx.shape = x_backward.shape\n",
        "        return x_forward\n",
        "\n",
        "    @staticmethod\n",
        "    def backward(ctx, grad_in):\n",
        "        return None, grad_in.sum_to_size(ctx.shape)\n",
        "\n",
        "\n",
        "class ClampWithGrad(torch.autograd.Function):\n",
        "    @staticmethod\n",
        "    def forward(ctx, input, min, max):\n",
        "        ctx.min = min\n",
        "        ctx.max = max\n",
        "        ctx.save_for_backward(input)\n",
        "        return input.clamp(min, max)\n",
        "\n",
        "    @staticmethod\n",
        "    def backward(ctx, grad_in):\n",
        "        input, = ctx.saved_tensors\n",
        "        return grad_in * (grad_in * (input - input.clamp(ctx.min, ctx.max)) >= 0), None, None\n",
        "\n",
        "replace_grad = ReplaceGrad.apply\n",
        "\n",
        "clamp_with_grad = ClampWithGrad.apply\n",
        "# clamp_with_grad = torch.clamp\n",
        "\n",
        "def vector_quantize(x, codebook):\n",
        "    d = x.pow(2).sum(dim=-1, keepdim=True) + codebook.pow(2).sum(dim=1) - 2 * x @ codebook.T\n",
        "    indices = d.argmin(-1)\n",
        "    x_q = F.one_hot(indices, codebook.shape[0]).to(d.dtype) @ codebook\n",
        "    return replace_grad(x_q, x)\n",
        "\n",
        "\n",
        "class Prompt(nn.Module):\n",
        "    def __init__(self, embed, weight=1., stop=float('-inf')):\n",
        "        super().__init__()\n",
        "        self.register_buffer('embed', embed)\n",
        "        self.register_buffer('weight', torch.as_tensor(weight))\n",
        "        self.register_buffer('stop', torch.as_tensor(stop))\n",
        "\n",
        "    def forward(self, input):\n",
        "        \n",
        "        input_normed = F.normalize(input.unsqueeze(1), dim=2)#(input / input.norm(dim=-1, keepdim=True)).unsqueeze(1)# \n",
        "        embed_normed = F.normalize((self.embed).unsqueeze(0), dim=2)#(self.embed / self.embed.norm(dim=-1, keepdim=True)).unsqueeze(0)#\n",
        "\n",
        "        dists = input_normed.sub(embed_normed).norm(dim=2).div(2).arcsin().pow(2).mul(2)\n",
        "        dists = dists * self.weight.sign()\n",
        "        \n",
        "        return self.weight.abs() * replace_grad(dists, torch.maximum(dists, self.stop)).mean()\n",
        "\n",
        "\n",
        "def parse_prompt(prompt):\n",
        "    vals = prompt.rsplit(':', 2)\n",
        "    vals = vals + ['', '1', '-inf'][len(vals):]\n",
        "    return vals[0], float(vals[1]), float(vals[2])\n",
        "\n",
        "def one_sided_clip_loss(input, target, labels=None, logit_scale=100):\n",
        "    input_normed = F.normalize(input, dim=-1)\n",
        "    target_normed = F.normalize(target, dim=-1)\n",
        "    logits = input_normed @ target_normed.T * logit_scale\n",
        "    if labels is None:\n",
        "        labels = torch.arange(len(input), device=logits.device)\n",
        "    return F.cross_entropy(logits, labels)\n",
        "\n",
        "class MakeCutouts(nn.Module):\n",
        "    def __init__(self, cut_size, cutn, cut_pow=1.):\n",
        "        super().__init__()\n",
        "        self.cut_size = cut_size\n",
        "        self.cutn = cutn\n",
        "        self.cut_pow = cut_pow\n",
        "\n",
        "        self.av_pool = nn.AdaptiveAvgPool2d((self.cut_size, self.cut_size))\n",
        "        self.max_pool = nn.AdaptiveMaxPool2d((self.cut_size, self.cut_size))\n",
        "\n",
        "    def set_cut_pow(self, cut_pow):\n",
        "        self.cut_pow = cut_pow\n",
        "\n",
        "    def forward(self, input):\n",
        "        sideY, sideX = input.shape[2:4]\n",
        "        max_size = min(sideX, sideY)\n",
        "        min_size = min(sideX, sideY, self.cut_size)\n",
        "        cutouts = []\n",
        "        cutouts_full = []\n",
        "        \n",
        "        min_size_width = min(sideX, sideY)\n",
        "        lower_bound = float(self.cut_size/min_size_width)\n",
        "        \n",
        "        for ii in range(self.cutn):\n",
        "            size = int(min_size_width*torch.zeros(1,).normal_(mean=.8, std=.3).clip(lower_bound, 1.)) # replace .5 with a result for 224 the default large size is .95\n",
        "          \n",
        "            offsetx = torch.randint(0, sideX - size + 1, ())\n",
        "            offsety = torch.randint(0, sideY - size + 1, ())\n",
        "            cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size]\n",
        "            cutouts.append(resample(cutout, (self.cut_size, self.cut_size)))\n",
        "\n",
        "        cutouts = torch.cat(cutouts, dim=0)\n",
        "\n",
        "        return clamp_with_grad(cutouts, 0, 1)\n",
        "\n",
        "\n",
        "def load_vqgan_model(config_path, checkpoint_path):\n",
        "    config = OmegaConf.load(config_path)\n",
        "    if config.model.target == 'taming.models.vqgan.VQModel':\n",
        "        model = vqgan.VQModel(**config.model.params)\n",
        "        model.eval().requires_grad_(False)\n",
        "        model.init_from_ckpt(checkpoint_path)\n",
        "    elif config.model.target == 'taming.models.cond_transformer.Net2NetTransformer':\n",
        "        parent_model = cond_transformer.Net2NetTransformer(**config.model.params)\n",
        "        parent_model.eval().requires_grad_(False)\n",
        "        parent_model.init_from_ckpt(checkpoint_path)\n",
        "        model = parent_model.first_stage_model\n",
        "    elif config.model.target == 'taming.models.vqgan.GumbelVQ':\n",
        "        model = vqgan.GumbelVQ(**config.model.params)\n",
        "        model.eval().requires_grad_(False)\n",
        "        model.init_from_ckpt(checkpoint_path)\n",
        "    else:\n",
        "        raise ValueError(f'unknown model type: {config.model.target}')\n",
        "    del model.loss\n",
        "    return model\n",
        "\n",
        "def resize_image(image, out_size):\n",
        "    ratio = image.size[0] / image.size[1]\n",
        "    area = min(image.size[0] * image.size[1], out_size[0] * out_size[1])\n",
        "    size = round((area * ratio)**0.5), round((area / ratio)**0.5)\n",
        "    return image.resize(size, Image.LANCZOS)\n",
        "\n",
        "class GaussianBlur2d(nn.Module):\n",
        "    def __init__(self, sigma, window=0, mode='reflect', value=0):\n",
        "        super().__init__()\n",
        "        self.mode = mode\n",
        "        self.value = value\n",
        "        if not window:\n",
        "            window = max(math.ceil((sigma * 6 + 1) / 2) * 2 - 1, 3)\n",
        "        if sigma:\n",
        "            kernel = torch.exp(-(torch.arange(window) - window // 2)**2 / 2 / sigma**2)\n",
        "            kernel /= kernel.sum()\n",
        "        else:\n",
        "            kernel = torch.ones([1])\n",
        "        self.register_buffer('kernel', kernel)\n",
        "\n",
        "    def forward(self, input):\n",
        "        n, c, h, w = input.shape\n",
        "        input = input.view([n * c, 1, h, w])\n",
        "        start_pad = (self.kernel.shape[0] - 1) // 2\n",
        "        end_pad = self.kernel.shape[0] // 2\n",
        "        input = F.pad(input, (start_pad, end_pad, start_pad, end_pad), self.mode, self.value)\n",
        "        input = F.conv2d(input, self.kernel[None, None, None, :])\n",
        "        input = F.conv2d(input, self.kernel[None, None, :, None])\n",
        "        return input.view([n, c, h, w])\n",
        "\n",
        "class EMATensor(nn.Module):\n",
        "    \"\"\"implmeneted by Katherine Crowson\"\"\"\n",
        "    def __init__(self, tensor, decay):\n",
        "        super().__init__()\n",
        "        self.tensor = nn.Parameter(tensor)\n",
        "        self.register_buffer('biased', torch.zeros_like(tensor))\n",
        "        self.register_buffer('average', torch.zeros_like(tensor))\n",
        "        self.decay = decay\n",
        "        self.register_buffer('accum', torch.tensor(1.))\n",
        "        self.update()\n",
        "    \n",
        "    @torch.no_grad()\n",
        "    def update(self):\n",
        "        if not self.training:\n",
        "            raise RuntimeError('update() should only be called during training')\n",
        "\n",
        "        self.accum *= self.decay\n",
        "        self.biased.mul_(self.decay)\n",
        "        self.biased.add_((1 - self.decay) * self.tensor)\n",
        "        self.average.copy_(self.biased)\n",
        "        self.average.div_(1 - self.accum)\n",
        "\n",
        "    def forward(self):\n",
        "        if self.training:\n",
        "            return self.tensor\n",
        "        return self.average\n",
        "  \n",
        "import io\n",
        "import base64\n",
        "def image_to_data_url(img, ext):  \n",
        "    img_byte_arr = io.BytesIO()\n",
        "    img.save(img_byte_arr, format=ext)\n",
        "    img_byte_arr = img_byte_arr.getvalue()\n",
        "    # ext = filename.split('.')[-1]\n",
        "    prefix = f'data:image/{ext};base64,'\n",
        "    return prefix + base64.b64encode(img_byte_arr).decode('utf-8')\n",
        " \n",
        "\n",
        "def update_random( seed, purpose ):\n",
        "  if seed == -1:\n",
        "    seed = random.seed()\n",
        "    seed = random.randrange(1,99999)\n",
        "    \n",
        "  print( f'Using seed {seed} for {purpose}')\n",
        "  random.seed(seed)\n",
        "  torch.manual_seed(seed)\n",
        "  np.random.seed(seed)\n",
        "\n",
        "  return seed\n",
        "\n",
        "def clear_memory():\n",
        "  gc.collect()\n",
        "  torch.cuda.empty_cache()\n",
        "  \n",
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KDvGKLWVKqj1",
        "cellView": "form"
      },
      "source": [
        "#@title Loss Module Definitions\n",
        "from typing import cast, Dict, Optional\n",
        "from kornia.augmentation.base import IntensityAugmentationBase2D\n",
        "\n",
        "class FixPadding(nn.Module):\n",
        "    \n",
        "    def __init__(self, module=None, threshold=1e-12, noise_frac=0.00 ):\n",
        "        super().__init__()\n",
        "\n",
        "        self.threshold = threshold\n",
        "        self.noise_frac = noise_frac\n",
        "\n",
        "        self.module = module\n",
        "\n",
        "    def forward(self,input):\n",
        "\n",
        "        dims = input.shape\n",
        "\n",
        "        if self.module is not None:\n",
        "            input = self.module(input + self.threshold)\n",
        "\n",
        "        light = input.new_empty(dims[0],1,1,1).uniform_(0.,2.)\n",
        "\n",
        "        mixed = input.view(*dims[:2],-1).sum(dim=1,keepdim=True)\n",
        "\n",
        "        black = mixed < self.threshold\n",
        "        black = black.view(-1,1,*dims[2:4]).type(torch.float)\n",
        "        black = kornia.filters.box_blur( black, (5,5) ).clip(0,0.1)/0.1\n",
        "\n",
        "        mean = input.view(*dims[:2],-1).sum(dim=2) / mixed.count_nonzero(dim=2)\n",
        "        mean = ( mean[:,:,None,None] * light ).clip(0,1)\n",
        "\n",
        "        fill = mean.expand(*dims)\n",
        "        if 0 < self.noise_frac:\n",
        "            rng = torch.get_rng_state()\n",
        "            fill = fill + torch.randn_like(mean) * self.noise_frac\n",
        "            torch.set_rng_state(rng)\n",
        "        \n",
        "        if self.module is not None:\n",
        "            input = input - self.threshold\n",
        "\n",
        "        return torch.lerp(input,fill,black)\n",
        "\n",
        "\n",
        "class MyRandomNoise(IntensityAugmentationBase2D):\n",
        "    def __init__(\n",
        "        self,\n",
        "        frac: float = 0.1,\n",
        "        return_transform: bool = False,\n",
        "        same_on_batch: bool = False,\n",
        "        p: float = 0.5,\n",
        "    ) -> None:\n",
        "        super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=1.0)\n",
        "        self.frac = frac\n",
        "\n",
        "    def __repr__(self) -> str:\n",
        "        return self.__class__.__name__ + f\"({super().__repr__()})\"\n",
        "\n",
        "    def generate_parameters(self, shape: torch.Size) -> Dict[str, torch.Tensor]:\n",
        "        noise = torch.FloatTensor(1).uniform_(0,self.frac)\n",
        "        \n",
        "        # generate pixel data without throwing off determinism of augs\n",
        "        rng = torch.get_rng_state()\n",
        "        noise = noise * torch.randn(shape)\n",
        "        torch.set_rng_state(rng)\n",
        "\n",
        "        return dict(noise=noise)\n",
        "\n",
        "    def apply_transform(\n",
        "        self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None\n",
        "    ) -> torch.Tensor:\n",
        "        return input + params['noise'].to(input.device)\n",
        "\n",
        "class MakeCutouts2(nn.Module):\n",
        "    def __init__(self, cut_size, cutn):\n",
        "        super().__init__()\n",
        "        self.cut_size = cut_size\n",
        "        self.cutn = cutn\n",
        "\n",
        "    def forward(self, input):\n",
        "        sideY, sideX = input.shape[2:4]\n",
        "        max_size = min(sideX, sideY)\n",
        "        min_size = min(sideX, sideY, self.cut_size)\n",
        "        cutouts = []\n",
        "        cutouts_full = []\n",
        "        \n",
        "        min_size_width = min(sideX, sideY)\n",
        "        lower_bound = float(self.cut_size/min_size_width)\n",
        "        \n",
        "        for ii in range(self.cutn):\n",
        "            size = int(min_size_width*torch.zeros(1,).normal_(mean=.8, std=.3).clip(lower_bound, 1.)) # replace .5 with a result for 224 the default large size is .95\n",
        "          \n",
        "            offsetx = torch.randint(0, sideX - size + 1, ())\n",
        "            offsety = torch.randint(0, sideY - size + 1, ())\n",
        "            cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size]\n",
        "            cutouts.append(cutout)\n",
        "        \n",
        "        return cutouts\n",
        "\n",
        "\n",
        "class MultiClipLoss(nn.Module):\n",
        "    def __init__(self, clip_models, text_prompt, normalize_prompt_weights, cutn, cut_pow=1., clip_weight=1., use_old_augs=False, simulate_old_cuts=False ):\n",
        "        super().__init__()\n",
        "\n",
        "        self.use_old_augs = use_old_augs\n",
        "        self.simulate_old_cuts = simulate_old_cuts \n",
        "\n",
        "        # Load Clip\n",
        "        self.perceptors = []\n",
        "        for cm in clip_models:\n",
        "          c = clip.load(cm[0], jit=False)[0].eval().requires_grad_(False).to(device)\n",
        "          self.perceptors.append( { 'res': c.visual.input_resolution, 'perceptor': c, 'weight': cm[1], 'prompts':[] } )        \n",
        "        self.perceptors.sort(key=lambda e: e['res'], reverse=True)\n",
        "        \n",
        "        # Make Cutouts\n",
        "        self.cut_sizes = list(set([p['res'] for p in self.perceptors]))\n",
        "        self.cut_sizes.sort( reverse=True )\n",
        "        \n",
        "        self.make_cuts = MakeCutouts2(self.cut_sizes[-1], cutn)\n",
        "\n",
        "        # Get Prompt Embedings\n",
        "        texts = [phrase.strip() for phrase in text_prompt.split(\"|\")]\n",
        "        if text_prompt == ['']:\n",
        "          texts = []\n",
        "\n",
        "        self.pMs = []\n",
        "\n",
        "        prompts_weight_sum = 0\n",
        "        parsed_prompts = []\n",
        "        for prompt in texts:\n",
        "          txt, weight, stop = parse_prompt(prompt)\n",
        "          parsed_prompts.append( [txt,weight,stop] )\n",
        "          prompts_weight_sum += max( weight, 0 )\n",
        "\n",
        "        for prompt in parsed_prompts:\n",
        "          txt, weight, stop = prompt\n",
        "          clip_token = clip.tokenize(txt).to(device)\n",
        "\n",
        "          if normalize_prompt_weights and 0 < prompts_weight_sum:\n",
        "              weight /= prompts_weight_sum\n",
        "\n",
        "          for p in self.perceptors:\n",
        "            embed = p['perceptor'].encode_text(clip_token).float()\n",
        "            embed_normed = F.normalize(embed.unsqueeze(0), dim=2)\n",
        "            p['prompts'].append({'embed_normed':embed_normed,'weight':torch.as_tensor(weight, device=device),'stop':torch.as_tensor(stop, device=device)})\n",
        "    \n",
        "        # Prep Augments\n",
        "        self.noise_fac = 0.1\n",
        "        self.normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],\n",
        "                                 std=[0.26862954, 0.26130258, 0.27577711])        \n",
        "        \n",
        "        self.augs = nn.Sequential(\n",
        "            K.RandomHorizontalFlip(p=0.5),\n",
        "            K.RandomSharpness(0.3,p=0.1),\n",
        "            FixPadding( nn.Sequential(\n",
        "                K.RandomAffine(degrees=30, translate=0.1, p=0.8, padding_mode='zeros'), # padding_mode=2\n",
        "                K.RandomPerspective(0.2,p=0.4, ),\n",
        "            )),\n",
        "            K.ColorJitter(hue=0.01, saturation=0.01, p=0.7),\n",
        "            K.RandomGrayscale(p=0.15), \n",
        "            MyRandomNoise(frac=self.noise_fac,p=1.),\n",
        "        )\n",
        "\n",
        "        self.clip_weight = clip_weight\n",
        "\n",
        "    def prepare_cuts(self,img):\n",
        "        cutouts = self.make_cuts(img)\n",
        "        cutouts_out = []\n",
        "            \n",
        "        rng = torch.get_rng_state()\n",
        "\n",
        "        for sz in self.cut_sizes:\n",
        "            cuts = [resample(c, (sz,sz)) for c in cutouts]\n",
        "            cuts = torch.cat(cuts, dim=0)\n",
        "            cuts = clamp_with_grad(cuts,0,1)\n",
        "\n",
        "            torch.set_rng_state(rng)\n",
        "            cuts = self.augs(cuts)\n",
        "            cuts = self.normalize(cuts)\n",
        "\n",
        "            cutouts_out.append(cuts)\n",
        "\n",
        "        return cutouts_out\n",
        "\n",
        "    def forward( self, i, img ):\n",
        "        cutouts = self.prepare_cuts( img )\n",
        "        loss = []\n",
        "        \n",
        "        current_cuts = None\n",
        "        currentres = 0\n",
        "        \n",
        "        for p in self.perceptors:\n",
        "            if currentres != p['res']:\n",
        "                currentres = p['res']\n",
        "                current_cuts = cutouts[self.cut_sizes.index( currentres )]\n",
        "\n",
        "            iii = p['perceptor'].encode_image(current_cuts).float()\n",
        "            input_normed = F.normalize(iii.unsqueeze(1), dim=2)\n",
        "            for prompt in p['prompts']:\n",
        "                dists = input_normed.sub(prompt['embed_normed']).norm(dim=2).div(2).arcsin().pow(2).mul(2)\n",
        "                dists = dists * prompt['weight'].sign()\n",
        "                l = prompt['weight'].abs() * replace_grad(dists, torch.maximum(dists, prompt['stop'])).mean()\n",
        "                loss.append(l * p['weight'])\n",
        "\n",
        "        return loss\n",
        "\n",
        "class MSEDecayLoss(nn.Module):\n",
        "    def __init__(self, init_weight, mse_decay_rate, mse_epoches, mse_quantize ):\n",
        "        super().__init__()\n",
        "      \n",
        "        self.init_weight = init_weight\n",
        "        self.has_init_image = False\n",
        "        self.mse_decay = init_weight / mse_epoches if init_weight else 0 \n",
        "        self.mse_decay_rate = mse_decay_rate\n",
        "        self.mse_weight = init_weight\n",
        "        self.mse_epoches = mse_epoches\n",
        "        self.mse_quantize = mse_quantize\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def set_target( self, z_tensor, model ):\n",
        "        z_tensor = z_tensor.detach().clone()\n",
        "        if self.mse_quantize:\n",
        "            z_tensor = vector_quantize(z_tensor.movedim(1, 3), model.quantize.embedding.weight).movedim(3, 1)#z.average\n",
        "        self.z_orig = z_tensor\n",
        "          \n",
        "    def forward( self, i, z ):\n",
        "        if self.is_active(i):\n",
        "            return F.mse_loss(z, self.z_orig) * self.mse_weight / 2\n",
        "        return 0\n",
        "        \n",
        "    def is_active(self, i):\n",
        "        if not self.init_weight:\n",
        "          return False\n",
        "        if i <= self.mse_decay_rate and not self.has_init_image:\n",
        "          return False\n",
        "        return True\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def step( self, i ):\n",
        "\n",
        "        if i % self.mse_decay_rate == 0 and i != 0 and i < self.mse_decay_rate * self.mse_epoches:\n",
        "            \n",
        "            if self.mse_weight - self.mse_decay > 0 and self.mse_weight - self.mse_decay >= self.mse_decay:\n",
        "              self.mse_weight -= self.mse_decay\n",
        "            else:\n",
        "              self.mse_weight = 0\n",
        "            print(f\"updated mse weight: {self.mse_weight}\")\n",
        "\n",
        "            return True\n",
        "\n",
        "        return False\n",
        "  \n",
        "class TVLoss(nn.Module):\n",
        "    def forward(self, input):\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",
        "        diff = x_diff**2 + y_diff**2 + 1e-8\n",
        "        return diff.mean(dim=1).sqrt().mean()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "agEFaX64bjdj",
        "cellView": "form"
      },
      "source": [
        "#@title Random Inits\n",
        "\n",
        "import torch\n",
        "import math\n",
        "\n",
        "def rand_perlin_2d(shape, res, fade = lambda t: 6*t**5 - 15*t**4 + 10*t**3):\n",
        "    delta = (res[0] / shape[0], res[1] / shape[1])\n",
        "    d = (shape[0] // res[0], shape[1] // res[1])\n",
        "    \n",
        "    grid = torch.stack(torch.meshgrid(torch.arange(0, res[0], delta[0]), torch.arange(0, res[1], delta[1])), dim = -1) % 1\n",
        "    angles = 2*math.pi*torch.rand(res[0]+1, res[1]+1)\n",
        "    gradients = torch.stack((torch.cos(angles), torch.sin(angles)), dim = -1)\n",
        "    \n",
        "    tile_grads = lambda slice1, slice2: gradients[slice1[0]:slice1[1], slice2[0]:slice2[1]].repeat_interleave(d[0], 0).repeat_interleave(d[1], 1)\n",
        "    dot = lambda grad, shift: (torch.stack((grid[:shape[0],:shape[1],0] + shift[0], grid[:shape[0],:shape[1], 1] + shift[1]  ), dim = -1) * grad[:shape[0], :shape[1]]).sum(dim = -1)\n",
        "    \n",
        "    n00 = dot(tile_grads([0, -1], [0, -1]), [0,  0])\n",
        "    n10 = dot(tile_grads([1, None], [0, -1]), [-1, 0])\n",
        "    n01 = dot(tile_grads([0, -1],[1, None]), [0, -1])\n",
        "    n11 = dot(tile_grads([1, None], [1, None]), [-1,-1])\n",
        "    t = fade(grid[:shape[0], :shape[1]])\n",
        "    return math.sqrt(2) * torch.lerp(torch.lerp(n00, n10, t[..., 0]), torch.lerp(n01, n11, t[..., 0]), t[..., 1])\n",
        "\n",
        "def rand_perlin_2d_octaves( desired_shape, octaves=1, persistence=0.5):\n",
        "    shape = torch.tensor(desired_shape)\n",
        "    shape = 2 ** torch.ceil( torch.log2( shape ) )\n",
        "    shape = shape.type(torch.int)\n",
        "\n",
        "    max_octaves = int(min(octaves,math.log(shape[0])/math.log(2), math.log(shape[1])/math.log(2)))\n",
        "    res = torch.floor( shape / 2 ** max_octaves).type(torch.int)\n",
        "\n",
        "    noise = torch.zeros(list(shape))\n",
        "    frequency = 1\n",
        "    amplitude = 1\n",
        "    for _ in range(max_octaves):\n",
        "        noise += amplitude * rand_perlin_2d(shape, (frequency*res[0], frequency*res[1]))\n",
        "        frequency *= 2\n",
        "        amplitude *= persistence\n",
        "    \n",
        "    return noise[:desired_shape[0],:desired_shape[1]]\n",
        "\n",
        "def rand_perlin_rgb( desired_shape, amp=0.1, octaves=6 ):\n",
        "  r = rand_perlin_2d_octaves( desired_shape, octaves )\n",
        "  g = rand_perlin_2d_octaves( desired_shape, octaves )\n",
        "  b = rand_perlin_2d_octaves( desired_shape, octaves )\n",
        "  rgb = ( torch.stack((r,g,b)) * amp + 1 ) * 0.5\n",
        "  return rgb.unsqueeze(0).clip(0,1).to(device)\n",
        "\n",
        "\n",
        "def pyramid_noise_gen(shape, octaves=5, decay=1.):\n",
        "    n, c, h, w = shape\n",
        "    noise = torch.zeros([n, c, 1, 1])\n",
        "    max_octaves = int(min(math.log(h)/math.log(2), math.log(w)/math.log(2)))\n",
        "    if octaves is not None and 0 < octaves:\n",
        "      max_octaves = min(octaves,max_octaves)\n",
        "    for i in reversed(range(max_octaves)):\n",
        "        h_cur, w_cur = h // 2**i, w // 2**i\n",
        "        noise = F.interpolate(noise, (h_cur, w_cur), mode='bicubic', align_corners=False)\n",
        "        noise += ( torch.randn([n, c, h_cur, w_cur]) / max_octaves ) * decay**( max_octaves - (i+1) )\n",
        "    return noise\n",
        "\n",
        "def rand_z(model, toksX, toksY):\n",
        "    e_dim = model.quantize.e_dim\n",
        "    n_toks = model.quantize.n_e\n",
        "    z_min = model.quantize.embedding.weight.min(dim=0).values[None, :, None, None]\n",
        "    z_max = model.quantize.embedding.weight.max(dim=0).values[None, :, None, None]\n",
        "\n",
        "    one_hot = F.one_hot(torch.randint(n_toks, [toksY * toksX], device=device), n_toks).float()\n",
        "    z = one_hot @ model.quantize.embedding.weight\n",
        "    z = z.view([-1, toksY, toksX, e_dim]).permute(0, 3, 1, 2)\n",
        "\n",
        "    return z\n",
        "\n",
        "\n",
        "def make_rand_init( mode, model, perlin_octaves, perlin_weight, pyramid_octaves, pyramid_decay, toksX, toksY, f ):\n",
        "\n",
        "  if mode == 'VQGAN ZRand':\n",
        "    return rand_z(model, toksX, toksY)\n",
        "  elif mode == 'Perlin Noise':\n",
        "    rand_init = rand_perlin_rgb((toksY * f, toksX * f), perlin_weight, perlin_octaves )\n",
        "    z, *_ = model.encode(rand_init * 2 - 1)\n",
        "    return z\n",
        "  elif mode == 'Pyramid Noise':\n",
        "    rand_init = pyramid_noise_gen( (1,3,toksY * f, toksX * f), pyramid_octaves, pyramid_decay).to(device)\n",
        "    rand_init = ( rand_init * 0.5 + 0.5 ).clip(0,1)\n",
        "    z, *_ = model.encode(rand_init * 2 - 1)\n",
        "    return z\n",
        "    \n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lLJmiZ8mRks_"
      },
      "source": [
        "# Make some Art!"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown # 1) Download Models\n",
        "#@markdown Download and rename the files to the model name and stick them in the root directory where you started jupyter\n",
        "\n",
        "#@markdown imagenet_1024 - 913 MB\n",
        "#@markdown [vqgan_imagenet_f16_1024.yaml](https://heibox.uni-heidelberg.de/d/8088892a516d4e3baf92/files/?p=%2Fconfigs%2Fmodel.yaml&dl=1)\n",
        "#@markdown <> [vqgan_imagenet_f16_1024.ckpt](https://heibox.uni-heidelberg.de/d/8088892a516d4e3baf92/files/?p=%2Fckpts%2Flast.ckpt&dl=1)\n",
        "\n",
        "#@markdown imagenet_16384 - 934 MB\n",
        "#@markdown [vqgan_imagenet_f16_16384.yaml](https://heibox.uni-heidelberg.de/d/a7530b09fed84f80a887/files/?p=%2Fconfigs%2Fmodel.yaml&dl=1)\n",
        "#@markdown <> [vqgan_imagenet_f16_16384.ckpt](https://heibox.uni-heidelberg.de/d/a7530b09fed84f80a887/files/?p=%2Fckpts%2Flast.ckpt&dl=1)\n",
        "\n",
        "#@markdown openimages_8192 - 359 MB\n",
        "#@markdown [vqgan_openimages_f16_8192.yaml](https://www.dropbox.com/s/zfstzj17xztmd8c/vqgan_openimages_f16_8192.yaml?dl=1)\n",
        "#@markdown <> [vqgan_openimages_f16_8192.ckpt](https://www.dropbox.com/s/216igtr2hod4jpg/vqgan_openimages_f16_8192.ckpt?dl=1)\n",
        "\n",
        "#@markdown gumbel_8192 - 359 MB\n",
        "#@markdown [gumbel_8192.yaml](https://heibox.uni-heidelberg.de/d/2e5662443a6b4307b470/files/?p=%2Fconfigs%2Fmodel.yaml&dl=1)\n",
        "#@markdown <> [gumbel_8192.ckpt](https://heibox.uni-heidelberg.de/d/2e5662443a6b4307b470/files/?p=%2Fckpts%2Flast.ckpt&dl=1)\n",
        "\n",
        "#@markdown faceshq - 3.70 GB\n",
        "#@markdown [faceshq.yaml](https://drive.google.com/uc?export=download&id=1fHwGx_hnBtC8nsq7hesJvs-Klv-P0gzT)\n",
        "#@markdown <> [faceshq.ckpt](https://app.koofr.net/content/links/a04deec9-0c59-4673-8b37-3d696fe63a5d/files/get/last.ckpt?path=%2F2020-11-13T21-41-45_faceshq_transformer%2Fcheckpoints%2Flast.ckpt)\n",
        "\n",
        "#@markdown coco - 7.85 GB\n",
        "#@markdown [coco.yaml](https://dl.nmkd.de/ai/clip/coco/coco.yaml)\n",
        "#@markdown <> [coco.ckpt](https://dl.nmkd.de/ai/clip/coco/coco.ckpt)\n",
        "\n",
        "#@markdown wikiart_1024 - 913 MB\n",
        "#@markdown [wikiart_1024.yaml](https://www.dropbox.com/s/pruzqc6unxzxo71/wikiart_1024.yaml?dl=1)\n",
        "#@markdown <> [wikiart_1024.ckpt](https://www.dropbox.com/s/o44ymgfay4ps0i9/wikiart_1024.ckpt?dl=1)\n",
        "\n",
        "#@markdown wikiart_16384 - 958 MB\n",
        "#@markdown [wikiart_16384.yaml](http://eaidata.bmk.sh/data/Wikiart_16384/wikiart_f16_16384_8145600.yaml)\n",
        "#@markdown <> [wikiart_16384.ckpt](http://eaidata.bmk.sh/data/Wikiart_16384/wikiart_f16_16384_8145600.ckpt)\n",
        "\n",
        "#@markdown sflckr - 3.97 GB\n",
        "#@markdown [sflckr.yaml](https://heibox.uni-heidelberg.de/d/73487ab6e5314cb5adba/files/?p=%2Fconfigs%2F2020-11-09T13-31-51-project.yaml&dl=1)\n",
        "#@markdown <> [sflckr.ckpt](https://heibox.uni-heidelberg.de/d/73487ab6e5314cb5adba/files/?p=%2Fcheckpoints%2Flast.ckpt&dl=1)\n",
        "\n",
        "#@markdown ade20k - 4.61 GB\n",
        "#@markdown [ade20k.yaml](https://www.dropbox.com/s/wdlz8e30a521hb5/ade20k.yaml?dl=1)\n",
        "#@markdown <> [ade20k.ckpt](https://www.dropbox.com/s/znqi0ovzp4wdjn9/ade20k.ckpt?dl=1)\n",
        "\n",
        "#@markdown ffhq - 9.2 GB\n",
        "#@markdown [ffhq.yaml](https://app.koofr.net/content/links/0fc005bf-3dca-4079-9d40-cdf38d42cd7a/files/get/2021-04-23T18-19-01-project.yaml?path=%2F2021-04-23T18-19-01_ffhq_transformer%2Fconfigs%2F2021-04-23T18-19-01-project.yaml&force)\n",
        "#@markdown <> [ffhq.ckpt](https://app.koofr.net/content/links/0fc005bf-3dca-4079-9d40-cdf38d42cd7a/files/get/last.ckpt?path=%2F2021-04-23T18-19-01_ffhq_transformer%2Fcheckpoints%2Flast.ckpt&force)\n",
        "\n",
        "#@markdown celebahq - 9.2 GB\n",
        "#@markdown [celebahq.yaml](https://app.koofr.net/content/links/6dddf083-40c8-470a-9360-a9dab2a94e96/files/get/2021-04-23T18-11-19-project.yaml?path=%2F2021-04-23T18-11-19_celebahq_transformer%2Fconfigs%2F2021-04-23T18-11-19-project.yaml&force)\n",
        "#@markdown <> [celebahq.ckpt](https://app.koofr.net/content/links/6dddf083-40c8-470a-9360-a9dab2a94e96/files/get/last.ckpt?path=%2F2021-04-23T18-11-19_celebahq_transformer%2Fcheckpoints%2Flast.ckpt&force)\n",
        "\n",
        "\n",
        "\n",
        "filename= os.listdir(os.getcwd())\n",
        "os.startfile(os.getcwd())"
      ],
      "metadata": {
        "cellView": "form",
        "id": "5PVbg4zq_XSF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nxcx0j6zmePl",
        "cellView": "form"
      },
      "source": [
        "#@title Set VQGAN Model Save Location\n",
        "#@markdown path to models, leave blank for working directory\n",
        "model_path = \"\" #@param {type: 'string'}\n",
        "vqgan_model = 'vqgan_imagenet_f16_16384'#@param [ \"vqgan_imagenet_f16_1024\", \"vqgan_imagenet_f16_16384\", \"vqgan_openimages_f8_8192\", \"coco\", \"faceshq\",\"wikiart_1024\", \"wikiart_16384\", \"sflckr\"]\n",
        "#@markdown Path to Save Images\n",
        "save_output_path = \"Art/\" #@param {type: 'string'}\n",
        "\n",
        "#@markdown When saving the images, how much should be included in the name?\n",
        "include_full_prompt_in_filename = False #@param {type:'boolean'}\n",
        "shortname_limit = 50 #@param {type: 'number'}\n",
        "filename_limit = 250 #@param {type: 'number'}\n",
        "\n",
        "if model_path:\n",
        "  !mkdir \"$model_path\"\n",
        "\n",
        "if save_output_path:\n",
        "  !mkdir \"$save_output_path\"\n",
        "\n",
        "save_output_path += \"/\" if not save_output_path.endswith('/') else \"\"\n",
        "model_path += \"/\" if not model_path.endswith('/') else \"\"\n",
        "loaded_model = None\n",
        "loaded_model_name = None\n",
        "\n",
        "def get_vqgan_model(image_model):\n",
        "    global loaded_model\n",
        "    global loaded_model_name\n",
        "    if loaded_model is None or loaded_model_name != image_model:\n",
        "\n",
        "        print(f'loading {image_model} vqgan checkpoint')\n",
        "\n",
        "        vqgan_config= os.getcwd() + model_path + vqgan_model + '.yaml'\n",
        "        vqgan_checkpoint= os.getcwd() + model_path + vqgan_model + '.ckpt'\n",
        "        print('vqgan_config',vqgan_config)\n",
        "        print('vqgan_checkpoint',vqgan_checkpoint)\n",
        "\n",
        "        model = load_vqgan_model(vqgan_config, vqgan_checkpoint).to(device)\n",
        "        if image_model == 'vqgan_openimages_f8_8192':\n",
        "            model.quantize.e_dim = 256\n",
        "            model.quantize.n_e = model.quantize.n_embed\n",
        "            model.quantize.embedding = model.quantize.embed\n",
        "\n",
        "        loaded_model = model\n",
        "        loaded_model_name = image_model\n",
        "\n",
        "    return loaded_model\n",
        "\n",
        "def slugify(value):\n",
        "    value = str(value)\n",
        "    value = re.sub(r':([-\\d.]+)', ' [\\\\1]', value)\n",
        "    value = re.sub(r'[|]','; ',value)\n",
        "    value = re.sub(r'[<>:\"/\\\\|?*]', ' ', value)\n",
        "    return value\n",
        "\n",
        "def get_filename(text, seed, i, ext):\n",
        "    if ( not include_full_prompt_in_filename ):\n",
        "        text = re.split(r'[|:;]',text, 1)[0][:shortname_limit]\n",
        "    text = slugify(text)\n",
        "\n",
        "    now = datetime.now()\n",
        "    t = now.strftime(\"%y%m%d%H%M\")\n",
        "    if i is not None:\n",
        "        data = f'; r{seed} i{i} {t}{ext}'\n",
        "    else:\n",
        "        data = f'; r{seed} {t}{ext}'\n",
        "\n",
        "    return text[:filename_limit-len(data)] + data\n",
        "\n",
        "def save_output(pil, text, seed, i):\n",
        "    fname = get_filename(text,seed,i,'.png')\n",
        "    pil.save(save_output_path + fname)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AUMYhUaEbnkL",
        "cellView": "form"
      },
      "source": [
        "#@title Set Display Rate\n",
        "#@markdown If `use_automatic_display_schedule` is enabled, the image will be output frequently at first, and then more spread out as time goes on. Turn this off if you want to specify the display rate yourself.\n",
        "use_automatic_display_schedule = True #@param {type:'boolean'}\n",
        "display_every = 50 #@param {type:'number'}\n",
        "\n",
        "def should_checkin(i):\n",
        "  if i == max_iter: \n",
        "    return True \n",
        "\n",
        "  if not use_automatic_display_schedule:\n",
        "    return i % display_every == 0\n",
        "\n",
        "  schedule = [[100,25],[500,50],[1000,100],[2000,200]]\n",
        "  for s in schedule:\n",
        "    if i <= s[0]:\n",
        "      return i % s[1] == 0\n",
        "  return i % 500 == 0\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BbqnA_STd9bp"
      },
      "source": [
        "Before generating, the rest of the setup steps must first be executed by pressing **`Runtime > Run All`**. This only needs to be done once."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YWGR8q9AT8uk",
        "cellView": "form"
      },
      "source": [
        "#@title Do the Run\n",
        "\n",
        "#@markdown What do you want to see?\n",
        "text_prompt = 'The Enchanted Garden by James Gurney'#@param {type:'string'}\n",
        "gen_seed = -1#@param {type:'number'}\n",
        "#@markdown - If you want to keep starting from the same point, set `gen_seed` to a positive number. `-1` will make it random every time. \n",
        "init_image = 'C:\\\\Users\\\\HostsServer\\\\Desktop\\\\AI-Image_gens\\\\VQGAN-CLIP-main\\\\geo.jpg'#@param {type:'string'}\n",
        "width = 300#@param {type:'number'}\n",
        "height = 300#@param {type:'number'}\n",
        "max_iter = 100 #@param {type:'number'}\n",
        "\n",
        "\n",
        "#@markdown There are different ways of generating the random starting point, when not using an init image. These influence how the image turns out. The default VQGAN ZRand is good, but some models and subjects may do better with perlin or pyramid noise.\n",
        "rand_init_mode = 'VQGAN ZRand'#@param [ \"VQGAN ZRand\", \"Perlin Noise\", \"Pyramid Noise\"]\n",
        "perlin_octaves = 7#@param {type:\"slider\", min:1, max:8, step:1}\n",
        "perlin_weight = 0.22#@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "pyramid_octaves = 5#@param {type:\"slider\", min:1, max:8, step:1}\n",
        "pyramid_decay = 0.99#@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "ema_val = 0.99\n",
        "\n",
        "#@markdown How many slices of the image should be sent to CLIP each iteration to score? Higher numbers are better, but cost more memory. If you are running into memory issues try lowering this value.\n",
        "cut_n = 64 #@param {type:'number'}\n",
        "\n",
        "#@markdown One clip model is good. Two is better? You may need to reduce the number of cuts to support having more than one CLIP model. CLIP is what scores the image against your prompt and each model has slightly different ideas of what things are.\n",
        "#@markdown - `ViT-B/32` is fast and good and what most people use to begin with\n",
        "\n",
        "clip_model = 'ViT-B/32' #@param [\"ViT-B/16\", \"ViT-B/32\", \"ViT-L/14\", \"RN50x64\", \"RN50x16\", \"RN50x4\"]\n",
        "clip_model2 ='None' #@param [\"None\",\"ViT-B/16\", \"ViT-B/32\", \"ViT-L/14\", \"RN50x64\", \"RN50x16\", \"RN50x4\"]\n",
        "if clip_model2 == \"None\":\n",
        "    clip_model2 = None \n",
        "clip1_weight = 0.5 #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "\n",
        "#@markdown Picking a different VQGAN model will impact how an image generates. Think of this as giving the generator a different set of brushes and paints to work with. CLIP is still the \"eyes\" and is judging the image against your prompt but using different brushes will make a different image.\n",
        "#@markdown - `vqgan_imagenet_f16_16384` is the default and what most people use\n",
        "#@markdown - this must match the same one you selected in the code block above this one\n",
        "vqgan_model = 'vqgan_imagenet_f16_16384'#@param [ \"vqgan_imagenet_f16_1024\", \"vqgan_imagenet_f16_16384\", \"vqgan_openimages_f8_8192\", \"coco\", \"faceshq\",\"wikiart_1024\", \"wikiart_16384\", \"sflckr\"]\n",
        "\n",
        "#@markdown Learning rates greatly impact how quickly an image can generate, or if an image can generate at all. The first learning rate is only for the first 50 iterations. The epoch rate is what is used after reaching the first mse epoch. \n",
        "#@markdown You can try lowering the epoch rate while raising the initial learning rate and see what happens\n",
        "learning_rate = 0.2#@param {type:'number'}\n",
        "learning_rate_epoch = 0.2#@param {type:'number'}\n",
        "#@markdown How much should we try to match the init image, or if no init image how much should we resist change after reaching the first epoch?\n",
        "mse_weight = 0.5 #@param {type:'number'}\n",
        "#@markdown Adding some TV may make the image blurrier but also helps to get rid of noise. A good value to try might be 0.1.\n",
        "tv_weight = 0.0 #@param {type:'number'}\n",
        "#@markdown Should the total weight of the text prompts stay in the same range, relative to other loss functions?\n",
        "normalize_prompt_weights = True #@param {type:'boolean'}\n",
        "\n",
        "#@markdown Enabling the EMA tensor will cause the image to be slower to generate but may help it be more cohesive.\n",
        "#@markdown This can also help keep the final image closer to the init image, if you are providing one.\n",
        "use_ema_tensor = False #@param {type:'boolean'}\n",
        "\n",
        "#@markdown If you want to generate a video of the run, you need to save the frames as you go. The more frequently you save, the longer the video but the slower it will take to generate.\n",
        "save_art_output = True #@param {type:'boolean'}\n",
        "save_frames_for_video = True #@param {type:'boolean'}\n",
        "save_frequency_for_video = 3  #@param {type:'number'}\n",
        "\n",
        "\n",
        "#@markdown  ----\n",
        "#@markdown  I'd love to see what you can make with my notebook. Tweet me your art [@remi_durant](https://twitter.com/remi_durant)!\n",
        "\n",
        "output_as_png = True\n",
        "\n",
        "print('Using device:', device)\n",
        "print('using prompts: ', text_prompt)\n",
        "\n",
        "clear_memory()\n",
        "\n",
        "!mkdir steps\n",
        "\n",
        "model = get_vqgan_model( vqgan_model )\n",
        "\n",
        "if clip_model2:\n",
        "  clip_models = [[clip_model, clip1_weight], [clip_model2, 1. - clip1_weight]]\n",
        "else:\n",
        "  clip_models = [[clip_model, 1.0]]\n",
        "print(clip_models)\n",
        "\n",
        "clip_loss = MultiClipLoss( clip_models, text_prompt, normalize_prompt_weights=normalize_prompt_weights, cutn=cut_n)\n",
        "\n",
        "seed = update_random( gen_seed, 'image generation')\n",
        "    \n",
        "# Make Z Init\n",
        "z = 0\n",
        "\n",
        "f = 2**(model.decoder.num_resolutions - 1)\n",
        "toksX, toksY = math.ceil( width / f), math.ceil( height / f)\n",
        "\n",
        "print(f'Outputing size: [{toksX*f}x{toksY*f}]')\n",
        "\n",
        "has_init_image = (init_image != \"\")\n",
        "if has_init_image:\n",
        "    if 'http' in init_image:\n",
        "      req = Request(init_image, headers={'User-Agent': 'Mozilla/5.0'})\n",
        "      img = Image.open(urlopen(req))\n",
        "    else:\n",
        "      img = Image.open(init_image)\n",
        "\n",
        "    pil_image = img.convert('RGB')\n",
        "    pil_image = pil_image.resize((toksX * f, toksY * f), Image.LANCZOS)\n",
        "    pil_image = TF.to_tensor(pil_image)\n",
        "    #if args.use_noise:\n",
        "    #  pil_image = pil_image + args.use_noise * torch.randn_like(pil_image) \n",
        "    z, *_ = model.encode(pil_image.to(device).unsqueeze(0) * 2 - 1)\n",
        "    del pil_image\n",
        "    del img\n",
        "\n",
        "else:\n",
        "    z = make_rand_init( rand_init_mode, model, perlin_octaves, perlin_weight, pyramid_octaves, pyramid_decay, toksX, toksY, f )\n",
        "    \n",
        "z = EMATensor(z, ema_val)\n",
        "\n",
        "opt = optim.Adam( z.parameters(), lr=learning_rate, weight_decay=0.00000000)\n",
        "\n",
        "mse_loss = MSEDecayLoss( mse_weight, mse_decay_rate=50, mse_epoches=5, mse_quantize=True )\n",
        "mse_loss.set_target( z.tensor, model )\n",
        "mse_loss.has_init_image = has_init_image\n",
        "\n",
        "tv_loss = TVLoss() \n",
        "\n",
        "\n",
        "losses = []\n",
        "mb = master_bar(range(1))\n",
        "gnames = ['losses']\n",
        "\n",
        "mb.names=gnames\n",
        "mb.graph_fig, axs = plt.subplots(1, 1) # For custom display\n",
        "mb.graph_ax = axs\n",
        "mb.graph_out = display.display(mb.graph_fig, display_id=True)\n",
        "\n",
        "## optimizer loop\n",
        "\n",
        "def synth(z, quantize=True, scramble=True):\n",
        "    z_q = 0\n",
        "    if quantize:\n",
        "      z_q = vector_quantize(z.movedim(1, 3), model.quantize.embedding.weight).movedim(3, 1)\n",
        "    else:\n",
        "      z_q = z.model\n",
        "\n",
        "    out = clamp_with_grad(model.decode(z_q).add(1).div(2), 0, 1)\n",
        "\n",
        "    return out\n",
        "\n",
        "@torch.no_grad()\n",
        "def checkin(i, z, out_pil, losses):\n",
        "    losses_str = ', '.join(f'{loss.item():g}' for loss in losses)\n",
        "    tqdm.write(f'i: {i}, loss: {sum(losses).item():g}, losses: {losses_str}')\n",
        "\n",
        "    display_format='png' if output_as_png else 'jpg'\n",
        "    pil_data = image_to_data_url(out_pil, display_format)\n",
        "    \n",
        "    display.display(display.HTML(f'<img src=\"{pil_data}\" />'))\n",
        "\n",
        "def should_save_for_video(i):\n",
        "    return save_frames_for_video and i % save_frequency_for_video\n",
        "\n",
        "def train(i):\n",
        "    global opt\n",
        "    global z \n",
        "    opt.zero_grad( set_to_none = True )\n",
        "\n",
        "    out = checkpoint( synth, z.tensor )\n",
        "\n",
        "    lossAll = []\n",
        "    lossAll += clip_loss( i,out )\n",
        "\n",
        "    if 0 < mse_weight:\n",
        "      msel = mse_loss(i,z.tensor)\n",
        "      if 0 < msel:\n",
        "        lossAll.append(msel)\n",
        "    \n",
        "    if 0 < tv_weight:\n",
        "      lossAll.append(tv_loss(out)*tv_weight)\n",
        "    \n",
        "    loss = sum(lossAll)\n",
        "    loss.backward()\n",
        "\n",
        "    if should_checkin(i) or should_save_for_video(i):\n",
        "        with torch.no_grad():\n",
        "            if use_ema_tensor:\n",
        "                out = synth( z.average )\n",
        "\n",
        "            pil = TF.to_pil_image(out[0].cpu())\n",
        "\n",
        "            if should_checkin(i):\n",
        "                checkin(i, z, pil, lossAll)\n",
        "                if save_art_output:\n",
        "                    save_output(pil, text_prompt, seed, i)\n",
        "                    \n",
        "            if should_save_for_video(i):\n",
        "                pil.save(f'steps/step{i//save_frequency_for_video:04}.png')\n",
        "                \n",
        "\n",
        "i = 0\n",
        "try:\n",
        "  with tqdm() as pbar:\n",
        "    while True and i <= max_iter:\n",
        "      \n",
        "      if i % 200 == 0:\n",
        "        clear_memory()\n",
        "\n",
        "      train(i)\n",
        "\n",
        "      with torch.no_grad():\n",
        "          if mse_loss.step(i):\n",
        "              print('Reseting optimizer at mse epoch')\n",
        "\n",
        "              if mse_loss.has_init_image and use_ema_tensor:\n",
        "                mse_loss.set_target(z.average,model)\n",
        "              else:\n",
        "                mse_loss.set_target(z.tensor,model)\n",
        "              \n",
        "              # Make sure not to spike loss when mse_loss turns on\n",
        "              if not mse_loss.is_active(i):\n",
        "                z.tensor = nn.Parameter(mse_loss.z_orig.clone())\n",
        "                z.tensor.requires_grad = True\n",
        "\n",
        "              if use_ema_tensor:\n",
        "                z = EMATensor(z.average, ema_val)\n",
        "              else:\n",
        "                z = EMATensor(z.tensor, ema_val)\n",
        "\n",
        "              opt = optim.Adam(z.parameters(), lr=learning_rate_epoch, weight_decay=0.00000000)\n",
        "\n",
        "      i += 1\n",
        "      pbar.update()\n",
        "\n",
        "except KeyboardInterrupt:\n",
        "    pass\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Iaw63lWBqsnj",
        "cellView": "form"
      },
      "source": [
        "#@title Make a Video of Your Last Run!\n",
        "#@markdown If you want to make a video, you must first enable `save_frames_for_video` during the run. Setting a higher frequency will make a longer video, and a higher framerate will make a shorter video.\n",
        "\n",
        "fps = 24 #@param{type:'number'}\n",
        "\n",
        "!mkdir \"video\"\n",
        "vname = os.getcwd()+\"\\\\video\\\\\"+get_filename(text_prompt,seed,None,'.mp4')\n",
        "\n",
        "!ffmpeg -y -v 1 -framerate $fps -i \"steps/step%04d.png\" -r $fps -vcodec libx264 -crf 32 -pix_fmt yuv420p \"$vname\"\n",
        "\n",
        "mp4 = open(vname,'rb').read()\n",
        "data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "display.display( display.HTML(f'<video controls><source src=\"{data_url}\" type=\"video/mp4\"></video>') )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @markdown # Clear GPU memory\n",
        "# @markdown  incase of constant out of memory issues run this code block\n",
        "torch.cuda.empty_cache()\n",
        "import gc\n",
        "gc.collect()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "QS1TD4tQlgzT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WZnqxw4uM_7h"
      },
      "source": [
        "# Extra Resources\n",
        "\n",
        "You may want to check out some of my other projects as well to get more insight into how the different parts of VQGAN+CLIP work together to generate an image:\n",
        "\n",
        "1. [Artist Studies](https://remidurant.com/artists) - Find good artists for your prompts\n",
        "2. [How CLIP \"sees\"](https://twitter.com/remi_durant/status/1460607677801897990?s=20)\n",
        "3. Art Styles and Movements \n",
        " - [VQGAN Imagenet16k + ViT-B/32](https://imgur.com/gallery/BZzXLHY)\n",
        " - [VQGAN Imagenet16k + ViT-B/16](https://imgur.com/gallery/w14XZFd)\n",
        " - [VQGAN Imagenet16k + RN50x16](https://imgur.com/gallery/Kd0WYfo)\n",
        " - [VQGAN Imagenet16k + RN50x4](https://imgur.com/gallery/PNd7zYp)\n",
        "\n",
        "\n",
        "There is also this great prompt exploration from @kingdomakrillic which showcases a lot of the words you can add to your prompt to push CLIP towards certain styles:\n",
        "- [CLIP + VQGAN Keyword Comparison](https://imgur.com/a/SnSIQRu)\n"
      ]
    }
  ]
}