Update app.py

app.py

@@ -181,7 +181,9 @@ with gr.Blocks(title="Kiri OCR - Streaming Demo", css=css, theme=gr.themes.Soft(
                         sources=["upload", "clipboard", "webcam"]
                     )
                     
-                    doc_btn = gr.Button("⚡ Stream Text", variant="primary")
+                    with gr.Row():
+                        doc_btn = gr.Button("⚡ Stream Text", variant="primary")
+                        doc_stop = gr.Button("⏹️ Stop", variant="secondary", visible=False)
                 
                 with gr.Column(scale=1):
                     # Annotated image updates in real-time
@@ -205,7 +207,9 @@ with gr.Blocks(title="Kiri OCR - Streaming Demo", css=css, theme=gr.themes.Soft(
                         type="pil",
                         sources=["upload", "clipboard"]
                     )
-                    line_btn = gr.Button("⚡ Stream Recognize", variant="primary")
+                    with gr.Row():
+                        line_btn = gr.Button("⚡ Stream Recognize", variant="primary")
+                        line_stop = gr.Button("⏹️ Stop", variant="secondary", visible=False)
                 
                 with gr.Column(scale=1):
                     line_output_text = gr.Textbox(
@@ -215,21 +219,54 @@ with gr.Blocks(title="Kiri OCR - Streaming Demo", css=css, theme=gr.themes.Soft(
                     )
             
     
-    stop_btn = gr.Button("⏹️ Stop Generation", variant="secondary")
+    # Toggle buttons visibility
+    def toggle_doc_buttons():
+        return gr.update(visible=False), gr.update(visible=True)
     
+    def reset_doc_buttons():
+        return gr.update(visible=True), gr.update(visible=False)
+
     doc_event = doc_btn.click(
+        fn=toggle_doc_buttons,
+        outputs=[doc_btn, doc_stop]
+    ).then(
         fn=process_document_stream,
         inputs=[doc_input],
         outputs=[doc_output_img, doc_output_text]
+    ).then(
+        fn=reset_doc_buttons,
+        outputs=[doc_btn, doc_stop]
+    )
+    
+    doc_stop.click(
+        fn=reset_doc_buttons,
+        outputs=[doc_btn, doc_stop],
+        cancels=[doc_event]
     )
     
+    def toggle_line_buttons():
+        return gr.update(visible=False), gr.update(visible=True)
+        
+    def reset_line_buttons():
+        return gr.update(visible=True), gr.update(visible=False)
+
     line_event = line_btn.click(
+        fn=toggle_line_buttons,
+        outputs=[line_btn, line_stop]
+    ).then(
         fn=recognize_line_stream,
         inputs=line_input,
         outputs=line_output_text
+    ).then(
+        fn=reset_line_buttons,
+        outputs=[line_btn, line_stop]
     )
     
-    stop_btn.click(fn=None, inputs=None, outputs=None, cancels=[doc_event, line_event])
+    line_stop.click(
+        fn=reset_line_buttons,
+        outputs=[line_btn, line_stop],
+        cancels=[line_event]
+    )
 
     gr.Markdown(
         """