app/model_loader.py

from __future__ import annotations

import logging
import os
from dataclasses import dataclass
from typing import Optional

import numpy as np
from huggingface_hub import snapshot_download

from .config import AppConfig
from .model_config import (
    FacedetectionModelConfig,
    SegmentationModelConfig,
    load_facedetection_config,
    load_segmentation_config,
)
from .ort_pool import OrtSessionPool, build_session


@dataclass
class LoadedModels:
    """Container for loaded model pools and typed model configs."""

    segmentation_pool: OrtSessionPool
    segmentation_config: SegmentationModelConfig
    facedetection_pool: Optional[OrtSessionPool]
    facedetection_config: Optional[FacedetectionModelConfig]
    textmask_pool: Optional[OrtSessionPool]
    cuda_stream_ptrs: Optional[dict[str, int]]
    cuda_stream_handles: Optional[dict[str, object]]


def _download_repo(
    repo_id: str,
    revision: Optional[str],
    token: Optional[str],
    local_files_only: bool,
    allow_patterns: list[str],
) -> str:
    """Download selected model files from a configured Hugging Face repo snapshot."""
    if not repo_id:
        raise ValueError("HF repo_id must be configured")
    return snapshot_download(
        repo_id=repo_id,
        revision=revision,
        token=token,
        allow_patterns=allow_patterns,
        local_files_only=local_files_only,
    )


def _create_pool(
    onnx_path: str,
    config: AppConfig,
    user_compute_stream: Optional[int] = None,
) -> OrtSessionPool:
    """Create an ONNX Runtime session pool on a single device."""
    session = build_session(
        onnx_path,
        device_id=config.gpu_id,
        use_cuda=config.use_cuda,
        user_compute_stream=user_compute_stream,
    )
    return OrtSessionPool([session], max_concurrency_per_gpu=config.max_concurrency_per_gpu)


def _warmup(pool: OrtSessionPool, input_name: str, output_names: list[str], input_size: int) -> None:
    """Warm the pool with a dummy inference to trigger lazy initialization."""
    dummy = np.zeros((1, 3, input_size, input_size), dtype=np.float32)
    pool.warmup_sync({input_name: dummy}, output_names)


def _normalize_device_id(device_id: Optional[int]) -> int:
    """Return a valid CUDA device index."""
    if isinstance(device_id, (list, tuple)):
        device_id = device_id[0] if device_id else 0
    try:
        return int(device_id) if device_id is not None else 0
    except (TypeError, ValueError):
        return 0


def _create_explicit_cuda_streams(
    config: AppConfig,
    logger: logging.Logger,
) -> tuple[dict[str, int], dict[str, object]]:
    """Create explicit non-default CUDA streams for model sessions."""
    if not config.use_cuda:
        return {}, {}

    try:
        import onnxruntime as ort
    except Exception:
        logger.warning("ONNX Runtime import failed while probing CUDA; continuing without explicit CUDA streams")
        return {}, {}

    if "CUDAExecutionProvider" not in ort.get_available_providers():
        logger.warning("CUDAExecutionProvider unavailable; continuing without explicit CUDA streams")
        return {}, {}

    try:
        import cupy as cp
    except Exception as exc:
        logger.warning(
            "CuPy unavailable (%s); continuing without explicit CUDA streams",
            exc,
        )
        return {}, {}

    device_id = _normalize_device_id(config.gpu_id)
    try:
        device_count = int(cp.cuda.runtime.getDeviceCount())
    except Exception as exc:
        logger.warning("CUDA runtime probe failed (%s); continuing without explicit CUDA streams", exc)
        return {}, {}

    if device_count <= 0:
        logger.warning("No CUDA devices found; continuing without explicit CUDA streams")
        return {}, {}
    if device_id < 0 or device_id >= device_count:
        logger.warning(
            "Configured gpu_id=%s out of range (device_count=%s); continuing without explicit CUDA streams",
            device_id,
            device_count,
        )
        return {}, {}

    try:
        with cp.cuda.Device(device_id):
            seg_stream = cp.cuda.Stream(non_blocking=True)
            face_stream = cp.cuda.Stream(non_blocking=True)
    except Exception as exc:
        logger.warning("Failed to create CUDA streams (%s); continuing without explicit CUDA streams", exc)
        return {}, {}

    stream_ptrs = {
        "segmentation": int(seg_stream.ptr),
        "facedetection": int(face_stream.ptr),
    }
    stream_handles = {
        "segmentation": seg_stream,
        "facedetection": face_stream,
    }
    return stream_ptrs, stream_handles


def load_models(config: AppConfig) -> LoadedModels:
    """Download, validate, and load model artifacts into runtime pools."""
    logger = logging.getLogger("trifecta")
    cuda_stream_ptrs, cuda_stream_handles = _create_explicit_cuda_streams(config, logger)
    if cuda_stream_ptrs:
        logger.info(
            "Using explicit CUDA streams: segmentation=%s facedetection=%s",
            cuda_stream_ptrs["segmentation"],
            cuda_stream_ptrs["facedetection"],
        )

    # Segmentation model is required.
    seg_snapshot_dir = _download_repo(
        repo_id=config.segmentation_repo_id,
        revision=config.segmentation_revision,
        token=config.segmentation_token,
        local_files_only=config.segmentation_local_files_only,
        allow_patterns=[
            config.segmentation_onnx_filename,
            config.segmentation_config_filename,
            config.textmask_onnx_filename,
        ],
    )
    seg_onnx_path = os.path.join(seg_snapshot_dir, config.segmentation_onnx_filename)
    seg_cfg_path = os.path.join(seg_snapshot_dir, config.segmentation_config_filename)

    if not os.path.exists(seg_onnx_path):
        raise FileNotFoundError(f"Missing segmentation model at {seg_onnx_path}")
    if not os.path.exists(seg_cfg_path):
        raise FileNotFoundError(f"Missing segmentation config at {seg_cfg_path}")

    seg_config = load_segmentation_config(seg_cfg_path, fallback_input_size=512)
    seg_pool = _create_pool(
        seg_onnx_path,
        config,
        user_compute_stream=cuda_stream_ptrs.get("segmentation"),
    )
    _warmup(seg_pool, seg_config.input_name, [seg_config.output_name], seg_config.input_size)

    # Face detection model is optional at startup, but required by /run_facedetection.
    facedetection_pool: Optional[OrtSessionPool] = None
    facedetection_config: Optional[FacedetectionModelConfig] = None
    try:
        fd_snapshot_dir = _download_repo(
            repo_id=config.facedetection_repo_id,
            revision=config.facedetection_revision,
            token=config.facedetection_token,
            local_files_only=config.facedetection_local_files_only,
            allow_patterns=[
                config.facedetection_onnx_filename,
                config.facedetection_config_filename,
            ],
        )
        fd_onnx_path = os.path.join(fd_snapshot_dir, config.facedetection_onnx_filename)
        fd_cfg_path = os.path.join(fd_snapshot_dir, config.facedetection_config_filename)

        if os.path.exists(fd_onnx_path):
            if not os.path.exists(fd_cfg_path):
                raise FileNotFoundError(f"Missing facedetection config at {fd_cfg_path}")

            # Do not default to a fixed size: if the config omits input_size (or sets it null/<=0),
            # we treat the model as dynamic-H/W and the request path will not force a resize.
            facedetection_config = load_facedetection_config(fd_cfg_path, fallback_input_size=None)
            facedetection_pool = _create_pool(
                fd_onnx_path,
                config,
                user_compute_stream=cuda_stream_ptrs.get("facedetection"),
            )
            # Warmup needs a concrete shape; this does not affect runtime request resizing logic.
            warmup_size = int(facedetection_config.input_size) if facedetection_config.input_size is not None else 224
            _warmup(
                facedetection_pool,
                facedetection_config.input_name,
                facedetection_config.output_name,
                warmup_size,
            )
    except Exception as exc:
        logger.warning("Face detection model was not initialized: %s", exc)
        facedetection_pool = None
        facedetection_config = None

    textmask_pool = None
    textmask_path = os.path.join(seg_snapshot_dir, config.textmask_onnx_filename)
    if os.path.exists(textmask_path):
        textmask_pool = _create_pool(
            textmask_path,
            config,
            user_compute_stream=cuda_stream_ptrs.get("segmentation"),
        )

    return LoadedModels(
        segmentation_pool=seg_pool,
        segmentation_config=seg_config,
        facedetection_pool=facedetection_pool,
        facedetection_config=facedetection_config,
        textmask_pool=textmask_pool,
        cuda_stream_ptrs=cuda_stream_ptrs or None,
        cuda_stream_handles=cuda_stream_handles or None,
    )