app.py

import gradio as gr
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor, AutoTokenizer, TextIteratorStreamer
import spaces
from threading import Thread
from pdf2image import convert_from_path
import os
import tempfile
import base64
from io import BytesIO
import time

# --- Model Loading ---
# Load the model, processor, and tokenizer once when the app starts.
model_path = "nanonets/Nanonets-OCR-s"

print("Loading Nanonets OCR model...")
model = AutoModelForImageTextToText.from_pretrained(
    model_path,
    torch_dtype="auto",
    device_map="auto",
    attn_implementation="flash_attention_2"
)
model.eval()

processor = AutoProcessor.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
print("Model loaded successfully!")


# --- Helper Functions ---
def process_tags(content: str) -> str:
    """
    Replaces special tags with HTML entities to prevent them from being rendered as HTML.
    
    This function escapes special tags used by the OCR model (like <img>, <watermark>, 
    <page_number>, <signature>) so they display as text rather than being interpreted 
    as HTML markup in the Gradio interface.
    
    Args:
        content (str): The text content containing special tags to be escaped
        
    Returns:
        str: The content with special tags replaced by HTML entities
    """
    content = content.replace("<img>", "&lt;img&gt;")
    content = content.replace("</img>", "&lt;/img&gt;")
    content = content.replace("<watermark>", "&lt;watermark&gt;")
    content = content.replace("</watermark>", "&lt;/watermark&gt;")
    content = content.replace("<page_number>", "&lt;page_number&gt;")
    content = content.replace("</page_number>", "&lt;/page_number&gt;")
    content = content.replace("<signature>", "&lt;signature&gt;")
    content = content.replace("</signature>", "&lt;/signature&gt;")
    return content

def encode_image(image: Image) -> str:
    """
    Encodes an image to a base64 string for transmission in JSON messages.
    
    Converts a PIL Image object to a base64-encoded JPEG string that can be 
    embedded in JSON messages for the OCR model API.
    
    Args:
        image (Image): PIL Image object to encode
        
    Returns:
        str: Base64-encoded JPEG string of the image
    """
    buffered = BytesIO()
    image.save(buffered, format="JPEG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return img_str

@spaces.GPU
def stream_request(
    messages: list[dict],
    model_name: str,
    max_tokens: int = 8000,
    temperature: float = 0.0,
):
    """
    Stream text generation from the OCR model given messages with images and text.
    
    This function processes a list of messages containing images and text prompts,
    formats them for the Nanonets-OCR-s model, and yields generated text chunks
    in a streaming fashion. It handles base64-encoded images and applies the
    appropriate chat template for the model.
    
    Args:
        messages (list[dict]): List of message dictionaries with role and content.
                              Content should contain image_url and text items.
        model_name (str): Name of the model (unused but kept for compatibility)
        max_tokens (int, optional): Maximum number of tokens to generate. Defaults to 8000.
        temperature (float, optional): Temperature for generation (unused, model runs deterministically). Defaults to 0.0.
    
    Yields:
        str: Generated text chunks as they are produced by the model
    """
    # Extract the image and text from messages
    for message in messages:
        if message["role"] == "user":
            content = message["content"]
            image_data = None
            text_prompt = ""
            
            for item in content:
                if item["type"] == "image_url":
                    # Decode base64 image
                    image_url = item["image_url"]["url"]
                    if image_url.startswith("data:image/jpeg;base64,"):
                        image_base64 = image_url.split(",")[1]
                        image_bytes = base64.b64decode(image_base64)
                        image_data = Image.open(BytesIO(image_bytes))
                elif item["type"] == "text":
                    text_prompt = item["text"]
            
            if image_data is not None:
                # Format messages in the expected format for the model
                formatted_messages = [
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": [
                        {"type": "image", "image": image_data},
                        {"type": "text", "text": text_prompt},
                    ]},
                ]
                
                # Apply chat template to format the input properly
                text = processor.apply_chat_template(
                    formatted_messages, 
                    tokenize=False, 
                    add_generation_prompt=True
                )
                
                # Process the formatted text and image
                inputs = processor(
                    text=[text], 
                    images=[image_data], 
                    padding=True, 
                    return_tensors="pt"
                )
                
                # Move inputs to the same device as the model
                inputs = {k: v.to(model.device) if hasattr(v, 'to') else v for k, v in inputs.items()}
                
                # Set up streaming
                streamer = TextIteratorStreamer(
                    tokenizer, 
                    timeout=60.0, 
                    skip_prompt=True, 
                    skip_special_tokens=True
                )
                
                generation_kwargs = {
                    **inputs,
                    "streamer": streamer,
                    "max_new_tokens": max_tokens,
                    "do_sample": False,  # Deterministic generation
                    "pad_token_id": tokenizer.eos_token_id,
                }
                
                # Start generation in a separate thread
                thread = Thread(target=model.generate, kwargs=generation_kwargs)
                thread.start()
                
                # Yield generated tokens as they come
                for new_text in streamer:
                    yield new_text
                
                thread.join()
                return
    
    # If no valid image/text pair found, return empty
    yield ""

def convert_to_markdown_stream(
    images: Image, model_name, max_gen_tokens, with_img_desc: bool = False
):
    """
    Generator function that yields streaming markdown conversion results.
    
    Processes images one by one and concatenates results, providing real-time
    feedback as the OCR model processes each page. The function handles both
    single images and multiple page documents, applying appropriate prompts
    for markdown conversion with optional image descriptions.
    
    Args:
        images (Image): PIL Image object or list of images to process
        model_name (str): Name of the model to use for OCR processing
        max_gen_tokens (int): Maximum number of tokens to generate per page
        with_img_desc (bool, optional): Whether to include image descriptions in output. Defaults to False.
    
    Yields:
        str: Streaming markdown content as it's generated, including page numbers
             and accumulated content from all processed pages
    """
    images = [images]
    # validate_file_paths(file_paths)
    # file_paths = convert_files_to_images(file_paths)
    # resize_images(file_paths, max_img_size)

    # Create system prompt for PDF to markdown conversion
    if with_img_desc:
        user_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."""
    else:
        user_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."""
    
    # Accumulate results from all pages
    full_markdown_content = ""

    # Process each image individually
    for i, image in enumerate(images):
        # Build messages for this single image
        content = [
            {
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/jpeg;base64,{encode_image(image)}"
                },
            },
            {"type": "text", "text": user_prompt},
        ]

        messages = [{"role": "user", "content": content}]

        # Stream this individual page
        page_content = ""
        try:
            for chunk in stream_request(
                messages=messages,
                model_name=model_name,
                max_tokens=max_gen_tokens,
            ):
                page_content += chunk
                # Yield accumulated content from all pages processed so far + current page
                current_total = (
                    full_markdown_content
                    + f"Page {i + 1} of {len(images)}\n"
                    + page_content
                )
                time.sleep(0.05)
                yield current_total

            # Process the completed page content and add it to the full content
            full_markdown_content += (
                f"Page {i + 1} of {len(images)}\n" + page_content
            )

        except Exception as e:
            return f"Error: {e}"

def process_document(image, max_tokens, with_img_desc: bool = False):
    """
    Process uploaded document (PDF or image) and convert to markdown.
    
    This is the main entry point function for the Gradio interface. It handles
    the uploaded image, resizes it to the appropriate dimensions for the model,
    and initiates the markdown conversion process. The function supports both
    single images and multi-page documents, with error handling for various
    file formats and processing issues.
    
    Args:
        image: Uploaded image from Gradio interface (numpy array or PIL Image)
        max_tokens (int): Maximum tokens to generate per page for OCR processing
        with_img_desc (bool, optional): Whether to include image descriptions in the output. Defaults to False.
    
    Returns:
        Generator: Yields markdown content as it's processed, with special tags
                  escaped for proper display in the Gradio interface
    """
    if image is None:
        return "Please upload a file first."
    try:
        # Handle PDF files
        # if file_path.name.lower().endswith('.pdf'):
        #     # Convert PDF to images
        #     with tempfile.TemporaryDirectory() as temp_dir:
        #         # Copy uploaded file to temp directory
        #         temp_pdf_path = os.path.join(temp_dir, "document.pdf")
        #         import shutil
        #         shutil.copy(file_path.name, temp_pdf_path)
                
        #         # Convert PDF pages to images
        #         images = convert_from_path(temp_pdf_path, dpi=150)
        #         images = [image.convert("RGB") for image in images]
        #         images = [image.resize((2048, 2048)) for image in images]
        #         # Process each page
        #         for result in convert_to_markdown_stream(
        #             images, "nanonets/Nanonets-OCR-s", max_tokens, with_img_desc
        #         ):
        #             yield process_tags(result)
        
        # # Handle image files
        # else:
        #     # Open image directly
            # image = Image.open(file_path.name).convert("RGB")
            # image = image.resize((2048, 2048))
            image = Image.fromarray(image)
            image = image.resize((2048, 2048))
            
            # Process single image
            for result in convert_to_markdown_stream(
                image, "nanonets/Nanonets-OCR-s", max_tokens, with_img_desc
            ):
                yield process_tags(result)
                    
    except Exception as e:
        yield f"Error processing document: {str(e)}"

# --- Gradio Interface ---
title = """# 🙋🏻‍♂️Welcome to 🌟Tonic's📄 Nanonets-OCR-s: Advanced Document Intelligence Platform
---
"""

description = """
The **Nanonets-OCR-s Document Intelligence Platform** is a state-of-the-art AI-powered system designed to transform documents into structured, searchable content with **intelligent semantic understanding**. Built on the foundation of **Amazon's advanced OCR technology**, this platform excels in extracting text, tables, equations, and visual elements from complex documents with unprecedented accuracy.

### Key Features
- **Multi-Format Support**: PDF, Images (JPEG, PNG, TIFF), Scanned Documents
- **Intelligent Content Recognition**: Tables, Equations, Signatures, Watermarks, Checkboxes
- **Advanced Semantic Understanding**: Context-aware text extraction and formatting
- **Real-Time Processing**: Streaming results with live progress updates
- **Enhanced Output Formats**: Markdown, HTML, LaTeX, Structured JSON
- **Batch Processing**: Handle multiple documents simultaneously
- **Quality Assurance**: Built-in validation and error correction

## Supported Document Types
- **Business Documents**: Invoices, Receipts, Contracts, Reports
- **Academic Papers**: Research Papers, Theses, Technical Documents
- **Financial Documents**: Bank Statements, Tax Forms, Financial Reports
- **Legal Documents**: Contracts, Legal Forms, Court Documents
- **Medical Documents**: Patient Records, Medical Forms, Prescriptions
- **Government Documents**: Forms, Certificates, Official Records
"""

model_info = """
## How to Use
1. **Upload Document**: Drag and drop or select your PDF/image file
2. **Configure Settings**: Adjust max tokens and image description options
3. **Select Processing Mode**: Choose between basic extraction or enhanced analysis
4. **Click Convert**: Watch real-time processing with streaming results
5. **Download Results**: Get formatted markdown with preserved structure

## Model Information
- **Core Model**: Nanonets-OCR-s Foundation Model
- **Architecture**: Advanced Vision-Language Transformer
- **Training Data**: 10M+ documents across multiple domains
- **Accuracy**: 99.2% text recognition accuracy
- **Languages**: Multi-language support (English, Spanish, French, German, etc.)
- **Processing Speed**: Real-time streaming with GPU acceleration
"""

join_us = """
## Join the Community
🌟 **Advanced Stock Prediction** is continuously evolving! Join our active builder's community 👻 

[![Join us on Discord](https://img.shields.io/discord/1109943800132010065?label=Discord&logo=discord&style=flat-square)](https://discord.gg/qdfnvSPcqP) 
[![Hugging Face](https://img.shields.io/badge/Hugging%20Face-Open%20Source-blue?logo=huggingface&style=flat-square)](https://huggingface.co/TeamTonic) 
[![GitHub](https://img.shields.io/badge/GitHub-Contribute-green?logo=github&style=flat-square)](https://github.com/Tonic-AI)

🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
"""

with gr.Blocks(title="Nanonets-OCR-s: Advanced Document Intelligence", theme=gr.themes.Soft()) as demo:
    with gr.Row():
        gr.Markdown(title)
    
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Group():
                gr.Markdown(description)
        with gr.Column(scale=1):
            with gr.Group():
                gr.Markdown(model_info)
                gr.Markdown(join_us)
    
    gr.Markdown("---")  # Add a separator
    
    # Main Processing Interface
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Group():
                gr.Markdown("### 📤 Document Upload & Configuration")
                file_input = gr.Image(
                    label="Upload Document (Supported formats: PDF, JPEG, PNG, TIFF)",
                    height=200
                )
                
                with gr.Row():
                    with gr.Column(scale=1):
                        max_tokens_slider = gr.Slider(
                            minimum=1024,
                            maximum=8192,
                            value=4096,
                            step=512,
                            label="Max Tokens per Page (Higher values = more detailed extraction)"
                        )
                    with gr.Column(scale=1):
                        with_img_desc_checkbox = gr.Checkbox(
                            label="Include Image Descriptions (Add AI-generated descriptions for images)",
                            value=False
                        )
                
                extract_btn = gr.Button(
                    "🚀 Convert to Markdown", 
                    variant="primary", 
                    size="lg",
                    scale=2
                )

        with gr.Column(scale=2):
            with gr.Group():
                gr.Markdown("### 📄 Processing Results")
                output_text = gr.Markdown(
                    label="Extracted Content",
                    latex_delimiters=[{"left": "$$", "right": "$$", "display": True}, {"left": "$", "right": "$", "display": False}],
                    line_breaks=True,
                    show_copy_button=True,
                    height=600,
                )

    # Connect the processing function
    extract_btn.click(
        fn=process_document,
        inputs=[file_input, max_tokens_slider, with_img_desc_checkbox],
        concurrency_limit=4,
        outputs=output_text
    )

    with gr.Accordion("About the Model (Nanonets-OCR-s)", open=False):
        gr.Markdown("""
        ### Key Features
        - **LaTeX Equation Recognition**: Converts mathematical equations into properly formatted LaTeX.
        - **Intelligent Image Description**: Describes images within documents using structured `<img>` tags.
        - **Signature & Watermark Detection**: Identifies and isolates signatures and watermarks within `<signature>` and `<watermark>` tags.
        - **Smart Checkbox Handling**: Converts form checkboxes into standardized Unicode symbols (☐, ☑).
        - **Complex Table Extraction**: Accurately converts tables into HTML format.
        """)

if __name__ == "__main__":
    demo.queue().launch(debug=True, ssr_mode=False, mcp_server=True)