Keep track of threads running per instance, to kill only those when switching tabs

app.py

@@ -161,9 +161,13 @@ def process_search_tta(
     sat_path: str | None,
     ground_path: str | None,
     taxonomy: str | None = None,
+    session_threads: list[threading.Thread] | None = None,
 ):
     """Run both TTA and non-TTA search episodes concurrently and stream both heat-maps."""
 
+    if session_threads is None:
+        session_threads = []
+
     # Disable Run button and clear image/status outputs, hide sliders, clear frame states
     yield (
         gr.update(interactive=False),
@@ -175,11 +179,23 @@ def process_search_tta(
         gr.update(visible=False),
         [],
         [],
+        session_threads,
     )
 
     # Bail early if satellite image missing
     if sat_path is None:
-        yield gr.update(interactive=True), gr.update(value=None), gr.update(value=None), gr.update(value="No satellite image provided.", visible=True), gr.update(value="", visible=True)
+        yield (
+            gr.update(interactive=True),
+            gr.update(value=None),
+            gr.update(value=None),
+            gr.update(value="No satellite image provided.", visible=True),
+            gr.update(value="", visible=True),
+            gr.update(visible=False),
+            gr.update(visible=False),
+            [],
+            [],
+            session_threads,
+        )
         return
 
     # Prepare PIL images
@@ -282,8 +298,8 @@ def process_search_tta(
         args=(False, gifs_dir_no, None, "no"),
         daemon=True,
     )
-    _register_thread(thread_tta)
-    _register_thread(thread_no)
+    # Track threads for this user session
+    session_threads.extend([thread_tta, thread_no])
     thread_tta.start()
     thread_no.start()
 
@@ -358,6 +374,7 @@ def process_search_tta(
                     gr.update(visible=False),
                     gr.update(),
                     gr.update(),
+                    session_threads,
                 )
                 prev_status_tta = status_tta
                 prev_status_no  = status_no
@@ -372,10 +389,8 @@ def process_search_tta(
 
     # Remove finished threads from global registry
     with _running_threads_lock:
-        if thread_tta in _running_threads:
-            _running_threads.remove(thread_tta)
-        if thread_no in _running_threads:
-            _running_threads.remove(thread_no)
+        # Clear session thread list
+        session_threads.clear()
 
     # Small delay to ensure last frame files are fully flushed
     time.sleep(0.2)
@@ -412,6 +427,7 @@ def process_search_tta(
         gr.update(visible=True, minimum=1, maximum=n_no,  value=n_no),
         frames_tta,
         frames_no,
+        session_threads,
     )
 
 
@@ -482,6 +498,7 @@ with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
 
             frames_state_tta = gr.State([])
             frames_state_no = gr.State([])
+            session_threads_state = gr.State([])
 
     # Slider callbacks (updates image when user drags slider)
     def _show_frame(idx: int, frames: list[str]):
@@ -532,8 +549,8 @@ with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
 
     run_btn.click(
         fn=process_search_tta,
-        inputs=[sat_input, ground_input, taxonomy_input],
-        outputs=[run_btn, display_img_tta, display_img_no_tta, status_tta, status_no_tta, slider_tta, slider_no, frames_state_tta, frames_state_no],
+        inputs=[sat_input, ground_input, taxonomy_input, session_threads_state],
+        outputs=[run_btn, display_img_tta, display_img_no_tta, status_tta, status_no_tta, slider_tta, slider_no, frames_state_tta, frames_state_no, session_threads_state],
     )
 
     # Footer to point out to model and data from app page.
@@ -561,14 +578,19 @@ if __name__ == "__main__":
 
         outputs_on_tab = [_cleanup_status]
 
-        def _on_tab_change(evt: gr.SelectData):
+        def _on_tab_change(evt: gr.SelectData, session_threads: list[threading.Thread]):
             # evt.value contains the name of the newly-selected tab.
             if evt.value == "Multimodal Inference":
-                _kill_running_threads()
+                # Stop only threads started in this session
+                for th in list(session_threads):
+                    if th is not None and th.is_alive():
+                        _stop_thread(th)
+                        th.join(timeout=1)
+                session_threads.clear()
                 return "Stopped running Search-TTA threads."
             return ""
 
-        tabs.select(_on_tab_change, outputs=outputs_on_tab)
+        tabs.select(_on_tab_change, inputs=[session_threads_state], outputs=outputs_on_tab)
 
     root.queue(max_size=15)
     root.launch(share=True)