Text2Video.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/pschaldenbrand/Text2Video/blob/main/Text2Video.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e7DCzf-EHzL_"
      },
      "source": [
        "# Towards Real-Time Text2Video via Model-Free, CLIP-Guided Pixel Optimization\n",
        "\n",
        "## Peter Schaldenbrand, Zhixuan Liu, Jean Oh\n",
        "\n",
        "The Robotics Institute, Carnegie Mellon University\n",
        "\n",
        "Questions and comments to Peter (pschalde at andrew dot cmu dot edu)\n",
        "\n",
        "### Directions\n",
        "1. Run the \"Pre Installation\" cell\n",
        "3. Edit the list of prompts in the last few cells to generate your animation\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1IlcVZ2V1xM0"
      },
      "source": [
        "# Pre Installation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "5dyyH781qzIC"
      },
      "outputs": [],
      "source": [
        "#@title Pre Installation {vertical-output: true}\n",
        "%cd /content/\n",
        "\n",
        "!pip install ftfy regex tqdm pytorch-lightning omegaconf                 &> /dev/null\n",
        "!apt install exempi                                                      &> /dev/null\n",
        "!pip install git+https://github.com/openai/CLIP.git --no-deps            &> /dev/null"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "hjt9T3ARukAg"
      },
      "outputs": [],
      "source": [
        "#@title Imports and Notebook Utilities {vertical-output: true}\n",
        "import os\n",
        "import io\n",
        "import PIL.Image, PIL.ImageDraw\n",
        "import base64\n",
        "import zipfile\n",
        "import json\n",
        "import requests\n",
        "import numpy as np\n",
        "import matplotlib.pylab as pl\n",
        "import glob\n",
        "from datetime import datetime\n",
        "from tqdm import tqdm_notebook as tqdm\n",
        "from IPython.display import Image, HTML, clear_output\n",
        "from tqdm import tqdm_notebook, tnrange\n",
        "from omegaconf import OmegaConf\n",
        "\n",
        "import torch\n",
        "import skimage\n",
        "import skimage.io\n",
        "import random\n",
        "import argparse\n",
        "import math\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "import requests\n",
        "from io import BytesIO\n",
        "\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import PIL\n",
        "from time import time\n",
        "\n",
        "device = torch.device('cuda')\n",
        "\n",
        "def imread(url, max_size=None, mode=None):\n",
        "  if url.startswith(('http:', 'https:')):\n",
        "    r = requests.get(url)\n",
        "    f = io.BytesIO(r.content)\n",
        "  else:\n",
        "    f = url\n",
        "  img = PIL.Image.open(f)\n",
        "  if max_size is not None:\n",
        "    img = img.resize((max_size, max_size))\n",
        "  if mode is not None:\n",
        "    img = img.convert(mode)\n",
        "  img = np.float32(img)/255.0\n",
        "  return img\n",
        "\n",
        "def np2pil(a):\n",
        "  if a.dtype in [np.float32, np.float64]:\n",
        "    a = np.uint8(np.clip(a, 0, 1)*255)\n",
        "  return PIL.Image.fromarray(a)\n",
        "\n",
        "def imwrite(f, a, fmt=None):\n",
        "  a = np.asarray(a)\n",
        "  if isinstance(f, str):\n",
        "    fmt = f.rsplit('.', 1)[-1].lower()\n",
        "    if fmt == 'jpg':\n",
        "      fmt = 'jpeg'\n",
        "    f = open(f, 'wb')\n",
        "  np2pil(a).save(f, fmt, quality=95)\n",
        "\n",
        "def imencode(a, fmt='jpeg'):\n",
        "  a = np.asarray(a)\n",
        "  if len(a.shape) == 3 and a.shape[-1] == 4:\n",
        "    fmt = 'png'\n",
        "  f = io.BytesIO()\n",
        "  imwrite(f, a, fmt)\n",
        "  return f.getvalue()\n",
        "\n",
        "def im2url(a, fmt='jpeg'):\n",
        "  encoded = imencode(a, fmt)\n",
        "  base64_byte_string = base64.b64encode(encoded).decode('ascii')\n",
        "  return 'data:image/' + fmt.upper() + ';base64,' + base64_byte_string\n",
        "\n",
        "def imshow(a, fmt='jpeg'):\n",
        "  display(Image(data=imencode(a, fmt)))\n",
        "\n",
        "\n",
        "def tile2d(a, w=None):\n",
        "  a = np.asarray(a)\n",
        "  if w is None:\n",
        "    w = int(np.ceil(np.sqrt(len(a))))\n",
        "  th, tw = a.shape[1:3]\n",
        "  pad = (w-len(a))%w\n",
        "  a = np.pad(a, [(0, pad)]+[(0, 0)]*(a.ndim-1), 'constant')\n",
        "  h = len(a)//w\n",
        "  a = a.reshape([h, w]+list(a.shape[1:]))\n",
        "  a = np.rollaxis(a, 2, 1).reshape([th*h, tw*w]+list(a.shape[4:]))\n",
        "  return a\n",
        "\n",
        "from torchvision import utils\n",
        "def show_img(img):\n",
        "    img = np.transpose(img, (1, 2, 0))\n",
        "    img = np.clip(img, 0, 1)\n",
        "    img = np.uint8(img * 254)\n",
        "    # img = np.repeat(img, 4, axis=0)\n",
        "    # img = np.repeat(img, 4, axis=1)\n",
        "    pimg = PIL.Image.fromarray(img, mode=\"RGB\")\n",
        "    imshow(pimg)\n",
        "\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "import os\n",
        "\n",
        "# torch.set_default_tensor_type('torch.cuda.FloatTensor')\n",
        "\n",
        "print(\"Torch version:\", torch.__version__)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Z-Wt7UjTi8Le"
      },
      "outputs": [],
      "source": [
        "#@title Load CLIP {vertical-output: true}\n",
        "\n",
        "# os.environ['CUDA_LAUNCH_BLOCKING'] = '1'\n",
        "\n",
        "import os\n",
        "import clip\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "import torchvision\n",
        "from torchvision import transforms\n",
        "\n",
        "# Load the model\n",
        "device = torch.device('cuda')\n",
        "model, preprocess = clip.load('ViT-B/32', device, jit=False)\n",
        "model_16, preprocess_16 = clip.load('ViT-B/16', device, jit=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mmGSyESqOT_y",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Some Functions {vertical-output: true}\n",
        "def pil_resize_long_edge_to(pil, trg_size):\n",
        "  short_w = pil.width < pil.height\n",
        "  ar_resized_long = (trg_size / pil.height) if short_w else (trg_size / pil.width)\n",
        "  resized = pil.resize((int(pil.width * ar_resized_long), int(pil.height * ar_resized_long)), PIL.Image.BICUBIC)\n",
        "  return resized\n",
        "\n",
        "def draw_text_on_image(img, text):\n",
        "    img = img.transpose((1,2,0))\n",
        "    img = PIL.Image.fromarray((img*255.).astype('uint8'), 'RGB')\n",
        "    \n",
        "    # Call draw Method to add 2D graphics in an image\n",
        "    I1 = PIL.ImageDraw.Draw(img)\n",
        "    font = PIL.ImageFont.truetype(r'/usr/share/fonts/truetype/humor-sans/Humor-Sans.ttf', 17) \n",
        "  \n",
        "    \n",
        "    # Add Text to an image\n",
        "    I1.text((5, 5), text, fill=(255, 255, 255), font=font)\n",
        "    \n",
        "    # Display edited image\n",
        "    # img.show()\n",
        "    \n",
        "    return np.array(img).transpose(2,0,1)/255.\n",
        "\n",
        "def get_image_augmentation(use_normalized_clip):\n",
        "    # augment_trans = transforms.Compose([\n",
        "    #     transforms.RandomPerspective(fill=1, p=1, distortion_scale=0.5),\n",
        "    #     transforms.RandomResizedCrop(224, scale=(0.7,0.9)),\n",
        "    # ])\n",
        "\n",
        "    # if use_normalized_clip:\n",
        "    #     augment_trans = transforms.Compose([\n",
        "    #     transforms.RandomPerspective(fill=1, p=1, distortion_scale=0.5),\n",
        "    #     transforms.RandomResizedCrop(224, scale=(0.7,0.9)),\n",
        "    #     # transforms.GaussianBlur((3,3)),\n",
        "    #     transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))\n",
        "    # ])\n",
        "    augment_trans = transforms.Compose([\n",
        "        # transforms.Resize(224),\n",
        "        transforms.Pad(20,40),\n",
        "        transforms.RandomPerspective(fill=1, p=1, distortion_scale=0.5),\n",
        "        transforms.RandomResizedCrop(224, scale=(0.7,0.999)),\n",
        "        transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))\n",
        "    ])\n",
        "    augment_trans_style = transforms.Compose([\n",
        "        transforms.Resize(256)\n",
        "    ])\n",
        "    \n",
        "    augment_change_clip = transforms.Compose([\n",
        "        transforms.Resize((224, 224)),\n",
        "        transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))\n",
        "    ]) \n",
        "    return augment_trans, augment_trans_style, augment_change_clip\n",
        "\n",
        "# Tensor and PIL utils\n",
        "\n",
        "def pil_loader(path):\n",
        "    with open(path, 'rb') as f:\n",
        "        img = PIL.Image.open(f)\n",
        "        return img.convert('RGB')\n",
        "\n",
        "def pil_loader_internet(url):\n",
        "    response = requests.get(url)\n",
        "    img = PIL.Image.open(BytesIO(response.content))\n",
        "    return img.convert('RGB')\n",
        "\n",
        "def tensor_resample(tensor, dst_size, mode='bilinear'):\n",
        "    return F.interpolate(tensor, dst_size, mode=mode, align_corners=False)\n",
        "\n",
        "def pil_resize_short_edge_to(pil, trg_size):\n",
        "    short_w = pil.width < pil.height\n",
        "    ar_resized_short = (trg_size / pil.width) if short_w else (trg_size / pil.height)\n",
        "    resized = pil.resize((int(pil.width * ar_resized_short), int(pil.height * ar_resized_short)), PIL.Image.BICUBIC)\n",
        "    return resized\n",
        "\n",
        "def pil_resize_long_edge_to(pil, trg_size):\n",
        "    short_w = pil.width < pil.height\n",
        "    ar_resized_long = (trg_size / pil.height) if short_w else (trg_size / pil.width)\n",
        "    resized = pil.resize((int(pil.width * ar_resized_long), int(pil.height * ar_resized_long)), PIL.Image.BICUBIC)\n",
        "    return resized\n",
        "\n",
        "def np_to_pil(npy):\n",
        "    return PIL.Image.fromarray(npy.astype(np.uint8))\n",
        "\n",
        "def pil_to_np(pil):\n",
        "    return np.array(pil)\n",
        "\n",
        "def tensor_to_np(tensor, cut_dim_to_3=True):\n",
        "    if len(tensor.shape) == 4:\n",
        "        if cut_dim_to_3:\n",
        "            tensor = tensor[0]\n",
        "        else:\n",
        "            return tensor.data.cpu().numpy().transpose((0, 2, 3, 1))\n",
        "    return tensor.data.cpu().numpy().transpose((1,2,0))\n",
        "\n",
        "def np_to_tensor(npy, space):\n",
        "    if space == 'vgg':\n",
        "        return np_to_tensor_correct(npy)\n",
        "    return (torch.Tensor(npy.astype(np.float) / 127.5) - 1.0).permute((2,0,1)).unsqueeze(0)\n",
        "\n",
        "def np_to_tensor_correct(npy):\n",
        "    pil = np_to_pil(npy)\n",
        "    transform = transforms.Compose([transforms.ToTensor(),\n",
        "        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
        "    ])\n",
        "    return transform(pil).unsqueeze(0)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Rsy1J5hZiqN9"
      },
      "outputs": [],
      "source": [
        "#@title more functions (style, gif)\n",
        "\n",
        "def load_init_canvas(path_or_url, width):\n",
        "    t = pil_loader(path_or_url) if os.path.exists(path_or_url) else pil_loader_internet(path_or_url)\n",
        "    canvas = np_to_tensor(pil_to_np(pil_resize_long_edge_to(t, width)), \"normal\").to(device)\n",
        "    canvas -= canvas.min()\n",
        "    canvas /= canvas.max()\n",
        "    canvas.requires_grad=True\n",
        "    return canvas\n",
        "\n",
        "def to_gif(canvases, fn='/animation.gif', duration=250):\n",
        "    #imgs = [PIL.Image.fromarray((img.transpose((1,2,0))*255.).astype(np.uint8)) for img in canvases]\n",
        "    imgs = []\n",
        "    for i in range(len(canvases)):\n",
        "      if True:\n",
        "          np_img = (np.clip(canvases[i], 0, 1).transpose((1,2,0))*255.).astype(np.uint8)\n",
        "\n",
        "          imgs.append(PIL.Image.fromarray(np_img))\n",
        "    # duration is the number of milliseconds between frames; this is 40 frames per second\n",
        "    # imgs[0].save(fn, save_all=True, append_images=imgs[1:], duration=50, loop=0)\n",
        "    imgs[0].save(fn, save_all=True, append_images=imgs[1:], duration=duration, loop=0)\n",
        "    \n",
        "import cv2\n",
        "def to_video(frames, fn=None, frame_rate=4):\n",
        "    #if fn is None: fn = '/content/drive/MyDrive/animations/{}.mp4'.format(time())\n",
        "    if fn is None: \n",
        "        import datetime\n",
        "        date_and_time = datetime.datetime.now()\n",
        "        run_name = '' + date_and_time.strftime(\"%m_%d__%H_%M_%S\")\n",
        "        fn = '/content/{}.mp4'.format(run_name)\n",
        "    h, w = frames[0].shape[1], frames[0].shape[2]\n",
        "    print(h,w)\n",
        "    _fourcc = cv2.VideoWriter_fourcc(*'MP4V')\n",
        "    # _fourcc = cv2.VideoWriter_fourcc(*'H264')\n",
        "    # _fourcc = cv2.VideoWriter_fourcc(*'XVID')\n",
        "    out = cv2.VideoWriter(fn, _fourcc, frame_rate, (w,h))\n",
        "    for frame in frames:\n",
        "        cv2_frame = np.clip(frame, a_min=0, a_max=1)\n",
        "        cv2_frame = (cv2_frame * 255.).astype(np.uint8).transpose((1,2,0))[:,:,::-1]\n",
        "        out.write(cv2_frame)\n",
        "    out.release()\n",
        "    return fn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "7wTGRvauVaFq"
      },
      "outputs": [],
      "source": [
        "#@title CycleGAN\n",
        "if not os.path.exists('/content/pytorch-CycleGAN-and-pix2pix'):\n",
        "    !git clone https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix.git &> /dev/null\n",
        "    os.chdir('/content/pytorch-CycleGAN-and-pix2pix/')\n",
        "    !pip install -r requirements.txt   &> /dev/null\n",
        "else:\n",
        "    os.chdir('/content/pytorch-CycleGAN-and-pix2pix/')\n",
        "\n",
        "\n",
        "parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n",
        "\n",
        "parser.add_argument('--dataroot', default='/content/cyclegan_dataset', help='path to images (should have subfolders trainA, trainB, valA, valB, etc)')\n",
        "parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment. It decides where to store samples and models')\n",
        "parser.add_argument('--use_wandb', action='store_true', help='use wandb')\n",
        "parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')\n",
        "parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')\n",
        "# parser.add_argument('--checkpoints_dir', type=str, default=cyclegan_models_dir, help='models are saved here')\n",
        "# model parameters\n",
        "parser.add_argument('--model', type=str, default='cycle_gan', help='chooses which model to use. [cycle_gan | pix2pix | test | colorization]')\n",
        "parser.add_argument('--input_nc', type=int, default=3, help='# of input image channels: 3 for RGB and 1 for grayscale')\n",
        "parser.add_argument('--output_nc', type=int, default=3, help='# of output image channels: 3 for RGB and 1 for grayscale')\n",
        "parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in the last conv layer')\n",
        "parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in the first conv layer')\n",
        "parser.add_argument('--netD', type=str, default='basic', help='specify discriminator architecture [basic | n_layers | pixel]. The basic model is a 70x70 PatchGAN. n_layers allows you to specify the layers in the discriminator')\n",
        "parser.add_argument('--netG', type=str, default='resnet_9blocks', help='specify generator architecture [resnet_9blocks | resnet_6blocks | unet_256 | unet_128]')\n",
        "parser.add_argument('--n_layers_D', type=int, default=3, help='only used if netD==n_layers')\n",
        "parser.add_argument('--norm', type=str, default='instance', help='instance normalization or batch normalization [instance | batch | none]')\n",
        "parser.add_argument('--init_type', type=str, default='normal', help='network initialization [normal | xavier | kaiming | orthogonal]')\n",
        "parser.add_argument('--init_gain', type=float, default=0.02, help='scaling factor for normal, xavier and orthogonal.')\n",
        "parser.add_argument('--no_dropout', action='store_true', help='no dropout for the generator')\n",
        "# dataset parameters\n",
        "parser.add_argument('--dataset_mode', type=str, default='unaligned', help='chooses how datasets are loaded. [unaligned | aligned | single | colorization]')\n",
        "parser.add_argument('--direction', type=str, default='AtoB', help='AtoB or BtoA')\n",
        "parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')\n",
        "parser.add_argument('--num_threads', default=1, type=int, help='# threads for loading data')\n",
        "# parser.add_argument('--batch_size', type=int, default=1, help='input batch size')\n",
        "parser.add_argument('--batch_size', type=int, default=4, help='input batch size')\n",
        "# parser.add_argument('--load_size', type=int, default=286, help='scale images to this size')\n",
        "# parser.add_argument('--crop_size', type=int, default=256, help='then crop to this size')\n",
        "parser.add_argument('--load_size', type=int, default=512, help='scale images to this size')\n",
        "# parser.add_argument('--crop_size', type=int, default=256, help='then crop to this size')\n",
        "parser.add_argument('--crop_size', type=int, default=128, help='then crop to this size')\n",
        "parser.add_argument('--max_dataset_size', type=int, default=float(\"inf\"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')\n",
        "parser.add_argument('--preprocess', type=str, default='resize_and_crop', help='scaling and cropping of images at load time [resize_and_crop | crop | scale_width | scale_width_and_crop | none]')\n",
        "parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation')\n",
        "parser.add_argument('--display_winsize', type=int, default=256, help='display window size for both visdom and HTML')\n",
        "# additional parameters\n",
        "parser.add_argument('--epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')\n",
        "parser.add_argument('--load_iter', type=int, default='0', help='which iteration to load? if load_iter > 0, the code will load models by iter_[load_iter]; otherwise, the code will load models by [epoch]')\n",
        "parser.add_argument('--verbose', action='store_true', help='if specified, print more debugging information')\n",
        "parser.add_argument('--suffix', default='', type=str, help='customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{load_size}')\n",
        "parser.add_argument('--display_freq', type=int, default=800, help='frequency of showing training results on screen')\n",
        "parser.add_argument('--display_ncols', type=int, default=4, help='if positive, display all images in a single visdom web panel with certain number of images per row.')\n",
        "parser.add_argument('--display_id', type=int, default=1, help='window id of the web display')\n",
        "parser.add_argument('--display_server', type=str, default=\"http://localhost\", help='visdom server of the web display')\n",
        "parser.add_argument('--display_env', type=str, default='main', help='visdom display environment name (default is \"main\")')\n",
        "parser.add_argument('--display_port', type=int, default=8097, help='visdom port of the web display')\n",
        "parser.add_argument('--update_html_freq', type=int, default=1000, help='frequency of saving training results to html')\n",
        "parser.add_argument('--print_freq', type=int, default=100, help='frequency of showing training results on console')\n",
        "parser.add_argument('--no_html', action='store_true', help='do not save intermediate training results to [opt.checkpoints_dir]/[opt.name]/web/')\n",
        "# network saving and loading parameters\n",
        "parser.add_argument('--save_latest_freq', type=int, default=5000, help='frequency of saving the latest results')\n",
        "parser.add_argument('--save_epoch_freq', type=int, default=5, help='frequency of saving checkpoints at the end of epochs')\n",
        "parser.add_argument('--save_by_iter', action='store_true', help='whether saves model by iteration')\n",
        "parser.add_argument('--continue_train', action='store_true', help='continue training: load the latest model')\n",
        "parser.add_argument('--epoch_count', type=int, default=1, help='the starting epoch count, we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>, ...')\n",
        "parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')\n",
        "# training parameters\n",
        "parser.add_argument('--n_epochs', type=int, default=100, help='number of epochs with the initial learning rate')\n",
        "parser.add_argument('--n_epochs_decay', type=int, default=100, help='number of epochs to linearly decay learning rate to zero')\n",
        "parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam')\n",
        "parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')\n",
        "parser.add_argument('--gan_mode', type=str, default='lsgan', help='the type of GAN objective. [vanilla| lsgan | wgangp]. vanilla GAN loss is the cross-entropy objective used in the original GAN paper.')\n",
        "parser.add_argument('--pool_size', type=int, default=50, help='the size of image buffer that stores previously generated images')\n",
        "parser.add_argument('--lr_policy', type=str, default='linear', help='learning rate policy. [linear | step | plateau | cosine]')\n",
        "parser.add_argument('--lr_decay_iters', type=int, default=50, help='multiply by a gamma every lr_decay_iters iterations')\n",
        "\n",
        "parser.add_argument('--lambda_A', type=float, default=10.0, help='weight for cycle loss (A -> B -> A)')\n",
        "parser.add_argument('--lambda_B', type=float, default=10.0, help='weight for cycle loss (B -> A -> B)')\n",
        "parser.add_argument('--lambda_identity', type=float, default=0.5, help='use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1')\n",
        "\n",
        "\n",
        "opt, _ = parser.parse_known_args()\n",
        "opt.isTrain = False\n",
        "from models import create_model\n",
        "# from data import create_dataset\n",
        "\n",
        "str_ids = opt.gpu_ids.split(',')\n",
        "opt.gpu_ids = []\n",
        "for str_id in str_ids:\n",
        "    id = int(str_id)\n",
        "    if id >= 0:\n",
        "        opt.gpu_ids.append(id)\n",
        "if len(opt.gpu_ids) > 0:\n",
        "    torch.cuda.set_device(opt.gpu_ids[0])\n",
        "    \n",
        "opt.epoch = 'beautiful' \n",
        "\n",
        "cyclegan_model = create_model(opt)\n",
        "\n",
        "opt.ngf, opt.ndf = 64, 64\n",
        "for name in cyclegan_model.model_names:\n",
        "    if isinstance(name, str):\n",
        "        save_filename = '%s_net_%s.pth' % (opt.epoch, name)\n",
        "        save_path = '/content/pytorch-CycleGAN-and-pix2pix/cyclegan_generator.pth'\n",
        "        if name == 'G_A':\n",
        "            net = getattr(cyclegan_model, 'net' + name)\n",
        "            if not os.path.exists(save_path):\n",
        "                #!wget -O cyclegan_generator.pth https://dl.boxcloud.com/d/1/b1!NNM62Iu7CAUniqipvBqxuFBwgzqaXUC50EIiSclFKWmfIgKp-LtxUp_YUqiz27LnKk005brY-D7GQjDD6R2JDkDB-jnMJ-50QMdn__dRkNeQCIpT8BXyawLJyNhv0vab0Mc2PjKjPSb2jASASijPqEDJ1Rq09ff2MikKVBjekuaGt_c1k23gMlEr2YkRTsvOUYv1k6HSfUApa_3ZCJbzRqbzUEOMOl_2phJv0orv36UF4dr0lojNsVp8fcVrFkyaZcEfGJI4eBPAxM18UnZ_7ellnX7GpLAA709rczmOBp23xrqIR7feCs-JrsBnMA7n1HOkaLQhH80plhmYt2fViRhGqdYrO4D9SQVIKWHTNdy5eiofKGNkCqoVD0yUHB-oC_FsFkJ2MtzMNQlg_jwTYhSUylbj7OPxFFZfHt9Z_4oeh2x1S8F6AFtmalgDnQYwPEGVIQMNhBr-MZEVEur9yHFOBF2NnYXrthHdNuV36sIIuxTJg_YmXlLOc5Wr4Ac9VyrIiOVsoIRT7uK4Rm3jSRvy6-0BoFN85omCP03LFAxrdzcFmRcS6D1j534-HYDoLYiLHHXY8-gf5xRPvFh9TSEM4-PMzilIWxAJOvBa_-oJ4JsNiUZbZOejCHiJL459Zgc7IU_MCcAvWtoNm3n06HjM9_MzxbGKOM_0xQ1kvr0i4rtqLNL6-NKmUBMTDe5Hgk8QAZ2dX3NamzQCI0eesCjOIDEXEB_wJHrZXSG3alh3qcqRrLUa2tBeLA-4o3SIZMg4H99GL76_g9ggb4gvEdcPnrF5TzU9el5enf2QjahxY9BmpvoAeBgTsbHm3RmGr8lSRR21K3WGsRkA0lfr2uhZnYk5NSSRzlW7FquvUV2hQLNRGaLmFXFL2i9rTQlA4dtOhyR4alazzboLTx5IGnwvw3yqcFu5rbGNPaQqpWQInV9Uf6hhpSEdB1bO_G_h0cegNJz-JWO9Nvn_wL-PDZ2vUZwopLzA09WG_9qpgTGf0ogJl345thnno1ydXr58Y-ZGu44a3lkLC0HtBRniz35AIeT9DymTzfSGYkaZ4pkxWbLviy7YEkfmHWyg_f0kZm7DSQqJPq4x5-cqouBwTvXMZYEFLRnwflfmXnHSGMO1UrjMWfRAGxsg1ncXE7nLHPKQaoMVsx3gu_0cGkWxr1_3T50vYK1j1rz07u7xlKTmnv59lhKGdoUdFZAEQG4myMcedzZF31B5A4AgkM6Y7knLOz8mORKH_R41IdPhAtJyVVsNev0daSywTHkHMqCIOkA9aoVGYXU8iSPiiIc0XzF1FQaN38jiE23Z6ppM4nNUgoeiAghKet5t-JolWJ3u0UG0VNJuA7pubik_zxAHkMh34G1Bd3k--hfc-Sfg_lcqcu5KXGBQdpKb5HYyrGVQpjGvNXoruQtRPA../download\n",
        "                !curl -L   https://cmu.box.com/shared/static/lxgfgfw9aqia8crfyg5jw0h1gkdv64mk --output cyclegan_generator.pth\n",
        "            net.load_state_dict(torch.load(save_path))\n",
        "\n",
        "os.chdir('/content')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "EVsPjRqxzvvn"
      },
      "outputs": [],
      "source": [
        "#@title Load Models\n",
        "import sys\n",
        "\n",
        "def load_generator_model(model_type, n=1, ngf=64, h=None, w=None, pretrained_model=None):\n",
        "\n",
        "    generate = lambda g, z : g(z)\n",
        "\n",
        "    if model_type == 'cyclegan':\n",
        "        z = torch.rand((n, 3, h, w), device=device)\n",
        "        gen = cyclegan_model.netG_A\n",
        "        generate = lambda g, z : (g(z)+ 1)/2\n",
        "    else:\n",
        "        class Nothin(nn.Module):\n",
        "            def __init__(self):\n",
        "                super(Nothin, self).__init__()\n",
        "            def forward(self, z):\n",
        "                return z\n",
        "        gen = Nothin()\n",
        "        z = torch.rand((n, 3, h, w), device=device)\n",
        "        generate = lambda g, z : g(z)\n",
        "\n",
        "    # for param in gen.parameters():\n",
        "    #     param.requires_grad = True\n",
        "    # gen.train()\n",
        "       \n",
        "    z.requires_grad = True\n",
        "    \n",
        "    for param in gen.parameters():\n",
        "        param.requires_grad = False\n",
        "    gen.eval()\n",
        "\n",
        "    return gen, z, generate\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "5kPSwRLeuUeB"
      },
      "outputs": [],
      "source": [
        "#@title Relaxed Earth Mover's Distance\n",
        "# From : https://github.com/futscdav/strotss/blob/master/strotss.py\n",
        "def pairwise_distances_cos(x, y):\n",
        "    x_norm = torch.sqrt((x**2).sum(1).view(-1, 1))\n",
        "    y_t = torch.transpose(y, 0, 1)\n",
        "    y_norm = torch.sqrt((y**2).sum(1).view(1, -1))\n",
        "    dist = 1.-torch.mm(x, y_t)/x_norm/y_norm\n",
        "    return dist\n",
        "\n",
        "def pairwise_distances_sq_l2(x, y):\n",
        "    x_norm = (x**2).sum(1).view(-1, 1)\n",
        "    y_t = torch.transpose(y, 0, 1)\n",
        "    y_norm = (y**2).sum(1).view(1, -1)\n",
        "    dist = x_norm + y_norm - 2.0 * torch.mm(x, y_t)\n",
        "    return torch.clamp(dist, 1e-5, 1e5)/x.size(1)\n",
        "def distmat(x, y, cos_d=True):\n",
        "    if cos_d:\n",
        "        M = pairwise_distances_cos(x, y)\n",
        "    else:\n",
        "        M = torch.sqrt(pairwise_distances_sq_l2(x, y))\n",
        "    return M\n",
        "def EMD(X, Y):\n",
        "    CX_M = distmat(X, Y, cos_d=True)\n",
        "\n",
        "    # if d==3: CX_M = CX_M + distmat(X, Y, cos_d=False)\n",
        "\n",
        "    m1, m1_inds = CX_M.min(1)\n",
        "    m2, m2_inds = CX_M.min(0)\n",
        "\n",
        "    remd = torch.max(m1.mean(), m2.mean())\n",
        "\n",
        "    return remd"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L6_I95FOqQ4M",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Generate Video Code Definition\n",
        "\n",
        "# Image Augmentation Transformation\n",
        "augment_trans, augment_trans_style, augment_change_clip = get_image_augmentation(True)\n",
        "\n",
        "def alter_z_noise(z, squish=4, noise_std=1.):\n",
        "    # Alter the params so the next image isn't exactly like the previous.\n",
        "    with torch.no_grad():\n",
        "        z /= squish\n",
        "        z += torch.randn(z.shape).to(device) * noise_std\n",
        "    return z\n",
        "\n",
        "def generate_video( prompts, # List of text prompts to use to generate media\n",
        "                    h=9*40,w=16*40,\n",
        "                    lr=.1,\n",
        "                    num_augs=4, \n",
        "                    model_type='cyclegan',   \n",
        "                    debug=True, display_prompt=True,\n",
        "                    frames_per_prompt=10, # Number of frames to dedicate to each prompt\n",
        "                    first_iter=300, # Number of optimization iterations for first first frame\n",
        "                    num_iter=50, # Optimization iterations for all but first frame\n",
        "                    z_unchanging_weight=3, # Weight to ensure z does not change at all * l1_loss(z, z_prev)\n",
        "                    z_noise_squish=4., # Amount to squish z by between frames\n",
        "                    carry_over_iter=17, # Which iteration of optimization to use as the start of the next frame\n",
        "                    encoding_comparison='cosine', # or \"emd\"\n",
        "                    n_samples=1):\n",
        "    \n",
        "    start_time, all_canvases = time(), []\n",
        "    \n",
        "    gen, z_for_next_frame, generate = load_generator_model(model_type, n=n_samples, ngf=666, h=h, w=w, pretrained_model=None)\n",
        "\n",
        "    # Optimizers\n",
        "    #optim, style_optim, z_optim = torch.optim.Adam([z], lr=lr), torch.optim.RMSprop([z], lr=lr), torch.optim.Adam([z], lr=lr)\n",
        "\n",
        "    content_loss, z_loss, styleloss_tot = 0, 0, 0\n",
        "    prev_z = None\n",
        "    image_features, image_features_16 = None, None\n",
        "    total_chunks = (len(prompts)-1) * 2*frames_per_prompt + frames_per_prompt\n",
        "    pbar = tqdm(total=total_chunks)\n",
        "\n",
        "    cosine_dist = lambda a, b: -1 * torch.cosine_similarity(a, b, dim=1)\n",
        "    encoding_compare = cosine_dist if encoding_comparison == 'cosine' else EMD \n",
        "    l1_loss = nn.L1Loss()\n",
        "\n",
        "    neg_prompt = \"Words and text.\"\n",
        "\n",
        "\n",
        "    for prompt_ind in range(len(prompts)):\n",
        "        prompt_now  = prompts[prompt_ind]\n",
        "        prompt_next = prompts[prompt_ind+1] if prompt_ind < len(prompts)-1 else None\n",
        "\n",
        "        with torch.no_grad():\n",
        "            text_features_now  = model.encode_text(clip.tokenize(prompt_now).to(device))\n",
        "            text_features_next = model.encode_text(clip.tokenize(prompt_next).to(device)) if prompt_next is not None else None\n",
        "            text_features_now_16  = model_16.encode_text(clip.tokenize(prompt_now).to(device))\n",
        "            text_features_next_16 = model_16.encode_text(clip.tokenize(prompt_next).to(device)) if prompt_next is not None else None\n",
        "            neg_text_features = model.encode_text(clip.tokenize(neg_prompt).to(device))\n",
        "            neg_text_features_16 = model_16.encode_text(clip.tokenize(neg_prompt).to(device))\n",
        "\n",
        "        tot_frames = frames_per_prompt*2 if prompt_ind < len(prompts)-1 else frames_per_prompt\n",
        "        for frame in range(tot_frames):\n",
        "            # Assign a weight to the current and next prompts\n",
        "            weight_now = 1 - (frame/(tot_frames))\n",
        "            weight_next = frame/(tot_frames)\n",
        "            if prompt_ind == (len(prompts) - 1): weight_now = 1.\n",
        "\n",
        "            # Alter the params so the next image isn't exactly like the previous.\n",
        "            z = alter_z_noise(z_for_next_frame, squish=z_noise_squish, noise_std=1.)\n",
        "            z.requires_grad = True\n",
        "\n",
        "            # Optimizers\n",
        "            optim, style_optim, z_optim = torch.optim.Adam([z], lr=lr), torch.optim.RMSprop([z], lr=lr), torch.optim.Adam([z], lr=lr)\n",
        "            \n",
        "            # Save features from previous frame\n",
        "            prev_image_features = image_features.detach() if image_features is not None else None\n",
        "            prev_image_features_16 = image_features_16.detach() if image_features_16 is not None else None\n",
        "\n",
        "            # Run the main optimization loop\n",
        "            iterations = first_iter if (prompt_ind==0 and frame==0) else num_iter\n",
        "            for t in range(iterations):\n",
        "\n",
        "                ''' Loss that just operates on z '''\n",
        "                ex_freq = 2 # Alternate between two clip models for robustness\n",
        "                z_optim.zero_grad()\n",
        "                loss = 0\n",
        "                im_batch = torch.cat([augment_trans(z) for n in range(num_augs)])\n",
        "                if t % ex_freq == 0:\n",
        "                    image_features_16 = model_16.encode_image(im_batch) \n",
        "                else:\n",
        "                    image_features = model.encode_image(im_batch)\n",
        "                for n in range(num_augs):\n",
        "                    # loss for clip features of z and text features (This and next prompt)\n",
        "                    if t % ex_freq == 0:\n",
        "                        loss += encoding_compare(text_features_now_16, image_features_16[n:n+1]) * weight_now\n",
        "                        loss -= encoding_compare(neg_text_features_16, image_features_16[n:n+1]) * weight_now\n",
        "                        if text_features_next_16 is not None: loss += encoding_compare(text_features_next_16, image_features_16[n:n+1]) * weight_next\n",
        "                    else:\n",
        "                        loss += encoding_compare(text_features_now, image_features[n:n+1]) * weight_now\n",
        "                        loss -= encoding_compare(neg_text_features, image_features[n:n+1]) * weight_now\n",
        "                        if text_features_next is not None: loss += encoding_compare(text_features_next, image_features[n:n+1]) * weight_next\n",
        "                    if prev_image_features is not None: \n",
        "                        # Loss to make sure that z doesn't change much\n",
        "                        if t % 4 == 0:\n",
        "                            loss += l1_loss(z, prev_z) * z_unchanging_weight\n",
        "                \n",
        "                loss.backward()\n",
        "                z_loss = loss.item()\n",
        "                z_optim.step()\n",
        "\n",
        "                if t == carry_over_iter-1:\n",
        "                    z_for_next_frame = z.detach().clone()\n",
        "\n",
        "            prev_z = z.detach().clone()\n",
        "            pbar.update(1)\n",
        "            gen.eval()\n",
        "            \n",
        "            with torch.no_grad():\n",
        "                if model_type=='cyclegan':\n",
        "                    z_norm = z.detach().clone()\n",
        "                    img = generate(gen, z_norm).detach().cpu().numpy()[0] \n",
        "                    # show_img(z.detach().cpu().numpy()[0])\n",
        "                else:\n",
        "                    img = generate(gen, z).detach().cpu().numpy()[0]\n",
        "                if display_prompt:\n",
        "                    img = draw_text_on_image(img, prompt_now)\n",
        "                all_canvases.append(img)\n",
        "                if frame % 4 == 0: print('Frame: ', len(all_canvases)), show_img(img)\n",
        "                \n",
        "    to_gif(all_canvases, fn='/animation.gif')\n",
        "    # from IPython.display import Image, display\n",
        "    # ipython_img = Image(open('/animation.gif','rb').read())\n",
        "    # display(ipython_img)\n",
        "    \n",
        "    # to_gif(all_canvases, fn='/content/drive/MyDrive/animations/{}.gif'.format(time()))\n",
        "    to_video(all_canvases, frame_rate=3)\n",
        "    fn = to_video(all_canvases, frame_rate=8)\n",
        "    return all_canvases, fn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3BTp7LfYFuIk",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title generate_video_wrapper\n",
        "def generate_video_wrapper(prompts, frames_per_prompt=10, h=9*40,w=16*40, fast=False,\n",
        "                           style_opt_iter=0, temperature=50, download=True, display_prompt=True):\n",
        "    \n",
        "    lr = .17 if fast else .1\n",
        "    num_iter = 10 if fast else 25\n",
        "    carry_over_iter = 9 if fast else 13\n",
        "    temperature = 0.5 * temperature if fast else temperature\n",
        "\n",
        "    z_unchanging_weight = 4 - (temperature/100) * 4\n",
        "    z_noise_squish = (temperature/100) * 4 + 2\n",
        "\n",
        "    all_canvases, fn = generate_video( prompts, # List of text prompts to use to generate media\n",
        "                    h=h,w=w,\n",
        "                    lr=lr,\n",
        "                    num_augs=4, \n",
        "                    debug=False, display_prompt=display_prompt,\n",
        "                    frames_per_prompt=frames_per_prompt, # Number of frames to dedicate to each prompt\n",
        "                    first_iter=50, # Number of optimization iterations for first first frame\n",
        "                    num_iter=num_iter, # Optimization iterations for all but first frame\n",
        "                    carry_over_iter=carry_over_iter,\n",
        "                    z_unchanging_weight=z_unchanging_weight, # Weight to ensure z does not change at all * l1_loss(z, z_prev)\n",
        "                    z_noise_squish=z_noise_squish, # Amount to squish z by between frames\n",
        "                    n_samples=1)\n",
        "    if download:\n",
        "        from google.colab import files\n",
        "        files.download(fn)\n",
        "    return all_canvases"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Make a Video :]"
      ],
      "metadata": {
        "id": "Ht5KUCWEuqEh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title # Options\n",
        "#@markdown Do you want to display the prompts on the video frames?\n",
        "display_prompt_on_frames = True #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown How much frame-to-frame differences should be encouraged (0 is near none, 100 is a lot)\n",
        "temperature = 30 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown Number of video frames per prompt\n",
        "frames_per_prompt = 60 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown Fast mode generates at 1-2 FPS at the cost of some quality\n",
        "fast_mode = True #@param {type:\"boolean\"}"
      ],
      "metadata": {
        "cellView": "form",
        "id": "lC-NtoeEUlYG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prompts = ['There is a car rusting at the bottom of the ocean on fifth avenue',\n",
        "           'A diver swims to the car and sits in the drivers seat',\n",
        "           'The diver glares at the facade of their old favorite pizza shop, now completely submerged',\n",
        "           'A school of fish passes by over the head of the diver',\n",
        "           'The moonlight shines through the water onto the diver as they swim down the street']\n",
        "canvases = generate_video_wrapper(prompts, frames_per_prompt=frames_per_prompt, fast=fast_mode,\n",
        "                                  temperature=temperature, display_prompt=display_prompt_on_frames)"
      ],
      "metadata": {
        "id": "xsPfy7t1UOj6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prompts = ['A sad frog sits on a stump in the forest',\n",
        "           'A magic speaker appears and plays the frog\\'s favorite song',\n",
        "           'The frog is suddenly dressed as a ballerina and dances',\n",
        "           'The ballerina frog does amazing acrobatic moves in the moonlit forest',\n",
        "           'The ballerina frog dances in the forest until the sun comes up']\n",
        "canvases = generate_video_wrapper(prompts, frames_per_prompt=frames_per_prompt, fast=fast_mode,\n",
        "                                  temperature=temperature, display_prompt=display_prompt_on_frames)"
      ],
      "metadata": {
        "id": "FvA5T4-qbo5t"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "1IlcVZ2V1xM0"
      ],
      "provenance": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}