CLIP+VQGAN_crowsonkb.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "CLIP+VQGAN_crowsonkb.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "mount_file_id": "1cRKYm9EoJ9-G5w6Pb3CkSWf6-DbKXprU",
      "authorship_tag": "ABX9TyM+OWHQPQ2Ls++dIC8Yimm6",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/olaviinha/NeuralTextToImage/blob/main/CLIP%2BVQGAN_crowsonkb.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pVxFPluEoUka"
      },
      "source": [
        "#<font face=\"Trebuchet MS\" size=\"6\">Big Sleep: CLIP+VQGAN <font color=\"#999\" size=\"4\">&nbsp;&nbsp;&nbsp;|&nbsp;&nbsp;&nbsp;</font><font color=\"#999\" size=\"4\">CLIP+VQGAN: Neural text-to-image</font><font color=\"#999\" size=\"4\">&nbsp;&nbsp;&nbsp;|&nbsp;&nbsp;&nbsp;</font><a href=\"https://github.com/olaviinha/NeuralImageGeneration\" target=\"_blank\"><font color=\"#999\" size=\"4\">Github</font></a>\n",
        "\n",
        "Big Sleep generates images from text input. It's originally a combination of [CLIP](https://github.com/openai/CLIP) by OpenAI and [BigGAN](https://arxiv.org/abs/1809.11096) by Andrew Brock et al., a concept introduced by [Ryan Murdock](https://github.com/rynmurdock) (Adobe) in his [original notebook](https://colab.research.google.com/drive/1NCceX2mbiKOSlAd_o7IU7nA9UskKN5WR?usp=sharing). This notebook is based on CLIP Guided Diffusion of a [VQGAN](https://arxiv.org/abs/2012.09841)+CLIP notebook by [Katherine Crowson](https://github.com/crowsonkb).\n",
        "\n",
        "<hr size=\"1\" color=\"#666\">\n",
        "\n",
        "### Tips:\n",
        "- Enter a simple string of text to `generate_image_of` field. Advanced usage: You may also 1) use a semicolon `;` as a separator to batch process multiple strings of texts to images in one go, 2) use a pipe `|` to train the image on multiple strings of text. If field is left empty, a random blog headline will be used.\n",
        "- Enter `output_dir` path relative to your Google Drive root, or leave blank to not save output anywhere outside this notebook. Each run of the _Sleep_ cell will **create a new subdirectory** under `output_dir`, under which all material will be saved.\n",
        "- `initial_image` and `target_image` define the \"shape\" and \"texture\" of the result image; best way to understand what they do is just to test them out. Field values may be URL addresses to images online or paths to images located in your Google Drive (enter path relative to Drive root).\n",
        "- `random_initial_photo` and `random_target_photo` will use random photographs from [Lorem Picsum](https://picsum.photos/)\n",
        "- You may find that over 400 iterations is usually a waste of time.\n",
        "- If you get _CUDA out of memory_ errors, reduce width and/or height. If it still happens (when you know it shouldn't), factory reset runtime). Colab is unable to generate images larger than ~0.5 MP, or 700x700 pixels. Some standard aspect ratios as maximum resolutions that Colab Pro can handle (from the original notebook):‎‎‎‎‏‏‎\n",
        "\n",
        " ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ 1:1 = 700 x 700‎‎‎‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎4:3 = 808 x 606‎‎‎‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ 16:9 = 928 x 522‎‎‎‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ 2:1 = 988 x 494‎‎‎‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ ‎‏‏‎ 1.66:1 = 903 x 544\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4D3YFbKGoSje",
        "cellView": "form"
      },
      "source": [
        "#@title #Setup\n",
        "#@markdown This cell needs to be run only once. It will mount your Google Drive and setup prerequisites.\n",
        "\n",
        "force_setup = False\n",
        "pip_packages = 'ftfy regex tqdm omegaconf pytorch-lightning kornia einops stegano python-xmp-toolkit imgtag imageio-ffmpeg transformers pillow==7.1.2'\n",
        "main_repository = ''\n",
        "\n",
        "import os\n",
        "from google.colab import output\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "%cd /content/\n",
        "\n",
        "# inhagcutils\n",
        "if not os.path.isfile('/content/inhagcutils.ipynb') or force_setup == True:\n",
        "  !pip -q install import-ipynb {pip_packages}\n",
        "  !apt -q install exempi\n",
        "  !curl -s -O https://raw.githubusercontent.com/olaviinha/inhagcutils/master/inhagcutils.ipynb\n",
        "import import_ipynb\n",
        "from inhagcutils import *\n",
        "\n",
        "# Mount Drive\n",
        "if not os.path.isdir('/content/drive') and force_setup == False:\n",
        "  from google.colab import drive\n",
        "  drive.mount('/content/drive')\n",
        "\n",
        "# Drive symlink\n",
        "if not os.path.isdir('/content/mydrive') and force_setup == False:\n",
        "  os.symlink('/content/drive/My Drive', '/content/mydrive')\n",
        "  drive_root_set = True\n",
        "drive_root = '/content/mydrive/'\n",
        "\n",
        "# !git clone {main_repository}\n",
        "\n",
        "#-----------------------------------\n",
        " \n",
        "!git clone https://github.com/openai/CLIP                 &> /dev/null\n",
        "!git clone https://github.com/CompVis/taming-transformers &> /dev/null\n",
        "\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 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",
        "# !pip install imageio-ffmpeg                               &> /dev/null\n",
        "# !pip install transformers                                 &> /dev/null\n",
        "\n",
        "downloaded_models = []\n",
        "\n",
        "#-----------------------------------\n",
        "\n",
        "# @title Carga de bibliotecas y definiciones\n",
        " \n",
        "import argparse\n",
        "import math\n",
        "from pathlib import Path\n",
        "import sys\n",
        "import requests\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",
        "from urllib.parse import urljoin, urlparse\n",
        "import itertools\n",
        "ImageFile.LOAD_TRUNCATED_IMAGES = True\n",
        "\n",
        "def keyboardInterruptHandler():\n",
        "  global dir_output, uniq_id\n",
        "  op(c.warn, 'Interrupted!', 'Cleaning up...')\n",
        "  remove_dirs([dir_output])\n",
        "  print('Run', uniq_id, 'directory and content removed:', dir_output)\n",
        "  sys.exit()\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",
        "def fix_path(path, add_slash=False):\n",
        "  if path.endswith('/'):\n",
        "    path = path #path[:-1]\n",
        "  if not path.endswith('/'):\n",
        "    path = path+\"/\"\n",
        "  if path.startswith('/') and add_slash == True:\n",
        "    path = path[1:]\n",
        "  return path\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)\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",
        "    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",
        "    global dir_progress, dir_output, max_iteraciones, id, display_save_every\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",
        "    if display_save_every:\n",
        "      display.display(display.Image('progress.png'))\n",
        "    save_img = dir_progress+'step_'+str(i).zfill(4)+'.png'\n",
        "    !cp progress.png {save_img}\n",
        "    # if i is max_iteraciones:\n",
        "    #   !cp progress.png {dir_output}{id}.png\n",
        "    # else:\n",
        "    prog_path = dir_progress.replace(drive_root, '')\n",
        "    prog_img = prog_path+'step_'+str(i).zfill(4)+'.png'\n",
        "    if display_save_every:\n",
        "      op(c.ok, '^Image saved as', prog_img + '\\n')\n",
        "    else:\n",
        "      op(c.ok, 'Image saved as', prog_img)\n",
        "\n",
        "def ascend_txt():\n",
        "    global i, dir_steps\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\"{dir_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",
        "def append_txt(file, content):\n",
        "  txt = open(txt_file, 'a+') \n",
        "  txt.writelines(content+'\\n')\n",
        "  txt.close();\n",
        "#-----------------------------------\n",
        "\n",
        "dir_tmp = '/content/tmp/'\n",
        "dir_steps = '/content/tmp/steps/'\n",
        "dir_initial = '/content/tmp/init/'\n",
        "dir_target = '/content/tmp/target/'\n",
        "create_dirs([dir_tmp, dir_steps, dir_initial, dir_target])\n",
        "\n",
        "output.clear()\n",
        "# !nvidia-smi\n",
        "op(c.ok, 'Setup finished.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Xj9rq4yHqiYH",
        "cellView": "form"
      },
      "source": [
        "#@markdown <br>\n",
        "\n",
        "#@markdown #S̛̞̩͎͓ ̦̤͉͚̏ ̧̠͋͘ͅl͕̞͕̝͗̐͘.̠̰̳̫̈́̚ ̡͉̼̩̬̈́̇͒͘ȩ̨͎͛̔͆͊̏͜ͅ.͕̩̹̠̕͜ ̛̦̦̮e̢͐͊͂̀̊ͅ ̜̙̝̊͋ ̬̝̱̱͗p̮̎̽̌\n",
        "\n",
        "#@markdown <br>\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "clean_dirs([dir_steps, dir_initial, dir_target])\n",
        "\n",
        "generate_image_of = \"\" #@param {type:\"string\"}\n",
        "width = 520 #@param {type:\"slider\", min:512, max:1024, step:2}\n",
        "height = 520 #@param {type:\"slider\", min:512, max:1024, step:2}\n",
        "output_dir = '' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown <hr color=\"#666\" size=\"1\"><br>\n",
        "\n",
        "#@markdown ## Advanced settings\n",
        "\n",
        "model = \"vqgan_imagenet_f16_16384\" #@param [\"vqgan_imagenet_f16_16384\", \"vqgan_imagenet_f16_1024\", \"wikiart_1024\", \"wikiart_16384\", \"coco\", \"faceshq\", \"sflckr\"]\n",
        "# model = \"vqgan_imagenet_f16_16384\" #@param [\"vqgan_imagenet_f16_16384\", \"wikiart_16384\"]\n",
        "\n",
        "iterations = 200 #@param {type:\"slider\", min:0, max:2000, step:100}\n",
        "save_every = 50 #@param {type:\"slider\", min:0, max:500, step:1}\n",
        "display_save_every = True #@param {type:\"boolean\"}\n",
        "create_video = False #@param {type:\"boolean\"}\n",
        "#@markdown **Specific images (URL or parth)**\n",
        "initial_image = '' #@param {type:\"string\"}\n",
        "target_image = '' #@param {type:\"string\"}\n",
        "#@markdown **Random images from Lorem Picsum**\n",
        "random_initial_photo = True #@param {type:\"boolean\"}\n",
        "random_target_photo = False #@param {type:\"boolean\"}\n",
        "# #@markdown **Random images from Google Image Search (requires Google Custom Search)**\n",
        "# google_api_key = '' #@param {type:\"string\"}\n",
        "# google_search_initial = False #@param {type:\"boolean\"}\n",
        "# google_search_target = False #@param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "save_all_steps = False\n",
        "remove_interrupted = True\n",
        "repetitions = 0\n",
        "\n",
        "#\n",
        "# --- Very advanced settings ---------------------------------\n",
        "\n",
        "# #@markdown <hr color=\"#666\" size=\"1\">\n",
        "# #@markdown <font size=\"1\">&nbsp;</font>\n",
        "\n",
        "# #@markdown ### Very advanced settings\n",
        "# repetitions = 0 #@param {type:\"slider\", min:0, max:20, step:1}\n",
        "# save_all_steps = False #@param {type:\"boolean\"}\n",
        "# remove_interrupted = True #@param {type:\"boolean\"}\n",
        "\n",
        "# --- //Very advanced settings -------------------------------\n",
        "#\n",
        "\n",
        "input_images = \"\"\n",
        "text = generate_image_of\n",
        "textos = generate_image_of\n",
        "modelo = model\n",
        "ancho = width\n",
        "alto = height\n",
        "imagen_inicial = initial_image\n",
        "imagenes_objetivo = target_image\n",
        "max_iteraciones = iterations\n",
        "intervalo_imagenes = save_every\n",
        "seed = -1\n",
        "display_model = model\n",
        "\n",
        "dlmsg = 'Downloading required model...\\n'\n",
        "if model is 'vqgan_imagenet_f16_1024' and 'vqgan_imagenet_f16_1024' not in downloaded_models:\n",
        "  op(c.title, dlmsg)\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",
        "  downloaded_models.append('vqgan_imagenet_f16_1024')\n",
        "  output.clear()\n",
        "if model is 'vqgan_imagenet_f16_16384' and 'vqgan_imagenet_f16_16384' not in downloaded_models:\n",
        "  op(c.title, dlmsg)\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",
        "  downloaded_models.append('vqgan_imagenet_f16_16384')\n",
        "  output.clear()\n",
        "if model is 'coco' and 'coco' not in downloaded_models:\n",
        "  op(c.title, dlmsg)\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",
        "  downloaded_models.append('coco')\n",
        "  output.clear()\n",
        "if model is 'faceshq' and 'faceshq' not in downloaded_models:\n",
        "  op(c.title, dlmsg)\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",
        "  downloaded_models.append('faceshq')\n",
        "  output.clear()\n",
        "if model is 'wikiart_1024 'and 'wikiart_1024' not in downloaded_models:\n",
        "  op(c.title, dlmsg)\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",
        "  downloaded_models.append('wikiart_1024')\n",
        "  output.clear()\n",
        "if model is 'wikiart_16384' and 'wikiart_16384' not in downloaded_models:\n",
        "  op(c.title, dlmsg)\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",
        "  downloaded_models.append('wikiart_16384')\n",
        "  output.clear()\n",
        "if model is 'sflckr' and 'sflckr' not in downloaded_models:\n",
        "  op(c.title, dlmsg)\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",
        "  downloaded_models.append('sflckr')\n",
        "  output.clear()\n",
        "\n",
        "op(c.title, 'Majik & wizardry...\\n')\n",
        "\n",
        "nombres_modelos = {\n",
        "  \"vqgan_imagenet_f16_16384\": \"ImageNet 16384\",\n",
        "  \"vqgan_imagenet_f16_1024\":\"ImageNet 1024\", \n",
        "  \"wikiart_1024\":\"WikiArt 1024\",\n",
        "  \"wikiart_16384\":\"WikiArt 16384\",\n",
        "  \"coco\":\"COCO-Stuff\",\n",
        "  \"faceshq\":\"FacesHQ\",\n",
        "  \"sflckr\":\"S-FLCKR\"\n",
        "}\n",
        "nombre_modelo = nombres_modelos[modelo]     \n",
        "\n",
        "if output_dir is '' or not output_dir:\n",
        "  drive_root = '/content/fauxdrive/'\n",
        "  output_dir = 'output'\n",
        "  create_dirs([drive_root, drive_root+output_dir])\n",
        "else:\n",
        "  drive_root = '/content/mydrive/'\n",
        "\n",
        "if text is not '':\n",
        "  if check_input_type(drive_root+text) is 'file':\n",
        "    input_txt = drive_root+text\n",
        "    with open(input_txt) as f:\n",
        "      texts = f.readlines()\n",
        "      texts = [x.strip() for x in texts] \n",
        "  elif \";\" in text:\n",
        "    texts = text.split(';')\n",
        "    texts = [text.strip() for text in texts]\n",
        "  else:\n",
        "    texts = [text]\n",
        "else:\n",
        "  text = requests.get('https://api.inha.asia/headline/').text\n",
        "  texts = [text]\n",
        "if repetitions > 0:\n",
        "  texts = list(itertools.chain.from_iterable(itertools.repeat(x, repetitions) for x in texts))\n",
        "\n",
        "def remove_params(url):\n",
        "  return urljoin(url, urlparse(url).path)\n",
        "\n",
        "if seed == -1:\n",
        "  seed = None\n",
        "\n",
        "# ------------------------------------------\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",
        "\n",
        "uniq_id = gen_id()\n",
        "repeat_index = 1\n",
        "\n",
        "for textos in texts:\n",
        "\n",
        "  imagen_inicial = initial_image\n",
        "  imagenes_objetivo = target_image\n",
        "\n",
        "  if random_initial_photo is True:\n",
        "    imagen_inicial = 'https://picsum.photos/'+str(width)+'/'+str(height)\n",
        "  if imagen_inicial == '' or not imagen_inicial:\n",
        "    imagen_inicial = None\n",
        "  else:\n",
        "    if check_input_type(imagen_inicial) is \"link\":\n",
        "      imagen_inicial = remove_params(imagen_inicial)\n",
        "      ext = path_ext(imagen_inicial) or '.jpg'\n",
        "      new_initial = dir_initial+'initial'+ext\n",
        "      !wget -O {new_initial} {imagen_inicial}\n",
        "      imagen_inicial = new_initial\n",
        "    else:\n",
        "      imagen_inicial = drive_root+imagen_inicial\n",
        "  \n",
        "  if random_target_photo is True:\n",
        "    imagenes_objetivo = 'https://picsum.photos/'+str(width)+'/'+str(height)\n",
        "  if imagenes_objetivo == '' or not imagenes_objetivo:\n",
        "    imagenes_objetivo = []\n",
        "  elif check_input_type(imagenes_objetivo) is \"link\":\n",
        "    imagenes_objetivo = imagenes_objetivo.split(\"|\")\n",
        "    for i, img in enumerate(imagenes_objetivo):\n",
        "      new_target = dir_target+path_leaf(remove_params(img))\n",
        "      !wget -O {new_target} {img}\n",
        "      imagenes_objetivo[i] = new_target\n",
        "  else:\n",
        "    if \";\" in imagenes_objetivo:\n",
        "      imgs = imagenes_objetivo.split(';')\n",
        "      imagenes_objetivo = [drive_root+image.strip() for image in imags]\n",
        "    else:\n",
        "      imagenes_objetivo = [drive_root+imagenes_objetivo]\n",
        "  if imagen_inicial or imagenes_objetivo != []:\n",
        "      input_images = True\n",
        "\n",
        "  print( imagenes_objetivo )\n",
        "\n",
        "  if save_all_steps is False:\n",
        "    clean_dirs([dir_steps])\n",
        "\n",
        "  display_text = textos\n",
        "  title = textos.split(\"|\")[0].title()\n",
        "  file_title = ''.join(e for e in title if e.isalnum())\n",
        "  textos = [frase.strip() for frase in textos.split(\"|\")]\n",
        "\n",
        "  id = uniq_id+'_'+file_title\n",
        "  if repetitions > 0:\n",
        "    id = uniq_id+'_'+str(repeat_index)+'_'+file_title\n",
        "\n",
        "  dir_output = fix_path(drive_root+output_dir)+id+'/'\n",
        "  dir_progress = dir_output+'progress/'\n",
        "  create_dirs([dir_output, dir_progress])\n",
        "\n",
        "  if save_all_steps is True:\n",
        "    dir_steps = dir_output+'steps/'\n",
        "    create_dirs([dir_steps])\n",
        "\n",
        "  if textos == ['']:\n",
        "    textos = []\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",
        "  )\n",
        "\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",
        "  # ------------------------------------------------\n",
        "\n",
        "  print( args.vqgan_config )\n",
        "  print( args.vqgan_checkpoint )\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",
        "  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",
        "    print( prompt )\n",
        "    path, weight, stop = parse_prompt(prompt)\n",
        "    print( path )\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",
        "  #-----------------------------------\n",
        "\n",
        "  txt_file = dir_output+uniq_id+'-info.txt'\n",
        "  txt = open(txt_file,'w') \n",
        "  params = [timestamp()+' '+uniq_id+'\\n\\n',\n",
        "            'model:         '+display_model+'\\n',\n",
        "            'iterations:    '+str(iterations)+'\\n',\n",
        "            'text:          '+text+'\\n',\n",
        "            'initial_image: '+initial_image+'\\n',\n",
        "            'target_image:  '+target_image+'\\n']\n",
        "  txt.writelines(params)\n",
        "  txt.close()\n",
        "\n",
        "  output.clear()\n",
        "  op(c.title, '\\nGenerating image of:', display_text)\n",
        "  op(c.title, 'Run ID:', uniq_id)\n",
        "  if repetitions > 0:\n",
        "    op(c.title, 'Repetition:', repeat_index)\n",
        "  op(c.okb, 'Sweet dreams.\\n')\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",
        "    if remove_interrupted: keyboardInterruptHandler()\n",
        "\n",
        "  last_step = dir_steps+str(i).zfill(4)+'.png'\n",
        "  fin_out = dir_output+file_title+'.png'\n",
        "  !cp {last_step} {fin_out}\n",
        "  display_fin = fin_out.replace(drive_root, '')\n",
        "  op(c.ok, 'Final image saved as', display_fin)\n",
        "\n",
        "  # ---------------------------------------------------------\n",
        "  # - video\n",
        "\n",
        "  if create_video is True:\n",
        "    op(c.title, '\\nGenerating video')\n",
        "\n",
        "    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",
        "    fps = 30\n",
        "    output_video = dir_output+file_title+'.mp4'\n",
        "\n",
        "    frames = []\n",
        "\n",
        "    for i in range(init_frame,last_frame): #\n",
        "      filename = f\"{dir_steps}/{i:04}.png\"\n",
        "      frames.append(Image.open(filename))\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', '13', '-preset', 'veryslow', output_video], stdin=PIPE)\n",
        "    for im in tqdm(frames):\n",
        "      im.save(p.stdin, 'PNG')\n",
        "    p.stdin.close()\n",
        "    p.wait()\n",
        "\n",
        "    #---------------\n",
        "\n",
        "    fin_vid = fin_out.replace('.png', '.mp4')\n",
        "    op(c.ok, 'Video saved as', fin_vid)\n",
        "\n",
        "  if repeat_index is repetitions:\n",
        "    repeat_index = 1\n",
        "  else:\n",
        "    repeat_index += 1\n",
        "\n",
        "op(c.title, '\\nFIN.')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}