Add Nanonets OCR model integration with Gradio interface

app.py

@@ -0,0 +1,116 @@
+import gradio as gr
+from PIL import Image
+from transformers import AutoTokenizer, AutoProcessor, AutoModelForImageTextToText
+import torch
+import spaces
+
+model_path = "nanonets/Nanonets-OCR-s"
+
+# Load model once at startup
+print("Loading Nanonets OCR model...")
+model = AutoModelForImageTextToText.from_pretrained(
+    model_path, 
+    torch_dtype="auto", 
+    device_map="auto", 
+    attn_implementation="flash_attention_2"
+)
+model.eval()
+
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+processor = AutoProcessor.from_pretrained(model_path)
+print("Model loaded successfully!")
+
+@spaces.GPU()
+def ocr_image_gradio(image, max_tokens=4096):
+    """Process image through Nanonets OCR model for Gradio interface"""
+    if image is None:
+        return "Please upload an image."
+    
+    try:
+        prompt = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the <img></img> tag; otherwise, add the image caption inside <img></img>. 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."""
+        
+        # Convert PIL image if needed
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+        
+        messages = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": [
+                {"type": "image", "image": image},
+                {"type": "text", "text": prompt},
+            ]},
+        ]
+        
+        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt")
+        inputs = inputs.to(model.device)
+        
+        with torch.no_grad():
+            output_ids = model.generate(**inputs, max_new_tokens=max_tokens, do_sample=False)
+            generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, output_ids)]
+            
+        output_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+        return output_text[0]
+        
+    except Exception as e:
+        return f"Error processing image: {str(e)}"
+
+# Create Gradio interface
+with gr.Blocks(title="Nanonets OCR Demo", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🔍 Nanonets OCR - Document Text Extraction")
+    gr.Markdown("Upload an image of a document to extract text, tables, equations, and more!")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_input = gr.Image(
+                label="Upload Document Image",
+                type="pil",
+                height=400
+            )
+            max_tokens_slider = gr.Slider(
+                minimum=1024,
+                maximum=8192,
+                value=4096,
+                step=512,
+                label="Max Tokens",
+                info="Maximum number of tokens to generate"
+            )
+            extract_btn = gr.Button("Extract Text", variant="primary", size="lg")
+            
+        with gr.Column(scale=2):
+            output_text = gr.Textbox(
+                label="Extracted Text",
+                lines=20,
+                max_lines=30,
+                show_copy_button=True,
+                placeholder="Extracted text will appear here..."
+            )
+    
+    # Example images section
+    gr.Markdown("## 📄 Try with example images:")
+    gr.Examples(
+        examples=[
+            ["examples/sample1.jpg"] if "examples/sample1.jpg" else None,
+            ["examples/sample2.png"] if "examples/sample2.png" else None,
+        ],
+        inputs=image_input,
+        label="Sample Documents"
+    )
+    
+    # Event handlers
+    extract_btn.click(
+        fn=ocr_image_gradio,
+        inputs=[image_input, max_tokens_slider],
+        outputs=output_text,
+        show_progress=True
+    )
+    
+    image_input.change(
+        fn=ocr_image_gradio,
+        inputs=[image_input, max_tokens_slider],
+        outputs=output_text,
+        show_progress=True
+    )
+
+if __name__ == "__main__":
+    demo.queue().launch()

requirements.txt

@@ -0,0 +1,3 @@
+transformers
+torch
+spaces