Moved model loading into thread, for better concurrency and speed

app.py

@@ -22,6 +22,7 @@ import torch
 from PIL import Image
 import json
 import copy
+import shutil
 import spaces   # integration with ZeroGPU on hf
 
 # Import configuration & RL / TTA utilities -------------------------------------------------
@@ -205,28 +206,36 @@ def process_search_tta(
         planner.gifs_path = save_dir
         return planner
 
-    # Prepare gif directories
+    # Directory paths for generated frames
     gifs_dir_tta = os.path.join(gifs_path, "with_tta")
-    gifs_dir_no = os.path.join(gifs_path, "no_tta")
-    for d in (gifs_dir_tta, gifs_dir_no):
-        os.makedirs(d, exist_ok=True)
-        # Clean previous pngs
-        for f in os.listdir(d):
-            # if f.endswith(".png"):
-            os.remove(os.path.join(d, f))
-
-    planner_tta = build_planner(True, gifs_dir_tta, clip_seg_tta_1)
-    planner_no  = build_planner(False, gifs_dir_no, clip_seg_tta_2)
-
-    # Launch both planners in background threads
-    thread_tta = threading.Thread(target=planner_tta.run_episode, args=(0,), daemon=True)
-    thread_no  = threading.Thread(target=planner_no.run_episode,  args=(0,), daemon=True)
-    thread_tta.start()
-    thread_no.start()
-
-    # Register threads so they can be cancelled when user switches tabs
+    gifs_dir_no  = os.path.join(gifs_path, "no_tta")
+
+    # Clean previous run's PNG frames before launching new planners
+    for _dir in (gifs_dir_tta, gifs_dir_no):
+        shutil.rmtree(_dir, ignore_errors=True)
+        os.makedirs(_dir, exist_ok=True)
+
+    def _planner_thread(enable_tta: bool, save_dir: str, clip_obj):
+        """Prepare directory, build planner, and run an episode (in one thread)."""
+        # Directory prepared by caller; just build planner and run episode
+        planner = build_planner(enable_tta, save_dir, clip_obj)
+        planner.run_episode(0)
+
+    # Launch both planners in background threads – preparation included
+    thread_tta = threading.Thread(
+        target=_planner_thread,
+        args=(True, gifs_dir_tta, clip_seg_tta_1),
+        daemon=True,
+    )
+    thread_no = threading.Thread(
+        target=_planner_thread,
+        args=(False, gifs_dir_no, clip_seg_tta_2),
+        daemon=True,
+    )
     _register_thread(thread_tta)
     _register_thread(thread_no)
+    thread_tta.start()
+    thread_no.start()
 
 
     sent_tta: set[str] = set()