Update app.py

app.py

@@ -38,9 +38,12 @@ def process_tags(content: str) -> str:
 def ocr_image_gradio_stream(image, max_tokens=4096):
     """Process image through Nanonets OCR model with streaming output"""
     if image is None:
-        yield "Please upload an image."
+        yield "Please upload an image.", gr.update(interactive=True)
         return
     
+    # Disable button at start of processing
+    yield "Processing image...", gr.update(interactive=False)
+    
     try:
         prompt = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes."""
         
@@ -84,10 +87,14 @@ def ocr_image_gradio_stream(image, max_tokens=4096):
         for new_text in streamer:
             generated_text += new_text
             processed_text = process_tags(generated_text)
-            yield processed_text
+            yield processed_text, gr.update(interactive=False)
+        
+        # Re-enable button when complete
+        yield processed_text, gr.update(interactive=True)
             
     except Exception as e:
-        yield f"Error processing image: {str(e)}"
+        error_msg = f"Error processing image: {str(e)}"
+        yield error_msg, gr.update(interactive=True)
 
 # Keep the original function for non-streaming use if needed
 @spaces.GPU()
@@ -185,14 +192,14 @@ with gr.Blocks(title="Nanonets OCR Demo") as demo:
     extract_btn.click(
         fn=ocr_image_gradio_stream,
         inputs=[image_input, max_tokens_slider],
-        outputs=output_text,
+        outputs=[output_text, extract_btn],
         show_progress=True
     )
     
     image_input.change(
         fn=ocr_image_gradio_stream,
         inputs=[image_input, max_tokens_slider],
-        outputs=output_text,
+        outputs=[output_text, extract_btn],
         show_progress=True
     )