import base64 import gc import json import os import re import time from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass from io import BytesIO from collections import OrderedDict from collections.abc import Mapping import threading from typing import Dict, List, Optional, Tuple import gradio as gr from PIL import Image try: import transformers hf_pipeline = transformers.pipeline # pragma: no cover AutoImageProcessor = getattr(transformers, "AutoImageProcessor", None) # pragma: no cover AutoModel = getattr(transformers, "AutoModel", None) # pragma: no cover AutoModelForCausalLM = getattr(transformers, "AutoModelForCausalLM", None) # pragma: no cover AutoModelForImageTextToText = getattr(transformers, "AutoModelForImageTextToText", None) # pragma: no cover AutoModelForSeq2SeqLM = getattr(transformers, "AutoModelForSeq2SeqLM", None) # pragma: no cover AutoModelForVision2Seq = getattr(transformers, "AutoModelForVision2Seq", None) # pragma: no cover AutoModelForVisionEncoderDecoder = getattr(transformers, "AutoModelForVisionEncoderDecoder", None) # pragma: no cover AutoModelForConditionalGeneration = getattr(transformers, "AutoModelForConditionalGeneration", None) # pragma: no cover AutoModelForDocumentQuestionAnswering = getattr(transformers, "AutoModelForDocumentQuestionAnswering", None) # pragma: no cover AutoTokenizer = getattr(transformers, "AutoTokenizer", None) # pragma: no cover AutoProcessor = getattr(transformers, "AutoProcessor", None) # pragma: no cover AutoConfig = getattr(transformers, "AutoConfig", None) # pragma: no cover TRANSFORMERS_IMPORT_ERROR = None # pragma: no cover except Exception as exc: # pragma: no cover hf_pipeline = None # pragma: no cover AutoImageProcessor = None # pragma: no cover AutoModel = None # pragma: no cover AutoModelForCausalLM = None # pragma: no cover AutoModelForImageTextToText = None # pragma: no cover AutoModelForSeq2SeqLM = None # pragma: no cover AutoModelForVision2Seq = None # pragma: no cover AutoModelForVisionEncoderDecoder = None # pragma: no cover AutoModelForConditionalGeneration = None # pragma: no cover AutoModelForDocumentQuestionAnswering = None # pragma: no cover AutoTokenizer = None # pragma: no cover AutoProcessor = None # pragma: no cover AutoConfig = None # pragma: no cover TRANSFORMERS_IMPORT_ERROR = repr(exc) try: import torch except Exception: # pragma: no cover torch = None # pragma: no cover try: import pytesseract except Exception: # pragma: no cover pytesseract = None try: import spaces except Exception: # pragma: no cover spaces = None try: from starlette.templating import Jinja2Templates _original_get_template = Jinja2Templates.get_template _original_template_response = Jinja2Templates.TemplateResponse def _normalize_template_name(template_name): if isinstance(template_name, dict): template_name = template_name.get("template_name", None) or template_name.get("name", None) if isinstance(template_name, (list, tuple)): template_name = template_name[0] if template_name else None if not template_name: template_name = "frontend/index.html" if template_name in {"index.html", "share.html"}: template_name = f"frontend/{template_name}" return template_name or "frontend/index.html" def _safe_get_template(self, name, *args, **kwargs): if isinstance(name, dict): return _original_get_template(self, _normalize_template_name(name), *args, **kwargs) return _original_get_template(self, name, *args, **kwargs) def _safe_template_response(self, name, context=None, *args, **kwargs): template_name = name request = kwargs.pop("request", None) if hasattr(name, "scope"): request = name template_name = context context = args[0] if args else None args = args[1:] if args else () elif request is None and isinstance(context, Mapping): request = context.get("request") if not request and args: request = args[0] if hasattr(args[0], "scope") else None if request is not None: args = args[1:] if isinstance(template_name, dict): template_name = _normalize_template_name(template_name) try: if context is None: context = {} elif isinstance(context, Mapping): context = dict(context) else: context = dict(context) except Exception: context = {} if isinstance(context, dict): context.setdefault( "config", { "body_css": {}, "title": "Hebrew OCR Document Comparator", "simple_description": "Private OCR comparator for Hebrew documents", "thumbnail": "", }, ) context.setdefault("gradio_api_info", {}) if request is None: return _original_template_response(self, template_name, context, *args, **kwargs) if args: return _original_template_response(self, request, template_name, context, *args, **kwargs) return _original_template_response(self, request, template_name, context=context, *args, **kwargs) Jinja2Templates.get_template = _safe_get_template Jinja2Templates.TemplateResponse = _safe_template_response except Exception: pass if spaces is not None: @spaces.GPU def _ensure_zero_gpu_lease() -> None: return None ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) REGISTRY_PATH = os.path.join(ROOT_DIR, "model_registry.json") DEFAULT_OCR_PROMPT = ( "You are a high-accuracy OCR engine for Hebrew documents that may contain printed and handwritten text. " "Return the exact document text in Hebrew, preserving line breaks and spacing. " "Do not add explanation or JSON." ) ZERO_GPU_MAX_WORKERS = 4 ZERO_GPU_SAFE_HEADROOM = 0.62 ZERO_GPU_MAX_HEADROOM = 0.82 ZERO_GPU_MIN_GUARDED_FREE_GB = 1.0 ZERO_GPU_MODEL_FALLBACK_GB_SAFE = 2.6 ZERO_GPU_MODEL_FALLBACK_GB_MAX = 2.0 ZERO_GPU_MAX_SAFE_SELECTED = 6 DEFAULT_MAX_TOKENS = 4096 DEFAULT_PRECHECK_TIMEOUT_SEC = 12 DIRECT_MODEL_CACHE_MAX = 1 LOCAL_PIPELINE_CACHE_MAX = 1 TESSERACT_EXECUTABLE_PATHS = ("/usr/bin/tesseract", "/usr/local/bin/tesseract", "/opt/conda/bin/tesseract", "/bin/tesseract") def _inference_device() -> int: if torch is not None and torch.cuda.is_available(): return 0 return -1 def _inference_device_label() -> str: if torch is not None and torch.cuda.is_available(): return "cuda" return "cpu" def _inference_dtype(): if torch is None: return None if hasattr(torch, "bfloat16"): return torch.bfloat16 if hasattr(torch, "float16"): return torch.float16 return None def _env_flag_true(name: str, default: bool = False) -> bool: raw = os.getenv(name) if raw is None: return default return raw.strip().lower() in {"1", "true", "yes", "on", "y"} def _allow_cpu_fallback() -> bool: return _env_flag_true("ALLOW_CPU_FALLBACK", default=False) def _strict_cuda_required() -> bool: return not _allow_cpu_fallback() def _ensure_cuda_available(context: str = "Inference") -> None: if torch is None: raise RuntimeError(f"{context} requires PyTorch, but PyTorch is not available in this Space.") if not torch.cuda.is_available(): raise RuntimeError( f"{context} requires CUDA, but no CUDA device is available. " "Run this Space on ZeroGPU and unset ALLOW_CPU_FALLBACK, or set ALLOW_CPU_FALLBACK=1." ) def _assert_model_on_cuda(model, model_id: str, context: str) -> None: if not _strict_cuda_required(): return if torch is None or not torch.cuda.is_available(): _ensure_cuda_available(context) if model is None: raise RuntimeError(f"{context}: loaded model is missing for {model_id}.") try: device_map = getattr(model, "hf_device_map", None) if isinstance(device_map, Mapping): non_cuda = { str(value) for value in device_map.values() if str(value) not in {"0", "cuda", "cuda:0", "cuda:1", "cuda:2", "cuda:3"} and str(value) != "None" } if non_cuda: raise RuntimeError(f"{context}: {model_id} model-map is not on CUDA (found: {sorted(non_cuda)}).") except Exception as exc: raise RuntimeError( f"{context}: {model_id} model-map validation failed for CUDA check ({type(exc).__name__}: {exc}). " "If you intentionally allow CPU execution, set ALLOW_CPU_FALLBACK=1." ) try: labels = {str(getattr(param, "device", "")) for param in model.parameters() if hasattr(param, "device")} if not labels: raise RuntimeError( f"{context}: {model_id} exposes no torch parameters to validate CUDA placement. " "If you intentionally allow CPU execution, set ALLOW_CPU_FALLBACK=1." ) non_cuda = [label for label in labels if not str(label).startswith("cuda") and str(label) != "0"] if non_cuda: raise RuntimeError( f"{context}: {model_id} is not on CUDA (tensor devices: {sorted(labels)}). " "If you intentionally allow CPU execution, set ALLOW_CPU_FALLBACK=1." ) except RuntimeError: raise except Exception as exc: raise RuntimeError( f"{context}: {model_id} model did not expose cuda-placed tensors ({type(exc).__name__}: {exc}). " "If you intentionally allow CPU execution, set ALLOW_CPU_FALLBACK=1." ) def _move_pipeline_to_cuda(pipeline, model_id: str, context: str): if torch is None or not torch.cuda.is_available(): return pipeline candidate_model = getattr(pipeline, "model", None) try: if hasattr(pipeline, "to"): pipeline = pipeline.to("cuda") elif candidate_model is not None: candidate_model = candidate_model.to("cuda") if hasattr(pipeline, "model"): pipeline.model = candidate_model else: raise RuntimeError("pipeline has no move target.") return pipeline except Exception as exc: raise RuntimeError( f"{context}: {model_id} failed to move pipeline to CUDA ({type(exc).__name__}: {exc}). " "If you intentionally allow CPU execution, set ALLOW_CPU_FALLBACK=1." ) def _inference_torch_kwargs_for_model(strict_cuda: bool = False) -> Tuple[dict, list]: kwargs_priority = [] dtype = _inference_dtype() if dtype is None: kwargs_priority.append({}) return dtype, kwargs_priority if strict_cuda and torch is not None and torch.cuda.is_available(): kwargs_priority.append({"torch_dtype": dtype, "low_cpu_mem_usage": True}) if dtype != getattr(torch, "float16", None): kwargs_priority.append({"torch_dtype": torch.float16, "low_cpu_mem_usage": True}) kwargs_priority.append({"torch_dtype": torch.float16}) kwargs_priority.append({"torch_dtype": dtype}) kwargs_priority.append({}) else: kwargs_priority.append({"torch_dtype": dtype, "device_map": "auto", "low_cpu_mem_usage": True}) if dtype != getattr(torch, "float16", None): kwargs_priority.append({"torch_dtype": torch.float16, "device_map": "auto", "low_cpu_mem_usage": True}) kwargs_priority.append({"torch_dtype": torch.float16}) kwargs_priority.append({"torch_dtype": dtype, "device_map": "auto"}) kwargs_priority.append({"torch_dtype": dtype}) kwargs_priority.append({"torch_dtype": torch.float16}) kwargs_priority.append({}) return dtype, kwargs_priority def _pipeline_cache_key(model_id: str, task: str, device: str, dtype: object, trust_remote_code: bool = False) -> Tuple[str, str, str, object, bool]: return (model_id, task, device, str(dtype), trust_remote_code) TRUST_REMOTE_CODE_MODELS = { "microsoft/Phi-4-multimodal-instruct", "PaddlePaddle/PaddleOCR-VL", "PaddlePaddle/PaddleOCR-VL-1.5", "PaddlePaddle/PaddleOCR-VL-1.6", "datalab-to/chandra-ocr-2", "datalab-to/surya-ocr-2", "ronylicha/gigapdf-ocr-hebrew", "liskcell/qunie-v7-mini", "0cve0/openmlkitocr", "waraja/tzefa-word-ocr-trocr", "liskcell/qunie-v7-pico", "cyttic/exp10-trocr-hebrew-matan-full", "cyttic/exp23-directfit-unfrozen", "cyttic/heb-verifier17-connected", "cyttic/exp26-composed1m", "deepseek-ai/deepseek-ocr", "deepseek-ai/deepseek-ocr-2", "coherelabs/aya-vision-8b", "coherelabs/aya-vision-32b", } TRUST_REMOTE_CODE_MODELS = {m.lower() for m in TRUST_REMOTE_CODE_MODELS} PIPELINE_TASK_HINTS = { "qwen/qwen3-vl-8b-instruct": "image-to-text", "qwen/qwen3-vl-4b-instruct": "image-to-text", "qwen/qwen3-vl-8b-thinking": "image-to-text", "qwen/qwen3-vl-4b-thinking": "image-to-text", "qwen/qwen3-vl-30b-a3b-instruct": "image-to-text", "qwen/qwen3-vl-30b-a3b-thinking": "image-to-text", "google/gemma-4-e4b-it": "image-to-text", "google/gemma-4-12b-it": "image-to-text", "google/gemma-4-26b-a4b-it": "image-to-text", "google/gemma-4-31b-it": "image-to-text", "ronylicha/gigapdf-ocr-hebrew": "image-to-text", "liskcell/qunie-v7-mini": "image-to-text", "0cve0/openmlkitocr": "image-to-text", "waraja/tzefa-word-ocr-trocr": "image-to-text", "liskcell/qunie-v7-pico": "image-to-text", "paddlepaddle/paddleocr-vl": "image-to-text", "paddlepaddle/paddleocr-vl-1.5": "image-to-text", "paddlepaddle/paddleocr-vl-1.6": "image-to-text", "datalab-to/chandra-ocr-2": "image-to-text", "datalab-to/surya-ocr-2": "image-to-text", "deepseek-ai/deepseek-ocr": "image-to-text", "deepseek-ai/deepseek-ocr-2": "image-to-text", "cyttic/exp10-trocr-hebrew-matan-full": "image-to-text", "cyttic/exp23-directfit-unfrozen": "image-to-text", "cyttic/heb-verifier17-connected": "image-to-text", "cyttic/exp26-composed1m": "image-to-text", } GATED_MODELS_REQUIRING_TOKEN = { "coherelabs/aya-vision-8b", "coherelabs/aya-vision-32b", } GATED_MODELS_REQUIRING_TOKEN = {m.lower() for m in GATED_MODELS_REQUIRING_TOKEN} NON_HF_OR_NON_TRANSFORMERS_MODELS = { "ronylicha/gigapdf-ocr-hebrew": ( "This checkpoint is not a standard Transformers OCR package (ONNX/RTen assets only)." ), "liskcell/qunie-v7-mini": ( "This checkpoint appears distributed as a non-standard GGUF/ONNX package and cannot be loaded through " "local Transformers in this Space." ), "liskcell/qunie-v7-pico": ( "This checkpoint appears distributed as a non-standard GGUF/ONNX package and cannot be loaded through " "local Transformers in this Space." ), "0cve0/openmlkitocr": ( "This repository is not structured as a standard Hugging Face Transformers OCR model." ), } MODEL_PRECHECK_CACHE: Dict[str, Tuple[bool, str]] = {} _MODEL_PRECHECK_LOCK = threading.Lock() _LOCAL_PIPELINE_CACHE: OrderedDict = OrderedDict() _LOCAL_DIRECT_CACHE: OrderedDict = OrderedDict() _LOCAL_PIPELINE_LOCK = threading.Lock() _LOCAL_DIRECT_LOCK = threading.Lock() def load_registry() -> List[dict]: with open(REGISTRY_PATH, "r", encoding="utf-8") as f: payload = json.load(f) return payload.get("models", []) def _safe_load_registry() -> List[dict]: models = load_registry() if not isinstance(models, list): raise RuntimeError("Invalid model_registry.json format: expected a list under 'models'.") for model in models: if not isinstance(model, dict): continue if "id" not in model: raise RuntimeError("Invalid registry entry: missing 'id'.") if "model_id" not in model and model.get("provider") != "tesseract": raise RuntimeError(f"Invalid registry entry '{model.get('id')}': missing 'model_id'.") return models def image_to_data_uri(image_bytes: bytes) -> str: return f"data:image/png;base64,{base64.b64encode(image_bytes).decode('ascii')}" def normalize_text_for_metrics(text: str) -> str: if text is None: return "" text = text.replace("\u200c", "") # remove Hebrew ligature marks text = re.sub(r"\s+", " ", text.strip()) return text def edit_distance(a: str, b: str) -> int: if not a: return len(b) if not b: return len(a) prev = list(range(len(b) + 1)) for i, ca in enumerate(a, 1): curr = [i] + [0] * len(b) for j, cb in enumerate(b, 1): cost = 0 if ca == cb else 1 curr[j] = min( prev[j] + 1, curr[j - 1] + 1, prev[j - 1] + cost, ) prev = curr return prev[-1] def compute_cer_wer(reference: str, hypothesis: str) -> Tuple[Optional[float], Optional[float]]: if not reference: return None, None ref_norm = normalize_text_for_metrics(reference) hyp_norm = normalize_text_for_metrics(hypothesis) ref_chars = ref_norm hyp_chars = hyp_norm ref_words = ref_norm.split(" ") if ref_norm else [] hyp_words = hyp_norm.split(" ") if hyp_norm else [] cer = edit_distance(ref_chars, hyp_chars) / max(1, len(ref_chars)) wer = edit_distance(" ".join(ref_words), " ".join(hyp_words)) / max(1, len(ref_words)) return cer, wer def parse_output(output) -> str: if output is None: return "" if isinstance(output, str): return output.strip() if isinstance(output, bytes): return output.decode("utf-8", errors="ignore").strip() if isinstance(output, list): if not output: return "" if len(output) == 1: return parse_output(output[0]) texts = [parse_output(item) for item in output if parse_output(item)] return "\n\n".join(texts) if isinstance(output, dict): for key in ("text", "generated_text", "answer", "output", "prediction"): if key in output and isinstance(output[key], str): return output[key].strip() if "choices" in output and isinstance(output["choices"], list): choice0 = output["choices"][0] if isinstance(choice0, dict) and "message" in choice0: msg = choice0["message"] if isinstance(msg, dict) and isinstance(msg.get("content"), str): return msg["content"].strip() return json.dumps(output, ensure_ascii=False) if hasattr(output, "choices"): choice = output.choices[0] if hasattr(choice, "message") and hasattr(choice.message, "content"): return (choice.message.content or "").strip() if isinstance(choice, dict) and isinstance(choice.get("message"), dict): return str(choice["message"].get("content", "")).strip() return str(output) def _sanitize_hf_token(hf_token: str) -> str: if hf_token: return hf_token.strip() return os.getenv("HF_TOKEN", "").strip() def _model_id_value(entry: dict) -> str: return str(entry.get("model_id", "") or "").strip() def _model_id_norm(entry: dict) -> str: return _model_id_value(entry).lower() def _model_requires_trust_remote_code(entry: dict) -> bool: if entry.get("trust_remote_code") is True: return True model_id = _model_id_norm(entry) if not model_id: return False if model_id in TRUST_REMOTE_CODE_MODELS: return True return False def _model_pipeline_task(entry: dict) -> str: task = (entry.get("pipeline_task") or entry.get("task") or "").strip().lower() if task and task != "auto": return task model_id = _model_id_norm(entry) return PIPELINE_TASK_HINTS.get(model_id, "auto") def _precheck_cache_key(entry: dict, hf_token: str) -> str: return f"{_model_id_norm(entry)}|{_sanitize_hf_token(hf_token) or 'anon'}" def _classify_precheck_exception(exc: Exception) -> str: text = str(exc).lower() if "gated repo" in text or "requires authentication" in text or "401" in text: return "Repository is gated. Provide a valid HF token with access." if "403" in text or "authorization" in text: return "Authorization failed for this repository. HF token may be missing or lacks access." if "qwen3_5" in text or "gemma4_unified" in text: return ( "Model architecture requires a newer Transformers runtime than this Space currently has. " "Rebuild after updating transformers from source." ) if "tokenizersbackend" in text: return ( "Tokenizer backend unavailable in this environment. Rebuild after dependency refresh " "(tokenizers / transformers)." ) if "404" in text or "not found" in text or "config.json" in text: return "Model config is not available in standard Transformers format." if "could not infer task" in text or "document-question-answering" in text: return "Model task/config mismatch for current local Transformers runtime." return _format_dependency_error(exc) def _get_model_precheck(entry: dict, hf_token: str) -> Tuple[bool, str]: provider = str(entry.get("provider", "local_transformer")).lower() if provider != "local_transformer": return True, "" model_id = _model_id_norm(entry) if not model_id: return False, "Missing model_id." if model_id in NON_HF_OR_NON_TRANSFORMERS_MODELS: return False, NON_HF_OR_NON_TRANSFORMERS_MODELS[model_id] cache_key = _precheck_cache_key(entry, hf_token) with _MODEL_PRECHECK_LOCK: cached = MODEL_PRECHECK_CACHE.get(cache_key) if cached is not None: return cached if AutoConfig is None: result = (False, "Transformers AutoConfig is not available in this runtime.") with _MODEL_PRECHECK_LOCK: MODEL_PRECHECK_CACHE[cache_key] = result return result token = _sanitize_hf_token(hf_token) or None try: # Lightweight preflight to reject incompatible repos before heavy pipeline/model loading. AutoConfig.from_pretrained( model_id, token=token, trust_remote_code=_model_requires_trust_remote_code(entry), ) except Exception as exc: # pragma: no cover result = (False, _classify_precheck_exception(exc)) with _MODEL_PRECHECK_LOCK: MODEL_PRECHECK_CACHE[cache_key] = result return result result = (True, "") with _MODEL_PRECHECK_LOCK: MODEL_PRECHECK_CACHE[cache_key] = result return result def _transformers_runtime_message() -> Optional[str]: if hf_pipeline is not None: return None if TRANSFORMERS_IMPORT_ERROR: return f"transformers import failed in this Space: {TRANSFORMERS_IMPORT_ERROR}" return "transformers is not installed in this Space." def _format_dependency_error(exc: Exception) -> str: text = str(exc) if "TokenizersBackend" in text and "not currently imported" in text: return ( "Tokenizer backend import failed. " "Rebuild after updating `tokenizers` and `transformers` in requirements." ) if "requires the following packages" in text and "addict" in text: return "Missing dependency: addict. Rebuild after adding `addict` to requirements." if "requires the following packages" in text and "torchvision" in text: return "Missing dependency: torchvision. Rebuild after adding `torchvision` to requirements." return text def _runtime_dependency_warning() -> Optional[str]: messages = [] if pytesseract is None: messages.append("pytesseract is not installed. Install with `pytesseract` in requirements.") if pytesseract is not None and not _tesseract_binary_available(): messages.append( "The tesseract binary is not available in PATH. Add tesseract packages to `apt.txt` and rebuild." ) if torch is None: messages.append("PyTorch is not installed; local models cannot run.") elif not getattr(torch, "cuda", None) or not torch.cuda.is_available(): messages.append( "PyTorch is installed but CUDA is not available in this runtime. " "For local ZeroGPU inference, rebuild with a CUDA-enabled PyTorch." ) return " | ".join(messages) if messages else None def _required_token_message(entry: dict, hf_token: str) -> Optional[str]: model_id = _model_id_norm(entry) if model_id in GATED_MODELS_REQUIRING_TOKEN and not _sanitize_hf_token(hf_token): return ( f"{entry.get('name', model_id)} appears to require an HF token in this space. " "Set HF_TOKEN to a token with access to the gated model." ) return None def _tesseract_binary_available() -> bool: import shutil for candidate in TESSERACT_EXECUTABLE_PATHS: if os.path.exists(candidate): return True if shutil.which("tesseract"): return True return False def _decode_image(image_bytes: bytes) -> Image.Image: image = Image.open(BytesIO(image_bytes)) if image.mode != "RGB": image = image.convert("RGB") return image def _load_image_bytes(image_file) -> bytes: if image_file is None: raise RuntimeError("No file was uploaded.") if isinstance(image_file, (list, tuple)): if not image_file: raise RuntimeError("No file was uploaded.") image_file = image_file[0] if isinstance(image_file, dict): image_file = ( image_file.get("path") or image_file.get("name") or image_file.get("url") or image_file.get("file_path") ) if isinstance(image_file, bytes): return image_file if hasattr(image_file, "read"): try: image_bytes = image_file.read() if image_bytes: return image_bytes except Exception as exc: raise RuntimeError(f"Could not read uploaded file object: {exc}") from exc raise RuntimeError("Uploaded file object is empty.") if isinstance(image_file, Image.Image): buffer = BytesIO() image_file.convert("RGB").save(buffer, format="PNG") return buffer.getvalue() if not isinstance(image_file, str): if hasattr(image_file, "name") and isinstance(getattr(image_file, "name"), str): image_file = getattr(image_file, "name") else: raise RuntimeError("Uploaded image is not a valid file path.") if not os.path.exists(image_file): raise RuntimeError("Uploaded file is missing from disk.") filename = os.path.basename(image_file).lower() if filename.endswith(".pdf"): raise RuntimeError( "PDF files are not supported yet in this Space build. Export the document to PNG/JPG and upload again." ) try: with open(image_file, "rb") as f: return f.read() except Exception as exc: raise RuntimeError(f"Could not read uploaded file: {exc}") from exc def _normalize_model_selection(selected_model_ids) -> List[str]: if selected_model_ids is None: return [] if isinstance(selected_model_ids, str): try: loaded = json.loads(selected_model_ids) if isinstance(loaded, list): selected_model_ids = loaded else: selected_model_ids = [selected_model_ids] except Exception: if "," in selected_model_ids: parts = [entry.strip() for entry in selected_model_ids.split(",") if entry.strip()] selected_model_ids = parts else: selected_model_ids = [selected_model_ids] elif isinstance(selected_model_ids, set): selected_model_ids = list(selected_model_ids) elif not isinstance(selected_model_ids, (list, tuple)): selected_model_ids = [str(selected_model_ids)] return [str(item) for item in selected_model_ids if item] def _build_local_pipeline(model_id: str, hf_token: str, task: str, trust_remote_code: bool = False): if hf_pipeline is None: raise RuntimeError(_transformers_runtime_message() or "transformers is not installed in this space.") token = _sanitize_hf_token(hf_token) or None device = _inference_device() strict_cuda = _strict_cuda_required() and torch is not None and torch.cuda.is_available() if strict_cuda: _ensure_cuda_available(f"Local pipeline load for {model_id}") base_kwargs = {"model": model_id, "token": token, "trust_remote_code": trust_remote_code} if device >= 0: base_kwargs["device"] = device if task and task != "auto": task_candidates = [task] else: task_candidates = [ "image-to-text", "image-text-to-text", "document-question-answering", ] seen_tasks = set() task_candidates = [t for t in task_candidates if t and not (t in seen_tasks or seen_tasks.add(t))] last_error = None preferred_dtype, dtype_chain = _inference_torch_kwargs_for_model(strict_cuda=strict_cuda) if preferred_dtype is not None: base_kwargs["torch_dtype"] = preferred_dtype for candidate_task in task_candidates: for dtype_kwargs in dtype_chain: cleaned_kwargs = {k: v for k, v in {**base_kwargs, **dtype_kwargs}.items() if v is not None} try: candidate_pipeline = hf_pipeline(candidate_task, **cleaned_kwargs) if strict_cuda: candidate_pipeline = _move_pipeline_to_cuda( candidate_pipeline, model_id, f"pipeline load for {candidate_task}", ) candidate_model = getattr(candidate_pipeline, "model", None) if candidate_model is not None and torch is not None and torch.cuda.is_available(): try: candidate_model = candidate_model.to("cuda") if hasattr(candidate_pipeline, "model"): candidate_pipeline.model = candidate_model except Exception: raise RuntimeError( f"Failed to move pipeline model for {model_id} to CUDA. " "This usually indicates ZeroGPU memory constraints." ) _assert_model_on_cuda(candidate_model, model_id, f"pipeline load for task {candidate_task}") return candidate_pipeline except Exception as exc: # pragma: no cover last_error = exc if task and task != "auto": return _build_local_pipeline(model_id, hf_token, "auto", trust_remote_code=trust_remote_code) raise RuntimeError(f"Failed to load local pipeline for {model_id}: {last_error}") def _load_direct_components(model_id: str, hf_token: str, trust_remote_code: bool = False): if AutoProcessor is None and AutoImageProcessor is None and AutoTokenizer is None: return None, None token = _sanitize_hf_token(hf_token) or None component = (model_id, trust_remote_code) with _LOCAL_DIRECT_LOCK: cached = _LOCAL_DIRECT_CACHE.get(component) if cached is not None: _LOCAL_DIRECT_CACHE.move_to_end(component) return cached processor = None processor_error: Optional[Exception] = None for processor_ctor in [AutoProcessor, AutoImageProcessor]: if processor_ctor is None: continue try: processor = processor_ctor.from_pretrained( model_id, token=token, trust_remote_code=trust_remote_code, ) break except Exception as exc: # pragma: no cover processor_error = exc if processor is None: if AutoTokenizer is not None: try: processor = AutoTokenizer.from_pretrained(model_id, token=token, trust_remote_code=trust_remote_code) except Exception as exc: # pragma: no cover if processor_error is None: processor_error = exc if processor is None: error_message = ( f"Failed to load processor for {model_id}: {processor_error}" if processor_error is not None else f"Failed to load processor for {model_id}: no compatible processor classes available." ) raise RuntimeError(error_message) strict_cuda = _strict_cuda_required() and torch is not None and torch.cuda.is_available() if strict_cuda: _ensure_cuda_available(f"Direct model load for {model_id}") if torch is None: raise RuntimeError("PyTorch is required for direct model loading.") model_errors = [] model = None preferred_dtype, model_kwargs_chain = _inference_torch_kwargs_for_model(strict_cuda=strict_cuda) if preferred_dtype is not None: base_kwargs = { "token": token, "trust_remote_code": trust_remote_code, } else: base_kwargs = { "token": token, "trust_remote_code": trust_remote_code, } model_classes = [ AutoModelForImageTextToText, AutoModelForVision2Seq, AutoModelForVisionEncoderDecoder, AutoModelForSeq2SeqLM, AutoModelForCausalLM, AutoModelForConditionalGeneration, AutoModelForDocumentQuestionAnswering, AutoModel, ] for model_class in model_classes: if model_class is None: continue try: for model_kwargs in model_kwargs_chain: candidate_kwargs = {k: v for k, v in {**base_kwargs, **model_kwargs}.items() if v is not None} try: model = model_class.from_pretrained( model_id, **candidate_kwargs, ) break except TypeError as exc: model_errors.append(f"{getattr(model_class, '__name__', str(model_class))}: {exc}") continue except Exception as exc: # pragma: no cover model_errors.append(f"{getattr(model_class, '__name__', str(model_class))}: {exc}") if model is not None: break except Exception as exc: # pragma: no cover model_errors.append(f"{getattr(model_class, '__name__', str(model_class))}: {exc}") if model is None: raise RuntimeError(f"Direct loading failed for {model_id}. " + " | ".join(model_errors)) if strict_cuda: try: model = model.to("cuda") except Exception as exc: raise RuntimeError(f"Failed to place {model_id} on CUDA: {exc}") _assert_model_on_cuda(model, model_id, "direct model load") elif torch.cuda.is_available(): model = model.to("cuda") with _LOCAL_DIRECT_LOCK: _LOCAL_DIRECT_CACHE[component] = (model, processor) if len(_LOCAL_DIRECT_CACHE) > DIRECT_MODEL_CACHE_MAX: _, evicted = _LOCAL_DIRECT_CACHE.popitem(last=False) try: model_ref, _ = evicted del model_ref except Exception: pass if torch is not None and torch.cuda.is_available(): torch.cuda.empty_cache() return model, processor def _direct_infer(model_id: str, image_bytes: bytes, prompt: str, hf_token: str, params: dict, trust_remote_code: bool = False) -> str: if torch is None: raise RuntimeError("PyTorch is required for direct model inference.") model, processor = _load_direct_components(model_id, hf_token, trust_remote_code=trust_remote_code) image = _decode_image(image_bytes) input_candidates = [] if prompt: input_candidates.extend( [ lambda: processor(text=prompt, images=image, return_tensors="pt"), lambda: processor(images=image, text=prompt, return_tensors="pt"), lambda: processor(prompt, image, return_tensors="pt"), lambda: processor(prompt, return_tensors="pt"), ] ) input_candidates.append(lambda: processor(images=image, return_tensors="pt")) prepared_inputs = None prep_error = None for builder in input_candidates: try: candidate = builder() if isinstance(candidate, Mapping) and candidate: prepared_inputs = {} for key, value in candidate.items(): if hasattr(value, "to"): prepared_inputs[key] = value.to(model.device) elif isinstance(value, (list, tuple)): prepared_inputs[key] = value if "images" in prepared_inputs and "pixel_values" not in prepared_inputs: prepared_inputs["pixel_values"] = prepared_inputs["images"] prepared_inputs.pop("images", None) if prepared_inputs: break except Exception as exc: # pragma: no cover prep_error = exc if prepared_inputs is None: raise RuntimeError(f"Could not prepare inputs for {model_id}: {prep_error}") gen_kwargs = {} if "max_new_tokens" in params: gen_kwargs["max_new_tokens"] = params["max_new_tokens"] if "temperature" in params: gen_kwargs["temperature"] = params["temperature"] if "top_p" in params: gen_kwargs["top_p"] = params["top_p"] if "top_k" in params: gen_kwargs["top_k"] = params["top_k"] if "do_sample" in params: gen_kwargs["do_sample"] = params["do_sample"] if "num_beams" in params: gen_kwargs["num_beams"] = params["num_beams"] if not gen_kwargs and hasattr(model, "generation_config"): try: gen_kwargs = { "max_new_tokens": getattr(model.generation_config, "max_new_tokens", None), "temperature": getattr(model.generation_config, "temperature", None), "top_p": getattr(model.generation_config, "top_p", None), "top_k": getattr(model.generation_config, "top_k", None), "num_beams": getattr(model.generation_config, "num_beams", None), } gen_kwargs = {k: v for k, v in gen_kwargs.items() if v is not None} except Exception: gen_kwargs = {} with torch.no_grad(): generated = model.generate(**prepared_inputs, **gen_kwargs) if isinstance(generated, torch.Tensor): decoded = processor.batch_decode(generated, skip_special_tokens=True) return parse_output(decoded) if isinstance(generated, (list, tuple)): return parse_output(generated) return parse_output(str(generated)) def _get_local_pipeline(model_id: str, hf_token: str, task: str, trust_remote_code: bool = False): device = "gpu" if (torch is not None and torch.cuda.is_available()) else "cpu" dtype = _inference_dtype() key = _pipeline_cache_key(model_id, task or "auto", device, dtype, trust_remote_code) with _LOCAL_PIPELINE_LOCK: cached = _LOCAL_PIPELINE_CACHE.get(key) if cached is not None: _LOCAL_PIPELINE_CACHE.move_to_end(key) return cached pipeline = _build_local_pipeline(model_id, hf_token, task, trust_remote_code=trust_remote_code) with _LOCAL_PIPELINE_LOCK: _LOCAL_PIPELINE_CACHE[key] = pipeline if len(_LOCAL_PIPELINE_CACHE) > LOCAL_PIPELINE_CACHE_MAX: _, evicted = _LOCAL_PIPELINE_CACHE.popitem(last=False) try: model = getattr(evicted, "model", None) if model is not None: del model del evicted except Exception: pass if torch is not None and torch.cuda.is_available(): torch.cuda.empty_cache() gc.collect() gc.collect() return pipeline def _clear_local_pipeline_cache() -> None: global _LOCAL_PIPELINE_CACHE with _LOCAL_PIPELINE_LOCK: pipelines = list(_LOCAL_PIPELINE_CACHE.values()) _LOCAL_PIPELINE_CACHE.clear() for pipeline in pipelines: try: if hasattr(pipeline, "model"): del pipeline.model if hasattr(pipeline, "processor"): del pipeline.processor if hasattr(pipeline, "tokenizer"): del pipeline.tokenizer except Exception: pass gc.collect() if torch is not None and torch.cuda.is_available(): torch.cuda.empty_cache() def _clear_direct_cache() -> None: global _LOCAL_DIRECT_CACHE with _LOCAL_DIRECT_LOCK: directs = list(_LOCAL_DIRECT_CACHE.values()) _LOCAL_DIRECT_CACHE.clear() for model, _ in directs: try: del model except Exception: pass gc.collect() if torch is not None and torch.cuda.is_available(): torch.cuda.empty_cache() def run_local_model(entry: dict, image_bytes: bytes, prompt: str, hf_token: str) -> str: if hf_pipeline is None: raise RuntimeError(_transformers_runtime_message() or "transformers is not installed in this space.") if torch is None: raise RuntimeError("PyTorch is required for local model inference.") model_id = entry["model_id"] task = _model_pipeline_task(entry) trust_remote_code = _model_requires_trust_remote_code(entry) if _strict_cuda_required(): _ensure_cuda_available(f"Local OCR execution for {model_id}") try: pipeline = _get_local_pipeline(model_id, hf_token, task, trust_remote_code=trust_remote_code) except Exception as exc: pipeline = None pipeline_error = exc else: pipeline_error = None params = sanitize_generation_params(entry.get("parameters", {})) image = _decode_image(image_bytes) if "extra_body" in params and isinstance(params["extra_body"], Mapping): body = params.pop("extra_body") if isinstance(body, Mapping): supported_body_keys = { "top_p", "top_k", "temperature", "do_sample", "num_beams", "repetition_penalty", "max_new_tokens", "max_length", "min_length", "enable_thinking", "seed", "return_dict_in_generate", "output_scores", } for key, value in body.items(): if isinstance(key, str) and key in supported_body_keys: params.setdefault(key, value) if pipeline is not None: try: if torch is not None and torch.cuda.is_available(): pipeline = _move_pipeline_to_cuda( pipeline, model_id, "pipeline execution", ) if _strict_cuda_required(): _assert_model_on_cuda(getattr(pipeline, "model", None), model_id, "pipeline execution") pipeline_model = getattr(pipeline, "model", None) if pipeline_model is not None and torch is not None and torch.cuda.is_available(): pipeline_model = pipeline_model.to("cuda") if hasattr(pipeline, "model"): pipeline.model = pipeline_model else: pipeline_model = getattr(pipeline, "model", None) if pipeline_model is not None and torch is not None and torch.cuda.is_available(): pipeline_model = pipeline_model.to("cuda") if hasattr(pipeline, "model"): pipeline.model = pipeline_model except Exception as exc: if not _strict_cuda_required(): pipeline_error = _format_dependency_error(exc) else: pipeline = None pipeline_error = _format_dependency_error(exc) calls = [] if pipeline is not None: if prompt: calls.extend( [ lambda: pipeline({"image": image, "text": prompt}, **params), lambda: pipeline({"image": image, "question": prompt}, **params), lambda: pipeline({"text": prompt, "image": image}, **params), lambda: pipeline({"images": image, "text": prompt}, **params), lambda: pipeline(image, text=prompt, **params), lambda: pipeline(image, question=prompt, **params), lambda: pipeline(image, **params), lambda: pipeline({"image": image}, **params), ] ) else: calls.extend([lambda: pipeline(image, **params), lambda: pipeline({"image": image}, **params)]) for call in calls: try: output = call() parsed = parse_output(output) if parsed: return parsed except Exception as exc: # pragma: no cover pipeline_error = _format_dependency_error(exc) break if pipeline is None: pipeline_error = pipeline_error or "Pipeline execution failed." try: direct_output = _direct_infer( model_id, image_bytes, prompt, hf_token, params, trust_remote_code=trust_remote_code, ) parsed = parse_output(direct_output) if parsed: return parsed except Exception as exc: # pragma: no cover if pipeline_error is None: pipeline_error = _format_dependency_error(exc) else: pipeline_error = RuntimeError(f"{pipeline_error}; direct fallback failed: {_format_dependency_error(exc)}") raise RuntimeError(f"Local model inference failed for {model_id}: {pipeline_error}") def run_tesseract(image_bytes: bytes, params: dict) -> str: if pytesseract is None: raise RuntimeError("pytesseract is not installed in this Space.") if not _tesseract_binary_available(): raise RuntimeError( "pytesseract is installed but the tesseract executable is not available in PATH. " "The Space should install tesseract via apt.txt, please confirm a clean rebuild." ) for candidate in TESSERACT_EXECUTABLE_PATHS: if os.path.exists(candidate): try: pytesseract.pytesseract.tesseract_cmd = candidate except Exception: pass image = Image.open(BytesIO(image_bytes)) lang = params.get("lang", "heb+eng") psm = params.get("psm", 6) oem = params.get("oem", 1) extra = params.get("extra_config", "").strip() cfg_bits = [f"--psm {int(psm)}", f"--oem {int(oem)}"] if extra: cfg_bits.append(extra) config = " ".join(cfg_bits).strip() return pytesseract.image_to_string(image, lang=lang, config=config).strip() @dataclass class RunnerResult: model_id: str label: str provider: str status: str output: str latency_sec: Optional[float] char_count: int cer: Optional[float] wer: Optional[float] notes: str def sanitize_generation_params(raw: dict) -> dict: params = {} if not raw: return {"max_new_tokens": DEFAULT_MAX_TOKENS} for k, v in raw.items(): if v is None: continue if k == "max_new_tokens": requested = int(v) params["max_new_tokens"] = max(DEFAULT_MAX_TOKENS, requested) elif k == "max_tokens": requested = int(v) params["max_new_tokens"] = max(DEFAULT_MAX_TOKENS, requested) elif k == "temperature": params["temperature"] = float(v) elif k == "top_p": params["top_p"] = float(v) elif k == "top_k": params["top_k"] = int(v) elif k == "repetition_penalty": params["repetition_penalty"] = float(v) elif k == "frequency_penalty": params["frequency_penalty"] = float(v) elif k == "presence_penalty": params["presence_penalty"] = float(v) elif k == "seed": params["seed"] = int(v) elif k == "extra_body": params["extra_body"] = v else: params[k] = v return params def run_single_model(entry: dict, image_bytes: bytes, data_uri: str, hf_token: str, timeout_sec: int = 120) -> RunnerResult: provider = entry.get("provider", "local_transformer").lower() if provider != "tesseract": provider = "local_transformer" model_id = entry["model_id"] label = entry.get("name", model_id) prompt = entry.get("prompt", DEFAULT_OCR_PROMPT) params = sanitize_generation_params(entry.get("parameters", {})) notes = entry.get("notes", "") runtime = _get_runtime_profile() runtime_tag = f"runtime: {_inference_device_label()} | provider: {runtime.get('provider')}" notes = f"{notes} [{runtime_tag}]" if notes else runtime_tag start = time.perf_counter() output = "" try: token_message = _required_token_message(entry, hf_token) if token_message: raise RuntimeError(token_message) if provider == "tesseract": output = run_tesseract(image_bytes, params) status = "ok" elif provider in {"hf_chat", "hf_image_to_text", "local_transformer", "local"}: output = run_local_model(entry, image_bytes, prompt, hf_token) status = "ok" else: raise RuntimeError(f"Unsupported provider '{provider}'") except Exception as exc: # pragma: no cover status = "error" output = f"{type(exc).__name__}: {exc}" duration = round(time.perf_counter() - start, 3) output = output or "" return RunnerResult( model_id=entry["id"], label=label, provider=provider, status=status, output=output, latency_sec=duration, char_count=len(output), cer=None, wer=None, notes=notes, ) def _get_runtime_profile() -> Dict[str, object]: profile = { "has_cuda": False, "torch_version": None, "torch_cuda_version": None, "torch_cuda_built": None, "gpu_name": None, "gpu_vram_gb": None, "gpu_free_vram_gb": None, "gpu_reserved_vram_gb": None, "gpu_allocated_vram_gb": None, "provider": "cpu", } env_hardware = os.getenv("SPACE_HARDWARE", "").lower() or os.getenv("HF_SPACE_HARDWARE", "").lower() if env_hardware: profile["provider"] = env_hardware try: import torch profile["torch_version"] = getattr(torch, "__version__", None) profile["torch_cuda_built"] = bool(getattr(getattr(torch, "version", None), "cuda", None)) profile["torch_cuda_version"] = getattr(getattr(torch, "version", None), "cuda", None) if torch.cuda.is_available(): profile["has_cuda"] = True if not profile["provider"] or profile["provider"] == "cpu": profile["provider"] = "cuda" props = torch.cuda.get_device_properties(0) profile["gpu_name"] = props.name profile["gpu_vram_gb"] = round(props.total_memory / (1024 ** 3), 1) try: free_mem, total_mem = torch.cuda.mem_get_info(0) profile["gpu_free_vram_gb"] = round(free_mem / (1024 ** 3), 1) profile["gpu_vram_gb"] = round(total_mem / (1024 ** 3), 1) except Exception: free_approx = max(0, props.total_memory - torch.cuda.memory_reserved(0)) profile["gpu_free_vram_gb"] = round(free_approx / (1024 ** 3), 1) profile["gpu_reserved_vram_gb"] = round(torch.cuda.memory_reserved(0) / (1024 ** 3), 1) profile["gpu_allocated_vram_gb"] = round(torch.cuda.memory_allocated(0) / (1024 ** 3), 1) provider = str(profile["provider"] or "").lower() normalized_name = str(props.name or "").lower() if "zero" in provider or "a10g" in normalized_name or "l4" in normalized_name: profile["provider"] = "zero_gpu" except Exception: pass return profile def _runtime_capacity_note(runtime: Dict[str, object]) -> str: if not runtime.get("has_cuda"): return "Runtime compute mode: CPU-only." if runtime.get("gpu_name"): total = runtime.get("gpu_vram_gb") free = runtime.get("gpu_free_vram_gb") allocated = runtime.get("gpu_allocated_vram_gb") reserved = runtime.get("gpu_reserved_vram_gb") if free is None: return f"Runtime compute mode: {runtime['gpu_name']} with {total}GB total VRAM." return ( f"Runtime compute mode: {runtime['gpu_name']} | total {total}GB | " f"free {free}GB | allocated {allocated}GB | reserved {reserved}GB." ) return "Runtime compute mode: GPU detected but profile unavailable." def _startup_cuda_audit() -> None: runtime = _get_runtime_profile() print( "[startup] Runtime: " f"torch={runtime.get('torch_version')} " f"cuda={runtime.get('torch_cuda_version')} " f"cuda_built={runtime.get('torch_cuda_built')}" ) if runtime.get("has_cuda"): provider = runtime.get("provider") if str(provider).lower() == "cpu": provider = "cuda" runtime["provider"] = provider print( "[startup] Compute backend: " f"{runtime.get('provider')} · {runtime.get('gpu_name')} · total {runtime.get('gpu_vram_gb')}GB · " f"free {runtime.get('gpu_free_vram_gb')}GB" ) else: print("[startup] WARNING: CUDA is not available in this runtime.") def _estimate_model_vram_gb(entry: dict, *, compute_mode: str) -> float: if not isinstance(entry, dict): return ZERO_GPU_MODEL_FALLBACK_GB_SAFE provider = (entry.get("provider") or "").lower() if provider == "tesseract": return 0.2 model_override = entry.get("estimated_vram_gb") if isinstance(model_override, (int, float)) and model_override > 0: return float(model_override) model_id = (entry.get("model_id", "") or "").lower() fallback = ZERO_GPU_MODEL_FALLBACK_GB_MAX if compute_mode == "max" else ZERO_GPU_MODEL_FALLBACK_GB_SAFE if "31b" in model_id or "32b" in model_id: return 22.0 if "30b_a3b" in model_id: return 20.0 if "26b" in model_id: return 16.0 if "12b" in model_id: return 8.0 if "8b" in model_id: return 6.0 if "4b" in model_id or "e4b" in model_id: return 4.0 if "a3b" in model_id: return 10.0 if "3b" in model_id: return 2.5 return fallback def _execution_plan(selected_entries: List[dict], compute_mode: str = "safe") -> Tuple[int, str]: runtime = _get_runtime_profile() selected_count = len(selected_entries) capacity_note = _runtime_capacity_note(runtime) if selected_count <= 1: return 1, capacity_note mode_is_max = compute_mode == "max" estimates = [_estimate_model_vram_gb(entry, compute_mode=compute_mode) for entry in selected_entries] estimated_gb = round(sum(estimates), 1) provider = runtime.get("provider") notes = [] max_workers = min(ZERO_GPU_MAX_WORKERS, selected_count) gpu_name = str(runtime.get("gpu_name", "")).lower() is_zero_gpu = ( provider == "zero_gpu" or "zero-a10g" in str(provider) or "a10g" in gpu_name or "l4" in gpu_name ) if runtime.get("has_cuda") and not is_zero_gpu: notes.append("GPU runtime detected; applying generic CUDA concurrency sizing.") if is_zero_gpu: if compute_mode in {"sequential", "safe"}: max_workers = 1 notes.append("Safe/Sequential mode runs one model at a time on ZeroGPU.") elif runtime.get("gpu_free_vram_gb") and runtime.get("gpu_vram_gb") and selected_count: # Keep a hard headroom + runtime overhead to avoid zeroGPU OOM/rate failures. headroom = float(runtime["gpu_free_vram_gb"]) - ZERO_GPU_MIN_GUARDED_FREE_GB headroom = max(0.0, headroom) headroom *= ZERO_GPU_MAX_HEADROOM if mode_is_max else ZERO_GPU_SAFE_HEADROOM sorted_estimates = sorted(estimates) fit_workers = 0 running = 0.0 for est in sorted_estimates: # Conservative per-run margin for HTTP/image payload/responses/runtime overhead. needed = est + 0.8 if fit_workers < ZERO_GPU_MAX_WORKERS and fit_workers + 1 <= selected_count and running + needed <= headroom: running += needed fit_workers += 1 else: break if fit_workers <= 0 and sorted_estimates: fit_workers = 1 if fit_workers < selected_count: notes.append( f"ZeroGPU {'max' if mode_is_max else 'safe'} mode: {fit_workers}/{selected_count} models can be " f"run in parallel with current free memory." ) max_workers = min(max_workers, fit_workers) elif runtime.get("gpu_vram_gb") and estimated_gb >= runtime["gpu_vram_gb"] * 0.7: notes.append( f"Estimated total selected model VRAM ({estimated_gb:.1f}GB) is above 70% of available GPU ({runtime['gpu_vram_gb']}GB)." ) max_workers = 1 if mode_is_max and max_workers > 1 and selected_count > 6: # Keep API call fanout bounded for many queued tasks. max_workers = min(max_workers, 2) notes.append("Max-compute mode capped at 2 concurrent workers when many models are selected.") if max_workers > ZERO_GPU_MAX_SAFE_SELECTED: max_workers = ZERO_GPU_MAX_SAFE_SELECTED if runtime.get("has_cuda") and estimated_gb and estimated_gb > 0: if estimated_gb > 24: notes.append( f"Estimated total model VRAM {estimated_gb:.1f}GB is high. Running sequentially to be conservative." ) max_workers = 1 if selected_count > ZERO_GPU_MAX_SAFE_SELECTED and is_zero_gpu: notes.append("Running many models at once increases timeout risk on ZeroGPU.") notes.append(capacity_note) return max_workers, " | ".join(notes) def run_comparison( image_file, selected_model_ids, ground_truth_text, hf_token, compute_mode, ) -> Tuple[str, str, str]: try: models = _safe_load_registry() except Exception as exc: # pragma: no cover return ( "Configuration error: failed to load model registry.", "[]", json.dumps({"error": str(exc)}, ensure_ascii=False), ) if not image_file: return "Upload an image first.", "[]", json.dumps({"error": "No image provided"}, ensure_ascii=False) try: image_bytes = _load_image_bytes(image_file) except Exception as exc: error_text = f"Failed to load uploaded image: {exc}" return error_text, "[]", json.dumps({"error": error_text}, ensure_ascii=False) selected_model_ids = _normalize_model_selection(selected_model_ids) if not selected_model_ids: selected_model_ids = [] results: List[RunnerResult] = [] selected = set(selected_model_ids) selected_entries = [] for entry in models: if entry["id"] not in selected: continue if not entry.get("enabled", True): continue precheck_ok, precheck_message = _get_model_precheck(entry, hf_token) if not precheck_ok: entry_provider = entry.get("provider", "local_transformer") if entry_provider != "tesseract": entry_provider = "local_transformer" results.append( RunnerResult( model_id=entry["id"], label=entry.get("name", entry["id"]), provider=entry_provider, status="unsupported", output=precheck_message, latency_sec=None, char_count=0, cer=None, wer=None, notes=entry.get("notes", ""), ) ) continue selected_entries.append(entry) if spaces is not None and selected_entries: if _strict_cuda_required(): _ensure_zero_gpu_lease() else: try: _ensure_zero_gpu_lease() except Exception: pass max_workers, execution_warning = _execution_plan(selected_entries, compute_mode=compute_mode) hf_token = (hf_token or os.getenv("HF_TOKEN") or "").strip() futures_map = {} with ThreadPoolExecutor(max_workers=max_workers) as executor: for entry in models: if not entry.get("enabled", True): entry_provider = entry.get("provider", "local_transformer") if entry_provider != "tesseract": entry_provider = "local_transformer" results.append( RunnerResult( model_id=entry["id"], label=entry.get("name", entry["id"]), provider=entry_provider, status="disabled", output="", latency_sec=None, char_count=0, cer=None, wer=None, notes=entry.get("notes", ""), ) ) continue if entry["id"] not in selected: entry_provider = entry.get("provider", "local_transformer") if entry_provider != "tesseract": entry_provider = "local_transformer" results.append( RunnerResult( model_id=entry["id"], label=entry.get("name", entry["id"]), provider=entry_provider, status="skipped", output="", latency_sec=None, char_count=0, cer=None, wer=None, notes=entry.get("notes", ""), ) ) continue if entry not in selected_entries: continue futures_map[executor.submit(run_single_model, entry, image_bytes, "", hf_token)] = entry for future in as_completed(futures_map): try: result = future.result() if ground_truth_text: cer, wer = compute_cer_wer(ground_truth_text, result.output) result.cer = cer result.wer = wer results.append(result) except Exception as exc: # pragma: no cover results.append( RunnerResult( model_id="__worker_error__", label="Worker failure", provider="local_transformer", status="error", output=f"Unexpected worker failure: {exc}", latency_sec=None, char_count=0, cer=None, wer=None, notes="", ) ) order_map = {model["id"]: idx for idx, model in enumerate(models)} results.sort(key=lambda r: order_map.get(r.model_id, 1_000_000)) summary = [] if execution_warning: summary.append(f"**Execution mode:** {execution_warning}") for result in results: if result.status == "ok": metrics = [] if result.cer is not None: metrics.append(f"CER {result.cer:.4f}") if result.wer is not None: metrics.append(f"WER {result.wer:.4f}") metric_str = " | ".join(metrics) if metrics else "N/A" summary.append( f"- **{result.label}**: {result.status} · {result.latency_sec:.2f}s · {result.char_count} chars · {metric_str}" ) else: summary.append(f"- **{result.label}**: {result.status} · {result.notes or result.output[:120]}") table_rows = [] for result in results: row_output = result.output if len(row_output) > 400: row_output = row_output[:397] + "..." table_rows.append( { "Model": result.label, "Provider": result.provider, "Status": result.status, "Time (s)": result.latency_sec, "Chars": result.char_count, "CER": result.cer, "WER": result.wer, "Output preview": row_output, "Notes": result.notes, } ) json_payload = [r.__dict__ for r in results] markdown = "# OCR comparison results\n\n" + ("\n".join(summary) if summary else "No models selected.") return ( markdown, json.dumps(table_rows, ensure_ascii=False, indent=2), json.dumps(json_payload, ensure_ascii=False, indent=2), ) def refresh_model_choices(): try: models = _safe_load_registry() except Exception: return [("Model registry is unavailable", "__registry_error__")], [] options = [(m["name"], m["id"]) for m in models if m.get("enabled", True)] # Keep CPU-only OCR baseline off by default so GPU comparison runs do not # unintentionally spend time on CPU-only inference. default_checked = [ m["id"] for m in models if m.get("enabled", True) and str(m.get("provider", "")).lower() != "tesseract" ] return options, default_checked def build_ui(): options, default_checked = refresh_model_choices() runtime_warning = _transformers_runtime_message() dependency_warning = _runtime_dependency_warning() if runtime_warning is None: runtime_warning = dependency_warning elif dependency_warning: runtime_warning = f"{runtime_warning}\n\n{dependency_warning}" with gr.Blocks(title="Hebrew/English OCR Model Comparison (Zero GPU)") as demo: gr.Markdown( """ # Private Zero-GPU HF Space: Hebrew OCR Comparator Upload one document image and run only the models you check. The benchmark is built for **Hebrew documents containing printed + handwritten text**. Unchecked models will be skipped and **not executed**. All OCR inference is executed locally in this Space on ZeroGPU (no external inference API calls). By default, CPU fallback is disabled so models that cannot stay on CUDA fail with clear errors. To allow CPU fallback (slower, for compatibility only), set `ALLOW_CPU_FALLBACK=1`. Tesseract is CPU-only and is disabled by default. Enable it only when you need the baseline check. """ ) if runtime_warning: gr.Markdown(f"### Runtime warning\n{runtime_warning}") with gr.Row(): with gr.Column(scale=2): image_input = gr.Image(type="filepath", label="Document image (printed and handwritten Hebrew, English optional)") hf_token = gr.Textbox( label="HF_TOKEN (optional)", type="password", value=os.getenv("HF_TOKEN", ""), info="Needed only for private/gated model downloads into this Space.", ) ground_truth = gr.Textbox( label="Ground truth (optional)", lines=5, placeholder="Paste exact expected text for CER/WER", ) with gr.Column(scale=3): compute_mode = gr.Radio( label="Compute mode", choices=[ ("Safe (recommended, avoid overrun)", "safe"), ("Max compute (faster)", "max"), ("Sequential only", "sequential"), ], value="safe", ) model_selector = gr.CheckboxGroup( label="Models to run", choices=options, value=default_checked, interactive=True, ) run_btn = gr.Button("Run selected models", variant="primary") with gr.Row(): status_md = gr.Markdown("## Results") results_md = gr.Markdown("") results_df = gr.JSON( label="Result rows (JSON)", ) results_rows_json = gr.Textbox( label="Result rows (JSON)", lines=12, interactive=False, ) raw_json = gr.Textbox( label="Full outputs (JSON)", lines=12, interactive=False, ) run_btn.click( fn=run_comparison, inputs=[image_input, model_selector, ground_truth, hf_token, compute_mode], outputs=[results_md, results_rows_json, raw_json], api_name="run_comparison", ) return demo if __name__ == "__main__": if not os.path.exists(REGISTRY_PATH): raise RuntimeError("Missing model_registry.json. Create it before running the app.") _startup_cuda_audit() if _strict_cuda_required(): _ensure_cuda_available("Application startup") demo = build_ui() port = int(os.getenv("PORT", "7860")) demo.launch(server_name="0.0.0.0", server_port=port)