pyttitools-PYTTI.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WksDx4_To_su"
      },
      "source": [
        "# PYTTI-TOOLS!\n",
        "\n",
        "## A brief history of this notebook\n",
        "\n",
        "The tools and techniques below were pioneered in 2021 by a diverse and distributed collection of amazingly talented ML practitioners, researchers, and artists. The short version of this history is that Katherine Crowson ([@RiversHaveWings](https://twitter.com/RiversHaveWings)) published a notebook inspired by work done by [@advadnoun](https://twitter.com/advadnoun). Katherine's notebook spawned a litany of variants, each with their own twist on the technique or adding a feature to someone else's work. Henry Rachootin ([@sportsracer48](https://twitter.com/sportsracer48)) collected several of the most interesting notebooks and stuck the important bits together with bublegum and scotch tape. Thus was born PYTTI, and there was much rejoicing in sportsracer48's patreon, where it was shared in closed beta for several months so sportsracer48 wouldn't get buried under tech support requests (or so he hoped).\n",
        "\n",
        "PYTTI rapidly gained a reputation as one of the most powerful tools available for generating CLIP-guided images. In late November, @sportsracer48 released the last version in his closed beta: the \"pytti 5 beta\" notebook. David Marx ([@DigThatData](https://twitter.com/DigThatData)) offered to help tidy up the mess a few weeks later, and sportsracer48 encouraged him to run wild with it. Henry didn't realize he'd been speaking with someone who had recently quit their job and had a lot of time on their hands, and David's contributions snowballed into [PYTTI-Tools](https://github.com/pytti-tools)!\n",
        "\n",
        "## How is PYTTI-Tools different from PYTTI 5 Beta?\n",
        "\n",
        "Right now, not very. The main user-visible changes are:\n",
        "\n",
        "* Local use is now a first-class citizen\n",
        "* PyTTI is installable and can be run as a CLI tool\n",
        "* Using PyTTI on the command line gives you magic powers\n",
        "* PyTTI supports tensorboard, meaning it also integrates with tools like MLFlow and WandB\n",
        "* Bug fixes and slightly saner code\n",
        "\n",
        "## Call to action!\n",
        "\n",
        "My hope is that rather than continuing to fork off messy notebooks with minor changes between them, pytti-tools will become a central hub for organizing and sharing related techniques for this kind of generative art in a way that will enable methods to be shared and combined more fluidly than the 2021 paradigm of doing everything in colab permitted. \n",
        "\n",
        "If you're interested in contributing (even if you aren't a coder and just have an idea for something to add to the documentation), please visit our issue tracker: https://github.com/pytti-tools/pytti-core/issues\n",
        "\n",
        "Please help me untangle this thing before it swallows me whole. Thanks!\n",
        "\n",
        "`--The Management`\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8LjPC_0vN87t"
      },
      "source": [
        "# Instructions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AW4VrGKU0VUp"
      },
      "source": [
        "`scenes:` Descriptions of scenes you want generated, separated by `||`. Each scene can contain multiple prompts, separated by `|`.\n",
        "\n",
        "*Example:* `Winter sunrise | icy landscape || Winter day | snowy skyline || Winter sunset | chilly air || Winter night | clear sky` would go through several winter scenes.\n",
        "\n",
        "**Advanced:** weight prompts with `description:weight`. Higher `weight` values will be prioritized by the optimizer, and negative `weight` values will remove the description from the image. The default weight is $1$. Weights can also be functions of $t$ to change over the course of an animation.\n",
        "\n",
        "*Example scene:* `blue sky:10|martian landscape|red sky:-1` would try to turn the martian sky blue.\n",
        "\n",
        "**Advanced:** stop prompts once the image matches them sufficiently with `description:weight:stop`. `stop` should be between $0$ and $1$ for positive prompts, or between $-1$ and $0$ for negative prompts. Lower `stop` values will have more effect on the image (remember that $-1<-0.5<0$). A prompt with a negative `weight` will often go haywire without a stop. Stops can also be functions of $t$ to change over the course of an animation.\n",
        "\n",
        "*Example scene:* `Feathered dinosaurs|birds:1:0.87|scales:-1:-.9|text:-1:-.9` Would try to make feathered dinosaurs, lightly like birds, without scales or text, but without making 'anti-scales' or 'anti-text.'\n",
        "\n",
        "#**NEW:**\n",
        "\n",
        "**Advanced:** Use `description:weight_mask description` with a text prompt as `mask`. The prompt will only be applied to areas of the image that match `mask description` according to CLIP.\n",
        "\n",
        "*Example scene:* `Khaleesi Daenerys Targaryen | mother of dragons | dragon:3_baby` would only apply the weight `dragon` to parts of the image that match `baby`, thus turning the babies that `mother` tends to make into dragons (hopefully).\n",
        "\n",
        "**Advanced:** Use `description:weight_[mask]` with a URL or path to an image, or a path to a .mp4 video to use as a `mask`. The prompt will only be applied to the masked (white) areas of the mask image. Use `description:weight_[-mask]` to apply the prompt to the black areas instead.\n",
        "\n",
        "*Example scene:* `sunlight:3_[mask.mp4]|midnight:3_[-mask.mp4]` Would apply `sunlight` in the white areas of `mask.mp4`, and `midnight` in the black areas.\n",
        "\n",
        "**Legacy:** Directional weights will still work as before, but they aren't as good as masks.\n",
        "\n",
        "**Advanced:** Use `[path or url]` as a prompt to add a semantic image prompt. This will be read by CLIP and understood as a near perfect text description of the image.\n",
        "\n",
        "*Example scene:* `[artist signature.png]:-1:-.95|[https://i.redd.it/ewpeykozy7e71.png]:3|fractal clouds|hole in the sky`\n",
        "\n",
        "---\n",
        "\n",
        "`scene_prefix:` text prepended to the beginning of each scene.\n",
        "\n",
        "*Example:* `Trending on Arstation|`\n",
        "\n",
        "`scene_suffix:` text appended to the end of each scene.\n",
        "\n",
        "*Example:* ` by James Gurney`\n",
        "\n",
        "`interpolation_steps:` number of steps to spend smoothly transitioning from the last scene at the start of each scene. $200$ is a good default. Set to $0$ to disable.\n",
        "\n",
        "`steps_per_scene:` total number of steps to spend rendering each scene. Should be at least `interpolation_steps`. This will indirectly control the total length of an animation.\n",
        "\n",
        "---\n",
        "#**NEW**: \n",
        "`direct_image_prompts:` paths or urls of images that you want your image to look like in a literal sense, along with `weight_mask` and `stop` values, separated by `|`.\n",
        "\n",
        "Apply masks to direct image prompts with `path or url of image:weight_path or url of mask` For video masks it must be a path to an mp4 file.\n",
        "\n",
        "**Legacy** latent image prompts are no more. They are now rolled into direct image prompts.\n",
        "\n",
        "---\n",
        "\n",
        "`init_image:` path or url of start image. Works well for creating a central focus.\n",
        "\n",
        "\n",
        "`direct_init_weight:` Defaults to $0$. Use the initial image as a direct image prompt. Equivalent to adding `init_image:direct_init_weight` as a `direct_image_prompt`. Supports weights, masks, and stops.\n",
        "\n",
        "`semantic_init_weight:` Defaults to $0$. Defaults to $0$. Use the initial image as a semantic image prompt. Equivalent to adding `[init_image]:direct_init_weight` as a prompt to each scene in `scenes`. Supports weights, masks, and stops. **IMPORTANT** since this is a semantic prompt, you still need to put the mask in `[` `]` to denote it as a path or url, otherwise it will be read as text instead of a file.\n",
        "\n",
        "---\n",
        "\n",
        "`width`, `height:` image size. Set one of these $-1$ to derive it from the aspect ratio of the init image.\n",
        "\n",
        "`pixel_size:` integer image scale factor. Makes the image bigger. Set to $1$ for VQGAN or face VRAM issues.\n",
        "\n",
        "`smoothing_weight:` makes the image smoother. Defaults to $0$ (no smoothing). Can also be negative for that deep fried look.\n",
        "\n",
        "`image_model:` select how your image will be represented.\n",
        "\n",
        "`vqgan_model:` select your VQGAN version (only for `image_model: VQGAN`)\n",
        "\n",
        "`random_initial_palette:` if checked, palettes will start out with random colors. Otherwise they will start out as grayscale. (only for `image_model: Limited Palette`)\n",
        "\n",
        "`palette_size:` number of colors in each palette. (only for `image_model: Limited Palette`)\n",
        "\n",
        "`palettes:` total number of palettes. The image will have `palette_size*palettes` colors total. (only for `image_model: Limited Palette`)\n",
        "\n",
        "`gamma:` relative gamma value. Higher values make the image darker and higher contrast, lower values make the image lighter and lower contrast. (only for `image_model: Limited Palette`). $1$ is a good default.\n",
        "\n",
        "`hdr_weight:` how strongly the optimizer will maintain the `gamma`. Set to $0$ to disable. (only for `image_model: Limited Palette`)\n",
        "\n",
        "`palette_normalization_weight:` how strongly the optimizer will maintain the palettes' presence in the image. Prevents the image from losing palettes. (only for `image_model: Limited Palette`)\n",
        "\n",
        "`show_palette:` check this box to see the palette each time the image is displayed. (only for `image_model: Limited Palette`)\n",
        "\n",
        "`target_pallete:` path or url of an image which the model will use to make the palette it uses.\n",
        "\n",
        "`lock_pallete:` force the model to use the initial palette (most useful from restore, but will force a grayscale image or a wonky palette otherwise).\n",
        "\n",
        "---\n",
        "\n",
        "`animation_mode:` select animation mode or disable animation.\n",
        "\n",
        "`sampling_mode:` how pixels are sampled during animation. `nearest` will keep the image sharp, but may look bad. `bilinear` will smooth the image out, and `bicubic` is untested :)\n",
        "\n",
        "`infill_mode:` select how new pixels should be filled if they come in from the edge.\n",
        "* mirror: reflect image over boundary\n",
        "* wrap: pull pixels from opposite side\n",
        "* black: fill with black \n",
        "* smear: sample closest pixel in image\n",
        "\n",
        "`pre_animation_steps:` number of steps to run before animation starts, to begin with a stable image. $250$ is a good default.\n",
        "\n",
        "`steps_per_frame:` number of steps between each image move. $50$ is a good default.\n",
        "\n",
        "`frames_per_second:` number of frames to render each second. Controls how $t$ is scaled.\n",
        "\n",
        "`direct_stabilization_weight: ` keeps the current frame as a direct image prompt. For `Video Source` this will use the current frame of the video as a direct image prompt. For `2D` and `3D` this will use the shifted version of the previous frame. Also supports masks: `weight_mask.mp4`.\n",
        "\n",
        "`semantic_stabilization_weight: ` keeps the current frame as a semantic image prompt. For `Video Source` this will use the current frame of the video as a direct image prompt. For `2D` and `3D` this will use the shifted version of the previous frame. Also supports masks: `weight_[mask.mp4]` or `weight_mask phrase`.\n",
        "\n",
        "`depth_stabilization_weight: ` keeps the depth model output somewhat consistent at a *VERY* steep performance cost. For `Video Source` this will use the current frame of the video as a semantic image prompt. For `2D` and `3D` this will use the shifted version of the previous frame. Also supports masks: `weight_mask.mp4`.\n",
        "\n",
        "`edge_stabilization_weight: ` keeps the images contours somewhat consistent at very little performance cost. For `Video Source` this will use the current frame of the video as a direct image prompt with a sobel filter. For `2D` and `3D` this will use the shifted version of the previous frame. Also supports masks: `weight_mask.mp4`.\n",
        "\n",
        "`flow_stabilization_weight: ` used for `animation_mode: 3D` and `Video Source` to prevent flickering. Comes with a slight performance cost for `Video Source`, and a great one for `3D`, due to implementation differences. Also supports masks: `weight_mask.mp4`. For video source, the mask should select the part of the frame you want to move, and the rest will be treated as a still background.\n",
        "\n",
        "---\n",
        "`video_path: ` path to mp4 file for `Video Source`\n",
        "\n",
        "`frame_stride` advance this many frames in the video for each output frame. This is surprisingly useful. Set to $1$ to render each frame. Video masks will also step at this rate.\n",
        "\n",
        "`reencode_each_frame: ` check this box to use each video frame as an `init_image` instead of warping each output frame into the init for the next. Cuts will still be detected and trigger a reencode.\n",
        "\n",
        "\n",
        "`flow_long_term_samples: ` Sample multiple frames into the past for consistent interpolation even with disocclusion, as described by [Manuel Ruder, Alexey Dosovitskiy, and Thomas Brox (2016)](https://arxiv.org/abs/1604.08610). Each sample is twice as far back in the past as the last, so the earliest sampled frame is $2^{\\text{long_term_flow_samples}}$ frames in the past. Set to $0$ to disable.\n",
        "\n",
        "---\n",
        "\n",
        "`translate_x:` horizontal image motion as a function of time $t$ in seconds.\n",
        "\n",
        "`translate_y:` vertical image motion as a function of time $t$ in seconds.\n",
        "\n",
        "`translate_z_3d:` forward image motion as a function of time $t$ in seconds. (only for `animation_mode:3D`)\n",
        "\n",
        "`rotate_3d:` image rotation as a quaternion $\\left[r,x,y,z\\right]$ as a function of time $t$ in seconds. (only for `animation_mode:3D`)\n",
        "\n",
        "`rotate_2d:` image rotation in degrees as a function of time $t$ in seconds. (only for `animation_mode:2D`)\n",
        "\n",
        "`zoom_x_2d:` horizontal image zoom as a function of time $t$ in seconds. (only for `animation_mode:2D`)\n",
        "\n",
        "`zoom_y_2d:` vertical image zoom as a function of time $t$ in seconds. (only for `animation_mode:2D`)\n",
        "\n",
        "`lock_camera:` check this box to prevent all scrolling or drifting. Makes for more stable 3D rotations. (only for `animation_mode:3D`)\n",
        "\n",
        "`field_of_view:` vertical field of view in degrees. (only for `animation_mode:3D`)\n",
        "\n",
        "`near_plane:` closest depth distance in pixels. (only for `animation_mode:3D`)\n",
        "\n",
        "`far_plane:` farthest depth distance in pixels. (only for `animation_mode:3D`)\n",
        "\n",
        "---\n",
        "\n",
        "`file_namespace:` output directory name.\n",
        "\n",
        "`allow_overwrite:` check to overwrite existing files in `file_namespace`.\n",
        "\n",
        "`display_every:` how many steps between each time the image is displayed in the notebook.\n",
        "\n",
        "`clear_every:` how many steps between each time notebook console is cleared.\n",
        "\n",
        "`display_scale:` image display scale in notebook. $1$ will show the image at full size. Does not affect saved images.\n",
        "\n",
        "`save_every:` how many steps between each time the image is saved. Set to `steps_per_frame` for consistent animation.\n",
        "\n",
        "`backups:` number of backups to keep (only the oldest backups are deleted). Large images make very large backups, so be warned. Set to `all` to save all backups. These are used for the `flow_long_term_samples` so be sure that this is at least $2^{\\text{flow_long_term_samples}}+1$ for `Video Source` mode.\n",
        "\n",
        "`show_graphs:` check this to see graphs of the loss values each time the image is displayed. Disable this for local runtimes.\n",
        "\n",
        "`approximate_vram_usage:` currently broken. Don't believe its lies.\n",
        "\n",
        "---\n",
        "\n",
        "`ViTB32, ViTB16, RN50, RN50x4:` select your CLIP models. These take a lot of VRAM.\n",
        "\n",
        "`learning_rate:` how quickly the image changes.\n",
        "\n",
        "`reset_lr_each_frame:` the optimizer will adaptively change the learning rate, so this will thwart it.\n",
        "\n",
        "`seed:` pseudorandom seed.\n",
        "\n",
        "---\n",
        "\n",
        "`cutouts:` number of cutouts. Reduce this to use less VRAM at the cost of quality and speed.\n",
        "\n",
        "`cut_pow:` should be positive. Large values shrink cutouts, making the image more detailed, small values expand the cutouts, making it more coherent. $1$ is a good default. $3$ or higher can cause crashes.\n",
        "\n",
        "`cutout_border:` should be between $0$ and $1$. Allows cutouts to poke out over the edges of the image by this fraction of the image size, allowing better detail around the edges of the image. Set to $0$ to disable. $0.25$ is a good default.\n",
        "\n",
        "`border_mode:` how to fill cutouts that stick out over the edge of the image. Match with `infill_mode` for consistent infill.\n",
        "\n",
        "* clamp: move cutouts back onto image\n",
        "* mirror: reflect image over boundary\n",
        "* wrap: pull pixels from opposite side\n",
        "* black: fill with black \n",
        "* smear: sample closest pixel in image"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xZt160ePEc8f"
      },
      "source": [
        "#Step 1: Setup\n",
        "Run the cells in this section once for each runtime, or after a factory reset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "CAGyDOe2o9AE"
      },
      "outputs": [],
      "source": [
        "#@title 1.3 Install everything else\n",
        "#@markdown Run this cell on a fresh runtime to install the libraries and modules.\n",
        "from os.path import exists as path_exists\n",
        "if path_exists('/content/drive/MyDrive/pytti_test'):\n",
        "  %cd /content/drive/MyDrive/pytti_test\n",
        "\n",
        "def flush_reqs():\n",
        "    !rm -r pytti-core\n",
        "\n",
        "def install_everything():\n",
        "    if path_exists('./pytti-core'):\n",
        "        try:\n",
        "            flush_reqs()\n",
        "        except Exception as ex:\n",
        "            logger.warning(\n",
        "                str(ex)\n",
        "            )\n",
        "            logger.warning(\n",
        "                \"A `pytti` folder already exists and could not be deleted.\"\n",
        "                \"If you encounter problems, try deleting that folder and trying again.\"\n",
        "                \"Please report this and any other issues here: \"\n",
        "                \"https://github.com/pytti-tools/pytti-notebook/issues/new\",\n",
        "                exc_info=True)\n",
        "\n",
        "    !git clone --recurse-submodules -j8 https://github.com/pytti-tools/pytti-core\n",
        "\n",
        "    !pip install kornia pytorch-lightning\n",
        "    !pip install jupyter gdown loguru einops PyGLM ftfy regex tqdm hydra-core exrex\n",
        "    !pip install seaborn adjustText bunch matplotlib-label-lines\n",
        "    \n",
        "    !pip install ./pytti-core/vendor/AdaBins\n",
        "    !pip install ./pytti-core/vendor/CLIP\n",
        "    !pip install ./pytti-core/vendor/GMA\n",
        "    !pip install ./pytti-core/vendor/taming-transformers\n",
        "    !pip install ./pytti-core\n",
        "\n",
        "    !mkdir -p images_out\n",
        "    !mkdir -p videos\n",
        "\n",
        "    from pytti.Notebook import change_tqdm_color\n",
        "    change_tqdm_color()\n",
        "\n",
        "try:\n",
        "    from adjustText import adjust_text\n",
        "    import pytti, torch\n",
        "    everything_installed = True\n",
        "except ModuleNotFoundError:\n",
        "    everything_installed = False\n",
        "\n",
        "force_install = False #@param{type:\"boolean\"}\n",
        "if not everything_installed or force_install:\n",
        "    install_everything()\n",
        "elif everything_installed:\n",
        "    from pytti.Notebook import change_tqdm_color\n",
        "    change_tqdm_color()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a9C9tARLzyzq"
      },
      "outputs": [],
      "source": [
        "#@title 1.1 Mount google drive (optional)\n",
        "#@markdown Mounting your drive is optional but recommended. You can even restore from google randomly\n",
        "#@markdown kicking you out if you mount your drive.\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive', force_remount = True)\n",
        "!mkdir -p /content/drive/MyDrive/pytti_test\n",
        "%cd /content/drive/MyDrive/pytti_test"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "RosI5DYxtjh0"
      },
      "outputs": [],
      "source": [
        "#@title 1.2 NVIDIA-SMI (optional)\n",
        "#@markdown View information about your runtime GPU.\n",
        "#@markdown Google will connect you to an industrial strength GPU, which is needed to run\n",
        "#@markdown this notebook. You can also disable error checking on your GPU to get some\n",
        "#@markdown more VRAM, at a marginal cost to stability. You will have to restart the runtime after\n",
        "#@markdown disabling it.\n",
        "enable_error_checking = False#@param {type:\"boolean\"}\n",
        "if enable_error_checking:\n",
        "  !nvidia-smi\n",
        "else:\n",
        "  !nvidia-smi\n",
        "  !nvidia-smi -i 0 -e 0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PcayyBJjE-qy"
      },
      "source": [
        "# Step 2: Configure Experiment\n",
        "\n",
        "Edit the parameters, or load saved parameters, then run the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "TiKD7os1pyXW"
      },
      "outputs": [],
      "source": [
        "#@title #2.1 Parameters:\n",
        "#@markdown ---\n",
        "from os.path import exists as path_exists\n",
        "if path_exists('/content/drive/MyDrive/pytti_test'):\n",
        "  %cd /content/drive/MyDrive/pytti_test\n",
        "  drive_mounted = True\n",
        "else:\n",
        "  drive_mounted = False\n",
        "try:\n",
        "  from pytti.Notebook import change_tqdm_color, get_last_file\n",
        "except ModuleNotFoundError:\n",
        "  if drive_mounted:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError('ERROR: please run setup (step 1.3).')\n",
        "  else:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError('WARNING: drive is not mounted.\\nERROR: please run setup (step 1.3).')\n",
        "change_tqdm_color()\n",
        "\n",
        "import glob, json, random, re, math\n",
        "try:\n",
        "  from bunch import Bunch\n",
        "except ModuleNotFoundError:\n",
        "  if drive_mounted:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError('ERROR: please run setup (step 1.3).')\n",
        "  else:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError('WARNING: drive is not mounted.\\nERROR: please run setup (step 1.3).')\n",
        "\n",
        "#these are used to make the defaults look pretty\n",
        "model_default = None\n",
        "random_seed = None\n",
        "all  = math.inf\n",
        "derive_from_init_aspect_ratio = -1\n",
        "\n",
        "def define_parameters():\n",
        "  locals_before = locals().copy()\n",
        "  #@markdown ###Prompts:\n",
        "  \n",
        "  scenes = \"deep space habitation ring made of glass | galactic nebula | wow! space is full of fractal creatures darting around everywhere like fireflies\"#@param{type:\"string\"}\n",
        "  scene_prefix = \"astrophotography #pixelart | image credit nasa | space full of cybernetic neon:3_galactic nebula | isometric pixelart by Sachin Teng | \"#@param{type:\"string\"}\n",
        "  scene_suffix = \"| satellite image:-1:-.95 | text:-1:-.95 | anime:-1:-.95 | watermark:-1:-.95 | backyard telescope:-1:-.95 | map:-1:-.95\"#@param{type:\"string\"}\n",
        "  interpolation_steps = 0#@param{type:\"number\"}\n",
        "  steps_per_scene =  60100#@param{type:\"raw\"}\n",
        "  #@markdown ---\n",
        "  #@markdown ###Image Prompts:\n",
        "  direct_image_prompts   = \"\"#@param{type:\"string\"}\n",
        "  #@markdown ---\n",
        "  #@markdown ###Initial image:\n",
        "  init_image = \"\"#@param{type:\"string\"}\n",
        "  direct_init_weight =  \"\"#@param{type:\"string\"}\n",
        "  semantic_init_weight = \"\"#@param{type:\"string\"}\n",
        "  #@markdown ---\n",
        "  #@markdown ###Image:\n",
        "  #@markdown Use `image_model` to select how the model will encode the image\n",
        "  image_model = \"Limited Palette\" #@param [\"VQGAN\", \"Limited Palette\", \"Unlimited Palette\"]\n",
        "\n",
        "  #@markdown image_model | description | strengths | weaknesses\n",
        "  #@markdown --- | -- | -- | --\n",
        "  #@markdown  VQGAN | classic VQGAN image | smooth images | limited datasets, slow, VRAM intesnsive \n",
        "  #@markdown  Limited Palette | pytti differentiable palette | fast,  VRAM scales with `palettes` | pixel images\n",
        "  #@markdown  Unlimited Palette | simple RGB optimization | fast, VRAM efficient | pixel images\n",
        "  \n",
        "  #@markdown The output image resolution will be `width` $\\times$ `pixel_size` by height $\\times$ `pixel_size` pixels.\n",
        "  #@markdown The easiest way to run out of VRAM is to select `image_model` VQGAN without reducing\n",
        "  #@markdown `pixel_size` to $1$.\n",
        "\n",
        "  #@markdown For `animation_mode: 3D` the minimum resoultion is about 450 by 400 pixels.\n",
        "  width =  180#@param {type:\"raw\"}\n",
        "  height =  112#@param {type:\"raw\"}\n",
        "  pixel_size = 4#@param{type:\"number\"}\n",
        "  smoothing_weight =  0.02#@param{type:\"number\"}\n",
        "  #@markdown `VQGAN` specific settings:\n",
        "  vqgan_model = \"sflckr\" #@param [\"imagenet\", \"coco\", \"wikiart\", \"sflckr\", \"openimages\"]\n",
        "  #@markdown `Limited Palette` specific settings:\n",
        "  random_initial_palette = False#@param{type:\"boolean\"}\n",
        "  palette_size = 6#@param{type:\"number\"}\n",
        "  palettes   = 9#@param{type:\"number\"}\n",
        "  gamma = 1#@param{type:\"number\"}\n",
        "  hdr_weight = 0.01#@param{type:\"number\"}\n",
        "  palette_normalization_weight = 0.2#@param{type:\"number\"}\n",
        "  show_palette = False #@param{type:\"boolean\"}\n",
        "  target_palette = \"\"#@param{type:\"string\"}\n",
        "  lock_palette = False #@param{type:\"boolean\"}\n",
        "  #@markdown ---\n",
        "  #@markdown ###Animation:\n",
        "  animation_mode = \"3D\" #@param [\"off\",\"2D\", \"3D\", \"Video Source\"]\n",
        "  sampling_mode = \"bicubic\" #@param [\"bilinear\",\"nearest\",\"bicubic\"]\n",
        "  infill_mode = \"wrap\" #@param [\"mirror\",\"wrap\",\"black\",\"smear\"]\n",
        "  pre_animation_steps =  100#@param{type:\"number\"}\n",
        "  steps_per_frame =  50#@param{type:\"number\"}\n",
        "  frames_per_second =  12#@param{type:\"number\"}\n",
        "  #@markdown ---\n",
        "  #@markdown ###Stabilization Weights:\n",
        "  direct_stabilization_weight = \"\"#@param{type:\"string\"}\n",
        "  semantic_stabilization_weight = \"\"#@param{type:\"string\"}\n",
        "  depth_stabilization_weight = \"\"#@param{type:\"string\"}\n",
        "  edge_stabilization_weight = \"\"#@param{type:\"string\"}\n",
        "  #@markdown `flow_stabilization_weight` is used for `animation_mode: 3D` and `Video Source`\n",
        "  flow_stabilization_weight = \"\"#@param{type:\"string\"}\n",
        "  #@markdown ---\n",
        "  #@markdown ###Video Tracking:\n",
        "  #@markdown Only for `animation_mode: Video Source`.\n",
        "  video_path = \"\"#@param{type:\"string\"}\n",
        "  frame_stride = 1#@param{type:\"number\"}\n",
        "  reencode_each_frame = True #@param{type:\"boolean\"}\n",
        "  flow_long_term_samples = 1#@param{type:\"number\"}\n",
        "  #@markdown ---\n",
        "  #@markdown ###Image Motion:\n",
        "  translate_x    = \"-1700*sin(radians(1.5))\" #@param{type:\"string\"}\n",
        "  translate_y    = \"0\" #@param{type:\"string\"}\n",
        "  #@markdown `..._3d` is only used in 3D mode.\n",
        "  translate_z_3d = \"(50+10*t)*sin(t/10*pi)**2\" #@param{type:\"string\"}\n",
        "  #@markdown `rotate_3d` *must* be a `[w,x,y,z]` rotation (unit) quaternion. Use `rotate_3d: [1,0,0,0]` for no rotation.\n",
        "  #@markdown [Learn more about rotation quaternions here](https://eater.net/quaternions).\n",
        "  rotate_3d      = \"[cos(radians(1.5)), 0, -sin(radians(1.5))/sqrt(2), sin(radians(1.5))/sqrt(2)]\"#@param{type:\"string\"}\n",
        "  #@markdown `..._2d` is only used in 2D mode.\n",
        "  rotate_2d      = \"5\" #@param{type:\"string\"}\n",
        "  zoom_x_2d      = \"0\" #@param{type:\"string\"}\n",
        "  zoom_y_2d      = \"0\" #@param{type:\"string\"}\n",
        "  #@markdown  3D camera (only used in 3D mode):\n",
        "  lock_camera   = True#@param{type:\"boolean\"}\n",
        "  field_of_view = 60#@param{type:\"number\"}\n",
        "  near_plane    = 1#@param{type:\"number\"}\n",
        "  far_plane     = 10000#@param{type:\"number\"}\n",
        "\n",
        "  #@markdown ---\n",
        "  #@markdown ###Output:\n",
        "  file_namespace = \"default\"#@param{type:\"string\"}\n",
        "  if file_namespace == '':\n",
        "    file_namespace = 'out'\n",
        "  allow_overwrite = False#@param{type:\"boolean\"}\n",
        "  base_name = file_namespace\n",
        "  if not allow_overwrite and path_exists(f'images_out/{file_namespace}'):\n",
        "    _, i = get_last_file(f'images_out/{file_namespace}', \n",
        "                         f'^(?P<pre>{re.escape(file_namespace)}\\\\(?)(?P<index>\\\\d*)(?P<post>\\\\)?_1\\\\.png)$')\n",
        "    if i == 0:\n",
        "      print(f\"WARNING: file_namespace {file_namespace} already has images from run 0\")\n",
        "    elif i is not None:\n",
        "      print(f\"WARNING: file_namespace {file_namespace} already has images from runs 0 through {i}\")\n",
        "  elif glob.glob(f'images_out/{file_namespace}/{base_name}_*.png'):\n",
        "    print(f\"WARNING: file_namespace {file_namespace} has images which will be overwritten\")\n",
        "  try:\n",
        "    del i\n",
        "    del _\n",
        "  except NameError:\n",
        "    pass\n",
        "  del base_name\n",
        "  display_every = steps_per_frame #@param{type:\"raw\"}\n",
        "  clear_every = 0 #@param{type:\"raw\"}\n",
        "  display_scale = 1#@param{type:\"number\"}\n",
        "  save_every = steps_per_frame #@param{type:\"raw\"}\n",
        "  backups =  2**(flow_long_term_samples+1)+1#this is used for video transfer, so don't lower it if that's what you're doing#@param {type:\"raw\"}\n",
        "  show_graphs = False #@param{type:\"boolean\"}\n",
        "  approximate_vram_usage = False#@param{type:\"boolean\"}\n",
        "\n",
        "  #@markdown ---\n",
        "  #@markdown ###Model:\n",
        "  #@markdown Quality settings from Dribnet's CLIPIT (https://github.com/dribnet/clipit).\n",
        "  #@markdown Selecting too many will use up all your VRAM and slow down the model.\n",
        "  #@markdown I usually use ViTB32, ViTB16, and RN50 if I get a A100, otherwise I just use ViT32B.\n",
        "\n",
        "  #@markdown quality | CLIP models\n",
        "  #@markdown --- | --\n",
        "  #@markdown  draft | ViTB32 \n",
        "  #@markdown  normal | ViTB32, ViTB16 \n",
        "  #@markdown  high | ViTB32, ViTB16, RN50\n",
        "  #@markdown  best | ViTB32, ViTB16, RN50x4\n",
        "  ViTB32 = True #@param{type:\"boolean\"}\n",
        "  ViTB16 = False #@param{type:\"boolean\"}\n",
        "  RN50 = False #@param{type:\"boolean\"}\n",
        "  RN50x4 = False #@param{type:\"boolean\"}\n",
        "  #@markdown the default learning rate is `0.1` for all the VQGAN models\n",
        "  #@markdown except openimages, which is `0.15`. For the palette modes the\n",
        "  #@markdown default is `0.02`. \n",
        "  learning_rate =  model_default#@param{type:\"raw\"}\n",
        "  reset_lr_each_frame = True#@param{type:\"boolean\"}\n",
        "  seed = random_seed #@param{type:\"raw\"}\n",
        "  #@markdown **Cutouts**:\n",
        "\n",
        "  #@markdown [Cutouts are how CLIP sees the image.](https://twitter.com/remi_durant/status/1460607677801897990)\n",
        "  cutouts =  40#@param{type:\"number\"}\n",
        "  cut_pow =  2#@param {type:\"number\"}\n",
        "  cutout_border =  .25#@param {type:\"number\"}\n",
        "  #@markdown NOTE: prompt masks (`promt:weight_[mask.png]`) will not work right on '`wrap`' or '`mirror`' mode.\n",
        "  border_mode = \"clamp\" #@param [\"clamp\",\"mirror\",\"wrap\",\"black\",\"smear\"]\n",
        "  \n",
        "  if seed is None:\n",
        "    seed = random.randint(-0x8000_0000_0000_0000, 0xffff_ffff_ffff_ffff)\n",
        "  locals_after = locals().copy()\n",
        "  for k in locals_before.keys():\n",
        "    del locals_after[k]\n",
        "  del locals_after['locals_before']\n",
        "  return locals_after\n",
        "\n",
        "params = Bunch(define_parameters())\n",
        "print(\"SETTINGS:\")\n",
        "print(json.dumps(params))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "lWlZ2Gocb2fF"
      },
      "outputs": [],
      "source": [
        "#@title 2.2 Load settings (optional)\n",
        "#@markdown copy the `SETTINGS:` output from the **Parameters** cell (tripple click to select the whole\n",
        "#@markdown line from `{'scenes'...` to `}`) and paste them in a note to save them for later.\n",
        "\n",
        "#@markdown Paste them here in the future to load those settings again. Running this cell with blank settings won't do anything.\n",
        "from os.path import exists as path_exists\n",
        "if path_exists('/content/drive/MyDrive/pytti_test'):\n",
        "  %cd /content/drive/MyDrive/pytti_test\n",
        "  drive_mounted = True\n",
        "else:\n",
        "  drive_mounted = False\n",
        "try:\n",
        "  from pytti.Notebook import *\n",
        "except ModuleNotFoundError:\n",
        "  if drive_mounted:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError('ERROR: please run setup (step 1.3).')\n",
        "  else:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError('WARNING: drive is not mounted.\\nERROR: please run setup (step 1.3).')\n",
        "change_tqdm_color()\n",
        "  \n",
        "import json, random\n",
        "try:\n",
        "  from bunch import Bunch\n",
        "except ModuleNotFoundError:\n",
        "  if drive_mounted:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError('ERROR: please run setup (step 1.3).')\n",
        "  else:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError('WARNING: drive is not mounted.\\nERROR: please run setup (step 1.3).')\n",
        "\n",
        "settings = \"\"#@param{type:\"string\"}\n",
        "#@markdown Check `random_seed` to overwrite the seed from the settings with a random one for some variation.\n",
        "random_seed = False #@param{type:\"boolean\"}\n",
        "\n",
        "if settings != '':\n",
        "  params = load_settings(settings, random_seed)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AsW_gLPgUMZc"
      },
      "outputs": [],
      "source": [
        "from pytti.workhorse import _main, TB_LOGDIR \n",
        "\n",
        "%load_ext tensorboard"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9vHDNnvYUMZc"
      },
      "outputs": [],
      "source": [
        "%tensorboard --logdir $TB_LOGDIR "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h4LDZYtOUMZd"
      },
      "outputs": [],
      "source": [
        "#@title 2.3 Run it!\n",
        "from omegaconf import OmegaConf\n",
        "cfg = OmegaConf.create(dict(params))\n",
        "\n",
        "# function wraps step 2.3 of the original p5 notebook\n",
        "_main(cfg)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AtekvTZxFNpf"
      },
      "source": [
        "# Step 3: Render video\n",
        "You can dowload from the notebook, but it's faster to download from your drive."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "RZH-r4yyShnX"
      },
      "outputs": [],
      "source": [
        "#@title 3.1 Render video\n",
        "from os.path import exists as path_exists\n",
        "if path_exists('/content/drive/MyDrive/pytti_test'):\n",
        "  %cd /content/drive/MyDrive/pytti_test\n",
        "  drive_mounted = True\n",
        "else:\n",
        "  drive_mounted = False\n",
        "try:\n",
        "  from pytti.Notebook import change_tqdm_color\n",
        "except ModuleNotFoundError:\n",
        "  if drive_mounted:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError('ERROR: please run setup (step 1.3).')\n",
        "  else:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError('WARNING: drive is not mounted.\\nERROR: please run setup (step 1.3).')\n",
        "change_tqdm_color()\n",
        "  \n",
        "from tqdm.notebook import tqdm\n",
        "import numpy as np\n",
        "from os.path import exists as path_exists\n",
        "from subprocess import Popen, PIPE\n",
        "from PIL import Image, ImageFile\n",
        "from os.path import splitext as split_file\n",
        "import glob\n",
        "from pytti.Notebook import get_last_file\n",
        "\n",
        "ImageFile.LOAD_TRUNCATED_IMAGES = True\n",
        "\n",
        "try:\n",
        "  params\n",
        "except NameError:\n",
        "  raise RuntimeError(\"ERROR: no parameters. Please run parameters (step 2.1).\")\n",
        "\n",
        "if not path_exists(f\"images_out/{params.file_namespace}\"):\n",
        "  if path_exists(f\"/content/drive/MyDrive\"):\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError(f\"ERROR: file_namespace: {params.file_namespace} does not exist.\")\n",
        "  else:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError(f\"WARNING: Drive is not mounted.\\nERROR: file_namespace: {params.file_namespace} does not exist.\")\n",
        "\n",
        "#@markdown The first run executed in `file_namespace` is number $0$, the second is number $1$, etc.\n",
        "\n",
        "latest = -1\n",
        "run_number = latest#@param{type:\"raw\"}\n",
        "if run_number == -1:\n",
        "  _, i = get_last_file(f'images_out/{params.file_namespace}', \n",
        "                       f'^(?P<pre>{re.escape(params.file_namespace)}\\\\(?)(?P<index>\\\\d*)(?P<post>\\\\)?_1\\\\.png)$')\n",
        "  run_number = i\n",
        "base_name = params.file_namespace if run_number == 0 else (params.file_namespace+f\"({run_number})\")\n",
        "tqdm.write(f'Generating video from {params.file_namespace}/{base_name}_*.png')\n",
        "\n",
        "all_frames = glob.glob(f'images_out/{params.file_namespace}/{base_name}_*.png')\n",
        "all_frames.sort(key = lambda s: int(split_file(s)[0].split('_')[-1]))\n",
        "print(f'found {len(all_frames)} frames matching images_out/{params.file_namespace}/{base_name}_*.png')\n",
        "\n",
        "start_frame = 0#@param{type:\"number\"}\n",
        "all_frames = all_frames[start_frame:]\n",
        "\n",
        "fps =  params.frames_per_second#@param{type:\"raw\"}\n",
        "\n",
        "total_frames = len(all_frames)\n",
        "\n",
        "if total_frames == 0:\n",
        "  #THIS IS NOT AN ERROR. This is the code that would\n",
        "  #make an error if something were wrong.\n",
        "  raise RuntimeError(f\"ERROR: no frames to render in images_out/{params.file_namespace}\")\n",
        "\n",
        "frames = []\n",
        "\n",
        "for filename in tqdm(all_frames):\n",
        "  frames.append(Image.open(filename))\n",
        "\n",
        "p = Popen(['ffmpeg', '-y', '-f', 'image2pipe', '-vcodec', 'png', '-r', str(fps), '-i', '-', '-vcodec', 'libx264', '-r', str(fps), '-pix_fmt', 'yuv420p', '-crf', '1', '-preset', 'veryslow', f\"videos/{base_name}.mp4\"], stdin=PIPE)\n",
        "for im in tqdm(frames):\n",
        "  im.save(p.stdin, 'PNG')\n",
        "p.stdin.close()\n",
        "\n",
        "print(\"Encoding video...\")\n",
        "p.wait()\n",
        "print(\"Video complete.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "t3EgqHKrSjZx"
      },
      "outputs": [],
      "source": [
        "#@title 3.1 Render video (concatenate all runs)\n",
        "from os.path import exists as path_exists\n",
        "if path_exists('/content/drive/MyDrive/pytti_test'):\n",
        "  %cd /content/drive/MyDrive/pytti_test\n",
        "  drive_mounted = True\n",
        "else:\n",
        "  drive_mounted = False\n",
        "try:\n",
        "  from pytti.Notebook import change_tqdm_color\n",
        "except ModuleNotFoundError:\n",
        "  if drive_mounted:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError('ERROR: please run setup (step 1.3).')\n",
        "  else:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError('WARNING: drive is not mounted.\\nERROR: please run setup (step 1.3).')\n",
        "change_tqdm_color()\n",
        "  \n",
        "from tqdm.notebook import tqdm\n",
        "import numpy as np\n",
        "from os.path import exists as path_exists\n",
        "from subprocess import Popen, PIPE\n",
        "from PIL import Image, ImageFile\n",
        "from os.path import splitext as split_file\n",
        "import glob\n",
        "from pytti.Notebook import get_last_file\n",
        "\n",
        "ImageFile.LOAD_TRUNCATED_IMAGES = True\n",
        "\n",
        "try:\n",
        "  params\n",
        "except NameError:\n",
        "  raise RuntimeError(\"ERROR: no parameters. Please run parameters (step 2.1).\")\n",
        "\n",
        "if not path_exists(f\"images_out/{params.file_namespace}\"):\n",
        "  if path_exists(f\"/content/drive/MyDrive\"):\n",
        "    raise RuntimeError(f\"ERROR: file_namespace: {params.file_namespace} does not exist.\")\n",
        "  else:\n",
        "    raise RuntimeError(f\"WARNING: Drive is not mounted.\\nERROR: file_namespace: {params.file_namespace} does not exist.\")\n",
        "\n",
        "#@markdown The first run executed in `file_namespace` is number $0$, the second is number $1$, etc.\n",
        "\n",
        "latest = -1\n",
        "run_number = latest\n",
        "if run_number == -1:\n",
        "  _, i = get_last_file(f'images_out/{params.file_namespace}', \n",
        "                       f'^(?P<pre>{re.escape(params.file_namespace)}\\\\(?)(?P<index>\\\\d*)(?P<post>\\\\)?_1\\\\.png)$')\n",
        "  run_number = i\n",
        "\n",
        "all_frames = []\n",
        "for i in range(run_number+1):\n",
        "  base_name = params.file_namespace if i == 0 else (params.file_namespace+f\"({i})\")\n",
        "  frames = glob.glob(f'images_out/{params.file_namespace}/{base_name}_*.png')\n",
        "  frames.sort(key = lambda s: int(split_file(s)[0].split('_')[-1]))\n",
        "  all_frames.extend(frames)\n",
        "\n",
        "start_frame = 0#@param{type:\"number\"}\n",
        "all_frames = all_frames[start_frame:]\n",
        "\n",
        "fps =  params.frames_per_second#@param{type:\"raw\"}\n",
        "\n",
        "total_frames = len(all_frames)\n",
        "\n",
        "if total_frames == 0:\n",
        "  #THIS IS NOT AN ERROR. This is the code that would\n",
        "  #make an error if something were wrong.\n",
        "  raise RuntimeError(f\"ERROR: no frames to render in images_out/{params.file_namespace}\")\n",
        "\n",
        "frames = []\n",
        "\n",
        "for filename in tqdm(all_frames):\n",
        "  frames.append(Image.open(filename))\n",
        "\n",
        "p = Popen(['ffmpeg', '-y', '-f', 'image2pipe', '-vcodec', 'png', '-r', str(fps), '-i', '-', '-vcodec', 'libx264', '-r', str(fps), '-pix_fmt', 'yuv420p', '-crf', '1', '-preset', 'veryslow', f\"videos/{base_name}.mp4\"], stdin=PIPE)\n",
        "for im in tqdm(frames):\n",
        "  im.save(p.stdin, 'PNG')\n",
        "p.stdin.close()\n",
        "\n",
        "print(\"Encoding video...\")\n",
        "p.wait()\n",
        "print(\"Video complete.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-qZ8c_-iZ0QM"
      },
      "outputs": [],
      "source": [
        "#@title 3.2 Download the last exported video\n",
        "from os.path import exists as path_exists\n",
        "if path_exists('/content/drive/MyDrive/pytti_test'):\n",
        "  %cd /content/drive/MyDrive/pytti_test\n",
        "\n",
        "try:\n",
        "  from pytti.Notebook import get_last_file\n",
        "except ModuleNotFoundError:\n",
        "  if drive_mounted:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError('ERROR: please run setup (step 1.3).')\n",
        "  else:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError('WARNING: drive is not mounted.\\nERROR: please run setup (step 1.3).')\n",
        "\n",
        "try:\n",
        "  params\n",
        "except NameError:\n",
        "  #THIS IS NOT AN ERROR. This is the code that would\n",
        "  #make an error if something were wrong.\n",
        "  raise RuntimeError(\"ERROR: please run parameters (step 2.1).\")\n",
        "\n",
        "from google.colab import files\n",
        "try:\n",
        "  base_name = params.file_namespace if run_number == 0 else (params.file_namespace+f\"({run_number})\")\n",
        "  filename = f'{base_name}.mp4'\n",
        "except NameError:\n",
        "  filename, i = get_last_file(f'videos', \n",
        "                       f'^(?P<pre>{re.escape(params.file_namespace)}\\\\(?)(?P<index>\\\\d*)(?P<post>\\\\)?\\\\.mp4)$')\n",
        "\n",
        "if path_exists(f'videos/{filename}'):\n",
        "  files.download(f\"videos/{filename}\")\n",
        "else:\n",
        "  if path_exists(f\"/content/drive/MyDrive\"):\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError(f\"ERROR: video videos/{filename} does not exist.\")\n",
        "  else:\n",
        "    #THIS IS NOT AN ERROR. This is the code that would\n",
        "    #make an error if something were wrong.\n",
        "    raise RuntimeError(f\"WARNING: Drive is not mounted.\\nERROR: video videos/{filename} does not exist.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AwwtSgibUMZe"
      },
      "source": [
        "# Batch Settings\n",
        "\n",
        "Be Advised: google may penalize you for sustained colab GPU utilization, even if you are a PRO+ subscriber. Tread lightly with batch runs, you don't wanna end up in GPU jail."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ut4HfS3qUMZe"
      },
      "source": [
        "FYI: the batch setting feature below may not work at present. We recommend using the CLI for batch jobs, see usage instructions at https://github.com/pytti-tools/pytti-core . The code below will probably be removed in the near future."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "if_Fdy_OFkjZ"
      },
      "source": [
        "# Batch Setings\n",
        "WARNING: If you use google colab (even with pro and pro+) GPUs for long enought google will throttle your account. Be careful with batch runs if you don't want to get kicked."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "vHHfWWqoSz35"
      },
      "outputs": [],
      "source": [
        "#@title batch settings\n",
        "\n",
        "# ngl... this probably doesn't work right now.\n",
        "\n",
        "from os.path import exists as path_exists\n",
        "if path_exists('/content/drive/MyDrive/pytti_test'):\n",
        "  %cd /content/drive/MyDrive/pytti_test\n",
        "  drive_mounted = True\n",
        "else:\n",
        "  drive_mounted = False\n",
        "try:\n",
        "  from pytti.Notebook import change_tqdm_color, save_batch\n",
        "except ModuleNotFoundError:\n",
        "  if drive_mounted:\n",
        "    raise RuntimeError('ERROR: please run setup (step 1).')\n",
        "  else:\n",
        "    raise RuntimeError('WARNING: drive is not mounted.\\nERROR: please run setup (step 1).')\n",
        "change_tqdm_color()\n",
        "\n",
        "try:\n",
        "  import exrex, random, glob\n",
        "except ModuleNotFoundError:\n",
        "  if drive_mounted:\n",
        "    raise RuntimeError('ERROR: please run setup (step 1).')\n",
        "  else:\n",
        "    raise RuntimeError('WARNING: drive is not mounted.\\nERROR: please run setup (step 1).')\n",
        "from numpy import arange\n",
        "import itertools\n",
        "\n",
        "def all_matches(s):\n",
        "  return list(exrex.generate(s))\n",
        "\n",
        "def dict_product(dictionary):\n",
        "  return [dict(zip(dictionary, x)) for x in itertools.product(*dictionary.values())]\n",
        "\n",
        "#these are used to make the defaults look pretty\n",
        "model_default = None\n",
        "random_seed = None\n",
        "\n",
        "def define_parameters():\n",
        "  locals_before = locals().copy()\n",
        "  scenes = [\"list\",\"your\",\"runs\"] #@param{type:\"raw\"}\n",
        "  scene_prefix = [\"all \",\" permutations \",\" are run \"] #@param{type:\"raw\"}\n",
        "  scene_suffix = [\" that\", \" makes\", \" 27\" ] #@param{type:\"raw\"}\n",
        "  interpolation_steps = [0] #@param{type:\"raw\"}\n",
        "  steps_per_scene = [300] #@param{type:\"raw\"}\n",
        "  direct_image_prompts = [\"\"] #@param{type:\"raw\"}\n",
        "  init_image = [\"\"] #@param{type:\"raw\"}\n",
        "  direct_init_weight = [\"\"] #@param{type:\"raw\"}\n",
        "  semantic_init_weight = [\"\"] #@param{type:\"raw\"}\n",
        "  image_model = [\"Limited Palette\"] #@param{type:\"raw\"}\n",
        "  width = [180] #@param{type:\"raw\"}\n",
        "  height = [112] #@param{type:\"raw\"}\n",
        "  pixel_size = [4] #@param{type:\"raw\"}\n",
        "  smoothing_weight = [0.05] #@param{type:\"raw\"}\n",
        "  vqgan_model = [\"sflckr\"] #@param{type:\"raw\"}\n",
        "  random_initial_palette = [False] #@param{type:\"raw\"}\n",
        "  palette_size = [9] #@param{type:\"raw\"}\n",
        "  palettes = [8] #@param{type:\"raw\"}\n",
        "  gamma = [1] #@param{type:\"raw\"}\n",
        "  hdr_weight = [1.0] #@param{type:\"raw\"}\n",
        "  palette_normalization_weight = [1.0] #@param{type:\"raw\"}\n",
        "  show_palette = [False] #@param{type:\"raw\"}\n",
        "  target_palette = [\"\"] #@param{type:\"raw\"}\n",
        "  lock_palette = [False] #@param{type:\"raw\"}\n",
        "  animation_mode = [\"off\"] #@param{type:\"raw\"}\n",
        "  sampling_mode = [\"bicubic\"] #@param{type:\"raw\"}\n",
        "  infill_mode = [\"wrap\"] #@param{type:\"raw\"}\n",
        "  pre_animation_steps = [100] #@param{type:\"raw\"}\n",
        "  steps_per_frame = [50] #@param{type:\"raw\"}\n",
        "  frames_per_second = [12] #@param{type:\"raw\"}\n",
        "  direct_stabilization_weight = [\"\"] #@param{type:\"raw\"}\n",
        "  semantic_stabilization_weight = [\"\"] #@param{type:\"raw\"}\n",
        "  depth_stabilization_weight = [\"\"] #@param{type:\"raw\"}\n",
        "  edge_stabilization_weight = [\"\"] #@param{type:\"raw\"}\n",
        "  flow_stabilization_weight = [\"\"] #@param{type:\"raw\"}\n",
        "  video_path = [\"\"] #@param{type:\"raw\"}\n",
        "  frame_stride = [1] #@param{type:\"raw\"}\n",
        "  reencode_each_frame = [True] #@param{type:\"raw\"}\n",
        "  flow_long_term_samples = [0] #@param{type:\"raw\"}\n",
        "  translate_x = [\"0\"] #@param{type:\"raw\"}\n",
        "  translate_y = [\"0\"] #@param{type:\"raw\"}\n",
        "  translate_z_3d = [\"0\"] #@param{type:\"raw\"}\n",
        "  rotate_3d = [\"[1,0,0,0]\"] #@param{type:\"raw\"}\n",
        "  rotate_2d = [\"0\"] #@param{type:\"raw\"}\n",
        "  zoom_x_2d = [\"0\"] #@param{type:\"raw\"}\n",
        "  zoom_y_2d = [\"0\"] #@param{type:\"raw\"}\n",
        "  lock_camera = [True] #@param{type:\"raw\"}\n",
        "  field_of_view = [60] #@param{type:\"raw\"}\n",
        "  near_plane = [1] #@param{type:\"raw\"}\n",
        "  far_plane = [10000] #@param{type:\"raw\"}\n",
        "  file_namespace = [\"Basic Batch\"] #@param{type:\"raw\"}\n",
        "  allow_overwrite = [False]\n",
        "  display_every = [50] #@param{type:\"raw\"}\n",
        "  clear_every = [0] #@param{type:\"raw\"}\n",
        "  display_scale = [1] #@param{type:\"raw\"}\n",
        "  save_every = [50] #@param{type:\"raw\"}\n",
        "  backups = [2] #@param{type:\"raw\"}\n",
        "  show_graphs = [False] #@param{type:\"raw\"}\n",
        "  approximate_vram_usage = [False] #@param{type:\"raw\"}\n",
        "  ViTB32 = [True] #@param{type:\"raw\"}\n",
        "  ViTB16 = [False] #@param{type:\"raw\"}\n",
        "  RN50 = [False] #@param{type:\"raw\"}\n",
        "  RN50x4 = [False] #@param{type:\"raw\"}\n",
        "  learning_rate = [None] #@param{type:\"raw\"}\n",
        "  reset_lr_each_frame = [True] #@param{type:\"raw\"}\n",
        "  seed = [None] #@param{type:\"raw\"}\n",
        "  cutouts = [40] #@param{type:\"raw\"}\n",
        "  cut_pow = [2] #@param{type:\"raw\"}\n",
        "  cutout_border = [0.25] #@param{type:\"raw\"}\n",
        "  border_mode = [\"clamp\"] #@param{type:\"raw\"}\n",
        "  locals_after = locals().copy()\n",
        "  for k in locals_before.keys():\n",
        "    del locals_after[k]\n",
        "  del locals_after['locals_before']\n",
        "  return locals_after\n",
        "\n",
        "param_dict = define_parameters()\n",
        "batch_list = dict_product(param_dict)\n",
        "namespace = batch_list[0]['file_namespace']\n",
        "if glob.glob(f'images_out/{namespace}/*.png'):\n",
        "  print(f\"WARNING: images_out/{namespace} contains images. Batch indicies may not match filenames unless restoring.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Z0hRb20yxxsc"
      },
      "outputs": [],
      "source": [
        "# @title Licensed under the MIT License\n",
        "# Copyleft (c) 2021 Henry Rachootin\n",
        "# Copyright (c) 2022 David Marx\n",
        "\n",
        "# Permission is hereby granted, free of charge, to any person obtaining a copy\n",
        "# of this software and associated documentation files (the \"Software\"), to deal\n",
        "# in the Software without restriction, including without limitation the rights\n",
        "# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
        "# copies of the Software, and to permit persons to whom the Software is\n",
        "# furnished to do so, subject to the following conditions:\n",
        "\n",
        "# The above copyright notice and this permission notice shall be included in\n",
        "# all copies or substantial portions of the Software.\n",
        "\n",
        "# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
        "# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
        "# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
        "# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
        "# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
        "# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\n",
        "# THE SOFTWARE."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "pytti 5 beta.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.9.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}