app.py

import spaces
import json
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 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.lib import colors
from reportlab.platypus import SimpleDocTemplate, Image as RLImage, Paragraph, Spacer
from reportlab.lib.units import inch
import uuid

# --- Constants and Model Setup ---
MAX_INPUT_TOKEN_LENGTH = 4096
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)

# --- 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()

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

# --- PDF Generation Functions ---
def generate_pdf(media_path, plain_text, font_size, line_spacing, alignment, image_size):
    """Generates a PDF document."""
    filename = f"output_{uuid.uuid4()}.pdf"
    doc = SimpleDocTemplate(
        filename,
        pagesize=A4,
        rightMargin=inch,
        leftMargin=inch,
        topMargin=inch,
        bottomMargin=inch
    )
    styles = getSampleStyleSheet()
    styles["Normal"].fontSize = int(font_size)
    styles["Normal"].leading = int(font_size) * line_spacing
    styles["Normal"].alignment = {
        "Left": 0,
        "Center": 1,
        "Right": 2,
        "Justified": 4
    }[alignment]

    story = []

    # Add image with size adjustment
    image_sizes = {
        "Small": (200, 200),
        "Medium": (400, 400),
        "Large": (600, 600)
    }
    img = RLImage(media_path, width=image_sizes[image_size][0], height=image_sizes[image_size][1])
    story.append(img)
    story.append(Spacer(1, 12))

    # Add plain text output
    text = Paragraph(plain_text, styles["Normal"])
    story.append(text)

    doc.build(story)
    return filename

# --- Core Application Logic ---
@spaces.GPU
def process_document_stream(model_name: str, image: Image.Image, max_new_tokens: int, font_size: str, line_spacing: float, alignment: str, image_size: str):
    """
    Main generator function for OCR task, also generating PDF for preview.
    """
    if image is None:
        yield "Please upload an image.", "Please upload an image.", None
        return

    # 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.", None
        return

    # Save image temporarily for PDF generation
    temp_image_path = f"temp_{uuid.uuid4()}.png"
    image.save(temp_image_path)

    # Prepare model inputs and streamer
    text_prompt = ocr_prompt
    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()

    # 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)
        # Generate PDF with current buffer
        pdf_file = generate_pdf(temp_image_path, buffer, font_size, line_spacing, alignment, image_size)
        yield buffer, buffer, pdf_file

    # Final PDF with complete output
    pdf_file = generate_pdf(temp_image_path, buffer, font_size, line_spacing, alignment, image_size)
    yield buffer, buffer, pdf_file

    # Clean up temporary image file
    if os.path.exists(temp_image_path):
        os.remove(temp_image_path)

# --- 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; }
    .download-btn { background-color: #35a6d6 !important; color: white !important; }
    .download-btn:hover { background-color: #22bcff !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 Extraction and PDF Generation
            </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")
                    font_size = gr.Dropdown(
                        choices=["8", "10", "12", "14", "16", "18", "20", "22", "24"],
                        value="16",
                        label="Font Size"
                    )
                    line_spacing = gr.Dropdown(
                        choices=[0.5, 1.0, 1.15, 1.5, 2.0, 2.5, 3.0],
                        value=1.5,
                        label="Line Spacing"
                    )
                    alignment = gr.Dropdown(
                        choices=["Left", "Center", "Right", "Justified"],
                        value="Justified",
                        label="Text Alignment"
                    )
                    image_size = gr.Dropdown(
                        choices=["Small", "Medium", "Large"],
                        value="Medium",
                        label="Image Size"
                    )

                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=13, 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"):
                        with gr.Accordion("(Formatted Result)", open=True):
                            markdown_output = gr.Markdown(label="Formatted Markdown")
                    with gr.Tab("📋 PDF Preview"):
                        pdf_output = gr.File(label="Download PDF", interactive=True)

        # 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, font_size, line_spacing, alignment, image_size],
            outputs=[raw_output_stream, markdown_output, pdf_output]
        )
        clear_btn.click(
            fn=clear_all_outputs,
            outputs=[image_input, raw_output_stream, markdown_output, pdf_output]
        )
    return demo

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