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",
        "tags": []
      },
      "outputs": [],
      "source": [
        "# @title # 📊 Check GPU Status\n",
        "\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": "cADf0Y3pk5Uj",
        "tags": []
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "# @title # 🛠️ Installations\n",
        "\n",
        "try: \n",
        "    import google.colab\n",
        "    local=False\n",
        "except:\n",
        "    local=True\n",
        "\n",
        "# local only additional dependencies\n",
        "if local:\n",
        "    %pip install pandas torch pillow beautifulsoup4 scipy toolz numpy lxml\n",
        "\n",
        "# dependencies for both colab and local\n",
        "%pip install yt-dlp python-tsp stability-sdk diffusers transformers ftfy accelerate omegaconf\n",
        "%pip install openai-whisper  panel prefetch_generator huggingface_hub ipywidgets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "cM8cux9b7F4v",
        "tags": []
      },
      "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",
        "try: \n",
        "    import google.colab\n",
        "    local=False\n",
        "except:\n",
        "    local=True\n",
        "\n",
        "if local:\n",
        "    mount_gdrive=False\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",
        "    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",
        "    if not local:\n",
        "        from google.colab import output\n",
        "        output.enable_custom_widget_manager()\n",
        "        \n",
        "    from huggingface_hub import notebook_login\n",
        "    notebook_login()"
      ]
    },
    {
      "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",
        "tags": []
      },
      "outputs": [],
      "source": [
        "\n",
        "import time\n",
        "import string\n",
        "from omegaconf import OmegaConf\n",
        "\n",
        "def sanitize_folder_name(fp):\n",
        "    outv = ''\n",
        "    whitelist = string.ascii_letters + string.digits + '-_'\n",
        "    for token in str(fp):\n",
        "        if token not in whitelist:\n",
        "            token = '-'\n",
        "        outv += token\n",
        "    return outv\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",
        "tags": []
      },
      "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",
        "import subprocess\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",
        "# TODO: add support for whisper API\n",
        "import whisper\n",
        "\n",
        "storyboard = OmegaConf.create()\n",
        "\n",
        "d_ = dict(\n",
        "    # all the underscore 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",
        ")\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",
        "\n",
        "d_.pop('_')\n",
        "storyboard.params = d_\n",
        "\n",
        "if not storyboard.params.audio_fpath:\n",
        "    storyboard.params.audio_fpath = None\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",
        "# 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",
        "        ytdl_prefix = \"DOWNLOADED__\"\n",
        "        ytdl_fname = f\"{str(root / ytdl_prefix)}%(title)s.%(ext)s\"\n",
        "\n",
        "        # force re-download for now... existing mp4 causes download step to be skipped.\n",
        "        # TODO: intelligently detect when to re-download or not.\n",
        "        # -> pretty sure this is mainly going to be changing the checkpointing sequence,\n",
        "        #    i.e. so we can compare the input video_url with whatever may already be on the storyboard\n",
        "        #if Path(ytdl_fname).exists():\n",
        "        #    Path(ytdl_fname).unlink()\n",
        "        \n",
        "        !yt-dlp -o \"{ytdl_fname}\" {video_url}\n",
        "\n",
        "        matched_files = root.glob(ytdl_prefix+\"*\")\n",
        "        most_recent_file = max(matched_files, key=os.path.getctime)\n",
        "        print(f\"downloaded: {most_recent_file}\")\n",
        "        ytdl_fname = most_recent_file\n",
        "        # new attribute cause why not\n",
        "        storyboard.params.downloaded_video_fpath = ytdl_fname\n",
        "\n",
        "        #audio_fpath = str( root / 'audio.aac' )\n",
        "        audio_fpath = str( root / 'audio.m4a' )\n",
        "        input_audio = ytdl_fname\n",
        "        #!ffmpeg -y -i \"{input_audio}\" -c:a aac {audio_fpath}\n",
        "        #!ffmpeg -y -i \"{input_audio}\" -vn -acodec copy {audio_fpath}\n",
        "        !ffmpeg -y -i \"{input_audio}\" -vn -c:a aac {audio_fpath}\n",
        "\n",
        "        storyboard.params.audio_fpath = audio_fpath\n",
        "    \n",
        "with open(storyboard_fname,'wb') as fp:\n",
        "    OmegaConf.save(config=storyboard, f=fp.name)\n",
        "\n",
        "\n",
        "###################################################\n",
        "# 💬 Transcribe and segment speech using whisper #\n",
        "###################################################\n",
        "\n",
        "# for video duration\n",
        "def get_audio_duration_seconds(audio_fpath):\n",
        "    outv = subprocess.run([\n",
        "        'ffprobe'\n",
        "        ,'-i',audio_fpath\n",
        "        ,'-show_entries', 'format=duration'\n",
        "        ,'-v','quiet'\n",
        "        ,'-of','csv=p=0'\n",
        "        ],\n",
        "        stdout=subprocess.PIPE\n",
        "        ).stdout.decode('utf-8')\n",
        "    return float(outv.strip())\n",
        "\n",
        "audio_fpath = Path(storyboard.params.audio_fpath)\n",
        "#!whisper --model large --word_timestamps True {storyboard.params.audio_fpath}\n",
        "!whisper --model large --word_timestamps True -o {str(root)} {storyboard.params.audio_fpath}\n",
        "\n",
        "# outputs text files as audio.* locally\n",
        "#with Path('audio.json').open() as f:\n",
        "with Path(root / 'audio.json').open() as f:\n",
        "    timings = json.load(f)\n",
        "\n",
        "\n",
        "# TODO: this doesn't need to be a function...\n",
        "def whisper_segments_to_vktrs_promptstarts(segments):\n",
        "    for rec in segments:\n",
        "        rec['ts'] = rec['start']\n",
        "        rec['prompt'] = rec['text']\n",
        "    return segments\n",
        "\n",
        "storyboard.prompt_starts = whisper_segments_to_vktrs_promptstarts(timings['segments'])\n",
        "\n",
        "# i don't think this is reliable unfortunately.\n",
        "#storyboard.params['video_duration'] = storyboard.prompt_starts[-1]['end']\n",
        "storyboard.params['video_duration'] = get_audio_duration_seconds(audio_fpath)\n",
        "\n",
        "prompt_starts = storyboard.prompt_starts\n",
        "\n",
        "### checkpoint the processing work we've done to this point\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",
        "# 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",
        "prompt_starts = OmegaConf.to_container(prompt_starts)\n",
        "df = pd.DataFrame(prompt_starts).rename(\n",
        "    columns={\n",
        "        'ts':'Timestamp (sec)',\n",
        "        'prompt':'Lyric',\n",
        "    }\n",
        ")[['start','end','Lyric']]\n",
        "\n",
        "if 'td' in df:\n",
        "    del df['td']\n",
        "\n",
        "df['override_prompt'] = ''\n",
        "\n",
        "df_pre = copy.deepcopy(df)\n",
        "outv = pn.widgets.Tabulator(df, formatters=tabulator_formatters)\n",
        "#if local: # really the issue isn't \"local\", it's VSCode\n",
        "#    outv = df\n",
        "outv\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3b5oOU8Tk5Uk"
      },
      "outputs": [],
      "source": [
        "# @markdown Run this cell to preview an alternative segmentation\n",
        "# @markdown of the lyrics. If you like what you see, run the next \n",
        "# @markdown cell to keep this segmentation, or skip\n",
        "# @markdown the next cell to keep your lyrics segmented as above.\n",
        "try_to_segment_further = True # @param {'type':'boolean'}\n",
        "\n",
        "def calculate_interword_gaps(segment):\n",
        "    end_prev = -1\n",
        "    gaps = []\n",
        "    for word in segment['words']:\n",
        "        if end_prev < 0:\n",
        "            end_prev = word['end']\n",
        "            continue \n",
        "        gap = word['start'] - end_prev\n",
        "        gaps.append(gap)\n",
        "        end_prev = word['end']\n",
        "    return gaps \n",
        "\n",
        "def trivial_subsegmentation(segment, threshold=0, gaps=None):\n",
        "    \"\"\"\n",
        "    split on gaps in detected vocal activity. \n",
        "    Contiguity = gap between adjacent tokens is less than the input threshold.\n",
        "    \"\"\"\n",
        "    if gaps is None:\n",
        "        gaps = calculate_interword_gaps(seg)\n",
        "    out_segments = []\n",
        "    this_segment = [seg['words'][0]]\n",
        "    for word, preceding_pause in zip(seg['words'][1:], gaps):\n",
        "        if preceding_pause <= threshold:\n",
        "            this_segment.append(word)\n",
        "        else:\n",
        "            out_segments.append(this_segment)\n",
        "            this_segment = [word]\n",
        "    out_segments.append(this_segment)\n",
        "\n",
        "    outv = [dict(\n",
        "        start=seg[0]['start'],\n",
        "        end=seg[-1]['end'],\n",
        "        text=''.join([w['word'] for w in seg]).strip(),\n",
        "    ) for seg in out_segments]\n",
        "\n",
        "    outv = whisper_segments_to_vktrs_promptstarts(outv)\n",
        "    return outv\n",
        "\n",
        "segments = timings['segments']\n",
        "segments_df = df\n",
        "if try_to_segment_further:\n",
        "    segments = []\n",
        "    for seg in timings['segments']:\n",
        "        segments.extend(trivial_subsegmentation(seg))\n",
        "\n",
        "    segments_df = pd.DataFrame(segments)\n",
        "pn.widgets.Tabulator(segments_df.rename(\n",
        "    columns={\n",
        "        'ts':'Timestamp (sec)',\n",
        "        'prompt':'Lyric',\n",
        "    })[['start', 'end','Lyric']], formatters=tabulator_formatters)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "NQmupEmIk5Ul"
      },
      "outputs": [],
      "source": [
        "# @markdown Run this cell (with the box ticked) to use the alternative version of the segmentation.\n",
        "# @markdown The \"alternative\" or \"further segmented\" version should be broken up into more scenes (or will be unchanged).\n",
        "\n",
        "use_new_segmentation = False # @param {'type':'boolean'}\n",
        "\n",
        "if use_new_segmentation:\n",
        "    prompt_starts=segments\n",
        "    storyboard.prompt_starts = prompt_starts\n",
        "\n",
        "    with open(storyboard_fname) as fp:\n",
        "        OmegaConf.save(config=storyboard, f=fp.name)\n",
        "\n",
        "#prompt_starts = OmegaConf.to_container(prompt_starts)\n",
        "df = pd.DataFrame(prompt_starts).rename(\n",
        "    columns={\n",
        "        'ts':'Timestamp (sec)',\n",
        "        'prompt':'Lyric',\n",
        "    }\n",
        ")[['start','end','Lyric']]\n",
        "\n",
        "# pretty sure can delete this..?\n",
        "if 'td' in df:\n",
        "    del df['td']\n",
        "\n",
        "df['override_prompt'] = ''\n",
        "\n",
        "df_pre = copy.deepcopy(df)\n",
        "outv = pn.widgets.Tabulator(df, formatters=tabulator_formatters)\n",
        "#if local: # really the issue isn't \"local\", it's VSCode\n",
        "#    outv = df\n",
        "outv"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_RTUFeyQqCfd"
      },
      "source": [
        "## $3.$ 🎬 Animate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Sh514DGj_sua",
        "tags": []
      },
      "outputs": [],
      "source": [
        "#####################################\n",
        "# @title ## 🎨 Generate init images\n",
        "#####################################\n",
        "\n",
        "\n",
        "import copy\n",
        "import datetime as dt\n",
        "from pathlib import Path\n",
        "import random\n",
        "import string\n",
        "import time\n",
        "import os\n",
        "import io\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 PIL import Image\n",
        "from tqdm.autonotebook import tqdm\n",
        "\n",
        "import torch\n",
        "from torch import autocast\n",
        "from diffusers import (\n",
        "    StableDiffusionImg2ImgPipeline,\n",
        "    StableDiffusionPipeline,\n",
        ")\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['start'] = float(df.loc[i,'start'])\n",
        "            rec['ts'] = rec['start']\n",
        "            rec['end'] = float(df.loc[i,'end'])\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",
        "\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",
        "########################################\n",
        "\n",
        "# misc utils\n",
        "\n",
        "def rand_str(n_char=5):\n",
        "    return ''.join(random.choice(string.ascii_lowercase) for i in range(n_char))\n",
        "\n",
        "def save_frame(\n",
        "    img: Image,\n",
        "    idx:int=0,\n",
        "    root_path=Path('./frames'),\n",
        "    name=None,\n",
        "):\n",
        "    root_path.mkdir(parents=True, exist_ok=True)\n",
        "    if name is None:\n",
        "        name = rand_str()\n",
        "    outpath = root_path / f\"{idx}-{name}.png\"\n",
        "    img.save(outpath)\n",
        "    return str(outpath)\n",
        "\n",
        "def get_image_sequence(idx, root, init_first=True):\n",
        "    root = Path(root)\n",
        "    images = (root / 'frames' ).glob(f'{idx}-*.png')\n",
        "    images = [str(fp) for fp in images]\n",
        "    if init_first:\n",
        "        init_image = None\n",
        "        images2 = []\n",
        "        for i, fp in enumerate(images):\n",
        "            if 'anchor' in fp:\n",
        "                init_image = fp\n",
        "            else:\n",
        "                images2.append(fp)\n",
        "        if not init_image:\n",
        "            try:\n",
        "                init_image, images2 = images2[0], images2[1:]\n",
        "                images = [init_image] + images2\n",
        "            except IndexError:\n",
        "                images = images2\n",
        "    return images\n",
        "\n",
        "def archive_images(idx, root, archive_root = None):\n",
        "    root = Path(root)\n",
        "    if archive_root is None:\n",
        "        archive_root = root / 'archive'\n",
        "    archive_root = Path(archive_root)\n",
        "    archive_root.mkdir(parents=True, exist_ok=True)\n",
        "    old_images = get_image_sequence(idx, root=root)\n",
        "    if not old_images:\n",
        "        return\n",
        "    print(f\"moving {len(old_images)} old images for scene {idx} to {archive_root}\")\n",
        "    for old_fp in old_images:\n",
        "        old_fp = Path(old_fp)\n",
        "        im_name = Path(old_fp.name)\n",
        "        new_path = archive_root / im_name\n",
        "        if new_path.exists():\n",
        "            im_name = f\"{im_name.stem}-{time.time()}{im_name.suffix}\"\n",
        "            new_path = archive_root / im_name\n",
        "        old_fp.rename(new_path)\n",
        "\n",
        "############################\n",
        "\n",
        "device = 'cuda'\n",
        "model_id = \"CompVis/stable-diffusion-v1-4\"\n",
        "download=True\n",
        "\n",
        "use_stability_api = workspace.use_stability_api\n",
        "model_dir = workspace.model_dir\n",
        "model_path= str(Path(model_dir) / 'huggingface' / 'diffusers')\n",
        "\n",
        "\n",
        "if 'get_image_for_prompt' not in locals():\n",
        "\n",
        "    if use_stability_api:\n",
        "        import warnings\n",
        "        from stability_sdk import client\n",
        "        import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation\n",
        "\n",
        "        # TODO: update this stuff to reflect updates to API/sdk\n",
        "        def get_image_for_prompt(prompt, max_retries=5, **kargs):\n",
        "            stability_api = client.StabilityInference(\n",
        "                key=os.environ['STABILITY_KEY'], \n",
        "                verbose=False,\n",
        "            )\n",
        "\n",
        "            # auto-retry if mitigation triggered\n",
        "            while max_retries:\n",
        "                try:\n",
        "                    answers = stability_api.generate(prompt=prompt, **kargs)\n",
        "                    response = process_response(answers)\n",
        "                    for img in response:\n",
        "                        yield img\n",
        "                    break # NB: this breaks us out of the while loop, not the for loop.\n",
        "\n",
        "                # TODO: better regen handling\n",
        "                except RuntimeError:\n",
        "                    print(\"runtime error\")\n",
        "                    max_retries -= 1\n",
        "                    warnings.warn(f\"mitigation triggered, retries remaining: {max_retries}\")\n",
        "\n",
        "\n",
        "        def process_response(answers):\n",
        "            for resp in answers:\n",
        "                for artifact in resp.artifacts:\n",
        "                    if artifact.finish_reason == generation.FILTER:\n",
        "                        warnings.warn(\n",
        "                            \"Your request activated the API's safety filters and could not be processed.\"\n",
        "                            \"Please modify the prompt and try again.\")\n",
        "                        raise RuntimeError\n",
        "                    if artifact.type == generation.ARTIFACT_IMAGE:\n",
        "                        img = Image.open(io.BytesIO(artifact.binary))\n",
        "                        yield img\n",
        "\n",
        "    else:\n",
        "\n",
        "        if download:\n",
        "            img2img = StableDiffusionImg2ImgPipeline.from_pretrained(\n",
        "                model_id,\n",
        "                revision=\"fp16\", \n",
        "                torch_dtype=torch.float16,\n",
        "                use_auth_token=True\n",
        "            )\n",
        "            img2img = img2img.to(device)\n",
        "            img2img.save_pretrained(model_path)\n",
        "        else:\n",
        "            img2img = StableDiffusionImg2ImgPipeline.from_pretrained(\n",
        "                model_path,\n",
        "                local_files_only=True\n",
        "            ).to(device)\n",
        "\n",
        "        text2img = StableDiffusionPipeline(\n",
        "            vae=img2img.vae,\n",
        "            text_encoder=img2img.text_encoder,\n",
        "            tokenizer=img2img.tokenizer,\n",
        "            unet=img2img.unet,\n",
        "            feature_extractor=img2img.feature_extractor,\n",
        "            scheduler=img2img.scheduler,\n",
        "            safety_checker=img2img.safety_checker,\n",
        "        )\n",
        "        text2img.enable_attention_slicing()\n",
        "        img2img.enable_attention_slicing()\n",
        "\n",
        "\n",
        "        def get_image_for_prompt_hf(\n",
        "            prompt,\n",
        "            **kwargs\n",
        "        ):\n",
        "            f = text2img if kwargs.get('image') is None else img2img\n",
        "            n_retries = 5\n",
        "            with autocast(device):\n",
        "                while n_retries > 0:\n",
        "                    n_retries-=1\n",
        "                    result = f(prompt, **kwargs)\n",
        "                    if not any(result.nsfw_content_detected):\n",
        "                        return result.images\n",
        "                    else:\n",
        "                        print(f\"nsfw content detectected. retries remaining: {n_retries}\")\n",
        "\n",
        "        def get_image_for_prompt(*args, **kargs):\n",
        "            if 'init_image' in kargs:\n",
        "                kargs['image'] = kargs.pop('init_image')\n",
        "            if 'start_schedule' in kargs:\n",
        "                kargs['strength'] = kargs.pop('start_schedule')\n",
        "            return get_image_for_prompt_hf(*args, **kargs)\n",
        "\n",
        "###############################\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 = \"retrofuturism\" # @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",
        "\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",
        "    print(\n",
        "        f\"\\n[{idx} | {rec['ts']}] - {lyric} - {prompt}\"\n",
        "    )\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",
        "tags": []
      },
      "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",
        "\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",
        "\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=3 # @param {type:'integer'}\n",
        "    , image_consistency=0.7 # @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=True\n",
        "#display_frames_as_we_get_them = storyboard.params.display_frames_as_we_get_them\n",
        "\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",
        "tags": []
      },
      "outputs": [],
      "source": [
        "# @title ## 📺 Compile your video and enjoy your animation!\n",
        "\n",
        "import shutil\n",
        "\n",
        "# to do: skip tsp if n_variations ==1\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",
        "# 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",
        "download_video = True # @param {type:'boolean'}\n",
        "# @markdown NB: Your video will probably download way faster from https://drive.google.com\n",
        "\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",
        "# TODO: make sure image is being written to correct location?\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",
        "\n",
        "\n",
        "#####################################\n",
        "\n",
        "import time\n",
        "\n",
        "import numpy as np\n",
        "from scipy.spatial.distance import pdist, squareform\n",
        "from itertools import cycle\n",
        "from python_tsp.exact import solve_tsp_dynamic_programming\n",
        "\n",
        "import textwrap\n",
        "from PIL import Image, ImageDraw, ImageFont\n",
        "\n",
        "def tsp_sort(frames):\n",
        "    frames_m = np.array([np.array(f).ravel() for f in frames])\n",
        "    dmat = pdist(frames_m, metric='cosine')\n",
        "    dmat = squareform(dmat)\n",
        "    permutation, _ = solve_tsp_dynamic_programming(dmat)\n",
        "    return permutation\n",
        "\n",
        "def add_caption2image(\n",
        "      image, \n",
        "      caption, \n",
        "      text_font='LiberationSans-Regular.ttf', \n",
        "      font_size=20,\n",
        "      fill_color=(255, 255, 255),\n",
        "      stroke_color=(0, 0, 0), #stroke_fill\n",
        "      stroke_width=2,\n",
        "      align='center',\n",
        "      ):\n",
        "    # via https://stackoverflow.com/a/59104505/819544\n",
        "    wrapper = textwrap.TextWrapper(width=50) \n",
        "    word_list = wrapper.wrap(text=caption) \n",
        "    caption_new = ''\n",
        "    for ii in word_list[:-1]:\n",
        "        caption_new = caption_new + ii + '\\n'\n",
        "    caption_new += word_list[-1]\n",
        "\n",
        "    draw = ImageDraw.Draw(image)\n",
        "\n",
        "    # Download the Font and Replace the font with the font file. \n",
        "    font = ImageFont.truetype(text_font, size=font_size)\n",
        "    w,h = draw.textsize(caption_new, font=font, stroke_width=stroke_width)\n",
        "    W,H = image.size\n",
        "    x,y = 0.5*(W-w),0.90*H-h\n",
        "    draw.text(\n",
        "        (x,y), \n",
        "        caption_new,\n",
        "        font=font,\n",
        "        fill=fill_color, \n",
        "        stroke_fill=stroke_color,\n",
        "        stroke_width=stroke_width,\n",
        "        align=align,\n",
        "    )\n",
        "\n",
        "    return image\n",
        "\n",
        "# prep everything...\n",
        "ffmpeg_cmd_script = \"\"\n",
        "for idx, rec in enumerate(storyboard.prompt_starts):\n",
        "    if 'frame_order' not in rec:\n",
        "        im_paths = get_image_sequence(idx, root)\n",
        "        # to do: persist the ordering in the storyboard\n",
        "        if optimal_ordering:\n",
        "            images = [Image.open(fp) for fp in im_paths]\n",
        "            frame_order = tsp_sort(images)\n",
        "            im_paths = [im_paths[j] for j in frame_order]\n",
        "            images = [images[j] for j in frame_order]\n",
        "        # TODO: actually persist frame order to storyboard...\n",
        "        rec['frame_order'] = im_paths\n",
        "    else:\n",
        "        im_paths = rec['frame_order']\n",
        "\n",
        "    images = [Image.open(fp) for fp in im_paths]\n",
        "\n",
        "    if add_caption:\n",
        "        new_paths = []\n",
        "        images_captioned = [add_caption2image(im, rec['prompt']) for im in images]\n",
        "        for fp, im in zip(im_paths, images_captioned):\n",
        "            fp = Path(fp)\n",
        "            #fp = fp.with_stem(fp.stem + '-captioned')\n",
        "            fp = fp.parent / 'captioned' / fp.name\n",
        "            fp.parent.mkdir(exist_ok=True, parents=True)\n",
        "            # TODO: just save a captioned version of every image while they're generating\n",
        "            im.save(fp)\n",
        "            new_paths.append(fp)\n",
        "        im_paths = new_paths\n",
        "    \n",
        "    frame_picker = cycle(im_paths)\n",
        "    for _ in range(rec.frames):\n",
        "        fpath = next(frame_picker)\n",
        "        ffmpeg_cmd_script += f\"file '{fpath}'\\nduration {1/fps}\\n\"\n",
        "    \n",
        "    with open(root/'scenes.txt', 'w') as f:\n",
        "        f.write(ffmpeg_cmd_script)\n",
        "\n",
        "\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",
        "    !cd {root}; ffmpeg -y -f concat -safe 0 -i {'scenes.txt'} -i {storyboard.params.audio_fpath} -r {storyboard.params.fps} -pix_fmt yuv420p -crf 15 -preset slow -vf scale={2*width}x{2*height}:flags=lanczos -shortest {storyboard.params.output_filename}\n",
        "else:\n",
        "    #!ffmpeg -y -f concat -i {root / 'scenes.txt'} -i {storyboard.params.audio_fpath} -r {storyboard.params.fps} -pix_fmt yuv420p -crf 15 -preset veryslow -shortest {storyboard.params.output_filename}\n",
        "    !cd {root}; ffmpeg -y -f concat -safe 0 -i {'scenes.txt'} -i {storyboard.params.audio_fpath} -r {storyboard.params.fps} -pix_fmt yuv420p -crf 15 -preset veryfast -shortest {storyboard.params.output_filename}\n",
        "\n",
        "\n",
        "# EASTER EGG FEATURE\n",
        "#  NB: only embed short videos\n",
        "embed_video_in_notebook = False\n",
        "\n",
        "output_filename = storyboard.params.output_filename\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": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.10.6"
    },
    "vscode": {
      "interpreter": {
        "hash": "81794d4967e6c3204c66dcd87b604927b115b27c00565d3d43f05ba2f3a2cb0d"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}