ERROR HANDLING GPU ABORT Hugging Face

app/ui_utils.py

@@ -397,63 +397,81 @@ def start_task_asr_ast(
         streaming_policy, alignatt_thr, waitk_lagging,
         exclude_sink_frames, xatt_scores_layer, hallucinations_detector
     )
-
+    try:
     # Loop over partial results
     # result_data: can be transcribed text OR an info message depending on status
-    for result_data, status, debug_info in task_generator:
-        
-        # 1. Default states for this iteration ('In Progress' mode)
-        # By default, lock config and allow stopping
-        start_btn = gr.update(visible=False)
-        stop_btn = gr.update(visible=True)
-        config_step = gr.update(interactive=False)
-        
-        # Status message and main text depend on return type
-        status_msg = gr.update(visible=True)
-        main_output = accumulated_text
-
-        # --- CASE 1: SUCCESS (New text segment) ---
-        if status == "success":
-            # result_data is the new text chunk here
-            partial_text = result_data
-            
-            # Update accumulator
-            accumulated_text += partial_text
-            main_output = accumulated_text
+        for result_data, status, debug_info in task_generator:
             
-            # Status message displays chunk info (e.g., timestamps)
-            status_msg = gr.update(visible=True, value=debug_info, elem_classes=[status])
-
-        # --- CASE 2: WARNING / INFO (System message) ---
-        elif status in ["warning", "info"]:
-            # result_data is the error or info message here
-            # Do not touch accumulated_text
-            status_msg = gr.update(visible=True, value=result_data, elem_classes=[status])
-
-
-        # --- CASE 3: DONE / ERROR---
-        elif status in ["done", "error"]:
-            logging.error(f"[ui] error ")
-            # Re-enable controls
-            is_streaming = is_active_stream(session_hash_code)
-            start_btn = gr.update(visible=is_streaming) # Show Start only if audio stream is active
-            stop_btn = gr.update(visible=False)
-            config_step = gr.update(interactive=True)
+            # 1. Default states for this iteration ('In Progress' mode)
+            # By default, lock config and allow stopping
+            start_btn = gr.update(visible=False)
+            stop_btn = gr.update(visible=True)
+            config_step = gr.update(interactive=False)
             
-            # result_data is the completion message
-            status_msg = gr.update(visible=True, value=result_data, elem_classes=[status])
+            # Status message and main text depend on return type
+            status_msg = gr.update(visible=True)
+            main_output = accumulated_text
 
-        # 2. Single dispatch to UI
-        # Expected order: [task_output, status_message_task, start_task_button, stop_task_button, config_step]
+            # --- CASE 1: SUCCESS (New text segment) ---
+            if status == "success":
+                # result_data is the new text chunk here
+                partial_text = result_data
+                
+                # Update accumulator
+                accumulated_text += partial_text
+                main_output = accumulated_text
+                
+                # Status message displays chunk info (e.g., timestamps)
+                status_msg = gr.update(visible=True, value=debug_info, elem_classes=[status])
+
+            # --- CASE 2: WARNING / INFO (System message) ---
+            elif status in ["warning", "info"]:
+                # result_data is the error or info message here
+                # Do not touch accumulated_text
+                status_msg = gr.update(visible=True, value=result_data, elem_classes=[status])
+
+
+            # --- CASE 3: DONE / ERROR---
+            elif status in ["done", "error"]:
+                logging.error(f"[ui] error ")
+                # Re-enable controls
+                is_streaming = is_active_stream(session_hash_code)
+                start_btn = gr.update(visible=is_streaming) # Show Start only if audio stream is active
+                stop_btn = gr.update(visible=False)
+                config_step = gr.update(interactive=True)
+                
+                # result_data is the completion message
+                status_msg = gr.update(visible=True, value=result_data, elem_classes=[status])
+
+            # 2. Single dispatch to UI
+            # Expected order: [task_output, status_message_task, start_task_button, stop_task_button, config_step]
+            yield (
+                main_output,
+                status_msg,
+                start_btn,
+                stop_btn,
+                config_step
+            )
+    except Exception as e:
+        # --- ERROR HANDLING (GPU ABORT / RUNTIME ERROR) ---
+        error_msg = str(e)
+        logging.error(f"Task Error for {session_hash_code}: {error_msg}", exc_info=True)
+
+        # Detect specific Hugging Face / GPU errors
+        if "GPU task aborted" in error_msg or "CUDA out of memory" in error_msg or "Device" in error_msg:
+            display_msg = f"🛑 **System Error:** GPU Task Aborted. The model may have run out of memory. ({error_msg})"
+        else:
+            display_msg = f"⚠️ **Task Error:** {error_msg}"
+
+        # Update UI to reflect the crash
         yield (
-            main_output,
-            status_msg,
-            start_btn,
-            stop_btn,
-            config_step
+            accumulated_text, # Keep what we managed to generate so far
+            gr.update(visible=True, value=display_msg, elem_classes=["error"]), # Show error in status box
+            gr.update(visible=True),  # Re-enable Start button to allow retry
+            gr.update(visible=False), # Hide Stop button
+            gr.update(interactive=True) # Unlock config
         )
 
-
 def stop_task_fn(session_hash_code):
     remove_active_task_flag_file(session_hash_code)
     yield "Task stopped by user."