Non-CUDA inference: CPU works after removing hardcoded .cuda(); MPS (Apple Silicon) crashes mid-generation

Title

Non-CUDA inference: CPU works after removing hardcoded .cuda(); MPS (Apple Silicon) crashes mid-generation

Body

Context

All three documented inference paths (Transformers, vLLM, SGLang) assume an NVIDIA CUDA device. modeling_unlimitedocr.py (the trust_remote_code file served from the Hugging Face repo) has ~17 hardcoded .cuda() calls and 3 torch.autocast("cuda", dtype=torch.bfloat16) blocks inside UnlimitedOCRForCausalLM.infer() / .infer_multi() (roughly lines 582, 1003–1070, 1238–1259 in the current revision).

Following the README's Transformers example on any machine without CUDA — e.g. an Apple Silicon Mac — fails immediately with:

AssertionError: Torch not compiled with CUDA enabled

We tested two non-CUDA paths on an Apple M1 Pro (16GB RAM, macOS) after patching every .cuda() → .to(self.device) and making the autocast device type follow self.device.type instead of being hardcoded:

1) CPU — works correctly

model.infer(...) ran end-to-end on CPU against a real scanned document page and produced correct output (titles, product codes, spec text all transcribed correctly with layout <|det|> tags matching the source). This confirms the forward pass itself has no CUDA-only kernels — no custom Triton/CUDA ops in the MoE routing, flash_attn is already conditionally imported behind is_flash_attn_2_available(), and the vision encoder defaults use_flash_attn=False.

Performance context (not a complaint — just so the tradeoff is clear): ~550s for a single dense page (~2,700 image/prefix tokens, ~500 output tokens, ~1 token/s decode). Understandably far from GPU/vLLM/SGLang throughput; this is about making CPU execution possible for local dev, not fast.

2) MPS (Apple Silicon GPU) — loads, then crashes during generation

We also tried moving the model to mps (model.eval().to("mps")) instead of CPU. Loading succeeded (model.device correctly reports mps:0, ~22s). However, calling model.infer(...) crashes partway through generation with no Python-catchable exception — no traceback, just the process dying mid-run (consistent with an unsupported MPS kernel hitting a hard abort rather than raising a normal error). We couldn't isolate which specific op triggers it without deeper access to the model internals (likely somewhere in the SAM/CLIP vision encoder's attention, or in the DeepSeekV2 MoE expert routing/gather — both have non-trivial indexing patterns that sometimes lack MPS kernel coverage in PyTorch).

What we're asking

1. Portability fix (small, we can PR this): accept a small, backward-compatible patch replacing hardcoded .cuda() with .to(self.device) and making autocast device-aware, so CPU inference at least works instead of hard-crashing on any non-CUDA machine. Zero behavior change for existing CUDA users.
2. Apple Silicon / MPS support (needs your side): since MPS crashes without a catchable exception, someone with deeper familiarity with the vision encoder and MoE routing internals would need to trace which op lacks MPS coverage. We don't have enough visibility to pinpoint the exact failure ourselves, but wanted to report it in case it's useful — full Apple Silicon GPU support (not just CPU fallback) would make local development/testing on Mac hardware genuinely usable rather than just "not crashing but very slow."

Happy to send a PR for (1) if that's welcome — it's a small, mechanical change. For (2), we're reporting what we observed in case it helps prioritize, but the actual fix is likely beyond what we can diagnose from the outside.