Update app.py

app.py

@@ -1,315 +1,458 @@
-import os
-import time
-import threading
-import gradio as gr
 import spaces
-import torch
-import numpy as np
-from PIL import Image
-import cv2
-from transformers import (
-    Qwen2_5_VLForConditionalGeneration,
-    Qwen2VLForConditionalGeneration,
-    Glm4vForConditionalGeneration,
-    AutoProcessor,
-    TextIteratorStreamer,
-)
-from qwen_vl_utils import process_vision_info
-
-# Constants for text generation
-MAX_MAX_NEW_TOKENS = 4096
-DEFAULT_MAX_NEW_TOKENS = 3584
-MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+import json
+import math
+import os
+import traceback
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Tuple
+import re
 
-# Load Camel-Doc-OCR-062825
-MODEL_ID_M = "prithivMLmods/Camel-Doc-OCR-062825"
-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()
+import fitz  # PyMuPDF
+import gradio as gr
+import requests
+from PIL import Image, ImageDraw, ImageFont
 
-# Load Qwen2.5-VL-3B-Instruct-abliterated
-MODEL_ID_X = "huihui-ai/Qwen2.5-VL-3B-Instruct-abliterated"
-processor_x = AutoProcessor.from_pretrained(MODEL_ID_X, trust_remote_code=True)
-model_x = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_ID_X,
-    trust_remote_code=True,
-    torch_dtype=torch.float16
-).to(device).eval()
+from model import load_model, inference_dots_ocr, inference_dolphin
 
-# Load Megalodon-OCR-Sync-0713
-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()
+# Constants
+MIN_PIXELS = 3136
+MAX_PIXELS = 11289600
+IMAGE_FACTOR = 28
 
-# Load GLM-4.1V-9B-Thinking
-MODEL_ID_S = "zai-org/GLM-4.1V-9B-Thinking"
-processor_s = AutoProcessor.from_pretrained(MODEL_ID_S, trust_remote_code=True)
-model_s = Glm4vForConditionalGeneration.from_pretrained(
-    MODEL_ID_S,
-    trust_remote_code=True,
-    torch_dtype=torch.float16
-).to(device).eval()
+# Prompts
+prompt = """Please output the layout information from the PDF image, including each layout element's bbox, its category, and the corresponding text content within the bbox.
 
-# Load DeepEyes-7B
-MODEL_ID_Y = "ChenShawn/DeepEyes-7B"
-processor_y = AutoProcessor.from_pretrained(MODEL_ID_Y, trust_remote_code=True)
-model_y = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_ID_Y,
-    trust_remote_code=True,
-    torch_dtype=torch.float16
-).to(device).eval()
+1. Bbox format: [x1, y1, x2, y2]
+2. Layout Categories: ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title']
+3. Text Extraction & Formatting Rules:
+    - Picture: Omit the text field
+    - Formula: format as LaTeX
+    - Table: format as HTML
+    - Others: format as Markdown
+4. Constraints:
+    - Use original text, no translation
+    - Sort elements by human reading order
+5. Final Output: Single JSON object
+"""
 
-def downsample_video(video_path):
-    """
-    Downsample a video to evenly spaced frames, returning each as a PIL image with its timestamp.
-    """
-    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
+# Load models at startup
+models = {
+    "dots.ocr": load_model("dots.ocr"),
+    "Dolphin": load_model("Dolphin")
+}
 
-@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 using the selected model for image input.
-    """
-    if model_name == "Camel-Doc-OCR-062825":
-        processor = processor_m
-        model = model_m
-    elif model_name == "Megalodon-OCR-Sync-0713":
-        processor = processor_t
-        model = model_t
-    elif model_name == "GLM-4.1V-9B-Thinking":
-        processor = processor_s
-        model = model_s
-    elif model_name == "DeepEyes-7B-Thinking":
-        processor = processor_y
-        model = model_y
-    elif model_name == "Qwen2.5-VL-3B-Instruct-abliterated":
-        processor = processor_x
-        model = model_x
-    else:
-        yield "Invalid model selected.", "Invalid model selected."
-        return
+# Global state for PDF handling
+pdf_cache = {
+    "images": [],
+    "current_page": 0,
+    "total_pages": 0,
+    "file_type": None,
+    "is_parsed": False,
+    "results": []
+}
 
-    if image is None:
-        yield "Please upload an image.", "Please upload an image."
-        return
+# Utility functions
+def round_by_factor(number: int, factor: int) -> int:
+    return round(number / factor) * factor
 
-    messages = [{
-        "role": "user",
-        "content": [
-            {"type": "image", "image": image},
-            {"type": "text", "text": text},
-        ]
-    }]
-    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=False,
-        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 = threading.Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        time.sleep(0.01)
-        yield buffer, buffer
+def smart_resize(height: int, width: int, factor: int = 28, min_pixels: int = 3136, max_pixels: int = 11289600):
+    if max(height, width) / min(height, width) > 200:
+        raise ValueError(f"Aspect ratio must be < 200, got {max(height, width) / min(height, width)}")
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = round_by_factor(height / beta, factor)
+        w_bar = round_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = round_by_factor(height * beta, factor)
+        w_bar = round_by_factor(width * beta, factor)
+    return h_bar, w_bar
 
-@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 using the selected model for video input.
-    """
-    if model_name == "Camel-Doc-OCR-062825":
-        processor = processor_m
-        model = model_m
-    elif model_name == "Megalodon-OCR-Sync-0713":
-        processor = processor_t
-        model = model_t
-    elif model_name == "GLM-4.1V-9B-Thinking":
-        processor = processor_s
-        model = model_s
-    elif model_name == "DeepEyes-7B-Thinking":
-        processor = processor_y
-        model = model_y
-    elif model_name == "Qwen2.5-VL-3B-Instruct-abliterated":
-        processor = processor_x
-        model = model_x
+def fetch_image(image_input, min_pixels: int = None, max_pixels: int = None):
+    if isinstance(image_input, str):
+        if image_input.startswith(("http://", "https://")):
+            response = requests.get(image_input)
+            image = Image.open(BytesIO(response.content)).convert('RGB')
+        else:
+            image = Image.open(image_input).convert('RGB')
+    elif isinstance(image_input, Image.Image):
+        image = image_input.convert('RGB')
     else:
-        yield "Invalid model selected.", "Invalid model selected."
-        return
+        raise ValueError(f"Invalid image input type: {type(image_input)}")
+    if min_pixels or max_pixels:
+        min_pixels = min_pixels or MIN_PIXELS
+        max_pixels = max_pixels or MAX_PIXELS
+        height, width = smart_resize(image.height, image.width, factor=IMAGE_FACTOR, min_pixels=min_pixels, max_pixels=max_pixels)
+        image = image.resize((width, height), Image.LANCZOS)
+    return image
 
-    if video_path is None:
-        yield "Please upload a video.", "Please upload a video."
-        return
+def load_images_from_pdf(pdf_path: str) -> List[Image.Image]:
+    images = []
+    try:
+        pdf_document = fitz.open(pdf_path)
+        for page_num in range(len(pdf_document)):
+            page = pdf_document.load_page(page_num)
+            mat = fitz.Matrix(2.0, 2.0)
+            pix = page.get_pixmap(matrix=mat)
+            img_data = pix.tobytes("ppm")
+            image = Image.open(BytesIO(img_data)).convert('RGB')
+            images.append(image)
+        pdf_document.close()
+    except Exception as e:
+        print(f"Error loading PDF: {e}")
+        return []
+    return images
 
-    frames = downsample_video(video_path)
-    messages = [
-        {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]},
-        {"role": "user", "content": [{"type": "text", "text": text}]}
-    ]
-    for frame in frames:
-        image, timestamp = frame
-        messages[1]["content"].append({"type": "text", "text": f"Frame {timestamp}:"})
-        messages[1]["content"].append({"type": "image", "image": image})
-    inputs = processor.apply_chat_template(
-        messages,
-        tokenize=True,
-        add_generation_prompt=True,
-        return_dict=True,
-        return_tensors="pt",
-        truncation=False,
-        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,
-        "do_sample": True,
-        "temperature": temperature,
-        "top_p": top_p,
-        "top_k": top_k,
-        "repetition_penalty": repetition_penalty,
+def draw_layout_on_image(image: Image.Image, layout_data: List[Dict]) -> Image.Image:
+    img_copy = image.copy()
+    draw = ImageDraw.Draw(img_copy)
+    colors = {
+        'Caption': '#FF6B6B', 'Footnote': '#4ECDC4', 'Formula': '#45B7D1', 'List-item': '#96CEB4',
+        'Page-footer': '#FFEAA7', 'Page-header': '#DDA0DD', 'Picture': '#FFD93D', 'Section-header': '#6C5CE7',
+        'Table': '#FD79A8', 'Text': '#74B9FF', 'Title': '#E17055'
     }
-    thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        buffer = buffer.replace("<|im_end|>", "")
-        time.sleep(0.01)
-        yield buffer, buffer
+    try:
+        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 12)
+    except Exception:
+        font = ImageFont.load_default()
+    try:
+        for item in layout_data:
+            if 'bbox' in item and 'category' in item:
+                bbox = item['bbox']
+                category = item['category']
+                color = colors.get(category, '#000000')
+                draw.rectangle(bbox, outline=color, width=2)
+                label = category
+                label_bbox = draw.textbbox((0, 0), label, font=font)
+                label_width = label_bbox[2] - label_bbox[0]
+                label_height = label_bbox[3] - label_bbox[1]
+                label_x = bbox[0]
+                label_y = max(0, bbox[1] - label_height - 2)
+                draw.rectangle([label_x, label_y, label_x + label_width + 4, label_y + label_height + 2], fill=color)
+                draw.text((label_x + 2, label_y + 1), label, fill='white', font=font)
+    except Exception as e:
+        print(f"Error drawing layout: {e}")
+    return img_copy
 
-# Define examples for image and video inference
-image_examples = [
-    ["explain the movie shot in detail.", "images/5.jpg"],
-    ["convert this page to doc [text] precisely for markdown.", "images/1.png"],
-    ["convert this page to doc [table] precisely for markdown.", "images/2.png"],
-    ["explain the movie shot in detail.", "images/3.png"],   
-    ["fill the correct numbers.", "images/4.png"]
-]
+def is_arabic_text(text: str) -> bool:
+    if not text:
+        return False
+    header_pattern = r'^#{1,6}\s+(.+)$'
+    paragraph_pattern = r'^(?!#{1,6}\s|!\[|```|\||\s*[-*+]\s|\s*\d+\.\s)(.+)$'
+    content_text = []
+    for line in text.split('\n'):
+        line = line.strip()
+        if not line:
+            continue
+        header_match = re.match(header_pattern, line, re.MULTILINE)
+        if header_match:
+            content_text.append(header_match.group(1))
+            continue
+        if re.match(paragraph_pattern, line, re.MULTILINE):
+            content_text.append(line)
+    if not content_text:
+        return False
+    combined_text = ' '.join(content_text)
+    arabic_chars = 0
+    total_chars = 0
+    for char in combined_text:
+        if char.isalpha():
+            total_chars += 1
+            if ('\u0600' <= char <= '\u06FF') or ('\u0750' <= char <= '\u077F') or ('\u08A0' <= char <= '\u08FF'):
+                arabic_chars += 1
+    return total_chars > 0 and (arabic_chars / total_chars) > 0.5
 
-video_examples = [
-    ["explain the video in detail.", "videos/b.mp4"],
-    ["explain the ad video in detail.", "videos/a.mp4"]
-]
+def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = 'text') -> str:
+    import base64
+    markdown_lines = []
+    try:
+        sorted_items = sorted(layout_data, key=lambda x: (x.get('bbox', [0, 0, 0, 0])[1], x.get('bbox', [0, 0, 0, 0])[0]))
+        for item in sorted_items:
+            category = item.get('category', '')
+            text = item.get(text_key, '')
+            bbox = item.get('bbox', [])
+            if category == 'Picture':
+                if bbox and len(bbox) == 4:
+                    try:
+                        x1, y1, x2, y2 = [max(0, int(x)) if i < 2 else min(image.width if i % 2 == 0 else image.height, int(x)) for i, x in enumerate(bbox)]
+                        if x2 > x1 and y2 > y1:
+                            cropped_img = image.crop((x1, y1, x2, y2))
+                            buffer = BytesIO()
+                            cropped_img.save(buffer, format='PNG')
+                            img_data = base64.b64encode(buffer.getvalue()).decode()
+                            markdown_lines.append(f"<image-card alt="Image" src="data:image/png;base64,{img_data}" ></image-card>\n")
+                        else:
+                            markdown_lines.append("<image-card alt="Image" src="Image region detected" ></image-card>\n")
+                    except Exception as e:
+                        print(f"Error processing image region: {e}")
+                        markdown_lines.append("<image-card alt="Image" src="Image detected" ></image-card>\n")
+                else:
+                    markdown_lines.append("<image-card alt="Image" src="Image detected" ></image-card>\n")
+            elif not text:
+                continue
+            elif category == 'Title':
+                markdown_lines.append(f"# {text}\n")
+            elif category == 'Section-header':
+                markdown_lines.append(f"## {text}\n")
+            elif category == 'Text':
+                markdown_lines.append(f"{text}\n")
+            elif category == 'List-item':
+                markdown_lines.append(f"- {text}\n")
+            elif category == 'Table':
+                if text.strip().startswith('<'):
+                    markdown_lines.append(f"{text}\n")
+                else:
+                    markdown_lines.append(f"**Table:** {text}\n")
+            elif category == 'Formula':
+                if text.strip().startswith('$') or '\\' in text:
+                    markdown_lines.append(f"$$  \n{text}\n  $$\n")
+                else:
+                    markdown_lines.append(f"**Formula:** {text}\n")
+            elif category == 'Caption':
+                markdown_lines.append(f"*{text}*\n")
+            elif category == 'Footnote':
+                markdown_lines.append(f"^{text}^\n")
+            elif category in ['Page-header', 'Page-footer']:
+                continue
+            else:
+                markdown_lines.append(f"{text}\n")
+            markdown_lines.append("")
+    except Exception as e:
+        print(f"Error converting to markdown: {e}")
+        return str(layout_data)
+    return "\n".join(markdown_lines)
 
-# Updated CSS with model choice highlighting
-css = """
-.submit-btn {
-    background-color: #2980b9 !important;
-    color: white !important;
-}
-.submit-btn:hover {
-    background-color: #3498db !important;
-}
-.canvas-output {
-    border: 2px solid #4682B4;
-    border-radius: 10px;
-    padding: 20px;
-}
-"""
-
-# Create the Gradio Interface
-with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
-    gr.Markdown("# **[Multimodal VLM 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]
-                    )
+def load_file_for_preview(file_path: str) -> Tuple[Optional[Image.Image], str]:
+    global pdf_cache
+    if not file_path or not os.path.exists(file_path):
+        return None, "No file selected"
+    file_ext = os.path.splitext(file_path)[1].lower()
+    try:
+        if file_ext == '.pdf':
+            images = load_images_from_pdf(file_path)
+            if not images:
+                return None, "Failed to load PDF"
+            pdf_cache.update({
+                "images": images,
+                "current_page": 0,
+                "total_pages": len(images),
+                "file_type": "pdf",
+                "is_parsed": False,
+                "results": []
+            })
+            return images[0], f"Page 1 / {len(images)}"
+        elif file_ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
+            image = Image.open(file_path).convert('RGB')
+            pdf_cache.update({
+                "images": [image],
+                "current_page": 0,
+                "total_pages": 1,
+                "file_type": "image",
+                "is_parsed": False,
+                "results": []
+            })
+            return image, "Page 1 / 1"
+        else:
+            return None, f"Unsupported file format: {file_ext}"
+    except Exception as e:
+        print(f"Error loading file: {e}")
+        return None, f"Error loading file: {str(e)}"
 
-            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)
+@spaces.GPU()
+def process_document(file_path, model_choice, max_tokens, min_pix, max_pix):
+    global pdf_cache
+    if not file_path:
+        return None, "Please upload a file first.", None
+    model, processor = models[model_choice]
+    image, page_info = load_file_for_preview(file_path)
+    if image is None:
+        return None, page_info, None
+    if pdf_cache["file_type"] == "pdf":
+        all_results = []
+        for i, img in enumerate(pdf_cache["images"]):
+            if model_choice == "dots.ocr":
+                raw_output = inference_dots_ocr(model, processor, img, prompt, max_tokens)
+                try:
+                    layout_data = json.loads(raw_output)
+                    processed_image = draw_layout_on_image(img, layout_data)
+                    markdown_content = layoutjson2md(img, layout_data)
+                    result = {
+                        'processed_image': processed_image,
+                        'markdown_content': markdown_content,
+                        'layout_result': layout_data
+                    }
+                except Exception:
+                    result = {
+                        'processed_image': img,
+                        'markdown_content': raw_output,
+                        'layout_result': None
+                    }
+            else:  # Dolphin
+                text = inference_dolphin(model, processor, img)
+                result = f"## Page {i+1}\n\n{text}" if text else "No text extracted"
+            all_results.append(result)
+        pdf_cache["results"] = all_results
+        pdf_cache["is_parsed"] = True
+        first_result = all_results[0]
+        if model_choice == "dots.ocr":
+            markdown_update = gr.update(value=first_result['markdown_content'], rtl=is_arabic_text(first_result['markdown_content']))
+            return first_result['processed_image'], markdown_update, first_result['layout_result']
+        else:
+            markdown_update = gr.update(value=first_result, rtl=is_arabic_text(first_result))
+            return None, markdown_update, None
+    else:
+        if model_choice == "dots.ocr":
+            raw_output = inference_dots_ocr(model, processor, image, prompt, max_tokens)
+            try:
+                layout_data = json.loads(raw_output)
+                processed_image = draw_layout_on_image(image, layout_data)
+                markdown_content = layoutjson2md(image, layout_data)
+                result = {
+                    'processed_image': processed_image,
+                    'markdown_content': markdown_content,
+                    'layout_result': layout_data
+                }
+            except Exception:
+                result = {
+                    'processed_image': image,
+                    'markdown_content': raw_output,
+                    'layout_result': None
+                }
+            pdf_cache["results"] = [result]
+        else:  # Dolphin
+            text = inference_dolphin(model, processor, image)
+            result = text if text else "No text extracted"
+            pdf_cache["results"] = [result]
+        pdf_cache["is_parsed"] = True
+        if model_choice == "dots.ocr":
+            markdown_update = gr.update(value=result['markdown_content'], rtl=is_arabic_text(result['markdown_content']))
+            return result['processed_image'], markdown_update, result['layout_result']
+        else:
+            markdown_update = gr.update(value=result, rtl=is_arabic_text(result))
+            return None, markdown_update, None
 
-        with gr.Column():
-            with gr.Column(elem_classes="canvas-output"):
-                gr.Markdown("## Output")
-                output = gr.Textbox(label="Raw Output Stream", interactive=False, lines=2)
-                with gr.Accordion("(Result.md)", open=False):                
-                    markdown_output = gr.Markdown(label="(Result.md)")
+def turn_page(direction: str) -> Tuple[Optional[Image.Image], str, Any, Optional[Image.Image], Optional[Dict]]:
+    global pdf_cache
+    if not pdf_cache["images"]:
+        return None, '<div class="page-info">No file loaded</div>', "No results yet", None, None
+    if direction == "prev":
+        pdf_cache["current_page"] = max(0, pdf_cache["current_page"] - 1)
+    elif direction == "next":
+        pdf_cache["current_page"] = min(pdf_cache["total_pages"] - 1, pdf_cache["current_page"] + 1)
+    index = pdf_cache["current_page"]
+    current_image_preview = pdf_cache["images"][index]
+    page_info_html = f'<div class="page-info">Page {index + 1} / {pdf_cache["total_pages"]}</div>'
+    if pdf_cache["is_parsed"] and index < len(pdf_cache["results"]):
+        result = pdf_cache["results"][index]
+        if isinstance(result, dict):  # dots.ocr
+            markdown_content = result.get('markdown_content', 'No content available')
+            processed_img = result.get('processed_image', None)
+            layout_json = result.get('layout_result', None)
+        else:  # Dolphin
+            markdown_content = result
+            processed_img = None
+            layout_json = None
+    else:
+        markdown_content = "Page not processed yet"
+        processed_img = None
+        layout_json = None
+    markdown_update = gr.update(value=markdown_content, rtl=is_arabic_text(markdown_content))
+    return current_image_preview, page_info_html, markdown_update, processed_img, layout_json
 
-            model_choice = gr.Radio(
-                choices=["Camel-Doc-OCR-062825", "GLM-4.1V-9B-Thinking", "Megalodon-OCR-Sync-0713", "DeepEyes-7B-Thinking", "Qwen2.5-VL-3B-Instruct-abliterated"],
-                label="Select Model",
-                value="Camel-Doc-OCR-062825"
-            )
-            gr.Markdown("**Model Info 💻** | [Report Bug](https://huggingface.co/spaces/prithivMLmods/Multimodal-OCR-Comparator/discussions)")
-            gr.Markdown("> Camel-Doc-OCR-062825 and Megalodon-OCR-Sync-0713 are both fine-tuned versions of the Qwen2.5-VL series focused on document retrieval, content extraction, analysis recognition, and excelling in OCR and visual document analysis tasks for structured and unstructured content—Camel-Doc-OCR-062825 leveraging the Qwen2.5-VL-7B-Instruct as its base, while Megalodon-OCR-Sync-0713 uses Qwen2.5-VL-3B-Instruct and is especially trained on diverse captioning datasets.")
-            gr.Markdown("> GLM-4.1V-9B-Thinking is a vision-language model (VLM) based on the GLM-4-9B-0414 foundation, with a strong emphasis on advanced reasoning capabilities, chain-of-thought inference, and robust bilingual (Chinese/English) performance on complex multimodal benchmarks.")
-            gr.Markdown("> DeepEyes-7B stands out for its agentic reinforcement learning approach, focusing on thinking with images for better visual reasoning, math problem-solving, and mitigating hallucination using Qwen2.5-VL-7B-Instruct as its foundation. Finally, Qwen2.5-VL-3B-Instruct-abliterated is part of the Qwen2.5-VL family, known for its versatile vision-language understanding and generation, serving as the foundational architecture for several of these fine-tuned vision-language and OCR models.")
-    
-    # Define the submit button actions
-    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, markdown_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, markdown_output])
+def create_gradio_interface():
+    css = """
+    .main-container { max-width: 1400px; margin: 0 auto; }
+    .header-text { text-align: center; color: #2c3e50; margin-bottom: 20px; }
+    .process-button { border: none !important; color: white !important; font-weight: bold !important; }
+    .process-button:hover { transform: translateY(-2px) !important; box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important; }
+    .info-box { border: 1px solid #dee2e6; border-radius: 8px; padding: 15px; margin: 10px 0; }
+    .page-info { text-align: center; padding: 8px 16px; border-radius: 20px; font-weight: bold; margin: 10px 0; }
+    .model-status { padding: 10px; border-radius: 8px; margin: 10px 0; text-align: center; font-weight: bold; }
+    .status-ready { background: #d1edff; color: #0c5460; border: 1px solid #b8daff; }
+    """
+    with gr.Blocks(theme=gr.themes.Soft(), css=css, title="Dots.OCR Demo") as demo:
+        gr.HTML("""
+        <div class="title" style="text-align: center">
+            <h1>🔍 Dot-OCR - Multilingual Document Text Extraction</h1>
+            <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;">
+                A state-of-the-art image/pdf-to-markdown vision language model for intelligent document processing
+            </p>
+            <div style="display: flex; justify-content: center; gap: 20px; margin: 15px 0;">
+                <a href="https://huggingface.co/rednote-hilab/dots.ocr" target="_blank" style="text-decoration: none; color: #2563eb; font-weight: 500;">
+                    📚 Hugging Face Model
+                </a>
+                <a href="https://github.com/rednote-hilab/dots.ocr/blob/master/assets/blog.md" target="_blank" style="text-decoration: none; color: #2563eb; font-weight: 500;">
+                    📝 Release Blog
+                </a>
+                <a href="https://github.com/rednote-hilab/dots.ocr" 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.File(
+                    label="Upload Image or PDF",
+                    file_types=[".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".pdf"],
+                    type="filepath"
+                )
+                image_preview = gr.Image(label="Preview", type="pil", interactive=False, height=300)
+                with gr.Row():
+                    prev_page_btn = gr.Button("◀ Previous", size="md")
+                    page_info = gr.HTML('<div class="page-info">No file loaded</div>')
+                    next_page_btn = gr.Button("Next ▶", size="md")
+                model_choice = gr.Radio(
+                    choices=["dots.ocr", "Dolphin"],
+                    label="Select Model",
+                    value="dots.ocr"
+                )
+                with gr.Accordion("Advanced Settings", open=False):
+                    max_new_tokens = gr.Slider(minimum=1000, maximum=32000, value=24000, step=1000, label="Max New Tokens")
+                    min_pixels = gr.Number(value=MIN_PIXELS, label="Min Pixels")
+                    max_pixels = gr.Number(value=MAX_PIXELS, label="Max Pixels")
+                process_btn = gr.Button("🚀 Process Document", variant="primary", elem_classes=["process-button"], size="lg")
+                clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
+            with gr.Column(scale=2):
+                with gr.Tabs():
+                    with gr.Tab("🖼️ Processed Image"):
+                        processed_image = gr.Image(label="Image with Layout Detection", type="pil", interactive=False, height=500)
+                    with gr.Tab("📝 Extracted Content"):
+                        markdown_output = gr.Markdown(value="Click 'Process Document' to see extracted content...", height=500)
+                    with gr.Tab("📋 Layout JSON"):
+                        json_output = gr.JSON(label="Layout Analysis Results", value=None)
+        
+        def handle_file_upload(file_path):
+            image, page_info = load_file_for_preview(file_path)
+            return image, page_info
+        
+        def clear_all():
+            global pdf_cache
+            pdf_cache = {"images": [], "current_page": 0, "total_pages": 0, "file_type": None, "is_parsed": False, "results": []}
+            return None, None, '<div class="page-info">No file loaded</div>', None, "Click 'Process Document' to see extracted content...", None
+        
+        file_input.change(handle_file_upload, inputs=[file_input], outputs=[image_preview, page_info])
+        prev_page_btn.click(lambda: turn_page("prev"), outputs=[image_preview, page_info, markdown_output, processed_image, json_output])
+        next_page_btn.click(lambda: turn_page("next"), outputs=[image_preview, page_info, markdown_output, processed_image, json_output])
+        process_btn.click(
+            process_document,
+            inputs=[file_input, model_choice, max_new_tokens, min_pixels, max_pixels],
+            outputs=[processed_image, markdown_output, json_output]
+        )
+        clear_btn.click(
+            clear_all,
+            outputs=[file_input, image_preview, page_info, processed_image, markdown_output, json_output]
+        )
+    return demo
 
 if __name__ == "__main__":
-    demo.queue(max_size=30).launch(share=True, mcp_server=True, ssr_mode=False, show_error=True)
+    demo = create_gradio_interface()
+    demo.queue(max_size=10).launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        debug=True,
+        show_error=True
+    )