S2_GAN_Art_Generator_(VQGAN,_CLIP,_Guided_Diffusion).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "S2 GAN Art Generator (VQGAN, CLIP, Guided Diffusion) ",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [
        "VyN-gqJKfLAu"
      ],
      "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_GAN_Art_Generator_(VQGAN%2C_CLIP%2C_Guided_Diffusion).ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VyN-gqJKfLAu"
      },
      "source": [
        "### Follow & tag me in your art, revision requests etc. on Twitter: [@somewheresy](https://twitter.com/somewheresy)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CppIQlPhhwhs"
      },
      "source": [
        "# LEGACY VERSION: [S2 VQGAN+CLIP Classic.ipynb](https://github.com/benjismith/somewhere-ml/blob/main/S2_VQGAN%2BCLIP_Classic.ipynb)\n",
        "\n",
        "The notebook you're currently using is a multimodal GAN art generator patched together from various ML notebooks for generative art (see license). The results of this notebook may be distinct from others in the space.\n",
        "\n",
        "This notebook is great for procedurally generating new images from a text prompt or input image. At Somewhere Systems we use this for everything from generative landscapes to materials design for 3D graphics. Consider checking out our work @ https://s2.lol and hiring us to demystify technology like AR, ML, etc. \n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VA1PHoJrRiK9",
        "cellView": "form"
      },
      "source": [
        "#@title MIT License\n",
        "# \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": {
        "cellView": "form",
        "id": "5wCNUGHyfyPO"
      },
      "source": [
        "#@markdown What GPU am I using?\n",
        "\n",
        "#@markdown V100 > P100 > everything else\n",
        "\n",
        "!nvidia-smi --query-gpu=gpu_name,gpu_bus_id,vbios_version --format=csv\n",
        "gpu_name = !nvidia-smi --query-gpu=gpu_name, --format=csv"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KqgIZS75GU-O",
        "cellView": "form"
      },
      "source": [
        "#@title Connect Google Drive (recommended)\n",
        "import os\n",
        "abs_root_path = \"/content\"\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 (recommended)\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 = \"AI_ART\" #@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 = \"ALL_DATASETS_TEST\" #@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",
        "import os \n",
        "!nvidia-smi\n",
        "print(\"Downloading CLIP...\")\n",
        "!git clone https://github.com/openai/CLIP                 &> /dev/null\n",
        "!git clone https://github.com/crowsonkb/guided-diffusion  &> /dev/null\n",
        "!pip install -e ./CLIP                                    &> /dev/null\n",
        "print(\"Installing library for guided diffusion...\")\n",
        "!pip install -e ./guided-diffusion                        &> /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",
        " \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",
        "!pip install taming-transformers                          &> /dev/null        \n",
        "                            \n",
        "\n",
        "print(\"Installing ISR for image upscaling...\")\n",
        "!pip install ISR &> /dev/null \n",
        "\n",
        "print(\"Installing ffmpeg for creating videos...\")\n",
        "!pip install imageio-ffmpeg &> /dev/null\n",
        "\n",
        "if not os.path.exists(abs_root_path + \"/vqgan-steps\"):\n",
        "    !mkdir \"vqgan-steps\" &> /dev/null  \n",
        "    print(\"No directory for VQGAN+CLIP image output found. Made directory: ~/vqgan-steps\")\n",
        "if not os.path.exists(abs_root_path + \"/diffusion-steps\"):\n",
        "    !mkdir \"diffusion-steps\" &> /dev/null  \n",
        "    print(\"No directory for CLIP-guided diffusion image output found. Made directory: ~/diffusion-steps\")\n",
        "!pip freeze > requirements.txt\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 Ensure you select a model you've downloaded in the parameters block\n",
        "\n",
        "#@markdown *New!* It's CLIP-guided diffusion, a method that takes a bit longer to produce good results but generally makes more \"realistic\" interpretations of the prompt. If you'd like to use this, make sure to download Katherine Crowson's diffusion model. \n",
        "diffusion = True #@param {type: \"boolean\"}\n",
        "\n",
        "#@markdown Models for VQGAN+CLIP method\n",
        "imagenet_1024 = False #@param {type:\"boolean\"}\n",
        "imagenet_16384 = True #@param {type:\"boolean\"}\n",
        "coco = False #@param {type:\"boolean\"}\n",
        "faceshq = False #@param {type:\"boolean\"}\n",
        "# wikiart_1024 = True #@param {type:\"boolean\"}\n",
        "wikiart_16384 = False #@param {type:\"boolean\"}\n",
        "sflckr = False #@param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "if not os.path.exists(abs_root_path + \"/models\"):\n",
        "          os.mkdir(abs_root_path + \"/models\")\n",
        "os.chdir(abs_root_path + \"/models\")\n",
        "print(\"changing to models subdirectory: \")\n",
        "!pwd\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\n",
        "if diffusion:\n",
        "  # Download the diffusion model\n",
        "  !curl -L -o 512x512_diffusion_uncond_finetune_008100.pt -C - 'https://the-eye.eu/public/AI/models/512x512_diffusion_unconditional_ImageNet/512x512_diffusion_uncond_finetune_008100.pt' #DIFFUSION 512\n",
        "  !curl -L -o 512x512_diffusion_uncond_finetune_008100.pt -C - 'https://the-eye.eu/public/AI/models/512x512_diffusion_unconditional_ImageNet/512x512_diffusion_uncond_finetune_008100.pt' #DIFFUSION 256"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EXMSuW2EQWsd",
        "cellView": "form"
      },
      "source": [
        "# @title Load libraries and definitions\n",
        "print(abs_root_path)\n",
        "os.chdir(abs_root_path)\n",
        "!pwd\n",
        "\n",
        "import argparse\n",
        "import math\n",
        "from pathlib import Path\n",
        "import io\n",
        "import sys\n",
        "\n",
        "\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",
        "sys.path.append('./CLIP')\n",
        "sys.path.append('./guided-diffusion')\n",
        "\n",
        "import clip\n",
        "from guided_diffusion.script_util import create_model_and_diffusion, model_and_diffusion_defaults\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)\n",
        "\n",
        "# Define necessary functions for CLIP guided diffusion\n",
        "\n",
        "def fetch(url_or_path):\n",
        "    if str(url_or_path).startswith('http://') or str(url_or_path).startswith('https://'):\n",
        "        r = requests.get(url_or_path)\n",
        "        r.raise_for_status()\n",
        "        fd = io.BytesIO()\n",
        "        fd.write(r.content)\n",
        "        fd.seek(0)\n",
        "        return fd\n",
        "    return open(url_or_path, 'rb')\n",
        "\n",
        "class MakeCutouts(nn.Module):\n",
        "    def __init__(self, cut_size, cutn, cut_pow=1.):\n",
        "        super().__init__()\n",
        "        self.cut_size = cut_size\n",
        "        self.cutn = cutn\n",
        "        self.cut_pow = cut_pow\n",
        "\n",
        "    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(F.adaptive_avg_pool2d(cutout, self.cut_size))\n",
        "        return torch.cat(cutouts)\n",
        "\n",
        "\n",
        "def spherical_dist_loss(x, y):\n",
        "    x = F.normalize(x, dim=-1)\n",
        "    y = F.normalize(y, dim=-1)\n",
        "    return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2)\n",
        "\n",
        "\n",
        "def tv_loss(input):\n",
        "    \"\"\"L2 total variation loss, as in Mahendran et al.\"\"\"\n",
        "    input = F.pad(input, (0, 1, 0, 1), 'replicate')\n",
        "    x_diff = input[..., :-1, 1:] - input[..., :-1, :-1]\n",
        "    y_diff = input[..., 1:, :-1] - input[..., :-1, :-1]\n",
        "    return (x_diff**2 + y_diff**2).mean([1, 2, 3])\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1tthw0YaispD"
      },
      "source": [
        "## Implementation tools"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "JzM5ySlHJ3Pw"
      },
      "source": [
        "#@markdown #Global Parameters\n",
        "os.chdir(abs_root_path)\n",
        "\n",
        "seed = -1#@param {type:\"number\"}\n",
        "\n",
        "display_frequency =  50#@param {type:\"number\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4lACzTFbJzuc",
        "cellView": "form"
      },
      "source": [
        "usingDiffusion = False;\n",
        "#@markdown #VQGAN+CLIP Parameters & Execution\n",
        "prompts = \"A duck sitting in a pond\" #@param {type:\"string\"}\n",
        "width =  512#@param {type:\"number\"}\n",
        "height =  512#@param {type:\"number\"}\n",
        "#@markdown Note: x4 and x16 models for CLIP may not work reliably on lower-memory machines\n",
        "clip_model = \"ViT-B/32\" #@param ['RN50', 'RN101', 'RN50x4', 'RN50x16', 'ViT-B/32','ViT-B/16']\n",
        "vqgan_model = \"vqgan_imagenet_f16_16384\" #@param [\"vqgan_imagenet_f16_16384\", \"vqgan_imagenet_f16_1024\", \"wikiart_16384\", \"coco\", \"faceshq\", \"sflckr\"]\n",
        "initial_image = \"\"#@param {type:\"string\"}\n",
        "target_images = \"\"#@param {type:\"string\"}\n",
        "max_iterations = 10000#@param {type:\"number\"}\n",
        "input_images = \"\"\n",
        "#@markdown ### Advanced VQGAN+CLIP Parameters \n",
        "\n",
        "vq_init_weight = 0.0#@param {type:\"number\"}\n",
        "vq_step_size = 0.1#@param {type:\"number\"}\n",
        "vq_cutn = 64#@param {type:\"number\"}\n",
        "vq_cutpow = 1.0#@param {type:\"number\"}\n",
        "\n",
        "model_names={\"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\"}\n",
        "model_name = model_names[vqgan_model]     \n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "with torch.no_grad():\n",
        "    torch.cuda.empty_cache()\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",
        "\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= clip_model,\n",
        "        vqgan_config=f'models/{vqgan_model}.yaml',\n",
        "        vqgan_checkpoint=f'models/{vqgan_model}.ckpt',\n",
        "        step_size=vq_step_size,\n",
        "        cutn=vq_cutn,\n",
        "        cut_pow=vq_cutpow,\n",
        "        display_freq=display_frequency,\n",
        "        seed=seed,\n",
        "    )\n",
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "notebook_name = \"VQGAN+CLIP\"\n",
        "print('Executing using VQGAN+CLIP method')\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\": notebook_name,\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\"vqgan-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\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZdlpRFL8UAlW",
        "cellView": "form"
      },
      "source": [
        "#@markdown # CLIP Guided Diffusion Parameters & Execution\n",
        "#@markdown ##### WARNING: This requires access to 16GB of VRAM reliably, so may not work for users not using Colab Pro/+\n",
        "usingDiffusion = True;\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "with torch.no_grad():\n",
        "    torch.cuda.empty_cache()\n",
        "\n",
        "prompt = \"A group of vultures decide the behavior of the stock market\\\"\"#@param {type:\"string\"}\n",
        "batch_size = 1#@param {type:\"number\"}\n",
        "#@markdown Note: x4 and x16 models for CLIP may not work reliably on lower-memory machines\n",
        "clip_model = \"ViT-B/32\" #@param ['RN50', 'RN101', 'RN50x4', 'RN50x16', 'ViT-B/32','ViT-B/16']\n",
        "    #@markdown Controls how much the image should look like the prompt.\n",
        "clip_guidance_scale = 1000#@param {type:\"number\"}  \n",
        "    #@markdown Controls the smoothness of the final output.\n",
        "tv_scale = 150#@param {type:\"number\"}             \n",
        "cutn = 32#@param {type:\"number\"}\n",
        "cut_pow = 0.5#@param {type:\"number\"}\n",
        "n_batches = 1#@param {type:\"number\"}\n",
        "    #@markdown This can be an URL or Colab local path and must be in quotes.\n",
        "init_image = None #@param {type:\"string\"}\n",
        "    #@markdown This needs to be between approx. 200 and 500 when using an init image.  \n",
        "    #@markdown Higher values make the output look more like the init.\n",
        "skip_timesteps = 0#@param {type:\"number\"}  \n",
        "\n",
        "      \n",
        "\n",
        "if seed == -1:\n",
        "    seed = None\n",
        "\n",
        "diff_image_size = 256 # size of image when using diffusion\n",
        "diff_image_size = int(diff_image_size)\n",
        "\n",
        "model_config = model_and_diffusion_defaults()\n",
        "model_config.update({\n",
        "    'attention_resolutions': '32, 16, 8',\n",
        "    'class_cond': False,\n",
        "    'diffusion_steps': 1000,\n",
        "    'rescale_timesteps': True,\n",
        "    'timestep_respacing': '1000',  # Modify this value to decrease the number of\n",
        "                                  # timesteps.\n",
        "    'image_size': 512,\n",
        "    'learn_sigma': True,\n",
        "    'noise_schedule': 'linear',\n",
        "    'num_channels': 256,\n",
        "    'num_head_channels': 64,\n",
        "    'num_res_blocks': 2,\n",
        "    'resblock_updown': True,\n",
        "    'use_fp16': True,\n",
        "    'use_scale_shift_norm': True,\n",
        "})\n",
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "print('Executing using CLIP guided diffusion method')\n",
        "if (prompt != None):\n",
        "  print('Using prompt: '+ prompt)\n",
        "\n",
        "print('Using device:', device)\n",
        "\n",
        "model, diffusion = create_model_and_diffusion(**model_config)\n",
        "model.load_state_dict(torch.load(abs_root_path + \"/models/\" + '512x512_diffusion_uncond_finetune_008100.pt', map_location='cpu'))\n",
        "model.requires_grad_(False).eval().to(device)\n",
        "for name, param in model.named_parameters():\n",
        "    if 'qkv' in name or 'norm' in name or 'proj' in name:\n",
        "        param.requires_grad_()\n",
        "if model_config['use_fp16']:\n",
        "    model.convert_to_fp16()\n",
        "\n",
        "clip_model = clip.load(clip_model, jit=False)[0].eval().requires_grad_(False).to(device)\n",
        "clip_size = clip_model.visual.input_resolution\n",
        "normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],\n",
        "                                std=[0.26862954, 0.26130258, 0.27577711])\n",
        "def do_run():\n",
        "    if seed is not None:\n",
        "        torch.manual_seed(seed)\n",
        "\n",
        "    text_embed = clip_model.encode_text(clip.tokenize(prompt).to(device)).float()\n",
        "\n",
        "    init = None\n",
        "    if init_image is not None:\n",
        "        init = Image.open(fetch(init_image)).convert('RGB')\n",
        "        init = init.resize((model_config['image_size'], model_config['image_size']), Image.LANCZOS)\n",
        "        init = TF.to_tensor(init).to(device).unsqueeze(0).mul(2).sub(1)\n",
        "\n",
        "    make_cutouts = MakeCutouts(clip_size, cutn, cut_pow)\n",
        "\n",
        "    cur_t = None\n",
        "\n",
        "    def cond_fn(x, t, y=None):\n",
        "        with torch.enable_grad():\n",
        "            x = x.detach().requires_grad_()\n",
        "            n = x.shape[0]\n",
        "            my_t = torch.ones([n], device=device, dtype=torch.long) * cur_t\n",
        "            out = diffusion.p_mean_variance(model, x, my_t, clip_denoised=False, model_kwargs={'y': y})\n",
        "            fac = diffusion.sqrt_one_minus_alphas_cumprod[cur_t]\n",
        "            x_in = out['pred_xstart'] * fac + x * (1 - fac)\n",
        "            clip_in = normalize(make_cutouts(x_in.add(1).div(2)))\n",
        "            image_embeds = clip_model.encode_image(clip_in).float().view([cutn, n, -1])\n",
        "            dists = spherical_dist_loss(image_embeds, text_embed.unsqueeze(0))\n",
        "            losses = dists.mean(0)\n",
        "            tv_losses = tv_loss(x_in)\n",
        "            loss = losses.sum() * clip_guidance_scale + tv_losses.sum() * tv_scale\n",
        "            return -torch.autograd.grad(loss, x)[0]\n",
        "\n",
        "    if model_config['timestep_respacing'].startswith('ddim'):\n",
        "        sample_fn = diffusion.ddim_sample_loop_progressive\n",
        "    else:\n",
        "        sample_fn = diffusion.p_sample_loop_progressive\n",
        "\n",
        "    for i in range(n_batches):\n",
        "        cur_t = diffusion.num_timesteps - skip_timesteps - 1\n",
        "\n",
        "        samples = sample_fn(\n",
        "            model,\n",
        "            (batch_size, 3, model_config['image_size'], model_config['image_size']),\n",
        "            clip_denoised=False,\n",
        "            model_kwargs={},\n",
        "            cond_fn=cond_fn,\n",
        "            progress=True,\n",
        "            skip_timesteps=skip_timesteps,\n",
        "            init_image=init,\n",
        "            randomize_class=True,\n",
        "        )\n",
        "        \n",
        "        for j, sample in enumerate(samples):\n",
        "            cur_t -= 1     \n",
        "            for k, image in enumerate(sample['pred_xstart']):\n",
        "                filename = f'diffusion-steps/{batch_size * j:05}.png'\n",
        "                TF.to_pil_image(image.add(1).div(2).clamp(0, 1)).save(filename)\n",
        "                if j % display_frequency == 0 or cur_t == -1:\n",
        "                    tqdm.write(f'Batch {i}, step {j}, output {k}:')\n",
        "                    print()      \n",
        "                    display.display(display.Image(filename))\n",
        "            \n",
        "                \n",
        "\n",
        "do_run()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g_65YqAYVgUX"
      },
      "source": [
        "## Upscale an image"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "qcm-oJK0wljA"
      },
      "source": [
        "#@title ISR (Image Super-Resolution) Setup\n",
        "!pip install 'h5py==2.10.0' -q -q -q --force-reinstall &> /dev/null \n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "66PgI94EgMbk"
      },
      "source": [
        "#@title ISR Parameters\n",
        "\n",
        "sourceImage = \"/content/drive/MyDrive/FolderName/ProjectName/vqgan-steps/0235.png\" #@param {type: \"string\"}\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4b-CSxBpgy23",
        "cellView": "form"
      },
      "source": [
        "#@title ISR Execution\n",
        "#@markdown Takes a while, but it's free and gets the job done.\n",
        "upscaling_src_img = Image.open(sourceImage) \n",
        "lr_img = np.array(upscaling_src_img)\n",
        "\n",
        "from ISR.models import RDN, RRDN \n",
        "rrdn = RRDN(weights='gans')                \n",
        "rdn = RDN(weights='noise-cancel')          \n",
        "print(\"De-noising...\") \n",
        "sr_img_denoised = rdn.predict(lr_img)\n",
        "Image.fromarray(sr_img_denoised)\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "02ZbcWw5YYnU"
      },
      "source": [
        "## Generate a video\n",
        "##### Consider playing with higher FPS rates if you used the CLIP-guided diffusion method!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mFo5vz0UYBrF",
        "cellView": "form"
      },
      "source": [
        "#@title Generate video using ffmpeg\n",
        "\n",
        "\n",
        "init_frame =  1#@param  {type: \"number\"}\n",
        "last_frame =  999#@param {type: \"number\"}\n",
        "\n",
        "min_fps =  120#@param {type: \"number\"}\n",
        "max_fps =  120#@param {type: \"number\"}\n",
        "\n",
        "total_frames = last_frame-init_frame\n",
        "\n",
        "# Desired video runtime in seconds\n",
        "length = 3 #@param {type: \"number\"}\n",
        "\n",
        "frames = []\n",
        "tqdm.write('Generating video...')\n",
        "\n",
        "for i in range(init_frame,last_frame): #\n",
        "    if usingDiffusion == False:\n",
        "      filename = f\"{abs_root_path}/vqgan-steps/{i:04}.png\"\n",
        "    elif usingDiffusion == True:\n",
        "      filename = f\"{abs_root_path}/diffusion-steps/{i:05}.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\" #@param {type: \"string\"}\n",
        "\n",
        "\n",
        "video_filename = listToString(video_filename).replace(\" \",\"_\")\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",
        "\n",
        "\n",
        "print(\"Video ready.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y0e8pHyJmi7s",
        "cellView": "form"
      },
      "source": [
        "# @title Download video\n",
        "from google.colab import files\n",
        "files.download(video_filename)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E8lvN6b0mb-b",
        "cellView": "form"
      },
      "source": [
        "# @title View video in browser\n",
        "#@markdown \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": []
    }
  ]
}