app.py

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)