""" LightOnOCR engine wrapper. Model files are cached through storage.py before Hugging Face / Transformers are initialized. On app restarts, existing files in the persistent cache are reused. """ from __future__ import annotations import base64 import io import json import os import shutil from typing import Any, Dict, Literal, Union import cv2 import numpy as np import requests from PIL import Image try: from .storage import ensure_dirs, get_env_overrides, is_model_cached except ImportError: from storage import ensure_dirs, get_env_overrides, is_model_cached ensure_dirs() os.environ.update(get_env_overrides()) DEFAULT_MODEL_ID = "PetaniHandal/LightOnOCR-2-1B-1025-ft-iam-handwriting-vllm" PROMPTS = { "handwriting": ( "Extract all handwritten and printed text from this document image. " "Preserve reading order, line breaks, tables, totals, and important labels. " "Return clean Markdown only." ), "document": ( "Perform OCR on this document. Preserve reading order, headings, tables, " "and key-value fields. Return clean Markdown only." ), "receipt": ( "Extract receipt or invoice content from this image. Preserve merchant, " "date, item rows, prices, totals, and notes. Return clean Markdown only." ), } def _hub_cache_dir() -> str: return os.environ.get("HUGGINGFACE_HUB_CACHE", os.path.join(os.environ["HF_HOME"], "hub")) class OCREngine: def __init__( self, preset: Literal["handwriting", "document", "receipt"] = "handwriting", model_id: str | None = None, endpoint_url: str | None = None, max_new_tokens: int = 4096, temperature: float = 0.1, top_p: float = 0.9, ): if preset not in PROMPTS: raise ValueError(f"Preset '{preset}' tidak dikenal. Pilihan: {list(PROMPTS)}") self.preset = preset self.model_id = model_id or os.getenv("LIGHTONOCR_MODEL_ID", DEFAULT_MODEL_ID) self.endpoint_url = endpoint_url or os.getenv("LIGHTONOCR_ENDPOINT_URL") self.max_new_tokens = max_new_tokens self.temperature = temperature self.top_p = top_p self._processor = None self._model = None self._device = None self._dtype = None if not self.endpoint_url: self._load_local_model() def _load_local_model(self) -> None: import torch try: from transformers import AutoProcessor, LightOnOcrForConditionalGeneration except ImportError: from transformers import AutoProcessor from transformers import AutoModelForSeq2SeqLM as LightOnOcrForConditionalGeneration device = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.bfloat16 if device == "cuda" else torch.float32 adapter_config = self._load_adapter_config() model_to_load = adapter_config.get("base_model_name_or_path", self.model_id) model_local_only = is_model_cached(model_to_load) adapter_local_only = is_model_cached(self.model_id) self._processor = AutoProcessor.from_pretrained( self.model_id, cache_dir=_hub_cache_dir(), local_files_only=adapter_local_only, trust_remote_code=True, ) try: self._model = LightOnOcrForConditionalGeneration.from_pretrained( model_to_load, cache_dir=_hub_cache_dir(), local_files_only=model_local_only, dtype=dtype, device_map="auto", attn_implementation="sdpa", trust_remote_code=True, ) except OSError as exc: if model_local_only and self._looks_like_missing_weights(exc): self._remove_partial_model_cache(model_to_load) self._model = LightOnOcrForConditionalGeneration.from_pretrained( model_to_load, cache_dir=_hub_cache_dir(), local_files_only=False, dtype=dtype, device_map="auto", attn_implementation="sdpa", trust_remote_code=True, ) else: raise if adapter_config: from peft import PeftModel try: self._model = PeftModel.from_pretrained( self._model, self.model_id, cache_dir=_hub_cache_dir(), local_files_only=adapter_local_only, ) except OSError as exc: if adapter_local_only and self._looks_like_missing_weights(exc): self._remove_partial_model_cache(self.model_id) self._model = PeftModel.from_pretrained( self._model, self.model_id, cache_dir=_hub_cache_dir(), local_files_only=False, ) else: raise self._model.eval() self._device = device self._dtype = dtype def _load_adapter_config(self) -> dict[str, Any]: from huggingface_hub import hf_hub_download from huggingface_hub.utils import EntryNotFoundError local_only = is_model_cached(self.model_id) try: path = hf_hub_download( repo_id=self.model_id, filename="adapter_config.json", cache_dir=_hub_cache_dir(), local_files_only=local_only, ) except EntryNotFoundError: return {} except OSError: if local_only: path = hf_hub_download( repo_id=self.model_id, filename="adapter_config.json", cache_dir=_hub_cache_dir(), local_files_only=False, ) else: raise with open(path, "r", encoding="utf-8") as file: return json.load(file) def _looks_like_missing_weights(self, exc: OSError) -> bool: message = str(exc) return "pytorch_model.bin" in message or "model.safetensors" in message def _remove_partial_model_cache(self, model_id: str) -> None: repo_dir = f"models--{model_id.replace('/', '--')}" cache_dir = os.path.join(os.environ["HF_HOME"], "hub", repo_dir) if os.path.isdir(cache_dir): shutil.rmtree(cache_dir, ignore_errors=True) def process_image( self, image: Union[str, np.ndarray, Image.Image], max_size: int = 1540, ) -> Dict[str, Any]: pil_image = self._to_pil_image(image) pil_image = self._resize_if_needed(pil_image, max_size) if self.endpoint_url: markdown_text = self._process_with_endpoint(pil_image) else: markdown_text = self._process_local(pil_image) return { "markdown_text": markdown_text.strip(), "images": [], "model_id": self.model_id, "runtime": "vLLM endpoint" if self.endpoint_url else "Transformers local", } def _process_local(self, image: Image.Image) -> str: import torch assert self._processor is not None assert self._model is not None messages = [{"role": "user", "content": [{"type": "image"}]}] text = self._processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, ) inputs = self._processor(text=[text], images=[image], return_tensors="pt").to(self._device) if "pixel_values" in inputs: inputs["pixel_values"] = inputs["pixel_values"].to(self._dtype) with torch.inference_mode(): output_ids = self._model.generate( **inputs, max_new_tokens=self.max_new_tokens, do_sample=self.temperature > 0, temperature=self.temperature, top_p=self.top_p, ) input_length = inputs["input_ids"].shape[1] return self._processor.tokenizer.decode( output_ids[0, input_length:], skip_special_tokens=True, ) def _process_with_endpoint(self, image: Image.Image) -> str: buffer = io.BytesIO() image.save(buffer, format="PNG") image_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8") payload = { "model": self.model_id, "messages": [ { "role": "user", "content": [ {"type": "text", "text": PROMPTS[self.preset]}, { "type": "image_url", "image_url": { "url": f"data:image/png;base64,{image_base64}", }, }, ], } ], "max_tokens": self.max_new_tokens, "temperature": self.temperature, "top_p": self.top_p, } response = requests.post(self.endpoint_url, json=payload, timeout=300) response.raise_for_status() data = response.json() return data["choices"][0]["message"]["content"] def _to_pil_image(self, image: Union[str, np.ndarray, Image.Image]) -> Image.Image: if isinstance(image, Image.Image): return image.convert("RGB") if isinstance(image, str): img = cv2.imread(image) if img is None: raise ValueError(f"Gambar tidak bisa dibaca: {image}") return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) if isinstance(image, np.ndarray): if image.ndim == 2: return Image.fromarray(image).convert("RGB") return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) raise TypeError(f"Tipe gambar tidak didukung: {type(image)!r}") def _resize_if_needed(self, image: Image.Image, max_size: int) -> Image.Image: longest = max(image.size) if longest <= max_size: return image scale = max_size / longest size = (int(image.width * scale), int(image.height * scale)) return image.resize(size, Image.Resampling.LANCZOS)