Fix UI not resetting after cancel - add reset handler

- Add reset_ui_after_cancel function to properly reset UI components
- Chain .then() after cancel button click to reset UI state
- Clear cancel_event flag for next generation
- Ensures Submit button becomes interactive again after cancellation

app.py

@@ -693,6 +693,16 @@ with gr.Blocks(title="LLM Inference with ZeroGPU") as demo:
         """Called by the cancel button, sets the global event."""
         cancel_event.set()
         print("Cancellation signal sent.")
+    
+    def reset_ui_after_cancel():
+        """Reset UI components after cancellation."""
+        cancel_event.clear()  # Clear the flag for next generation
+        print("UI reset after cancellation.")
+        return {
+            txt: gr.update(interactive=True),
+            submit_btn: gr.update(interactive=True),
+            cancel_btn: gr.update(visible=False),
+        }
 
     # Event for submitting text via Enter key or Submit button
     submit_event = txt.submit(
@@ -707,10 +717,13 @@ with gr.Blocks(title="LLM Inference with ZeroGPU") as demo:
     )
 
     # Event for the "Cancel" button.
-    # It calls the flag-setting function and cancels the long-running submit_event.
+    # It sets the cancel flag, cancels the submit event, then resets the UI.
     cancel_btn.click(
         fn=set_cancel_flag,
         cancels=[submit_event]
+    ).then(
+        fn=reset_ui_after_cancel,
+        outputs=ui_components
     )
 
     # Listeners for updating the duration estimate