improved prefetch

yyahav · Oct 9, 2022 · 2c416eb · 2c416eb
1 parent 7f3e2a1
commit 2c416eb
Showing 1 changed file with 9 additions and 4 deletions.
diff --git 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",