Copy_of_VQGAN+CLIP_(z+quantize_method_con_augmentations,_interfaz_amigable).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Copy of VQGAN+CLIP (z+quantize method con augmentations, interfaz amigable).ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/meefs/art/blob/master/Copy_of_VQGAN%2BCLIP_(z%2Bquantize_method_con_augmentations%2C_interfaz_amigable).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": [
        "# Genera imágenes desde frases de texto con VQGAN y CLIP (método z+quantize con augmentations).\n",
        "\n",
        "Notebook hecho por Katherine Crowson (https://github.com/crowsonkb, https://twitter.com/RiversHaveWings). El método BigGAN+CLIP original fue hecho por https://twitter.com/advadnoun. Traducido y agregadas explicaciones, y modificaciones por Eleiber#8347, y la interfaz amigable fue hecha gracias a Abulafia#3734.\n",
        "\n",
        "Para un tutorial detallado de cómo usarlo, recomiendo [visitar este artículo](https://tuscriaturas.miraheze.org/wiki/Ayuda:Crear_imágenes_con_VQGAN+CLIP), hecho por Jakeukalane#2767 y Avengium (Ángel)#3715\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": "TkUfzT60ZZ9q",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3dbf6099-e7e7-4e8e-b3c2-c9d3792a71a9"
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Fri Oct 22 15:42:44 2021       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 470.74       Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla P100-PCIE...  Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   34C    P0    26W / 250W |      0MiB / 16280MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wSfISAhyPmyp",
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "de9885fe-a05d-4f4c-8478-2a549be41d65"
      },
      "source": [
        "# @title Instalación de bibliotecas\n",
        "# @markdown Esta celda tardará un poco porque tiene que descargar varias librerías\n",
        " \n",
        "print(\"Descargando CLIP...\")\n",
        "!git clone https://github.com/openai/CLIP                 &> /dev/null\n",
        " \n",
        "print(\"Instalando bibliotecas de Python para IA...\")\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",
        "!pip install wget                                         &> /dev/null\n",
        " \n",
        "print(\"Instalando bibliotecas para manejo de metadatos...\")\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",
        " \n",
        "print(\"Instalando bibliotecas de Python para creación de vídeos...\")\n",
        "!pip install imageio-ffmpeg                               &> /dev/null\n",
        "!mkdir steps\n",
        "print(\"Instalación finalizada.\")"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Descargando CLIP...\n",
            "Instalando bibliotecas de Python para IA...\n",
            "Instalando bibliotecas para manejo de metadatos...\n",
            "Instalando bibliotecas de Python para creación de vídeos...\n",
            "Instalación finalizada.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FhhdWrSxQhwg"
      },
      "source": [
        "#@title Selección de modelos a descargar\n",
        "#@markdown Por defecto, el notebook descarga el modelo 16384 de ImageNet. Existen otros que no se descargan por defecto, ya que sería en vano si no los vas a usar, así que si quieres usarlos, simplemente selecciona los modelos a descargar.\n",
        "\n",
        "imagenet_1024 = False #@param {type:\"boolean\"}\n",
        "imagenet_16384 = False #@param {type:\"boolean\"}\n",
        "gumbel_8192 = False #@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",
        "ade20k = False #@param {type:\"boolean\"}\n",
        "ffhq = False #@param {type:\"boolean\"}\n",
        "celebahq = False #@param {type:\"boolean\"}\n",
        "\n",
        "if imagenet_1024:\n",
        "  !curl -L -o vqgan_imagenet_f16_1024.yaml -C - 'https://heibox.uni-heidelberg.de/d/8088892a516d4e3baf92/files/?p=%2Fconfigs%2Fmodel.yaml&dl=1' #ImageNet 1024\n",
        "  !curl -L -o vqgan_imagenet_f16_1024.ckpt -C - 'https://heibox.uni-heidelberg.de/d/8088892a516d4e3baf92/files/?p=%2Fckpts%2Flast.ckpt&dl=1'  #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",
        "if gumbel_8192:\n",
        "  !curl -L -o gumbel_8192.yaml -C - 'https://heibox.uni-heidelberg.de/d/2e5662443a6b4307b470/files/?p=%2Fconfigs%2Fmodel.yaml&dl=1' #Gumbel 8192\n",
        "  !curl -L -o gumbel_8192.ckpt -C - 'https://heibox.uni-heidelberg.de/d/2e5662443a6b4307b470/files/?p=%2Fckpts%2Flast.ckpt&dl=1' #Gumbel 8192\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://eaidata.bmk.sh/data/Wikiart_16384/wikiart_f16_16384_8145600.yaml' #WikiArt 16384\n",
        "  !curl -L -o wikiart_16384.ckpt -C - 'http://eaidata.bmk.sh/data/Wikiart_16384/wikiart_f16_16384_8145600.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\n",
        "if ade20k:\n",
        "  !curl -L -o ade20k.yaml -C - 'https://static.miraheze.org/intercriaturaswiki/b/bf/Ade20k.txt' #ADE20K\n",
        "  !curl -L -o ade20k.ckpt -C - 'https://app.koofr.net/content/links/0f65c2cd-7102-4550-a2bd-07fd383aac9e/files/get/last.ckpt?path=%2F2020-11-20T21-45-44_ade20k_transformer%2Fcheckpoints%2Flast.ckpt' #ADE20K\n",
        "if ffhq:\n",
        "  !curl -L -o ffhq.yaml -C - 'https://app.koofr.net/content/links/0fc005bf-3dca-4079-9d40-cdf38d42cd7a/files/get/2021-04-23T18-19-01-project.yaml?path=%2F2021-04-23T18-19-01_ffhq_transformer%2Fconfigs%2F2021-04-23T18-19-01-project.yaml&force' #FFHQ\n",
        "  !curl -L -o ffhq.ckpt -C - 'https://app.koofr.net/content/links/0fc005bf-3dca-4079-9d40-cdf38d42cd7a/files/get/last.ckpt?path=%2F2021-04-23T18-19-01_ffhq_transformer%2Fcheckpoints%2Flast.ckpt&force' #FFHQ\n",
        "if celebahq:\n",
        "  !curl -L -o celebahq.yaml -C - 'https://app.koofr.net/content/links/6dddf083-40c8-470a-9360-a9dab2a94e96/files/get/2021-04-23T18-11-19-project.yaml?path=%2F2021-04-23T18-11-19_celebahq_transformer%2Fconfigs%2F2021-04-23T18-11-19-project.yaml&force' #CelebA-HQ\n",
        "  !curl -L -o celebahq.ckpt -C - 'https://app.koofr.net/content/links/6dddf083-40c8-470a-9360-a9dab2a94e96/files/get/last.ckpt?path=%2F2021-04-23T18-11-19_celebahq_transformer%2Fcheckpoints%2Flast.ckpt&force' #CelebA-HQ"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ixo-kVb1q-3S",
        "outputId": "6df1ab28-9167-4ff0-abdb-94bb8a429c9d"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EXMSuW2EQWsd",
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 368
        },
        "outputId": "2135cb70-bed2-492d-a208-48da6c69d93b"
      },
      "source": [
        "# @title Carga de bibliotecas y definiciones\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    # metadatos \n",
        "from libxmp import *         # metadatos\n",
        "import libxmp                # metadatos\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",
        "    elif config.model.target == 'taming.models.vqgan.GumbelVQ':\n",
        "        model = vqgan.GumbelVQ(**config.model.params)\n",
        "        print(config.model.params)\n",
        "        model.eval().requires_grad_(False)\n",
        "        model.init_from_ckpt(checkpoint_path)\n",
        "    else:\n",
        "        raise ValueError(f'unknown model type: {config.model.target}')\n",
        "    del model.loss\n",
        "    return model\n",
        " \n",
        " \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",
        "def download_img(img_url):\n",
        "    try:\n",
        "        return wget.download(img_url,out=\"input.jpg\")\n",
        "    except:\n",
        "        return\n"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "error",
          "ename": "ModuleNotFoundError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-5-23f3325cda1b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0momegaconf\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mOmegaConf\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mPIL\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mImage\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 13\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mtaming\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodels\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mcond_transformer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvqgan\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     14\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mtorch\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mnn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptim\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'taming'",
            "",
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0;32m\nNOTE: If your import is failing due to a missing package, you can\nmanually install dependencies using either !pip or !apt.\n\nTo view examples of installing some common dependencies, click the\n\"Open Examples\" button below.\n\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n"
          ],
          "errorDetails": {
            "actions": [
              {
                "action": "open_url",
                "actionText": "Open Examples",
                "url": "/notebooks/snippets/importing_libraries.ipynb"
              }
            ]
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1tthw0YaispD"
      },
      "source": [
        "## Herramientas para la ejecución:\n",
        "Principalmente lo que tendrás que modificar será `textos:`, ahí puedes colocar el o los textos que  quieres generar (separados con `|` ). Es una lista porque puedes poner más de un texto, y así la IA trate de 'mezclar' las imágenes, dándole la misma prioridad a ambos textos.\n",
        "\n",
        "Para usar una imagen inicial al modelo, solo debes subir un archivo al entorno del Colab (en la sección a la izquierda), y luego modificas `imagen_inicial:` poniendo el nombre exacto del archivo. Ejemplo: `sample.png`\n",
        "\n",
        "También puedes modificar el modelo cambiando las línea que dice `modelo:`. Actualmente están disponibles 1024, 16384, WikiArt, S-FLCKR y COCO-Stuff. Para activarlos tienes que haberlos descargado primero, y luego puedes simplemente seleccionarlo.\n",
        "\n",
        "También puedes usar `imagenes_objetivo`, que es básicamente ponerle una o más imágenes que la IA tomará como \"meta\", cumpliendo la misma función que ponerle un texto. Para poner más de una tienes que usar `|` como separador."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZdlpRFL8UAlW"
      },
      "source": [
        "#@title Parámetros\n",
        "textos = \"a photo-realistic painting of the self by Picaso\\\"\" #@param {type:\"string\"}\n",
        "ancho =  480#@param {type:\"number\"}\n",
        "alto =  480#@param {type:\"number\"}\n",
        "modelo = \"ffhq\" #@param [\"vqgan_imagenet_f16_16384\", \"vqgan_imagenet_f16_1024\", \"wikiart_1024\", \"wikiart_16384\", \"coco\", \"faceshq\", \"sflckr\", \"ade20k\", \"ffhq\", \"celebahq\", \"gumbel_8192\"]\n",
        "intervalo_imagenes =  25#@param {type:\"number\"}\n",
        "imagen_inicial = None#@param {type:\"string\"}\n",
        "imagenes_objetivo = None#@param {type:\"string\"}\n",
        "seed = 158758#@param {type:\"number\"}\n",
        "max_iteraciones = -1#@param {type:\"number\"}\n",
        "input_images = \"\"\n",
        "\n",
        "nombres_modelos={\"vqgan_imagenet_f16_16384\": 'ImageNet 16384',\"vqgan_imagenet_f16_1024\":\"ImageNet 1024\", \n",
        "                 \"wikiart_1024\":\"WikiArt 1024\", \"wikiart_16384\":\"WikiArt 16384\", \"coco\":\"COCO-Stuff\", \"faceshq\":\"FacesHQ\", \"sflckr\":\"S-FLCKR\", \"ade20k\":\"ADE20K\", \"ffhq\":\"FFHQ\", \"celebahq\":\"CelebA-HQ\", \"gumbel_8192\": \"Gumbel 8192\"}\n",
        "nombre_modelo = nombres_modelos[modelo]     \n",
        "\n",
        "if modelo == \"gumbel_8192\":\n",
        "    is_gumbel = True\n",
        "else:\n",
        "    is_gumbel = False\n",
        "\n",
        "if seed == -1:\n",
        "    seed = None\n",
        "if imagen_inicial == \"None\":\n",
        "    imagen_inicial = None\n",
        "elif imagen_inicial and imagen_inicial.lower().startswith(\"http\"):\n",
        "    imagen_inicial = download_img(imagen_inicial)\n",
        "\n",
        "\n",
        "if imagenes_objetivo == \"None\" or not imagenes_objetivo:\n",
        "    imagenes_objetivo = []\n",
        "else:\n",
        "    imagenes_objetivo = imagenes_objetivo.split(\"|\")\n",
        "    imagenes_objetivo = [image.strip() for image in imagenes_objetivo]\n",
        "\n",
        "if imagen_inicial or imagenes_objetivo != []:\n",
        "    input_images = True\n",
        "\n",
        "textos = [frase.strip() for frase in textos.split(\"|\")]\n",
        "if textos == ['']:\n",
        "    textos = []\n",
        "\n",
        "\n",
        "args = argparse.Namespace(\n",
        "    prompts=textos,\n",
        "    image_prompts=imagenes_objetivo,\n",
        "    noise_prompt_seeds=[],\n",
        "    noise_prompt_weights=[],\n",
        "    size=[ancho, alto],\n",
        "    init_image=imagen_inicial,\n",
        "    init_weight=0.,\n",
        "    clip_model='ViT-B/32',\n",
        "    vqgan_config=f'{modelo}.yaml',\n",
        "    vqgan_checkpoint=f'{modelo}.ckpt',\n",
        "    step_size=0.1,\n",
        "    cutn=64,\n",
        "    cut_pow=1.,\n",
        "    display_freq=intervalo_imagenes,\n",
        "    seed=seed,\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g7EDme5RYCrt"
      },
      "source": [
        "#@title Hacer la ejecución...\n",
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "print('Using device:', device)\n",
        "if textos:\n",
        "    print('Using texts:', textos)\n",
        "if imagenes_objetivo:\n",
        "    print('Using image prompts:', imagenes_objetivo)\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",
        "if is_gumbel:\n",
        "    e_dim = model.quantize.embedding_dim\n",
        "else:\n",
        "    e_dim = model.quantize.e_dim\n",
        "\n",
        "f = 2**(model.decoder.num_resolutions - 1)\n",
        "make_cutouts = MakeCutouts(cut_size, args.cutn, cut_pow=args.cut_pow)\n",
        "if is_gumbel:\n",
        "    n_toks = model.quantize.n_embed\n",
        "else:\n",
        "    n_toks = model.quantize.n_e\n",
        "\n",
        "toksX, toksY = args.size[0] // f, args.size[1] // f\n",
        "sideX, sideY = toksX * f, toksY * f\n",
        "if is_gumbel:\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",
        "    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",
        "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",
        "    if is_gumbel:\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_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",
        "    if is_gumbel:\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",
        "    \n",
        "    return clamp_with_grad(model.decode(z_q).add(1).div(2), 0, 1)\n",
        "\n",
        "def add_xmp_data(nombrefichero):\n",
        "    imagen = ImgTag(filename=nombrefichero)\n",
        "    imagen.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",
        "        imagen.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",
        "        imagen.xmp.append_array_item(libxmp.consts.XMP_NS_DC, 'title', 'None', {\"prop_array_is_ordered\":True, \"prop_value_is_array\":True})\n",
        "    imagen.xmp.append_array_item(libxmp.consts.XMP_NS_DC, 'i', str(i), {\"prop_array_is_ordered\":True, \"prop_value_is_array\":True})\n",
        "    imagen.xmp.append_array_item(libxmp.consts.XMP_NS_DC, 'model', nombre_modelo, {\"prop_array_is_ordered\":True, \"prop_value_is_array\":True})\n",
        "    imagen.xmp.append_array_item(libxmp.consts.XMP_NS_DC, 'seed',str(seed) , {\"prop_array_is_ordered\":True, \"prop_value_is_array\":True})\n",
        "    imagen.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",
        "    #    imagen.xmp.append_array_item(libxmp.consts.XMP_NS_DC, 'Prompt' ,frases, {\"prop_array_is_ordered\":True, \"prop_value_is_array\":True})\n",
        "    imagen.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\": nombre_modelo,\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_iteraciones:\n",
        "                break\n",
        "            i += 1\n",
        "            pbar.update()\n",
        "except KeyboardInterrupt:\n",
        "    pass"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "02ZbcWw5YYnU"
      },
      "source": [
        "## Genera un vídeo con los resultados\n",
        "\n",
        "Si quieres generar un vídeo con los frames, solo haz click abajo. Puedes modificar el número de FPS, el frame inicial, el último frame, etc."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mFo5vz0UYBrF"
      },
      "source": [
        "init_frame = 1 #Este es el frame donde el vídeo empezará\n",
        "last_frame = i #Puedes cambiar i a el número del último frame que quieres generar. 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 #Tiempo deseado del vídeo en segundos\n",
        "\n",
        "frames = []\n",
        "tqdm.write('Generando 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",
        "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",
        "\n",
        "print(\"El vídeo está siendo ahora comprimido, espera...\")\n",
        "p.wait()\n",
        "print(\"El vídeo está listo\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "E8lvN6b0mb-b"
      },
      "source": [
        "# @title Ver vídeo en el navegador\n",
        "# @markdown Este proceso puede tardar un poco más. Si no quieres esperar, descárgalo ejecutando la celda siguiente en vez de usar esta celda.\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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "Y0e8pHyJmi7s"
      },
      "source": [
        "# @title Descargar vídeo\n",
        "from google.colab import files\n",
        "files.download(\"video.mp4\")"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}