Update app.py

app.py

@@ -48,45 +48,121 @@ async def lifespan(app: FastAPI):
 app = FastAPI(title="DeepSeek-OCR-2 API", version="2.0.0", lifespan=lifespan)
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
+# ─── CPU monkey-patch context manager ────────────────────────────────────────
+from contextlib import contextmanager
+
+@contextmanager
+def force_cpu():
+    """
+    DeepSeek-OCR-2's model.infer() hardcodes .cuda() even when no GPU is present.
+    This context manager temporarily replaces all CUDA-moving calls with no-ops
+    so the model runs on CPU without modification.
+    """
+    # Save originals
+    _tensor_cuda   = torch.Tensor.cuda
+    _module_cuda   = torch.nn.Module.cuda
+    _tensor_to     = torch.Tensor.to
+    _module_to     = torch.nn.Module.to
+
+    # Tensor.cuda() → return self (stay on CPU)
+    def _noop_tensor_cuda(self, device=None, *args, **kwargs):
+        return self
+
+    # Module.cuda() → return self
+    def _noop_module_cuda(self, device=None):
+        return self
+
+    # Tensor.to("cuda") / to(device) → stay on CPU; allow dtype casts
+    def _safe_tensor_to(self, *args, **kwargs):
+        filtered = [
+            a for a in args
+            if not (isinstance(a, (str, torch.device)) and "cuda" in str(a))
+        ]
+        kwargs.pop("device", None)
+        if filtered or kwargs:
+            try:
+                return _tensor_to(self, *filtered, **kwargs)
+            except Exception:
+                return self
+        return self
+
+    # Module.to("cuda") → stay on CPU; allow dtype casts
+    def _safe_module_to(self, *args, **kwargs):
+        filtered = [
+            a for a in args
+            if not (isinstance(a, (str, torch.device)) and "cuda" in str(a))
+        ]
+        kwargs.pop("device", None)
+        if filtered or kwargs:
+            try:
+                return _module_to(self, *filtered, **kwargs)
+            except Exception:
+                return self
+        return self
+
+    torch.Tensor.cuda     = _noop_tensor_cuda
+    torch.nn.Module.cuda  = _noop_module_cuda
+    torch.Tensor.to       = _safe_tensor_to
+    torch.nn.Module.to    = _safe_module_to
+
+    try:
+        yield
+    finally:
+        torch.Tensor.cuda    = _tensor_cuda
+        torch.nn.Module.cuda = _module_cuda
+        torch.Tensor.to      = _tensor_to
+        torch.nn.Module.to   = _module_to
+
+
 # ─── Core OCR inference ───────────────────────────────────────────────────────
 def run_ocr(pil_image: Image.Image, mode: str = "free") -> str:
-    """Run DeepSeek-OCR-2 on a PIL image, return extracted text."""
+    """
+    Run DeepSeek-OCR-2 on a PIL image and return extracted text.
+    Works on both CPU (HF free tier) and GPU.
+    """
     prompt_text = (
         "Convert the document to markdown."
         if mode == "markdown"
         else "Please OCR the image and return all text exactly."
     )
+
     with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
         tmp_path = tmp.name
         pil_image.save(tmp_path, format="PNG")
+
     try:
         if hasattr(model, "infer"):
             with tempfile.TemporaryDirectory() as out_dir:
-                result = model.infer(
-                    tokenizer,
-                    prompt=f"<image>\n{prompt_text}",
-                    image_file=tmp_path,
-                    output_path=out_dir,
-                    base_size=1024,
-                    image_size=768,
-                    crop_mode=True,
-                    save_results=False,
-                )
+                # force_cpu() patches .cuda() → no-op so model.infer() works on CPU
+                with force_cpu():
+                    result = model.infer(
+                        tokenizer,
+                        prompt=f"<image>\n{prompt_text}",
+                        image_file=tmp_path,
+                        output_path=out_dir,
+                        base_size=1024,
+                        image_size=768,
+                        crop_mode=True,
+                        save_results=False,
+                    )
             if isinstance(result, dict):
                 return result.get("text", str(result))
             return str(result) if result else ""
 
-        # fallback: standard generate()
+        # ── Fallback: standard generate() if model.infer() is not available ──
         messages = [{"role": "user", "content": [
             {"type": "image", "image": tmp_path},
             {"type": "text",  "text": prompt_text},
         ]}]
-        text_in = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs  = tokenizer(text_in, return_tensors="pt")
+        text_in = tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        inputs = tokenizer(text_in, return_tensors="pt")
         with torch.no_grad():
             out = model.generate(**inputs, max_new_tokens=1024, do_sample=False)
         new_ids = out[:, inputs["input_ids"].shape[1]:]
         return tokenizer.decode(new_ids[0], skip_special_tokens=True)
+
     finally:
         os.unlink(tmp_path)
 
@@ -816,4 +892,4 @@ function cp(id){
 function setProgress(pct){document.getElementById('prog').style.width=pct+'%';}
 </script>
 </body>
-</html>"""
+</html>"""