Update app.py

app.py

@@ -3,7 +3,7 @@ from PIL import Image
 from transformers import AutoTokenizer, AutoProcessor, AutoModelForImageTextToText, TextIteratorStreamer
 import torch
 import spaces
-import threading
+from threading import Thread
 
 model_path = "nanonets/Nanonets-OCR-s"
 
@@ -62,40 +62,37 @@ def ocr_image_gradio_stream(image, max_tokens=4096):
         
         # Set up streaming
         streamer = TextIteratorStreamer(
-            tokenizer=tokenizer,
-            skip_prompt=True,
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=True
+            tokenizer, 
+            timeout=60.0, 
+            skip_prompt=True, 
+            skip_special_tokens=True
         )
         
         generation_kwargs = {
             **inputs,
             "max_new_tokens": max_tokens,
             "do_sample": False,
-            "streamer": streamer,
+            "streamer": streamer
         }
         
         # Start generation in a separate thread
-        generation_thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
-        generation_thread.start()
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
         
-        # Stream the output
-        partial_output = ""
-        for new_token in streamer:
-            partial_output += new_token
-            processed_output = process_tags(partial_output)
-            yield processed_output
+        # Stream the results
+        generated_text = ""
+        for new_text in streamer:
+            generated_text += new_text
+            processed_text = process_tags(generated_text)
+            yield processed_text
             
-        # Ensure thread completes
-        generation_thread.join()
-        
     except Exception as e:
         yield f"Error processing image: {str(e)}"
 
-# Non-streaming version as fallback
+# Keep the original function for non-streaming use if needed
 @spaces.GPU()
 def ocr_image_gradio(image, max_tokens=4096):
-    """Process image through Nanonets OCR model for Gradio interface"""
+    """Process image through Nanonets OCR model for Gradio interface (non-streaming)"""
     if image is None:
         return "Please upload an image."
     
@@ -148,9 +145,6 @@ with gr.Blocks(title="Nanonets OCR Demo") as demo:
                 💻 GitHub Repository
             </a>
         </div>
-        <p style="font-size: 0.9em; color: #10b981; font-weight: 500;">
-            ✨ Now with streaming output and support for 4 concurrent uploads!
-        </p>
     </div>
     """)
     
@@ -171,16 +165,9 @@ with gr.Blocks(title="Nanonets OCR Demo") as demo:
             )
             extract_btn = gr.Button("Extract Text", variant="primary", size="lg")
             
-            gr.Markdown("""
-            **💡 Tips:**
-            - Upload supports concurrent processing of up to 4 images
-            - Results stream in real-time as they're generated
-            - Automatic processing starts when you upload an image
-            """)
-            
         with gr.Column(scale=2):
             output_text = gr.Markdown(
-                label="Streaming model prediction",
+                label="Formatted model prediction",
                 latex_delimiters=[
                     {"left": "$$", "right": "$$", "display": True},
                     {"left": "$", "right": "$", "display": False},
@@ -194,7 +181,7 @@ with gr.Blocks(title="Nanonets OCR Demo") as demo:
                 show_copy_button=True,
             )
     
-    # Event handlers with streaming
+    # Event handlers - Updated to use streaming
     extract_btn.click(
         fn=ocr_image_gradio_stream,
         inputs=[image_input, max_tokens_slider],
@@ -240,14 +227,4 @@ for downstream processing by Large Language Models (LLMs).
 """)
 
 if __name__ == "__main__":
-    # Configure for concurrent processing with streaming support
-    demo.queue(
-        max_size=1000,  # Maximum queue size
-        default_concurrency_limit=4,  # Allow 4 concurrent requests
-        status_update_rate=0.1,  # Update status every 100ms for better streaming experience
-    ).launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True,
-        share=False
-    )
+    demo.queue().launch()