LatentVision_rynmurdock.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "LatentVision_rynmurdock.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "mount_file_id": "1C_dnmhsdDYkOMhXCn6oL0gIUe2eoHABO",
      "authorship_tag": "ABX9TyPZH9/WXdqVcFQfDstksby3",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/olaviinha/NeuralTextToImage/blob/main/LatentVision_rynmurdock.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pVxFPluEoUka"
      },
      "source": [
        "#<font face=\"Trebuchet MS\" size=\"6\">Big Sleep: Latent vision <font color=\"#999\" size=\"4\">&nbsp;&nbsp;&nbsp;|&nbsp;&nbsp;&nbsp;</font><font color=\"#999\" size=\"4\">CLIP+VQGAN: Neural text-to-image</font><font color=\"#999\" size=\"4\">&nbsp;&nbsp;&nbsp;|&nbsp;&nbsp;&nbsp;</font><a href=\"https://github.com/olaviinha/NeuralImageGeneration\" target=\"_blank\"><font color=\"#999\" size=\"4\">Github</font></a>\n",
        "\n",
        "Big Sleep generates images from text input. It's originally a combination of [CLIP](https://github.com/openai/CLIP) by OpenAI and [BigGAN](https://arxiv.org/abs/1809.11096) by Andrew Brock et al., a concept introduced by [Ryan Murdock](https://github.com/rynmurdock) in his [original notebook](https://colab.research.google.com/drive/1NCceX2mbiKOSlAd_o7IU7nA9UskKN5WR?usp=sharing). This notebook is based on another similar implementation (CLIP+[VQGAN](https://arxiv.org/abs/2012.09841)) by the same author. It generates 2 images on each run.\n",
        "\n",
        "<hr size=\"1\" color=\"#666\">\n",
        "\n",
        "### Tips\n",
        "- Enter a simple string of text to `generate_image_of` field. You may also use a semicolon `;` as a separator to batch process multiple strings of texts to images in one go, and/or pipe `|` to train the image on multiple strings of text. If field is left empty, a random blog headline will be used.\n",
        "- Enter `output_dir` path relative to your Google Drive root, or leave blank to not save output anywhere outside this notebook. Each run of the _Sleep_ cell will **create a new subdirectory** under `output_dir`, under which all material will be saved.\n",
        "- In vast majority of cases, over 400 iterations seems to be a waste of time.\n",
        "- Setup cell will say you need to restart runtime. You can ignore it and not restart runtime."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4D3YFbKGoSje",
        "cellView": "form"
      },
      "source": [
        "#@title #Setup\n",
        "#@markdown This cell needs to be run only once. It will mount your Google Drive and setup prerequisites.\n",
        " \n",
        "force_setup = False\n",
        "pip_packages = 'kornia ftfy regex tqdm einops omegaconf==2.0.0 pytorch-lightning==1.0.8'\n",
        "main_repository = ''\n",
        " \n",
        "import os\n",
        "from google.colab import output\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "%cd /content/\n",
        " \n",
        "# inhagcutils\n",
        "if not os.path.isfile('/content/inhagcutils.ipynb') and force_setup == False:\n",
        "  !pip -q install import-ipynb {pip_packages}\n",
        "  !pip uninstall torchtext --yes\n",
        "  !curl -s -O https://raw.githubusercontent.com/olaviinha/inhagcutils/master/inhagcutils.ipynb\n",
        "import import_ipynb\n",
        "from inhagcutils import *\n",
        " \n",
        "# Mount Drive\n",
        "if not os.path.isdir('/content/drive') and force_setup == False:\n",
        "  from google.colab import drive\n",
        "  drive.mount('/content/drive')\n",
        " \n",
        "# Drive symlink\n",
        "if not os.path.isdir('/content/mydrive') and force_setup == False:\n",
        "  os.symlink('/content/drive/My Drive', '/content/mydrive')\n",
        "  drive_root_set = True\n",
        "drive_root = '/content/mydrive/'\n",
        " \n",
        "!git clone https://github.com/openai/CLIP.git\n",
        "!git clone https://github.com/CompVis/taming-transformers.git\n",
        " \n",
        "dir_tmp = '/content/tmp/'\n",
        "dir_steps1 = '/content/tmp/steps_1/'\n",
        "dir_steps2 = '/content/tmp/steps_2/'\n",
        "dir_initial = '/content/tmp/init/'\n",
        "dir_target = '/content/tmp/target/'\n",
        "create_dirs([dir_tmp, dir_steps1, dir_steps2,   dir_initial, dir_target])\n",
        " \n",
        "#-----\n",
        " \n",
        "import torch\n",
        "import numpy as np\n",
        "import torchvision\n",
        "import torchvision.transforms.functional as TF\n",
        "import kornia\n",
        " \n",
        "import PIL\n",
        "from PIL import ImageFile, Image\n",
        "import matplotlib.pyplot as plt\n",
        " \n",
        "import os\n",
        "import random\n",
        "import imageio\n",
        "from IPython import display\n",
        "from IPython.core.interactiveshell import InteractiveShell\n",
        "InteractiveShell.ast_node_interactivity = \"all\"\n",
        " \n",
        "import glob\n",
        "\n",
        "import itertools\n",
        "# from skimage import img_as_ubyte\n",
        "from subprocess import Popen, PIPE\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "#--\n",
        " \n",
        "from CLIP import clip\n",
        "perceptor, preprocess = clip.load('ViT-B/32', jit=False)\n",
        "perceptor.eval()\n",
        "clip.available_models()\n",
        "perceptor.visual.input_resolution\n",
        "scaler = 1\n",
        " \n",
        "im_shape = [512, 512, 3]\n",
        "sideX, sideY, channels = im_shape\n",
        "batch_size = 2\n",
        " \n",
        " \n",
        "#--\n",
        "\n",
        "def normalize8(I):\n",
        "  mn = I.min()\n",
        "  mx = I.max()\n",
        "  mx -= mn\n",
        "  I = ((I - mn)/mx) * 255\n",
        "  return I.astype(np.uint8)\n",
        " \n",
        "def displ(img, idx, i, dir, pre_scaled=True):\n",
        "  global dir_steps, dir_progress, save_all_steps, iterations\n",
        "  # dir = path_leaf(dir)+'_'+str(idx+1)+'/'\n",
        "  img = np.array(img)[:,:,:]\n",
        "  img = np.transpose(img, (1, 2, 0))\n",
        "  if not pre_scaled:\n",
        "    img = scale(img, 48*4, 32*4)\n",
        "  imgarr = np.array(img)\n",
        "  imgarr = normalize8(imgarr)\n",
        "  imageio.imwrite(dir + str(i).zfill(4) + '.png', imgarr)\n",
        "  if i == iterations:\n",
        "    imageio.imwrite(dir + str(i).zfill(4) + '.png', imgarr)\n",
        "  return display.Image(str(3)+'.png')\n",
        " \n",
        "def gallery(array, ncols=2):\n",
        "    nindex, height, width, intensity = array.shape\n",
        "    nrows = nindex//ncols\n",
        "    assert nindex == nrows*ncols\n",
        "    # want result.shape = (height*nrows, width*ncols, intensity)\n",
        "    result = (array.reshape(nrows, ncols, height, width, intensity)\n",
        "              .swapaxes(1,2)\n",
        "              .reshape(height*nrows, width*ncols, intensity))\n",
        "    return result\n",
        " \n",
        "def card_padded(im, to_pad=3):\n",
        "  return np.pad(np.pad(np.pad(im, [[1,1], [1,1], [0,0]],constant_values=0), [[2,2], [2,2], [0,0]],constant_values=1),\n",
        "            [[to_pad,to_pad], [to_pad,to_pad], [0,0]],constant_values=0)\n",
        " \n",
        "def get_all(img):\n",
        "  print('get all')\n",
        "  img = np.transpose(img, (0,2,3,1))\n",
        "  cards = np.zeros((img.shape[0], sideX+12, sideY+12, 3))\n",
        "  for i in range(len(img)):\n",
        "    cards[i] = card_padded(img[i])\n",
        "  print(img.shape)\n",
        "  cards = gallery(cards)\n",
        "  imageio.imwrite(str(3) + '.png', np.array(cards))\n",
        "  return display.Image(str(3)+'.png')\n",
        "  \n",
        " \n",
        "#-----\n",
        " \n",
        "%cd /content/taming-transformers\n",
        " \n",
        "!mkdir -p logs/vqgan_imagenet_f16_16384/checkpoints\n",
        "!mkdir -p logs/vqgan_imagenet_f16_16384/configs\n",
        " \n",
        "if len(os.listdir('logs/vqgan_imagenet_f16_16384/checkpoints/')) == 0:\n",
        "  !wget 'https://heibox.uni-heidelberg.de/f/867b05fc8c4841768640/?dl=1' -O 'logs/vqgan_imagenet_f16_16384/checkpoints/last.ckpt' \n",
        "  !wget 'https://heibox.uni-heidelberg.de/f/274fb24ed38341bfa753/?dl=1' -O 'logs/vqgan_imagenet_f16_16384/configs/model.yaml' \n",
        " \n",
        " \n",
        " \n",
        "# !cp /content/drive/MyDrive/vqgan_imagenet_f16_16384-20210325T002625Z-001.zip /content/vq.zip\n",
        "# !unzip /content/vq.zip -d /content/taming-transformers/logs/\n",
        " \n",
        "import sys\n",
        "sys.path.append(\".\")\n",
        " \n",
        "DEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        " \n",
        "import yaml\n",
        "import torch\n",
        "from omegaconf import OmegaConf\n",
        "from taming.models.vqgan import VQModel\n",
        " \n",
        "def load_config(config_path, display=False):\n",
        "  config = OmegaConf.load(config_path)\n",
        "  if display:\n",
        "    print(yaml.dump(OmegaConf.to_container(config)))\n",
        "  return config\n",
        " \n",
        "def load_vqgan(config, ckpt_path=None):\n",
        "  model = VQModel(**config.model.params)\n",
        "  if ckpt_path is not None:\n",
        "    sd = torch.load(ckpt_path, map_location=\"cpu\")[\"state_dict\"]\n",
        "    missing, unexpected = model.load_state_dict(sd, strict=False)\n",
        "  return model.eval()\n",
        " \n",
        "def preprocess_vqgan(x):\n",
        "  x = 2.*x - 1.\n",
        "  return x\n",
        " \n",
        "def custom_to_pil(x):\n",
        "  x = x.detach().cpu()\n",
        "  x = torch.clamp(x, -1., 1.)\n",
        "  x = (x + 1.)/2.\n",
        "  x = x.permute(1,2,0).numpy()\n",
        "  x = (255*x).astype(np.uint8)\n",
        "  x = Image.fromarray(x)\n",
        "  if not x.mode == \"RGB\":\n",
        "    x = x.convert(\"RGB\")\n",
        "  return x\n",
        " \n",
        "def reconstruct_with_vqgan(x, model):\n",
        "  # could also use model(x) for reconstruction but use explicit encoding and decoding here\n",
        "  z, _, [_, _, indices] = model.encode(x)\n",
        "  print(f\"VQGAN: latent shape: {z.shape[2:]}\")\n",
        "  xrec = model.decode(z)\n",
        "  return xrec\n",
        " \n",
        "class Pars(torch.nn.Module):\n",
        "  def __init__(self):\n",
        "    super(Pars, self).__init__()\n",
        "    self.normu = .5*torch.randn(batch_size, 256, sideX//16, sideY//16).cuda()\n",
        "    self.normu = torch.nn.Parameter(torch.sinh(1.9*torch.arcsinh(self.normu)))\n",
        "  def forward(self):\n",
        "    return self.normu.clip(-6, 6)\n",
        "      \n",
        "def model(x):\n",
        "  o_i2 = x\n",
        "  o_i3 = model16384.post_quant_conv(o_i2)\n",
        "  i = model16384.decoder(o_i3)\n",
        "  return i\n",
        " \n",
        "config16384 = load_config(\"logs/vqgan_imagenet_f16_16384/configs/model.yaml\", display=False)\n",
        "model16384 = load_vqgan(config16384, ckpt_path=\"logs/vqgan_imagenet_f16_16384/checkpoints/last.ckpt\").to(DEVICE)\n",
        " \n",
        " \n",
        " \n",
        "output.clear()\n",
        "# !nvidia-smi\n",
        "op(c.ok, 'Setup finished.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EBuZ6jhBo1ro",
        "cellView": "form"
      },
      "source": [
        "#@markdown <br>\n",
        " \n",
        "#@markdown #S̛̞̩͎͓ ̦̤͉͚̏ ̧̠͋͘ͅl͕̞͕̝͗̐͘.̠̰̳̫̈́̚ ̡͉̼̩̬̈́̇͒͘ȩ̨͎͛̔͆͊̏͜ͅ.͕̩̹̠̕͜ ̛̦̦̮e̢͐͊͂̀̊ͅ ̜̙̝̊͋ ̬̝̱̱͗p̮̎̽̌\n",
        " \n",
        "#@markdown <br>\n",
        " \n",
        "torch.cuda.empty_cache()\n",
        "clean_dirs([dir_steps1, dir_steps2, dir_initial, dir_target])\n",
        "generate_image_of = \"\" #@param {type:\"string\"}\n",
        "output_dir = '' #@param {type:\"string\"}\n",
        " \n",
        "#@markdown <hr color=\"#666\" size=\"1\">\n",
        "#@markdown <font size=\"1\">&nbsp;</font>\n",
        " \n",
        "#@markdown ### Advanced settings\n",
        " \n",
        "iterations = 400 #@param {type:\"slider\", min:0, max:2000, step:100}\n",
        "save_every = 50 #@param {type:\"slider\", min:0, max:500, step:1}\n",
        "display_save_every = True #@param {type:\"boolean\"}\n",
        "create_video = False #@param {type:\"boolean\"}\n",
        " \n",
        "save_all_steps = False\n",
        "remove_interrupted = True\n",
        "repetitions = 0\n",
        " \n",
        "#\n",
        "# --- Very advanced settings ---------------------------------\n",
        "\n",
        "# #@markdown <hr color=\"#666\" size=\"1\">\n",
        "# #@markdown <font size=\"1\">&nbsp;</font>\n",
        " \n",
        "# #@markdown ### Very advanced settings\n",
        "# repetitions = 0 #@param {type:\"slider\", min:0, max:20, step:1}\n",
        "# save_all_steps = False #@param {type:\"boolean\"}\n",
        "# remove_interrupted = True #@param {type:\"boolean\"}\n",
        "\n",
        "# --- //Very advanced settings -------------------------------\n",
        "#\n",
        " \n",
        "\n",
        " \n",
        "text = generate_image_of\n",
        "iterations = iterations+1\n",
        "\n",
        "if save_all_steps is True:\n",
        "  dir_steps1 = dir_output+'steps_1/'\n",
        "  dir_steps1 = dir_output+'steps_2/'\n",
        "  create_dirs([dir_steps1, dir_steps2])\n",
        " \n",
        "if output_dir is '' or not output_dir:\n",
        "  drive_root = '/content/fauxdrive/'\n",
        "  output_dir = 'output'\n",
        "  create_dirs([drive_root, output_dir])\n",
        "else:\n",
        "  drive_root = '/content/mydrive/'\n",
        " \n",
        "if text is not '':\n",
        "  if check_input_type(drive_root+text) is 'file':\n",
        "    input_txt = drive_root+text\n",
        "    with open(input_txt) as f:\n",
        "      texts = f.readlines()\n",
        "      texts = [x.strip() for x in texts] \n",
        "  elif \";\" in text:\n",
        "    texts = text.split(';')\n",
        "    texts = [text.strip() for text in texts]\n",
        "  else:\n",
        "    texts = [text]\n",
        "else:\n",
        "  text = requests.get('https://api.inha.asia/headline/').text\n",
        "  texts = [text]\n",
        "if repetitions > 0:\n",
        "  texts = list(itertools.chain.from_iterable(itertools.repeat(x, repetitions) for x in texts))\n",
        " \n",
        "uniq_id = gen_id()\n",
        "repeat_index = 1\n",
        " \n",
        "#---------------\n",
        " \n",
        "def keyboardInterruptHandler():\n",
        "  global dir_output, uniq_id\n",
        "  op(c.warn, 'Interrupted!', 'Cleaning up...')\n",
        "  remove_dirs([dir_output])\n",
        "  print('Run', uniq_id, 'directory and content removed:', dir_output)\n",
        "  sys.exit()\n",
        "\n",
        "def augment(into, cutn=32):\n",
        "  into = torch.nn.functional.pad(into, (sideX//2, sideX//2, sideX//2, sideX//2), mode='constant', value=0)\n",
        "  into = augs(into)\n",
        "  p_s = []\n",
        "  for ch in range(cutn):\n",
        "    size = int(torch.normal(1.2, .3, ()).clip(.43, 1.9) * sideX)\n",
        "    if ch > cutn - 4:\n",
        "      size = int(sideX*1.4)\n",
        "    offsetx = torch.randint(0, int(sideX*2 - size), ())\n",
        "    offsety = torch.randint(0, int(sideX*2 - size), ())\n",
        "    apper = into[:, :, offsetx:offsetx + size, offsety:offsety + size]\n",
        "    apper = torch.nn.functional.interpolate(apper, (int(224*scaler), int(224*scaler)), mode='bilinear', align_corners=True)\n",
        "    p_s.append(apper)\n",
        "  into = torch.cat(p_s, 0)\n",
        "  into = into + up_noise*torch.rand((into.shape[0], 1, 1, 1)).cuda()*torch.randn_like(into, requires_grad=False)\n",
        "  return into\n",
        "\n",
        "def checkin(loss, i, xtype, display_image):\n",
        "  global up_noise, dir_progress1, dir_progress2, dir_steps1, dir_steps2\n",
        "  if xtype is 'step':\n",
        "    xdir = dir_steps1.replace('_1/', '')\n",
        "  if xtype is 'progress':\n",
        "    xdir = dir_progress1.replace('_1/', '')\n",
        "  with torch.no_grad():\n",
        "    alnot = model(lats()).float()\n",
        "    alnot = augment((((alnot).clip(-1, 1) + 1) / 2), cutn=1)\n",
        "    # for allls in alnot.cpu():\n",
        "    #   displ(allls)\n",
        "    #   display.display(display.Image(str(3)+'.png'))\n",
        "    alnot = (model(lats()).cpu().clip(-1, 1) + 1) / 2\n",
        "    for idx, allls in enumerate(alnot.cpu()):\n",
        "      dir = xdir.replace(path_leaf(xdir), path_leaf(xdir)+'_'+str(idx+1)+'/')\n",
        "      displ(allls, idx, i, dir)\n",
        "      if display_image is True: \n",
        "        display.display(display.Image(dir+str(i).zfill(4)+'.png'))\n",
        "        op(c.ok, '^Iteration image saved as:', dir.replace(drive_root, '')+str(i).zfill(4)+'.png\\n')\n",
        "      elif xtype is 'progress':\n",
        "        op(c.ok, 'Iteration image saved as:', dir.replace(drive_root, '')+str(i).zfill(4)+'.png\\n')\n",
        " \n",
        "def ascend_txt():\n",
        "  global up_noise\n",
        "  out = model(lats())\n",
        "  into = augment((out.clip(-1, 1) + 1) / 2)\n",
        "  into = nom(into)\n",
        "  iii = perceptor.encode_image(into)\n",
        "  t_x = torch.cosine_similarity(t_not, iii, -1)\n",
        "  all_s = torch.cosine_similarity(t, iii, -1)\n",
        "  return [3*t_x, -10*all_s]\n",
        "  \n",
        "def train(i):\n",
        "  global dec, up_noise\n",
        "  loss1 = ascend_txt()\n",
        "  loss = loss1[0] + loss1[1]\n",
        "  loss = loss.mean()\n",
        "  optimizer.zero_grad()\n",
        "  loss.backward()\n",
        "  optimizer.step()\n",
        "  if i > 400:\n",
        "    for g in optimizer.param_groups:\n",
        "      g['lr'] *= .995\n",
        "      g['lr'] = max(g['lr'], .1)\n",
        "    dec *= .995\n",
        "  if torch.abs(lats()).max() > 5:\n",
        "    for g in optimizer.param_groups:\n",
        "      g['weight_decay'] = dec\n",
        "  else:\n",
        "    for g in optimizer.param_groups:\n",
        "      g['weight_decay'] = 0\n",
        "  # if itt % 100 == 0:\n",
        "  #   checkin(loss1)\n",
        "  if save_all_steps is True or create_video is True or i == iterations-1:\n",
        "    checkin(loss1, i, 'step', display_image=False)\n",
        "  if i % save_every == 0:\n",
        "    checkin(loss1, i, 'progress', display_image=display_save_every)\n",
        " \n",
        "def loop():\n",
        "  global itt, iterations\n",
        "  try:\n",
        "    for asatreat in range(iterations):\n",
        "      train(itt)\n",
        "      itt+=1\n",
        "  # pbar = tqdm(total=iterations)\n",
        "  #   for i in range(10):\n",
        "  #     while True:\n",
        "  #       train(itt)\n",
        "  #       if itt == iterations:\n",
        "  #         break\n",
        "  #       itt += 1\n",
        "  #       tqdm.update()\n",
        "  except KeyboardInterrupt:\n",
        "    if remove_interrupted: keyboardInterruptHandler()\n",
        "    \n",
        "#---------------\n",
        " \n",
        "for text_input in texts:\n",
        " \n",
        "  torch.cuda.empty_cache()\n",
        " \n",
        "  display_text = text_input\n",
        "  title = text_input.split(\"|\")[0].title()\n",
        "  file_title = ''.join(e for e in title if e.isalnum())\n",
        " \n",
        "  id = uniq_id+'_'+file_title\n",
        "  if repetitions > 0:\n",
        "    id = uniq_id+'_'+str(repeat_index)+'_'+file_title\n",
        " \n",
        "  dir_output = fix_path(drive_root+output_dir)+id+'/'\n",
        "  dir_progress1 = dir_output+'progress_1/'\n",
        "  dir_progress2 = dir_output+'progress_2/'\n",
        "  create_dirs([dir_output, dir_progress1, dir_progress2])\n",
        " \n",
        "  text_not = '''disconnected, confusing, incoherent, watermarks, text, writing'''\n",
        " \n",
        "  dec = .1\n",
        "  lats = Pars().cuda()\n",
        "  mapper = [lats.normu]\n",
        "  optimizer = torch.optim.AdamW([{'params': mapper, 'lr': .5}], weight_decay=dec)\n",
        "  eps = 0\n",
        "  tx = clip.tokenize(text_input)\n",
        "  t = perceptor.encode_text(tx.cuda()).detach().clone()\n",
        "  t_not = clip.tokenize(text_not)\n",
        "  t_not = perceptor.encode_text(t_not.cuda()).detach().clone()\n",
        "  nom = torchvision.transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))\n",
        " \n",
        "  augs = torch.nn.Sequential(\n",
        "      torchvision.transforms.RandomHorizontalFlip(),\n",
        "      torchvision.transforms.RandomAffine(24, (.1, .1), fill=0)\n",
        "  ).cuda()\n",
        " \n",
        "  up_noise = .11\n",
        "  itt = 0\n",
        " \n",
        "  output.clear()\n",
        "  op(c.title, '\\nGenerating image of', display_text)\n",
        "  op(c.title, 'Run ID:', uniq_id)\n",
        "  if repetitions > 0:\n",
        "    op(c.title, 'Repetition:', repeat_index)\n",
        "  op(c.okb, 'Sweet dreams.\\n')\n",
        " \n",
        "  with torch.no_grad():\n",
        "    al = (model(lats()).cpu().clip(-1, 1) + 1) / 2\n",
        "    for allls in al:\n",
        "      # displ(allls[:3])\n",
        "      print('\\n')\n",
        " \n",
        "  loop()\n",
        "\n",
        "  last_step1 = dir_steps1+str(iterations-1).zfill(4)+'.png'\n",
        "  last_step2 = dir_steps2+str(iterations-1).zfill(4)+'.png'\n",
        "  fin_out1 = dir_output+file_title+'_1.png'\n",
        "  fin_out2 = dir_output+file_title+'_2.png'\n",
        "  !cp {last_step1} {fin_out1}\n",
        "  !cp {last_step2} {fin_out2}\n",
        "  display_fin1 = fin_out1.replace(drive_root, '')\n",
        "  display_fin2 = fin_out2.replace(drive_root, '')\n",
        "  op(c.ok, '\\nFinal images saved as')\n",
        "  print('-', display_fin1)\n",
        "  print('-', display_fin2)\n",
        "\n",
        "  if create_video is True:\n",
        "    op(c.title, '\\nGenerating video')\n",
        "\n",
        "    init_frame = 1\n",
        "    last_frame = iterations-1\n",
        "\n",
        "    fps = 30\n",
        "    output_video1 = dir_output+file_title+'_1.mp4'\n",
        "    output_video2 = dir_output+file_title+'_2.mp4'\n",
        "\n",
        "    frames1 = []\n",
        "    for i in range(init_frame,last_frame): #\n",
        "      filename1 = f\"{dir_steps1}/{i:04}.png\"\n",
        "      frames1.append(Image.open(filename1))\n",
        "    p = Popen(['ffmpeg', '-y', '-f', 'image2pipe', '-vcodec', 'png', '-r', str(fps), '-i', '-', '-vcodec', 'libx264', '-r', str(fps), '-pix_fmt', 'yuv420p', '-crf', '13', '-preset', 'veryslow', output_video1], stdin=PIPE)\n",
        "    for im in tqdm(frames1):\n",
        "      im.save(p.stdin, 'PNG')\n",
        "    p.stdin.close()\n",
        "    p.wait()\n",
        "    fin_vid1 = fin_out1.replace('.png', '.mp4')\n",
        "\n",
        "    frames2 = []\n",
        "    for i in range(init_frame,last_frame): #\n",
        "      filename2 = f\"{dir_steps2}/{i:04}.png\"\n",
        "      frames2.append(Image.open(filename2))\n",
        "    p = Popen(['ffmpeg', '-y', '-f', 'image2pipe', '-vcodec', 'png', '-r', str(fps), '-i', '-', '-vcodec', 'libx264', '-r', str(fps), '-pix_fmt', 'yuv420p', '-crf', '13', '-preset', 'veryslow', output_video2], stdin=PIPE)\n",
        "    for im in tqdm(frames2):\n",
        "      im.save(p.stdin, 'PNG')\n",
        "    p.stdin.close()\n",
        "    p.wait()\n",
        "    fin_vid2 = fin_out2.replace('.png', '.mp4')\n",
        "\n",
        "    op(c.ok, 'Videos saved as')\n",
        "    print('-', fin_vid1)\n",
        "    print('-', fin_vid2)\n",
        "\n",
        "  if repeat_index is repetitions:\n",
        "    repeat_index = 1\n",
        "  else:\n",
        "    repeat_index += 1\n",
        "\n",
        "op(c.title, '\\nFIN.')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}