Handle GeneratorExit by yielding UI reset and returning cleanly

app.py

@@ -665,6 +665,15 @@ with gr.Blocks(title="LLM Inference with ZeroGPU") as demo:
                     chat: response_chunk[0],
                     dbg: response_chunk[1],
                 }
+        except GeneratorExit:
+            # Handle GeneratorExit by resetting UI and returning cleanly
+            print("Generation cancelled by user.")
+            yield {
+                txt: gr.update(interactive=True),
+                submit_btn: gr.update(interactive=True),
+                cancel_btn: gr.update(visible=False),
+            }
+            return  # Don't re-raise, just return cleanly
         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.
@@ -674,7 +683,7 @@ 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.
+            # Reset UI for normal completion or errors
             print("Resetting UI state.")
             yield {
                 txt: gr.update(interactive=True),