S2_VQGAN+CLIP_Classic.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "S2 VQGAN+CLIP Classic.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "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/benjismith/somewhere-ml/blob/main/S2_VQGAN%2BCLIP_Classic.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": [
        "# Generate images from text phrases with VQGAN and CLIP (z+quantize method with augmentations)\n",
        "\n",
        "[How to use VQGAN+CLIP](https://docs.google.com/document/d/1Lu7XPRKlNhBQjcKr8k8qRzUzbBW7kzxb5Vu72GMRn2E/edit)\n",
        "\n",
        "The original idea behind CLIP came from this article:\n",
        "\n",
        "https://medium.com/@blaisea/physiognomys-new-clothes-f2d4b59fdd6a\n",
        "\n",
        "This notebook has been translated by [twitter.com/somewheresy](https://twitter.com/somewheresy) to English and has some improvements added like Google Drive support and some basic project management. Feel free to contact me with any revisions/suggestions.\n",
        "\n",
        "The last version of this notebook was authored by Katherine Crowson (https://github.com/crowsonkb, https://twitter.com/RiversHaveWings). The original BigGAN + CLIP method was demonstrated by https://twitter.com/advadnoun. Translated (to Spanish) and added explanations, and modifications by Eleiber # 8347, and the friendly interface was made thanks to Abulafia # 3734.\n",
        "\n",
        "For a detailed tutorial on how to use it, I recommend [visiting this article](https://tuscriaturas.miraheze.org/w/index.php?title=Ayuda:Generar_imágenes_con_VQGAN%2BCLIP/English), made by Jakeukalane # 2767 and Avengium (Ángel) # 3715.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VA1PHoJrRiK9",
        "cellView": "form"
      },
      "source": [
        "# @title Licensed under the MIT License\n",
        "\n",
        "# Copyright (c) 2021 Katherine Crowson\n",
        "\n",
        "# Permission is hereby granted, free of charge, to any person obtaining a copy\n",
        "# of this software and associated documentation files (the \"Software\"), to deal\n",
        "# in the Software without restriction, including without limitation the rights\n",
        "# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
        "# copies of the Software, and to permit persons to whom the Software is\n",
        "# furnished to do so, subject to the following conditions:\n",
        "\n",
        "# The above copyright notice and this permission notice shall be included in\n",
        "# all copies or substantial portions of the Software.\n",
        "\n",
        "# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
        "# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
        "# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
        "# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
        "# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
        "# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\n",
        "# THE SOFTWARE."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KqgIZS75GU-O",
        "cellView": "form"
      },
      "source": [
        "#@title Google Drive Integration (optional)\n",
        "#@markdown To connect Google Drive, set `root_path` to the relative drive folder path you want outputs to be saved to if you already made a directory, then execute this cell. Leaving the field blank or just not running this will have outputs save to the runtime temp storage.\n",
        "import os\n",
        "root_path = \"\" #@param {type: \"string\"}\n",
        "abs_root_path = \"/content\"\n",
        "if len(root_path) > 0:\n",
        "    abs_root_path = abs_root_path + \"/drive/MyDrive/\" + root_path\n",
        "\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "def ensureProperRootPath():\n",
        "    if len(abs_root_path) > 0:\n",
        "        os.chdir(abs_root_path) # Changes directory to absolute root path\n",
        "        print(\"Root path check: \")\n",
        "        !pwd\n",
        "\n",
        "ensureProperRootPath()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pY2XJV_eGXUh",
        "cellView": "form"
      },
      "source": [
        "#@title Make a new folder & set root path to that folder (optional)\n",
        "#@markdown Saves a step if you don't have a folder in your Google Drive for this. Makes one, sets the root_path to that new folder. You can name it whatever you'd like:\n",
        "\n",
        "folder_name = \"FolderName\" #@param {type: \"string\"}\n",
        "abs_root_path = \"/content\"\n",
        "if len(folder_name) > 0:\n",
        "    path_tmp = abs_root_path + \"/drive/MyDrive/\" + folder_name\n",
        "    if not os.path.exists(path_tmp):\n",
        "        os.mkdir(path_tmp)\n",
        "    abs_root_path = path_tmp\n",
        "\n",
        "print(\"Created folder & set root path to: \" + abs_root_path)\n",
        "\n",
        "#@markdown Make & assign path to a project subfolder (optional)\n",
        "\n",
        "project_name = \"ProjectName\" #@param {type: \"string\"}\n",
        "if len(project_name) > 0:\n",
        "      path_tmp = abs_root_path + \"/\" + project_name\n",
        "      if not os.path.exists(path_tmp):\n",
        "          os.mkdir(path_tmp)\n",
        "      abs_root_path = path_tmp\n",
        "print(\"Created project subfolder & set root path to: \" + abs_root_path)\n",
        "\n",
        "ensureProperRootPath()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wSfISAhyPmyp",
        "cellView": "form"
      },
      "source": [
        "# @title Setup, Installing Libraries\n",
        "# @markdown This cell might take some time due to installing several libraries.\n",
        "\n",
        "!nvidia-smi\n",
        "print(\"Downloading CLIP...\")\n",
        "!git clone https://github.com/openai/CLIP                 &> /dev/null\n",
        " \n",
        "print(\"Installing Python Libraries for AI\")\n",
        "!git clone https://github.com/CompVis/taming-transformers &> /dev/null\n",
        "!pip install ftfy regex tqdm omegaconf pytorch-lightning  &> /dev/null\n",
        "!pip install kornia                                       &> /dev/null\n",
        "!pip install einops                                       &> /dev/null\n",
        "print(\"Installing transformers library...\")\n",
        "!pip install transformers                                 &> /dev/null\n",
        "print(\"Installing taming.models...\")   \n",
        "!pip install taming.models                                &> /dev/null\n",
        " \n",
        "print(\"Installing libraries for managing metadata...\")\n",
        "!pip install stegano                                      &> /dev/null\n",
        "!apt install exempi                                       &> /dev/null\n",
        "!pip install python-xmp-toolkit                           &> /dev/null\n",
        "!pip install imgtag                                       &> /dev/null\n",
        "!pip install pillow==7.1.2                                &> /dev/null\n",
        "%reload_ext autoreload\n",
        "%autoreload                  &> /dev/null\n",
        " \n",
        "print(\"Installing ffmpeg for creating videos...\")\n",
        "!pip install imageio-ffmpeg &> /dev/null\n",
        "!mkdir steps\n",
        "print(\"Installation finished.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FhhdWrSxQhwg",
        "cellView": "form"
      },
      "source": [
        "#@title Selection of models to download\n",
        "#@markdown By default, the notebook downloads Model 16384 from ImageNet. There are others such as ImageNet 1024, COCO-Stuff, WikiArt 1024, WikiArt 16384, FacesHQ or S-FLCKR, which are not downloaded by default, since it would be in vain if you are not going to use them, so if you want to use them, simply select the models to download.\n",
        "\n",
        "#@markdown WARNING: \n",
        "#@markdown Not all datasets are licensed for commercial use (i.e. selling your artwork as an NFT).\n",
        "\n",
        "\n",
        "#@markdown Datasets you can use for non-commercial purposes:\n",
        "imagenet_1024 = False #@param {type:\"boolean\"} \n",
        "imagenet_16384 = True #@param {type:\"boolean\"}\n",
        "coco = False #@param {type:\"boolean\"}\n",
        "wikiart_16384 = False #@param {type:\"boolean\"}\n",
        "#@markdown Datasets you can use for commercial purposes:\n",
        "faceshq = False #@param {type:\"boolean\"}\n",
        "sflckr = False #@param {type:\"boolean\"}\n",
        "\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",
        "  !curl -L -o vqgan_imagenet_f16_1024.ckpt -C - 'https://heibox.uni-heidelberg.de/f/140747ba53464f49b476/?dl=1' \n",
        "  !curl -L -o vqgan_imagenet_f16_1024.yaml -C - 'https://heibox.uni-heidelberg.de/f/6ecf2af6c658432c8298/?dl=1' \n",
        "  \n",
        "if imagenet_16384:\n",
        "  # !curl -L -o vqgan_imagenet_f16_16384.yaml -C - 'http://mirror.io.community/blob/vqgan/vqgan_imagenet_f16_16384.yaml' #ImageNet 16384\n",
        "  # !curl -L -o vqgan_imagenet_f16_16384.ckpt -C - 'http://mirror.io.community/blob/vqgan/vqgan_imagenet_f16_16384.ckpt' #ImageNet 16384\n",
        "  !curl -L -o vqgan_imagenet_f16_16384.ckpt -C - 'https://heibox.uni-heidelberg.de/f/867b05fc8c4841768640/?dl=1'\n",
        "  !curl -L -o vqgan_imagenet_f16_16384.yaml -C - 'https://heibox.uni-heidelberg.de/f/274fb24ed38341bfa753/?dl=1'\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",
        "  !curl -L -o wikiart_16384.ckpt -C - 'http://eaidata.bmk.sh/data/Wikiart_16384/wikiart_f16_16384_8145600.ckpt'\n",
        "  !curl -L -o wikiart_16384.yaml -C - 'http://eaidata.bmk.sh/data/Wikiart_16384/wikiart_f16_16384_8145600.yaml'\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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EXMSuW2EQWsd",
        "cellView": "form"
      },
      "source": [
        "# @title Loading libraries and definitions\n",
        " \n",
        "import argparse\n",
        "import math\n",
        "from pathlib import Path\n",
        "import sys\n",
        " \n",
        "sys.path.append('./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 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",
        "from imgtag import ImgTag    # metadata\n",
        "from libxmp import *         # metadata\n",
        "import libxmp                # metadata\n",
        "from stegano import lsb\n",
        "import json\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",
        "        self.augs = nn.Sequential(\n",
        "            K.RandomHorizontalFlip(p=0.5),\n",
        "            # K.RandomSolarize(0.01, 0.01, p=0.7),\n",
        "            K.RandomSharpness(0.3,p=0.4),\n",
        "            K.RandomAffine(degrees=30, translate=0.1, p=0.8, padding_mode='border'),\n",
        "            K.RandomPerspective(0.2,p=0.4),\n",
        "            K.ColorJitter(hue=0.01, saturation=0.01, p=0.7))\n",
        "        self.noise_fac = 0.1\n",
        " \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",
        "        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.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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1tthw0YaispD"
      },
      "source": [
        "## Implementation tools\n",
        "Mainly what you will have to modify will be `texts:`, there you can place the text (s) 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 `initial_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 1024, 16384, WikiArt, 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 using a text input. To put more than one you have to use `|` as a separator."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZdlpRFL8UAlW",
        "cellView": "form"
      },
      "source": [
        "#@title Parameters\n",
        "prompts = \"prompt\" #@param {type:\"string\"}\n",
        "width =  512#@param {type:\"number\"}\n",
        "height =  512#@param {type:\"number\"}\n",
        "model = \"vqgan_imagenet_f16_16384\" #@param [\"vqgan_imagenet_f16_16384\", \"vqgan_imagenet_f16_1024\", \"wikiart_16384\", \"coco\", \"faceshq\", \"sflckr\"]\n",
        "display_frequency =  50#@param {type:\"number\"}\n",
        "initial_image = \"None\"#@param {type:\"string\"}\n",
        "target_images = \"None\"#@param {type:\"string\"}\n",
        "seed = -1#@param {type:\"number\"}\n",
        "max_iterations = -1#@param {type:\"number\"}\n",
        "input_images = \"\"\n",
        "\n",
        "model_names={\"vqgan_imagenet_f16_16384\": 'ImageNet 16384',\"vqgan_imagenet_f16_1024\":\"ImageNet 1024\", \n",
        "                  \"wikiart_16384\":\"WikiArt 16384\", \"coco\":\"COCO-Stuff\", \"faceshq\":\"FacesHQ\", \"sflckr\":\"S-FLCKR\"}\n",
        "model_name = model_names[model]     \n",
        "\n",
        "if seed == -1:\n",
        "    seed = None\n",
        "if initial_image == \"None\":\n",
        "    initial_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",
        "if initial_image or target_images != []:\n",
        "    input_images = True\n",
        "\n",
        "prompts = [frase.strip() for frase in prompts.split(\"|\")]\n",
        "if prompts == ['']:\n",
        "    prompts = []\n",
        "\n",
        "\n",
        "args = argparse.Namespace(\n",
        "    prompts=prompts,\n",
        "    image_prompts=target_images,\n",
        "    noise_prompt_seeds=[],\n",
        "    noise_prompt_weights=[],\n",
        "    size=[width, height],\n",
        "    init_image=initial_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=64,\n",
        "    cut_pow=1.,\n",
        "    display_freq=display_frequency,\n",
        "    seed=seed,\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g7EDme5RYCrt",
        "cellView": "form"
      },
      "source": [
        "#@title Execution\n",
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "print('Using device:', device)\n",
        "if prompts:\n",
        "    print('Using text prompt:', prompts)\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",
        "\n",
        "cut_size = perceptor.visual.input_resolution\n",
        "e_dim = model.quantize.e_dim\n",
        "f = 2**(model.decoder.num_resolutions - 1)\n",
        "make_cutouts = MakeCutouts(cut_size, args.cutn, cut_pow=args.cut_pow)\n",
        "n_toks = model.quantize.n_e\n",
        "toksX, toksY = args.size[0] // f, args.size[1] // f\n",
        "sideX, sideY = toksX * f, toksY * f\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",
        "\n",
        "if args.init_image:\n",
        "    pil_image = Image.open(args.init_image).convert('RGB')\n",
        "    pil_image = pil_image.resize((sideX, sideY), Image.LANCZOS)\n",
        "    z, *_ = model.encode(TF.to_tensor(pil_image).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",
        "    z = z.view([-1, toksY, toksX, e_dim]).permute(0, 3, 1, 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",
        "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 = resize_image(Image.open(path).convert('RGB'), (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",
        "    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 add_xmp_data(nombrefichero):\n",
        "    image = ImgTag(filename=nombrefichero)\n",
        "    image.xmp.append_array_item(libxmp.consts.XMP_NS_DC, 'creator', 'VQGAN+CLIP', {\"prop_array_is_ordered\":True, \"prop_value_is_array\":True})\n",
        "    if args.prompts:\n",
        "        image.xmp.append_array_item(libxmp.consts.XMP_NS_DC, 'title', \" | \".join(args.prompts), {\"prop_array_is_ordered\":True, \"prop_value_is_array\":True})\n",
        "    else:\n",
        "        image.xmp.append_array_item(libxmp.consts.XMP_NS_DC, 'title', 'None', {\"prop_array_is_ordered\":True, \"prop_value_is_array\":True})\n",
        "    image.xmp.append_array_item(libxmp.consts.XMP_NS_DC, 'i', str(i), {\"prop_array_is_ordered\":True, \"prop_value_is_array\":True})\n",
        "    image.xmp.append_array_item(libxmp.consts.XMP_NS_DC, 'model', model_name, {\"prop_array_is_ordered\":True, \"prop_value_is_array\":True})\n",
        "    image.xmp.append_array_item(libxmp.consts.XMP_NS_DC, 'seed',str(seed) , {\"prop_array_is_ordered\":True, \"prop_value_is_array\":True})\n",
        "    image.xmp.append_array_item(libxmp.consts.XMP_NS_DC, 'input_images',str(input_images) , {\"prop_array_is_ordered\":True, \"prop_value_is_array\":True})\n",
        "    #for frases in args.prompts:\n",
        "    #    image.xmp.append_array_item(libxmp.consts.XMP_NS_DC, 'Prompt' ,frases, {\"prop_array_is_ordered\":True, \"prop_value_is_array\":True})\n",
        "    image.close()\n",
        "\n",
        "def add_stegano_data(filename):\n",
        "    data = {\n",
        "        \"title\": \" | \".join(args.prompts) if args.prompts else None,\n",
        "        \"notebook\": \"VQGAN+CLIP\",\n",
        "        \"i\": i,\n",
        "        \"model\": model_name,\n",
        "        \"seed\": str(seed),\n",
        "        \"input_images\": input_images\n",
        "    }\n",
        "    lsb.hide(filename, json.dumps(data)).save(filename)\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",
        "    add_stegano_data('progress.png')\n",
        "    add_xmp_data('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",
        "\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",
        "    filename = f\"steps/{i:04}.png\"\n",
        "    imageio.imwrite(filename, np.array(img))\n",
        "    add_stegano_data(filename)\n",
        "    add_xmp_data(filename)\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",
        "    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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "02ZbcWw5YYnU"
      },
      "source": [
        "## Generate a video with the results\n",
        "\n",
        "If you want to generate a video with the images as frames, just click below. You can modify the number of FPS, the initial frame, the last frame, etc."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mFo5vz0UYBrF",
        "cellView": "form"
      },
      "source": [
        "#@title Generate video using ffmpeg\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 = 30\n",
        "\n",
        "total_frames = last_frame-init_frame\n",
        "\n",
        "length = 15 # Desired video runtime in seconds\n",
        "\n",
        "frames = []\n",
        "tqdm.write('Generating video...')\n",
        "for i in range(init_frame,last_frame): #\n",
        "    filename = f\"steps/{i:04}.png\"\n",
        "    frames.append(Image.open(filename))\n",
        "\n",
        "#fps = last_frame/10\n",
        "fps = np.clip(total_frames/length,min_fps,max_fps)\n",
        "\n",
        "# Names the video after the prompt if there is one, if not, defaults to video.mp4\n",
        "def listToString(s): \n",
        "    \n",
        "    # initialize an empty string\n",
        "    str1 = \"\" \n",
        "    \n",
        "    # traverse in the string  \n",
        "    for ele in s: \n",
        "        str1 += ele  \n",
        "    \n",
        "    # return string  \n",
        "    return str1 \n",
        "        \n",
        "\n",
        "video_filename = \"video\"\n",
        "\n",
        "if len(prompts) > 0:\n",
        "    video_filename = listToString(prompts).replace(\" \",\"_\")\n",
        "    video_filename = video_filename[:120]\n",
        "    print(\"Video filename: \"+ video_filename)\n",
        "\n",
        "video_filename = video_filename + \".mp4\"\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_filename], stdin=PIPE)\n",
        "for im in tqdm(frames):\n",
        "    im.save(p.stdin, 'PNG')\n",
        "p.stdin.close()\n",
        "\n",
        "print(\"Compressing video...\")\n",
        "p.wait()\n",
        "print(\"Video ready\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E8lvN6b0mb-b",
        "cellView": "form"
      },
      "source": [
        "# @title View video in browser\n",
        "# @markdown This process may take a little longer. If you don't want to wait, download it by executing the next cell instead of using this cell.\n",
        "mp4 = open(video_filename,'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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y0e8pHyJmi7s"
      },
      "source": [
        "# @title Download video\n",
        "from google.colab import files\n",
        "files.download(video_filename)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}