Converting to DeepSeek-OCR-2

app.py

@@ -14,23 +14,25 @@ import numpy as np
 import base64
 from io import StringIO, BytesIO
 
-MODEL_NAME = 'deepseek-ai/DeepSeek-OCR'
+MODEL_NAME = 'deepseek-ai/DeepSeek-OCR-2'
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
 model = AutoModel.from_pretrained(MODEL_NAME, _attn_implementation='flash_attention_2', torch_dtype=torch.bfloat16, trust_remote_code=True, use_safetensors=True)
 model = model.eval().cuda()
 
 MODEL_CONFIGS = {
-    "Gundam": {"base_size": 1024, "image_size": 640, "crop_mode": True},
-    "Tiny": {"base_size": 512, "image_size": 512, "crop_mode": False},
-    "Small": {"base_size": 640, "image_size": 640, "crop_mode": False},
-    "Base": {"base_size": 1024, "image_size": 1024, "crop_mode": False},
-    "Large": {"base_size": 1280, "image_size": 1280, "crop_mode": False}
+    "Default": {"base_size": 1024, "image_size": 768, "crop_mode": True},
+    "Quality": {"base_size": 1280, "image_size": 960, "crop_mode": True},
+    "Fast": {"base_size": 1024, "image_size": 640, "crop_mode": True},
+    "No Crop": {"base_size": 1024, "image_size": 768, "crop_mode": False},
+    "Small": {"base_size": 768, "image_size": 512, "crop_mode": False},
 }
 
 TASK_PROMPTS = {
     "📋 Markdown": {"prompt": "<image>\n<|grounding|>Convert the document to markdown.", "has_grounding": True},
     "📝 Free OCR": {"prompt": "<image>\nFree OCR.", "has_grounding": False},
+    "🖼️ OCR Image": {"prompt": "<image>\n<|grounding|>OCR this image.", "has_grounding": True},
+    "📊 Parse Figure": {"prompt": "<image>\nParse the figure.", "has_grounding": False},
     "📍 Locate": {"prompt": "<image>\nLocate <|ref|>text<|/ref|> in the image.", "has_grounding": True},
     "🔍 Describe": {"prompt": "<image>\nDescribe this image in detail.", "has_grounding": False},
     "✏️ Custom": {"prompt": "", "has_grounding": False}
@@ -97,6 +99,8 @@ def clean_output(text, include_images=False):
         else:
             text = re.sub(rf'(?m)^[^\n]*{re.escape(match[0])}[^\n]*\n?', '', text)
     
+    text = text.replace('\\coloneqq', ':=').replace('\\eqqcolon', '=:')
+    
     return text.strip()
 
 def embed_images(markdown, crops):
@@ -109,12 +113,12 @@ def embed_images(markdown, crops):
         markdown = markdown.replace(f'**[Figure {i + 1}]**', f'\n\n![Figure {i + 1}](data:image/png;base64,{b64})\n\n', 1)
     return markdown
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=90)
 def process_image(image, mode, task, custom_prompt):
     if image is None:
-        return " Error Upload image", "", "", None, []
+        return "Error: Upload an image", "", "", None, []
     if task in ["✏️ Custom", "📍 Locate"] and not custom_prompt.strip():
-        return "Enter prompt", "", "", None, []
+        return "Please enter a prompt", "", "", None, []
     
     if image.mode in ('RGBA', 'LA', 'P'):
         image = image.convert('RGB')
@@ -140,8 +144,16 @@ def process_image(image, mode, task, custom_prompt):
     stdout = sys.stdout
     sys.stdout = StringIO()
     
-    model.infer(tokenizer=tokenizer, prompt=prompt, image_file=tmp.name, output_path=out_dir,
-                base_size=config["base_size"], image_size=config["image_size"], crop_mode=config["crop_mode"])
+    model.infer(
+        tokenizer=tokenizer, 
+        prompt=prompt, 
+        image_file=tmp.name, 
+        output_path=out_dir,
+        base_size=config["base_size"], 
+        image_size=config["image_size"], 
+        crop_mode=config["crop_mode"],
+        save_results=False
+    )
     
     result = '\n'.join([l for l in sys.stdout.getvalue().split('\n') 
                         if not any(s in l for s in ['image:', 'other:', 'PATCHES', '====', 'BASE:', '%|', 'torch.Size'])]).strip()
@@ -151,7 +163,7 @@ def process_image(image, mode, task, custom_prompt):
     shutil.rmtree(out_dir, ignore_errors=True)
     
     if not result:
-        return "No text", "", "", None, []
+        return "No text detected", "", "", None, []
     
     cleaned = clean_output(result, False)
     markdown = clean_output(result, True)
@@ -168,7 +180,7 @@ def process_image(image, mode, task, custom_prompt):
     
     return cleaned, markdown, result, img_out, crops
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=90)
 def process_pdf(path, mode, task, custom_prompt, page_num):
     doc = fitz.open(path)
     total_pages = len(doc)
@@ -184,7 +196,7 @@ def process_pdf(path, mode, task, custom_prompt, page_num):
 
 def process_file(path, mode, task, custom_prompt, page_num):
     if not path:
-        return "Error Upload file", "", "", None, []
+        return "Error: Upload a file", "", "", None, []
     if path.lower().endswith('.pdf'):
         return process_pdf(path, mode, task, custom_prompt, page_num)
     else:
@@ -192,9 +204,9 @@ def process_file(path, mode, task, custom_prompt, page_num):
 
 def toggle_prompt(task):
     if task == "✏️ Custom":
-        return gr.update(visible=True, label="Custom Prompt", placeholder="Add <|grounding|> for boxes")
+        return gr.update(visible=True, label="Custom Prompt", placeholder="Add <|grounding|> for bounding boxes")
     elif task == "📍 Locate":
-        return gr.update(visible=True, label="Text to Locate", placeholder="Enter text")
+        return gr.update(visible=True, label="Text to Locate", placeholder="Enter text to locate")
     return gr.update(visible=False)
 
 def select_boxes(task):
@@ -233,12 +245,12 @@ def update_page_selector(file_path):
                         label=f"Select Page (1-{page_count})")
     return gr.update(visible=False)
 
-with gr.Blocks(title="DeepSeek-OCR") as demo:
+with gr.Blocks(title="DeepSeek-OCR-2") as demo:
     gr.Markdown("""
-    # 🚀 DeepSeek-OCR Demo
-    **Convert documents to markdown, extract raw text, and locate specific content with bounding boxes. It takes 20~ sec for markdown and 3~ sec for locate task examples. Check the info at the bottom of the page for more information.**
+    # 🚀 DeepSeek-OCR-2 Demo
+    **Convert documents to markdown, extract text, parse figures, and locate specific content with bounding boxes.** 
     
-    **Hope this tool was helpful! If so, a quick like ❤️ would mean a lot :)**
+    **If this tool was helpful, please consider giving it a like ❤️!**
     """)
     
     with gr.Row():
@@ -246,7 +258,7 @@ with gr.Blocks(title="DeepSeek-OCR") as demo:
             file_in = gr.File(label="Upload Image or PDF", file_types=["image", ".pdf"], type="filepath")
             input_img = gr.Image(label="Input Image", type="pil", height=300)
             page_selector = gr.Number(label="Select Page", value=1, minimum=1, step=1, visible=False)
-            mode = gr.Dropdown(list(MODEL_CONFIGS.keys()), value="Gundam", label="Mode")
+            mode = gr.Dropdown(list(MODEL_CONFIGS.keys()), value="Default", label="Mode")
             task = gr.Dropdown(list(TASK_PROMPTS.keys()), value="📋 Markdown", label="Task")
             prompt = gr.Textbox(label="Prompt", lines=2, visible=False)
             btn = gr.Button("Extract", variant="primary", size="lg")
@@ -266,8 +278,8 @@ with gr.Blocks(title="DeepSeek-OCR") as demo:
     
     gr.Examples(
         examples=[
-            ["examples/ocr.jpg", "Gundam", "📋 Markdown", ""],
-            ["examples/reachy-mini.jpg", "Gundam", "📍 Locate", "Robot"]
+            ["examples/ocr.jpg", "Default", "📋 Markdown", ""],
+            ["examples/reachy-mini.jpg", "Default", "📍 Locate", "Robot"]
         ],
         inputs=[input_img, mode, task, prompt],
         cache_examples=False
@@ -276,18 +288,28 @@ with gr.Blocks(title="DeepSeek-OCR") as demo:
     with gr.Accordion("ℹ️ Info", open=False):
         gr.Markdown("""
         ### Modes
-        - **Gundam**: 1024 base + 640 tiles with cropping - Best balance
-        - **Tiny**: 512×512, no crop - Fastest
-        - **Small**: 640×640, no crop - Quick
-        - **Base**: 1024×1024, no crop - Standard
-        - **Large**: 1280×1280, no crop - Highest quality
+        - **Default**: 1024 base + 768 tiles with cropping - Recommended for most use cases
+        - **Quality**: 1280 base + 960 tiles with cropping - Higher quality, slower
+        - **Fast**: 1024 base + 640 tiles with cropping - Faster processing
+        - **No Crop**: 1024 base + 768 tiles without cropping - Single image processing
+        - **Small**: 768 base + 512 tiles without cropping - Fastest, lower quality
         
         ### Tasks
-        - **Markdown**: Convert document to structured markdown (grounding ✅)
-        - **Free OCR**: Simple text extraction
-        - **Locate**: Find specific things in image (grounding ✅)
+        - **Markdown**: Convert document to structured markdown with layout detection (grounding ✅)
+        - **Free OCR**: Simple text extraction without layout
+        - **OCR Image**: OCR for general images with grounding (grounding ✅)
+        - **Parse Figure**: Parse figures and charts in documents
+        - **Locate**: Find and highlight specific text/elements in image (grounding ✅)
         - **Describe**: General image description
-        - **Custom**: Your own prompt (add `<|grounding|>` for boxes)
+        - **Custom**: Your own prompt (add `<|grounding|>` for bounding boxes)
+        
+        Document: <image>\\n<|grounding|>Convert the document to markdown.
+        Free OCR: <image>\\nFree OCR.
+        Other Image: <image>\\n<|grounding|>OCR this image.
+        Parse Figure: <image>\\nParse the figure.
+        Describe: <image>\\nDescribe this image in detail.
+        Locate: <image>\\nLocate <|ref|>text<|/ref|> in the image.
+        ```
         """)
     
     file_in.change(load_image, [file_in, page_selector], [input_img])
@@ -301,7 +323,7 @@ with gr.Blocks(title="DeepSeek-OCR") as demo:
             return process_file(file_path, mode, task, custom_prompt, int(page_num))
         if image is not None:
             return process_image(image, mode, task, custom_prompt)
-        return "Error uploading file or image", "", "", None, []
+        return "Error: Upload a file or image", "", "", None, []
 
     submit_event = btn.click(run, [input_img, file_in, mode, task, prompt, page_selector],
                              [text_out, md_out, raw_out, img_out, gallery])