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."""
    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."""
    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.
    
    Args:
        messages: List of message dictionaries with role and content
        model_name: Name of the model (unused but kept for compatibility)
        max_tokens: Maximum number of tokens to generate
        temperature: Temperature for generation (unused, model runs deterministically)
    
    Yields:
        str: Generated text chunks
    """
    # 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
    """
    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.
    
    Args:
        file_path: Path to uploaded file
        max_tokens: Maximum tokens per page
    
    Returns:
        Generator yielding markdown content
    """
    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 ---
with gr.Blocks(title="PDF to Markdown Converter", theme=gr.themes.Soft()) as demo:
    gr.HTML("""
    <div class="title" style="text-align: center">
        <h1>📄 Nanonets-OCR-s: PDF & Image to Markdown Converter</h1>
        <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;">
            Powered by <strong>Nanonets-OCR-s</strong>, A model for transforming documents into structured markdown with intelligent content recognition and semantic tagging.
        </p>
        <div style="display: flex; justify-content: center; gap: 20px; margin: 15px 0;">
            <a href="https://huggingface.co/nanonets/Nanonets-OCR-s" target="_blank" style="text-decoration: none; color: #2563eb; font-weight: 500;">
                📚 Hugging Face Model
            </a>
            <a href="https://nanonets.com/research/nanonets-ocr-s/" target="_blank" style="text-decoration: none; color: #2563eb; font-weight: 500;">
                📝 Release Blog
            </a>
            <a href="https://github.com/NanoNets/docext" target="_blank" style="text-decoration: none; color: #2563eb; font-weight: 500;">
                💻 GitHub Repository
            </a>
        </div>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.Image(
                label="Upload Image Document",
                height=200
            )
            max_tokens_slider = gr.Slider(
                minimum=1024,
                maximum=8192,
                value=4096,
                step=512,
                label="Max Tokens per Page",
                info="Maximum number of new tokens to generate for each page."
            )
            with_img_desc_checkbox = gr.Checkbox(
                label="Include Image Description",
                value=False,
                info="If enabled, the model will include a description of the image in the output. If no image is present, use with_img_desc=False."
            )
            extract_btn = gr.Button("Convert to Markdown", variant="primary", size="lg")

        with gr.Column(scale=2):
            output_text = gr.Markdown(
                label="Formatted Model Prediction",
                latex_delimiters=[{"left": "$$", "right": "$$", "display": True}, {"left": "$", "right": "$", "display": False}],
                line_breaks=True,
                show_copy_button=True,
                height=600,
            )

    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)