Update app.py

app.py

@@ -1,229 +1,167 @@
+import os
+import random
+import uuid
+import json
+import time
+import asyncio
+from threading import Thread
+
 import gradio as gr
 import spaces
 import torch
-from gradio_pdf import PDF
-from pdf2image import convert_from_path
-from PIL import Image
-from transformers import AutoModelForImageTextToText, AutoProcessor, AutoTokenizer
-
-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",
+import numpy as np
+from PIL import Image, ImageOps
+# import cv2  # not needed anymore
+
+from transformers import (
+    Qwen2_5_VLForConditionalGeneration,
+    AutoProcessor,
+    TextIteratorStreamer,
 )
-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"""
+from transformers.image_utils import load_image
+
+# Optional docling imports (unused now but kept for easy re-enable)
+# from docling_core.types.doc import DoclingDocument, DocTagsDocument
+
+import re
+import ast
+import html
+
+# ---------------------------
+# Constants & device
+# ---------------------------
+MAX_MAX_NEW_TOKENS     = 2048
+DEFAULT_MAX_NEW_TOKENS = 1024
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+# ---------------------------
+# Load ONLY Typhoon OCR 20B
+# ---------------------------
+MODEL_ID = "scb10x/typhoon-ocr-20b"  # <- 20B model
+processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
+).to(device).eval()
+
+# ---------------------------
+# (Optional) image helpers
+# ---------------------------
+def add_random_padding(image, min_percent=0.1, max_percent=0.10):
+    image = image.convert("RGB")
+    width, height = image.size
+    pad_w_percent = random.uniform(min_percent, max_percent)
+    pad_h_percent = random.uniform(min_percent, max_percent)
+    pad_w = int(width * pad_w_percent)
+    pad_h = int(height * pad_h_percent)
+    corner_pixel = image.getpixel((0, 0))
+    padded_image = ImageOps.expand(image, border=(pad_w, pad_h, pad_w, pad_h), fill=corner_pixel)
+    return padded_image
+
+def normalize_values(text, target_max=500):
+    def normalize_list(values):
+        max_value = max(values) if values else 1
+        return [round((v / max_value) * target_max) for v in values]
+
+    def process_match(match):
+        num_list = ast.literal_eval(match.group(0))
+        normalized = normalize_list(num_list)
+        return "".join([f"<loc_{num}>" for num in normalized])
+
+    pattern = r"\[([\d\.\s,]+)\]"
+    return re.sub(pattern, process_match, text)
+
+# ---------------------------
+# Image generation only
+# ---------------------------
+@spaces.GPU
+def generate_image(
+    text: str,
+    image: Image.Image,
+    max_new_tokens: int = 2048,
+    temperature: float = 0.1,
+    top_p: float = 0.9,
+    top_k: int = 50,
+    repetition_penalty: float = 1.2,
+):
+    """Generate OCR/vision response for a single image with Typhoon OCR 20B."""
     if image is None:
-        return "Please upload an image."
+        yield "Please upload an image."
+        return
 
-    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)
+    images = [image]
 
     messages = [
-        {"role": "system", "content": "You are a helpful assistant."},
         {
             "role": "user",
-            "content": [
-                {"type": "image", "image": image},
-                {"type": "text", "text": prompt},
-            ],
-        },
+            "content": [{"type": "image"} for _ in images] + [
+                {"type": "text", "text": text}
+            ]
+        }
     ]
-
-    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,
-            repetition_penalty=1.25,
-        )
-        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
+    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = processor(text=prompt, images=images, return_tensors="pt").to(device)
+
+    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = {
+        **inputs,
+        "streamer": streamer,
+        "max_new_tokens": max_new_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "top_k": top_k,
+        "repetition_penalty": repetition_penalty,
+    }
+
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text.replace("<|im_end|>", "")
+        yield buffer
+
+# ---------------------------
+# Minimal UI (Image only)
+# ---------------------------
+css = """
+.submit-btn {
+    background-color: #2980b9 !important;
+    color: white !important;
+}
+.submit-btn:hover {
+    background-color: #3498db !important;
+}
+"""
+
+with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
+    gr.Markdown("# **Typhoon OCR 20B**")
+
+    with gr.Row():
+        with gr.Column():
+            image_query = gr.Textbox(label="Query Input", placeholder="e.g., \"OCR the image\" or task instruction…")
+            image_upload = gr.Image(type="pil", label="Image")
+            image_submit = gr.Button("Submit", elem_classes="submit-btn")
+
+            with gr.Accordion("Advanced options", open=False):
+                max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
+                temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.1)
+                top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9)
+                top_k = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50)
+                repetition_penalty = gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2)
+
+        # Right column: ONLY output (no model info, no radios)
+        with gr.Column():
+            output = gr.Textbox(label="Output", interactive=False, lines=12, scale=2)
+
+    image_submit.click(
+        fn=generate_image,
+        inputs=[image_query, image_upload, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
+        outputs=output
     )
-    return output_text[0]
-
-
-@spaces.GPU()
-def ocr_pdf_gradio(pdf_path, max_tokens=4096, progress=gr.Progress()):
-    """Process each page of a PDF through Nanonets OCR model"""
-    if pdf_path is None:
-        return "Please upload a PDF file."
-
-    # Convert PDF to images
-    progress(0, desc="Converting PDF to images...")
-    pdf_images = convert_from_path(pdf_path)
-
-    # Process each page
-    all_text = []
-    total_pages = len(pdf_images)
-
-    for i, image in enumerate(pdf_images):
-        progress(
-            (i + 1) / total_pages, desc=f"Processing page {i + 1}/{total_pages}..."
-        )
-        page_text = ocr_image_gradio(image, max_tokens)
-        all_text.append(f"--- PAGE {i + 1} ---\n{page_text}\n")
-
-    # Combine results
-    combined_text = "\n".join(all_text)
-    return combined_text
-
-
-# Create Gradio interface
-with gr.Blocks(title="Nanonets OCR Demo") as demo:
-    # Replace simple markdown with styled HTML header that includes resources
-    gr.HTML("""
-    <div class="title" style="text-align: center">
-        <h1>🔍 Nanonets OCR - Document Text Extraction</h1>
-        <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;">
-            A state-of-the-art image-to-markdown OCR model for intelligent document processing
-        </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.Tabs() as tabs:
-        # Image tab
-        with gr.TabItem("Image OCR"):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    image_input = gr.Image(
-                        label="Upload Document Image", type="pil", height=400
-                    )
-                    image_max_tokens = gr.Slider(
-                        minimum=1024,
-                        maximum=8192,
-                        value=4096,
-                        step=512,
-                        label="Max Tokens",
-                        info="Maximum number of tokens to generate",
-                    )
-                    image_extract_btn = gr.Button(
-                        "Extract Text", variant="primary", size="lg"
-                    )
-
-                with gr.Column(scale=2):
-                    image_output_text = gr.Textbox(
-                        label="Extracted Text",
-                        lines=20,
-                        show_copy_button=True,
-                        placeholder="Extracted text will appear here...",
-                    )
-
-        # PDF tab
-        with gr.TabItem("PDF OCR"):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    pdf_input = PDF(label="Upload PDF Document", height=400)
-                    pdf_max_tokens = gr.Slider(
-                        minimum=1024,
-                        maximum=8192,
-                        value=4096,
-                        step=512,
-                        label="Max Tokens per Page",
-                        info="Maximum number of tokens to generate for each page",
-                    )
-                    pdf_extract_btn = gr.Button(
-                        "Extract PDF Text", variant="primary", size="lg"
-                    )
-
-                with gr.Column(scale=2):
-                    pdf_output_text = gr.Textbox(
-                        label="Extracted Text (All Pages)",
-                        lines=20,
-                        show_copy_button=True,
-                        placeholder="Extracted text will appear here...",
-                    )
-
-    # Event handlers for Image tab
-    image_extract_btn.click(
-        fn=ocr_image_gradio,
-        inputs=[image_input, image_max_tokens],
-        outputs=image_output_text,
-        show_progress=True,
-    )
-
-    image_input.change(
-        fn=ocr_image_gradio,
-        inputs=[image_input, image_max_tokens],
-        outputs=image_output_text,
-        show_progress=True,
-    )
-
-    # Event handlers for PDF tab
-    pdf_extract_btn.click(
-        fn=ocr_pdf_gradio,
-        inputs=[pdf_input, pdf_max_tokens],
-        outputs=pdf_output_text,
-        show_progress=True,
-    )
-
-    # Add model information section
-    with gr.Accordion("About Nanonets-OCR-s", open=False):
-        gr.Markdown("""
-        ## Nanonets-OCR-s
-        
-        Nanonets-OCR-s is a powerful, state-of-the-art image-to-markdown OCR model that goes far beyond traditional text extraction. 
-        It transforms documents into structured markdown with intelligent content recognition and semantic tagging, making it ideal 
-        for downstream processing by Large Language Models (LLMs).
-        
-        ### Key Features
-        
-        - **LaTeX Equation Recognition**: Automatically converts mathematical equations and formulas into properly formatted LaTeX syntax. 
-          It distinguishes between inline ($...$) and display ($$...$$) equations.
-        
-        - **Intelligent Image Description**: Describes images within documents using structured `<img>` tags, making them digestible 
-          for LLM processing. It can describe various image types, including logos, charts, graphs and so on, detailing their content, 
-          style, and context.
-        
-        - **Signature Detection & Isolation**: Identifies and isolates signatures from other text, outputting them within a `<signature>` tag. 
-          This is crucial for processing legal and business documents.
-        
-        - **Watermark Extraction**: Detects and extracts watermark text from documents, placing it within a `<watermark>` tag.
-        
-        - **Smart Checkbox Handling**: Converts form checkboxes and radio buttons into standardized Unicode symbols (☐, ☑, ☒) 
-          for consistent and reliable processing.
-        
-        - **Complex Table Extraction**: Accurately extracts complex tables from documents and converts them into both markdown 
-          and HTML table formats.
-        """)
 
 if __name__ == "__main__":
-    demo.queue().launch(ssr_mode=False)
+    demo.queue(max_size=30).launch(share=True, mcp_server=True, ssr_mode=False, show_error=True)