arnabkar101/dotsocr-space

1import os2import tempfile3import gradio as gr4import pymupdf  # <- PyMuPDF (was: import fitz)5import torch6from transformers import AutoModelForCausalLM, AutoProcessor7from qwen_vl_utils import process_vision_info8 9MODEL_ID = "rednote-hilab/dots.ocr"  # or your fork/fine-tune10 11 12def load_model(model_id: str):13    """Load model/processor while gracefully falling back if flash attention is unavailable."""14    try:15        return (16            AutoModelForCausalLM.from_pretrained(17                model_id,18                trust_remote_code=True,19                attn_implementation="flash_attention_2",20                torch_dtype=torch.bfloat16,21                device_map="auto",22            ),23            AutoProcessor.from_pretrained(model_id, trust_remote_code=True),24        )25    except (ValueError, RuntimeError) as exc:26        # Flash attention wheels might be unavailable on CPU/initial builds; retry with default attention27        print(f"Falling back to default attention due to: {exc}")28        return (29            AutoModelForCausalLM.from_pretrained(30                model_id,31                trust_remote_code=True,32                torch_dtype=torch.bfloat16,33                device_map="auto",34            ),35            AutoProcessor.from_pretrained(model_id, trust_remote_code=True),36        )37 38 39model, processor = load_model(MODEL_ID)40 41PROMPT = """... (your layout JSON prompt) ..."""42 43 44def pdf_to_pngs(pdf_path: str, dpi: int = 200) -> list[str]:45    """46    Render each PDF page to a PNG file using PyMuPDF (pymupdf).47 48    Returns a list of temp file paths (1 per page).49    """50    doc = pymupdf.open(pdf_path)51    try:52        out_paths: list[str] = []53        zoom = dpi / 72.054        mat = pymupdf.Matrix(zoom, zoom)55        for i in range(doc.page_count):56            page = doc.load_page(i)57            pix = page.get_pixmap(matrix=mat)58            tmp_path = os.path.join(tempfile.gettempdir(), f"page_{i}.png")59            pix.save(tmp_path)60            out_paths.append(tmp_path)61        return out_paths62    finally:63        doc.close()64 65 66def run(pdf_or_img, prompt):67    # Accept PDF or image inputs and expand PDFs into per-page PNGs.68    paths = [pdf_or_img]69    if str(pdf_or_img).lower().endswith(".pdf"):70        paths = pdf_to_pngs(pdf_or_img)71 72    results = []73    for img in paths:74        messages = [75            {76                "role": "user",77                "content": [78                    {"type": "image", "image": img},79                    {"type": "text", "text": prompt or PROMPT},80                ],81            }82        ]83 84        # Build chat template text and collect vision inputs85        text = processor.apply_chat_template(86            messages,87            tokenize=False,88            add_generation_prompt=True,89        )90        imgs, vids = process_vision_info(messages)91 92        inputs = processor(93            text=[text],94            images=imgs,95            videos=vids,96            padding=True,97            return_tensors="pt",98        )99 100        # Move to available accelerator101        if torch.cuda.is_available():102            inputs = {k: (v.cuda() if hasattr(v, "cuda") else v) for k, v in inputs.items()}103        elif torch.backends.mps.is_available():104            # Optional: support Apple Silicon MPS105            inputs = {k: (v.to("mps") if hasattr(v, "to") and isinstance(v, torch.Tensor) else v) for k, v in inputs.items()}106 107        with torch.inference_mode():108            out_ids = model.generate(**inputs, max_new_tokens=24000)109            # Trim the prompt tokens110            trimmed = [o[len(i):] for i, o in zip(inputs["input_ids"], out_ids)]111            text_out = processor.batch_decode(112                trimmed,113                skip_special_tokens=True,114                clean_up_tokenization_spaces=False,115            )[0]116 117        results.append(text_out)118 119    return results if len(results) > 1 else results[0]120 121 122demo = gr.Interface(123    fn=run,124    inputs=[gr.File(label="PDF or image"), gr.Textbox(lines=6, value=PROMPT)],125    outputs=gr.JSON(),126    title="DotsOCR PDF/Layout Parser",127)128 129if __name__ == "__main__":130    demo.launch()131