Update app.py

app.py

@@ -10,10 +10,13 @@ import re
 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 qwen_vl_utils import process_vision_info
 
-from model import load_model, inference_dots_ocr, inference_dolphin
-
+# JavaScript for theme refresh
 js_func = """
 function refresh() {
     const url = new URL(window.location);
@@ -29,7 +32,7 @@ MIN_PIXELS = 3136
 MAX_PIXELS = 11289600
 IMAGE_FACTOR = 28
 
-# Prompts
+# Prompt for dots.ocr
 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,6 +48,77 @@ prompt = """Please output the layout information from the PDF image, including e
 5. Final Output: Single JSON object
 """
 
+# Model loading functions (from model.py)
+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 (from model.py)
+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  # Removed temperature=0.1 to fix the 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"),
@@ -217,7 +291,7 @@ def layoutjson2md(image: Image.Image, layout_data: List[Dict], text_key: str = '
                     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")
+                    markdown_lines.append(f"$$    \n{text}\n    $$\n")
                 else:
                     markdown_lines.append(f"**Formula:** {text}\n")
             elif category == 'Caption':
@@ -463,4 +537,4 @@ def create_gradio_interface():
 
 if __name__ == "__main__":
     demo = create_gradio_interface()
-    demo.queue(max_size=30).launch(share=False, debug=True, show_error=True)
+    demo.queue(max_size=30).launch(share=True, debug=True, show_error=True)