Update app.py

app.py

@@ -6,152 +6,59 @@ import traceback
 from io import BytesIO
 from typing import Any, Dict, List, Optional, Tuple
 import re
-import warnings
 
 import fitz  # PyMuPDF
 import gradio as gr
 import requests
 import torch
-from PIL import Image, ImageDraw, ImageFont
-from transformers import AutoModelForCausalLM, AutoProcessor, VisionEncoderDecoderModel
 from huggingface_hub import snapshot_download
+from PIL import Image, ImageDraw, ImageFont
 from qwen_vl_utils import process_vision_info
-
-# Suppress the FutureWarning for cleaner output (optional)
-warnings.filterwarnings(
-    "ignore",
-    category=FutureWarning,
-    message="Both `num_logits_to_keep` and `logits_to_keep` are set"
-)
-
-# JavaScript for theme refresh
-js_func = """
-function refresh() {
-    const url = new URL(window.location);
-    if (url.searchParams.get('__theme') !== 'dark') {
-        url.searchParams.set('__theme', 'dark');
-        window.location.href = url.href;
-    }
-}
-"""
+from transformers import AutoModelForCausalLM, AutoProcessor, Qwen2_5_VLForConditionalGeneration
 
 # Constants
 MIN_PIXELS = 3136
 MAX_PIXELS = 11289600
 IMAGE_FACTOR = 28
 
-# Prompt for dots.ocr
+# 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.
 
 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']
+
+2. Layout Categories: The possible categories are ['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
-"""
+    - Picture: For the 'Picture' category, the text field should be omitted.
+    - Formula: Format its text as LaTeX.
+    - Table: Format its text as HTML.
+    - All Others (Text, Title, etc.): Format their text as Markdown.
 
-# Model loading functions
-def load_model(model_name):
-    if model_name == "dots.ocr":
-        model_id = "rednote-hilab/dots.ocr"
-        model_path = "./models/dots-ocr-local"
-        snapshot_download(
-            repo_id=model_id,
-            local_dir=model_path,
-            local_dir_use_symlinks=False,
-        )
-        model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            attn_implementation="flash_attention_2",
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-            trust_remote_code=True
-        )
-        processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
-    elif model_name == "Dolphin":
-        model_id = "ByteDance/Dolphin"
-        processor = AutoProcessor.from_pretrained(model_id)
-        model = VisionEncoderDecoderModel.from_pretrained(model_id)
-        model.eval()
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        model.to(device)
-        model = model.half()  # Use half precision
-    else:
-        raise ValueError(f"Unknown model: {model_name}")
-    return model, processor
-
-# Inference functions
-def inference_dots_ocr(model, processor, image, prompt, max_new_tokens):
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image", "image": image},
-                {"type": "text", "text": prompt}
-            ]
-        }
-    ]
-    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    image_inputs, video_inputs = process_vision_info(messages)
-    inputs = processor(
-        text=[text],
-        images=image_inputs,
-        videos=video_inputs,
-        padding=True,
-        return_tensors="pt",
-    )
-    inputs = inputs.to(model.device)
-    with torch.no_grad():
-        generated_ids = model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            do_sample=False  # Temperature removed previously to fix another warning
-        )
-    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
-    output_text = processor.batch_decode(
-        generated_ids_trimmed,
-        skip_special_tokens=True,
-        clean_up_tokenization_spaces=False
-    )
-    return output_text[0] if output_text else ""
-
-def inference_dolphin(model, processor, image):
-    pixel_values = processor(image, return_tensors="pt").pixel_values.to(model.device).half()
-    generated_ids = model.generate(pixel_values)
-    generated_text = processor.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    return generated_text
-
-# Load models at startup
-models = {
-    "dots.ocr": load_model("dots.ocr"),
-    "Dolphin": load_model("Dolphin")
-}
+4. Constraints:
+    - The output text must be the original text from the image, with no translation.
+    - All layout elements must be sorted according to human reading order.
 
-# Global state for PDF handling
-pdf_cache = {
-    "images": [],
-    "current_page": 0,
-    "total_pages": 0,
-    "file_type": None,
-    "is_parsed": False,
-    "results": []
-}
+5. Final Output: The entire output must be a single JSON object.
+"""
 
-# Utility functions
+# Utility Functions
 def round_by_factor(number: int, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
     return round(number / factor) * factor
 
-def smart_resize(height: int, width: int, factor: int = 28, min_pixels: int = 3136, max_pixels: int = 11289600):
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int = 28,
+    min_pixels: int = 3136,
+    max_pixels: int = 11289600,
+):
+    """Rescales the image so that dimensions are divisible by 'factor', within pixel range, maintaining aspect ratio."""
     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)
@@ -163,6 +70,7 @@ def smart_resize(height: int, width: int, factor: int = 28, min_pixels: int = 31
     return h_bar, w_bar
 
 def fetch_image(image_input, min_pixels: int = None, max_pixels: int = None):
+    """Fetch and process an image."""
     if isinstance(image_input, str):
         if image_input.startswith(("http://", "https://")):
             response = requests.get(image_input)
@@ -173,20 +81,29 @@ def fetch_image(image_input, min_pixels: int = None, max_pixels: int = None):
         image = image_input.convert('RGB')
     else:
         raise ValueError(f"Invalid image input type: {type(image_input)}")
-    if min_pixels or max_pixels:
+    
+    if min_pixels is not None or max_pixels is not None:
         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)
+        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
 
 def load_images_from_pdf(pdf_path: str) -> List[Image.Image]:
+    """Load images from PDF file."""
     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)
+            mat = fitz.Matrix(2.0, 2.0)  # Increase resolution
             pix = page.get_pixmap(matrix=mat)
             img_data = pix.tobytes("ppm")
             image = Image.open(BytesIO(img_data)).convert('RGB')
@@ -198,66 +115,43 @@ def load_images_from_pdf(pdf_path: str) -> List[Image.Image]:
     return images
 
 def draw_layout_on_image(image: Image.Image, layout_data: List[Dict]) -> Image.Image:
+    """Draw layout bounding boxes on 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'
+        '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'
     }
+    
     try:
-        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 12)
-    except Exception:
-        font = ImageFont.load_default()
-    try:
+        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 12) or ImageFont.load_default()
         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)
+                label_width, label_height = label_bbox[2] - label_bbox[0], label_bbox[3] - label_bbox[1]
+                
+                label_x, label_y = bbox[0], 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
 
-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
-
 def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = 'text') -> str:
+    """Convert layout JSON to markdown format."""
     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]))
@@ -265,23 +159,21 @@ def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = '
             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')
+            
+            if category == 'Picture' and bbox and len(bbox) == 4:
+                try:
+                    x1, y1, x2, y2 = [max(0, int(x1)), max(0, int(y1)), min(image.width, int(x2)), min(image.height, int(y2))]
+                    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](data:image/png;base64,{img_data})\n")
+                    else:
+                        markdown_lines.append("![Image](Image region detected)\n")
+                except Exception as e:
+                    print(f"Error processing image region: {e}")
+                    markdown_lines.append("![Image](Image detected)\n")
             elif not text:
                 continue
             elif category == 'Title':
@@ -293,15 +185,9 @@ def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = '
             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")
+                markdown_lines.append(f"{text}\n" if text.strip().startswith('<') else f"**Table:** {text}\n")
             elif category == 'Formula':
-                if text.strip().startswith('$') or '\\' in text:
-                    markdown_lines.append(f"$$    \n{text}\n    $$\n")  # Fixed formatting, removed extra spaces
-                else:
-                    markdown_lines.append(f"**Formula:** {text}\n")
+                markdown_lines.append(f"$$\n{text}\n$$\n" if text.strip().startswith('$') or '\\' in text else f"**Formula:** {text}\n")
             elif category == 'Caption':
                 markdown_lines.append(f"*{text}*\n")
             elif category == 'Footnote':
@@ -314,37 +200,122 @@ def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = '
     except Exception as e:
         print(f"Error converting to markdown: {e}")
         return str(layout_data)
+    
     return "\n".join(markdown_lines)
 
+# Load Models
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load dot.ocr
+model_id = "rednote-hilab/dots.ocr"
+model_path = "./models/dots-ocr-local"
+snapshot_download(repo_id=model_id, local_dir=model_path, local_dir_use_symlinks=False)
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    attn_implementation="flash_attention_2",
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True
+)
+processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
+
+# 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()
+
+# 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()
+
+# Model Dictionary
+model_dict = {
+    "dot.ocr": {"model": model, "processor": processor, "process_layout": True},
+    "Camel-Doc-OCR-062825": {"model": model_m, "processor": processor_m, "process_layout": False},
+    "Megalodon-OCR-Sync-0713": {"model": model_t, "processor": processor_t, "process_layout": False},
+}
+
+# Global State
+pdf_cache = {"images": [], "current_page": 0, "total_pages": 0, "file_type": None, "is_parsed": False, "results": []}
+
+@spaces.GPU()
+def inference(model, processor, image: Image.Image, prompt: str, max_new_tokens: int = 24000) -> str:
+    """Run inference on an image with the given prompt using the specified model and processor."""
+    try:
+        messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": prompt}]}]
+        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        image_inputs, video_inputs = process_vision_info(messages)
+        inputs = processor(text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt").to(device)
+        
+        with torch.no_grad():
+            generated_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False, temperature=0.1)
+        generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+        output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        return output_text[0] if output_text else ""
+    except Exception as e:
+        print(f"Error during inference: {e}")
+        traceback.print_exc()
+        return f"Error during inference: {str(e)}"
+
+def process_image(
+    image: Image.Image,
+    model,
+    processor,
+    process_layout: bool,
+    min_pixels: Optional[int] = None,
+    max_pixels: Optional[int] = None
+) -> Dict[str, Any]:
+    """Process a single image with the specified model and processor."""
+    try:
+        if min_pixels is not None or max_pixels is not None:
+            image = fetch_image(image, min_pixels=min_pixels, max_pixels=max_pixels)
+        
+        raw_output = inference(model, processor, image, prompt)
+        result = {'original_image': image, 'raw_output': raw_output, 'processed_image': image, 'layout_result': None, 'markdown_content': raw_output}
+        
+        if process_layout:
+            try:
+                layout_data = json.loads(raw_output)
+                result['layout_result'] = layout_data
+                result['processed_image'] = draw_layout_on_image(image, layout_data)
+                result['markdown_content'] = layoutjson2md(image, layout_data, text_key='text')
+            except json.JSONDecodeError:
+                print("Failed to parse JSON output, using raw output")
+            except Exception as e:
+                print(f"Error processing layout: {e}")
+        
+        return result
+    except Exception as e:
+        print(f"Error processing image: {e}")
+        traceback.print_exc()
+        return {'original_image': image, 'raw_output': str(e), 'processed_image': image, 'layout_result': None, 'markdown_content': str(e)}
+
 def load_file_for_preview(file_path: str) -> Tuple[Optional[Image.Image], str]:
+    """Load file for preview (supports PDF and images)."""
     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": []
-            })
+            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": []
-            })
+            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}"
@@ -352,108 +323,28 @@ def load_file_for_preview(file_path: str) -> Tuple[Optional[Image.Image], str]:
         print(f"Error loading file: {e}")
         return None, f"Error loading file: {str(e)}"
 
-@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
-
 def turn_page(direction: str) -> Tuple[Optional[Image.Image], str, Any, Optional[Image.Image], Optional[Dict]]:
+    """Navigate through PDF pages and update outputs."""
     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)
+
+    pdf_cache["current_page"] = max(0, pdf_cache["current_page"] - 1) if direction == "prev" else 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"]):
+
+    markdown_content, processed_img, layout_json = "Page not processed yet", None, None
+    if pdf_cache["is_parsed"] and index < len(pdf_cache["results"]) and pdf_cache["results"][index]:
         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
+        markdown_content = result.get('markdown_content') or result.get('raw_output', 'No content available')
+        processed_img = result.get('processed_image')
+        layout_json = result.get('layout_result')
+
+    return current_image_preview, page_info_html, markdown_content, processed_img, layout_json
 
 def create_gradio_interface():
+    """Create the Gradio interface."""
     css = """
     .main-container { max-width: 1400px; margin: 0 auto; }
     .header-text { text-align: center; color: #2c3e50; margin-bottom: 20px; }
@@ -471,75 +362,102 @@ def create_gradio_interface():
     .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="bethecloud/storj_theme", css=css, js=js_func) as demo:
+
+    with gr.Blocks(theme="bethecloud/storj_theme", css=css, title="Dot●OCR Comparator") as demo:
         gr.HTML("""
         <div class="title" style="text-align: center">
-            <h1>Dot<span style="color: red;">●</span><strong></strong>OCR vs Dolphin🐬</h1>
+            <h1>Dot<span style="color: red;">●</span>OCR Comparator</h1>
             <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;">
                 Advanced vision-language model for image/PDF to markdown document processing
             </p>
         </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"
+                model_choice = gr.Radio(
+                    choices=["dot.ocr", "Camel-Doc-OCR-062825", "Megalodon-OCR-Sync-0713"],
+                    label="Select Model",
+                    value="dot.ocr"
                 )
+                file_input = gr.File(label="Upload Image or PDF", file_types=[".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".pdf"], type="filepath")
+                with gr.Row():
+                    examples = gr.Examples(
+                        examples=["examples/sample_image1.png", "examples/sample_image2.png", "examples/sample_pdf.pdf"],
+                        inputs=file_input,
+                        label="Example Documents"
+                    )
                 image_preview = gr.Image(label="Preview", type="pil", interactive=False, height=300)
                 with gr.Row():
-                    prev_page_btn = gr.Button("⬅ Previous", size="md")
+                    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"
-                )
+                    next_page_btn = gr.Button("Next ▶", size="md")
                 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 Document", variant="secondary")
+                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"):
+                    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"):
+                    with gr.Tab("📝 Extracted Content"):
                         markdown_output = gr.Markdown(value="Click 'Process Document' to see extracted content...", height=500)
-                    with gr.Tab("⏲ Layout JSON"):
+                    with gr.Tab("📋 Layout JSON"):
                         json_output = gr.JSON(label="Layout Analysis Results", value=None)
-                with gr.Row():
-                    examples = gr.Examples(
-                        examples=["examples/sample_image1.png", "examples/sample_image2.png", "examples/sample_pdf.pdf"],
-                        inputs=file_input,
-                        label="Example Documents"
-                    )
+
+        def process_document(file_path, model_choice, max_tokens, min_pix, max_pix):
+            """Process the uploaded document with the selected model."""
+            global pdf_cache
+            if not file_path:
+                return None, "Please upload a file first.", None
+            if model_choice not in model_dict:
+                return None, "Invalid model selected", None
+            
+            selected_model = model_dict[model_choice]["model"]
+            selected_processor = model_dict[model_choice]["processor"]
+            process_layout = model_dict[model_choice]["process_layout"]
+            
+            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, all_markdown = [], []
+                for i, img in enumerate(pdf_cache["images"]):
+                    result = process_image(img, selected_model, selected_processor, process_layout, int(min_pix) if min_pix else None, int(max_pix) if max_pix else None)
+                    all_results.append(result)
+                    if result.get('markdown_content'):
+                        all_markdown.append(f"## Page {i+1}\n\n{result['markdown_content']}")
+                pdf_cache["results"] = all_results
+                pdf_cache["is_parsed"] = True
+                first_result = all_results[0]
+                return first_result['processed_image'], "\n\n---\n\n".join(all_markdown), first_result['layout_result']
+            else:
+                result = process_image(image, selected_model, selected_processor, process_layout, int(min_pix) if min_pix else None, int(max_pix) if max_pix else None)
+                pdf_cache["results"] = [result]
+                pdf_cache["is_parsed"] = True
+                return result['processed_image'], result['markdown_content'] or "No content extracted", result['layout_result']
+
         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]
-        )
+        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 = create_gradio_interface()
-    demo.queue(max_size=30).launch(debug=True, show_error=True)  # Removed share=True
+    demo.queue(max_size=50).launch(share=False, debug=True, show_error=True)