import asyncio import json import math import os import platform import secrets import tempfile from collections import defaultdict, deque from time import monotonic from typing import Any, Deque, DefaultDict, Optional from pathlib import Path import numpy as np from fastapi import Depends, FastAPI, Form, HTTPException, Request, UploadFile, status from fastapi.middleware.cors import CORSMiddleware from fastapi.security import APIKeyHeader from PIL import Image # Lazy import DeepSeek-OCR dependencies (only load when needed) _torch = None _transformers = None def _get_torch(): global _torch if _torch is None: try: import torch _torch = torch except ImportError: raise RuntimeError( "torch is not installed. Install with: pip install torch" ) return _torch def _get_transformers(): global _transformers if _transformers is None: try: from transformers import AutoModel, AutoTokenizer _transformers = (AutoModel, AutoTokenizer) except ImportError: raise RuntimeError( "transformers is not installed. Install with: pip install transformers" ) return _transformers # Import llm_splitter (works as module or direct import) try: from llm_splitter import call_llm_splitter except ImportError: # Fallback for relative import try: from .llm_splitter import call_llm_splitter except ImportError: # If llm_splitter doesn't exist, define a stub async def call_llm_splitter(*args, **kwargs): raise NotImplementedError("llm_splitter not available") ALLOWED_CONTENT_TYPES = { "image/jpeg", "image/png", "image/webp", } MAX_UPLOAD_BYTES = int(os.getenv("MAX_UPLOAD_BYTES", str(5 * 1024 * 1024))) RATE_LIMIT_REQUESTS = int(os.getenv("RATE_LIMIT_REQUESTS", "30")) RATE_LIMIT_WINDOW_SECONDS = float(os.getenv("RATE_LIMIT_WINDOW_SECONDS", "60")) # Allow API key to be optional for development (security risk in production!) SERVICE_API_KEY = os.getenv("SERVICE_API_KEY", "dev-key-change-in-production") REQUIRE_API_KEY = os.getenv("REQUIRE_API_KEY", "false").lower() == "true" API_KEY_HEADER_NAME = "X-API-Key" MAX_CHILD_LINES = 500 MAX_JSON_DEPTH = 4 MAX_JSON_STRING_LENGTH = 512 MAX_JSON_DICT_KEYS = 50 MAX_JSON_LIST_ITEMS = 100 # DeepSeek-OCR Model Configuration - Maximum Quality Settings for M4 Mac (Apple Silicon) MODEL_NAME = "deepseek-ai/DeepSeek-OCR" # Detect Apple Silicon (M1/M2/M3/M4) - use MPS if available, otherwise CPU IS_APPLE_SILICON = platform.machine() == "arm64" USE_GPU = os.getenv("USE_GPU", "true").lower() == "true" and not IS_APPLE_SILICON # M4 uses MPS, not CUDA USE_MPS = IS_APPLE_SILICON # Use Metal Performance Shaders on Apple Silicon # Maximum quality settings (larger = better, slower = more accurate) BASE_SIZE = int(os.getenv("DEEPSEEK_BASE_SIZE", "1280")) # Maximum quality: 1280 (not light!) IMAGE_SIZE = int(os.getenv("DEEPSEEK_IMAGE_SIZE", "1280")) # Maximum quality: 1280 (not light!) CROP_MODE = os.getenv("DEEPSEEK_CROP_MODE", "true").lower() == "true" # True for best accuracy app = FastAPI() # Add CORS middleware to allow frontend requests app.add_middleware( CORSMiddleware, allow_origins=["*"], # In production, replace with specific origins allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Initialize DeepSeek-OCR model _ocr_model = None _ocr_tokenizer = None _model_lock = asyncio.Lock() def _patch_deepseek_model_for_m4(): """ Patch DeepSeek-OCR model code to fix LlamaFlashAttention2 import error on M4 Mac. This is needed because transformers 4.57.1 doesn't have LlamaFlashAttention2, but DeepSeek-OCR's model code tries to import it. """ from pathlib import Path cache_dir = Path.home() / ".cache" / "huggingface" model_files = list(cache_dir.glob("**/modeling_deepseekv2.py")) if not model_files: return # Model not downloaded yet, will patch on first load model_file = model_files[0] # Check if already patched try: with open(model_file, 'r') as f: content = f.read() if "LlamaFlashAttention2 = LlamaAttention" in content: return # Already patched except: pass # Original import pattern original_import = """from transformers.models.llama.modeling_llama import ( LlamaAttention, LlamaFlashAttention2 )""" # Patched version with fallback patched_import = """from transformers.models.llama.modeling_llama import ( LlamaAttention, ) # Patch for M4 Mac: LlamaFlashAttention2 not available in transformers 4.57.1 # Use LlamaAttention as fallback when flash attention unavailable try: from transformers.models.llama.modeling_llama import LlamaFlashAttention2 except ImportError: # Fallback: Use LlamaAttention when flash attention not available LlamaFlashAttention2 = LlamaAttention""" try: if original_import in content: # Create backup backup_file = model_file.with_suffix('.py.backup') try: with open(backup_file, 'w') as f: f.write(content) except: pass # Apply patch content = content.replace(original_import, patched_import) with open(model_file, 'w') as f: f.write(content) print(f"✅ Patched DeepSeek model for M4 Mac compatibility") except Exception as e: print(f"⚠️ Could not patch model file: {e}") async def get_ocr_model(): """Lazy load DeepSeek-OCR model with M4 Mac compatibility patching""" global _ocr_model, _ocr_tokenizer if _ocr_model is None or _ocr_tokenizer is None: async with _model_lock: if _ocr_model is None or _ocr_tokenizer is None: # Patch DeepSeek model code for M4 Mac compatibility BEFORE loading _patch_deepseek_model_for_m4() # Lazy import dependencies AutoModel, AutoTokenizer = _get_transformers() torch = _get_torch() print(f"Loading DeepSeek-OCR model (MAXIMUM QUALITY): {MODEL_NAME}") print(f" - Base size: {BASE_SIZE} (maximum quality, not light version!)") print(f" - Image size: {IMAGE_SIZE} (maximum quality, not light version!)") print(f" - Crop mode: {CROP_MODE} (best accuracy)") _ocr_tokenizer = AutoTokenizer.from_pretrained( MODEL_NAME, trust_remote_code=True ) # Load model with Apple Silicon (M4) optimized settings # M4 Mac: Use SDPA (not flash_attention_2) - flash attention doesn't work on Apple Silicon load_kwargs = { "trust_remote_code": True, "use_safetensors": False, # Avoid safetensors issues on M4 } # Force SDPA attention for Apple Silicon compatibility # This avoids LlamaFlashAttention2 import errors on M4 Mac if IS_APPLE_SILICON: load_kwargs["_attn_implementation"] = "sdpa" print(" - Using SDPA attention (Apple Silicon/M4 optimized)") else: # For non-Apple Silicon, let model choose pass try: _ocr_model = AutoModel.from_pretrained(MODEL_NAME, **load_kwargs) except Exception as e: error_msg = str(e) print(f"⚠️ Model load error: {error_msg}") # If still fails, try minimal config if "LlamaFlashAttention2" in error_msg or "flash" in error_msg.lower(): print(" - Retrying with explicit SDPA attention...") load_kwargs_minimal = { "trust_remote_code": True, "use_safetensors": False, "_attn_implementation": "sdpa", # Force SDPA } _ocr_model = AutoModel.from_pretrained(MODEL_NAME, **load_kwargs_minimal) else: raise _ocr_model = _ocr_model.eval() # Handle device placement for M4 Mac (Apple Silicon) if USE_MPS and torch.backends.mps.is_available(): _ocr_model = _ocr_model.to("mps") print(" - DeepSeek-OCR loaded on Apple Silicon GPU (MPS/M4)") elif USE_GPU and torch.cuda.is_available(): _ocr_model = _ocr_model.cuda().to(torch.bfloat16) print(" - DeepSeek-OCR loaded on NVIDIA GPU") else: print(" - DeepSeek-OCR loaded on CPU") return _ocr_model, _ocr_tokenizer async def run_deepseek_ocr( image_path: str, prompt: str = "\n<|grounding|>Convert the document to markdown with preserved layout.", use_grounding: bool = True ) -> dict: """ Run DeepSeek-OCR on an image file with advanced grounding support. Genius enhancement: Uses grounding prompts for better structure extraction and layout preservation, following DeepSeek-OCR best practices. """ model, tokenizer = await get_ocr_model() output_path = tempfile.mkdtemp() try: # Maximum quality inference - best OCR quality settings result = model.infer( tokenizer, prompt=prompt, image_file=image_path, output_path=output_path, base_size=BASE_SIZE, # 1280 = maximum quality (not light version!) image_size=IMAGE_SIZE, # 1280 = maximum quality (not light version!) crop_mode=CROP_MODE, # True = best accuracy for complex documents save_results=False, test_compress=False, # False = maximum quality, no compression ) # Parse result - DeepSeek-OCR returns structured markdown output ocr_text = result if isinstance(result, str) else str(result) # Genius parsing: Extract structured lines from markdown with better layout awareness lines = _parse_deepseek_output(ocr_text) return { "text": ocr_text, "lines": lines, } except Exception as e: print(f"DeepSeek-OCR error: {e}") import traceback traceback.print_exc() raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"OCR processing failed: {str(e)}", ) finally: # Cleanup temp directory try: import shutil if os.path.exists(output_path): shutil.rmtree(output_path) except: pass def _parse_deepseek_output(ocr_text: str) -> list: """ Genius parser: Extract structured lines from DeepSeek-OCR markdown output. Preserves layout, handles tables, lists, and structured content. """ lines = [] text_lines = ocr_text.split('\n') y_offset = 0 line_height = 24 # Estimated line height in pixels for line_idx, line in enumerate(text_lines): stripped = line.strip() if not stripped: # Empty lines still take space y_offset += line_height // 2 continue # Remove markdown formatting but preserve text structure # Handle markdown tables (| separated) if '|' in stripped and stripped.count('|') >= 2: # Table row - split by | and process each cell cells = [cell.strip() for cell in stripped.split('|') if cell.strip()] for cell_idx, cell in enumerate(cells): if cell: lines.append({ "bbox": [ cell_idx * 200, # Approximate x position y_offset, (cell_idx + 1) * 200, y_offset + line_height ], "text": cell, "conf": 0.95, }) y_offset += line_height # Handle markdown lists (-, *, 1., etc.) elif stripped.startswith(('-', '*', '+')) or (len(stripped) > 2 and stripped[1] == '.'): # List item - remove list marker text = stripped.lstrip('-*+').lstrip('0123456789.').strip() if text: lines.append({ "bbox": [40, y_offset, 1000, y_offset + line_height], "text": text, "conf": 0.95, }) y_offset += line_height # Handle headers (# ## ###) elif stripped.startswith('#'): header_level = len(stripped) - len(stripped.lstrip('#')) text = stripped.lstrip('#').strip() if text: # Headers are typically larger header_height = line_height + (header_level * 4) lines.append({ "bbox": [0, y_offset, 1000, y_offset + header_height], "text": text, "conf": 0.95, }) y_offset += header_height # Regular text line else: # Estimate width based on text length (rough approximation) estimated_width = min(len(stripped) * 8, 1000) # ~8px per char average lines.append({ "bbox": [0, y_offset, estimated_width, y_offset + line_height], "text": stripped, "conf": 0.95, }) y_offset += line_height return lines api_key_header = APIKeyHeader(name=API_KEY_HEADER_NAME, auto_error=False) _rate_limit_lock = asyncio.Lock() _request_log: DefaultDict[str, Deque[float]] = defaultdict(deque) def ensure_upload_is_safe(file: UploadFile) -> None: # Check content type from header content_type = (file.content_type or "").lower() # Also check file extension as fallback (browsers sometimes send application/octet-stream) filename = (file.filename or "").lower() extension = filename.split('.')[-1] if '.' in filename else "" allowed_extensions = {'jpg', 'jpeg', 'png', 'webp'} # Allow if content type matches OR extension matches content_type_valid = content_type in ALLOWED_CONTENT_TYPES extension_valid = extension in allowed_extensions if not content_type_valid and not extension_valid: raise HTTPException( status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, detail=f"Unsupported file type. Content-Type: {content_type}, Extension: {extension}. Allowed: {', '.join(ALLOWED_CONTENT_TYPES)}", ) file.file.seek(0, os.SEEK_END) size = file.file.tell() file.file.seek(0) if size > MAX_UPLOAD_BYTES: raise HTTPException( status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, detail="Uploaded file exceeds size limit", ) async def verify_api_key(api_key: Optional[str] = Depends(api_key_header)) -> str: # Skip API key verification in development mode if not REQUIRE_API_KEY: return api_key or SERVICE_API_KEY # Enforce API key in production if not api_key or not secrets.compare_digest(api_key, SERVICE_API_KEY): raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key", ) return api_key async def enforce_rate_limit( request: Request, api_key: str = Depends(verify_api_key) ) -> None: if RATE_LIMIT_REQUESTS <= 0: return identifier = api_key or (request.client.host if request.client else "anonymous") now = monotonic() async with _rate_limit_lock: window = _request_log[identifier] while window and now - window[0] > RATE_LIMIT_WINDOW_SECONDS: window.popleft() if len(window) >= RATE_LIMIT_REQUESTS: raise HTTPException( status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail="Rate limit exceeded", ) window.append(now) def _decode_image(file: UploadFile) -> Image.Image: """Decode uploaded image file to PIL Image""" data = file.file.read() if not data: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Uploaded file is empty", ) # Save to temp file for DeepSeek-OCR with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file: tmp_file.write(data) tmp_path = tmp_file.name try: img = Image.open(tmp_path).convert("RGB") return img, tmp_path except Exception as e: os.unlink(tmp_path) raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Unable to decode image: {str(e)}", ) async def load_img(file: UploadFile): ensure_upload_is_safe(file) file.file.seek(0) img, img_path = _decode_image(file) return img, img_path def _parse_json_field(name: str, raw: str, expected_type: type) -> Any: try: value = json.loads(raw) except json.JSONDecodeError as exc: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid {name} payload", ) from exc if not isinstance(value, expected_type): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"{name} must be a {expected_type.__name__}", ) return value def _validate_safe_json(value: Any, name: str, depth: int = 0) -> None: if depth > MAX_JSON_DEPTH: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"{name} is too deeply nested", ) if isinstance(value, dict): if len(value) > MAX_JSON_DICT_KEYS: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"{name} has too many keys", ) for key, item in value.items(): if not isinstance(key, str) or len(key) > 64: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"{name} contains an invalid key", ) _validate_safe_json(item, f"{name}.{key}", depth + 1) return if isinstance(value, list): if len(value) > MAX_JSON_LIST_ITEMS: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"{name} has too many entries", ) for idx, item in enumerate(value): _validate_safe_json(item, f"{name}[{idx}]", depth + 1) return if isinstance(value, str): if len(value) > MAX_JSON_STRING_LENGTH: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"{name} contains an oversized string", ) if any(ord(ch) < 32 and ch not in (9, 10, 13) for ch in value): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"{name} contains control characters", ) return if isinstance(value, bool) or value is None: return if isinstance(value, (int, float)): if isinstance(value, float) and not math.isfinite(value): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"{name} must contain finite numbers", ) return raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"{name} contains an unsupported value type", ) def _sanitize_label(name: str, value: str) -> str: if not isinstance(value, str): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"{name} must be a string", ) trimmed = value.strip() if not trimmed: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"{name} cannot be empty", ) if len(trimmed) > 128: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"{name} is too long", ) if any(ord(ch) < 32 for ch in trimmed): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"{name} contains invalid characters", ) return trimmed def _parse_parent_bbox(raw: str, width: int, height: int) -> list[float]: values = _parse_json_field("parent_bbox", raw, list) if len(values) != 4: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="parent_bbox must have four values", ) coords: list[float] = [] for value in values: try: coord = float(value) except (TypeError, ValueError) as exc: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="parent_bbox must contain numeric values", ) from exc if not math.isfinite(coord): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="parent_bbox must contain finite coordinates", ) coords.append(coord) x1, y1, x2, y2 = coords if x2 <= x1 or y2 <= y1: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="parent_bbox coordinates are invalid", ) if x1 < 0 or y1 < 0 or x2 > width or y2 > height: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="parent_bbox is outside the image bounds", ) return coords def _parse_settings(raw: str) -> dict: settings = _parse_json_field("settings", raw, dict) if len(settings) > 50: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="settings payload is too large", ) _validate_safe_json(settings, "settings") return settings def _parse_rules(raw: str) -> list: rules = _parse_json_field("rules", raw, list) if len(rules) > 100: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="rules payload is too large", ) for idx, rule in enumerate(rules): if not isinstance(rule, dict): raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="rules entries must be objects", ) _validate_safe_json(rule, f"rules[{idx}]") return rules @app.post("/ocr") async def ocr_page( file: UploadFile, _: None = Depends(enforce_rate_limit), ): """OCR endpoint using DeepSeek-OCR""" img, img_path = await load_img(file) try: # Save PIL image to temporary file for DeepSeek-OCR with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file: img.save(tmp_file, 'JPEG', quality=95) tmp_img_path = tmp_file.name try: # Use grounding prompt for better structure extraction result = await run_deepseek_ocr( tmp_img_path, prompt="\n<|grounding|>Convert the document to markdown with preserved layout.", use_grounding=True ) return result except Exception as e: # Log the error but don't crash - return a helpful error message error_msg = str(e) print(f"OCR processing error: {error_msg}") # Check if it's a model loading issue if "matplotlib" in error_msg or "torchvision" in error_msg or "ImportError" in error_msg: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=f"OCR model dependencies missing: {error_msg}. Please install required packages." ) elif "Connection" in error_msg or "timeout" in error_msg.lower(): raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=f"OCR service temporarily unavailable: {error_msg}" ) else: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"OCR processing failed: {error_msg}" ) finally: if os.path.exists(tmp_img_path): os.unlink(tmp_img_path) finally: if os.path.exists(img_path): os.unlink(img_path) @app.post("/split") async def split( file: UploadFile, parent_bbox: str = Form(...), splitter: str = Form(...), schemaType: str = Form(...), settings: str = Form("{}"), rules: str = Form("[]"), _: None = Depends(enforce_rate_limit), ): """Split endpoint - uses DeepSeek-OCR for region extraction""" img, img_path = await load_img(file) try: width, height = img.size # Save image for DeepSeek-OCR with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file: img.save(tmp_file, 'JPEG', quality=95) tmp_img_path = tmp_file.name try: parent_box = _parse_parent_bbox(parent_bbox, width, height) x1, y1, x2, y2 = parent_box # Crop image to parent bbox crop_img = img.crop((int(x1), int(y1), int(x2), int(y2))) crop_path = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg').name crop_img.save(crop_path, 'JPEG', quality=95) try: # Use DeepSeek-OCR with grounding prompt for better structured extraction prompt = "\n<|grounding|>Convert the document region to markdown with preserved layout." ocr_result = await run_deepseek_ocr(crop_path, prompt=prompt, use_grounding=True) # Parse OCR result to extract lines child_lines = ocr_result.get("lines", []) # Adjust bboxes to parent coordinate space for line in child_lines: bbox = line["bbox"] line["bbox"] = [ bbox[0] + x1, bbox[1] + y1, bbox[2] + x1, bbox[3] + y1, ] line["blockType"] = "text" if len(child_lines) > MAX_CHILD_LINES: child_lines = child_lines[:MAX_CHILD_LINES] sanitized_splitter = _sanitize_label("splitter", splitter) sanitized_schema = _sanitize_label("schemaType", schemaType) parsed_settings = _parse_settings(settings) parsed_rules = _parse_rules(rules) raw_text = "\n".join([l["text"] for l in child_lines]) text_truncated = False if len(raw_text) > 5000: raw_text = raw_text[:5000] text_truncated = True llm_input = { "schemaType": sanitized_schema, "splitter": sanitized_splitter, "page": {"width": width, "height": height}, "parentBox": parent_box, "rawText": raw_text, "ocrLines": child_lines, "rawTextTruncated": text_truncated, "ocrLinesTruncated": len(child_lines) >= MAX_CHILD_LINES, "settings": parsed_settings, "rules": parsed_rules, } try: llm_result = await call_llm_splitter(llm_input) except ValueError as exc: raise HTTPException( status_code=status.HTTP_502_BAD_GATEWAY, detail=str(exc), ) from exc return llm_result finally: if os.path.exists(crop_path): os.unlink(crop_path) finally: if os.path.exists(tmp_img_path): os.unlink(tmp_img_path) finally: if os.path.exists(img_path): os.unlink(img_path) if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8080)