Update app.py

app.py

@@ -6,8 +6,9 @@ import traceback
 from io import BytesIO
 from typing import Any, Dict, List, Optional, Tuple
 import re
+from threading import Thread
+import time
 
-import fitz
 import gradio as gr
 import requests
 import torch
@@ -17,15 +18,15 @@ from transformers import (
     AutoProcessor,
     TextIteratorStreamer,
 )
-from qwen_vl_utils import process_vision_info
 
 # Constants
 MIN_PIXELS = 3136
 MAX_PIXELS = 11289600
 IMAGE_FACTOR = 28
+MAX_INPUT_TOKEN_LENGTH = 4096
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Prompts
+# Prompt for Layout Analysis
 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]
@@ -45,190 +46,74 @@ prompt = """Please output the layout information from the PDF image, including e
 5. Final Output: The entire output must be a single JSON object.
 """
 
-# Load models
+# Load Models
 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
+    MODEL_ID_M, trust_remote_code=True, torch_dtype=torch.float16
 ).to(device).eval()
 
 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
+    MODEL_ID_T, trust_remote_code=True, torch_dtype=torch.float16
 ).to(device).eval()
 
 MODEL_ID_C = "nanonets/Nanonets-OCR-s"
 processor_c = AutoProcessor.from_pretrained(MODEL_ID_C, trust_remote_code=True)
 model_c = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_ID_C,
-    trust_remote_code=True,
-    torch_dtype=torch.float16
+    MODEL_ID_C, trust_remote_code=True, torch_dtype=torch.float16
 ).to(device).eval()
 
 MODEL_ID_G = "echo840/MonkeyOCR"
 SUBFOLDER = "Recognition"
 processor_g = AutoProcessor.from_pretrained(
-    MODEL_ID_G,
-    trust_remote_code=True,
-    subfolder=SUBFOLDER
+    MODEL_ID_G, trust_remote_code=True, subfolder=SUBFOLDER
 )
 model_g = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_ID_G,
-    trust_remote_code=True,
-    subfolder=SUBFOLDER,
-    torch_dtype=torch.float16
+    MODEL_ID_G, trust_remote_code=True, subfolder=SUBFOLDER, torch_dtype=torch.float16
 ).to(device).eval()
 
-# Utility functions
-def round_by_factor(number: int, factor: int) -> int:
-    return round(number / factor) * factor
-
-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 too extreme: {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):
-    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 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
-
-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 images
-
-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'
-    }
-    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
 
+# Utility functions
 def is_arabic_text(text: str) -> bool:
+    """Check if text contains mostly Arabic characters."""
     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)
+    # Simplified check for Arabic characters in the given text
     arabic_chars = 0
     total_chars = 0
-    for char in combined_text:
+    for char in text:
         if char.isalpha():
             total_chars += 1
-            if ('\u0600' <= char <= '\u06FF') or ('\u0750' <= char <= '\u077F') or ('\u08A0' <= char <= '\u08FF'):
+            if '\u0600' <= char <= '\u06FF':
                 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
     from io import BytesIO
     markdown_lines = []
     try:
+        # Sort items by reading order (top to bottom, left to right)
         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")
+                        x1, y1, x2, y2 = [int(coord) for coord in bbox]
+                        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")
                     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")
+                        markdown_lines.append("![Image](Image region detected)\n")
             elif not text:
                 continue
             elif category == 'Title':
@@ -239,311 +124,173 @@ def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = '
                 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 == 'Table' and text.strip().startswith('<'):
+                markdown_lines.append(f"{text}\n")
+            elif category == 'Formula' and (text.strip().startswith('$') or '\\' in text):
+                markdown_lines.append(f"$$\n{text}\n$$\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")
+            elif category not in ['Page-header', 'Page-footer']:
                 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 f"### Error converting to Markdown\n\n```\n{str(layout_data)}\n```"
     return "\n".join(markdown_lines)
 
-# PDF handling state
-pdf_cache = {
-    "images": [],
-    "current_page": 0,
-    "total_pages": 0,
-    "file_type": None,
-    "is_parsed": False,
-    "results": []
-}
 
 @spaces.GPU
-def inference(model_name: str, image: Image.Image, prompt: str, max_new_tokens: int = 1024) -> str:
-    try:
-        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 == "Nanonets-OCR-s":
-            processor = processor_c
-            model = model_c
-        elif model_name == "MonkeyOCR-Recognition":
-            processor = processor_g
-            model = model_g
-        else:
-            raise ValueError(f"Invalid model selected: {model_name}")
-
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image", "image": image},
-                    {"type": "text", "text": prompt}
-                ]
-            }
-        ]
-
-        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = processor(text=[text], images=[image], 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 ""
+def generate_and_process(model_name: str, image: Image.Image, max_new_tokens: int):
+    """
+    Generates a response using streaming, then processes the final output.
+    Yields updates for the raw stream, final markdown, and JSON output.
+    """
+    if image is None:
+        yield "Please upload an image.", "Please upload an image.", None
+        return
 
-    except Exception as e:
-        print(f"Error during inference: {e}")
-        traceback.print_exc()
-        return f"Error during inference: {str(e)}"
-
-def process_image(
-    model_name: str,
-    image: Image.Image,
-    min_pixels: Optional[int] = None,
-    max_pixels: Optional[int] = None
-) -> Dict[str, Any]:
-    try:
-        if min_pixels or max_pixels:
-            image = fetch_image(image, min_pixels=min_pixels, max_pixels=max_pixels)
-        raw_output = inference(model_name, image, prompt)
-        result = {
-            'original_image': image,
-            'raw_output': raw_output,
-            'processed_image': image,
-            'layout_result': None,
-            'markdown_content': None
-        }
-        try:
-            json_match = re.search(r'```json\s*([\s\S]+?)\s*```', raw_output)
-            json_str = json_match.group(1) if json_match else raw_output
-            layout_data = json.loads(json_str)
-            result['layout_result'] = layout_data
-            try:
-                processed_image = draw_layout_on_image(image, layout_data)
-                result['processed_image'] = processed_image
-            except Exception as e:
-                print(f"Error drawing layout: {e}")
-            try:
-                markdown_content = layoutjson2md(image, layout_data, text_key='text')
-                result['markdown_content'] = markdown_content
-            except Exception as e:
-                print(f"Error generating markdown: {e}")
-                result['markdown_content'] = raw_output
-        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: {e}")
-        traceback.print_exc()
-        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 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()
+    # 1. Select Model and Processor
+    if model_name == "Camel-Doc-OCR-062825":
+        processor, model = processor_m, model_m
+    elif model_name == "Megalodon-OCR-Sync-0713":
+        processor, model = processor_t, model_t
+    elif model_name == "Nanonets-OCR-s":
+        processor, model = processor_c, model_c
+    elif model_name == "MonkeyOCR-Recognition":
+        processor, model = processor_g, model_g
+    else:
+        yield "Invalid model selected.", "Invalid model selected.", None
+        return
+
+    # 2. Prepare inputs for the model
+    messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt}]}]
+    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=True,
+        max_length=MAX_INPUT_TOKEN_LENGTH
+    ).to(device)
+
+    # 3. Stream the generation
+    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    buffer = ""
+    # Initial placeholder yield
+    yield buffer, "⏳ Generating response...", None
+
+    for new_text in streamer:
+        buffer += new_text
+        buffer = buffer.replace("<|im_end|>", "")
+        time.sleep(0.01)  # Small delay for smoother streaming
+        yield buffer, "⏳ Generating response...", None
+
+    # 4. Process the final buffer content
     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}"
+        json_match = re.search(r'```json\s*([\s\S]+?)\s*```', buffer)
+        json_str = json_match.group(1) if json_match else buffer
+        layout_data = json.loads(json_str)
+        
+        markdown_content = layoutjson2md(image, layout_data)
+        
+        # Final yield with all processed content
+        yield buffer, markdown_content, layout_data
+
+    except json.JSONDecodeError:
+        error_msg = "❌ Failed to parse JSON from model output."
+        yield buffer, error_msg, {"error": "JSONDecodeError", "raw_output": buffer}
     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, 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>'
-    markdown_content = "Page not processed yet"
-    processed_img = None
-    layout_json = None
-    if pdf_cache["is_parsed"] and index < len(pdf_cache["results"]) and pdf_cache["results"][index]:
-        result = pdf_cache["results"][index]
-        markdown_content = result.get('markdown_content') or result.get('raw_output', 'No content available')
-        processed_img = result.get('processed_image', None)
-        layout_json = result.get('layout_result', None)
-    if is_arabic_text(markdown_content):
-        markdown_update = gr.update(value=markdown_content, rtl=True)
-    else:
-        markdown_update = markdown_content
-    return current_image_preview, page_info_html, markdown_update, processed_img, layout_json
+        error_msg = f"❌ An error occurred during post-processing: {e}"
+        yield buffer, error_msg, {"error": str(e), "raw_output": buffer}
+
 
 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;} 
+        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; }
+        background-color: darkblue !important; transform: translateY(-2px) !important; 
+        box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important; 
+    }
     """
     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
+                Advanced vision-language model for image to markdown document processing
             </p>
         </div>
         """)
+        
+        # Keep track of the uploaded image
+        image_state = gr.State(None)
+
         with gr.Row():
+            # Left column - Input and controls
             with gr.Column(scale=1):
                 model_choice = gr.Radio(
                     choices=["Camel-Doc-OCR-062825", "MonkeyOCR-Recognition", "Nanonets-OCR-s", "Megalodon-OCR-Sync-0713"],
                     label="Select Model",
                     value="Camel-Doc-OCR-062825"
                 )
-                file_input = gr.File(
-                    label="Upload Image or PDF",
-                    file_types=[".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".pdf"],
-                    type="filepath"
+                file_input = gr.Image(
+                    label="Upload Image",
+                    type="pil",
+                    sources=['upload']
                 )
-                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")
                 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")
+
+            # Right column - Results
             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)
+                        output_stream = gr.Textbox(label="Raw Output Stream", interactive=False, lines=10, show_copy_button=True)
+                        with gr.Accordion("(Formatted Result)", open=True):
+                            markdown_output = gr.Markdown(label="Formatted Result (Result.md)")
+                    
                     with gr.Tab("📋 Layout JSON"):
-                        json_output = gr.JSON(label="Layout Analysis Results", value=None)
-        def process_document(model_name, file_path, max_tokens, min_pix, max_pix):
-            global pdf_cache
-            try:
-                if not file_path:
-                    return None, "Please upload a file first.", None
-                load_file_for_preview(file_path)
-                if pdf_cache["file_type"] == "pdf":
-                    all_results = []
-                    all_markdown = []
-                    for i, img in enumerate(pdf_cache["images"]):
-                        result = process_image(model_name, img, min_pixels=int(min_pix) if min_pix else None, max_pixels=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]
-                    combined_markdown = "\n\n---\n\n".join(all_markdown)
-                    if is_arabic_text(combined_markdown):
-                        markdown_update = gr.update(value=combined_markdown, rtl=True)
-                    else:
-                        markdown_update = combined_markdown
-                    return first_result['processed_image'], markdown_update, first_result['layout_result']
-                else:
-                    result = process_image(model_name, pdf_cache["images"][0], min_pixels=int(min_pix) if min_pix else None, max_pixels=int(max_pix) if max_pix else None)
-                    pdf_cache["results"] = [result]
-                    pdf_cache["is_parsed"] = True
-                    content = result['markdown_content'] or "No content extracted"
-                    if is_arabic_text(content):
-                        markdown_update = gr.update(value=content, rtl=True)
-                    else:
-                        markdown_update = content
-                    return result['processed_image'], markdown_update, result['layout_result']
-            except Exception as e:
-                error_msg = f"Error processing document: {str(e)}"
-                print(error_msg)
-                traceback.print_exc()
-                return None, error_msg, None
-        def handle_file_upload(file_path):
-            if not file_path:
-                return None, "No file loaded"
-            image, page_info = load_file_for_preview(file_path)
-            return image, page_info
+                        json_output = gr.JSON(label="Layout Analysis Results (JSON)", value=None)
+
+        # Event Handlers
+        def handle_file_upload(image):
+            """Store the uploaded image in the state."""
+            return image
+
         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=[model_choice, file_input, 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])
+            """Clear all data and reset the interface."""
+            return None, None, "Click 'Process Document' to see extracted content...", None, None
+
+        file_input.upload(handle_file_upload, inputs=[file_input], outputs=[image_state])
+        
+        process_btn.click(
+            generate_and_process,
+            inputs=[model_choice, image_state, max_new_tokens],
+            outputs=[output_stream, markdown_output, json_output]
+        )
+
+        clear_btn.click(
+            clear_all,
+            outputs=[file_input, image_state, markdown_output, json_output, output_stream]
+        )
+
     return demo
 
 if __name__ == "__main__":
     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)
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True, show_error=True)