Aphantasia.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Aphantasia.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.9"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "toWe1IoH7X35"
      },
      "source": [
        "# Text to Image tool\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), [Hannu Toyryla](https://twitter.com/htoyryla), [@eduwatch2](https://twitter.com/eduwatch2) for ideas.\n",
        "\n",
        "## Features \n",
        "* generates massive detailed imagery, a la deepdream \n",
        "* high resolution (up to 12K on RTX 3090)\n",
        "* directly parameterized with [FFT](https://github.com/greentfrapp/lucent/blob/master/lucent/optvis/param/spatial.py) [Fourier] or DWT [wavelets], without pretrained GANs\n",
        "* various CLIP models (including multi-language from [SBERT](https://sbert.net))\n",
        "* starting/resuming process from saved FFT parameters or from an image\n",
        "* using [aesthetic loss](https://github.com/LAION-AI/aesthetic-predictor) \n",
        "* complex requests:\n",
        "  * image and/or text as main prompts {with weights)  \n",
        "   (composition similarity controlled with [LPIPS](https://github.com/richzhang/PerceptualSimilarity) loss)\n",
        "  * extra text prompts for image style and to subtract (suppress) topics\n",
        "  * criteria inversion (show \"the opposite\")\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.  \n",
        "Or you can upload any image to start from - that is, to \"stylize\" that image.  \n",
        "Resolution settings below will be ignored in this case. "
      ]
    },
    {
      "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",
        "import os\n",
        "import io\n",
        "import time\n",
        "from math import exp\n",
        "import random\n",
        "import imageio\n",
        "import numpy as np\n",
        "import PIL\n",
        "from base64 import b64encode\n",
        "# import moviepy, moviepy.editor\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 sentence_transformers\n",
        "from sentence_transformers import SentenceTransformer\n",
        "!pip install kornia\n",
        "import kornia\n",
        "!pip install lpips\n",
        "import lpips\n",
        "\n",
        "!pip install PyWavelets==1.1.1\n",
        "!pip install git+https://github.com/fbcotter/pytorch_wavelets\n",
        "import pywt\n",
        "from pytorch_wavelets import DWTForward, DWTInverse\n",
        "# from pytorch_wavelets import DTCWTForward, DTCWTInverse\n",
        "\n",
        "%cd /content\n",
        "!pip install git+https://github.com/eps696/aphantasia\n",
        "from aphantasia.image import to_valid_rgb, fft_image, img2fft, dwt_image, img2dwt, init_dwt, dwt_scale\n",
        "from aphantasia.utils import slice_imgs, derivat, pad_up_to, aesthetic_model, basename, img_list, img_read, plot_text, txt_clean, checkout, old_torch\n",
        "from aphantasia import transforms\n",
        "from aphantasia.progress_bar import ProgressIPy as ProgressBar\n",
        "\n",
        "clear_output()\n",
        "\n",
        "resume = False #@param {type:\"boolean\"}\n",
        "if resume:\n",
        "  resumed = files.upload()\n",
        "  resumed_filename = list(resumed)[0]\n",
        "  resumed_bytes = list(resumed.values())[0]\n",
        "\n",
        "def makevid(seq_dir, size=None):\n",
        "  out_sequence = seq_dir + '/%04d.jpg'\n",
        "  out_video = seq_dir + '.mp4'\n",
        "  print('.. generating video ..')\n",
        "  !ffmpeg -y -v warning -i $out_sequence -crf 20 $out_video\n",
        "  # moviepy.editor.ImageSequenceClip(img_list(seq_dir), fps=25).write_videofile(out_video, verbose=False)\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]\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": "markdown",
      "metadata": {
        "id": "CbJ9K4Cq8MtB"
      },
      "source": [
        "Type some `text` and/or upload some image to start.  \n",
        "Describe `style`, which you'd like to apply to the imagery.  \n",
        "Put to `subtract` the topics, which you would like to avoid in the result.  \n",
        "All inputs understand syntax with weights: `goodprompt:1 | alsogoodprompt:1 | badprompt:-0.5`.  \n",
        "`invert` the whole criteria, if you want to see \"the totally opposite\".\n",
        "\n",
        "Options for non-English languages (use only one of them!):  \n",
        "`multilang` = use multi-language model, trained with ViT  \n",
        "`translate` = use Google translate (works with any visual model)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JUvpdy8BWGuM",
        "cellView": "form"
      },
      "source": [
        "#@title Input\n",
        "\n",
        "text = \"\" #@param {type:\"string\"}\n",
        "style = \"\" #@param {type:\"string\"}\n",
        "subtract = \"\" #@param {type:\"string\"}\n",
        "multilang = False #@param {type:\"boolean\"}\n",
        "translate = False #@param {type:\"boolean\"}\n",
        "invert = False #@param {type:\"boolean\"}\n",
        "upload_image = 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_image:\n",
        "  uploaded = files.upload()\n",
        "\n",
        "workdir = '_out'\n",
        "tempdir = os.path.join(workdir, '%s-%s' % (txt_clean(text)[:50], txt_clean(style)[:50]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Settings"
      ],
      "metadata": {
        "id": "7xEoTPlyk1SD"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f3Sj0fxmtw6K"
      },
      "source": [
        "If needed, select visual `model` (results do vary!). Dual (ViT-B/32 + ViT-B/16) usually works best. Otherwise, I prefer ViT for consistency (and it's the only native multi-language option).  \n",
        "`align` option is about composition. `uniform` looks most adequate, `overscan` can make semi-seamless tileable texture.  \n",
        "`aesthetic` enforces overall cuteness (try various values!). May be set negative.  \n",
        "`use_wavelets` for DWT encoding instead of FFT. Select `wave` method if needed.  \n",
        "`aug_transform` applies some augmentations, inhibiting image fragmentation & \"graffiti\" printing (slower, yet recommended).  \n",
        "`sync` value adds LPIPS loss between the output and input image (if there's one), allowing to \"redraw\" it with controlled similarity.  \n",
        "Decrease `samples` if you face OOM (it's the main RAM eater).  \n",
        "\n",
        "Select optimizer: `_custom` options are more stable but noisy; pure `adam` is softer, but may spill some colored blur on longer runs.  \n",
        "Setting `steps` much higher (1000-..) will elaborate details and make tones smoother, but may start throwing texts like graffiti.  \n",
        "Tune `decay` (compositional softness) and `sharpness`, `colors` (saturation) and `contrast` as needed.  \n",
        "\n",
        "Experimental tricks:  \n",
        "`aug_noise` augmentation, `macro` (from 0 to 1) and `progressive_grow` ([read more](https://github.com/eps696/aphantasia/issues/2)) may boost bigger forms, making composition less disperse.  \n",
        "`enhance` boosts training consistency (of simultaneous samples) and steps progress. good start is 0.1~0.2.  "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nq0wA-wc-P-s",
        "cellView": "form"
      },
      "source": [
        "sideX = 1280 #@param {type:\"integer\"}\n",
        "sideY = 720 #@param {type:\"integer\"}\n",
        "#@markdown > Config\n",
        "model = 'dual' #@param ['dual', 'ViT-B/16', 'ViT-B/32', 'RN101', 'RN50x16', 'RN50x4', 'RN50']\n",
        "align = 'uniform' #@param ['central', 'uniform', 'overscan']\n",
        "aesthetic =  1.#@param {type:\"number\"}\n",
        "use_wavelets = False #@param {type:\"boolean\"}\n",
        "wave = 'coif1' #@param ['db2', 'db3', 'db4', 'coif1', 'coif2', 'coif3', 'coif4']\n",
        "aug_transform = True #@param {type:\"boolean\"}\n",
        "sync =  0.4 #@param {type:\"number\"}\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",
        "#@markdown > Training\n",
        "steps = 150 #@param {type:\"integer\"}\n",
        "samples = 200 #@param {type:\"integer\"}\n",
        "learning_rate = .05 #@param {type:\"number\"}\n",
        "optimizer = 'adam_custom' #@param ['adam', 'adam_custom', 'adamw', 'adamw_custom']\n",
        "save_freq = 1 #@param {type:\"integer\"}\n",
        "#@markdown > Tricks\n",
        "aug_noise = 0. #@param {type:\"number\"}\n",
        "enhance = 0. #@param {type:\"number\"}\n",
        "macro = 0.4 #@param {type:\"number\"}\n",
        "progressive_grow = False #@param {type:\"boolean\"}\n",
        "if multilang: model = 'ViT-B/32' # sbert model is trained with ViT\n",
        "diverse = -enhance\n",
        "expand = abs(enhance)\n",
        "\n",
        "if model == 'dual':\n",
        "  dualmod = 2\n",
        "  model = 'ViT-B/32'\n",
        "else:\n",
        "  dualmod = None\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Generate"
      ],
      "metadata": {
        "id": "GP8RWuMFj8LO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title run\n",
        "\n",
        "!rm -rf $tempdir\n",
        "os.makedirs(tempdir, exist_ok=True)\n",
        "\n",
        "shape = [1, 3, sideY, sideX]\n",
        "\n",
        "if use_wavelets:\n",
        "  if resume is True:\n",
        "    if os.path.splitext(resumed_filename)[1].lower()[1:] in ['jpg','png','tif','bmp']:\n",
        "      img_in = imageio.imread(resumed_bytes)\n",
        "      init_pt = img2dwt(img_in, wave, 1.5)\n",
        "      scale = dwt_scale(init_pt, sharpness)\n",
        "      for i in range(len(init_pt)-1):\n",
        "        init_pt[i+1] /= (scale[i] * 10)\n",
        "      sideY, sideX = img_in.shape[0], img_in.shape[1]\n",
        "    else:\n",
        "      init_pt = torch.load(io.BytesIO(resumed_bytes))\n",
        "      # init_pt = [y.detach().cuda() for y in init_pt]\n",
        "  else:\n",
        "    init_pt, _, _, _ = init_dwt(None, shape, wave, colors)\n",
        "\n",
        "  shape = [1, 3, sideY, sideX]\n",
        "  params, image_f, _ = dwt_image(shape, wave, sharpness, colors, init_pt)\n",
        "\n",
        "else:\n",
        "  if resume is True:\n",
        "    if os.path.splitext(resumed_filename)[1].lower()[1:] in ['jpg','png','tif','bmp']:\n",
        "      img_in = imageio.imread(resumed_bytes)\n",
        "      init_pt = img2fft(img_in, 1.5, 1.5) * 0.1\n",
        "    else:\n",
        "      init_pt = torch.load(io.BytesIO(resumed_bytes))\n",
        "      if isinstance(init_pt, list): init_pt = init_pt[0]\n",
        "  #  init_pt = init_pt.cuda()\n",
        "    sideY, sideX = init_pt.shape[2], (init_pt.shape[3]-1)*2\n",
        "  else:\n",
        "    params_shape = [1, 3, sideY, sideX//2+1, 2]\n",
        "    init_pt = torch.randn(*params_shape) * 0.01\n",
        "\n",
        "  shape = [1, 3, sideY, sideX]\n",
        "  params, image_f, _ = fft_image(shape, 1, decay, init_pt)\n",
        "\n",
        "image_f = to_valid_rgb(image_f, colors = colors)\n",
        "\n",
        "if progressive_grow is True:\n",
        "  lr1 = learning_rate * 2\n",
        "  lr0 = lr1 * 0.01\n",
        "else:\n",
        "  lr0 = learning_rate\n",
        "\n",
        "if optimizer.lower() == 'adamw':\n",
        "  optimr = torch.optim.AdamW(params, lr0, weight_decay=0.01, amsgrad=True)\n",
        "elif optimizer.lower() == 'adamw_custom':\n",
        "  optimr = torch.optim.AdamW(params, lr0, weight_decay=0.01, amsgrad=True, betas=(.0,.999))\n",
        "elif optimizer.lower() == 'adam':\n",
        "  optimr = torch.optim.Adam(params, lr0)\n",
        "else: # adam_custom\n",
        "  optimr = torch.optim.Adam(params, lr0, betas=(.0,.999))\n",
        "\n",
        "if len(subtract) > 0:\n",
        "  samples = int(samples * 0.75)\n",
        "print(' using %d samples,' % samples, optimizer, 'optimizer')\n",
        "\n",
        "model_clip, _ = clip.load(model, jit=old_torch())\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",
        "  samples = int(samples * xmem[model])\n",
        "\n",
        "if multilang is True:\n",
        "    model_lang = SentenceTransformer('clip-ViT-B-32-multilingual-v1').cuda()\n",
        "\n",
        "if dualmod is not None: # second is vit-16\n",
        "  model_clip2, _ = clip.load('ViT-B/16', jit=old_torch())\n",
        "  samples = int(samples * 0.25)\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",
        "def enc_text(txt, model_clip=model_clip):\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",
        "    if multilang is True:\n",
        "      emb = model_lang.encode([subtxt], convert_to_tensor=True, show_progress_bar=False)\n",
        "    else:\n",
        "      emb = model_clip.encode_text(clip.tokenize(subtxt).cuda())\n",
        "    embs.append([emb.detach().clone(), wt])\n",
        "  return embs\n",
        "\n",
        "if diverse != 0:\n",
        "  samples = int(samples * 0.5)\n",
        "if sync > 0 and upload_image:\n",
        "  samples = int(samples * 0.6)\n",
        "        \n",
        "sign = 1. if invert is True else -1.\n",
        "if aug_transform is True:\n",
        "  trform_f = transforms.transforms_fast\n",
        "  samples = int(samples * 0.95)\n",
        "else:\n",
        "  trform_f = transforms.normalize()\n",
        "\n",
        "if upload_image:\n",
        "  in_img = list(uploaded.values())[0]\n",
        "  print(' image:', list(uploaded)[0])\n",
        "  img_in = torch.from_numpy(imageio.imread(in_img).astype(np.float32)/255.).unsqueeze(0).permute(0,3,1,2).cuda()[:,:3,:,:]\n",
        "  in_sliced = slice_imgs([img_in], samples, modsize, transforms.normalize(), align)[0]\n",
        "  img_enc = model_clip.encode_image(in_sliced).detach().clone()\n",
        "  if dualmod is not None:\n",
        "    img_enc2 = model_clip2.encode_image(in_sliced).detach().clone()\n",
        "  if sync > 0:\n",
        "    align = 'overscan'\n",
        "    sim_loss = lpips.LPIPS(net='vgg', verbose=False).cuda()\n",
        "    sim_size = [sideY//4, sideX//4]\n",
        "    img_in = F.interpolate(img_in, sim_size).float()\n",
        "    # img_in = F.interpolate(img_in, (sideY, sideX)).float()\n",
        "  else:\n",
        "    del img_in\n",
        "  del in_sliced; torch.cuda.empty_cache()\n",
        "\n",
        "if len(text) > 0:\n",
        "  print(' main topic:', text)\n",
        "  if translate:\n",
        "    text = translator.translate(text, dest='en').text\n",
        "    print(' translated to:', text) \n",
        "  txt_enc = enc_text(text)\n",
        "  if dualmod is not None:\n",
        "    txt_enc2 = enc_text(text, model_clip2)\n",
        "\n",
        "if len(style) > 0:\n",
        "  print(' style:', style)\n",
        "  if translate:\n",
        "    style = translator.translate(style, dest='en').text\n",
        "    print(' translated to:', style) \n",
        "  style_enc = enc_text(style)\n",
        "  if dualmod is not None:\n",
        "    style_enc2 = enc_text(style, model_clip2)\n",
        "\n",
        "if len(subtract) > 0:\n",
        "  print(' without:', subtract)\n",
        "  if translate:\n",
        "    subtract = translator.translate(subtract, dest='en').text\n",
        "    print(' translated to:', subtract) \n",
        "  not_enc = enc_text(subtract)\n",
        "  if dualmod is not None:\n",
        "    not_enc2 = enc_text(subtract, model_clip2)\n",
        "\n",
        "if multilang is True: del model_lang\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 checkout(num):\n",
        "  with torch.no_grad():\n",
        "    img = image_f(contrast=contrast).cpu().numpy()[0]\n",
        "    # empirical tone mapping\n",
        "    if sync > 0 and upload_image:\n",
        "      img = img **1.3 \n",
        "    # if sharpness != 0:\n",
        "      # img = img ** (1 + sharpness/2.)\n",
        "  save_img(img, os.path.join(tempdir, '%04d.jpg' % num))\n",
        "  outpic.clear_output()\n",
        "  with outpic:\n",
        "    display(Image('result.jpg'))\n",
        "\n",
        "prev_enc = 0\n",
        "def train(i):\n",
        "  loss = 0\n",
        "  noise = aug_noise * torch.randn(1, 1, *params[0].shape[2:4], 1).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=macro)[0]\n",
        "\n",
        "  if len(text) > 0: # input text\n",
        "    txt_enc_    = txt_enc2      if dualmod is not None and i in dualmod_nums else txt_enc\n",
        "  if len(style) > 0:\n",
        "    style_enc_  = style_enc2    if dualmod is not None and i in dualmod_nums else style_enc\n",
        "  if upload_image:\n",
        "    img_enc_    = img_enc2      if dualmod is not None and i in dualmod_nums else img_enc\n",
        "  if len(subtract) > 0:\n",
        "    not_enc_    = not_enc2      if dualmod is not None and i in dualmod_nums else not_enc\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",
        "\n",
        "  if len(text) > 0: # input text\n",
        "    for enc, wt in txt_enc_:\n",
        "      loss += sign * wt * torch.cosine_similarity(enc, out_enc, dim=-1).mean()\n",
        "  if len(style) > 0: # input text - style\n",
        "    for enc, wt in style_enc_:\n",
        "      loss += sign * wt * torch.cosine_similarity(enc, out_enc, dim=-1).mean()\n",
        "  if len(subtract) > 0: # subtract text\n",
        "    for enc, wt in not_enc_:\n",
        "      loss -= sign * wt * torch.cosine_similarity(enc, out_enc, dim=-1).mean()\n",
        "  if upload_image:\n",
        "    loss += sign * 0.7 * torch.cosine_similarity(img_enc_, out_enc, dim=-1).mean()\n",
        "  if sync > 0 and upload_image: # image composition sync\n",
        "    prog_sync = (steps - i) / steps \n",
        "    loss += prog_sync * sync * sim_loss(F.interpolate(img_out, sim_size).float(), img_in, normalize=True).squeeze()\n",
        "  if sharpness != 0 and not use_wavelets: # mode = scharr|sobel|default\n",
        "    loss -= sharpness * derivat(img_out, mode='sobel')\n",
        "    # loss -= sharpness * derivat(img_sliced, mode='scharr')\n",
        "  if diverse != 0:\n",
        "    img_sliced = slice_imgs([image_f(noise)], samples, modsize, trform_f, align, macro=macro)[0]\n",
        "    out_enc2 = model_clip_.encode_image(img_sliced)\n",
        "    loss += diverse * torch.cosine_similarity(out_enc, out_enc2, dim=-1).mean()\n",
        "    del out_enc2; torch.cuda.empty_cache()\n",
        "  if expand > 0:\n",
        "    global prev_enc\n",
        "    if i > 0:\n",
        "      loss += expand * torch.cosine_similarity(out_enc, prev_enc, dim=-1).mean()\n",
        "    prev_enc = out_enc.detach()\n",
        "  if aesthetic != 0 and aest_ is not None:\n",
        "    loss -= 0.001 * aesthetic * aest_(out_enc).mean()\n",
        "  del img_out, img_sliced, out_enc; torch.cuda.empty_cache()\n",
        "\n",
        "  if progressive_grow is True:\n",
        "    lr_cur = lr0 + (i / steps) * (lr1 - lr0)\n",
        "    for g in optimr.param_groups: \n",
        "      g['lr'] = lr_cur\n",
        "\n",
        "  optimr.zero_grad()\n",
        "  loss.backward()\n",
        "  optimr.step()\n",
        "  \n",
        "  if i % save_freq == 0:\n",
        "    checkout(i // save_freq)\n",
        "\n",
        "outpic = ipy.Output()\n",
        "outpic\n",
        "\n",
        "pbar = ProgressBar(steps)\n",
        "for i in range(steps):\n",
        "  train(i)\n",
        "  _ = pbar.upd()\n",
        "\n",
        "HTML(makevid(tempdir))\n",
        "torch.save(params, tempdir + '.pt')\n",
        "files.download(tempdir + '.pt')\n",
        "files.download(tempdir + '.mp4')\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "81P3Vl6ijnQJ"
      },
      "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.\n"
      ],
      "metadata": {
        "id": "s1_ppReXVwmN"
      }
    }
  ]
}