algoryn/dots-ocr-idcard
0
1"""Dots.OCR Model Loader2 3This module handles downloading and loading the Dots.OCR model using4Hugging Face's `snapshot_download`. It centralizes device selection,5dtype configuration, model initialization, and safe fallbacks.6 7Why this exists:8- Keep model lifecycle and I/O concerns isolated from API/business logic.9- Provide safe CPU defaults, optional CUDA acceleration, and optional10 FlashAttention2 when compatible and explicitly enabled.11 12Key environment variables:13- DOTS_OCR_REPO_ID: HF repo to download (default: "rednote-hilab/dots.ocr").14- DOTS_OCR_LOCAL_DIR: Local cache directory for `snapshot_download`.15- DOTS_OCR_DEVICE: One of {"cpu", "cuda", "auto"}. "auto" prefers CUDA.16- DOTS_OCR_MAX_NEW_TOKENS: Max generated tokens per request.17- DOTS_OCR_FLASH_ATTENTION: "1" to attempt FlashAttention2 when compatible.18- DOTS_OCR_MIN_PIXELS / DOTS_OCR_MAX_PIXELS: Image size bounds pre-inference.19- DOTS_OCR_PROMPT: Optional default transcription prompt.20 21Usage: call `load_model()` once, then `extract_text(image)` per request.22"""23 24import os25import logging26import torch27from typing import Optional, Tuple, Dict, Any28from pathlib import Path29 30from huggingface_hub import snapshot_download31from transformers import AutoModelForCausalLM, AutoProcessor32from PIL import Image33 34# Configure logging35logger = logging.getLogger(__name__)36 37# Environment variable configuration38#39# These env vars make runtime behavior tunable without code changes. Defaults are40# conservative to favor stability on CPU-only platforms; performance features41# are opt-in and gated by compatibility checks.42REPO_ID = os.getenv("DOTS_OCR_REPO_ID", "rednote-hilab/dots.ocr")43LOCAL_DIR = os.getenv("DOTS_OCR_LOCAL_DIR", "/data/models/dots-ocr")44DEVICE_CONFIG = os.getenv("DOTS_OCR_DEVICE", "auto") # "auto" prefers CUDA if available45MAX_NEW_TOKENS = int(os.getenv("DOTS_OCR_MAX_NEW_TOKENS", "2048"))46USE_FLASH_ATTENTION = os.getenv("DOTS_OCR_FLASH_ATTENTION", "0") == "1" # opt-in47MIN_PIXELS = int(os.getenv("DOTS_OCR_MIN_PIXELS", "3136")) # 56x56 lower bound48MAX_PIXELS = int(os.getenv("DOTS_OCR_MAX_PIXELS", "11289600")) # 3360x3360 upper bound49CUSTOM_PROMPT = os.getenv("DOTS_OCR_PROMPT")50 51# Default transcription prompt for faithful text extraction.52# Keep terse to reduce bias; we want faithful extraction, not translation or formatting.53DEFAULT_PROMPT = (54 "Transcribe all visible text in the image in the original language. "55 "Do not translate. Preserve natural reading order. Output plain text only."56)57 58 59class DotsOCRModelLoader:60 """Handles Dots.OCR model downloading, loading, and inference.61 62 Encapsulates model lifecycle (download, init, device placement), preprocessing,63 and a narrow inference surface for OCR. Exposes a minimal API and maintains a64 single global instance via helpers below.65 """66 67 def __init__(self):68 """Initialize the model loader.69 70 Heavyweight work is deferred until `load_model()` so that constructing this71 class is cheap. The default prompt is captured from env, if provided.72 """73 self.model = None74 self.processor = None75 self.device = None76 self.dtype = None77 self.local_dir = None78 self.prompt = CUSTOM_PROMPT or DEFAULT_PROMPT79 80 def _determine_device_and_dtype(self) -> Tuple[str, torch.dtype]:81 """Pick device and dtype based on availability and configuration.82 83 Rules:84 - Respect explicit "cpu" or "cuda" when valid.85 - "auto" selects CUDA when available, else CPU.86 - Use bfloat16 on CUDA for throughput; float32 on CPU for correctness.87 """88 if DEVICE_CONFIG == "cpu":89 device = "cpu"90 dtype = torch.float3291 elif DEVICE_CONFIG == "cuda" and torch.cuda.is_available():92 device = "cuda"93 dtype = torch.bfloat1694 elif DEVICE_CONFIG == "auto":95 if torch.cuda.is_available():96 device = "cuda"97 dtype = torch.bfloat1698 else:99 device = "cpu"100 dtype = torch.float32101 else:102 # Fallback to CPU if CUDA requested but not available103 logger.warning(f"CUDA requested but not available, falling back to CPU")104 device = "cpu"105 dtype = torch.float32106 107 logger.info(f"Selected device: {device}, dtype: {dtype}")108 return device, dtype109 110 def _download_model(self) -> str:111 """Download the model using `snapshot_download` and ensure cache dir exists.112 113 Returns the resolved local path for deterministic, offline-friendly loading.114 Raises `RuntimeError` on failure.115 """116 logger.info(f"Downloading model from {REPO_ID} to {LOCAL_DIR}")117 118 try:119 # Ensure local directory exists120 Path(LOCAL_DIR).mkdir(parents=True, exist_ok=True)121 122 # Download model snapshot123 local_path = snapshot_download(124 repo_id=REPO_ID,125 local_dir=LOCAL_DIR,126 )127 128 logger.info(f"Model downloaded successfully to {local_path}")129 return local_path130 131 except Exception as e:132 logger.error(f"Failed to download model: {e}")133 raise RuntimeError(f"Model download failed: {e}")134 135 def _can_use_flash_attn(self) -> bool:136 """Check whether FlashAttention2 can be enabled safely.137 138 Requires all of:139 - DOTS_OCR_FLASH_ATTENTION toggle is set.140 - `flash_attn` is importable.141 - dtype is fp16/bf16 per library support.142 """143 if not USE_FLASH_ATTENTION:144 return False145 try:146 # Import check avoids runtime error from Transformers if not installed147 import flash_attn # type: ignore # noqa: F401148 except Exception:149 logger.warning(150 "flash_attn package not installed; disabling FlashAttention2"151 )152 return False153 # FlashAttention2 supports fp16/bf16 only (see HF docs)154 return self.dtype in (torch.float16, torch.bfloat16)155 156 def load_model(self) -> None:157 """Load the Dots.OCR model and processor.158 159 Steps:160 1) Determine device/dtype161 2) Download snapshot if missing162 3) Load `AutoProcessor`163 4) Configure attention/device mapping164 5) Instantiate model and place on target device165 """166 try:167 # Determine device and dtype168 self.device, self.dtype = self._determine_device_and_dtype()169 170 # Download model if not already present171 self.local_dir = self._download_model()172 173 # Load processor174 logger.info("Loading processor...")175 self.processor = AutoProcessor.from_pretrained(176 self.local_dir, trust_remote_code=True177 )178 179 # Load model with appropriate configuration180 model_kwargs = {181 "dtype": self.dtype, # NOTE: `torch_dtype` is deprecated upstream182 "trust_remote_code": True,183 }184 185 # Add device-specific configurations186 if self.device == "cuda":187 # Prefer FlashAttention2 when truly available; otherwise SDPA188 if self._can_use_flash_attn():189 model_kwargs["attn_implementation"] = "flash_attention_2"190 logger.info("Using flash attention 2")191 else:192 model_kwargs["attn_implementation"] = "sdpa"193 logger.info(194 "Using SDPA attention (flash-attn unavailable or disabled)"195 )196 197 # Use device_map for automatic GPU memory management198 model_kwargs["device_map"] = "auto"199 else:200 # For CPU, don't use device_map201 model_kwargs["device_map"] = None202 203 logger.info("Loading model...")204 self.model = AutoModelForCausalLM.from_pretrained(205 self.local_dir, **model_kwargs206 )207 208 # Move model to device if not using device_map209 if self.device == "cpu" or model_kwargs.get("device_map") is None:210 self.model = self.model.to(self.device)211 212 logger.info(f"Model loaded successfully on {self.device}")213 214 except Exception as e:215 logger.error(f"Failed to load model: {e}")216 raise RuntimeError(f"Model loading failed: {e}")217 218 def _preprocess_image(self, image: Image.Image) -> Image.Image:219 """Preprocess image to meet model requirements.220 221 - Normalize to RGB222 - Constrain pixel count within [MIN_PIXELS, MAX_PIXELS]223 - Snap dimensions to multiples of 28 to satisfy backbone constraints224 """225 # Convert to RGB if necessary226 if image.mode != "RGB":227 image = image.convert("RGB")228 229 # Calculate current pixel count230 width, height = image.size231 current_pixels = width * height232 233 # Resize if necessary to meet pixel requirements234 if current_pixels < MIN_PIXELS:235 # Scale up to meet minimum pixel requirement236 scale_factor = (MIN_PIXELS / current_pixels) ** 0.5237 new_width = int(width * scale_factor)238 new_height = int(height * scale_factor)239 image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)240 logger.info(241 f"Scaled up image from {width}x{height} to {new_width}x{new_height}"242 )243 244 elif current_pixels > MAX_PIXELS:245 # Scale down to meet maximum pixel requirement246 scale_factor = (MAX_PIXELS / current_pixels) ** 0.5247 new_width = int(width * scale_factor)248 new_height = int(height * scale_factor)249 image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)250 logger.info(251 f"Scaled down image from {width}x{height} to {new_width}x{new_height}"252 )253 254 # Ensure dimensions are divisible by 28 (common requirement for vision models)255 width, height = image.size256 new_width = ((width + 27) // 28) * 28257 new_height = ((height + 27) // 28) * 28258 259 if new_width != width or new_height != height:260 image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)261 logger.info(262 f"Adjusted image dimensions to be divisible by 28: {new_width}x{new_height}"263 )264 265 return image266 267 @torch.inference_mode()268 def extract_text(self, image: Image.Image, prompt: Optional[str] = None) -> str:269 """Extract text from an image using the loaded model.270 271 Builds a single-turn chat message with the image and a transcription prompt,272 applies the model's chat template, and decodes deterministically (no sampling).273 """274 if self.model is None or self.processor is None:275 raise RuntimeError("Model not loaded. Call load_model() first.")276 277 try:278 # Preprocess image279 processed_image = self._preprocess_image(image)280 281 # Use provided prompt or default282 text_prompt = prompt or self.prompt283 284 # Prepare messages for the model285 messages = [286 {287 "role": "user",288 "content": [289 {"type": "image", "image": processed_image},290 {"type": "text", "text": text_prompt},291 ],292 }293 ]294 295 # Apply chat template (preserves special tokens/formatting expected by model)296 text = self.processor.apply_chat_template(297 messages, tokenize=False, add_generation_prompt=True298 )299 300 # Process vision information (required for some models)301 try:302 from qwen_vl_utils import process_vision_info303 304 image_inputs, video_inputs = process_vision_info(messages)305 except ImportError:306 # Fallback if qwen_vl_utils not available307 logger.warning("qwen_vl_utils not available, using basic processing")308 image_inputs = [processed_image]309 video_inputs = []310 311 # Prepare inputs312 inputs = self.processor(313 text=[text],314 images=image_inputs,315 videos=video_inputs,316 padding=True,317 return_tensors="pt",318 ).to(self.device)319 320 # Generate text deterministically (temperature=0, do_sample=False)321 output_ids = self.model.generate(322 **inputs,323 max_new_tokens=MAX_NEW_TOKENS,324 do_sample=False,325 temperature=0.0,326 pad_token_id=self.processor.tokenizer.eos_token_id,327 )328 329 # Decode only newly generated tokens (strip prompt tokens)330 trimmed = [331 out[len(inp) :] for inp, out in zip(inputs.input_ids, output_ids)332 ]333 decoded = self.processor.batch_decode(334 trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False335 )336 337 return decoded[0] if decoded else ""338 339 except Exception as e:340 logger.error(f"Text extraction failed: {e}")341 raise RuntimeError(f"Text extraction failed: {e}")342 343 def is_loaded(self) -> bool:344 """Return True when both model and processor are initialized."""345 return self.model is not None and self.processor is not None346 347 def get_model_info(self) -> Dict[str, Any]:348 """Get diagnostic information about the loaded model and configuration."""349 return {350 "device": self.device,351 "dtype": str(self.dtype),352 "local_dir": self.local_dir,353 "repo_id": REPO_ID,354 "max_new_tokens": MAX_NEW_TOKENS,355 "use_flash_attention": USE_FLASH_ATTENTION,356 "prompt": self.prompt,357 "is_loaded": self.is_loaded(),358 }359 360 361# Global model instance362_model_loader: Optional[DotsOCRModelLoader] = None363 364 365def get_model_loader() -> DotsOCRModelLoader:366 """Get the global model loader instance."""367 global _model_loader368 if _model_loader is None:369 _model_loader = DotsOCRModelLoader()370 return _model_loader371 372 373def load_model() -> None:374 """Load the Dots.OCR model."""375 loader = get_model_loader()376 loader.load_model()377 378 379def extract_text(image: Image.Image, prompt: Optional[str] = None) -> str:380 """Extract text from an image using the loaded model."""381 loader = get_model_loader()382 if not loader.is_loaded():383 raise RuntimeError("Model not loaded. Call load_model() first.")384 return loader.extract_text(image, prompt)385 386 387def is_model_loaded() -> bool:388 """Check if the model is loaded and ready."""389 loader = get_model_loader()390 return loader.is_loaded()391 392 393def get_model_info() -> Dict[str, Any]:394 """Get information about the loaded model."""395 loader = get_model_loader()396 return loader.get_model_info()397