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, )