ACMI_Public_API_x_VQGAN+CLIP_(with_pooling).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 2,
  "metadata": {
    "colab": {
      "name": "ACMI Public API x VQGAN+CLIP (with pooling)",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Generate images from ACMI Public API metadata with VQGAN and CLIP (z+quantize method).\n",
        "\n",
        "Modified by the [ACMI Labs](https://labs.acmi.net.au) team as an example use of our pubilc collection API: [api.acmi.net.au](https://api.acmi.net.au)\n",
        "\n",
        "See it on GitHub: https://github.com/ACMILabs/ACMI-API-x-VQGAN-CLIP\n",
        "\n",
        "## Credits\n",
        "\n",
        "Source: https://github.com/nerdyrodent/VQGAN-CLIP\n",
        "\n",
        "Originally made by Katherine Crowson (https://github.com/crowsonkb, https://twitter.com/RiversHaveWings). The original BigGAN+CLIP method was by https://twitter.com/advadnoun.\n",
        " Added some explanations and modifications by Eleiber#8347, pooling trick by Crimeacs#8222 (https://twitter.com/EarthML1) and the GUI was made with the help of Abulafia#3734.\n"
      ],
      "metadata": {
        "id": "CppIQlPhhwhs"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "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"
      ],
      "outputs": [],
      "metadata": {
        "cellView": "form",
        "id": "-nnPDoDFCp6n"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "!nvidia-smi"
      ],
      "outputs": [],
      "metadata": {
        "id": "TkUfzT60ZZ9q"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "!git clone https://github.com/openai/CLIP\n",
        "# !pip install taming-transformers\n",
        "!git clone https://github.com/CompVis/taming-transformers.git\n",
        "!pip install ftfy regex tqdm omegaconf pytorch-lightning\n",
        "!pip install kornia\n",
        "!pip install imageio-ffmpeg   \n",
        "!pip install einops          \n",
        "!mkdir steps"
      ],
      "outputs": [],
      "metadata": {
        "id": "VA1PHoJrRiK9"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "By default, the notebook downloads the 1024 and 16384 models from ImageNet. There are others like COCO-Stuff, WikiArt or S-FLCKR, which are heavy, and if you are not going to use them it would be useless to download them, so if you want to use them, simply remove the numerals at the beginning of the lines depending on the model you want (the model name is at the end of the lines)."
      ],
      "metadata": {
        "id": "nTg77tNuF7Og"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "#@title Selection of models to download\n",
        "#@markdown By default, the notebook downloads the 1024 and 16384 models from ImageNet. There are others like COCO-Stuff, WikiArt 1024, WikiArt 16384, FacesHQ or S-FLCKR, which are heavy, and if you are not going to use them it would be pointless to download them, so if you want to use them, simply select the models to download.\n",
        "\n",
        "imagenet_1024 = True #@param {type:\"boolean\"}\n",
        "imagenet_16384 = True #@param {type:\"boolean\"}\n",
        "coco = False #@param {type:\"boolean\"}\n",
        "faceshq = False #@param {type:\"boolean\"}\n",
        "wikiart_1024 = False #@param {type:\"boolean\"}\n",
        "wikiart_16384 = False #@param {type:\"boolean\"}\n",
        "sflckr = False #@param {type:\"boolean\"}\n",
        "openimages_8192 = False #@param {type:\"boolean\"}\n",
        "\n",
        "if imagenet_1024:\n",
        "  !curl -L -o vqgan_imagenet_f16_1024.yaml -C - 'http://mirror.io.community/blob/vqgan/vqgan_imagenet_f16_1024.yaml' #ImageNet 1024\n",
        "  !curl -L -o vqgan_imagenet_f16_1024.ckpt -C - 'http://mirror.io.community/blob/vqgan/vqgan_imagenet_f16_1024.ckpt'  #ImageNet 1024\n",
        "if imagenet_16384:\n",
        "  !curl -L -o vqgan_imagenet_f16_16384.yaml -C - 'https://heibox.uni-heidelberg.de/d/a7530b09fed84f80a887/files/?p=%2Fconfigs%2Fmodel.yaml&dl=1' #ImageNet 16384\n",
        "  !curl -L -o vqgan_imagenet_f16_16384.ckpt -C - 'https://heibox.uni-heidelberg.de/d/a7530b09fed84f80a887/files/?p=%2Fckpts%2Flast.ckpt&dl=1' #ImageNet 16384\n",
        "\n",
        "if openimages_8192:\n",
        "  !curl -L -o vqgan_openimages_f16_8192.yaml -C - 'https://heibox.uni-heidelberg.de/d/2e5662443a6b4307b470/files/?p=%2Fconfigs%2Fmodel.yaml&dl=1' #ImageNet 16384\n",
        "  !curl -L -o vqgan_openimages_f16_8192.ckpt -C - 'https://heibox.uni-heidelberg.de/d/2e5662443a6b4307b470/files/?p=%2Fckpts%2Flast.ckpt&dl=1' #ImageNet 16384\n",
        "\n",
        "if coco:\n",
        "  !curl -L -o coco.yaml -C - 'https://dl.nmkd.de/ai/clip/coco/coco.yaml' #COCO\n",
        "  !curl -L -o coco.ckpt -C - 'https://dl.nmkd.de/ai/clip/coco/coco.ckpt' #COCO\n",
        "if faceshq:\n",
        "  !curl -L -o faceshq.yaml -C - 'https://drive.google.com/uc?export=download&id=1fHwGx_hnBtC8nsq7hesJvs-Klv-P0gzT' #FacesHQ\n",
        "  !curl -L -o faceshq.ckpt -C - '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' #FacesHQ\n",
        "if wikiart_1024: \n",
        "  !curl -L -o wikiart_1024.yaml -C - 'http://mirror.io.community/blob/vqgan/wikiart.yaml' #WikiArt 1024\n",
        "  !curl -L -o wikiart_1024.ckpt -C - 'http://mirror.io.community/blob/vqgan/wikiart.ckpt' #WikiArt 1024\n",
        "if wikiart_16384: \n",
        "  !curl -L -o wikiart_16384.yaml -C - 'http://mirror.io.community/blob/vqgan/wikiart_16384.yaml' #WikiArt 16384\n",
        "  !curl -L -o wikiart_16384.ckpt -C - 'http://mirror.io.community/blob/vqgan/wikiart_16384.ckpt' #WikiArt 16384\n",
        "if sflckr:\n",
        "  !curl -L -o sflckr.yaml -C - 'https://heibox.uni-heidelberg.de/d/73487ab6e5314cb5adba/files/?p=%2Fconfigs%2F2020-11-09T13-31-51-project.yaml&dl=1' #S-FLCKR\n",
        "  !curl -L -o sflckr.ckpt -C - 'https://heibox.uni-heidelberg.de/d/73487ab6e5314cb5adba/files/?p=%2Fcheckpoints%2Flast.ckpt&dl=1' #S-FLCKR"
      ],
      "outputs": [],
      "metadata": {
        "id": "FhhdWrSxQhwg"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "# @title Load libraries and variables\n",
        "\n",
        "import argparse\n",
        "import math\n",
        "from pathlib import Path\n",
        "import sys\n",
        "\n",
        "sys.path.insert(1, './taming-transformers')\n",
        "from IPython import display\n",
        "from base64 import b64encode\n",
        "from omegaconf import OmegaConf\n",
        "from PIL import Image\n",
        "from taming.models import cond_transformer, vqgan\n",
        "import taming.modules \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",
        "\n",
        "from CLIP import clip\n",
        "import kornia.augmentation as K\n",
        "import numpy as np\n",
        "import imageio\n",
        "from PIL import ImageFile, Image\n",
        "ImageFile.LOAD_TRUNCATED_IMAGES = True\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",
        "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",
        "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",
        "        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 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",
        "\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.augs = nn.Sequential(\n",
        "            # K.RandomHorizontalFlip(p=0.5),\n",
        "            # K.RandomVerticalFlip(p=0.5),\n",
        "            # K.RandomSolarize(0.01, 0.01, p=0.7),\n",
        "            # K.RandomSharpness(0.3,p=0.4),\n",
        "            # K.RandomResizedCrop(size=(self.cut_size,self.cut_size), scale=(0.1,1),  ratio=(0.75,1.333), cropping_mode='resample', p=0.5),\n",
        "            # K.RandomCrop(size=(self.cut_size,self.cut_size), p=0.5),\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",
        "\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",
        "\n",
        "            # cutout = transforms.Resize(size=(self.cut_size, self.cut_size))(input)\n",
        "            \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.vqgan.GumbelVQ':\n",
        "        model = vqgan.GumbelVQ(**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.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",
        "\n"
      ],
      "outputs": [],
      "metadata": {
        "id": "EXMSuW2EQWsd"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# ACMI Public API\n",
        "Let's get a random object from the ACMI Public API and use its metadata as the input text, or alternatively if you'd like to specify the ID of the object you'd like to use, fill in the `work_id` field in the parameters below."
      ],
      "metadata": {
        "id": "iNm1EBnXJUP6"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "#@title ACMI API\n",
        "#@markdown If you'd like to specify an individual ACMI Work ID to use\n",
        "#@markdown add it below, else we'll choose a random Work.\n",
        "\n",
        "#@markdown Hint: you can find the ACMI Work ID in the URL from a page on our website.\n",
        "\n",
        "#@markdown e.g. Mad Max is ID `113980`\n",
        "#@markdown https://www.acmi.net.au/works/113980--mad-max/\n",
        "\n",
        "import random\n",
        "import requests\n",
        "\n",
        "from bs4 import BeautifulSoup\n",
        "\n",
        "works = []\n",
        "work_id = '' #@param {type:\"string\"}\n",
        "\n",
        "if work_id:\n",
        "  # Get a single work if the ID is set\n",
        "  work = requests.get(f'https://api.acmi.net.au/works/{work_id}/').json()\n",
        "else:\n",
        "  # Let's get 50 works (5 pages of results)\n",
        "  for page in range(1, 6):\n",
        "    print(f'Getting page {page} of ACMI Works...')\n",
        "    response = requests.get(f'https://api.acmi.net.au/works/?page={page + 1}').json()\n",
        "    works += response.get('results', [])\n",
        "\n",
        "  # Select a random work\n",
        "  print(f'Selecting a random work from {len(works)} ACMI works...')\n",
        "  work = random.choice(works)\n",
        "\n",
        "\n",
        "def clean(input_string):\n",
        "  \"\"\"\n",
        "  Strip all problematic characters and return a clean string.\n",
        "  \"\"\"\n",
        "  # Convert html to plain text\n",
        "  soup = BeautifulSoup(input_string)\n",
        "  output_string = soup.get_text()\n",
        "\n",
        "  # Remove extra spaces\n",
        "  output_string = output_string.strip()\n",
        "\n",
        "  # Remove troublesome characters for this Notebook\n",
        "  bad_characters = [\n",
        "    ':',\n",
        "    '|',\n",
        "    '\\n',\n",
        "    '[',\n",
        "    ']',\n",
        "  ]\n",
        "  for character in bad_characters:\n",
        "    output_string = output_string.replace(character, '')\n",
        "\n",
        "  return output_string\n",
        "\n",
        "\n",
        "# Use the title and brief description fields from our work\n",
        "# as input for this Notebook\n",
        "work_metadata_strings = []\n",
        "if work['title']:\n",
        "  work_metadata_strings.append(\n",
        "    clean(work['title'])\n",
        "  )\n",
        "if work['brief_description']:\n",
        "  clean_brief_description = clean(work['brief_description'])\n",
        "  # Split it into groups of seven words\n",
        "  space_separated_list = clean_brief_description.split()\n",
        "  NUMBER_OF_WORDS = 7\n",
        "  seven_word_groups = [\n",
        "    ' '.join(space_separated_list[x:x+NUMBER_OF_WORDS]) for x in range(\n",
        "      0,\n",
        "      len(space_separated_list),\n",
        "      NUMBER_OF_WORDS,\n",
        "    )\n",
        "  ]\n",
        "  work_metadata_strings += seven_word_groups\n",
        "\n",
        "# Join our metadata together with the '|' symbol for this Notebook\n",
        "input_string_from_acmi_api = '|'.join(work_metadata_strings)\n",
        "\n",
        "print(f'ACMI API input string from https://api.acmi.net.au/works/{work[\"id\"]}/\\n{input_string_from_acmi_api}')"
      ],
      "outputs": [],
      "metadata": {
        "id": "pqbvAiRzK3dJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Settings for this run:\n",
        "Mainly what you will have to modify will be `texts:`, there you can place the text or texts you want to generate (separated with `|`). It is a list because you can put more than one text, and so the AI ​​tries to 'mix' the images, giving the same priority to both texts.\n",
        "\n",
        "To use an initial image to the model, you just have to upload a file to the Colab environment (in the section on the left), and then modify `init_image:` putting the exact name of the file. Example: `sample.png`\n",
        "\n",
        "You can also modify the model by changing the lines that say `model:`. Currently ImageNet 1024, ImageNet 16384, WikiArt 1024, WikiArt 16384, S-FLCKR and COCO-Stuff are available. To activate them you have to have downloaded them first, and then you can simply select it.\n",
        "\n",
        "You can also use `target_images`, which is basically putting one or more images on it that the AI ​​will take as a \"target\", fulfilling the same function as putting text on it. To put more than one you have to use `|` as a separator."
      ],
      "metadata": {
        "id": "1tthw0YaispD"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "#@title Parameters\n",
        "texts = input_string_from_acmi_api #@param {type:\"string\"}\n",
        "width =  600#@param {type:\"number\"}\n",
        "height = 600#@param {type:\"number\"}\n",
        "model = \"vqgan_imagenet_f16_16384\" #@param [\"vqgan_imagenet_f16_16384\", \"vqgan_imagenet_f16_1024\", \"vqgan_openimages_f16_8192\", \"wikiart_1024\", \"wikiart_16384\", \"coco\", \"faceshq\", \"sflckr\"]\n",
        "images_interval =  50#@param {type:\"number\"}\n",
        "init_image = \"\"#@param {type:\"string\"}\n",
        "target_images = \"\"#@param {type:\"string\"}\n",
        "seed = 42#@param {type:\"number\"}\n",
        "max_iterations = 200#@param {type:\"number\"}\n",
        "\n",
        "model_names={\"vqgan_imagenet_f16_16384\": 'ImageNet 16384',\"vqgan_imagenet_f16_1024\":\"ImageNet 1024\", 'vqgan_openimages_f16_8192':'OpenImages 8912',\n",
        "                 \"wikiart_1024\":\"WikiArt 1024\", \"wikiart_16384\":\"WikiArt 16384\", \"coco\":\"COCO-Stuff\", \"faceshq\":\"FacesHQ\", \"sflckr\":\"S-FLCKR\"}\n",
        "name_model = model_names[model]     \n",
        "\n",
        "if seed == -1:\n",
        "    seed = None\n",
        "if init_image == \"None\":\n",
        "    init_image = None\n",
        "if target_images == \"None\" or not target_images:\n",
        "    target_images = []\n",
        "else:\n",
        "    target_images = target_images.split(\"|\")\n",
        "    target_images = [image.strip() for image in target_images]\n",
        "\n",
        "texts = [phrase.strip() for phrase in texts.split(\"|\")]\n",
        "if texts == ['']:\n",
        "    texts = []\n",
        "\n",
        "\n",
        "args = argparse.Namespace(\n",
        "    prompts=texts,\n",
        "    image_prompts=target_images,\n",
        "    noise_prompt_seeds=[],\n",
        "    noise_prompt_weights=[],\n",
        "    size=[width, height],\n",
        "    init_image=init_image,\n",
        "    init_weight=0.,\n",
        "    clip_model='ViT-B/32',\n",
        "    vqgan_config=f'{model}.yaml',\n",
        "    vqgan_checkpoint=f'{model}.ckpt',\n",
        "    step_size=0.1,\n",
        "    cutn=32,\n",
        "    cut_pow=1.,\n",
        "    display_freq=images_interval,\n",
        "    seed=seed,\n",
        ")"
      ],
      "outputs": [],
      "metadata": {
        "id": "ZdlpRFL8UAlW"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "#@title Actually do the run...\n",
        "from urllib.request import urlopen\n",
        "\n",
        "from PIL.PngImagePlugin import PngImageFile, PngInfo\n",
        "\n",
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "print('Using device:', device)\n",
        "if texts:\n",
        "    print('Using texts:', texts)\n",
        "if target_images:\n",
        "    print('Using image prompts:', target_images)\n",
        "if args.seed is None:\n",
        "    seed = torch.seed()\n",
        "else:\n",
        "    seed = args.seed\n",
        "torch.manual_seed(seed)\n",
        "print('Using seed:', seed)\n",
        "\n",
        "model = load_vqgan_model(args.vqgan_config, args.vqgan_checkpoint).to(device)\n",
        "perceptor = clip.load(args.clip_model, jit=False)[0].eval().requires_grad_(False).to(device)\n",
        "# clock=deepcopy(perceptor.visual.positional_embedding.data)\n",
        "# perceptor.visual.positional_embedding.data = clock/clock.max()\n",
        "# perceptor.visual.positional_embedding.data=clamp_with_grad(clock,0,1)\n",
        "\n",
        "cut_size = perceptor.visual.input_resolution\n",
        "\n",
        "f = 2**(model.decoder.num_resolutions - 1)\n",
        "make_cutouts = MakeCutouts(cut_size, args.cutn, cut_pow=args.cut_pow)\n",
        "\n",
        "toksX, toksY = args.size[0] // f, args.size[1] // f\n",
        "sideX, sideY = toksX * f, toksY * f\n",
        "\n",
        "if args.vqgan_checkpoint == 'vqgan_openimages_f16_8192.ckpt':\n",
        "    e_dim = 256\n",
        "    n_toks = model.quantize.n_embed\n",
        "    z_min = model.quantize.embed.weight.min(dim=0).values[None, :, None, None]\n",
        "    z_max = model.quantize.embed.weight.max(dim=0).values[None, :, None, None]\n",
        "else:\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",
        "# 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",
        "# normalize_imagenet = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
        "#                                            std=[0.229, 0.224, 0.225])\n",
        "\n",
        "if args.init_image:\n",
        "    if 'http' in args.init_image:\n",
        "      img = Image.open(urlopen(args.init_image))\n",
        "    else:\n",
        "      img = Image.open(args.init_image)\n",
        "    pil_image = img.convert('RGB')\n",
        "    pil_image = pil_image.resize((sideX, sideY), Image.LANCZOS)\n",
        "    pil_tensor = TF.to_tensor(pil_image)\n",
        "    z, *_ = model.encode(pil_tensor.to(device).unsqueeze(0) * 2 - 1)\n",
        "else:\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",
        "    if args.vqgan_checkpoint == 'vqgan_openimages_f16_8192.ckpt':\n",
        "        z = one_hot @ model.quantize.embed.weight\n",
        "    else:\n",
        "        z = one_hot @ 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_orig = z.clone()\n",
        "z.requires_grad_(True)\n",
        "opt = optim.Adam([z], lr=args.step_size)\n",
        "\n",
        "normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],\n",
        "                                  std=[0.26862954, 0.26130258, 0.27577711])\n",
        "\n",
        "\n",
        "\n",
        "pMs = []\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 = Image.open(path)\n",
        "    pil_image = img.convert('RGB')\n",
        "    img = resize_image(pil_image, (sideX, sideY))\n",
        "    batch = make_cutouts(TF.to_tensor(img).unsqueeze(0).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 synth(z):\n",
        "    if args.vqgan_checkpoint == 'vqgan_openimages_f16_8192.ckpt':\n",
        "        z_q = vector_quantize(z.movedim(1, 3), model.quantize.embed.weight).movedim(3, 1)\n",
        "    else:\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",
        "@torch.no_grad()\n",
        "def checkin(i, 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",
        "    out = synth(z)\n",
        "    TF.to_pil_image(out[0].cpu()).save('progress.png')\n",
        "    display.display(display.Image('progress.png'))\n",
        "\n",
        "def ascend_txt():\n",
        "    global i\n",
        "    out = synth(z)\n",
        "    iii = perceptor.encode_image(normalize(make_cutouts(out))).float()\n",
        "    \n",
        "    result = []\n",
        "\n",
        "    if args.init_weight:\n",
        "        # result.append(F.mse_loss(z, z_orig) * args.init_weight / 2)\n",
        "        result.append(F.mse_loss(z, torch.zeros_like(z_orig)) * ((1/torch.tensor(i*2 + 1))*args.init_weight) / 2)\n",
        "    for prompt in pMs:\n",
        "        result.append(prompt(iii))\n",
        "    img = np.array(out.mul(255).clamp(0, 255)[0].cpu().detach().numpy().astype(np.uint8))[:,:,:]\n",
        "    img = np.transpose(img, (1, 2, 0))\n",
        "    image_path = f'./steps/{str(i)}.png'\n",
        "    imageio.imwrite(image_path, np.array(img))\n",
        "\n",
        "    # Write image metadata\n",
        "    image = PngImageFile(image_path)\n",
        "    metadata = PngInfo()\n",
        "    metadata.add_text('Author', 'ACMI - Your museum of screen culture')\n",
        "    metadata.add_text('Description', f'Image generated using VQGAN+CLIP from metadata: {input_string_from_acmi_api}')\n",
        "    metadata.add_text('Software', 'https://github.com/ACMILabs/ACMI-API-x-VQGAN-CLIP')\n",
        "    metadata.add_text('Source', f'https://www.acmi.net.au/works/{work[\"id\"]}--{work[\"slug\"]}/')\n",
        "    image.save(image_path, pnginfo=metadata)\n",
        "\n",
        "    return result\n",
        "\n",
        "def train(i):\n",
        "    opt.zero_grad()\n",
        "    lossAll = ascend_txt()\n",
        "    if i % args.display_freq == 0:\n",
        "        checkin(i, lossAll)\n",
        "       \n",
        "    loss = sum(lossAll)\n",
        "    loss.backward()\n",
        "    opt.step()\n",
        "    with torch.no_grad():\n",
        "        z.copy_(z.maximum(z_min).minimum(z_max))\n",
        "\n",
        "i = 0\n",
        "try:\n",
        "    with tqdm() as pbar:\n",
        "        while True:\n",
        "            train(i)\n",
        "            if i == max_iterations:\n",
        "                break\n",
        "            i += 1\n",
        "            pbar.update()\n",
        "except KeyboardInterrupt:\n",
        "    pass"
      ],
      "outputs": [],
      "metadata": {
        "id": "JX56bq4rEKIp",
        "cellView": "form"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "#@title Generate a video with the result\n",
        "\n",
        "init_frame = 1 #This is the frame where the video will start\n",
        "last_frame = i #You can change i to the number of the last frame you want to generate. It will raise an error if that number of frames does not exist.\n",
        "\n",
        "min_fps = 10\n",
        "max_fps = 60\n",
        "\n",
        "total_frames = last_frame-init_frame\n",
        "\n",
        "length = 15 #Desired time of the video in seconds\n",
        "\n",
        "frames = []\n",
        "tqdm.write('Generating video...')\n",
        "for i in range(init_frame,last_frame): #\n",
        "    frames.append(Image.open(\"./steps/\"+ str(i) +'.png'))\n",
        "\n",
        "#fps = last_frame/10\n",
        "fps = np.clip(total_frames/length,min_fps,max_fps)\n",
        "\n",
        "from subprocess import Popen, PIPE\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', 'video.mp4'], stdin=PIPE)\n",
        "for im in tqdm(frames):\n",
        "    im.save(p.stdin, 'PNG')\n",
        "p.stdin.close()\n",
        "p.wait()\n",
        "mp4 = open('video.mp4','rb').read()\n",
        "data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "display.HTML(\"\"\"\n",
        "<video width=400 controls>\n",
        "      <source src=\"%s\" type=\"video/mp4\">\n",
        "</video>\n",
        "\"\"\" % data_url)"
      ],
      "outputs": [],
      "metadata": {
        "id": "gmK0k5zQeT5u"
      }
    }
  ]
}