Upload app.py

app.py

@@ -1,143 +1,362 @@
-import base64
-from io import BytesIO
-import json
-import os
-from openai import OpenAI
-from dotenv import load_dotenv
-from typhoon_ocr import prepare_ocr_messages
-import gradio as gr
-from PIL import Image
-
-load_dotenv()
-
-openai = OpenAI(base_url=os.getenv("TYPHOON_BASE_URL"), api_key=os.getenv("TYPHOON_API_KEY"))
-
-theme = gr.themes.Soft(
-    primary_hue=gr.themes.Color(
-        c50="#f7f7fd",
-        c100="#dfdef8",
-        c200="#c4c1f2",
-        c300="#a29eea",
-        c400="#8f8ae6",
-        c500="#756fe0",
-        c600="#635cc1",
-        c700="#4f4a9b",
-        c800="#433f83",
-        c900="#302d5e",
-        c950="#302d5e",
-    ),
-    secondary_hue="rose",
-    neutral_hue="stone",
-)
-
-def process_pdf(pdf_or_image_file, task_type, page_number):
-    if pdf_or_image_file is None:
-        return None, "No file uploaded"
-    
-    orig_filename = pdf_or_image_file.name
-    
-    try:
-        # Use the new simplified function to prepare OCR messages with page number
-        messages = prepare_ocr_messages(
-            pdf_or_image_path=orig_filename,
-            task_type=task_type,
-            target_image_dim=1800,
-            target_text_length=8000,
-            page_num=page_number if page_number else 1
-        )
-        
-        # Extract the image from the message content for display
-        image_url = messages[0]["content"][1]["image_url"]["url"]
-        image_base64 = image_url.replace("data:image/png;base64,", "")
-        image_pil = Image.open(BytesIO(base64.b64decode(image_base64)))
-        
-        # Send messages to OpenAI compatible API
-        response = openai.chat.completions.create(
-            model=os.getenv("TYPHOON_OCR_MODEL"),
-            messages=messages,
-            max_tokens=16384,
-            extra_body={
-                "repetition_penalty": 1.2,
-                "temperature": 0.1,
-                "top_p": 0.6,
-            },
-        )
-        text_output = response.choices[0].message.content
-        
-        # Try to parse the output assuming it is a Python dictionary containing 'natural_text'
-        try:
-            json_data = json.loads(text_output)
-            markdown_out = json_data.get('natural_text', "").replace("<figure>", "").replace("</figure>", "")
-        except Exception as e:
-            markdown_out = f"⚠️ Could not extract `natural_text` from output.\nError: {str(e)}"
-        
-        return image_pil, markdown_out
-    
-    except Exception as e:
-        return None, f"Error processing file: {str(e)}"
-
-
-# Build the Gradio UI.
-with gr.Blocks(theme=theme) as demo:
-    title = gr.HTML("""
-    <h1>Typhoon OCR</h1>
-    <ul>
-        <li>🤗 <b>Model weights</b>: <a href="https://huggingface.co/scb10x/typhoon-ocr-7b" target="_blank">https://huggingface.co/scb10x/typhoon-ocr-7b</a></li>
-    </ul>
-    <br />
-    <details>
-        <summary><strong>Disclaimer</strong></summary>
-        The responses generated by this Artificial Intelligence (AI) system are autonomously constructed and do not necessarily reflect the views or positions of the developing organizations, their affiliates, or any of their employees. These AI-generated responses do not represent those of the organizations. The organizations do not endorse, support, sanction, encourage, verify, or agree with the comments, opinions, or statements generated by this AI. The information produced by this AI is not intended to malign any religion, ethnic group, club, organization, company, individual, anyone, or anything. It is not the intent of the organizations to malign any group or individual. The AI operates based on its programming and training data and its responses should not be interpreted as the explicit intent or opinion of the organizations.
-    </details>
-    <br />
-    <details>
-        <summary><strong>Terms of use</strong></summary>
-        By using this service, users are required to agree to the following terms: The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. Vision language models are prone to hallucinations to a greater extent compared to text-only LLMs.
-    </details>
-    <br />
-    <details>
-        <summary><strong>License</strong></summary>
-        This project utilizes certain datasets and checkpoints that are subject to their respective original licenses. Users must comply with all terms and conditions of these original licenses. The content of this project itself is licensed under the Apache license 2.0.
-    </details>
-""")
-    with gr.Row():
-        with gr.Column(scale=1):
-            # Update file_types to accept PDF as well as common image formats.
-            pdf_input = gr.File(label="📄 Upload Image file or PDF file", file_types=[".pdf", ".png", ".jpg", ".jpeg"])
-            
-            with gr.Group(elem_classes=["task-background"]):
-                task_dropdown = gr.Radio(["default", "structure"], label="🎯 Select Task", value="default")
-                gr.HTML("""
-                <p><b>default</b>: This mode works for most cases and is recommended for files without a clear template such as infographics.</p>
-                <p><b>structure</b>: This mode offers improved performance for complex layout documents such as those containing images, tables and forms.</p>
-                <p>We recommend trying both and see which one works better for your use case.</p>
-                """, elem_classes=["task-dropdown-info"])
-                demo.css = """
-                .task-background {
-                    background: var(--block-background-fill) !important;
-                    
-                }
-                .task-background > * {
-                    background: var(--block-background-fill) !important;
-                }
-                .task-dropdown-info {
-                    padding: 0 16px;
-                    font-size: 12px;
-                }
-                """
-            page_number = gr.Number(label="📄 Page Number (for PDFs only)", value=1, minimum=1, step=1)
-            run_button = gr.Button("🚀 Run")
-            image_output = gr.Image(label="📸 Preview Image", type="pil")
-        with gr.Column(scale=2):
-            markdown_output = gr.Markdown(label='Markdown Result', show_label=True)
-
-    
-    # Connect the UI inputs to the processing function.
-    run_button.click(
-        fn=process_pdf,
-        inputs=[pdf_input, task_dropdown, page_number],
-        outputs=[image_output, markdown_output]
-    )
-
-# Launch the Gradio demo (temporary public share for 72 hours)
-demo.launch(share=False)
+import os
+import random
+import uuid
+import json
+import time
+import asyncio
+from threading import Thread
+
+import gradio as gr
+import spaces
+import torch
+import numpy as np
+from PIL import Image, ImageOps
+import cv2
+
+from transformers import (
+    Qwen2VLForConditionalGeneration,
+    Qwen2_5_VLForConditionalGeneration,
+    AutoModelForVision2Seq,
+    AutoProcessor,
+    TextIteratorStreamer,
+)
+from transformers.image_utils import load_image
+
+from docling_core.types.doc import DoclingDocument, DocTagsDocument
+
+import re
+import ast
+import html
+
+# Constants for text generation
+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 Adun/typhoon_ocr-7B-v1.4
+MODEL_ID_M = "Adun/typhoon_ocr-7B-v1.4"
+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()
+
+
+# Load typhoon-ocr-7b
+MODEL_ID_L = "scb10x/typhoon-ocr-7b"
+processor_l = AutoProcessor.from_pretrained(MODEL_ID_L, trust_remote_code=True)
+model_l = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    MODEL_ID_L,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
+).to(device).eval()
+
+
+# Preprocessing functions for SmolDocling-256M
+def add_random_padding(image, min_percent=0.1, max_percent=0.10):
+    """Add random padding to an image based on its size."""
+    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))  # Top-left corner
+    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):
+    """Normalize numerical values in text to a target maximum."""
+    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,]+)\]"
+    normalized_text = re.sub(pattern, process_match, text)
+    return normalized_text
+
+def downsample_video(video_path):
+    """Downsample a video to evenly spaced frames, returning PIL images with timestamps."""
+    vidcap = cv2.VideoCapture(video_path)
+    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = vidcap.get(cv2.CAP_PROP_FPS)
+    frames = []
+    frame_indices = np.linspace(0, total_frames - 1, 10, dtype=int)
+    for i in frame_indices:
+        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
+        success, image = vidcap.read()
+        if success:
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            pil_image = Image.fromarray(image)
+            timestamp = round(i / fps, 2)
+            frames.append((pil_image, timestamp))
+    vidcap.release()
+    return frames
+
+@spaces.GPU
+def generate_image(model_name: str, text: str, image: Image.Image,
+                   max_new_tokens: int = 1024,
+                   temperature: float = 0.6,
+                   top_p: float = 0.9,
+                   top_k: int = 50,
+                   repetition_penalty: float = 1.2):
+    """Generate responses for image input using the selected model."""
+    # Model selection
+    if model_name == "Adun/Typhoon-OCR-7B-1.4":
+        processor = processor_m
+        model = model_m
+    # elif model_name == "MonkeyOCR-Recognition":
+    #     processor = processor_g
+    #     model = model_g
+    # elif model_name == "SmolDocling-256M-preview":
+    #     processor = processor_x
+    #     model = model_x
+    elif model_name == "Typhoon-OCR-7B":
+        processor = processor_l
+        model = model_l
+    else:
+        yield "Invalid model selected."
+        return
+
+    if image is None:
+        yield "Please upload an image."
+        return
+
+    # Prepare images as a list (single image for image inference)
+    images = [image]
+
+    # # SmolDocling-256M specific preprocessing
+    # if model_name == "SmolDocling-256M-preview":
+    #     if "OTSL" in text or "code" in text:
+    #         images = [add_random_padding(img) for img in images]
+    #     if "OCR at text at" in text or "Identify element" in text or "formula" in text:
+    #         text = normalize_values(text, target_max=500)
+
+    # Unified message structure for all models
+    messages = [
+        {
+            "role": "user",
+            "content": [{"type": "image"} for _ in images] + [
+                {"type": "text", "text": text}
+            ]
+        }
+    ]
+    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = processor(text=prompt, images=images, return_tensors="pt").to(device)
+
+    # Generation with streaming
+    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()
+
+    # Stream output and collect full response
+    buffer = ""
+    full_output = ""
+    for new_text in streamer:
+        full_output += new_text
+        buffer += new_text.replace("<|im_end|>", "")
+        yield buffer
+
+    # SmolDocling-256M specific postprocessing
+    # if model_name == "SmolDocling-256M-preview":
+    #     cleaned_output = full_output.replace("<end_of_utterance>", "").strip()
+    #     if any(tag in cleaned_output for tag in ["<doctag>", "<otsl>", "<code>", "<chart>", "<formula>"]):
+    #         if "<chart>" in cleaned_output:
+    #             cleaned_output = cleaned_output.replace("<chart>", "<otsl>").replace("</chart>", "</otsl>")
+    #             cleaned_output = re.sub(r'(<loc_500>)(?!.*<loc_500>)<[^>]+>', r'\1', cleaned_output)
+    #         doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([cleaned_output], images)
+    #         doc = DoclingDocument.load_from_doctags(doctags_doc, document_name="Document")
+    #         markdown_output = doc.export_to_markdown()
+    #         yield f"**MD Output:**\n\n{markdown_output}"
+    #     else:
+    #         yield cleaned_output
+
+@spaces.GPU
+def generate_video(model_name: str, text: str, video_path: str,
+                   max_new_tokens: int = 1024,
+                   temperature: float = 0.6,
+                   top_p: float = 0.9,
+                   top_k: int = 50,
+                   repetition_penalty: float = 1.2):
+    """Generate responses for video input using the selected model."""
+    # Model selection
+    if model_name == "Adun/typhoon_ocr-7B-v1.4":
+        processor = processor_m
+        model = model_m
+    # elif model_name == "MonkeyOCR-Recognition":
+    #     processor = processor_g
+    #     model = model_g
+    # elif model_name == "SmolDocling-256M-preview":
+    #     processor = processor_x
+    #     model = model_x
+    elif model_name == "Typhoon-OCR-7B":
+        processor = processor_l
+        model = model_l
+    else:
+        yield "Invalid model selected."
+        return
+
+    if video_path is None:
+        yield "Please upload a video."
+        return
+
+    # Extract frames from video
+    frames = downsample_video(video_path)
+    images = [frame for frame, _ in frames]
+
+    # # SmolDocling-256M specific preprocessing
+    # if model_name == "SmolDocling-256M-preview":
+    #     if "OTSL" in text or "code" in text:
+    #         images = [add_random_padding(img) for img in images]
+    #     if "OCR at text at" in text or "Identify element" in text or "formula" in text:
+    #         text = normalize_values(text, target_max=500)
+
+    # Unified message structure for all models
+    messages = [
+        {
+            "role": "user",
+            "content": [{"type": "image"} for _ in images] + [
+                {"type": "text", "text": text}
+            ]
+        }
+    ]
+    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = processor(text=prompt, images=images, return_tensors="pt").to(device)
+
+    # Generation with streaming
+    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()
+
+    # Stream output and collect full response
+    buffer = ""
+    full_output = ""
+    for new_text in streamer:
+        full_output += new_text
+        buffer += new_text.replace("<|im_end|>", "")
+        yield buffer
+
+    # # SmolDocling-256M specific postprocessing
+    # if model_name == "SmolDocling-256M-preview":
+    #     cleaned_output = full_output.replace("<end_of_utterance>", "").strip()
+    #     if any(tag in cleaned_output for tag in ["<doctag>", "<otsl>", "<code>", "<chart>", "<formula>"]):
+    #         if "<chart>" in cleaned_output:
+    #             cleaned_output = cleaned_output.replace("<chart>", "<otsl>").replace("</chart>", "</otsl>")
+    #             cleaned_output = re.sub(r'(<loc_500>)(?!.*<loc_500>)<[^>]+>', r'\1', cleaned_output)
+    #         doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([cleaned_output], images)
+    #         doc = DoclingDocument.load_from_doctags(doctags_doc, document_name="Document")
+    #         markdown_output = doc.export_to_markdown()
+    #         yield f"**MD Output:**\n\n{markdown_output}"
+    #     else:
+    #         yield cleaned_output
+
+# Define examples for image and video inference
+image_examples = [
+    ["OCR the image", "images/2.jpg"],
+    ["Convert this page to docling", "images/1.png"],
+    ["Convert this page to docling", "images/3.png"],
+    ["Convert chart to OTSL.", "images/4.png"],
+    ["Convert code to text", "images/5.jpg"],
+    ["Convert this table to OTSL.", "images/6.jpg"],
+    ["Convert formula to late.", "images/7.jpg"],
+]
+
+video_examples = [
+    ["Explain the video in detail.", "videos/1.mp4"],
+    ["Explain the video in detail.", "videos/2.mp4"]
+]
+
+css = """
+.submit-btn {
+    background-color: #2980b9 !important;
+    color: white !important;
+}
+.submit-btn:hover {
+    background-color: #3498db !important;
+}
+"""
+
+# Create the Gradio Interface
+with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
+    gr.Markdown("# **[Multimodal OCR](https://huggingface.co/collections/prithivMLmods/multimodal-implementations-67c9982ea04b39f0608badb0)**")
+    with gr.Row():
+        with gr.Column():
+            with gr.Tabs():
+                with gr.TabItem("Image Inference"):
+                    image_query = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
+                    image_upload = gr.Image(type="pil", label="Image")
+                    image_submit = gr.Button("Submit", elem_classes="submit-btn")
+                    gr.Examples(
+                        examples=image_examples,
+                        inputs=[image_query, image_upload]
+                    )
+                with gr.TabItem("Video Inference"):
+                    video_query = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
+                    video_upload = gr.Video(label="Video")
+                    video_submit = gr.Button("Submit", elem_classes="submit-btn")
+                    gr.Examples(
+                        examples=video_examples,
+                        inputs=[video_query, video_upload]
+                    )
+            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.6)
+                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)
+        with gr.Column():
+            output = gr.Textbox(label="Output", interactive=False, lines=3, scale=2)
+            model_choice = gr.Radio(
+                choices=["Adun/typhoon_ocr-7B-v1.4", "Typhoon-OCR-7B"],
+                label="Select Model",
+                value="Nanonets-OCR-s"
+            )
+            
+           #gr.Markdown("**Model Info 💻** | [Report Bug](https://huggingface.co/spaces/prithivMLmods/Multimodal-OCR2/discussions)")
+
+            # gr.Markdown("> [SmolDocling-256M](https://huggingface.co/ds4sd/SmolDocling-256M-preview): SmolDocling is a multimodal Image-Text-to-Text model designed for efficient document conversion. It retains Docling's most popular features while ensuring full compatibility with Docling through seamless support for DoclingDocuments.")
+            # gr.Markdown("> [Nanonets-OCR-s](https://huggingface.co/nanonets/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.")
+            # gr.Markdown("> [MonkeyOCR-Recognition](https://huggingface.co/echo840/MonkeyOCR): MonkeyOCR adopts a Structure-Recognition-Relation (SRR) triplet paradigm, which simplifies the multi-tool pipeline of modular approaches while avoiding the inefficiency of using large multimodal models for full-page document processing.")
+            gr.Markdown("> [Typhoon-OCR-7B-1.4 finetuned by Aun](https://huggingface.co/scb10x/typhoon-ocr-7b): A bilingual document parsing model built specifically for real-world documents in Thai and English inspired by models like olmOCR based on Qwen2.5-VL-Instruction. Extracts and interprets embedded text (e.g., chart labels, captions) in Thai or English.")
+            gr.Markdown("> [Typhoon-OCR-7B](https://huggingface.co/scb10x/typhoon-ocr-7b): A bilingual document parsing model built specifically for real-world documents in Thai and English inspired by models like olmOCR based on Qwen2.5-VL-Instruction. Extracts and interprets embedded text (e.g., chart labels, captions) in Thai or English.")
+            gr.Markdown(">⚠️note: all the models in space are not guaranteed to perform well in video inference use cases.")  
+            
+    image_submit.click(
+        fn=generate_image,
+        inputs=[model_choice, image_query, image_upload, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
+        outputs=output
+    )
+    video_submit.click(
+        fn=generate_video,
+        inputs=[model_choice, video_query, video_upload, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
+        outputs=output
+    )
+
+if __name__ == "__main__":
+    demo.queue(max_size=30).launch(share=True, mcp_server=True, ssr_mode=False, show_error=True)