Video_Killed_The_Radio_Star_Defusion.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mgXxoDhMAiti"
      },
      "source": [
        "# $ \\text{Video Killed The Radio Star}$ $\\color{red}{...Diffusion}$\n",
        "\n",
        "Notebook by David Marx ([@DigThatData](https://twitter.com/digthatdata))\n",
        "\n",
        "Shared under MIT license\n",
        "\n",
        "\n",
        "# $\\text{FAQ}$\n",
        "\n",
        "**What is this?**\n",
        "\n",
        "Point this notebook at a youtube url and it'll make a music video for you.\n",
        "\n",
        "**How does this animation technique work?**\n",
        "\n",
        "For each text prompt you provide, the notebook will...\n",
        "\n",
        "1. Generate an image based on that text prompt (using stable diffusion)\n",
        "2. Use the generated image as the `init_image` to recombine with the text prompt to generate variations similar to the first image. This produces a sequence of extremely similar images based on the original text prompt\n",
        "3. Images are then intelligently reordered to find the smoothest animation sequence of those frames\n",
        "3. This image sequence is then repeated to pad out the animation duration as needed\n",
        "\n",
        "The technique demonstrated in this notebook was inspired by a [video](https://www.youtube.com/watch?v=WJaxFbdjm8c) created by Ben Gillin.\n",
        "\n",
        "**How are lyrics transcribed?**\n",
        "\n",
        "This notebook uses openai's recently released 'whisper' model for performing automatic speech recognition. \n",
        "OpenAI was kind enough to offer several different sizes of this model which each have their own pros and cons. \n",
        "This notebook uses the largest whisper model for transcribing the actual lyrics. Additionally, we use the \n",
        "smallest model for performing the lyric segmentation. Neither of these models is perfect, but the results \n",
        "so far seem pretty decent.\n",
        "\n",
        "The first draft of this notebook relied on subtitles from youtube videos to determine timing, which was\n",
        "then aligned with user-provided lyrics. Youtube's automated captions are powerful and I'll update the\n",
        "notebook shortly to leverage those again, but for the time being we're just using whisper for everything\n",
        "and not referencing user-provided captions at all.\n",
        "\n",
        "**Something didn't work quite right in the transcription process. How do fix the timing or the actual lyrics?**\n",
        "\n",
        "The notebook is divided into several steps. Between each step, a \"storyboard\" file is updated. If you want to\n",
        "make modifications, you can edit this file directly and those edits should be reflected when you next load the\n",
        "file. Depending on what you changed and what step you run next, your changes may be ignored or even overwritten.\n",
        "Still playing with different solutions here.\n",
        "\n",
        "**Can I provide my own images to 'bring to life' and associate with certain lyrics/sequences?**\n",
        "\n",
        "Yes, you can! As described above: you just need to modify the storyboard. Will describe this functionality in\n",
        "greater detail after the implementation stabilizes a bit more.\n",
        "\n",
        "**This gave me an idea and I'd like to use just a part of your process here. What's the best way to reuse just some of the machinery you've developed here?**\n",
        "\n",
        "Most of the functionality in this notebook has been offloaded to library I published to pypi called `vktrs`. I strongly encourage you to import anything you need \n",
        "from there rather than cutting and pasting function into a notebook. Similarly, if you have ideas for improvements, please don't hesitate to submit a PR!\n",
        "\n",
        "**How can I support your work or work like it?**\n",
        "\n",
        "This notebook was made possible thanks to ongoing support from [stability.ai](https://stability.ai/). The best way to support my work is to share it with your friends, [report bugs](https://github.com/dmarx/video-killed-the-radio-star/issues/new), [suggest features](https://github.com/dmarx/video-killed-the-radio-star/discussions) or to donate to open source non-profits :) "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sM147HP4kAdY"
      },
      "source": [
        "## $0.$ Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ZnTe8clZuZuj"
      },
      "outputs": [],
      "source": [
        "# @title # 📊 Check GPU Status\n",
        "\n",
        "try:\n",
        "    from vktrs.utils import gpu_info\n",
        "except:\n",
        "    import pandas as pd\n",
        "    import subprocess\n",
        "    \n",
        "    def gpu_info():\n",
        "        outv = subprocess.run([\n",
        "            'nvidia-smi',\n",
        "                # these lines concatenate into a single query string\n",
        "                '--query-gpu='\n",
        "                'timestamp,'\n",
        "                'name,'\n",
        "                'utilization.gpu,'\n",
        "                'utilization.memory,'\n",
        "                'memory.used,'\n",
        "                'memory.free,'\n",
        "                ,\n",
        "            '--format=csv'\n",
        "            ],\n",
        "            stdout=subprocess.PIPE).stdout.decode('utf-8')\n",
        "\n",
        "        header, rec = outv.split('\\n')[:-1]\n",
        "        return pd.DataFrame({' '.join(k.strip().split('.')).capitalize():v for k,v in zip(header.split(','), rec.split(','))}, index=[0]).T\n",
        "\n",
        "gpu_info()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "oPbeyWtesAoh"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "\n",
        "# @title # 🛠️ Installations\n",
        "!apt-get install cargo\n",
        "!pip install vktrs[api,hf]\n",
        "\n",
        "!pip install git+https://github.com/openai/whisper@v20230314\n",
        "\n",
        "# these are only needed for hf\n",
        "!pip install \"ipywidgets>=7,<8\"\n",
        "!sudo apt -qq install git-lfs\n",
        "!git config --global credential.helper store\n",
        "\n",
        "!pip install panel prefetch_generator\n",
        "\n",
        "!pip uninstall -y protobuf\n",
        "!pip install protobuf==4.22.1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "cM8cux9b7F4v"
      },
      "outputs": [],
      "source": [
        "# @title # 🔑 Provide your API Key\n",
        "# @markdown Running this cell will prompt you to enter your API Key below. \n",
        "\n",
        "# @markdown To get your API key, visit https://beta.dreamstudio.ai/membership\n",
        "\n",
        "# @markdown ---\n",
        "\n",
        "# @markdown A note on security best practices: **don't publish your API key.**\n",
        "\n",
        "# @markdown We're using a form field designed for sensitive data like passwords.\n",
        "# @markdown This notebook does not save your API key in the notebook itself,\n",
        "# @markdown but instead loads your API Key into the colab environment. This way,\n",
        "# @markdown you can make changes to this notebook and share it without concern\n",
        "# @markdown that you might accidentally share your API Key. \n",
        "# @markdown \n",
        "\n",
        "use_stability_api = False # @param {type:'boolean'}\n",
        "mount_gdrive = True # @param {type:'boolean'}\n",
        "\n",
        "import os\n",
        "from pathlib import Path\n",
        "import time\n",
        "\n",
        "from omegaconf import OmegaConf\n",
        "\n",
        "\n",
        "os.environ['XDG_CACHE_HOME'] = os.environ.get(\n",
        "    'XDG_CACHE_HOME',\n",
        "    str(Path('~/.cache').expanduser())\n",
        ")\n",
        "if mount_gdrive:\n",
        "    from google.colab import drive\n",
        "    drive.mount('/content/drive')\n",
        "    Path('/content/drive/MyDrive/AI/models/.cache/').mkdir(parents=True, exist_ok=True) \n",
        "    # This rm+ln solution is not great. Be careful not to run this locally. \n",
        "    # Low risk, but could be annoying    \n",
        "    !rm -rf /root/.cache\n",
        "    !ln -sf /content/drive/MyDrive/AI/models/.cache/ /root/\n",
        "    # Following line will be sufficient pending merge of https://github.com/openai/whisper/pull/257\n",
        "    os.environ['XDG_CACHE_HOME']='/content/drive/MyDrive/AI/models/.cache'\n",
        "\n",
        "model_dir_str=str(Path(os.environ['XDG_CACHE_HOME']))\n",
        "proj_root_str = '${active_project}'\n",
        "if mount_gdrive:\n",
        "    proj_root_str = '/content/drive/MyDrive/AI/VideoKilledTheRadioStar/${active_project}'\n",
        "\n",
        "\n",
        "# notebook config\n",
        "cfg = OmegaConf.create({\n",
        "    'active_project':str(time.time()),\n",
        "    'project_root':proj_root_str,\n",
        "    'gdrive_mounted':mount_gdrive,\n",
        "    'use_stability_api':use_stability_api,\n",
        "    'model_dir':model_dir_str,\n",
        "    'output_dir':'${active_project}/frames'\n",
        "})\n",
        "\n",
        "with open('config.yaml','w') as fp:\n",
        "    OmegaConf.save(config=cfg, f=fp.name)\n",
        "\n",
        "###################\n",
        "\n",
        "if use_stability_api:\n",
        "    import os, getpass\n",
        "    os.environ['STABILITY_KEY'] = getpass.getpass('Enter your API Key')\n",
        "else:\n",
        "    try:\n",
        "        from google.colab import output\n",
        "        output.enable_custom_widget_manager()\n",
        "    except ImportError:\n",
        "        # assume local use\n",
        "        pass\n",
        "    \n",
        "    from huggingface_hub import notebook_login\n",
        "\n",
        "    # to do: if gdrive mounted, check for API token... somewhere on drive?\n",
        "    # looks like we should be able to find the token through an environment variable\n",
        "    notebook_login()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rt9Mu97fk_bp"
      },
      "source": [
        "## $1.$ 📋 Set Project Name (create/resume)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "s-9xjgy0iHhS"
      },
      "outputs": [],
      "source": [
        "\n",
        "import time\n",
        "from vktrs.utils import sanitize_folder_name\n",
        "from omegaconf import OmegaConf\n",
        "\n",
        "project_name = '' # @param {type:'string'}\n",
        "if not project_name:\n",
        "    project_name = str(time.time())\n",
        "\n",
        "project_name = sanitize_folder_name(project_name)\n",
        "\n",
        "workspace = OmegaConf.load('config.yaml')\n",
        "workspace.active_project = project_name\n",
        "\n",
        "with open('config.yaml','w') as fp:\n",
        "    OmegaConf.save(config=workspace, f=fp.name)\n",
        "\n",
        "# @markdown To create a new project, enter a unique project name.\n",
        "# @markdown If you leave `project_name` blank, the current unix timestamp will be used\n",
        "# @markdown  (seconds since 1970-01-01 00:00).\n",
        "\n",
        "# @markdown If you use the name of an existing project, the workspace will switch to that project.\n",
        "\n",
        "# @markdown Non-alphanumeric characters (excluding '-' and '_') will be replaced with hyphens.\n",
        "\n",
        "\n",
        "# reset workspace\n",
        "if 'df' in locals():\n",
        "    del df\n",
        "if 'df_regen' in locals():\n",
        "    del df_regen"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7eTPNhcBomtL"
      },
      "source": [
        "## $2.$ 🔊 Infer speech from audio"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "9zT0u4-q_fMF"
      },
      "outputs": [],
      "source": [
        "from omegaconf import OmegaConf\n",
        "from pathlib import Path\n",
        "\n",
        "workspace = OmegaConf.load('config.yaml')\n",
        "use_stability_api = workspace.use_stability_api\n",
        "model_dir = workspace.model_dir\n",
        "\n",
        "root = workspace.project_root\n",
        "root = Path(root)\n",
        "root.mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "\n",
        "import copy\n",
        "import datetime as dt\n",
        "import gc\n",
        "from itertools import chain, cycle\n",
        "import json\n",
        "import os\n",
        "import re\n",
        "import string\n",
        "from subprocess import Popen, PIPE\n",
        "import textwrap\n",
        "import time\n",
        "import warnings\n",
        "\n",
        "from IPython.display import display\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import panel as pn\n",
        "from tqdm.autonotebook import tqdm\n",
        "\n",
        "import tokenizations\n",
        "import webvtt\n",
        "import whisper\n",
        "\n",
        "from vktrs.utils import remove_punctuation\n",
        "from vktrs.utils import get_audio_duration_seconds\n",
        "from vktrs.youtube import (\n",
        "    YoutubeHelper,\n",
        "    parse_timestamp,\n",
        "    vtt_to_token_timestamps,\n",
        "    srv2_to_token_timestamps,\n",
        ")\n",
        "\n",
        "storyboard = OmegaConf.create()\n",
        "\n",
        "d_ = dict(\n",
        "    # all this does is make it so each of the following lines can be preceded with a comma\n",
        "    # otw the first parameter would be offset from the other in the colab form\n",
        "    _=\"\"\n",
        "\n",
        "    , video_url = 'https://www.youtube.com/watch?v=REojIUxX4rw' # @param {type:'string'}\n",
        "    , audio_fpath = '' # @param {type:'string'}\n",
        "    , whisper_seg = True # @param {type:'boolean'}\n",
        ")\n",
        "d_.pop('_')\n",
        "storyboard.params = d_\n",
        "\n",
        "if not storyboard.params.audio_fpath:\n",
        "    storyboard.params.audio_fpath = None\n",
        "\n",
        "\n",
        "# @markdown `video_url` - URL of a youtube video to download as a source for audio and potentially for text transcription as well.\n",
        "\n",
        "# @markdown `audio_fpath` - Optionally provide an audio file instead of relying on a youtube download. Name it something other than 'audio.mp3', \n",
        "# @markdown                 otherwise it might get overwritten accidentally.\n",
        "\n",
        "# @markdown `whisper_seg` - Whether or not to use openai's whisper model for lyric segmentation. This is currently the only option, but that will change in a few days.\n",
        "\n",
        "\n",
        "storyboard_fname = root / 'storyboard.yaml'\n",
        "with open(storyboard_fname,'wb') as fp:\n",
        "    OmegaConf.save(config=storyboard, f=fp.name)\n",
        "\n",
        "\n",
        "###############################\n",
        "# Download audio from youtube #\n",
        "###############################\n",
        "\n",
        "video_url = storyboard.params.video_url\n",
        "\n",
        "if video_url:\n",
        "    # check if user provided an audio filepath (or we already have one from youtube) before attempting to download\n",
        "    if storyboard.params.get('audio_fpath') is None:\n",
        "        helper = YoutubeHelper(\n",
        "            video_url,\n",
        "            ydl_opts = {\n",
        "                'outtmpl':{'default':str( root / f\"ytdlp_content.%(ext)s\" )},\n",
        "                'writeautomaticsub':True,\n",
        "                'subtitlesformat':'srv2/vtt'\n",
        "                },\n",
        "        )\n",
        "\n",
        "        # estimate video end\n",
        "        video_duration = dt.timedelta(seconds=helper.info['duration'])\n",
        "        storyboard.params['video_duration'] = video_duration.total_seconds()\n",
        "\n",
        "        audio_fpath = str( root / 'audio.mp3' )\n",
        "        input_audio = helper.info['requested_downloads'][-1]['filepath']\n",
        "        !ffmpeg -y -i \"{input_audio}\" -acodec libmp3lame {audio_fpath}\n",
        "\n",
        "        # to do: write audio and subtitle paths/meta to storyboard\n",
        "        storyboard.params.audio_fpath = audio_fpath\n",
        "\n",
        "        if False:\n",
        "            subtitle_format = helper.info['requested_subtitles']['en']['ext']\n",
        "            subtitle_fpath = helper.info['requested_subtitles']['en']['filepath']\n",
        "\n",
        "            if subtitle_format == 'srv2':\n",
        "                with open(subtitle_fpath, 'r') as f:\n",
        "                    srv2_xml = f.read() \n",
        "                token_start_times = srv2_to_token_timestamps(srv2_xml)\n",
        "                # to do: handle timedeltas...\n",
        "                #storyboard.params.token_start_times = token_start_times\n",
        "\n",
        "            elif subtitle_format == 'vtt':\n",
        "                captions = webvtt.read(subtitle_fpath)\n",
        "                token_start_times = vtt_to_token_timestamps(captions)\n",
        "                # to do: handle timedeltas...\n",
        "                #storyboard.params.token_start_times = token_start_times\n",
        "\n",
        "            # If unable to download supported subtitles, force use whisper\n",
        "            else:\n",
        "                storyboard.params.whisper_seg = True\n",
        "\n",
        "\n",
        "# estimate video end\n",
        "if storyboard.params.get('video_duration') is None:\n",
        "    # estimate duration from audio file\n",
        "    audio_fpath = storyboard.params['audio_fpath']\n",
        "    storyboard.params['video_duration'] = get_audio_duration_seconds(audio_fpath)\n",
        "\n",
        "if storyboard.params.get('video_duration') is None:\n",
        "    raise RuntimeError('unable to determine audio duration. was a video url or path to a file supplied?')\n",
        "\n",
        "# force use\n",
        "storyboard.params.whisper_seg = True\n",
        "\n",
        "with open(storyboard_fname,'wb') as fp:\n",
        "    OmegaConf.save(config=storyboard, f=fp.name)\n",
        "\n",
        "whisper_seg = storyboard.params.whisper_seg\n",
        "\n",
        "###################################################\n",
        "# 💬 Transcribe and segment speech using whisper #\n",
        "###################################################\n",
        "\n",
        "# handle OOM... or try to, anyway\n",
        "if 'hf_helper' in locals():\n",
        "    del hf_helper.img2img\n",
        "    del hf_helper.text2img\n",
        "    del hf_helper\n",
        "\n",
        "\n",
        "if whisper_seg:\n",
        "    from vktrs.asr import (\n",
        "        #whisper_lyrics,\n",
        "        #whisper_transcribe,\n",
        "        #whisper_align,\n",
        "        whisper_transmit_meta_across_alignment,\n",
        "        whisper_segment_transcription,\n",
        "    )\n",
        "\n",
        "    #prompt_starts = whisper_lyrics(audio_fpath=storyboard.params.audio_fpath)\n",
        "\n",
        "    audio_fpath = storyboard.params.audio_fpath\n",
        "    #whispers = whisper_transcribe(audio_fpath)\n",
        "\n",
        "    # to do: dropdown selectors\n",
        "    segmentation_model = 'tiny'\n",
        "    transcription_model = 'large'\n",
        "\n",
        "    storyboard.params.whisper = dict(\n",
        "        segmentation_model = segmentation_model\n",
        "        ,transcription_model = transcription_model\n",
        "    )\n",
        "\n",
        "    whispers = {\n",
        "        #'tiny':None, # 5.83 s\n",
        "        #'large':None # 3.73 s\n",
        "    }\n",
        "    # accelerated runtime required for whisper\n",
        "    # to do: pypi package for whisper\n",
        "\n",
        "    # to do: use transcripts we've already built if we have them\n",
        "    #scripts = storyboard.params.whisper.get('transcriptions')\n",
        "    \n",
        "    for k in set([segmentation_model, transcription_model]):\n",
        "        #if k in scripts:\n",
        "\n",
        "        options = whisper.DecodingOptions(\n",
        "            language='en',\n",
        "        )\n",
        "        # to do: be more proactive about cleaning up these models when we're done with them\n",
        "        model = whisper.load_model(k).to('cuda')\n",
        "        start = time.time()\n",
        "        print(f\"Transcribing audio with whisper-{k}\")\n",
        "        \n",
        "        # to do: calling transcribe like this unnecessarily re-processes audio each time.\n",
        "        whispers[k] = model.transcribe(audio_fpath) # re-processes audio each time, ~10s overhead?\n",
        "        print(f\"elapsed: {time.time()-start}\")\n",
        "        del model\n",
        "        gc.collect()\n",
        "    \n",
        "    #######################\n",
        "    # save transcriptions #\n",
        "    #######################\n",
        "\n",
        "    transcriptions = {}\n",
        "    transcription_root = root / 'whispers'\n",
        "    transcription_root.mkdir(parents=True, exist_ok=True)\n",
        "    writer = whisper.utils.get_writer(output_format='vtt', output_dir=transcription_root) # output dir doesn't do anything...?\n",
        "    for k in whispers:\n",
        "        outpath = str( transcription_root / f\"{k}.vtt\" )\n",
        "        transcriptions[k] = outpath\n",
        "        with open(outpath,'w') as f:\n",
        "            # to do: upstream PR to control verbosity\n",
        "            writer.write_result(\n",
        "                whispers[k],\n",
        "                file=f,\n",
        "            )\n",
        "    storyboard.params.whisper.transcriptions = transcriptions\n",
        "\n",
        "    #tiny2large, large2tiny, whispers_tokens = whisper_align(whispers)\n",
        "    # sanitize and tokenize\n",
        "    whispers_tokens = {}\n",
        "    for k in whispers:\n",
        "        whispers_tokens[k] = [\n",
        "        remove_punctuation(tok) for tok in whispers[k]['text'].split()\n",
        "        ]\n",
        "\n",
        "    # align sequences\n",
        "    tiny2large, large2tiny = tokenizations.get_alignments(\n",
        "        whispers_tokens[segmentation_model], #whispers_tokens['tiny'],\n",
        "        whispers_tokens[transcription_model] #whispers_tokens['large']\n",
        "    )\n",
        "    #return tiny2large, large2tiny, whispers_tokens\n",
        "\n",
        "    token_large_index_segmentations = whisper_transmit_meta_across_alignment(\n",
        "        whispers,\n",
        "        large2tiny,\n",
        "        whispers_tokens,\n",
        "    )\n",
        "    prompt_starts = whisper_segment_transcription(\n",
        "        token_large_index_segmentations,\n",
        "    )\n",
        "\n",
        "\n",
        "    ### checkpoint the processing work we've done to this point\n",
        "\n",
        "    prompt_starts_copy = copy.deepcopy(prompt_starts)\n",
        "    \n",
        "    # to do: deal with timedeltas in asr.py and yt.py\n",
        "    for rec in prompt_starts_copy:\n",
        "        for k,v in list(rec.items()):\n",
        "            if isinstance(v, dt.timedelta):\n",
        "                rec[k] = v.total_seconds()\n",
        "    \n",
        "    storyboard.prompt_starts = prompt_starts_copy\n",
        "\n",
        "    with open(storyboard_fname) as fp:\n",
        "        OmegaConf.save(config=storyboard, f=fp.name)\n",
        "\n",
        "###############################\n",
        "# Review/Modify transcription #\n",
        "###############################\n",
        "\n",
        "# @markdown ---\n",
        "# @markdown NB: When this cell finishes running, a table will appear\n",
        "# @markdown at the bottom of the output window. This table is editable\n",
        "# @markdown and can be used to correct errors in the transcription.\n",
        "# @markdown\n",
        "# @markdown additionally, the `override_prompt` field can be used to provide an \n",
        "# @markdown alternative text prompt for image generation. If this feature is\n",
        "# @markdown used, both the lyric and the theme prompt (which you will specify \n",
        "# @markdown in the cell that follows this) will be ignored. If you want to use\n",
        "# @markdown an `override_prompt` and also want to stay on theme, you will have \n",
        "# @markdown to append the desired `theme_prompt` to the end of the \n",
        "# @markdown `override_prompt` manually.\n",
        "\n",
        "\n",
        "# https://panel.holoviz.org/reference/widgets/Tabulator.html\n",
        "pn.extension('tabulator') # I don't know that specifying 'tabulator' here is even necessary...\n",
        "\n",
        "tabulator_formatters = {\n",
        "    'bool': {'type': 'tickCross'}\n",
        "}\n",
        "\n",
        "# reset workspace\n",
        "if 'df_regen' in locals():\n",
        "    del df_regen\n",
        "\n",
        "df = pd.DataFrame(prompt_starts).rename(\n",
        "    columns={\n",
        "        'ts':'Timestamp (sec)',\n",
        "        'prompt':'Lyric',\n",
        "    }\n",
        ")\n",
        "\n",
        "if 'td' in df:\n",
        "    del df['td']\n",
        "\n",
        "df['override_prompt'] = ''\n",
        "\n",
        "df_pre = copy.deepcopy(df)\n",
        "pn.widgets.Tabulator(df, formatters=tabulator_formatters)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_RTUFeyQqCfd"
      },
      "source": [
        "## $3.$ 🎬 Animate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Sh514DGj_sua"
      },
      "outputs": [],
      "source": [
        "import copy\n",
        "import datetime as dt\n",
        "from pathlib import Path\n",
        "import random\n",
        "import string\n",
        "\n",
        "from bokeh.models.widgets.tables import (\n",
        "    NumberFormatter, \n",
        "    BooleanFormatter,\n",
        "    CheckboxEditor,\n",
        ")\n",
        "import numpy as np\n",
        "from omegaconf import OmegaConf, DictConfig\n",
        "import pandas as pd\n",
        "import panel as pn\n",
        "import PIL\n",
        "from tqdm.autonotebook import tqdm\n",
        "\n",
        "from vktrs.tsp import (\n",
        "    tsp_permute_frames,\n",
        "    batched_tsp_permute_frames,\n",
        ")\n",
        "\n",
        "from vktrs.utils import (\n",
        "    add_caption2image,\n",
        "    save_frame,\n",
        "    remove_punctuation,\n",
        "    get_image_sequence,\n",
        "    archive_images,\n",
        ")\n",
        "\n",
        "# to do: is there a way to check if this is in the env already?\n",
        "pn.extension('tabulator')\n",
        "\n",
        "# this processes optional edits to the transcription (above) \n",
        "if ('prompt_starts' in locals()) \\\n",
        "and ('df_pre' in locals()):\n",
        "    if isinstance(prompt_starts, DictConfig):\n",
        "        prompt_starts = OmegaConf.to_container(prompt_starts)\n",
        "    # update prompt_starts if any changes were made above\n",
        "    if not np.all(df_pre.values == df.values):\n",
        "        df_pre = copy.deepcopy(df)\n",
        "        for i, rec in enumerate(prompt_starts):\n",
        "            rec['ts'] = float(df.loc[i,'Timestamp (sec)'])\n",
        "            rec['prompt'] = df.loc[i,'Lyric']\n",
        "            rec['override_prompt'] = df.loc[i,'override_prompt']\n",
        "        \n",
        "        # ...actually, I think the above code might not do anything\n",
        "        # probably need to checkpoint prompt_starts into the storyboard on disk.\n",
        "        # let's do that here just to be safe.    \n",
        "        workspace = OmegaConf.load('config.yaml')\n",
        "        root = Path(workspace.project_root)\n",
        "\n",
        "        storyboard_fname = root / 'storyboard.yaml'\n",
        "        storyboard = OmegaConf.load(storyboard_fname)\n",
        "\n",
        "        storyboard.prompt_starts = prompt_starts\n",
        "        with open(storyboard_fname) as fp:\n",
        "            OmegaConf.save(config=storyboard, f=fp.name)\n",
        "\n",
        "\n",
        "#####################################\n",
        "# @title ## 🎨 Generate init images\n",
        "#####################################\n",
        "\n",
        "workspace = OmegaConf.load('config.yaml')\n",
        "root = Path(workspace.project_root)\n",
        "\n",
        "storyboard_fname = root / 'storyboard.yaml'\n",
        "storyboard = OmegaConf.load(storyboard_fname)\n",
        "\n",
        "prompt_starts = storyboard.prompt_starts\n",
        "use_stability_api = workspace.use_stability_api\n",
        "model_dir = workspace.model_dir\n",
        "\n",
        "if use_stability_api:\n",
        "    from vktrs.api import get_image_for_prompt\n",
        "elif 'hf_helper' not in locals():\n",
        "    from vktrs.hf import HfHelper\n",
        "    # this needs to not be in the same cell as the login.\n",
        "    # some sort of stupid race condition.\n",
        "    try:\n",
        "        hf_helper = HfHelper(\n",
        "            download=False,\n",
        "            model_path=str(Path(model_dir) / 'huggingface' / 'diffusers')\n",
        "        )\n",
        "    except:\n",
        "        hf_helper = HfHelper(\n",
        "            download=True,\n",
        "            model_path=str(Path(model_dir) / 'huggingface' / 'diffusers')\n",
        "        )\n",
        "\n",
        "    # I give up.\n",
        "    def get_image_for_prompt(*args, **kargs):\n",
        "        # ugly hotfix. todo: unify function signatures...\n",
        "        if 'init_image' in kargs:\n",
        "            img = kargs.pop('init_image')\n",
        "            if img:\n",
        "                kargs['image'] = img\n",
        "        result = hf_helper.get_image_for_prompt(*args, **kargs)\n",
        "        return result.images\n",
        "\n",
        "\n",
        "def get_variations_w_init(prompt, init_image, **kargs):\n",
        "    return list(get_image_for_prompt(prompt=prompt, init_image=init_image, **kargs))\n",
        "\n",
        "def get_close_variations_from_prompt(prompt, n_variations=2, image_consistency=.7):\n",
        "    \"\"\"\n",
        "    prompt: a text prompt\n",
        "    n_variations: total number of images to return\n",
        "    image_consistency: float in [0,1], controls similarity between images generated by the prompt.\n",
        "                        you can think of this as controlling how much \"visual vibration\" there will be.\n",
        "                        - 0=regenerate each iandely identical\n",
        "    \"\"\"\n",
        "    images = list(get_image_for_prompt(prompt))\n",
        "    for _ in range(n_variations - 1):\n",
        "        img = get_variations_w_init(prompt, images[0], start_schedule=(1-image_consistency))[0]\n",
        "        images.append(img)\n",
        "    return images\n",
        "\n",
        "\n",
        "d_ = dict(\n",
        "    _=''\n",
        "    , theme_prompt = \"by ralph steadman and radiohead,  so detailed, beautiful, amazing, provocative\" # @param {type:'string'}\n",
        "    , height = 512 # @param {type:'integer'}\n",
        "    , width = 512 # @param {type:'integer'}\n",
        "    , display_frames_as_we_get_them = True # @param {type:'boolean'}\n",
        ")\n",
        "d_.pop('_')\n",
        "\n",
        "regenerate_all_init_images = False # @param {type:'boolean'}\n",
        "\n",
        "prompt_lag = True # @param {type:'boolean'}\n",
        "\n",
        "# @markdown `theme_prompt` - Text that will be appended to the end of each lyric, useful for e.g. applying a consistent aesthetic style\n",
        "\n",
        "# @markdown `display_frames_as_we_get_them` - Displaying frames will make the notebook slightly slower\n",
        "\n",
        "# regenerate all images if the theme prompt has changed or user specifies\n",
        "\n",
        "# @markdown `prompt_lag` - Extend prompt with lyrics from previous frame. Can improve temporal consistency of narrative. \n",
        "# @markdown  Especially useful for lyrics segmented into short prompts.\n",
        "\n",
        "if d_['theme_prompt'] != storyboard.params.get('theme_prompt'):\n",
        "    regenerate_all_init_images = True\n",
        "\n",
        "storyboard.params.update(d_)\n",
        "\n",
        "if regenerate_all_init_images:\n",
        "    for i, rec in enumerate(prompt_starts):\n",
        "        rec['frame0_fpath'] = None\n",
        "        archive_images(i, root=root)\n",
        "    print(\"archival process complete\")\n",
        "\n",
        "# anchor images will be regenerated if there's no associated frame0_fpath\n",
        "# regenerate specific images if\n",
        "# * manually tagged by user in df_regen\n",
        "# * associated fpath doesn't exist (i.e. deleted)\n",
        "if 'df_regen' in locals():\n",
        "    for i, _ in df_regen.iterrows():\n",
        "        rec = prompt_starts[i]\n",
        "        regen = not _['keep']\n",
        "        if rec.get('frame0_fpath') is None:\n",
        "            regen = True\n",
        "        elif not Path(rec['frame0_fpath']).exists():\n",
        "            regen=True\n",
        "        if regen:\n",
        "            rec['frame0_fpath'] = None\n",
        "            rec['prompt'] = df_regen.loc[i, 'Lyric']\n",
        "            rec['override_prompt'] = df_regen.loc[i, 'override_prompt']\n",
        "            print(rec)\n",
        "            archive_images(i, root=root)\n",
        "    print(\"archival process complete\")\n",
        "\n",
        "\n",
        "theme_prompt = storyboard.params.theme_prompt\n",
        "display_frames_as_we_get_them = storyboard.params.display_frames_as_we_get_them\n",
        "height = storyboard.params.height\n",
        "width = storyboard.params.width\n",
        "\n",
        "proj_name = workspace.active_project\n",
        "\n",
        "print(\"Ensuring each prompt has an associated image\")\n",
        "for idx, rec in enumerate(prompt_starts):\n",
        "    lyric = rec['prompt']\n",
        "    prompt = f\"{lyric}, {theme_prompt}\"\n",
        "    override = rec.get('override_prompt','').strip()\n",
        "    if override:\n",
        "        print('override prompt detected')\n",
        "        prompt = override\n",
        "    print(\n",
        "        f\"\\n[{idx} | {rec['ts']}] - {lyric} - {prompt}\"\n",
        "    )\n",
        "    \n",
        "    if prompt_lag and (idx > 0):\n",
        "        rec_prev = prompt_starts[idx -1]\n",
        "        prev_prompt = rec_prev.get('override_prompt','').strip()\n",
        "        if not prev_prompt:\n",
        "            prev_prompt = rec_prev['prompt']\n",
        "        prompt = f\"{prev_prompt}, {prompt}\"\n",
        "    if rec.get('frame0_fpath') is None:\n",
        "        init_image = list(get_image_for_prompt(\n",
        "              prompt,\n",
        "              height=height,\n",
        "              width=width,\n",
        "              )\n",
        "          )[0]\n",
        "    # this shouldn't be necessary, but is a consequence of\n",
        "    # the globbing thing we're doing atm\n",
        "    if 'anchor' not in str(rec.get('frame0_fpath')):\n",
        "        rec['frame0_fpath'] = save_frame(\n",
        "            init_image,\n",
        "            idx,\n",
        "            root_path = root / 'frames',\n",
        "            name='anchor',\n",
        "            )\n",
        "\n",
        "        if display_frames_as_we_get_them:\n",
        "            print(lyric)\n",
        "            display(init_image)\n",
        "\n",
        "\n",
        "##############\n",
        "# checkpoint #\n",
        "##############\n",
        "\n",
        "prompt_starts_copy = copy.deepcopy(prompt_starts)\n",
        "storyboard.prompt_starts = prompt_starts_copy\n",
        "\n",
        "with open(storyboard_fname) as fp:\n",
        "    OmegaConf.save(config=storyboard, f=fp.name)\n",
        "\n",
        "\n",
        "###############\n",
        "# flag regens #\n",
        "###############\n",
        "\n",
        "# @markdown ---\n",
        "# @markdown NB: When this cell finishes running, a table will appear at the bottom of the output window. This table is editable and can be used to correct errors in the transcription (see above).\n",
        "\n",
        "# @markdown Additionally, this table can be used to trigger regeneration of\n",
        "# @markdown images you don't want to keep. On the far left of the table, you\n",
        "# @markdown you should see a `keep` column that defaults to \"true\". Double \n",
        "# @markdown clicking this value should flip it to \"false\". Rerunning this cell\n",
        "# @markdown will regenerate the `init_image` for all scenes where `keep=false`.\n",
        "# @markdown Images that are flagged for regeneration will be moved to the\n",
        "# @markdown project's `archive` folder.\n",
        "\n",
        "# @markdown Image regeneration can also be triggered by deleting the image from \n",
        "# @markdown the `frames` folder.\n",
        "\n",
        "\n",
        "df_regen = pd.DataFrame(prompt_starts)\n",
        "if 'override_prompt' not in df_regen:\n",
        "    df_regen['override_prompt'] = ''\n",
        "\n",
        "df_regen = df_regen[['ts','prompt','override_prompt']].rename(\n",
        "    columns={\n",
        "        'ts':'Timestamp (sec)',\n",
        "        'prompt':'Lyric',\n",
        "    }\n",
        ")\n",
        "\n",
        "\n",
        "\n",
        "df_regen['keep'] = True\n",
        "\n",
        "# move the \"keep\" column to the front\n",
        "df_regen= df_regen[['keep', 'Timestamp (sec)', 'Lyric', 'override_prompt']]\n",
        "\n",
        "pn.widgets.Tabulator(\n",
        "    df_regen,\n",
        "    formatters={'bool': BooleanFormatter()},\n",
        "    editors={'bool':CheckboxEditor()}\n",
        "    )\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "4AgMgvYusAo4"
      },
      "outputs": [],
      "source": [
        "\n",
        "# @title ## 🚀 Generate animation frames\n",
        "\n",
        "###################\n",
        "# improved resume #\n",
        "###################\n",
        "\n",
        "import copy\n",
        "import datetime as dt\n",
        "from itertools import cycle\n",
        "from pathlib import Path\n",
        "\n",
        "from omegaconf import OmegaConf\n",
        "from PIL import Image\n",
        "from vktrs.utils import (\n",
        "    add_caption2image,\n",
        "    get_image_sequence,\n",
        ")\n",
        "\n",
        "\n",
        "workspace = OmegaConf.load('config.yaml')\n",
        "root = Path(workspace.project_root)\n",
        "\n",
        "storyboard_fname = root / 'storyboard.yaml'\n",
        "storyboard = OmegaConf.load(storyboard_fname)\n",
        "\n",
        "if not 'prompt_starts' in locals():\n",
        "    prompt_starts = OmegaConf.to_container(storyboard.prompt_starts)\n",
        "else:\n",
        "    ##########################\n",
        "    # checkpoint any changes #\n",
        "    ##########################\n",
        "    prompt_starts_copy = copy.deepcopy(prompt_starts)\n",
        "\n",
        "    storyboard.prompt_starts = prompt_starts_copy\n",
        "\n",
        "    with open(storyboard_fname) as fp:\n",
        "        OmegaConf.save(config=storyboard, f=fp.name)\n",
        "\n",
        "\n",
        "#################################################\n",
        "# Math                                          #\n",
        "#                                               #\n",
        "#    This block computes how many frames are    #\n",
        "#    needed for each segment based on the start #\n",
        "#    times for each prompt                      #\n",
        "#################################################\n",
        "\n",
        "# to do: \n",
        "# * make this more portable and add to vktrs lib\n",
        "\n",
        "fps = 12 # @param {type:'integer'}\n",
        "storyboard.params.fps = fps\n",
        "\n",
        "ifps = 1/fps\n",
        "\n",
        "# estimate video end\n",
        "video_duration = storyboard.params['video_duration']\n",
        "\n",
        "# dummy prompt for last scene duration\n",
        "prompt_starts = OmegaConf.to_container(storyboard.prompt_starts)\n",
        "prompt_starts.append({'ts':video_duration})\n",
        "\n",
        "# make sure we respect the duration of the previous phrase\n",
        "frame_start=0\n",
        "prompt_starts[0]['anim_start']=frame_start\n",
        "for i, rec in enumerate(prompt_starts[1:], start=1):\n",
        "    rec_prev = prompt_starts[i-1]\n",
        "    k=0\n",
        "    while (rec_prev['anim_start'] + k*ifps) < rec['ts']:\n",
        "        k+=1\n",
        "    k-=1\n",
        "    rec_prev['frames'] = k\n",
        "    rec_prev['anim_duration'] = k*ifps\n",
        "    frame_start+=k*ifps\n",
        "    rec['anim_start']=frame_start\n",
        "\n",
        "# drop the dummy frame\n",
        "prompt_starts = prompt_starts[:-1]\n",
        "\n",
        "# to do: given a 0 duration prompt, assume its duration is captured in the next prompt \n",
        "#        and guesstimate a corrected prompt start time and duration \n",
        "\n",
        "\n",
        "##############\n",
        "# checkpoint #\n",
        "##############\n",
        "\n",
        "prompt_starts_copy = copy.deepcopy(prompt_starts)\n",
        "\n",
        "storyboard.prompt_starts = prompt_starts_copy\n",
        "\n",
        "with open(storyboard_fname) as fp:\n",
        "    OmegaConf.save(config=storyboard, f=fp.name)\n",
        "\n",
        "\n",
        "##################################\n",
        "# Generate animation frames #\n",
        "##################################\n",
        "\n",
        "d_ = dict(\n",
        "    _=''\n",
        "    , n_variations=5 # @param {type:'integer'}\n",
        "    , image_consistency=0.8 # @param {type:\"slider\", min:0, max:1, step:0.01}  \n",
        "    , max_video_duration_in_seconds = 300 # @param {type:'integer'}\n",
        ")\n",
        "d_.pop('_')\n",
        "\n",
        "\n",
        "# @markdown `fps` - Frames-per-second of generated animations\n",
        "\n",
        "# @markdown `n_variations` - How many unique variations to generate for a given text prompt. This determines the frequency of the visual \"pulsing\" effect\n",
        "\n",
        "# @markdown `image_consistency` - controls similarity between images generated by the prompt.\n",
        "# @markdown - 0: ignore the init image\n",
        "# @markdown - 1: true as possible to the init image\n",
        "\n",
        "# @markdown `max_video_duration_in_seconds` - Early stopping if you don't want to generate a video the full duration of the provided audio. Default = 5min.\n",
        "\n",
        "\n",
        "storyboard.params.update(d_)\n",
        "storyboard.params.max_frames = storyboard.params.fps * storyboard.params.max_video_duration_in_seconds\n",
        "\n",
        "# to do: compute and report unique of image generations\n",
        "\n",
        "display_frames_as_we_get_them = storyboard.params.display_frames_as_we_get_them\n",
        "image_consistency = storyboard.params.image_consistency\n",
        "max_frames = storyboard.params.max_frames\n",
        "\n",
        "n_variations = storyboard.params.n_variations\n",
        "theme_prompt = storyboard.params.get('theme_prompt')\n",
        "\n",
        "\n",
        "# load init_images and generate variations as needed\n",
        "# to do: request multiple images in single request\n",
        "print(\"Fetching variations\")\n",
        "for idx, rec in enumerate(prompt_starts):\n",
        "    new_images = []\n",
        "    images_fpaths = get_image_sequence(idx, root=root)\n",
        "    curr_variation_count = len(images_fpaths)\n",
        "    print(curr_variation_count)\n",
        "    if curr_variation_count < n_variations:\n",
        "        # to do: \n",
        "        # * prompt lag\n",
        "        lyric = rec['prompt']\n",
        "        prompt = f\"{lyric}, {theme_prompt}\"\n",
        "        if rec.get('override_prompt'):\n",
        "            prompt = rec['override_prompt']\n",
        "\n",
        "        init_image = Image.open(rec['frame0_fpath'])\n",
        "        # next line is here to permit user to specify more variations for a specific entry\n",
        "        tot_variations = rec.get('n_variations', n_variations)\n",
        "        tot_variations = min(tot_variations, rec['frames']) # don't generate variations we won't use\n",
        "        tot_variations -= curr_variation_count  # only generate variations we still need\n",
        "        for _ in range(tot_variations):\n",
        "            img = get_variations_w_init(prompt, init_image, start_schedule=(1-image_consistency))[0]\n",
        "            save_frame(\n",
        "                img,\n",
        "                idx,\n",
        "                root_path= root / 'frames',\n",
        "            )\n",
        "            if display_frames_as_we_get_them:\n",
        "                display(img)\n",
        "\n",
        "\n",
        "##############\n",
        "# checkpoint #\n",
        "##############\n",
        "\n",
        "prompt_starts_copy = copy.deepcopy(prompt_starts)\n",
        "\n",
        "storyboard.prompt_starts = prompt_starts_copy\n",
        "\n",
        "# to do: deal with these td objects\n",
        "with open(storyboard_fname) as fp:\n",
        "    OmegaConf.save(config=storyboard, f=fp.name)\n",
        "\n",
        "# @markdown ---\n",
        "# @markdown Running this cell will generate as many variation frames as required \n",
        "# @markdown per `n_variations`. To trigger regeneration of images that didn't\n",
        "# @markdown generate correctly (e.g. because a nsfw classifier was triggered),\n",
        "# @markdown just delete those images."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "cEwFI6kA_2SH"
      },
      "outputs": [],
      "source": [
        "# @title ## 🎞️ Compile your video!\n",
        "\n",
        "import shutil\n",
        "from subprocess import Popen, PIPE\n",
        "\n",
        "from pathlib import Path\n",
        "from PIL import Image\n",
        "from itertools import cycle\n",
        "\n",
        "from omegaconf import OmegaConf\n",
        "from tqdm.autonotebook import tqdm\n",
        "\n",
        "try:\n",
        "    from prefetch_generator import BackgroundGenerator\n",
        "except:\n",
        "    !pip install prefetch_generator\n",
        "    from prefetch_generator import BackgroundGenerator\n",
        "\n",
        "from vktrs.tsp import (\n",
        "    tsp_permute_frames,\n",
        "    batched_tsp_permute_frames,\n",
        ")\n",
        "\n",
        "from vktrs.utils import (\n",
        "    add_caption2image,\n",
        "    get_image_sequence,\n",
        "    save_frame,\n",
        "    remove_punctuation,\n",
        ")\n",
        "\n",
        "\n",
        "# reload config\n",
        "workspace = OmegaConf.load('config.yaml')\n",
        "root = Path(workspace.project_root)\n",
        "\n",
        "storyboard_fname = root / 'storyboard.yaml'\n",
        "storyboard = OmegaConf.load(storyboard_fname)\n",
        "\n",
        "########################\n",
        "# rendering parameters #\n",
        "########################\n",
        "\n",
        "output_filename = 'output.mp4' # @param {type:'string'}\n",
        "add_caption = True # @param {type:'boolean'}\n",
        "optimal_ordering = True # @param {type:'boolean'}\n",
        "upscale = False # @param {type:'boolean'}\n",
        "\n",
        "# @markdown `add_caption` - Whether or not to overlay the prompt text on the image\n",
        "\n",
        "# @markdown `optimal_ordering` - Intelligently permutes animation frames to provide a smoother animation.\n",
        "\n",
        "# @markdown  `upscale`: Naively (lanczos interpolation) upscale video 2x. This can be a way to force\n",
        "# @markdown  services like youtube to deliver your video without mangling it with compression\n",
        "# @markdown  artifacts. Thanks [@gandamu_ml](https://twitter.com/gandamu_ml) for this trick!\n",
        "\n",
        "\n",
        "# this parameter is currently not exposed in the form\n",
        "max_variations_per_opt_pass = 15\n",
        "\n",
        "if optimal_ordering:\n",
        "    opt_batch_size = min(storyboard.params.n_variations, max_variations_per_opt_pass)\n",
        "\n",
        "# I think it might be more efficient to write the video to the local disk first, then move it\n",
        "# afterwards, rather than writing into google drive\n",
        "final_output_filename = str( root / output_filename )\n",
        "storyboard.params.output_filename = final_output_filename\n",
        "\n",
        "# to do: move/duplicate fps computations here (?)\n",
        "fps = storyboard.params.fps\n",
        "input_audio = storyboard.params.audio_fpath\n",
        "\n",
        "#####################################\n",
        "\n",
        "# helper function for readability\n",
        "def process_sequence(idx):\n",
        "    im_paths = get_image_sequence(idx, root)\n",
        "    images = [Image.open(fp) for fp in im_paths]\n",
        "    \n",
        "    if add_caption:\n",
        "        rec = prompt_starts[idx]\n",
        "        images = [add_caption2image(im, rec['prompt']) for im in images]\n",
        "\n",
        "    # to do: persist the ordering in the storyboard\n",
        "    if optimal_ordering:\n",
        "        images = batched_tsp_permute_frames(\n",
        "            images,\n",
        "            #max_variations_per_opt_pass\n",
        "            opt_batch_size\n",
        "        )\n",
        "    return images\n",
        "\n",
        "############################\n",
        "\n",
        "cmd_in = ['ffmpeg', '-y', '-f', 'image2pipe', '-vcodec', 'png', '-r', str(fps), '-i', '-']\n",
        "cmd_out = ['-vcodec', 'libx264', '-r', str(fps), '-pix_fmt', 'yuv420p', '-crf', '1', '-preset', 'veryslow', '-shortest', output_filename]\n",
        "\n",
        "if input_audio:\n",
        "  cmd_in += ['-i', str(input_audio)]\n",
        "\n",
        "# NB: it might be more efficient to perform this upscaling step as a \n",
        "# separate step after compiling the video frames\n",
        "if upscale:\n",
        "    height=storyboard.params.height\n",
        "    width=storyboard.params.width\n",
        "    cmd_out = ['-vf', f'scale={2*width}x{2*height}:flags=lanczos'] + cmd_out\n",
        "\n",
        "\n",
        "cmd = cmd_in + cmd_out\n",
        "\n",
        "prompt_starts = storyboard.prompt_starts\n",
        "batch_gen = BackgroundGenerator(\n",
        "    [(idx, rec, process_sequence(idx))\n",
        "        for idx, rec in enumerate(prompt_starts)]\n",
        "    ,max_prefetch=2)\n",
        "\n",
        "p = Popen(cmd, stdin=PIPE)\n",
        "for idx, rec, batch in tqdm(batch_gen, total=len(prompt_starts)): \n",
        "    frame_factory = cycle(batch)\n",
        "    k = 0\n",
        "    while k < rec['frames']:\n",
        "        im = next(frame_factory)\n",
        "        im.save(p.stdin, 'PNG')\n",
        "        k+=1\n",
        "p.stdin.close()\n",
        "\n",
        "print(\"Encoding video...\")\n",
        "p.wait()\n",
        "\n",
        "if output_filename != final_output_filename:\n",
        "    print(f\"Local video compilation complete. Moving video to: {final_output_filename}\")\n",
        "    shutil.move(output_filename, final_output_filename)\n",
        "print(\"Video complete.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "NmUedPP4iHhX"
      },
      "outputs": [],
      "source": [
        "# @title ## 📺 Enjoy your animation!\n",
        "\n",
        "# to do: Merge with 'compile' cell?\n",
        "\n",
        "output_filename = storyboard.params.output_filename\n",
        "\n",
        "download_video = True # @param {type:'boolean'}\n",
        "compress_video = False # @param {type:'boolean'}\n",
        "\n",
        "# @markdown Compressing to `*.tar.gz`` format can reduce filesize, which in turn reduces\n",
        "# @markdown your download time. You may need to install additional software\n",
        "# @markdown to \"decompress\" the file after downloading to view your video.\n",
        "\n",
        "# @markdown NB: Your video will probably download way faster from https://drive.google.com\n",
        "\n",
        "#  NB: only embed short videos\n",
        "embed_video_in_notebook = False\n",
        "\n",
        "if compress_video:\n",
        "    uncompressed_fname = output_filename\n",
        "    output_filename = f\"{output_filename}.tar.gz\"\n",
        "    print(f\"Compressing to: {output_filename}\")\n",
        "    !tar -czvf {output_filename} {uncompressed_fname}\n",
        "\n",
        "if download_video:\n",
        "    from google.colab import files\n",
        "    files.download(output_filename)\n",
        "\n",
        "if embed_video_in_notebook:\n",
        "    from IPython.display import display, Video\n",
        "    display(Video(output_filename, embed=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aVu_TleBiHhY"
      },
      "source": [
        "# ⚖️ I put on my robe and lawyer hat\n",
        "\n",
        "### Notebook license\n",
        "\n",
        "This notebook and the accompanying [git repository](https://github.com/dmarx/video-killed-the-radio-star/) and its contents are shared under the MIT license.\n",
        "\n",
        "<!-- Note to self: lawyers should really be forced to use some sort of markup or pseudocode to eliminate ambiguity \n",
        "\n",
        "...oh shit, if laws were actually described in code, we could just run queries against it\n",
        "-->\n",
        "\n",
        "```\n",
        "MIT License\n",
        "\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 all\n",
        "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 THE\n",
        "SOFTWARE.\n",
        "```\n",
        "\n",
        "### DreamStudio API TOS\n",
        "\n",
        "The default behavior of this notebook uses the [DreamStudio](https://beta.dreamstudio.ai/) API to generate images. Users of the DreamStudio API are subject to the DreamStudio usage terms: https://beta.dreamstudio.ai/terms-of-service\n",
        "\n",
        "### Stable Diffusion\n",
        "\n",
        "As of the date of this writing (2022-09-29), all publicly available model checkpoints are subject to the restrictions of the Open RAIL license: https://huggingface.co/spaces/CompVis/stable-diffusion-license. \n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3.9.2 64-bit",
      "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.9.2"
    },
    "vscode": {
      "interpreter": {
        "hash": "81794d4967e6c3204c66dcd87b604927b115b27c00565d3d43f05ba2f3a2cb0d"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}