!pip install -q "transformers==4.46.3" "tokenizers==0.20.3" einops addict easydict pillow gradio

import os, time, torch, gradio as gr
from PIL import Image
from transformers import AutoModel, AutoTokenizer
from datetime import datetime
import random
import string

model_id = "deepseek-ai/DeepSeek-OCR"
tok = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
if tok.pad_token is None and tok.eos_token is not None:
    tok.pad_token = tok.eos_token
model = AutoModel.from_pretrained(
    model_id,
    trust_remote_code=True,
    use_safetensors=True,
    attn_implementation="eager"
).to(dtype=torch.bfloat16, device="cuda").eval()

def new_run_dir(base="/content/runs"):
    os.makedirs(base, exist_ok=True)
    ts = datetime.now().strftime("%Y%m%d-%H%M%S")
    rid = ''.join(random.choices(string.ascii_lowercase + string.digits, k=5))
    path = os.path.join(base, f"run_{ts}_{rid}")
    os.makedirs(path)
    return path

def gr_ocr(image, mode, custom_prompt, base_size, image_size, crop_mode):
    img = image.convert("RGB")
    if max(img.size) > 2000:
        s = 2000 / max(img.size)
        img = img.resize((int(img.width*s), int(img.height*s)), Image.LANCZOS)
    
    run_dir = new_run_dir()
    img_path_proc = os.path.join(run_dir, "input.png")
    img.save(img_path_proc, optimize=True)

    if mode == "Custom Prompt" and custom_prompt.strip():
        prompt = custom_prompt.strip()
    else:
        prompt = DEMO_MODES[mode]["prompt"]
    
    t0 = time.time()
    try:
        with torch.inference_mode():
            _ = model.infer(
                tok,
                prompt=prompt,
                image_file=img_path_proc,
                output_path=run_dir,
                base_size=base_size,
                image_size=image_size,
                crop_mode=crop_mode,
                save_results=True,
                test_compress=True
            )
    except ZeroDivisionError:
        print(" [Patched] Division by zero in compression ratio (valid_img_tokens==0). Ignored.")
    
    dt = time.time() - t0
    
    result_file = os.path.join(run_dir, "result.mmd")
    if not os.path.exists(result_file):
        result_file = os.path.join(run_dir, "result.txt")
    
    result = "[No text extracted]"
    if os.path.exists(result_file):
        with open(result_file, "r", encoding="utf-8") as f:
            result = f.read().strip() or "[No text extracted]"
            
    boxed_path = os.path.join(run_dir, "result_with_boxes.jpg")
    boxed_img = Image.open(boxed_path) if os.path.exists(boxed_path) else None
    
    stats = f"""
**{dt:.1f}s** | Image: {img.size[0]}×{img.size[1]} px  
**Output directory:** {run_dir}  
"""
    return result, stats, boxed_img

# (此处省略 DEMO_MODES 字典的定义，原文太长，功能是预设不同的 prompt)
DEMO_MODES = {
    "Document ➔ Markdown": {
        "prompt": "<image>\n<|grounding|>Convert the document to markdown.",
        "desc": "Extracts full document text as Markdown, preserving structure (headings, tables, lists, etc.).",
        "base_size": 1024, "image_size": 640, "crop_mode": False
    },
    "Chart Deep Parsing": {
        "prompt": "<image>\nParse all charts and tables. Extract data as HTML tables.",
        "desc": "Extracts tabular/chart data into HTML tables.",
        "base_size": 1024, "image_size": 640, "crop_mode": False
    },
    # ... 其他模式
    "Custom Prompt": {
        "prompt": "",
        "desc": "Provide your own custom prompt for flexible OCR or parsing.",
        "base_size": 1024, "image_size": 640, "crop_mode": False
    }
}


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    <div style="text-align:center;">
        <h1>DeepSeek-OCR Demo</h1>
    </div>
    """)
    with gr.Row():
        with gr.Column():
            # ... 此处省略左侧输入组件（图片上传、模式选择等）的 Gradio 代码
            image_input = gr.Image(type="pil", label="Upload Document/Image", height=350)
            mode = gr.Radio(choices=list(DEMO_MODES.keys()), value="Document ➔ Markdown", label="Select Capability to Test")
            desc = gr.Markdown(DEMO_MODES["Document ➔ Markdown"]["desc"])
            custom_prompt = gr.Textbox(label="Custom Prompt", visible=False)
            process_btn = gr.Button("Process Image", variant="primary")
        with gr.Column():
            # ... 此处省略右侧输出组件（文本框、状态、带框图片）的 Gradio 代码
            ocr_output = gr.Textbox(label="Extracted Content", lines=22, show_copy_button=True)
            status_out = gr.Markdown("_Process an image to see stats and output dir_")
            boxed_output = gr.Image(label="Result with Bounding Boxes", type="pil")

    def update_mode(selected):
        d = DEMO_MODES[selected]
        return d["desc"], gr.update(visible=selected=="Custom Prompt"), d["base_size"], d["image_size"], d["crop_mode"]
    
    mode.change(update_mode, inputs=mode, outputs=[desc, custom_prompt, base_size, image_size, crop_mode])
    
    process_btn.click(
        gr_ocr,
        inputs=[image_input, mode, custom_prompt, base_size, image_size, crop_mode],
        outputs=[ocr_output, status_out, boxed_output]
    )

demo.launch(share=True, debug=True)