From 2c416eb644192dc1916b4db3374d148e0571fc1a Mon Sep 17 00:00:00 2001 From: David Marx Date: Sun, 9 Oct 2022 12:03:34 -0700 Subject: [PATCH] improved prefetch --- Video_Killed_The_Radio_Star_Defusion.ipynb | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/Video_Killed_The_Radio_Star_Defusion.ipynb b/Video_Killed_The_Radio_Star_Defusion.ipynb index 9d64f32..dfbe566 100644 --- a/Video_Killed_The_Radio_Star_Defusion.ipynb +++ b/Video_Killed_The_Radio_Star_Defusion.ipynb @@ -1029,7 +1029,7 @@ "from subprocess import Popen, PIPE\n", "\n", "from omegaconf import OmegaConf\n", - "from prefetch_generator import background\n", + "from prefetch_generator import BackgroundGenerator\n", "from tqdm.autonotebook import tqdm\n", "\n", "from vktrs.tsp import (\n", @@ -1088,7 +1088,6 @@ "#####################################\n", "\n", "# helper function for readability\n", - "@background(max_prefetch=-1)\n", "def process_sequence(idx):\n", " im_paths = get_image_sequence(idx, root)\n", " images = [Image.open(fp) for fp in im_paths]\n", @@ -1123,9 +1122,15 @@ "\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 in enumerate(tqdm(storyboard.prompt_starts)): \n", - " frame_factory = cycle(process_sequence(idx))\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",