src/inference/model_loader.py

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path


@dataclass(frozen=True)
class LoadedModelBundle:
    model: object
    input_shape: tuple[int, int, int, int]
    config_path: Path
    weights_path: Path
    model_dir: Path


@dataclass(frozen=True)
class ModelArtifactPaths:
    model_dir: Path
    config_path: Path
    weights_path: Path


def _choose_latest(paths: list[Path]) -> Path:
    if not paths:
        raise FileNotFoundError("No candidate model files were found.")
    return sorted(paths, key=lambda p: p.stat().st_mtime_ns)[-1]


def resolve_model_artifacts(model_dir: Path) -> ModelArtifactPaths:
    model_dir = Path(model_dir).expanduser().resolve()
    if not model_dir.exists():
        raise FileNotFoundError(f"Model directory not found: {model_dir}")

    config_path = model_dir / "config.json"
    if not config_path.exists():
        raise FileNotFoundError(f"Missing model config: {config_path}")

    candidates: list[Path] = []
    candidates.extend(sorted(model_dir.glob("*.final.weights.h5")))

    named_best = model_dir / "best_model_dynamic.weights.h5"
    if named_best.exists():
        candidates.append(named_best)

    candidates.extend(sorted(model_dir.glob("*.weights.h5")))

    # Deduplicate while preserving order.
    seen = set()
    deduped: list[Path] = []
    for path in candidates:
        key = str(path.resolve())
        if key in seen:
            continue
        seen.add(key)
        deduped.append(path)

    if not deduped:
        raise FileNotFoundError(
            "No weights file found in model directory. "
            "Expected something like '*.final.weights.h5' or '*.weights.h5'."
        )

    weights_path = _choose_latest(deduped)
    return ModelArtifactPaths(
        model_dir=model_dir,
        config_path=config_path,
        weights_path=weights_path,
    )


def _read_model_config(config_path: Path) -> dict:
    payload = json.loads(config_path.read_text())
    input_shape = payload.get("INPUT_SHAPE")
    if not input_shape:
        raise RuntimeError(f"INPUT_SHAPE missing in model config: {config_path}")
    return payload


def _should_force_cpu() -> bool:
    raw = os.environ.get("INFER_FORCE_CPU", "0").strip().lower()
    return raw in {"1", "true", "yes", "on"}


@lru_cache(maxsize=4)
def _load_model_cached(
    model_dir_str: str,
    config_mtime_ns: int,
    weights_mtime_ns: int,
) -> LoadedModelBundle:
    del config_mtime_ns, weights_mtime_ns

    if _should_force_cpu():
        os.environ.setdefault("CUDA_VISIBLE_DEVICES", "-1")

    from tensorflow.keras import mixed_precision

    from .model_defs import ModelBuildConfig, build_dynamic_model

    artifacts = resolve_model_artifacts(Path(model_dir_str))
    payload = _read_model_config(artifacts.config_path)

    # Force float32 inference graph for CPU compatibility.
    mixed_precision.set_global_policy("float32")

    cfg = ModelBuildConfig(
        input_shape=tuple(int(x) for x in payload["INPUT_SHAPE"]),
        base_filters=int(payload.get("BASE_FILTERS", 8)),
        sam_heads=int(payload.get("SAM_HEADS", 2)),
        l2_reg=float(payload.get("L2_REG", 0.0)),
    )

    model = build_dynamic_model(cfg)
    model.load_weights(str(artifacts.weights_path))

    return LoadedModelBundle(
        model=model,
        input_shape=cfg.input_shape,
        config_path=artifacts.config_path,
        weights_path=artifacts.weights_path,
        model_dir=artifacts.model_dir,
    )


def load_model_bundle(model_dir: Path) -> LoadedModelBundle:
    artifacts = resolve_model_artifacts(model_dir)
    return _load_model_cached(
        str(artifacts.model_dir),
        artifacts.config_path.stat().st_mtime_ns,
        artifacts.weights_path.stat().st_mtime_ns,
    )