Upload ocr_service.py

ocr_service.py

@@ -133,7 +133,9 @@ def _download_and_patch_model_locally(model_id: str, revision: str) -> str:
     Download DeepSeek-OCR to a local dir, patch for CPU:
       - remove hardcoded .cuda()
       - force float32 (strip .bfloat16() / .to(torch.bfloat16))
-      - disable torch.autocast("cuda", ...) blocks
+    
+    Minimal patcher that avoids indentation issues by NOT touching autocast blocks.
+    On CPU, torch.autocast is auto-disabled anyway, so we leave it alone.
     
     Return local path for from_pretrained(...).
     
@@ -153,71 +155,53 @@ def _download_and_patch_model_locally(model_id: str, revision: str) -> str:
     print(f"  ✅ Downloaded to: {local_dir}")
     local_dir = Path(local_dir)
     
-    def ensure_imports(src: str) -> str:
-        """Ensure torch and contextlib are imported"""
-        if "import torch" not in src and "from torch" not in src:
-            src = "import torch\n" + src
-        if "import contextlib" not in src:
-            # Add after torch import
-            if "import torch" in src:
-                src = src.replace("import torch", "import torch\nimport contextlib")
-            elif "from torch" in src:
-                src = src.replace("from torch", "import contextlib\nfrom torch")
-            else:
-                # Add at top if no torch import
-                lines = src.split('\n')
-                for i, line in enumerate(lines):
-                    if line.startswith("import ") or line.startswith("from "):
-                        continue
-                    else:
-                        lines.insert(i, "import contextlib")
-                        break
-                src = '\n'.join(lines)
-        return src
-    
-    def patch_text(src: str) -> str:
-        """Patch text to remove CUDA/BF16 and force FP32"""
-        # A) Remove hardcoded CUDA device moves
-        src = src.replace(".unsqueeze(-1).cuda()", ".unsqueeze(-1)")
-        src = src.replace("input_ids.unsqueeze(0).cuda()", "input_ids.unsqueeze(0)")
-        src = src.replace("(images_crop.cuda(), images_ori.cuda())", "(images_crop, images_ori)")
-        src = src.replace("images_seq_mask = images_seq_mask.unsqueeze(0).cuda()", 
+    def patch_file(p: Path):
+        """Minimal patch - only string replacements, no indentation changes"""
+        txt0 = p.read_text(encoding="utf-8")
+        txt = txt0
+        
+        # A) Remove hardcoded CUDA device moves (CPU-safe)
+        txt = txt.replace(".unsqueeze(-1).cuda()", ".unsqueeze(-1)")
+        txt = txt.replace("input_ids.unsqueeze(0).cuda()", "input_ids.unsqueeze(0)")
+        txt = txt.replace("(images_crop.cuda(), images_ori.cuda())", "(images_crop, images_ori)")
+        txt = txt.replace("images_seq_mask = images_seq_mask.unsqueeze(0).cuda()", 
                           "images_seq_mask = images_seq_mask.unsqueeze(0)")
-        src = src.replace("input_ids.unsqueeze(0).cuda().shape[1]", 
+        txt = txt.replace("input_ids.unsqueeze(0).cuda().shape[1]", 
                           "input_ids.unsqueeze(0).shape[1]")
         
-        # B) Force FP32 (strip BF16 casts)
-        src = re.sub(r"\.bfloat16\(\)", ".float()", src)
-        src = re.sub(r"\.to\(\s*torch\.bfloat16\s*\)", ".to(torch.float32)", src)
-        src = re.sub(r"\.to\(\s*dtype\s*=\s*torch\.bfloat16\s*\)", ".to(dtype=torch.float32)", src)
+        # B) Force FP32 (fix BF16 vs FP32), pattern-safe (no newlines/indentation changes)
+        txt = re.sub(r"\.bfloat16\(\)", ".float()", txt)
+        txt = re.sub(r"\.to\(\s*torch\.bfloat16\s*\)", ".to(torch.float32)", txt)
+        txt = re.sub(r"\.to\(\s*dtype\s*=\s*torch\.bfloat16\s*\)", ".to(dtype=torch.float32)", txt)
         
-        # C) Disable CUDA autocast blocks
-        src = ensure_imports(src)
-        # Match both torch.autocast("cuda", ...) and torch.autocast(device_type="cuda", ...)
-        src = re.sub(r'torch\.autocast\(\s*["\']cuda["\']\s*,[^)]*\)', 
-                     "contextlib.nullcontext()", src)
-        src = re.sub(r'torch\.autocast\(\s*device_type\s*=\s*["\']cuda["\'][^)]*\)', 
-                     "contextlib.nullcontext()", src)
+        # Note: We do NOT touch torch.autocast() blocks - on CPU they're auto-disabled
+        # and modifying them risks breaking indentation/syntax
         
-        return src
+        if txt != txt0:
+            p.write_text(txt, encoding="utf-8")
+            print(f"  ✅ Patched CPU/FP32: {p.name}")
+        else:
+            print(f"  ℹ️ Already CPU/FP32-safe: {p.name}")
     
     # Patch both files where they may appear
     targets = list(local_dir.rglob("modeling_deepseekocr.py")) + \
               list(local_dir.rglob("deepencoder.py"))
     
     if not targets:
-        raise RuntimeError("Could not locate DeepSeek-OCR source files to patch")
+        raise RuntimeError("Could not find DeepSeek-OCR source files to patch")
     
-    for p in targets:
-        print(f"  🔍 Found file: {p.name}")
-        txt = p.read_text(encoding="utf-8")
-        new = patch_text(txt)
-        
-        if new != txt:
-            p.write_text(new, encoding="utf-8")
-            print(f"  ✅ Patched for CPU+FP32: {p.name}")
-        else:
-            print(f"  ℹ️  Already CPU-patched: {p.name}")
+    for f in targets:
+        print(f"  🔍 Found file: {f.name}")
+        patch_file(f)
+    
+    # Optional: compile check to catch syntax errors early
+    try:
+        import py_compile
+        for f in targets:
+            py_compile.compile(str(f), doraise=True)
+        print(f"  ✅ Syntax check passed for {len(targets)} file(s)")
+    except py_compile.PyCompileError as e:
+        raise RuntimeError(f"Syntax check failed after patch: {e}")
     
     return str(local_dir)