IllusTrip3D.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "IllusTrip3D.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [
        "f3Sj0fxmtw6K",
        "YdVubN0vb3TU"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "toWe1IoH7X35"
      },
      "source": [
        "# IllusTrip: Text to Video 3D\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/pixel ops from [Lucent](https://github.com/greentfrapp/lucent). \n",
        "3D part by [deKxi](https://twitter.com/deKxi), based on [AdaBins](https://github.com/shariqfarooq123/AdaBins) depth.  \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 **multiple sentences** (e.g. illustrating lyrics or poems)\n",
        "* makes **videos**, evolving with pan/zoom/rotate motion\n",
        "* works with [inverse FFT](https://github.com/greentfrapp/lucent/blob/master/lucent/optvis/param/spatial.py) representation of the image or **directly with RGB** pixels (no GANs involved)\n",
        "* generates massive detailed textures (a la deepdream), **unlimited resolution**\n",
        "* optional **depth** processing for 3D look\n",
        "* can start/resume from an image\n",
        "* various CLIP models, dual mode\n",
        "* employs [aesthetic loss](https://github.com/LAION-AI/aesthetic-predictor) \n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QytcEMSKBtN-"
      },
      "source": [
        "**Run the cell below after each session restart**\n",
        "\n",
        "Ensure that you're given Tesla T4/P4/P100 GPU. With Tesla K80 it will be ~8x slower."
      ]
    },
    {
      "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",
        "work_dir = '/content/illustrip'\n",
        "import os\n",
        "os.makedirs(work_dir, exist_ok=True)\n",
        "%cd $work_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",
        "from easydict import EasyDict as edict\n",
        "a = edict()\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision\n",
        "from torchvision import transforms as T\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 sentence_transformers\n",
        "from sentence_transformers import SentenceTransformer\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",
        "!pip install git+https://github.com/eps696/aphantasia\n",
        "from aphantasia.image import to_valid_rgb, fft_image, rfft2d_freqs, img2fft, pixel_image, un_rgb\n",
        "from aphantasia.utils import basename, file_list, img_list, img_read, txt_clean, plot_text\n",
        "from aphantasia.utils import slice_imgs, derivat, pad_up_to, slerp, aesthetic_model, intrl, checkout, sim_func, latent_anima\n",
        "from aphantasia import transforms\n",
        "from aphantasia.progress_bar import ProgressIPy as ProgressBar\n",
        "\n",
        "%cd $work_dir\n",
        "!git clone https://github.com/eps696/aphantasia --recursive\n",
        "work_dir = os.path.join(work_dir, 'aphantasia')\n",
        "%cd $work_dir\n",
        "from depth import depth\n",
        "# !wget https://github.com/eps696/aphantasia/blob/master/mask.jpg?raw=true -O mask.jpg\n",
        "depth_mask_file = os.path.join(work_dir, 'depth', 'mask.jpg')\n",
        "%cd /content\n",
        "\n",
        "clear_output()\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",
        "  char_len = len(basename(img_list(seq_dir)[0]))\n",
        "  out_sequence = seq_dir + '/%0{}d.jpg'.format(char_len)\n",
        "  out_video = seq_dir + '.mp4'\n",
        "  print('.. generating video ..')\n",
        "  !ffmpeg -y -v warning -i $out_sequence -crf 18 $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))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JUvpdy8BWGuM",
        "cellView": "form"
      },
      "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 an image (resolution settings below will be ignored in this case): \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",
        "  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",
        "  print(' styles:', text_file, len(styles), 'lines')\n",
        "else:\n",
        "  styles = [style]\n",
        "\n",
        "resume = False #@param {type:\"boolean\"}\n",
        "if resume:\n",
        "  print('Upload file to resume from')\n",
        "  resumed = files.upload()\n",
        "  resumed_filename = list(resumed)[0]\n",
        "  resumed_bytes = list(resumed.values())[0]\n",
        "\n",
        "assert len(texts) > 0 and len(texts[0]) > 0, 'No input text[s] found!'\n",
        "tempdir = os.path.join(work_dir, workname)\n",
        "os.makedirs(tempdir, exist_ok=True)\n",
        "print('main dir', tempdir)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PQFGziYKtHSa"
      },
      "source": [
        "**`content`** (what to draw) is your primary input; **`style`** (how to draw) is optional, if you want to separate such descriptions.  \n",
        "If you load text file[s], the imagery will interpolate from line to line (ensure equal line counts for content and style lists, for their accordance).  \n",
        "All text inputs understand syntax with weights, like `good prompt :1 | also good prompt :1 | bad prompt :-0.5` (within one line).  "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Uti6XrqiQumf",
        "cellView": "form"
      },
      "source": [
        "#@title Google Drive [optional]\n",
        "\n",
        "#@markdown Run this cell, if you want to store results on your Google Drive.\n",
        "using_GDrive = True#@param{type:\"boolean\"}\n",
        "if using_GDrive:\n",
        "  import os\n",
        "  from google.colab import drive\n",
        "\n",
        "  if not os.path.isdir('/G/MyDrive'): \n",
        "      drive.mount('/G', force_remount=True)\n",
        "  gdir = '/G/MyDrive'\n",
        "\n",
        "  tempdir = os.path.join(gdir, 'illustrip', workname)\n",
        "  os.makedirs(tempdir, exist_ok=True)\n",
        "  print('main dir', tempdir)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "64mlBCAYeOrB",
        "cellView": "form"
      },
      "source": [
        "#@title Main settings\n",
        "\n",
        "sideX = 1280 #@param {type:\"integer\"}\n",
        "sideY = 720 #@param {type:\"integer\"}\n",
        "steps = 200 #@param {type:\"integer\"}\n",
        "frame_step = 100 #@param {type:\"integer\"}\n",
        "#@markdown > Config\n",
        "method = 'RGB' #@param ['FFT', 'RGB']\n",
        "model = 'dual' #@param ['dual', 'ViT-B/16', 'ViT-B/32', 'RN101', 'RN50x16', 'RN50x4', 'RN50']\n",
        "aesthetic =  1#@param {type:\"number\"}\n",
        "\n",
        "# Default settings\n",
        "if method == 'RGB':\n",
        "  align = 'overscan'\n",
        "  colors = 2.3\n",
        "  contrast = 1.2\n",
        "  sharpness = -1.\n",
        "  aug_noise = 0.\n",
        "  smooth = False\n",
        "else:\n",
        "  align = 'uniform'\n",
        "  colors = 1.8\n",
        "  contrast = 1.1\n",
        "  sharpness = 1.\n",
        "  aug_noise = 2. # only for FFT\n",
        "  smooth = True\n",
        "interpolate_topics = True\n",
        "style_power = 1.\n",
        "samples = 200\n",
        "save_step = 1\n",
        "learning_rate = 0.1\n",
        "optimizer = 'adam_custom'\n",
        "aug_transform = 'fast'\n",
        "similarity_function = 'mixed'\n",
        "macro = 0.4\n",
        "enforce = 0.\n",
        "expand = 0.\n",
        "zoom = 0.012\n",
        "shift = 10\n",
        "rotate = 0.8\n",
        "distort = 0.3\n",
        "animate_them = True\n",
        "sample_decrease = 1.\n",
        "DepthStrength = 0.\n",
        "save_depth = False\n",
        "\n",
        "if model == 'dual':\n",
        "  dualmod = 2\n",
        "  model = 'ViT-B/32'\n",
        "else:\n",
        "  dualmod = None\n",
        "\n",
        "print(' loading CLIP model..')\n",
        "model_clip, _ = clip.load(model, jit=False)\n",
        "modsize = model_clip.visual.input_resolution\n",
        "xmem = {'ViT-B/16':0.25, 'RN50':0.5, 'RN50x4':0.16, 'RN50x16':0.06, 'RN101':0.33}\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",
        "\n",
        "if method == 'RGB':\n",
        "\n",
        "  if resume:\n",
        "    img_in = imageio.imread(resumed_bytes)\n",
        "    # params_tmp = torch.Tensor(img_in).permute(2,0,1).unsqueeze(0).float().cuda()\n",
        "    params_tmp = 3.3 * un_rgb(img_in, colors=2.)\n",
        "    sideY, sideX = img_in.shape[0], img_in.shape[1]\n",
        "  else:\n",
        "    params_tmp = torch.randn(1, 3, sideY, sideX).cuda() # * 0.01\n",
        "\n",
        "else: # FFT\n",
        "\n",
        "  if resume:\n",
        "    if os.path.splitext(resumed_filename)[1].lower()[1:] in ['jpg','png','tif','bmp']:\n",
        "      img_in = imageio.imread(resumed_bytes)\n",
        "      params_tmp = img2fft(img_in, 1.5, 1.) * 2.\n",
        "    else:\n",
        "      params_tmp = torch.load(io.BytesIO(resumed_bytes))\n",
        "      if isinstance(params_tmp, list): params_tmp = params_tmp[0]\n",
        "    params_tmp = params_tmp.cuda()\n",
        "    sideY, sideX = params_tmp.shape[2], (params_tmp.shape[3]-1)*2\n",
        "  else:\n",
        "    params_shape = [1, 3, sideY, sideX//2+1, 2]\n",
        "    params_tmp = torch.randn(*params_shape).cuda() * 0.01\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JIWNmmd5uuSn"
      },
      "source": [
        "**`FFT`** method uses inverse FFT representation of the image. It allows flexible motion, but is either blurry (if smoothed) or noisy (if not).  \n",
        "**`RGB`** method directly optimizes image pixels (without FFT parameterization). It's more clean and stable, when zooming in.  \n",
        "There are few choices for CLIP `model` (results do vary!). Dual (ViT-B/32 + ViT-B/16) usually works best. Otherwise, I prefer ViT's for consistency.  \n",
        "`aesthetic` enforces overall cuteness (try various values!). May be  negative.  \n",
        "\n",
        "**`steps`** defines the length of animation per text line (multiply it to the inputs line count to get total video duration in frames).  \n",
        "`frame_step` sets frequency of the changes in animation (how many frames between motion keypoints).  \n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f3Sj0fxmtw6K"
      },
      "source": [
        "## Other settings [optional]"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P88_xcpAIXlq",
        "cellView": "form"
      },
      "source": [
        "#@title Run this cell to override settings, if needed\n",
        "#@markdown [to roll back defaults, run \"Main settings\" cell again]\n",
        "\n",
        "style_power = 1. #@param {type:\"number\"}\n",
        "overscan = True #@param {type:\"boolean\"}\n",
        "align = 'overscan' if overscan else 'uniform'\n",
        "interpolate_topics = True #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown > Look\n",
        "colors = 2.3 #@param {type:\"number\"}\n",
        "contrast = 1.2 #@param {type:\"number\"}\n",
        "sharpness =  0#@param {type:\"number\"}\n",
        "\n",
        "#@markdown > Training\n",
        "samples = 200 #@param {type:\"integer\"}\n",
        "save_step = 1 #@param {type:\"integer\"}\n",
        "learning_rate = 0.1 #@param {type:\"number\"}\n",
        "optimizer = 'adam_custom' #@param ['adam', 'adam_custom', 'adamw', 'adamw_custom']\n",
        "\n",
        "#@markdown > Tricks\n",
        "aug_transform = 'fast' #@param ['fast', 'custom', 'elastic', 'none']\n",
        "aug_noise = 0. #@param {type:\"number\"}\n",
        "macro = 0.3 #@param {type:\"number\"}\n",
        "enforce = 0. #@param {type:\"number\"}\n",
        "expand = 0. #@param {type:\"number\"}\n",
        "similarity_function = 'mixed' #@param ['cossim', 'spherical', 'mixed', 'angular', 'dot']\n",
        "\n",
        "#@markdown > Motion\n",
        "zoom = 0.012 #@param {type:\"number\"}\n",
        "shift = 10 #@param {type:\"number\"}\n",
        "rotate = 0.8 #@param {type:\"number\"}\n",
        "distort = 0.3 #@param {type:\"number\"}\n",
        "animate_them = True #@param {type:\"boolean\"}\n",
        "smooth = True #@param {type:\"boolean\"}\n",
        "if method == 'RGB': smooth = False\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QYrJTb8xDm9C"
      },
      "source": [
        "`style_power` controls overall strength of the style descriptions, comparing to the main input.  \n",
        "`overscan` provides better frame coverage (needed for RGB method).  \n",
        "`interpolate_topics` changes the subjects smoothly, otherwise they're switched by cut, making sharper transitions.  \n",
        "\n",
        "Decrease **`samples`** if you face OOM (it's the main RAM eater), or just to speed up the process (with the cost of quality).  \n",
        "`save_step` defines how many optimization steps are taken between saved frames, \n",
        "set it >1 for stronger image processing.  \n",
        "**`learning_rate`** is the main driver! Decrease it for softer and lighter effects, increase for more powerful but noisy processing.  \n",
        "Select optimizer: `_custom` options are more stable but noisy; pure `adam` is softer, but may spill some colored blur on longer videos.  \n",
        "Negative `sharpness` may also reduce noisiness.  \n",
        "\n",
        "Experimental tricks:  \n",
        "`aug_transform` applies some augmentations, which quite radically change the output of this method (and slow down the process).  \n",
        "Try yourself to see which is good for your case. `aug_noise` augmentation [FFT only!] seems to enhance optimization with transforms.  \n",
        "`macro` boosts bigger forms.  \n",
        "`enforce` adds more details by enforcing similarity between two parallel samples.  \n",
        "`expand` boosts diversity (up to irrelevant) by enforcing difference between prev/next samples.  \n",
        "\n",
        "Motion section:\n",
        "`shift` is in pixels, `rotate` in degrees. The values will be used as limits, if you mark `animate_them`.  \n",
        "\n",
        "`smooth` reduces blinking, but induces motion blur with subtle screen-fixed patterns (valid only for FFT method, disabled for RGB).  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YdVubN0vb3TU"
      },
      "source": [
        "## Add 3D depth [optional]"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vl-rm1Nm03lK",
        "cellView": "form"
      },
      "source": [
        "# pretrained models: Nyu is much better but Kitti is an option too\n",
        "depth_model = 'nyu' # @ param [\"nyu\",\"kitti\"]\n",
        "DepthStrength = 0.01 #@param{type:\"number\"}\n",
        "save_depth = False #@param{type:\"boolean\"}\n",
        "MaskBlurAmt = 33 \n",
        "size = (sideY,sideX)\n",
        "\n",
        "#@markdown NB: depth computing may double rendering time. \n",
        "\n",
        "#@markdown Originally by [deKxi](https://twitter.com/deKxi)\n",
        "\n",
        "%cd $work_dir\n",
        "\n",
        "if DepthStrength != 0:\n",
        "\n",
        "  if not os.path.exists(\"AdaBins_nyu.pt\"):\n",
        "    !gdown https://drive.google.com/uc?id=1lvyZZbC9NLcS8a__YPcUP7rDiIpbRpoF\n",
        "    if not os.path.exists('AdaBins_nyu.pt'):\n",
        "      !wget https://www.dropbox.com/s/tayczpcydoco12s/AdaBins_nyu.pt\n",
        "  # if depth_model=='kitti' and not os.path.exists(os.path.join(workdir_depth, \"pretrained/AdaBins_kitti.pt\")):\n",
        "    # !gdown https://drive.google.com/uc?id=1HMgff-FV6qw1L0ywQZJ7ECa9VPq1bIoj\n",
        "\n",
        "  if save_depth:\n",
        "    depthdir = os.path.join(tempdir, 'depth')\n",
        "    os.makedirs(depthdir, exist_ok=True)\n",
        "    print('depth dir', depthdir)\n",
        "  else:\n",
        "    depthdir = None\n",
        "\n",
        "  depth_infer, depth_mask = depth.init_adabins(size=size, model_path='AdaBins_nyu.pt', mask_path=depth_mask_file)\n",
        "\n",
        "  def depth_transform(img_t, depth_infer, depth_mask, size, depthX=0, scale=1., shift=[0,0], colors=1, depth_dir=None, save_num=0):\n",
        "    size2 = [s//2 for s in size]\n",
        "    if not isinstance(scale, float): scale = float(scale[0])\n",
        "    # d X/Y define the origin point of the depth warp, effectively a \"3D pan zoom\", [-1..1]\n",
        "    # plus = look ahead, minus = look aside\n",
        "    dX = 100. * shift[0] / size[1]\n",
        "    dY = 100. * shift[1] / size[0]\n",
        "    # dZ = movement direction: 1 away (zoom out), 0 towards (zoom in), 0.5 stay\n",
        "    dZ = 0.5 + 32. * (scale-1) \n",
        "    img = depth.depthwarp(img_t, depth_infer, depth_mask, size2, depthX, [dX,dY], dZ, save_path=depth_dir, save_num=save_num)\n",
        "    return img\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YZFuwNux8oEg"
      },
      "source": [
        "## Generate"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nq0wA-wc-P-s",
        "cellView": "form"
      },
      "source": [
        "#@title Generate\n",
        "\n",
        "if aug_transform == 'elastic':\n",
        "  trform_f = transforms.transforms_elastic\n",
        "  sample_decrease *= 0.95\n",
        "elif aug_transform == 'custom':\n",
        "  trform_f = transforms.transforms_custom\n",
        "  sample_decrease *= 0.95\n",
        "elif aug_transform == 'fast':\n",
        "  trform_f = transforms.transforms_fast\n",
        "  sample_decrease *= 0.95\n",
        "  print(' using fast aug transforms')\n",
        "else:\n",
        "  trform_f = transforms.normalize()\n",
        "\n",
        "if enforce != 0:\n",
        "  sample_decrease *= 0.5\n",
        "\n",
        "samples = int(samples * sample_decrease)\n",
        "print(' using %s method, %d samples, %s optimizer' % (method, samples, optimizer))\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",
        "# Encode inputs\n",
        "count = 0 # max count of texts and styles\n",
        "if translate:\n",
        "  texts = [tr.text for tr in translator.translate(texts)]\n",
        "key_txt_encs = [enc_text(txt) for txt in texts]\n",
        "if dualmod is not None:\n",
        "  key_txt_encs2 = [enc_text(txt, model_clip2) for txt in texts]\n",
        "count = max(count, len(key_txt_encs))\n",
        "if translate:\n",
        "  styles = [tr.text for tr in translator.translate(styles)]\n",
        "key_styl_encs = [enc_text(style) for style in styles]\n",
        "if dualmod is not None:\n",
        "  key_styl_encs2 = [enc_text(style, model_clip2) for style in styles]\n",
        "count = max(count, len(key_styl_encs))\n",
        "assert count > 0, \"No inputs found!\"\n",
        "\n",
        "# !rm -rf $tempdir\n",
        "# os.makedirs(tempdir, exist_ok=True)\n",
        "\n",
        "glob_steps = count * steps # saving\n",
        "if glob_steps == frame_step: frame_step = glob_steps // 2 # otherwise no motion\n",
        "\n",
        "outpic = ipy.Output()\n",
        "outpic\n",
        "\n",
        "params_tmp = params_tmp.detach()\n",
        "\n",
        "# animation controls\n",
        "if animate_them:\n",
        "  if method == 'RGB':\n",
        "    m_scale = latent_anima([1], glob_steps, frame_step, uniform=True, cubic=True, start_lat=[-0.3])\n",
        "    m_scale = 1 + (m_scale + 0.3) * zoom # only zoom in\n",
        "  else:\n",
        "    m_scale = latent_anima([1], glob_steps, frame_step, uniform=True, cubic=True, start_lat=[0.6])\n",
        "    m_scale = 1 - (m_scale-0.6) * zoom # ping pong\n",
        "  m_shift = latent_anima([2], glob_steps, frame_step, uniform=True, cubic=True, start_lat=[0.5,0.5])\n",
        "  m_angle = latent_anima([1], glob_steps, frame_step, uniform=True, cubic=True, start_lat=[0.5])\n",
        "  m_shear = latent_anima([1], glob_steps, frame_step, uniform=True, cubic=True, start_lat=[0.5])\n",
        "  m_shift = (m_shift-0.5) * shift   * abs(m_scale-1.) / zoom\n",
        "  m_angle = (m_angle-0.5) * rotate  * abs(m_scale-1.) / zoom\n",
        "  m_shear = (m_shear-0.5) * distort * abs(m_scale-1.) / zoom\n",
        "\n",
        "def get_encs(encs, num):\n",
        "  cnt = len(encs)\n",
        "  if cnt == 0: return []\n",
        "  enc_1 = encs[min(num,   cnt-1)]\n",
        "  enc_2 = encs[min(num+1, cnt-1)]\n",
        "  if not interpolate_topics: return [enc_1] * steps\n",
        "  enc_pairs = []\n",
        "  for i in range(steps):\n",
        "    enc1_step = []\n",
        "    if enc_1 is not None:\n",
        "      for enc, wt in enc_1:\n",
        "        enc1_step.append([enc, wt * (steps-i)/steps])\n",
        "    enc2_step = []\n",
        "    if enc_2 is not None:\n",
        "      for enc, wt in enc_2:\n",
        "        enc2_step.append([enc, wt * i/steps])\n",
        "    enc_pairs.append(enc1_step + enc2_step)\n",
        "  return enc_pairs\n",
        "\n",
        "def frame_transform(img, size, angle, shift, scale, shear):\n",
        "  img = T.functional.affine(img, angle, shift, scale, shear, fill=0, interpolation=T.InterpolationMode.BILINEAR)\n",
        "  img = T.functional.center_crop(img, size) # on 1.8+ also pads\n",
        "  return img\n",
        "\n",
        "prev_enc = 0\n",
        "def process(num):\n",
        "  global params_tmp, opt_state, params, image_f, optimizer, pbar\n",
        "\n",
        "  txt_encs  = get_encs(key_txt_encs,  num)\n",
        "  styl_encs = get_encs(key_styl_encs, num)\n",
        "  if dualmod is not None:\n",
        "    txt_encs2  = get_encs(key_txt_encs2,  num)\n",
        "    styl_encs2 = get_encs(key_styl_encs2, num)\n",
        "    txt_encs  = intrl(txt_encs,  txt_encs2,  dualmod)\n",
        "    styl_encs = intrl(styl_encs, styl_encs2, dualmod)\n",
        "    del txt_encs2, styl_encs2\n",
        "        \n",
        "  if len(texts)  > 0: print(' ref text: ',  texts[min(num, len(texts)-1)][:80])\n",
        "  if len(styles) > 0: print(' ref style: ', styles[min(num, len(styles)-1)][:80])\n",
        "\n",
        "  for ii in range(steps):\n",
        "    glob_step = num * steps + ii # saving/transforming\n",
        "\n",
        "    # get encoded inputs\n",
        "    txt_enc  = txt_encs[ii % len(txt_encs)]   if len(txt_encs)  > 0 else None\n",
        "    styl_enc = styl_encs[ii % len(styl_encs)] if len(styl_encs) > 0 else None\n",
        "    \n",
        "    model_clip_ = model_clip2 if dualmod is not None and ii in dualmod_nums else model_clip\n",
        "    if aesthetic != 0:\n",
        "      aest_ = aest2 if dualmod is not None and ii in dualmod_nums else aest\n",
        "\n",
        "    ### animation: transform frame, reload params\n",
        "\n",
        "    h, w = sideY, sideX\n",
        "    \n",
        "    # transform frame for motion\n",
        "    scale =       m_scale[glob_step]    if animate_them else 1-zoom\n",
        "    trans = tuple(m_shift[glob_step])   if animate_them else [0, shift]\n",
        "    angle =       m_angle[glob_step][0] if animate_them else rotate\n",
        "    shear =       m_shear[glob_step][0] if animate_them else distort\n",
        "\n",
        "    if method == 'RGB':\n",
        "      if DepthStrength != 0:\n",
        "        params_tmp = depth_transform(params_tmp, depth_infer, depth_mask, size, DepthStrength, scale, trans, colors, depthdir, glob_step)\n",
        "      params_tmp = frame_transform(params_tmp, (h,w), angle, trans, scale, shear)\n",
        "      params, image_f, _ = pixel_image([1,3,h,w], resume=params_tmp)\n",
        "      img_tmp = None\n",
        "\n",
        "    else: # FFT\n",
        "      if type(params_tmp) is not torch.complex64:\n",
        "        params_tmp = torch.view_as_complex(params_tmp)\n",
        "      img_tmp = torch.fft.irfftn(params_tmp, s=(h,w), norm='ortho')\n",
        "      if DepthStrength != 0:\n",
        "        img_tmp = depth_transform(img_tmp, depth_infer, depth_mask, size, DepthStrength, scale, trans, colors, depthdir, glob_step)\n",
        "      img_tmp = frame_transform(img_tmp, (h,w), angle, trans, scale, shear)\n",
        "      params_tmp = torch.fft.rfftn(img_tmp, s=[h,w], dim=[2,3], norm='ortho')\n",
        "      params_tmp = torch.view_as_real(params_tmp)\n",
        "\n",
        "      params, image_f, _ = fft_image([1,3,h,w], resume=params_tmp, sd=1.)\n",
        "\n",
        "    image_f = to_valid_rgb(image_f, colors=colors)\n",
        "    del img_tmp\n",
        "\n",
        "    if optimizer.lower() == 'adamw':\n",
        "      optimr = torch.optim.AdamW(params, learning_rate, weight_decay=0.01, amsgrad=True)\n",
        "    elif optimizer.lower() == 'adamw_custom':\n",
        "      optimr = torch.optim.AdamW(params, learning_rate, weight_decay=0.01, amsgrad=True, betas=(.0,.999))\n",
        "    elif optimizer.lower() == 'adam':\n",
        "      optimr = torch.optim.Adam(params, learning_rate)\n",
        "    else: # adam_custom\n",
        "      optimr = torch.optim.Adam(params, learning_rate, betas=(.0,.999))\n",
        "\n",
        "    if smooth is True and num + ii > 0:\n",
        "      optimr.load_state_dict(opt_state)\n",
        "\n",
        "    ### optimization\n",
        "\n",
        "    for ss in range(save_step):\n",
        "      loss = 0\n",
        "\n",
        "      noise = aug_noise * (torch.rand(1, 1, *params[0].shape[2:4], 1)-0.5).cuda() if aug_noise > 0 else 0.\n",
        "      img_out = image_f(noise, fixcontrast=resume)\n",
        "      img_sliced = slice_imgs([img_out], samples, modsize, trform_f, align, macro)[0]\n",
        "      out_enc = model_clip_.encode_image(img_sliced)\n",
        "\n",
        "      if aesthetic != 0 and model in ['ViT-B/32', 'ViT-B/16', 'ViT-L/14'] and aest_ is not None:\n",
        "        loss -= 0.001 * aesthetic * aest_(out_enc).mean()\n",
        "\n",
        "      if method == 'RGB': # empirical hack\n",
        "        loss += abs(img_out.mean((2,3)) - 0.45).mean() # fix brightness\n",
        "        loss += abs(img_out.std((2,3)) - 0.17).sum() # fix contrast\n",
        "\n",
        "      if txt_enc is not None:\n",
        "        for enc, wt in txt_enc:\n",
        "          loss -= wt * sim_func(enc, out_enc, similarity_function)\n",
        "      if styl_enc is not None:\n",
        "        for enc, wt in styl_enc:\n",
        "          loss -= style_power * wt * sim_func(enc, out_enc, similarity_function)\n",
        "      if sharpness != 0: # mode = scharr|sobel|naive\n",
        "        loss -= sharpness * derivat(img_out, mode='naive')\n",
        "      if enforce != 0:\n",
        "        img_sliced = slice_imgs([image_f(noise, fixcontrast=resume)], 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, similarity_function)\n",
        "        del out_enc2; torch.cuda.empty_cache()\n",
        "      if expand > 0:\n",
        "        global prev_enc\n",
        "        if ii > 0:\n",
        "          loss += expand * sim_func(prev_enc, out_enc, similarity_function)\n",
        "        prev_enc = out_enc.detach().clone()\n",
        "      del img_out, img_sliced, out_enc; torch.cuda.empty_cache()\n",
        "\n",
        "      optimr.zero_grad()\n",
        "      loss.backward()\n",
        "      optimr.step()\n",
        "    \n",
        "    ### save params & frame\n",
        "\n",
        "    params_tmp = params[0].detach().clone()\n",
        "    if smooth is True:\n",
        "      opt_state = optimr.state_dict()\n",
        "\n",
        "    with torch.no_grad():\n",
        "      img_t = image_f(contrast=contrast, fixcontrast=resume)[0].permute(1,2,0)\n",
        "      img_np = torch.clip(img_t*255, 0, 255).cpu().numpy().astype(np.uint8)\n",
        "    imageio.imsave(os.path.join(tempdir, '%05d.jpg' % glob_step), img_np, quality=95)\n",
        "    shutil.copy(os.path.join(tempdir, '%05d.jpg' % glob_step), 'result.jpg')\n",
        "    outpic.clear_output()\n",
        "    with outpic:\n",
        "      display(Image('result.jpg'))\n",
        "    del img_t\n",
        "    pbar.upd()\n",
        "\n",
        "  params_tmp = params[0].detach().clone()\n",
        "\n",
        "outpic = ipy.Output()\n",
        "outpic\n",
        "\n",
        "pbar = ProgressBar(glob_steps)\n",
        "for i in range(count):\n",
        "  process(i)\n",
        "\n",
        "HTML(makevid(tempdir))\n",
        "files.download(tempdir + '.mp4')\n",
        "\n",
        "if save_depth and DepthStrength != 0:\n",
        "  HTML(makevid(depthdir))\n",
        "  files.download(depthdir + '.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Consider [post-processing by ESRGAN](https://colab.research.google.com/drive/1CKwPQacNjmNA_ydY6Be7xwFNJJWJCm_T) with [\"painting\" models](https://archive.org/details/hr-painting-upscaling) by Peter Baylies.  \n",
        "`pre_downscale` option helps to clean up pixel noise from these methods."
      ],
      "metadata": {
        "id": "KgXHJz1kXBbW"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BsehQRircaw7"
      },
      "source": [
        "If video is not auto-downloaded after generation (for whatever reason), run this cell to do that:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m0BZVNi8cZUP"
      },
      "source": [
        "files.download(tempdir + '.mp4')\n",
        "if save_depth and DepthStrength != 0:\n",
        "  files.download(depthdir + '.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}