diff --git a/README.md b/README.md index a7ff052..5aceac5 100644 --- a/README.md +++ b/README.md @@ -307,3 +307,21 @@ We would like to thank [Deepseek-OCR](https://github.com/deepseek-ai/DeepSeek-OC primaryClass={cs.CV}, url={https://arxiv.org/abs/2606.23050}, } + +## Local inference with Transformers (CUDA / Apple Silicon / CPU) + +The SGLang path above requires CUDA. To run the model locally without a CUDA +GPU (for example on an Apple Silicon Mac), use the Transformers-based path: + +```shell +hf download baidu/Unlimited-OCR --local-dir ./Unlimited-OCR-local +python patch_model_for_local.py ./Unlimited-OCR-local +python infer_transformers.py --model_dir ./Unlimited-OCR-local \ + --image_dir ./my_pages --output_dir ./outputs +``` + +The patch step is required: the released modeling file hardcodes CUDA calls, +and on the MPS backend the image-embedding injection must use positional +assignment instead of `masked_scatter_`, which silently corrupts the visual +tokens there (the model then returns empty output with no error). The patch +asserts on the exact released code and is a behavioral no-op on CUDA. diff --git a/infer_transformers.py b/infer_transformers.py new file mode 100644 index 0000000..8701565 --- /dev/null +++ b/infer_transformers.py @@ -0,0 +1,173 @@ +"""Local single-process inference with Hugging Face Transformers. + +The SGLang path in ``infer.py`` requires CUDA. This script runs the model +directly through ``transformers`` on CUDA, Apple Silicon (MPS), or CPU, one +image at a time. It expects a local snapshot that has been prepared with +``patch_model_for_local.py`` (which removes the hardcoded CUDA calls and fixes +a silent image-embedding corruption on MPS; see that script's docstring). + +Usage: + hf download baidu/Unlimited-OCR --local-dir ./Unlimited-OCR-local + python patch_model_for_local.py ./Unlimited-OCR-local + python infer_transformers.py --model_dir ./Unlimited-OCR-local \ + --image_dir ./my_pages --output_dir ./outputs + +Two input modes mirror ``infer.py``: ``--image_dir`` sends every image found +under a directory, ``--pdf`` converts each page of a PDF first. +""" + +import argparse +import os +import pathlib +import tempfile +import time + +PROMPT = "document parsing." +NO_REPEAT_NGRAM_SIZE = 35 +NGRAM_WINDOW = 128 +PDF_DPI = 300 +IMAGE_MODES = { + # matches the gundam / base presets used by the SGLang path + "gundam": {"base_size": 1024, "image_size": 640, "crop_mode": True}, + "base": {"base_size": 1024, "image_size": 1024, "crop_mode": False}, +} + + +def pick_device(requested: str) -> str: + import torch + + if requested != "auto": + return requested + if torch.cuda.is_available(): + return "cuda" + if torch.backends.mps.is_available(): + return "mps" + return "cpu" + + +def check_patched(model_dir: str) -> None: + modeling = pathlib.Path(model_dir) / "modeling_unlimitedocr.py" + if modeling.exists() and ".cuda()" in modeling.read_text(encoding="utf-8"): + raise SystemExit( + f"{modeling} still contains hardcoded CUDA calls. Run\n" + f" python patch_model_for_local.py {model_dir}\n" + "first (required for MPS/CPU; a no-op change on CUDA)." + ) + + +def pdf_to_images(pdf_path: str, dpi: int = PDF_DPI) -> list: + import fitz + + doc = fitz.open(pdf_path) + tmp_dir = tempfile.mkdtemp(prefix="pdf_ocr_") + image_paths = [] + mat = fitz.Matrix(dpi / 72, dpi / 72) + for i, page in enumerate(doc): + out_path = os.path.join(tmp_dir, f"page_{i + 1:04d}.png") + page.get_pixmap(matrix=mat).save(out_path) + image_paths.append(out_path) + doc.close() + return image_paths + + +def collect_images(image_dir: str) -> list: + exts = (".png", ".jpg", ".jpeg", ".webp", ".bmp") + image_files = [] + for root, _, files in os.walk(image_dir): + for name in files: + if name.lower().endswith(exts): + image_files.append(os.path.join(root, name)) + return sorted(image_files) + + +def load_model(model_dir: str, device: str): + import torch + from transformers import AutoModel, AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True) + model = AutoModel.from_pretrained( + model_dir, + trust_remote_code=True, + use_safetensors=True, + torch_dtype=torch.bfloat16, + attn_implementation="eager", + low_cpu_mem_usage=True, + ) + return tokenizer, model.eval().to(device) + + +def run(args) -> None: + device = pick_device(args.device) + check_patched(args.model_dir) + if device == "mps": + os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") + + if args.pdf: + images = pdf_to_images(args.pdf) + elif args.image_dir: + images = collect_images(args.image_dir) + else: + raise SystemExit("either --image_dir or --pdf is required") + if not images: + raise SystemExit("no input images found") + + os.makedirs(args.output_dir, exist_ok=True) + print(f"device={device}, images={len(images)}, image_mode={args.image_mode}") + + tokenizer, model = load_model(args.model_dir, device) + mode = IMAGE_MODES[args.image_mode] + + import torch + + total_tokens = 0 + wall_start = time.time() + for i, image_path in enumerate(images, 1): + name = os.path.splitext(os.path.basename(image_path))[0] + t0 = time.time() + with torch.no_grad(): + text = model.infer( + tokenizer, + prompt=PROMPT, + image_file=image_path, + output_path=args.output_dir, + base_size=mode["base_size"], + image_size=mode["image_size"], + crop_mode=mode["crop_mode"], + no_repeat_ngram_size=NO_REPEAT_NGRAM_SIZE, + ngram_window=NGRAM_WINDOW, + max_length=args.max_length, + save_results=False, + eval_mode=True, + ) + elapsed = time.time() - t0 + n_tokens = len(tokenizer(text, add_special_tokens=False)["input_ids"]) + total_tokens += n_tokens + out_file = os.path.join(args.output_dir, f"{name}.md") + with open(out_file, "w", encoding="utf-8") as f: + f.write(text) + print(f" [{i}/{len(images)}] {name}: {n_tokens} tokens, " + f"{elapsed:.1f}s ({n_tokens / max(elapsed, 1e-6):.1f} tok/s)") + + wall = time.time() - wall_start + print(f"done: {len(images)} image(s), {total_tokens} tokens, {wall:.1f}s") + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Local Transformers inference (CUDA / Apple Silicon / CPU).", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--model_dir", required=True, + help="Local snapshot prepared by patch_model_for_local.py") + parser.add_argument("--image_dir", default="", help="Directory of images") + parser.add_argument("--pdf", default="", help="PDF file to convert and parse") + parser.add_argument("--output_dir", default="./outputs") + parser.add_argument("--image_mode", choices=tuple(IMAGE_MODES), default="gundam") + parser.add_argument("--device", choices=("auto", "cuda", "mps", "cpu"), + default="auto") + parser.add_argument("--max_length", type=int, default=8192) + return parser.parse_args() + + +if __name__ == "__main__": + run(parse_args()) diff --git a/patch_model_for_local.py b/patch_model_for_local.py new file mode 100644 index 0000000..f72cc93 --- /dev/null +++ b/patch_model_for_local.py @@ -0,0 +1,106 @@ +"""Patch a local baidu/Unlimited-OCR snapshot for device-agnostic inference. + +The released ``modeling_unlimitedocr.py`` hardcodes CUDA (``.cuda()`` calls and +``torch.autocast("cuda", ...)``), so the model cannot run on Apple Silicon (MPS) +or CPU at all. Replacing ``.cuda()`` with ``.to(device)`` is NOT sufficient: on +the MPS backend ``torch.Tensor.masked_scatter_`` silently mis-scatters when its +mask is a stride-0 broadcast view (which the ``.unsqueeze(-1)`` at the injection +call produces) or when its source tensor is non-contiguous. Either trigger +scrambles the injected image embeddings and makes the model emit an immediate +end-of-sequence token (empty output, no error); the broadcast mask is the one +live at this call site, since the source comes from ``torch.cat`` and is +contiguous. This script therefore applies two kinds of edits: + + 1. Replace the image-embedding injection with explicit positional assignment, + which is mathematically identical and correct on CUDA, MPS, and CPU. + 2. Make every hardcoded ``.cuda()`` / ``autocast("cuda")`` follow the model's + actual device (a no-op change on CUDA machines). + +Usage: + hf download baidu/Unlimited-OCR --local-dir ./Unlimited-OCR-local + python patch_model_for_local.py ./Unlimited-OCR-local + +The script asserts on the exact released code before editing, so it fails +loudly instead of mis-patching if the upstream file changes, and it is a no-op +when the snapshot is already patched. +""" + +import pathlib +import sys + +INJECTION_BEFORE = ( + " inputs_embeds[idx].masked_scatter_(" + "images_seq_mask[idx].unsqueeze(-1).cuda(), images_in_this_batch)" +) +INJECTION_AFTER = """\ + # masked_scatter_ silently mis-scatters on the MPS backend + # when its mask is a stride-0 broadcast view (as the + # .unsqueeze(-1) here produced) or its source tensor is + # non-contiguous; explicit positional assignment computes + # the same result and is correct on CUDA, MPS, and CPU. + _feat = images_in_this_batch.to( + device=inputs_embeds.device, dtype=inputs_embeds.dtype + ) + _pos = ( + images_seq_mask[idx] + .to(inputs_embeds.device) + .bool() + .nonzero(as_tuple=True)[0] + ) + inputs_embeds[idx, _pos] = _feat""" + +DEVICE_REPLACEMENTS = [ + # (pattern, replacement, expected occurrences) + (".unsqueeze(0).cuda().shape[1]", ".unsqueeze(0).shape[1]", 4), + ("input_ids.unsqueeze(0).cuda()", "input_ids.unsqueeze(0).to(self.device)", 3), + ("images_seq_mask.unsqueeze(0).cuda()", + "images_seq_mask.unsqueeze(0).to(self.device)", 3), + ("images_crop.cuda()", "images_crop.to(self.device)", 2), + ("images_ori.cuda()", "images_ori.to(self.device)", 3), + ("dummy_crop.cuda()", "dummy_crop.to(self.device)", 1), + ('with torch.autocast("cuda", dtype=torch.bfloat16):', + 'with torch.autocast(self.device.type, dtype=torch.bfloat16, ' + 'enabled=(self.device.type == "cuda")):', 3), +] + + +def patch_file(path: pathlib.Path) -> None: + src = path.read_text(encoding="utf-8") + + if "inputs_embeds[idx, _pos] = _feat" in src and ".cuda()" not in src: + print(f"{path.name}: already patched, nothing to do.") + return + + assert "class UnlimitedOCRForCausalLM" in src, "unexpected file contents" + + assert INJECTION_BEFORE in src, ( + "expected released injection line not found; the upstream file may " + "have changed -- refusing to patch blindly" + ) + src = src.replace(INJECTION_BEFORE, INJECTION_AFTER) + + for pattern, replacement, expected in DEVICE_REPLACEMENTS: + found = src.count(pattern) + assert found == expected, ( + f"expected {expected} occurrence(s) of {pattern!r}, found {found}" + ) + src = src.replace(pattern, replacement) + + assert ".cuda()" not in src, "unreplaced .cuda() call remains" + path.write_text(src, encoding="utf-8") + print(f"{path.name}: patched (1 injection fix, " + f"{sum(n for _, _, n in DEVICE_REPLACEMENTS)} device edits).") + + +def main() -> None: + if len(sys.argv) != 2: + sys.exit("usage: python patch_model_for_local.py ") + model_dir = pathlib.Path(sys.argv[1]) + target = model_dir / "modeling_unlimitedocr.py" + if not target.exists(): + sys.exit(f"{target} not found -- is this an Unlimited-OCR snapshot?") + patch_file(target) + + +if __name__ == "__main__": + main() diff --git a/tests/test_local_patch.py b/tests/test_local_patch.py new file mode 100644 index 0000000..33a6986 --- /dev/null +++ b/tests/test_local_patch.py @@ -0,0 +1,137 @@ +"""Unit tests for the local-inference patch (patch_model_for_local.py). + +Covers the things the patch must guarantee: + 1. The positional-assignment replacement computes exactly what the released + masked_scatter_ injection computes (checked against CPU, where + masked_scatter_ is correct), for contiguous AND non-contiguous sources. + 2. The same equivalence for the injection call's exact real layout: a + [B, T, H] destination written through a batch-index view, the [T, 1] + stride-0 broadcast mask produced by .unsqueeze(-1), and a + torch.cat-built (contiguous) source. + 3. On the MPS backend, the replacement matches the CPU reference. (The + motivation is that raw masked_scatter_ mis-scatters there under two + independent conditions: a stride-0 broadcast mask -- the trigger that + actually fires at the injection call -- or a non-contiguous source. + The tests assert our replacement is right rather than asserting the + upstream op is wrong, so they stay green if PyTorch fixes the + underlying issue.) + +Run: pytest tests/test_local_patch.py +""" + +import torch + + +def _reference_injection(dst, mask, src): + """The released behavior: in-place masked_scatter_ (computed on CPU).""" + out = dst.clone().cpu() + out.masked_scatter_(mask.cpu().unsqueeze(-1), src.cpu().contiguous()) + return out + + +def _patched_injection(dst, mask, src): + """The patched behavior: explicit positional assignment (any device).""" + out = dst.clone() + feat = src.to(device=out.device, dtype=out.dtype) + pos = mask.to(out.device).bool().nonzero(as_tuple=True)[0] + out[pos] = feat + return out + + +def _make_case(seq_len=32, hidden=16, n_img=10, noncontiguous=False, device="cpu"): + torch.manual_seed(0) + dst = torch.randn(seq_len, hidden, device=device) + mask = torch.zeros(seq_len, dtype=torch.bool, device=device) + mask[3:3 + n_img] = True + if noncontiguous: + src = torch.randn(hidden, n_img, device=device).T # transposed view + assert not src.is_contiguous() + else: + src = torch.randn(n_img, hidden, device=device) + return dst, mask, src + + +def test_positional_assignment_matches_masked_scatter_contiguous(): + dst, mask, src = _make_case(noncontiguous=False) + assert torch.equal(_patched_injection(dst, mask, src), + _reference_injection(dst, mask, src)) + + +def test_positional_assignment_matches_masked_scatter_noncontiguous(): + dst, mask, src = _make_case(noncontiguous=True) + assert torch.equal(_patched_injection(dst, mask, src), + _reference_injection(dst, mask, src)) + + +def test_patched_injection_correct_on_mps(): + if not torch.backends.mps.is_available(): + import pytest + pytest.skip("MPS not available") + for noncontiguous in (False, True): + dst, mask, src = _make_case(noncontiguous=noncontiguous, device="mps") + got = _patched_injection(dst, mask, src).cpu() + want = _reference_injection(dst, mask, src) + assert torch.allclose(got, want), ( + f"patched injection wrong on MPS (noncontiguous={noncontiguous})" + ) + + +def _make_real_call_case(batch=2, seq_len=32, hidden=16, device="cpu", + dtype=torch.float32): + """The injection call's exact layout: a [B, T, H] batch tensor written + through a batch-index view, a [B, T] mask whose row the released code + broadcasts via .unsqueeze(-1) (stride 0 across the hidden dim), and + per-row sources built with torch.cat like images_in_this_batch + (contiguous by construction).""" + torch.manual_seed(0) + emb = torch.randn(batch, seq_len, hidden, device=device, dtype=dtype) + mask = torch.zeros(batch, seq_len, dtype=torch.bool, device=device) + mask[0, 3:3 + 6] = True + mask[1, 10:10 + 4] = True + srcs = [ + torch.cat([torch.randn(4, hidden, device=device, dtype=dtype), + torch.randn(2, hidden, device=device, dtype=dtype)], dim=0), + torch.cat([torch.randn(3, hidden, device=device, dtype=dtype), + torch.randn(1, hidden, device=device, dtype=dtype)], dim=0), + ] + for src in srcs: + assert src.is_contiguous() + return emb, mask, srcs + + +def _reference_real_call(emb, mask, srcs): + """The released behavior, computed on CPU: per-row in-place + masked_scatter_ with the stride-0 broadcast mask from .unsqueeze(-1).""" + out = emb.clone().cpu() + for idx, src in enumerate(srcs): + out[idx].masked_scatter_(mask[idx].cpu().unsqueeze(-1), src.cpu()) + return out + + +def _patched_real_call(emb, mask, srcs): + """The patched behavior, exactly as patch_model_for_local.py writes it.""" + out = emb.clone() + for idx, src in enumerate(srcs): + _feat = src.to(device=out.device, dtype=out.dtype) + _pos = mask[idx].to(out.device).bool().nonzero(as_tuple=True)[0] + out[idx, _pos] = _feat + return out + + +def test_positional_assignment_matches_broadcast_mask_layout(): + emb, mask, srcs = _make_real_call_case() + assert torch.equal(_patched_real_call(emb, mask, srcs), + _reference_real_call(emb, mask, srcs)) + + +def test_broadcast_mask_layout_correct_on_mps(): + if not torch.backends.mps.is_available(): + import pytest + pytest.skip("MPS not available") + for dtype in (torch.float32, torch.bfloat16): + emb, mask, srcs = _make_real_call_case(device="mps", dtype=dtype) + got = _patched_real_call(emb, mask, srcs).cpu().float() + want = _reference_real_call(emb, mask, srcs).float() + assert torch.allclose(got, want), ( + f"patched injection wrong on MPS for the real call layout ({dtype})" + )