inference/src/ax_doclayout/model.py

# Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
"""Load a verified immutable Hub bundle and detect layout on RGB pages."""

from __future__ import annotations

import importlib
import shutil
from collections.abc import Callable
from pathlib import Path
from tempfile import TemporaryDirectory

from ax_doclayout.artifact import MODEL_ID, ArtifactReceipt
from ax_doclayout.contract import (
    HUB_COMMIT_RE,
    Element,
    ModelIdentity,
    PageImage,
    PageStructure,
)
from ax_doclayout.errors import (
    ArtifactError,
    ArtifactUnavailableError,
    AXDocLayoutError,
    BackendUnavailableError,
    InferenceError,
    InvalidPageImageError,
    InvalidRevisionError,
)
from ax_doclayout.hub_release import PINNED_HUB_RECEIPT, qualify_published_bundle
from ax_doclayout.postprocess import page_structure_from_elements

BACKENDS = ("pytorch", "onnx")
_PredictElements = Callable[[PageImage], tuple[Element, ...]]


def _download_bundle(directory: Path, *, receipt: ArtifactReceipt) -> None:
    """Materialize regular files from the Hub cache before strict verification."""

    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        raise BackendUnavailableError("huggingface-hub is not installed") from None
    try:
        for item in receipt.artifact_files:
            cached = hf_hub_download(
                repo_id=receipt.model_id,
                filename=item.name,
                revision=receipt.model_revision,
                token=False,
            )
            # Hub cache entries may be symlinks. Only the copied, hash-verified
            # regular files are handed to artifact validation and the backend.
            shutil.copyfile(cached, directory / item.name)
    except Exception:  # noqa: BLE001 - redact third-party download error details
        raise ArtifactUnavailableError("hub_bundle_download_failed") from None


def _load_backend(directory: Path, *, backend: str) -> _PredictElements:
    """Load all backend state before the temporary verified bundle is removed."""

    if backend == "pytorch":
        from ax_doclayout._local_pytorch import LocalPytorchBundle

        return LocalPytorchBundle(directory).predict_elements

    import onnxruntime as ort

    from ax_doclayout._cuda_onnx import create_cuda_session
    from ax_doclayout._local_onnx import LocalOnnxBundle

    session = create_cuda_session(directory / "model.onnx", ort=ort)
    runner = LocalOnnxBundle(
        directory, session_run=lambda feeds: session.run(None, dict(feeds)),
    )

    def predict_elements(page: PageImage) -> tuple[Element, ...]:
        return runner.predict(page).elements

    return predict_elements


class AXDocLayout:
    """A trusted Hub artifact with explicit CUDA ONNX or CPU PyTorch inference."""

    def __init__(self, model_id: str, revision: str, backend: str) -> None:
        self.model_id = model_id
        self.revision = revision
        self.backend = backend
        self._predict_elements: _PredictElements | None = None
        self._identity: ModelIdentity | None = None

    @classmethod
    def from_pretrained(
        cls,
        model_id: str,
        *,
        revision: str,
        backend: str | None = None,
    ) -> AXDocLayout:
        if not isinstance(revision, str) or HUB_COMMIT_RE.fullmatch(revision) is None:
            raise InvalidRevisionError(
                "revision must be a 40-character immutable Hub commit; "
                "'main', a tag, or 'latest' is not accepted"
            )
        if not isinstance(backend, str) or backend not in BACKENDS:
            raise BackendUnavailableError(
                "backend must be given explicitly as one of pytorch, onnx"
            )
        receipt = PINNED_HUB_RECEIPT
        if not isinstance(model_id, str) or model_id != MODEL_ID:
            raise ArtifactError("artifact_model_id_invalid")
        if receipt is None or revision != receipt.model_revision:
            raise ArtifactUnavailableError("no trusted artifact for this revision")
        try:
            dependencies = ("torch", "safetensors") if backend == "pytorch" else ("onnxruntime",)
            for dependency in dependencies:
                try:
                    importlib.import_module(dependency)
                except (ImportError, OSError):
                    raise BackendUnavailableError("requested backend is not installed") from None
            with TemporaryDirectory(prefix="ax-doclayout-hub-") as temporary:
                directory = Path(temporary)
                _download_bundle(directory, receipt=receipt)
                qualified = qualify_published_bundle(
                    model_id=model_id, revision=revision, backend=backend,
                    receipt=receipt, bundle=directory,
                )
                if qualified is None:
                    raise ArtifactUnavailableError("no trusted artifact for this revision")
                predictor = _load_backend(qualified, backend=backend)
        except AXDocLayoutError:
            raise
        except Exception:  # noqa: BLE001 - keep public errors content-free
            raise ArtifactError("hub_backend_load_failed") from None
        model = cls(model_id, revision, backend)
        model._predict_elements = predictor
        model._identity = ModelIdentity(
            id=receipt.model_id, revision=receipt.model_revision,
            artifact_sha256=receipt.artifact_sha256.removeprefix("sha256:"),
        )
        return model

    def predict(self, page: PageImage) -> PageStructure:
        if self._predict_elements is None or self._identity is None:
            raise ArtifactUnavailableError("no loaded artifact; construct through from_pretrained")
        if not isinstance(page, PageImage):
            raise InvalidPageImageError("page must be a PageImage")
        try:
            elements = self._predict_elements(page)
            return page_structure_from_elements(
                page=page, elements=elements, identity=self._identity,
            )
        except AXDocLayoutError:
            raise
        except Exception:  # noqa: BLE001 - keep public errors content-free
            raise InferenceError("page_inference_failed") from None