app.py

import spaces
import json
import math
import os
import traceback
from io import BytesIO
from typing import Any, Dict, List, Optional, Tuple
import re
import time
from threading import Thread
import uuid

import gradio as gr
import requests
import torch
from PIL import Image

from transformers import (
    Qwen2_5_VLForConditionalGeneration,
    AutoProcessor,
    TextIteratorStreamer,
)

from reportlab.lib.pagesizes import A4
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.platypus import SimpleDocTemplate, Image as RLImage, Paragraph, Spacer
from reportlab.lib.units import inch

# --- Constants and Model Setup ---
MAX_INPUT_TOKEN_LENGTH = 4096
# Note: The following line correctly falls back to CPU if CUDA is not available.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
print("torch.__version__ =", torch.__version__)
print("torch.version.cuda =", torch.version.cuda)
print("cuda available:", torch.cuda.is_available())
print("cuda device count:", torch.cuda.device_count())
if torch.cuda.is_available():
    print("current device:", torch.cuda.current_device())
    print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))

print("Using device:", device)

# --- Prompts for Different Tasks ---
ocr_prompt = "Perform precise OCR on the image. Extract all text content, maintaining the original structure, paragraphs, and tables as formatted markdown."

# --- Model Loading ---
MODEL_ID_M = "prithivMLmods/Camel-Doc-OCR-080125"
processor_m = AutoProcessor.from_pretrained(MODEL_ID_M, trust_remote_code=True)
model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_M, trust_remote_code=True, torch_dtype=torch.float16
).to(device).eval()

MODEL_ID_T = "prithivMLmods/Megalodon-OCR-Sync-0713"
processor_t = AutoProcessor.from_pretrained(MODEL_ID_T, trust_remote_code=True)
model_t = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_T, trust_remote_code=True, torch_dtype=torch.float16
).to(device).eval()

MODEL_ID_C = "nanonets/Nanonets-OCR-s"
processor_c = AutoProcessor.from_pretrained(MODEL_ID_C, trust_remote_code=True)
model_c = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_C, trust_remote_code=True, torch_dtype=torch.float16
).to(device).eval()

MODEL_ID_G = "echo840/MonkeyOCR"
SUBFOLDER = "Recognition"
processor_g = AutoProcessor.from_pretrained(
    MODEL_ID_G, trust_remote_code=True, subfolder=SUBFOLDER
)
model_g = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_G, trust_remote_code=True, subfolder=SUBFOLDER, torch_dtype=torch.float16
).to(device).eval()

MODEL_ID_I = "allenai/olmOCR-7B-0725"
processor_i = AutoProcessor.from_pretrained(MODEL_ID_I, trust_remote_code=True)
model_i = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID_I, trust_remote_code=True, torch_dtype=torch.float16
).to(device).eval()


# --- PDF Generation Utility Function ---
def generate_pdf(image: Image.Image, text_content: str, font_size: int, line_spacing: float, alignment: str, image_size: str) -> str:
    """
    Generates a PDF document with the input image and extracted text.
    """
    if image is None or not text_content:
        raise gr.Error("Cannot generate PDF. Image or text content is missing.")

    filename = f"/tmp/output_{uuid.uuid4()}.pdf"
    doc = SimpleDocTemplate(
        filename,
        pagesize=A4,
        rightMargin=inch,
        leftMargin=inch,
        topMargin=inch,
        bottomMargin=inch
    )
    styles = getSampleStyleSheet()
    style_normal = styles["Normal"]
    style_normal.fontSize = int(font_size)
    style_normal.leading = int(font_size) * line_spacing
    style_normal.alignment = {
        "Left": 0,
        "Center": 1,
        "Right": 2,
        "Justified": 4
    }[alignment]

    story = []

    # Handle Image
    # Convert PIL image to a format reportlab can use without saving to disk
    img_buffer = BytesIO()
    image.save(img_buffer, format='PNG')
    img_buffer.seek(0)
    
    # Image size settings
    page_width, _ = A4
    available_width = page_width - 2 * inch
    image_widths = {
        "Small": available_width * 0.3,
        "Medium": available_width * 0.6,
        "Large": available_width * 0.9,
    }
    img = RLImage(img_buffer, width=image_widths[image_size], height=image.height * (image_widths[image_size]/image.width))
    story.append(img)
    story.append(Spacer(1, 12))

    # Handle Text - Replace markdown with spaces for PDF
    # A simple replacement for basic markdown, for more complex cases a proper parser would be needed
    cleaned_text = text_content.replace("# ", "").replace("## ", "").replace("*", "")
    text_paragraphs = cleaned_text.split('\n')
    
    for para in text_paragraphs:
        if para.strip():
            story.append(Paragraph(para, style_normal))

    doc.build(story)
    return filename


# --- Core Application Logic ---
@spaces.GPU
def process_document_stream(model_name: str, image: Image.Image, max_new_tokens: int):
    """
    Main generator function that handles OCR tasks.
    """
    if image is None:
        yield "Please upload an image.", "Please upload an image."
        return

    # 1. Set prompt for OCR
    text_prompt = ocr_prompt

    # 2. Select model and processor
    if model_name == "Camel-Doc-OCR-080125": processor, model = processor_m, model_m
    elif model_name == "Megalodon-OCR-Sync-0713": processor, model = processor_t, model_t
    elif model_name == "Nanonets-OCR-s": processor, model = processor_c, model_c
    elif model_name == "MonkeyOCR-Recognition": processor, model = processor_g, model_g
    elif model_name == "olmOCR-7B-0725": processor, model = processor_i, model_i
    else:
        yield "Invalid model selected.", "Invalid model selected."
        return

    # 3. Prepare model inputs and streamer
    messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": text_prompt}]}]
    prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = processor(text=[prompt_full], images=[image], return_tensors="pt", padding=True, truncation=True, max_length=MAX_INPUT_TOKEN_LENGTH).to(device)
    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
    generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    # 4. Stream raw output to the UI in real-time
    buffer = ""
    for new_text in streamer:
        buffer += new_text
        buffer = buffer.replace("<|im_end|>", "")
        time.sleep(0.01)
        yield buffer , "⏳ Processing..."

    # 5. Yield the final result for both raw and formatted outputs
    yield buffer, buffer


# --- Gradio UI Definition ---
def create_gradio_interface():
    """Builds and returns the Gradio web interface."""
    css = """
    .main-container { max-width: 1400px; margin: 0 auto; }
    .process-button { border: none !important; color: white !important; font-weight: bold !important; background-color: blue !important;}
    .process-button:hover { background-color: darkblue !important; transform: translateY(-2px) !important; box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important; }
    """
    with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
        gr.HTML("""
        <div class="title" style="text-align: center">
            <h1>Tiny VLMs Lab🧪</h1>
            <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;">
                Advanced Vision-Language Model for Image Content and Layout Extraction
            </p>
        </div>
        """)

        with gr.Row():
            # Left Column (Inputs)
            with gr.Column(scale=1):
                model_choice = gr.Dropdown(
                    choices=["Camel-Doc-OCR-080125",
                             "MonkeyOCR-Recognition",
                             "olmOCR-7B-0725",
                             "Nanonets-OCR-s",
                             "Megalodon-OCR-Sync-0713"
                            ],
                    label="Select Model",
                    value="Nanonets-OCR-s"
                )
                image_input = gr.Image(label="Upload Image", type="pil", sources=['upload'])
                with gr.Accordion("Advanced Settings", open=False):
                    max_new_tokens = gr.Slider(minimum=512, maximum=8192, value=4096, step=256, label="Max New Tokens")
                    gr.Markdown("### PDF Export Settings")
                    font_size = gr.Dropdown(choices=["8", "10", "12", "14", "16", "18"], value="12", label="Font Size")
                    line_spacing = gr.Dropdown(choices=[1.0, 1.15, 1.5, 2.0], value=1.15, label="Line Spacing")
                    alignment = gr.Dropdown(choices=["Left", "Center", "Right", "Justified"], value="Left", label="Text Alignment")
                    image_size = gr.Dropdown(choices=["Small", "Medium", "Large"], value="Medium", label="Image Size in PDF")


                process_btn = gr.Button("🚀 Process Document", variant="primary", elem_classes=["process-button"], size="lg")
                clear_btn = gr.Button("🗑️ Clear All", variant="secondary")

            # Right Column (Outputs)
            with gr.Column(scale=2):
                with gr.Tabs() as tabs:
                    with gr.Tab("📝 Extracted Content"):
                        raw_output_stream = gr.Textbox(label="Raw Model Output Stream", interactive=False, lines=15, show_copy_button=True)
                        with gr.Row():
                            examples = gr.Examples(
                                examples=["examples/1.png", "examples/2.png", "examples/3.png", "examples/4.png", "examples/5.png"],
                                inputs=image_input,
                                label="Examples"
                        )
                        gr.Markdown("[Report-Bug💻](https://huggingface.co/spaces/prithivMLmods/OCR-Comparator/discussions)")
                    
                    with gr.Tab("📰 README.md"):
                        markdown_output = gr.Markdown(label="Formatted Markdown")

                    with gr.Tab("📋 PDF Preview"):
                        pdf_output_file = gr.File(label="Generated PDF Document", interactive=False)
                        generate_pdf_btn = gr.Button("📄 Generate PDF", variant="primary")


        # Event Handlers
        def clear_all_outputs():
            return None, "Raw output will appear here.", "Formatted results will appear here.", None

        process_btn.click(
            fn=process_document_stream,
            inputs=[model_choice,
                    image_input,
                    max_new_tokens],
            outputs=[raw_output_stream,
                     markdown_output]
        )
        
        generate_pdf_btn.click(
            fn=generate_pdf,
            inputs=[image_input, markdown_output, font_size, line_spacing, alignment, image_size],
            outputs=[pdf_output_file]
        )

        clear_btn.click(
            clear_all_outputs,
            outputs=[image_input,
                     raw_output_stream,
                     markdown_output,
                     pdf_output_file]
        )
    return demo

if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.queue(max_size=50).launch(share=True, ssr_mode=False, show_error=True)