Kills thread when changed tab

app.py

@@ -49,6 +49,24 @@ def _stop_thread(thread: threading.Thread):
         # If it returned >1, cleanup and fail safe
         ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(tid), None)
 
+# ──────────── Thread Registry for Cleanup on Tab Switch ─────────────
+_running_threads: list[threading.Thread] = []
+_running_threads_lock = threading.Lock()
+
+def _register_thread(th: threading.Thread):
+    """Record a newly started worker thread so we can later cancel it."""
+    with _running_threads_lock:
+        _running_threads.append(th)
+
+def _kill_running_threads():
+    """Stop all worker threads that are still alive."""
+    with _running_threads_lock:
+        for t in list(_running_threads):
+            _stop_thread(t)
+        # Clear list regardless of alive status
+        _running_threads.clear()
+
+
 # CHANGE ME!
 POLL_INTERVAL = 1.0 # For visualization
 
@@ -206,6 +224,11 @@ def process_search_tta(
     thread_tta.start()
     thread_no.start()
 
+    # Register threads so they can be cancelled when user switches tabs
+    _register_thread(thread_tta)
+    _register_thread(thread_no)
+
+
     sent_tta: set[str] = set()
     sent_no:  set[str] = set()
     last_tta = None
@@ -278,6 +301,13 @@ def process_search_tta(
                 _stop_thread(th)
                 th.join(timeout=1)
 
+    # 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)
+
     # One last scan after both threads have finished to catch any frame
     # that may have been written just before termination but after the last
     # polling iteration.
@@ -441,12 +471,29 @@ with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
 # if def main
 if __name__ == "__main__":
 
-    # Combine backup and current demos into a single tabbed interface (Backup first).
+    # Build UI with explicit Tabs so we can detect tab selection and clean up
     from app_multimodal_inference import demo as multimodal_demo
 
-    combined_demo = gr.TabbedInterface(
-        [demo, multimodal_demo],
-        ["Search-TTA", "Multimodal Inference"]
-    )
-    combined_demo.queue(max_size=15)
-    combined_demo.launch(share=True)
+    with gr.Blocks() as root:
+        with gr.Tabs() as tabs:
+            with gr.TabItem("Search-TTA"):
+                demo.render()
+            with gr.TabItem("Multimodal Inference"):
+                multimodal_demo.render()
+
+        # Hidden textbox purely to satisfy Gradio's need for an output component.
+        _cleanup_status = gr.Textbox(visible=False)
+
+        outputs_on_tab = [_cleanup_status]
+
+        def _on_tab_change(evt: gr.SelectData):
+            # evt.value contains the name of the newly-selected tab.
+            if evt.value == "Multimodal Inference":
+                _kill_running_threads()
+                return "Stopped running Search-TTA threads."
+            return ""
+
+        tabs.select(_on_tab_change, outputs=outputs_on_tab)
+
+    root.queue(max_size=15)
+    root.launch(share=True)