From ac21ea58b872ba81f3b2a64ba0f03ee63cc9f3db Mon Sep 17 00:00:00 2001 From: David Marx Date: Sat, 8 Oct 2022 15:45:58 -0700 Subject: [PATCH] light refactoring (#87) --- VERSION | 2 +- Video_Killed_The_Radio_Star_Defusion.ipynb | 192 ++++++++------------- vktrs/utils.py | 61 ++++++- 3 files changed, 125 insertions(+), 130 deletions(-) diff --git a/VERSION b/VERSION index 845639e..9faa1b7 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -0.1.4 +0.1.5 diff --git a/Video_Killed_The_Radio_Star_Defusion.ipynb b/Video_Killed_The_Radio_Star_Defusion.ipynb index a122a7c..6ab2deb 100644 --- a/Video_Killed_The_Radio_Star_Defusion.ipynb +++ b/Video_Killed_The_Radio_Star_Defusion.ipynb @@ -232,7 +232,7 @@ "id": "rt9Mu97fk_bp" }, "source": [ - "## $1.$ 📋 Create New Project" + "## $1.$ 📋 Create New Project (or set name to resume)" ] }, { @@ -253,21 +253,13 @@ "# @markdown Non-alphanumeric characters (excluding '-' and '_') will be replaced with hyphens.\n", "\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", - "import string\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 = sanitize_folder_name(project_name)\n", "\n", "workspace = OmegaConf.load('config.yaml')\n", @@ -627,6 +619,8 @@ " 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", @@ -642,7 +636,6 @@ " df_pre = copy.deepcopy(df)\n", " for i, rec in enumerate(prompt_starts):\n", " rec['ts'] = float(df.loc[i,'Timestamp (sec)'])\n", - " #rec['td'] = dt.timedelta(rec['ts'])\n", " rec['td'] = rec['ts']\n", " rec['prompt'] = df.loc[i,'Lyric']\n", " \n", @@ -660,8 +653,9 @@ " OmegaConf.save(config=storyboard, f=fp.name)\n", "\n", "\n", + "#####################################\n", "# @title ## 🎨 Generate init images\n", - "\n", + "#####################################\n", "\n", "workspace = OmegaConf.load('config.yaml')\n", "root = Path(workspace.project_root)\n", @@ -714,62 +708,28 @@ " return images\n", "\n", "\n", - "# to do: move this to utils\n", - "def get_image_sequence(idx, root=root, init_first=True):\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", - "\n", - "def archive_images(idx, 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)\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", "d_ = dict(\n", " _=''\n", " , theme_prompt = \"deepdream, mural designed by a local artist\" # @param {type:'string'}\n", - "\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", ")\n", + "d_.pop('_')\n", "\n", "regenerate_all_init_images = False # @param {type:'boolean'}\n", "\n", + "prompt_lag = False # @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", @@ -778,14 +738,13 @@ "if regenerate_all_init_images:\n", " for i, rec in enumerate(prompt_starts):\n", " rec['frame0_fpath'] = None\n", - " archive_images(i)\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", - "# to do: some sort of uniqueness + archival system linked to theme prompt\n", "if 'df_regen' in locals():\n", " for i, _ in df_regen.iterrows():\n", " rec = prompt_starts[i]\n", @@ -796,10 +755,9 @@ " regen=True\n", " if regen:\n", " rec['frame0_fpath'] = None\n", - " archive_images(i)\n", + " archive_images(i, root=root)\n", " print(\"archival process complete\")\n", "\n", - "prompt_lag = False # @param {type:'boolean'}\n", "\n", "\n", "theme_prompt = storyboard.params.theme_prompt\n", @@ -832,9 +790,7 @@ " rec['frame0_fpath'] = save_frame(\n", " init_image,\n", " idx,\n", - " #root_path=Path('./frames') / proj_name,\n", - " #name=proj_name, ## to do.... uh... i dunno\n", - " root_path = root / 'frames', # to do: this field should accept a string as well\n", + " root_path = root / 'frames',\n", " name='anchor',\n", " )\n", "\n", @@ -842,8 +798,10 @@ " print(lyric)\n", " display(init_image)\n", "\n", - "########################\n", - "# update config\n", + "\n", + "##############\n", + "# checkpoint #\n", + "##############\n", "\n", "prompt_starts_copy = copy.deepcopy(prompt_starts)\n", "\n", @@ -858,10 +816,10 @@ "with open(storyboard_fname) as fp:\n", " OmegaConf.save(config=storyboard, f=fp.name)\n", "\n", - "######################################################\n", - "\n", - "# flag regens in the table\n", "\n", + "###############\n", + "# flag regens #\n", + "###############\n", "\n", "df_regen = pd.DataFrame(prompt_starts)[['ts','prompt']].rename(\n", " columns={\n", @@ -905,7 +863,11 @@ "\n", "from omegaconf import OmegaConf\n", "from PIL import Image\n", - "from vktrs.utils import add_caption2image\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", @@ -915,52 +877,35 @@ "\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", + " # to do: this should be rendered unnecessary before this branch is merged\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", - "# to do: move to utils\n", - "def get_image_sequence(idx, root=root, init_first=True):\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", - "\n", - "########################\n", - "# update config\n", - "\n", - "prompt_starts_copy = copy.deepcopy(prompt_starts)\n", - "\n", - "# to do: this should be rendered unnecessary before this branch is merged\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", + " 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", - "### This block computes how many frames are needed for each segment\n", - "### based on the start times for each prompt\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", + "# * don't write timedelta objects into the prompt_starts... yeesh\n", "\n", "fps = 12 # @param {type:'integer'}\n", "storyboard.params.fps = fps\n", @@ -971,8 +916,7 @@ "video_duration = storyboard.params['video_duration']\n", "\n", "# dummy prompt for last scene duration\n", - "#prompt_starts = OmegaConf.to_container(prompt_starts, resolve=True)\n", - "prompt_starts = OmegaConf.to_container(storyboard.prompt_starts) # I don't think I need to resolve here..\n", + "prompt_starts = OmegaConf.to_container(storyboard.prompt_starts)\n", "for rec in prompt_starts:\n", " rec['td'] = dt.timedelta(seconds=rec['td'])\n", "prompt_starts.append({'td':dt.timedelta(seconds=video_duration)})\n", @@ -1004,7 +948,9 @@ "# and guesstimate a corrected prompt start time and duration \n", "\n", "\n", - "### checkpoint the processing work we've done to this point\n", + "##############\n", + "# checkpoint #\n", + "##############\n", "\n", "prompt_starts_copy = copy.deepcopy(prompt_starts)\n", "\n", @@ -1019,9 +965,9 @@ " OmegaConf.save(config=storyboard, f=fp.name)\n", "\n", "\n", - "\n", - "## 🚀 Generate animation frames\n", - "\n", + "##################################\n", + "# Generate animation frames #\n", + "##################################\n", "\n", "d_ = dict(\n", " _=''\n", @@ -1042,12 +988,10 @@ "# @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", - "\n", "storyboard.params.update(d_)\n", "storyboard.params.max_frames = storyboard.params.fps * storyboard.params.max_video_duration_in_seconds\n", "\n", - "print(f\"Max total frames: {storyboard.params.max_frames}\")\n", - "#print(f\"Max API requests: {int(max_frames/repeat)}\")\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", @@ -1062,7 +1006,7 @@ "print(\"Fetching variations\")\n", "for idx, rec in enumerate(prompt_starts):\n", " new_images = []\n", - " images_fpaths = get_image_sequence(idx)\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", @@ -1084,8 +1028,10 @@ " if display_frames_as_we_get_them:\n", " display(img)\n", "\n", - "########################\n", - "# update config\n", + "\n", + "##############\n", + "# checkpoint #\n", + "##############\n", "\n", "prompt_starts_copy = copy.deepcopy(prompt_starts)\n", "\n", @@ -1113,6 +1059,7 @@ "source": [ "# @title ## 🎞️ Compile your video!\n", "\n", + "import shutil\n", "from subprocess import Popen, PIPE\n", "\n", "from omegaconf import OmegaConf\n", @@ -1125,10 +1072,12 @@ "\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", @@ -1158,22 +1107,18 @@ "# this parameter is currently not exposed in the form\n", "max_variations_per_opt_pass = 15\n", "\n", - "\n", "if optimal_ordering:\n", " opt_batch_size = min(storyboard.params.n_variations, max_variations_per_opt_pass)\n", "\n", + "\n", "#####################################\n", "# video parameters\n", "\n", - "\n", - "#output_filename = str( root / output_filename )\n", - "#storyboard.params.output_filename = output_filename\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", - "\n", "# to do: move/duplicate fps computations here (?)\n", "fps = storyboard.params.fps\n", "input_audio = storyboard.params.audio_fpath\n", @@ -1229,7 +1174,6 @@ "\n", "if output_filename != final_output_filename:\n", " print(f\"Local video compilation complete. Moving video to: {final_output_filename}\")\n", - " import shutil\n", " shutil.move(output_filename, final_output_filename)\n", "print(\"Video complete.\")" ] diff --git a/vktrs/utils.py b/vktrs/utils.py index 304babe..8a82b0c 100644 --- a/vktrs/utils.py +++ b/vktrs/utils.py @@ -7,6 +7,7 @@ import pandas as pd from PIL import Image, ImageDraw, ImageFont + def gpu_info(): outv = subprocess.run([ 'nvidia-smi', @@ -40,11 +41,24 @@ def get_audio_duration_seconds(audio_fpath): return float(outv.strip()) +def rand_str(n_char=5): + return ''.join(random.choice(string.ascii_lowercase) for i in range(n_char)) + + def remove_punctuation(s): # https://stackoverflow.com/a/266162/819544 return s.translate(str.maketrans('', '', string.punctuation)) +def sanitize_folder_name(fp): + outv = '' + whitelist = string.ascii_letters + string.digits + '-_' + for token in str(fp): + if token not in whitelist: + token = '-' + outv += token + return outv + def add_caption2image( image, @@ -84,10 +98,6 @@ def add_caption2image( return image -def rand_str(n_char=5): - return ''.join(random.choice(string.ascii_lowercase) for i in range(n_char)) - - def save_frame( img: Image, idx:int=0, @@ -99,4 +109,45 @@ def save_frame( name = rand_str() outpath = root_path / f"{idx}-{name}.png" img.save(outpath) - return str(outpath) \ No newline at end of file + return str(outpath) + + +def get_image_sequence(idx, root, init_first=True): + root = Path(root) + images = (root / 'frames' ).glob(f'{idx}-*.png') + images = [str(fp) for fp in images] + if init_first: + init_image = None + images2 = [] + for i, fp in enumerate(images): + if 'anchor' in fp: + init_image = fp + else: + images2.append(fp) + if not init_image: + try: + init_image, images2 = images2[0], images2[1:] + images = [init_image] + images2 + except IndexError: + images = images2 + return images + + +def archive_images(idx, root, archive_root = None): + root = Path(root) + if archive_root is None: + archive_root = root / 'archive' + archive_root = Path(archive_root) + archive_root.mkdir(parents=True, exist_ok=True) + old_images = get_image_sequence(idx, root=root) + if not old_images: + return + print(f"moving {len(old_images)} old images for scene {idx} to {archive_root}") + for old_fp in old_images: + old_fp = Path(old_fp) + im_name = Path(old_fp.name) + new_path = archive_root / im_name + if new_path.exists(): + im_name = f"{im_name.stem}-{time.time()}{im_name.suffix}" + new_path = archive_root / im_name + old_fp.rename(new_path) \ No newline at end of file