app.py

import io
import os
import tempfile
import time
import uuid
import cv2
import gradio as gr
import pymupdf
import spaces
import torch
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoModelForCausalLM, AutoProcessor, VisionEncoderDecoderModel
from huggingface_hub import snapshot_download
from qwen_vl_utils import process_vision_info
from utils.utils import prepare_image, parse_layout_string, process_coordinates, ImageDimensions
from utils.markdown_utils import MarkdownConverter

# Define device
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load dot.ocr model
dot_ocr_model_id = "rednote-hilab/dots.ocr"
dot_ocr_model = AutoModelForCausalLM.from_pretrained(
    dot_ocr_model_id,
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True
)
dot_ocr_processor = AutoProcessor.from_pretrained(
    dot_ocr_model_id,
    trust_remote_code=True
)

# Load Dolphin model
dolphin_model_id = "ByteDance/Dolphin"
dolphin_processor = AutoProcessor.from_pretrained(dolphin_model_id)
dolphin_model = VisionEncoderDecoderModel.from_pretrained(dolphin_model_id)
dolphin_model.eval()
dolphin_model.to(device)
dolphin_model = dolphin_model.half()
dolphin_tokenizer = dolphin_processor.tokenizer

# Constants
MIN_PIXELS = 3136
MAX_PIXELS = 11289600
IMAGE_FACTOR = 28

# 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: 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: 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.

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.

5. Final Output: The entire output must be a single JSON object.
"""

# 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,
):
    """Rescales the image so that the following conditions are met:
    1. Both dimensions (height and width) are divisible by 'factor'.
    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
    3. The aspect ratio of the image is maintained as closely as possible.
    """
    if max(height, width) / min(height, width) > 200:
        raise ValueError(
            f"absolute aspect ratio must be smaller than 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

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)
            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:
        raise ValueError(f"Invalid image input type: {type(image_input)}")
    
    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
        )
        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 = pymupdf.open(pdf_path)
        for page_num in range(len(pdf_document)):
            page = pdf_document.load_page(page_num)
            mat = pymupdf.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')
            images.append(image)
        pdf_document.close()
    except Exception as e:
        print(f"Error loading PDF: {e}")
        return []
    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'
    }
    
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 12)
    except Exception:
        font = 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)
            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)
    return img_copy

def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = 'text') -> str:
    """Convert layout JSON to markdown format"""
    import base64
    from io import BytesIO
    
    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 = bbox
                        x1, y1 = max(0, int(x1)), max(0, int(y1))
                        x2, y2 = 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")
                else:
                    markdown_lines.append("![Image](Image detected)\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)

# Global state variables
pdf_cache = {
    "images": [],
    "current_page": 0,
    "total_pages": 0,
    "file_type": None,
    "is_parsed": False,
    "results": []
}

@spaces.GPU()
def dot_ocr_inference(image: Image.Image, prompt: str, max_new_tokens: int = 24000) -> str:
    """Run inference on an image with the given prompt using dot.ocr model"""
    try:
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {"type": "text", "text": prompt}
                ]
            }
        ]
        text = dot_ocr_processor.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        image_inputs, video_inputs = process_vision_info(messages)
        inputs = dot_ocr_processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        inputs = inputs.to(device)
        with torch.no_grad():
            generated_ids = dot_ocr_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 = dot_ocr_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 dot.ocr inference: {e}")
        return f"Error during inference: {str(e)}"

def process_image_dot_ocr(image: Image.Image, min_pixels: Optional[int] = None, max_pixels: Optional[int] = None) -> Dict[str, Any]:
    """Process a single image with the dot.ocr model"""
    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 = dot_ocr_inference(image, prompt)
        result = {
            'original_image': image,
            'raw_output': raw_output,
            'processed_image': image,
            'layout_result': None,
            'markdown_content': None
        }
        try:
            layout_data = json.loads(raw_output)
            result['layout_result'] = layout_data
            processed_image = draw_layout_on_image(image, layout_data)
            result['processed_image'] = processed_image
            markdown_content = layoutjson2md(image, layout_data, text_key='text')
            result['markdown_content'] = markdown_content
        except json.JSONDecodeError:
            print("Failed to parse JSON output, using raw output")
            result['markdown_content'] = raw_output
        return result
    except Exception as e:
        print(f"Error processing image with dot.ocr: {e}")
        return {
            'original_image': image,
            'raw_output': f"Error processing image: {str(e)}",
            'processed_image': image,
            'layout_result': None,
            'markdown_content': f"Error processing image: {str(e)}"
        }

def process_all_pages_dot_ocr(file_path, min_pixels, max_pixels):
    """Process all pages of a document with dot.ocr model"""
    if file_path.lower().endswith('.pdf'):
        images = load_images_from_pdf(file_path)
    else:
        images = [Image.open(file_path).convert('RGB')]
    results = []
    for img in images:
        result = process_image_dot_ocr(img, min_pixels, max_pixels)
        results.append(result)
    return results

# Dolphin model functions
@spaces.GPU()
def dolphin_model_chat(prompt, image):
    """Process an image or batch of images with the given prompt(s) using Dolphin model"""
    is_batch = isinstance(image, list)
    if not is_batch:
        images = [image]
        prompts = [prompt]
    else:
        images = image
        prompts = prompt if isinstance(prompt, list) else [prompt] * len(images)
    batch_inputs = dolphin_processor(images, return_tensors="pt", padding=True)
    batch_pixel_values = batch_inputs.pixel_values.half().to(device)
    prompts = [f"<s>{p} <Answer/>" for p in prompts]
    batch_prompt_inputs = dolphin_tokenizer(
        prompts,
        add_special_tokens=False,
        return_tensors="pt"
    )
    batch_prompt_ids = batch_prompt_inputs.input_ids.to(device)
    batch_attention_mask = batch_prompt_inputs.attention_mask.to(device)
    outputs = dolphin_model.generate(
        pixel_values=batch_pixel_values,
        decoder_input_ids=batch_prompt_ids,
        decoder_attention_mask=batch_attention_mask,
        min_length=1,
        max_length=4096,
        pad_token_id=dolphin_tokenizer.pad_token_id,
        eos_token_id=dolphin_tokenizer.eos_token_id,
        use_cache=True,
        bad_words_ids=[[dolphin_tokenizer.unk_token_id]],
        return_dict_in_generate=True,
        do_sample=False,
        num_beams=1,
        repetition_penalty=1.1
    )
    sequences = dolphin_tokenizer.batch_decode(outputs.sequences, skip_special_tokens=False)
    results = []
    for i, sequence in enumerate(sequences):
        cleaned = sequence.replace(prompts[i], "").replace("<pad>", "").replace("</s>", "").strip()
        results.append(cleaned)
    if not is_batch:
        return results[0]
    return results

def process_element_batch_dolphin(elements, prompt, max_batch_size=16):
    """Process elements of the same type in batches for Dolphin model"""
    results = []
    batch_size = min(len(elements), max_batch_size)
    for i in range(0, len(elements), batch_size):
        batch_elements = elements[i:i+batch_size]
        crops_list = [elem["crop"] for elem in batch_elements]
        prompts_list = [prompt] * len(crops_list)
        batch_results = dolphin_model_chat(prompts_list, crops_list)
        for j, result in enumerate(batch_results):
            elem = batch_elements[j]
            results.append({
                "label": elem["label"],
                "bbox": elem["bbox"],
                "text": result.strip(),
                "reading_order": elem["reading_order"],
            })
    return results

def process_page_dolphin(image_path):
    """Process a single page with Dolphin model"""
    pil_image = Image.open(image_path).convert("RGB")
    layout_output = dolphin_model_chat("Parse the reading order of this document.", pil_image)
    padded_image, dims = prepare_image(pil_image)
    recognition_results = process_elements_dolphin(layout_output, padded_image, dims)
    return recognition_results

def process_elements_dolphin(layout_results, padded_image, dims):
    """Parse all document elements for Dolphin model"""
    layout_results = parse_layout_string(layout_results)
    text_elements = []
    table_elements = []
    figure_results = []
    previous_box = None
    reading_order = 0
    for bbox, label in layout_results:
        try:
            x1, y1, x2, y2, orig_x1, orig_y1, orig_x2, orig_y2, previous_box = process_coordinates(
                bbox, padded_image, dims, previous_box
            )
            cropped = padded_image[y1:y2, x1:x2]
            if cropped.size > 0 and (cropped.shape[0] > 3 and cropped.shape[1] > 3):
                if label == "fig":
                    try:
                        pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
                        buffered = io.BytesIO()
                        pil_crop.save(buffered, format="PNG")
                        img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
                        figure_results.append(
                            {
                                "label": label,
                                "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
                                "text": img_base64,
                                "reading_order": reading_order,
                            }
                        )
                    except Exception as e:
                        print(f"Error encoding figure to base64: {e}")
                        figure_results.append(
                            {
                                "label": label,
                                "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
                                "text": "",
                                "reading_order": reading_order,
                            }
                        )
                else:
                    pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
                    element_info = {
                        "crop": pil_crop,
                        "label": label,
                        "bbox": [orig_x1, orig_y1, orig_x2, orig_y2],
                        "reading_order": reading_order,
                    }
                    if label == "tab":
                        table_elements.append(element_info)
                    else:
                        text_elements.append(element_info)
            reading_order += 1
        except Exception as e:
            print(f"Error processing bbox with label {label}: {str(e)}")
            continue
    recognition_results = figure_results.copy()
    if text_elements:
        text_results = process_element_batch_dolphin(text_elements, "Read text in the image.")
        recognition_results.extend(text_results)
    if table_elements:
        table_results = process_element_batch_dolphin(table_elements, "Parse the table in the image.")
        recognition_results.extend(table_results)
    recognition_results.sort(key=lambda x: x.get("reading_order", 0))
    return recognition_results

def generate_markdown(recognition_results):
    """Generate markdown from recognition results for Dolphin model"""
    converter = MarkdownConverter()
    return converter.convert(recognition_results)

def convert_all_pdf_pages_to_images(file_path, target_size=896):
    """Convert all pages of a PDF to images for Dolphin model"""
    if file_path is None:
        return []
    try:
        file_ext = os.path.splitext(file_path)[1].lower()
        if file_ext == '.pdf':
            doc = pymupdf.open(file_path)
            image_paths = []
            for page_num in range(len(doc)):
                page = doc[page_num]
                rect = page.rect
                scale = target_size / max(rect.width, rect.height)
                mat = pymupdf.Matrix(scale, scale)
                pix = page.get_pixmap(matrix=mat)
                img_data = pix.tobytes("png")
                pil_image = Image.open(io.BytesIO(img_data))
                with tempfile.NamedTemporaryFile(suffix=f"_page_{page_num}.png", delete=False) as tmp_file:
                    pil_image.save(tmp_file.name, "PNG")
                    image_paths.append(tmp_file.name)
            doc.close()
            return image_paths
        else:
            converted_path = convert_to_image(file_path, target_size)
            return [converted_path] if converted_path else []
    except Exception as e:
        print(f"Error converting PDF pages to images: {e}")
        return []

def convert_to_image(file_path, target_size=896, page_num=0):
    """Convert input file to image format for Dolphin model"""
    if file_path is None:
        return None
    try:
        file_ext = os.path.splitext(file_path)[1].lower()
        if file_ext == '.pdf':
            doc = pymupdf.open(file_path)
            if page_num >= len(doc):
                page_num = 0
            page = doc[page_num]
            rect = page.rect
            scale = target_size / max(rect.width, rect.height)
            mat = pymupdf.Matrix(scale, scale)
            pix = page.get_pixmap(matrix=mat)
            img_data = pix.tobytes("png")
            pil_image = Image.open(io.BytesIO(img_data))
            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
                pil_image.save(tmp_file.name, "PNG")
                doc.close()
                return tmp_file.name
        else:
            pil_image = Image.open(file_path).convert("RGB")
            w, h = pil_image.size
            if max(w, h) > target_size:
                if w > h:
                    new_w, new_h = target_size, int(h * target_size / w)
                else:
                    new_w, new_h = int(w * target_size / h), target_size
                pil_image = pil_image.resize((new_w, new_h), Image.Resampling.LANCZOS)
            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
                pil_image.save(tmp_file.name, "PNG")
                return tmp_file.name
    except Exception as e:
        print(f"Error converting file to image: {e}")
        return file_path

def process_all_pages_dolphin(file_path):
    """Process all pages of a document with Dolphin model"""
    image_paths = convert_all_pdf_pages_to_images(file_path)
    per_page_results = []
    for image_path in image_paths:
        try:
            original_image = Image.open(image_path).convert('RGB')
            recognition_results = process_page_dolphin(image_path)
            markdown_content = generate_markdown(recognition_results)
            placeholder_text = "Layout visualization not available for Dolphin model"
            processed_image = create_placeholder_image(placeholder_text, size=(original_image.width, original_image.height))
            per_page_results.append({
                'original_image': original_image,
                'processed_image': processed_image,
                'markdown_content': markdown_content,
                'layout_result': recognition_results
            })
        except Exception as e:
            print(f"Error processing page: {e}")
            per_page_results.append({
                'original_image': Image.new('RGB', (100, 100), color='white'),
                'processed_image': create_placeholder_image("Error processing page", size=(100, 100)),
                'markdown_content': f"Error processing page: {str(e)}",
                'layout_result': None
            })
        finally:
            if os.path.exists(image_path):
                os.remove(image_path)
    return per_page_results

def create_placeholder_image(text, size=(400, 200)):
    """Create a placeholder image with text"""
    img = Image.new('RGB', size, color='white')
    draw = ImageDraw.Draw(img)
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 16)
    except Exception:
        font = ImageFont.load_default()
    draw.text((10, 10), text, fill='black', font=font)
    return img

# Gradio interface functions
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": []
            })
            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)}"

def turn_page(direction: str) -> Tuple[Optional[Image.Image], str, str, Optional[Image.Image], Optional[Dict]]:
    """Navigate through PDF pages and update all relevant outputs."""
    global pdf_cache
    if not pdf_cache["images"]:
        return None, "No file loaded", "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"Page {index + 1} / {pdf_cache['total_pages']}"
    if pdf_cache["is_parsed"] and index < len(pdf_cache["results"]):
        result = pdf_cache["results"][index]
        processed_img = result['processed_image']
        markdown_content = result['markdown_content'] or "No content available"
        layout_json = result['layout_result']
    else:
        processed_img = None
        markdown_content = "Page not processed yet"
        layout_json = None
    return current_image_preview, page_info_html, markdown_content, processed_img, layout_json

def process_document(model_choice, file_path, max_tokens, min_pix, max_pix):
    """Process the uploaded document with the selected model"""
    global pdf_cache
    try:
        if not file_path:
            return None, "Please upload a file first.", None
        if model_choice == "dot.ocr":
            results = process_all_pages_dot_ocr(file_path, min_pix, max_pix)
        elif model_choice == "Dolphin":
            results = process_all_pages_dolphin(file_path)
        else:
            raise ValueError("Invalid model choice")
        pdf_cache["results"] = results
        pdf_cache["is_parsed"] = True
        first_result = results[0]
        if model_choice == "dot.ocr":
            processed_img = first_result['processed_image']
            markdown_content = first_result['markdown_content']
            layout_json = first_result['layout_result']
        else:
            processed_img = first_result['processed_image']
            markdown_content = first_result['markdown_content']
            layout_json = first_result['layout_result']
        return processed_img, markdown_content, layout_json
    except Exception as e:
        error_msg = f"Error processing document: {str(e)}"
        print(error_msg)
        return None, error_msg, None

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; }
    .process-button { 
        border: none !important; 
        color: white !important; 
        font-weight: bold !important; 
        background-color: blue !important;} 
    .process-button:hover { 
        background-color: darkblue !important;
        transform: translateY(-2px) !important; 
        box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important; }
    .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="bethecloud/storj_theme", css=css) as demo:
        gr.HTML("""
        <div class="title" style="text-align: center">
            <h1>Dot<span style="color: red;">●</span><strong></strong>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):
                model_choice = gr.Radio(
                    choices=["dot.ocr", "Dolphin"],
                    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")
                    page_info = gr.HTML("No file loaded")
                    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",
                        info="Maximum number of tokens to generate"
                    )
                    min_pixels = gr.Number(
                        value=MIN_PIXELS,
                        label="Min Pixels",
                        info="Minimum image resolution"
                    )
                    max_pixels = gr.Number(
                        value=MAX_PIXELS,
                        label="Max Pixels", 
                        info="Maximum image resolution"
                    )
                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
                        )

        # Event handlers
        file_input.change(
            lambda file_path: load_file_for_preview(file_path),
            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=[model_choice, file_input, max_new_tokens, min_pixels, max_pixels],
            outputs=[processed_image, markdown_output, json_output]
        )

        clear_btn.click(
            lambda: (None, None, "No file loaded", None, "Click 'Process Document' to see extracted content...", None),
            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=10).launch(share=False, debug=True, show_error=True)