Illustra.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Illustra.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "toWe1IoH7X35"
      },
      "source": [
        "# Illustra: Multi-text to Image\n",
        "\n",
        "Part of [Aphantasia](https://github.com/eps696/aphantasia) suite, made by Vadim Epstein [[eps696](https://github.com/eps696)]  \n",
        "Based on [CLIP](https://github.com/openai/CLIP) + FFT from [Lucent](https://github.com/greentfrapp/lucent).  \n",
        "thanks to [Ryan Murdock](https://twitter.com/advadnoun), [Jonathan Fly](https://twitter.com/jonathanfly), [@eduwatch2](https://twitter.com/eduwatch2) for ideas.\n",
        "\n",
        "## Features \n",
        "* **continuously processes phrase lists** (e.g. illustrating lyrics)\n",
        "* generates massive detailed high res imagery, a la deepdream\n",
        "* directly parameterized with [FFT](https://github.com/greentfrapp/lucent/blob/master/lucent/optvis/param/spatial.py) (no pretrained GANs)\n",
        "* various CLIP models, dual mode\n",
        "* saving/loading FFT snapshots to resume processing\n",
        "* separate text prompt for image style\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QytcEMSKBtN-"
      },
      "source": [
        "**Run the cell below after each session restart**\n",
        "\n",
        "Mark `resume` and upload `.pt` file, if you're resuming from the saved snapshot."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "etzxXVZ_r-Nf",
        "cellView": "form"
      },
      "source": [
        "#@title General setup\n",
        "\n",
        "!pip install ftfy transformers\n",
        "!pip install gputil ffpb \n",
        "\n",
        "# !apt-get -qq install ffmpeg\n",
        "from google.colab import drive\n",
        "drive.mount('/G', force_remount=True)\n",
        "gdir = '/G/MyDrive/'\n",
        "%cd $gdir\n",
        "root_dir = 'illustra'\n",
        "import os\n",
        "root_dir = os.path.join(gdir, root_dir)\n",
        "os.makedirs(root_dir, exist_ok=True)\n",
        "%cd $root_dir\n",
        "\n",
        "import os\n",
        "import io\n",
        "import time\n",
        "import math\n",
        "import random\n",
        "import imageio\n",
        "import numpy as np\n",
        "import PIL\n",
        "from base64 import b64encode\n",
        "import shutil\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision\n",
        "from torch.autograd import Variable\n",
        "\n",
        "from IPython.display import HTML, Image, display, clear_output\n",
        "from IPython.core.interactiveshell import InteractiveShell\n",
        "InteractiveShell.ast_node_interactivity = \"all\"\n",
        "import ipywidgets as ipy\n",
        "from google.colab import output, files\n",
        "output.enable_custom_widget_manager()\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "\n",
        "!pip install git+https://github.com/openai/CLIP.git --no-deps\n",
        "import clip\n",
        "!pip install kornia\n",
        "import kornia\n",
        "!pip install lpips\n",
        "import lpips\n",
        "!pip install PyWavelets==1.1.1\n",
        "!pip install git+https://github.com/fbcotter/pytorch_wavelets\n",
        "\n",
        "%cd /content\n",
        "!pip install git+https://github.com/eps696/aphantasia\n",
        "from aphantasia.image import to_valid_rgb, fft_image\n",
        "from aphantasia.utils import slice_imgs, derivat, basename, file_list, img_list, img_read, txt_clean, checkout, old_torch, save_cfg, sim_func, aesthetic_model\n",
        "from aphantasia import transforms\n",
        "from aphantasia.progress_bar import ProgressIPy as ProgressBar\n",
        "\n",
        "clear_output()\n",
        "\n",
        "def read_pt(file):\n",
        "  return torch.load(file).cuda()\n",
        "\n",
        "def save_img(img, fname=None):\n",
        "  img = np.array(img)[:,:,:]\n",
        "  img = np.transpose(img, (1,2,0))  \n",
        "  img = np.clip(img*255, 0, 255).astype(np.uint8)\n",
        "  if fname is not None:\n",
        "    imageio.imsave(fname, np.array(img))\n",
        "    imageio.imsave('result.jpg', np.array(img))\n",
        "\n",
        "def makevid(seq_dir, size=None):\n",
        "  out_sequence = seq_dir + '/%05d.jpg'\n",
        "  out_video = seq_dir + '.mp4'\n",
        "  print('.. generating video ..')\n",
        "  !ffmpeg -y -v warning -i $out_sequence -crf 20 $out_video\n",
        "  data_url = \"data:video/mp4;base64,\" + b64encode(open(out_video,'rb').read()).decode()\n",
        "  wh = '' if size is None else 'width=%d height=%d' % (size, size)\n",
        "  return \"\"\"<video %s controls><source src=\"%s\" type=\"video/mp4\"></video>\"\"\" % (wh, data_url)\n",
        "\n",
        "# Hardware check\n",
        "!ln -sf /opt/bin/nvidia-smi /usr/bin/nvidia-smi\n",
        "import GPUtil as GPU\n",
        "gpu = GPU.getGPUs()[0] # XXX: only one GPU on Colab and isn’t guaranteed\n",
        "!nvidia-smi -L\n",
        "print(\"GPU RAM {0:.0f}MB | Free {1:.0f}MB)\".format(gpu.memoryTotal, gpu.memoryFree))\n",
        "print('\\nDone!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Load inputs\n",
        "\n",
        "#@markdown **Content** (either type a text string, or upload a text file):\n",
        "content = \"\" #@param {type:\"string\"}\n",
        "upload_texts = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown **Style** (either type a text string, or upload a text file):\n",
        "style = \"\" #@param {type:\"string\"}\n",
        "upload_styles = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown For non-English languages use Google translation:\n",
        "translate = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown Resume from the saved snapshot (resolution settings below will be ignored in this case): \n",
        "resume = False #@param {type:\"boolean\"}\n",
        "\n",
        "if translate:\n",
        "  !pip3 install googletrans==3.1.0a0\n",
        "  clear_output()\n",
        "  from googletrans import Translator\n",
        "  translator = Translator()\n",
        "\n",
        "if upload_texts:\n",
        "  print('Upload main text file')\n",
        "  uploaded = files.upload()\n",
        "  text_file = list(uploaded)[0]\n",
        "  texts = list(uploaded.values())[0].decode().split('\\n')\n",
        "  texts = [tt.strip() for tt in texts if len(tt.strip())>0 and tt[0] != '#']\n",
        "  if translate:\n",
        "    texts = [translator.translate(txt, dest='en').text for txt in texts]\n",
        "  print(' main text:', text_file, len(texts), 'lines')\n",
        "  workname = txt_clean(basename(text_file))\n",
        "else:\n",
        "  texts = [content]\n",
        "  workname = txt_clean(content)[:44]\n",
        "\n",
        "if upload_styles:\n",
        "  print('Upload styles text file')\n",
        "  uploaded = files.upload()\n",
        "  text_file = list(uploaded)[0]\n",
        "  styles = list(uploaded.values())[0].decode().split('\\n')\n",
        "  styles = [tt.strip() for tt in styles if len(tt.strip())>0 and tt[0] != '#']\n",
        "  if translate:\n",
        "    styles = [translator.translate(txt, dest='en').text for txt in styles]\n",
        "  print(' styles:', text_file, len(styles), 'lines')\n",
        "else:\n",
        "  styles = [style]\n",
        "\n",
        "if resume:\n",
        "  print('Upload snapshot to resume from')\n",
        "  resumed = files.upload()\n",
        "  params_pt = list(resumed.values())[0]\n",
        "  params_pt = torch.load(io.BytesIO(params_pt))\n",
        "  if isinstance(params_pt, list): params_pt = params_pt[0]\n",
        "\n",
        "assert len(texts) > 0 or len(styles) > 0, 'No input text[s] found!'\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "eA2NSRf7bznV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**`content`** (what to draw) is your primary input; **`style`** (how to draw) is optional, if you want to separate such descriptions.  \n",
        "All text inputs understand syntax with weights, like `good prompt :1 | also good prompt :1 | bad prompt :-0.5` (within one line).  "
      ],
      "metadata": {
        "id": "vFSCG0U5cNuw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Main settings\n",
        "\n",
        "sideX = 1280 #@param {type:\"integer\"}\n",
        "sideY = 720 #@param {type:\"integer\"}\n",
        "duration = 10#@param {type:\"integer\"}\n",
        "#@markdown > Config\n",
        "model = 'dual' #@param ['dual', 'ViT-B/16', 'ViT-B/32']\n",
        "aesthetic =  1#@param {type:\"number\"}\n",
        "\n",
        "# Default settings\n",
        "steps = 150\n",
        "samples = 200\n",
        "show_freq = 10\n",
        "align = 'uniform'\n",
        "decay = 1.5\n",
        "contrast = 1.1\n",
        "colors = 1.8\n",
        "sharpness = 0.\n",
        "aug_noise = 0.\n",
        "learning_rate = 0.05\n",
        "optimizer = 'adam'\n",
        "aug_transform = 'fast'\n",
        "macro = 0.4\n",
        "enforce = 0.\n",
        "loop = True\n",
        "keep = 1.\n",
        "fps = 25\n",
        "sample_decrease = 1.\n",
        "\n",
        "if resume:\n",
        "  sideY = params_pt.shape[2]\n",
        "  sideX = (params_pt.shape[3] - 1) * 2\n",
        "\n",
        "if model == 'dual':\n",
        "  dualmod = 2\n",
        "  model = 'ViT-B/32'\n",
        "else:\n",
        "  dualmod = None\n",
        "\n",
        "model_clip, _ = clip.load(model, jit=False)\n",
        "try:\n",
        "  modsize = model_clip.visual.input_resolution\n",
        "except:\n",
        "  modsize = 336 if '336' in model else 224\n",
        "model_clip = model_clip.eval().cuda()\n",
        "xmem = {'ViT-B/16':0.25, 'ViT-L/14':0.04}\n",
        "if model in xmem.keys():\n",
        "  sample_decrease *= xmem[model]\n",
        "\n",
        "if dualmod is not None: # second is vit-16\n",
        "  model_clip2, _ = clip.load('ViT-B/16', jit=False)\n",
        "  sample_decrease *= 0.23\n",
        "  dualmod_nums = list(range(steps))[dualmod::dualmod]\n",
        "  print(' dual model every %d step' % dualmod)\n",
        "\n",
        "if aesthetic != 0 and model in ['ViT-B/32', 'ViT-B/16', 'ViT-L/14']:\n",
        "  aest = aesthetic_model(model).cuda()\n",
        "  if dualmod is not None:\n",
        "    aest2 = aesthetic_model('ViT-B/16').cuda()\n",
        "    \n",
        "%cd $work_dir\n",
        "clear_output()\n",
        "print(' using CLIP model', model if dualmod is None else 'dual')\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "XscTs4pXcrc2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Set the desired video resolution and `duration` (in sec).  \n",
        "Select CLIP visual `model` (results do vary!). `dual` (ViT-B/32 + ViT-B/16) usually works best.  \n",
        "`aesthetic` enforces overall cuteness (try various values!). May be  negative.  \n"
      ],
      "metadata": {
        "id": "2rBTuCmHjw3s"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Other settings [optional]"
      ],
      "metadata": {
        "id": "awtjCSPOkS4V"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Run this cell to override settings, if needed\n",
        "#@markdown [to roll back defaults, run \"Main settings\" cell again]\n",
        "\n",
        "#@markdown > Video\n",
        "loop = True #@param {type:\"boolean\"}\n",
        "keep = 1. #@param {type:\"number\"}\n",
        "fps = 25 #@param {type:\"integer\"}\n",
        "\n",
        "#@markdown > Look\n",
        "decay = 1.5 #@param {type:\"number\"}\n",
        "colors = 1.8 #@param {type:\"number\"}\n",
        "contrast =  1.1 #@param {type:\"number\"}\n",
        "sharpness = 0. #@param {type:\"number\"}\n",
        "\n",
        "#@markdown > Training\n",
        "steps = 150 #@param {type:\"integer\"}\n",
        "samples = 200 #@param {type:\"integer\"}\n",
        "show_freq = 10 #@param {type:\"integer\"}\n",
        "learning_rate = 0.05 #@param {type:\"number\"}\n",
        "optimizer = 'adam' #@param ['adam', 'adamw']\n",
        "\n",
        "#@markdown > Tricks\n",
        "aug_transform = 'fast' #@param ['fast', 'custom', 'elastic', 'none']\n",
        "macro = 0.4 #@param {type:\"number\"}\n",
        "aug_noise = 0. #@param {type:\"number\"}\n",
        "enforce = 0. #@param {type:\"number\"}\n",
        "overscan = False #@param {type:\"boolean\"}\n",
        "align = 'overscan' if overscan else 'uniform'\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "YI7iabKgkWtP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f3Sj0fxmtw6K"
      },
      "source": [
        "If `loop` - the inputs are looped [if there are fewer entries in one file]. Otherwise keeps the last one.\n",
        "`keep` parameter controls how well the next line/image follows the previous. 0 means it's randomly initiated, the higher - the stricter it will keep the original composition. Safe values are 1~2 (much higher numbers may cause the imagery getting stuck).  \n",
        "\n",
        "Tune `decay` (softness) and `sharpness`, `colors` (saturation) and `contrast` as needed.  \n",
        "\n",
        "`steps` defines the number of iterations per one input. 100~150 is usually enough.  \n",
        "Decrease **`samples`** if you face OOM (it's the main VRAM eater).  \n",
        "`show_freq` controls preview frequency (doesn't affect the results; one can set it higher to speed up the process).  \n",
        "**`learning_rate`** is the main driver! Decrease it for softer imagery, increase for more powerful processing.  \n",
        "\n",
        "`aug_transform` applies some augmentations, inhibiting image fragmentation & \"graffiti\" printing (slower, yet recommended).  \n",
        "`macro` (from 0 to 1) boosts bigger forms.  \n",
        "`aug_noise` augmentation can make the image less dispersed.  \n",
        "`enforce` may boost training consistency (of simultaneous samples). good start is 0.1~0.2.  \n",
        "`overscan` provides more uniform frame coverage [up to semi-seamless tileable texture].  \n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Generate"
      ],
      "metadata": {
        "id": "AT-E8jY3raml"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nq0wA-wc-P-s",
        "cellView": "form"
      },
      "source": [
        "#@title run\n",
        "\n",
        "work_dir = os.path.join(root_dir, workname)\n",
        "if dualmod is None: work_dir += '-%s' % model.replace('/','').replace('-','') \n",
        "if enforce != 0:    work_dir += '-e%.2g' % enforce\n",
        "os.makedirs(work_dir, exist_ok=True)\n",
        "print('main dir', work_dir)\n",
        "\n",
        "if 'elastic' in aug_transform:\n",
        "  trform_f = transforms.transforms_elastic\n",
        "elif 'custom' in aug_transform:\n",
        "  trform_f = transforms.transforms_custom\n",
        "elif 'fast' in aug_transform:\n",
        "  trform_f = transforms.transforms_fast\n",
        "else:\n",
        "  trform_f = transforms.normalize()\n",
        "  sample_decrease *= 1.05\n",
        "sample_decrease *= 0.95\n",
        "if enforce != 0:\n",
        "  sample_decrease *= 0.5\n",
        "\n",
        "samples = int(samples * sample_decrease)\n",
        "\n",
        "def enc_text(txt, model_clip=model_clip):\n",
        "  if txt is None or len(txt)==0: return None\n",
        "  embs = []\n",
        "  for subtxt in txt.split('|'):\n",
        "    if ':' in subtxt:\n",
        "      [subtxt, wt] = subtxt.split(':')\n",
        "      wt = float(wt)\n",
        "    else: wt = 1.\n",
        "    emb = model_clip.encode_text(clip.tokenize(subtxt).cuda()[:77])\n",
        "    embs.append([emb.detach().clone(), wt])\n",
        "  return embs\n",
        "\n",
        "def pick_(list_, num_):\n",
        "  cnt = len(list_)\n",
        "  if cnt == 0: return None\n",
        "  num = num_ % cnt if loop is True else min(num_, cnt-1)\n",
        "  return list_[num]\n",
        "\n",
        "count = 0\n",
        "\n",
        "txt_encs = [enc_text(txt) for txt in texts] \n",
        "if dualmod is not None:\n",
        "  txt_encs2 = [enc_text(txt, model_clip2) for txt in texts]\n",
        "count = max(count, len(txt_encs))\n",
        "\n",
        "styl_encs = [enc_text(style) for style in styles]\n",
        "if dualmod is not None:\n",
        "  styl_encs2 = [enc_text(style, model_clip2) for style in styles]\n",
        "count = max(count, len(styl_encs))\n",
        "    \n",
        "assert count > 0, \"No inputs found!\"\n",
        "\n",
        "outpic = ipy.Output()\n",
        "outpic\n",
        "\n",
        "def train(num, i):\n",
        "  loss = 0\n",
        "  noise = aug_noise * (torch.rand(1, 1, *params[0].shape[2:4], 1)-0.5).cuda() if aug_noise > 0 else None\n",
        "  img_out = image_f(noise)\n",
        "  img_sliced = slice_imgs([img_out], samples, modsize, trform_f, align, macro)[0]\n",
        "  \n",
        "  if len(texts) > 0:\n",
        "    txt_enc   = pick_(txt_encs2, num)  if dualmod is not None and i in dualmod_nums else pick_(txt_encs, num)\n",
        "  if len(styles) > 0:\n",
        "    style_enc = pick_(styl_encs2, num) if dualmod is not None and i in dualmod_nums else pick_(styl_encs, num)\n",
        "  model_clip_ = model_clip2 if dualmod is not None and i in dualmod_nums else model_clip\n",
        "  if aesthetic != 0:\n",
        "    aest_     = aest2       if dualmod is not None and i in dualmod_nums else aest\n",
        "\n",
        "  out_enc = model_clip_.encode_image(img_sliced)\n",
        "  if aesthetic != 0 and aest_ is not None:\n",
        "    loss -= 0.001 * aesthetic * aest_(out_enc).mean()\n",
        "  if len(texts) > 0 and txt_enc is not None: # input text - main topic\n",
        "    for enc, wt in txt_enc:\n",
        "      loss -= wt * sim_func(enc, out_enc, 'cossim')\n",
        "  if len(styles) > 0 and style_enc is not None: # input text - style\n",
        "    for enc, wt in style_enc:\n",
        "      loss -= wt * sim_func(enc, out_enc, 'cossim')\n",
        "  if sharpness != 0: # scharr|sobel|naiv\n",
        "    loss -= sharpness * derivat(img_out, mode='naiv')\n",
        "  if enforce != 0:\n",
        "    img_sliced = slice_imgs([image_f(noise_)], samples, modsize, trform_f, align, macro)[0]\n",
        "    out_enc2 = model_clip_.encode_image(img_sliced)\n",
        "    loss -= enforce * sim_func(out_enc, out_enc2, 'cossim')\n",
        "    del out_enc2\n",
        "\n",
        "  del img_out, img_sliced, out_enc\n",
        "  assert not isinstance(loss, int), ' Loss not defined, check inputs'\n",
        "  \n",
        "  optimr.zero_grad()\n",
        "  loss.backward()\n",
        "  optimr.step()\n",
        "\n",
        "  if i % show_freq == 0:\n",
        "    with torch.no_grad():\n",
        "      img = image_f(contrast=contrast).cpu().numpy()[0]\n",
        "    save_img(img, os.path.join(tempdir, '%04d.jpg' % (i // show_freq)))\n",
        "    outpic.clear_output()\n",
        "    with outpic:\n",
        "      display(Image('result.jpg'))\n",
        "    del img\n",
        "    _ = pbar.upd()\n",
        "\n",
        "for num in range(count):\n",
        "  shape = [1, 3, sideY, sideX]\n",
        "  global params\n",
        "\n",
        "  if num == 0:\n",
        "    resume_cur = params_pt if resume else None\n",
        "  else:\n",
        "    opt_state = optimr.state_dict()\n",
        "    param_ = params[0].detach()\n",
        "    resume_cur = [keep * param_ / (param_.max() - param_.min())]\n",
        "\n",
        "  params, image_f, sz = fft_image(shape, 0.08, decay, resume_cur)\n",
        "  if sz is not None: [sideY, sideX] = sz\n",
        "  image_f = to_valid_rgb(image_f, colors = colors)\n",
        "\n",
        "  if optimizer.lower() == 'adamw':\n",
        "     optimr = torch.optim.AdamW(params, learning_rate, weight_decay=0.01, betas=(.0,.999), amsgrad=True)\n",
        "  else:\n",
        "    optimr = torch.optim.Adam(params, learning_rate, betas=(.0, .999))\n",
        "  if num > 0: optimr.load_state_dict(opt_state)\n",
        "\n",
        "  out_names = []\n",
        "  if len(texts)  > 0: out_names += [txt_clean(pick_(texts, num))[:32]]\n",
        "  if len(styles) > 0: out_names += [txt_clean(pick_(styles, num))[:32]]\n",
        "  out_name = '-'.join(out_names)\n",
        "  if count > 1: out_name = '%03d-' % (num+1) + out_name\n",
        "  print(out_name)\n",
        "  tempdir = os.path.join(work_dir, out_name)\n",
        "  os.makedirs(tempdir, exist_ok=True)\n",
        "\n",
        "  pbar = ProgressBar(steps // show_freq)\n",
        "  for i in range(steps):\n",
        "    train(num, i)\n",
        "\n",
        "  file_out = os.path.join(work_dir, '%s-%d.jpg' % (out_name, steps))\n",
        "  _ = shutil.copy(img_list(tempdir)[-1], file_out)\n",
        "  _ = os.system('ffmpeg -v warning -y -i %s\\%%04d.jpg \"%s.mp4\"' % (tempdir, os.path.join(work_dir, out_name)))\n",
        "  torch.save(params[0], '%s.pt' % os.path.join(work_dir, out_name))\n",
        "\n",
        "vsteps = int(duration * fps / count)\n",
        "tempdir = os.path.join(work_dir, '_final')\n",
        "!rm -rf $tempdir\n",
        "os.makedirs(tempdir, exist_ok=True)\n",
        "\n",
        "print(' rendering complete piece')\n",
        "ptfiles = file_list(work_dir, 'pt')\n",
        "pbar = ProgressBar(vsteps * len(ptfiles))\n",
        "for px in range(len(ptfiles)):\n",
        "  params1 = read_pt(ptfiles[px])\n",
        "  params2 = read_pt(ptfiles[(px+1) % len(ptfiles)])\n",
        "\n",
        "  params, image_f, _ = fft_image([1, 3, sideY, sideX], resume=params1, sd=1., decay_power=decay)\n",
        "  image_f = to_valid_rgb(image_f, colors = colors)\n",
        "\n",
        "  for i in range(vsteps):\n",
        "    with torch.no_grad():\n",
        "      x = i/vsteps # math.sin(1.5708 * i/vsteps)\n",
        "      img = image_f((params2 - params1) * x, contrast=contrast)[0].permute(1,2,0)\n",
        "      img = torch.clip(img*255, 0, 255).cpu().numpy().astype(np.uint8)\n",
        "    imageio.imsave(os.path.join(tempdir, '%05d.jpg' % (px * vsteps + i)), img)\n",
        "    _ = pbar.upd()\n",
        "\n",
        "HTML(makevid(tempdir))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rARq2n6LgVVp",
        "cellView": "form"
      },
      "source": [
        "#@markdown Run this, if you want to make another video from the directory with saved snapshots [leave it empty to pick up the current ones]\n",
        "\n",
        "saved_dir = '' #@param {type:\"string\"}\n",
        "duration =  12 #@param {type:\"integer\"}\n",
        "fps = 25 #@param {type:\"integer\"}\n",
        "\n",
        "if len(saved_dir) > 0:\n",
        "  work_dir = saved_dir\n",
        "tempdir = os.path.join(work_dir, '_final')\n",
        "!rm -rf $tempdir\n",
        "os.makedirs(tempdir, exist_ok=True)\n",
        "\n",
        "print(' re-rendering final piece')\n",
        "ptfiles = file_list(work_dir, 'pt')\n",
        "vsteps = int(duration * fps / (len(ptfiles)))\n",
        "\n",
        "ptest = torch.load(ptfiles[0])\n",
        "if isinstance(ptest, list): ptest = ptest[0]\n",
        "shape = [*ptest.shape[:3], (ptest.shape[3]-1)*2]\n",
        "\n",
        "pbar = ProgressBar(vsteps * len(ptfiles))\n",
        "for px in range(len(ptfiles)):\n",
        "  params1 = read_pt(ptfiles[px])\n",
        "  params2 = read_pt(ptfiles[(px+1) % len(ptfiles)])\n",
        "\n",
        "  params, image_f, _ = fft_image(shape, resume=params1, decay_power=decay)\n",
        "  image_f = to_valid_rgb(image_f, colors = colors)\n",
        "\n",
        "  for i in range(vsteps):\n",
        "    with torch.no_grad():\n",
        "      img = image_f((params2 - params1) * math.sin(1.5708 * i/vsteps))[0].permute(1,2,0)\n",
        "      img = torch.clip(img*255, 0, 255).cpu().numpy().astype(np.uint8)\n",
        "    imageio.imsave(os.path.join(tempdir, '%05d.jpg' % (px * vsteps + i)), img)\n",
        "    _ = pbar.upd()\n",
        "\n",
        "HTML(makevid(tempdir))\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}