AI_Art_Machine.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "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/JosephBless/Deep-Learning-Colab/blob/main/AI_Art_Machine.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CppIQlPhhwhs"
      },
      "source": [
        "# ART MACHINE: Put in text, get AI art.\n",
        "\n",
        "This notebook is by Hillel Wayne ([website](https://www.hillelwayne.com/), [twitter](https://twitter.com/hillelogram)). It's based on [this notebook](https://colab.research.google.com/drive/15UwYDsnNeldJFHJ9NdgYBYeo6xPmSelP) by Katherine Crowson ([github](https://github.com/crowsonkb), [twitter](https://twitter.com/RiversHaveWings)), simplified to make it more accessible to nonprogrammers. The original technique was discovered by https://twitter.com/advadnoun.\n",
        "\n",
        "If you want to learn more about AI-generated art, [this](https://ml.berkeley.edu/blog/posts/clip-art/) is a great introduction. It lists other generators at the end.\n",
        "\n",
        "Want to share this? Short url is https://is.gd/artmachine\n",
        "\n",
        "**Licensing:** Anything you make with the art machine is yours to do whatever you want with it. (Disclaimer: I Am Not A Lawyer)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8OMcm6JwR2Zz"
      },
      "source": [
        "# @title Licensed under the MIT License\n",
        "\n",
        "# Copyright (c) 2021 Katherine Crowson\n",
        "\n",
        "# Permission is hereby granted, free of charge, to any person obtaining a copy\n",
        "# of this software and associated documentation files (the \"Software\"), to deal\n",
        "# in the Software without restriction, including without limitation the rights\n",
        "# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
        "# copies of the Software, and to permit persons to whom the Software is\n",
        "# furnished to do so, subject to the following conditions:\n",
        "\n",
        "# The above copyright notice and this permission notice shall be included in\n",
        "# all copies or substantial portions of the Software.\n",
        "\n",
        "# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
        "# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
        "# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
        "# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
        "# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
        "# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\n",
        "# THE SOFTWARE.\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P966kx616GlK"
      },
      "source": [
        "\n",
        "\n",
        "## How to use\n",
        "\n",
        "(You need to be logged into a Google account for this to work)\n",
        "\n",
        "1. Set your prompt [here](#scrollTo=TnMw4FrN6JeB).\n",
        "2. In the menu bar on the top of the webpage, go to \"Runtime\", then \"Run All\". The first run will take a while to start because the notebook will need to download some important files. *Note*: if you're on mobile, you can find the \"run all\" button in the sandwich menu.\n",
        "3. Images will appear [here](#scrollTo=oxRlmwZ9As3X).\n",
        "4. If you want to download the results, go [here](#scrollTo=2JR7TTOIcBr7) when it's done.\n",
        "5. For future runs, use \"Restart and Run All\" instead of \"Run All\". Runs after the first will be MUCH faster to start.\n",
        "\n",
        "#### Prompt Ideas\n",
        "\n",
        "Here are some things that worked for me:\n",
        "\n",
        "* Simple scenes: \"A cabin by the lake\", \"sunset over mountains\"\n",
        "* Things that aren't visual: \"Dubstep\", \"Sad\"\n",
        "* Media with distinct visuals: \"Homestar Runner\", \"Factorio\"\n",
        "* Something poetic: \"old gods and older machines\", \"a doorway into a place without sunlight\"\n",
        "\n",
        "Here are some that other people have made that I really like: [Robot Birthday](https://twitter.com/thristhart/status/1422734774884397056), [the eerie bliss and torture of solitude](https://twitter.com/dynamic_proxy/status/1423034901264310273), [voxel utopia in cyclades greece](https://twitter.com/_TheodoreM_/status/1422727156820545538), and [ghost chasing the life out of a planet](https://twitter.com/SnoPhort/status/1422721217564266497).\n",
        "\n",
        "**Flavors**: There are two \"flavors\" of art generator: \"default\" and \"cumin\". They have different aesthetics and produce different kinds of art for the same prompt. If you don't like what you get from one, try the other!\n",
        "\n",
        "**Keywords**: You can put style keywords onto the end of your prompt to get different results.\n",
        "\n",
        "Some keywords that work well are \"studio ghibli\", \"trending on artstation\", \"rendered in unreal engine\", \"drawn really good\", \"chibi\", and \"watercolor\". Experiment!\n",
        "\n",
        "**Style Images**: Under advanced options is `style_url`. If you put a link to an image here, it will try to generate an image with the same style. For example, if you put [this](https://upload.wikimedia.org/wikipedia/commons/a/a4/Piet_Mondriaan%2C_1930_-_Mondrian_Composition_II_in_Red%2C_Blue%2C_and_Yellow.jpg) url in the bar, the prompt will look more like a Piet Mondrian composition.\n",
        "\n",
        "**Multiple Prompts**: If you separate prompts with a pipe (`|`), the art machine will try them all simultaneously. For example, `bat man` will be similar to the superhero, while `bat | man` will get you something closer to a cross between a bat and a man.\n",
        "\n",
        "   When using multiple prompts, you can \"weight\" them by putting `:n` after the prompt. For example, `bat | man:3` will tell the art machine to focus 3 times as hard on `man` than on `bat`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TkUfzT60ZZ9q"
      },
      "source": [
        "#@title Instance Size { vertical-output: true, display-mode: \"form\" }\n",
        "x = !nvidia-smi\n",
        "if \"MiB / 15\" in x[9]:\n",
        "  print(\"Google gave you a big instance! Image size can go up to 700!\")\n",
        "else:\n",
        "  print(\"Google gave you a small instance. Image size can go up to 400.\")\n",
        "\n",
        "# TODO add a test here"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TnMw4FrN6JeB"
      },
      "source": [
        "#@title Parameters { run: \"auto\", display-mode: \"form\" }\n",
        "#@markdown ### Basic options:\n",
        "#@markdown ↓ The prompt to you want to try.\n",
        "\n",
        "p = \"chariot of fire\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ↓ How big you want the image to be. This can be up to 700, but if it's smaller, your images will appear *much* faster.\n",
        "image_size =   400#@param {type:\"integer\"}\n",
        "\n",
        "#@markdown ↓ If checked, will automatically bundle and download the images when the engine is done.\n",
        "\n",
        "download_after_finishing = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown ### Advanced Options\n",
        "#@markdown *You don't need to change any of these things, but you can if you want.*\n",
        "\n",
        "\n",
        "\n",
        "#@markdown ↓ How many steps you want the generator to run. Best between 400 and 1000 for default weirdness.\n",
        "\n",
        "total_steps = 400 #@param {type:\"integer\"}\n",
        "\n",
        "#@markdown ↓ How many steps you want between images. Use a lower number if you want to make videos.\n",
        "\n",
        "steps_per_image = 50 #@param {type:\"integer\"}\n",
        "\n",
        "#@markdown ↓ Different flavors give you very different art for the same prompt.\n",
        "\n",
        "flavor = 'thyme' #@param [\"cumin\", \"rosewater\", \"oregano\", \"thyme\"]\n",
        "\n",
        "#@markdown ↓ If nonempty, the generator will try to mimick the style of the image. Should be a url. Leave empty if you just want a prompt.\n",
        "\n",
        "style_url = \"\" #@param {type:\"string\"}\n",
        "#@markdown ↓ How fast and sloppy the AI should be. 2 is a good balance.\n",
        "\n",
        "weirdness = 2 #@param {type:\"slider\", min:1, max:11, step:1}\n",
        "#@markdown ↓ Changing this number will give you slightly different results for the same prompt.\n",
        "seed = 1 #@param {type:\"integer\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZdlpRFL8UAlW"
      },
      "source": [
        "#@title low-level config you don't need to worry about { display-mode: \"form\" }\n",
        "#@markdown You can click \"show code\" if you want to mess with it, but you shouldn't need to.\n",
        "\n",
        "# TODO why is this even argparse this makes no sense!\n",
        "import argparse\n",
        "args = argparse.Namespace(\n",
        "    prompts=p.split(\"|\"),\n",
        "    image_prompts=[],\n",
        "    noise_prompt_seeds=[],\n",
        "    noise_prompt_weights=[],\n",
        "    tv_weight=0.1,\n",
        "    step_size=0.05*(weirdness if weirdness != 11 else 22),\n",
        "    weight_decay=0.,\n",
        "    cutn=64,\n",
        "    cut_pow=1.,\n",
        "    display_freq=steps_per_image,\n",
        "    total_iterations=total_steps,\n",
        "    seed=seed,\n",
        ")\n",
        "\n",
        "if style_url:\n",
        "  args.image_prompts = [style_url]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wSfISAhyPmyp"
      },
      "source": [
        "#@title A bunch of setup stuff { display-mode: \"form\" }\n",
        "#@markdown This will run *much* faster on the second time.\n",
        "%%capture\n",
        "%%bash\n",
        "\n",
        "\n",
        "pip install ftfy regex tqdm omegaconf pytorch-lightning imageio-ffmpeg einops kornia transformers taming-transformers git+https://github.com/openai/CLIP.git\n",
        "curl -C - -o vqgan_imagenet_f16_1024.yaml -L 'https://heibox.uni-heidelberg.de/d/8088892a516d4e3baf92/files/?p=%2Fconfigs%2Fmodel.yaml&dl=1'\n",
        "curl -C - -o vqgan_imagenet_f16_1024.ckpt -L 'https://heibox.uni-heidelberg.de/d/8088892a516d4e3baf92/files/?p=%2Fckpts%2Flast.ckpt&dl=1'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EXMSuW2EQWsd"
      },
      "source": [
        "#@title More, different setup { display-mode: \"form\" }\n",
        "import math\n",
        "import io\n",
        "from pathlib import Path\n",
        "import sys\n",
        "import time\n",
        "\n",
        "from IPython import display\n",
        "from omegaconf import OmegaConf\n",
        "from PIL import Image\n",
        "import requests\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, trange\n",
        "import kornia.augmentation as K\n",
        "\n",
        "\n",
        "import warnings\n",
        "warnings.simplefilter(\"ignore\") # Avoid spookin people for Cumin\n",
        "\n",
        "from clip import clip\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",
        "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 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",
        "replace_grad = ReplaceGrad.apply\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",
        "\n",
        "clamp_with_grad = ClampWithGrad.apply\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",
        "        input_normed = F.normalize(input.unsqueeze(1), dim=2)\n",
        "        embed_normed = F.normalize(self.embed.unsqueeze(0), dim=2)\n",
        "        dists = input_normed.sub(embed_normed).norm(dim=2).div(2).arcsin().pow(2).mul(2)\n",
        "        dists = dists * self.weight.sign()\n",
        "        return self.weight.abs() * replace_grad(dists, torch.maximum(dists, self.stop)).mean()\n",
        "\n",
        "\n",
        "def fetch(url_or_path):\n",
        "    if str(url_or_path).startswith('http://') or str(url_or_path).startswith('https://'):\n",
        "        headers={\"user-agent\": \"AIArtMachineBot/0.0 (https://is.gd/aiartmachine; h@hillelwayne.com) generic-library/0.0\"}\n",
        "        r = requests.get(url_or_path, headers=headers)\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(':', 3)\n",
        "        vals = [vals[0] + ':' + vals[1], *vals[2:]]\n",
        "    else:\n",
        "        vals = [prompt]\n",
        "    vals = vals + ['', '1', '-inf'][len(vals):]\n",
        "    return vals[0], float(vals[1]), float(vals[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()\n",
        "\n",
        "\n",
        "class MakeCutoutsDefault(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",
        "    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",
        "        for _ in range(self.cutn):\n",
        "            size = int(torch.rand([])**self.cut_pow * (max_size - min_size) + min_size)\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",
        "        return clamp_with_grad(torch.cat(cutouts, dim=0), 0, 1)\n",
        "\n",
        "class MakeCutoutsCumin(nn.Module):\n",
        "    \"\"\"from https://colab.research.google.com/drive/1ZAus_gn2RhTZWzOWUpPERNC0Q8OhZRTZ\"\"\"\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.augs = nn.Sequential(\n",
        "            K.RandomAffine(degrees=15, translate=0.1, p=0.7, padding_mode='border'),\n",
        "            K.RandomPerspective(0.7,p=0.7),\n",
        "            K.ColorJitter(hue=0.1, saturation=0.1, p=0.7),\n",
        "            K.RandomErasing((.1, .4), (.3, 1/.3), same_on_batch=True, p=0.7),\n",
        "\n",
        ")\n",
        "        self.noise_fac = 0.1\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 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",
        "\n",
        "        for _ in range(self.cutn):\n",
        "            cutout = (self.av_pool(input) + self.max_pool(input))/2\n",
        "            cutouts.append(cutout)\n",
        "        batch = self.augs(torch.cat(cutouts, dim=0))\n",
        "        if self.noise_fac:\n",
        "            facs = batch.new_empty([self.cutn, 1, 1, 1]).uniform_(0, self.noise_fac)\n",
        "            batch = batch + facs * torch.randn_like(batch)\n",
        "        return batch\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",
        "    else:\n",
        "        raise ValueError(f'unknown model type: {config.model.target}')\n",
        "    del model.encoder, model.loss\n",
        "    return model\n",
        "\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",
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "\n",
        "vqgan_config='vqgan_imagenet_f16_1024.yaml'\n",
        "vqgan_checkpoint='vqgan_imagenet_f16_1024.ckpt'\n",
        "model = load_vqgan_model(vqgan_config, vqgan_checkpoint).to(device)\n",
        "perceptor = clip.load(\"ViT-B/32\", jit=False)[0].eval().requires_grad_(False).to(device)\n",
        "\n",
        "cut_size = perceptor.visual.input_resolution\n",
        "e_dim = model.quantize.e_dim\n",
        "f = 2**(model.decoder.num_resolutions - 1)\n",
        "flavordict = {\n",
        "    \"default\": MakeCutoutsDefault,\n",
        "    \"cumin\": MakeCutoutsCumin,\n",
        "    \"rosewater\": MakeCutoutsDefault,\n",
        "    \"oregano\": MakeCutoutsDefault,\n",
        "    \"thyme\": MakeCutoutsCumin,\n",
        "\n",
        "}\n",
        "\n",
        "def uses_zq():\n",
        "  return flavor in {\"oregano\", \"thyme\"}\n",
        "\n",
        "make_cutouts = flavordict[flavor](cut_size, args.cutn, cut_pow=args.cut_pow)\n",
        "n_toks = model.quantize.n_e\n",
        "toksX, toksY = image_size // f, image_size // f\n",
        "sideX, sideY = toksX * f, toksY * f\n",
        "\n",
        "if args.seed is not None:\n",
        "    torch.manual_seed(args.seed)\n",
        "\n",
        "\"\"\"\n",
        "Oregano and Thyme are based on z+quantize, from\n",
        "https://colab.research.google.com/drive/1wkF67ThUz37T2_oPIuSwuO4e_-0vjaLs\n",
        "They use a completely different means of guiding VQGAN vs codebook sampling.\n",
        "Best way to handle this rn is to branch on the logic. A smarter, less lazy\n",
        "person would have instead made Task objects with injectable Flavor objects, but I am\n",
        "neither smart nor less lazy.\n",
        "\"\"\"\n",
        "\n",
        "if uses_zq():\n",
        "  oh = F.one_hot(torch.randint(n_toks, [toksX * toksY], device=device), n_toks).float()\n",
        "  z = oh @ model.quantize.embedding.weight\n",
        "  z = z.view([-1, toksY, toksX, e_dim]).permute(0, 3, 1, 2)\n",
        "  z = torch.rand_like(z)*2\n",
        "  z.requires_grad_(True) #Does this slow down basic operations?\n",
        "\n",
        "  opt = optim.AdamW([z], lr=args.step_size, weight_decay=args.weight_decay)\n",
        "else:\n",
        "  logits = torch.randn([toksY * toksX, n_toks], device=device, requires_grad=True)\n",
        "\n",
        "  opt = optim.AdamW([logits], lr=args.step_size, weight_decay=args.weight_decay)\n",
        "\n",
        "normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],\n",
        "                                 std=[0.26862954, 0.26130258, 0.27577711])\n",
        "\n",
        "pMs = []\n",
        "\n",
        "from re import sub\n",
        "\n",
        "# I hate using underscores for names, but this way is better for people using utf-8\n",
        "p_filename = sub(\"\\W\", \"\", p.lower().replace(\" \",\"_\"))\n",
        "\n",
        "### make the output path for dumping the images\n",
        "img_path = Path(\"img\") / \"grouped\" / p_filename\n",
        "all_path = Path(\"img/all\")\n",
        "img_path.mkdir(parents = True, exist_ok = True)\n",
        "all_path.mkdir(parents = True, exist_ok = True)\n",
        "\n",
        "for prompt in args.prompts:\n",
        "    txt, weight, stop = parse_prompt(prompt)\n",
        "    embed = perceptor.encode_text(clip.tokenize(txt).to(device)).float()\n",
        "    pMs.append(Prompt(embed, weight, stop).to(device))\n",
        "\n",
        "for prompt in args.image_prompts:\n",
        "    path, weight, stop = parse_prompt(prompt)\n",
        "    img = resize_image(Image.open(fetch(path)).convert('RGB'), (sideX, sideY))\n",
        "    batch = make_cutouts(TF.to_tensor(img)[None].to(device))\n",
        "    embed = perceptor.encode_image(normalize(batch)).float()\n",
        "    pMs.append(Prompt(embed, weight, stop).to(device))\n",
        "\n",
        "for seed, weight in zip(args.noise_prompt_seeds, args.noise_prompt_weights):\n",
        "    gen = torch.Generator().manual_seed(seed)\n",
        "    embed = torch.empty([1, perceptor.visual.output_dim]).normal_(generator=gen)\n",
        "    pMs.append(Prompt(embed, weight).to(device))\n",
        "\n",
        "def z_synth():\n",
        "    z_q = vector_quantize(z.movedim(1, 3), model.quantize.embedding.weight).movedim(3, 1)\n",
        "    return clamp_with_grad(model.decode(z_q).add(1).div(2), 0, 1)\n",
        "\n",
        "def logit_synth(one_hot):\n",
        "    z = one_hot @ model.quantize.embedding.weight\n",
        "    z = z.view([-1, toksY, toksX, e_dim]).permute(0, 3, 1, 2)\n",
        "    return clamp_with_grad(model.decode(z).add(1).div(2), 0, 1)\n",
        "\n",
        "@torch.no_grad()\n",
        "def checkin(i, losses):\n",
        "    tqdm.write(f'iterations: {i}, prompt: {p}')\n",
        "    if uses_zq():\n",
        "      out = z_synth()\n",
        "    else:\n",
        "      one_hot = F.one_hot(logits.argmax(1), n_toks).to(logits.dtype)\n",
        "      out = logit_synth(one_hot)\n",
        "    out_img = TF.to_pil_image(out[0].cpu())\n",
        "    out_img.save('progress.png')\n",
        "    out_img.save(all_path / f\"{p_filename}-{i:0=4}.png\")\n",
        "    out_img.save(img_path / f\"{i:0=4}.png\")\n",
        "\n",
        "    display.display(display.Image('progress.png'))\n",
        "\n",
        "def ascend_txt():\n",
        "    opt.zero_grad(set_to_none=True)\n",
        "    if uses_zq():\n",
        "      out = z_synth()\n",
        "    else:\n",
        "      probs = logits.softmax(1)\n",
        "      one_hot = F.one_hot(probs.multinomial(1)[..., 0], n_toks).to(logits.dtype)\n",
        "      one_hot = replace_grad(one_hot, probs)\n",
        "      out = logit_synth(one_hot)\n",
        "    iii = perceptor.encode_image(normalize(make_cutouts(out))).float()\n",
        "\n",
        "    result = []\n",
        "\n",
        "    if args.tv_weight:\n",
        "        result.append(tv_loss(out) * args.tv_weight / 4)\n",
        "\n",
        "    for prompt in pMs:\n",
        "        result.append(prompt(iii))\n",
        "\n",
        "    return result\n",
        "\n",
        "def train(i):\n",
        "    opt.zero_grad(set_to_none=True)\n",
        "    lossAll = ascend_txt()\n",
        "    if i % args.display_freq == 0:\n",
        "        checkin(i, lossAll)\n",
        "    loss = sum(lossAll)\n",
        "    loss.backward()\n",
        "    opt.step()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EPEZ5OjQSZLh"
      },
      "source": [
        "# Output"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oxRlmwZ9As3X"
      },
      "source": [
        "#@title Images will appear here! { display-mode: \"form\" }\n",
        "\n",
        "i = 1\n",
        "print(\"NOTE: First image will look random. This is normal.\")\n",
        "try:\n",
        "    with trange(args.total_iterations) as pbar:\n",
        "        start = time.perf_counter()\n",
        "        while i < args.total_iterations:\n",
        "            pbar.update()\n",
        "            train(i)\n",
        "            if i == 10:\n",
        "              end = time.perf_counter()\n",
        "              print(f\"It will take about {int((end - start) * args.total_iterations // (60 * 10))} minutes to complete all {args.total_iterations} iterations.\")\n",
        "            i += 1\n",
        "except KeyboardInterrupt:\n",
        "    pass\n",
        "except RuntimeError:\n",
        "    print(\"ERROR! ERROR! ERROR!\")\n",
        "    print(\"The image size you chose was too big!\")\n",
        "else:\n",
        "  print(\"Final image.\")\n",
        "  checkin(i, 0)\n",
        "\n",
        "print(\"All done!\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2JR7TTOIcBr7"
      },
      "source": [
        "#@title Download images { display-mode: \"form\" }\n",
        "#@markdown The button will appear *after* the generator finishes.\n",
        "from google.colab import files\n",
        "import ipywidgets as widgets\n",
        "from shutil import make_archive\n",
        "\n",
        "def download_prompt():\n",
        "  make_archive(p_filename, \"zip\", base_dir = img_path)\n",
        "  files.download(f\"{p_filename}.zip\")\n",
        "\n",
        "if download_after_finishing:\n",
        "  download_prompt()\n",
        "\n",
        "def download_all(_):\n",
        "  make_archive(\"all_images\", \"zip\", base_dir = all_path)\n",
        "  print(\"ZIP file made. If the download prompt doesn't appear, you can download it manually in the 'files' sidebar (left of screen, fourth icon)\")\n",
        "  files.download('all_images.zip')\n",
        "\n",
        "button = widgets.Button(\n",
        "    description='Download All',\n",
        "    disabled=False,\n",
        "    button_style='success', # 'success', 'info', 'warning', 'danger' or ''\n",
        "    tooltip='Click me',\n",
        "    icon='check' # (FontAwesome names without the `fa-` prefix)\n",
        ")\n",
        "button.on_click(download_all)\n",
        "button"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2GgPOygcAmw5"
      },
      "source": [
        "#@title Make a Video! {  run: \"auto\", display-mode: \"form\" }\n",
        "#@markdown The button will appear *after* the generator finishes.\n",
        "# Taken from https://colab.research.google.com/drive/1ZAus_gn2RhTZWzOWUpPERNC0Q8OhZRTZ\n",
        "import numpy as np\n",
        "from subprocess import Popen, PIPE\n",
        "from base64 import b64encode\n",
        "\n",
        "video_length_in_seconds =  5#@param {type:\"integer\"}\n",
        "display.clear_output()\n",
        "\n",
        "def download_video(_):\n",
        "  print(\"If the download prompt doesn't appear, you can download the video manually in the 'files' sidebar (left of screen, fourth icon)\")\n",
        "  files.download(f\"{p_filename}.mp4\")\n",
        "\n",
        "def make_video(_):\n",
        "  min_fps = 10\n",
        "  max_fps = 60\n",
        "\n",
        "  length = video_length_in_seconds #Desired time of the video in seconds\n",
        "\n",
        "  frames = []\n",
        "  print('Generating video...')\n",
        "  for snapshot in sorted(img_path.iterdir()): #\n",
        "      frames += [Image.open(snapshot)]*1 # Make things a wee bit smoother\n",
        "\n",
        "  fps = np.clip(len(frames)/length,min_fps,max_fps)\n",
        "\n",
        "  p = Popen(['ffmpeg', '-y', '-f', 'image2pipe', '-vcodec', 'png', '-r', str(fps), '-i', '-', '-vcodec', 'libx264', '-r', str(fps), '-pix_fmt', 'yuv420p', '-crf', '17', '-preset', 'veryslow', f'{p_filename}.mp4'], stdin=PIPE)\n",
        "  for im in tqdm(frames):\n",
        "      im.save(p.stdin, 'PNG')\n",
        "  p.stdin.close()\n",
        "  p.wait()\n",
        "  print('Attaching video...')\n",
        "  mp4 = open(f'{p_filename}.mp4','rb').read()\n",
        "  data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "  display.display(display.HTML(\"\"\"\n",
        "<video width={} controls>\n",
        "      <source src=\"{}\" type=\"video/mp4\">\n",
        "</video>\n",
        "\"\"\".format(image_size, data_url)))\n",
        "\n",
        "  download_video_button = widgets.Button(\n",
        "    description='Download Video',\n",
        "    disabled=False,\n",
        "    button_style='success', # 'success', 'info', 'warning', 'danger' or ''\n",
        "    tooltip='Click me',\n",
        "    icon='check' # (FontAwesome names without the `fa-` prefix)\n",
        "  )\n",
        "  download_video_button.on_click(download_video)\n",
        "  display.display(download_video_button)\n",
        "\n",
        "\n",
        "video_button = widgets.Button(\n",
        "    description='Make a video',\n",
        "    disabled=False,\n",
        "    button_style='success', # 'success', 'info', 'warning', 'danger' or ''\n",
        "    tooltip='Click me',\n",
        "    icon='check' # (FontAwesome names without the `fa-` prefix)\n",
        ")\n",
        "\n",
        "\n",
        "video_button.on_click(make_video)\n",
        "video_button"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4y7CHnkwsKEy"
      },
      "source": [
        "# Learn More!\n",
        "\n",
        "If this interested you and you want to learn more about AI art, here are some links!\n",
        "\n",
        "### Resources\n",
        "\n",
        "* [Alien Dreams: An Emerging Art Scene](https://ml.berkeley.edu/blog/posts/clip-art/): first \"mainstream\" introduction to AI Art.\n",
        "* [The Illustrated VQGAN](https://ljvmiranda921.github.io/notebook/2021/08/08/clip-vqgan/): The neural network generating images is VQGAN. This is a description of how it works! Fairly technical.\n",
        "\n",
        "### Notebooks\n",
        "\n",
        "There's other notebooks you can try, too! Some of them use different styles, or give you different features to play with.\n",
        "\n",
        "* [Create Realistic AI-Generated Images](https://colab.research.google.com/drive/1wkF67ThUz37T2_oPIuSwuO4e_-0vjaLs): Also designed for beginners, this one lets you pick a starting initial image to work from.\n",
        "\n",
        "* [VQGAN+CLIP](https://colab.research.google.com/drive/1_4Jl0a7WIJeqy5LTjPJfZOwMZopG5C-W): The first notebook to use VQGAN, which is now the most popular AI generator to use. Mostly reimplemented as the herb flavors.\n",
        "* [Semantic Style Transfer](https://colab.research.google.com/drive/17AqhaKLZmmUA27aNSc6fJYMR9uypeIci?usp=sharing): Can do \"style transfers\", where you make one image stylized like another image.\n",
        "* [The Big Sleep](https://colab.research.google.com/drive/1NCceX2mbiKOSlAd_o7IU7nA9UskKN5WR). The very first CLIP-based notebook, starting the whole trend. Uses a different image generator and gives very different pictures.\n",
        "* [CLIP-Guided Diffusion](https://colab.research.google.com/drive/12a_Wrfi2_gwwAuN3VvMTwVMz9TfqctNj): One of the newest innovations in CLIP tech. Uses a different image generator.\n",
        "* [CLIPIT PixelDraw](https://colab.research.google.com/github/dribnet/clipit/blob/master/demos/PixelDrawer.ipynb): AI-generated pixel art!\n",
        "\n",
        "### Communities\n",
        "\n",
        "* [/r/deepdream](https://www.reddit.com/r/deepdream/): A subreddit where people post their AI art. Some really hardcore artists in here producing amazing stuff.\n",
        "* [EleutherAI](https://www.eleuther.ai/): A \"grassroots collective\" of AI researchers. To my understanding there's a lot of AI artists in their discord, but I haven't checked this myself.\n"
      ]
    }
  ]
}