Fix GeneratorExit handling by re-raising after UI reset

app.py

@@ -1,6 +1,7 @@
 import os
 import time
 import gc
+import sys
 import threading
 from itertools import islice
 from datetime import datetime
@@ -669,14 +670,13 @@ with gr.Blocks(title="LLM Inference with ZeroGPU") as demo:
         except GeneratorExit:
             # Handle Gradio's cancellation signal gracefully
             print("Generation cancelled by user.")
-            # Immediately reset UI state before returning
+            # Immediately reset UI state, then re-raise to properly close generator
             yield {
                 txt: gr.update(interactive=True),
                 submit_btn: gr.update(interactive=True),
                 cancel_btn: gr.update(visible=False),
             }
-            ui_reset_done = True
-            return
+            raise  # Re-raise GeneratorExit to properly close the generator
         except Exception as e:
             print(f"An error occurred during generation: {e}")
             # If an error happens, add it to the chat history to inform the user.
@@ -687,8 +687,8 @@ with gr.Blocks(title="LLM Inference with ZeroGPU") as demo:
             yield {chat: error_history}
         finally:
             # 3. ALWAYS reset the UI to an "idle" state upon completion, error, or cancellation.
-            # Only do this if we haven't already reset in the except block
-            if not ui_reset_done:
+            # Skip if GeneratorExit was raised (already handled above)
+            if not isinstance(sys.exc_info()[1], GeneratorExit):
                 print("Resetting UI state.")
                 yield {
                     txt: gr.update(interactive=True),