Fix UI state reset after generation cancellation

app.py

@@ -656,6 +656,8 @@ with gr.Blocks(title="LLM Inference with ZeroGPU") as demo:
             cancel_btn: gr.update(visible=True),
         }
 
+        ui_reset_done = False
+
         try:
             # 2. Call the backend and stream updates
             backend_args = [user_msg, chat_history] + list(args)
@@ -667,7 +669,13 @@ with gr.Blocks(title="LLM Inference with ZeroGPU") as demo:
         except GeneratorExit:
             # Handle Gradio's cancellation signal gracefully
             print("Generation cancelled by user.")
-            # Don't yield anything here - let the finally block handle UI reset
+            # Immediately reset UI state before returning
+            yield {
+                txt: gr.update(interactive=True),
+                submit_btn: gr.update(interactive=True),
+                cancel_btn: gr.update(visible=False),
+            }
+            ui_reset_done = True
             return
         except Exception as e:
             print(f"An error occurred during generation: {e}")
@@ -679,12 +687,14 @@ 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.
-            print("Resetting UI state.")
-            yield {
-                txt: gr.update(interactive=True),
-                submit_btn: gr.update(interactive=True),
-                cancel_btn: gr.update(visible=False),
-            }
+            # Only do this if we haven't already reset in the except block
+            if not ui_reset_done:
+                print("Resetting UI state.")
+                yield {
+                    txt: gr.update(interactive=True),
+                    submit_btn: gr.update(interactive=True),
+                    cancel_btn: gr.update(visible=False),
+                }
 
     def set_cancel_flag():
         """Called by the cancel button, sets the global event."""