""" Unified embedding interface for faces and objects. - Faces: InsightFace (ArcFace/Buffalo_L) — or reuse from Immich - Objects: SigLIP (Vision Transformer via transformers) - Caching: Disk-based cache avoids recomputation on reruns """ import importlib import logging import os import time import warnings from contextlib import contextmanager from pathlib import Path import cv2 import numpy as np from PIL import Image from .cache import get_cache logger = logging.getLogger(__name__) @contextmanager def _suppress_output(): """Suppress stdout/stderr at the file-descriptor level, silencing C extension noise.""" devnull_fd = os.open(os.devnull, os.O_WRONLY) saved_out, saved_err = os.dup(1), os.dup(2) try: os.dup2(devnull_fd, 1) os.dup2(devnull_fd, 2) yield finally: try: os.dup2(saved_out, 1) finally: os.dup2(saved_err, 2) os.close(devnull_fd) os.close(saved_out) os.close(saved_err) # Lazy-loaded singletons _insightface_app = None _insightface_loaded = False _siglip_model = None _siglip_processor = None _siglip_loaded = False def _is_force_cpu() -> bool: """Check if CPU mode is forced via environment variable.""" return os.getenv("FORCE_CPU", "").lower() in ("true", "1", "yes") def _preload_cuda_libs() -> None: """Preload CUDA/cuDNN DLLs so onnxruntime-gpu registers CUDAExecutionProvider. Starting with onnxruntime-gpu 1.19+, CUDA/cuDNN libraries are no longer bundled inside the ORT package — they come from the nvidia-* pip packages. preload_dlls() locates them automatically via site-packages discovery. """ try: import onnxruntime if hasattr(onnxruntime, "preload_dlls"): onnxruntime.preload_dlls(cuda=True, cudnn=True) logger.debug("Preloaded CUDA/cuDNN DLLs for onnxruntime-gpu") else: logger.debug("onnxruntime.preload_dlls() not available (ORT < 1.21)") except Exception as e: logger.warning(f"Failed to preload CUDA/cuDNN DLLs: {e}") # ============================================================================= # InsightFace (Faces) # ============================================================================= def get_insightface_app(): """Singleton for InsightFace app with automatic GPU/CPU fallback.""" global _insightface_app, _insightface_loaded if _insightface_loaded: return _insightface_app _insightface_loaded = True ctx_id = -1 insightface_home = os.environ.get("INSIGHTFACE_HOME", os.path.expanduser("~/.insightface")) try: import onnxruntime as ort from insightface.app import FaceAnalysis # Preload CUDA/cuDNN DLLs before any ORT InferenceSession is created. # Silently no-ops on ROCm/Intel builds where preload_dlls() is absent. _preload_cuda_libs() # Disk cache check — lets the user know whether a download is coming buffalo_path = Path(insightface_home) / "models" / "buffalo_l" if buffalo_path.exists() and any(buffalo_path.iterdir()): logger.info("InsightFace Buffalo_L: found in model cache") else: logger.info("InsightFace Buffalo_L: not cached — downloading now (~300 MB)") # Get providers, excluding TensorRT to avoid noisy errors providers = [p for p in ort.get_available_providers() if p != "TensorrtExecutionProvider"] logger.debug(f"ONNX providers available: {providers}") gpu_providers = { "CUDAExecutionProvider", "ROCmExecutionProvider", "MPSExecutionProvider", "CoreMLExecutionProvider", "OpenVINOExecutionProvider", } has_gpu_provider = bool(gpu_providers & set(providers)) ctx_id = -1 if _is_force_cpu() else (0 if has_gpu_provider else -1) # For OpenVINO EP, inject device_type from env var (default CPU; set GPU for Intel Arc/iGPU) has_openvino = "OpenVINOExecutionProvider" in providers if has_openvino: openvino_device = os.getenv("OPENVINO_DEVICE", "CPU") providers = [ ("OpenVINOExecutionProvider", {"device_type": openvino_device}) if p == "OpenVINOExecutionProvider" else p for p in providers ] logger.debug(f"OpenVINO EP: device_type={openvino_device}") if not has_gpu_provider and not _is_force_cpu(): logger.warning( "No GPU execution provider found — running InsightFace on CPU. " "Ensure the container has GPU access and the correct variant image is used " "(gpu for NVIDIA, rocm for AMD, intel for Intel Arc/iGPU)." ) if ctx_id < 0: device_str = "CPU" elif has_openvino: device_str = f"OpenVINO ({os.getenv('OPENVINO_DEVICE', 'CPU')})" else: device_str = "GPU" logger.info(f"InsightFace Buffalo_L: loading into memory on {device_str}...") t0 = time.time() with _suppress_output(): _insightface_app = FaceAnalysis(name="buffalo_l", root=insightface_home, providers=providers) _insightface_app.prepare(ctx_id=ctx_id, det_size=(640, 640)) logger.info(f"InsightFace Buffalo_L: ready on {device_str} ({time.time() - t0:.1f}s)") return _insightface_app except ImportError: logger.error("InsightFace not installed!") return None except Exception as e: logger.error(f"Failed to load InsightFace: {e}") if ctx_id == 0: logger.warning("InsightFace GPU load failed — retrying on CPU...") try: from insightface.app import FaceAnalysis t0 = time.time() with _suppress_output(): _insightface_app = FaceAnalysis( name="buffalo_l", root=insightface_home, providers=["CPUExecutionProvider"], ) _insightface_app.prepare(ctx_id=-1, det_size=(640, 640)) logger.info(f"InsightFace Buffalo_L: ready on CPU (fallback, {time.time() - t0:.1f}s)") return _insightface_app except Exception as ex: logger.error(f"InsightFace CPU fallback failed: {ex}") return None def get_face_embedding(img_pil: Image.Image) -> np.ndarray | None: """Get embedding of the largest face in a PIL image.""" app = get_insightface_app() if not app: return None try: # InsightFace expects BGR cv2 image img_bgr = cv2.cvtColor(np.asarray(img_pil), cv2.COLOR_RGB2BGR) # Suppress scikit-image FutureWarning from InsightFace's face_align.py with warnings.catch_warnings(): warnings.filterwarnings("ignore", message=".*estimate.*is deprecated", category=FutureWarning) faces = app.get(img_bgr) if not faces: return None # Return embedding of largest face largest = max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1])) return largest.embedding except Exception as e: logger.error(f"Error getting face embedding: {e}") return None # ============================================================================= # SigLIP (Objects) # ============================================================================= def get_siglip_model(): """Singleton for SigLIP model and processor with GPU auto-detection.""" global _siglip_model, _siglip_processor, _siglip_loaded if _siglip_loaded: return _siglip_model, _siglip_processor _siglip_loaded = True try: import warnings import torch from transformers import AutoImageProcessor, SiglipVisionModel model_name = "google/siglip-base-patch16-224" # Disk cache check — path derived from model_name using HuggingFace's slug convention hf_home = os.environ.get("HF_HOME", os.path.join(os.path.expanduser("~"), ".cache", "huggingface")) cache_slug = "models--" + model_name.replace("/", "--") model_cache = Path(hf_home) / "hub" / cache_slug if model_cache.exists() and any(model_cache.iterdir()): logger.info(f"SigLIP {model_name}: found in model cache") else: logger.info(f"SigLIP {model_name}: not cached — downloading now (~380 MB)") logger.info(f"SigLIP {model_name}: loading into memory...") t0 = time.time() with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=FutureWarning) warnings.filterwarnings("ignore", message=".*use_fast.*") _siglip_processor = AutoImageProcessor.from_pretrained(model_name, use_fast=True) _siglip_model = SiglipVisionModel.from_pretrained(model_name) _siglip_model.eval() # Move to GPU if available (ROCm builds expose torch.cuda.is_available() == True) if not _is_force_cpu(): if torch.cuda.is_available(): _siglip_model = _siglip_model.cuda() device_name = "CUDA GPU" elif hasattr(torch, "xpu") and torch.xpu.is_available(): _siglip_model = _siglip_model.to("xpu") device_name = "Intel XPU" elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): _siglip_model = _siglip_model.to("mps") device_name = "Apple MPS" else: device_name = "CPU" else: device_name = "CPU (FORCE_CPU)" logger.info(f"SigLIP {model_name}: ready on {device_name} ({time.time() - t0:.1f}s)") return _siglip_model, _siglip_processor except ImportError as e: logger.error(f"transformers/torch not installed: {e}") return None, None except Exception as e: logger.error(f"Failed to load SigLIP: {e}") return None, None def get_object_embedding(img_pil: Image.Image) -> np.ndarray | None: """Get 768-dim SigLIP embedding for an image.""" model, processor = get_siglip_model() if model is None: return None try: import torch inputs = processor(images=img_pil, return_tensors="pt") device = next(model.parameters()).device inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): outputs = model(**inputs) return outputs.pooler_output.squeeze().cpu().numpy() except Exception as e: logger.error(f"Error getting object embedding: {e}") return None def get_object_embeddings_batch(images: list[Image.Image]) -> list[np.ndarray | None]: """Get SigLIP embeddings for a batch of images (GPU-efficient).""" model, processor = get_siglip_model() if model is None: return [None] * len(images) try: import torch inputs = processor(images=images, return_tensors="pt", padding=True) device = next(model.parameters()).device inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): outputs = model(**inputs) embeddings = outputs.pooler_output.cpu().numpy() return [embeddings[i] for i in range(len(embeddings))] except Exception as e: logger.error(f"Error in batch embedding: {e}") # Fall back to individual computation return [get_object_embedding(img) for img in images] # ============================================================================= # Unified Interface with Caching # ============================================================================= def get_embedding( img_pil: Image.Image, entity_type: str = "face", asset_id: str | None = None, immich_embedding: np.ndarray | None = None, ) -> np.ndarray | None: """Get embedding for an image based on entity type. Priority: 1. Pre-fetched Immich embedding (if provided) 2. Disk cache (if enabled and asset_id provided) 3. Local model computation (InsightFace or SigLIP) Args: img_pil: The image to embed entity_type: 'face' or 'object' asset_id: Optional asset ID for cache lookup immich_embedding: Optional pre-fetched embedding from Immich API """ from .config import Config use_cache = Config.ENABLE_CACHE and asset_id is not None cache = get_cache(Config.CACHE_DIR) if use_cache else None # Use a single consistent cache key per model so lookups and stores always match. # "immich" was previously used as the face key on the lookup path but "insightface" # on the store path — meaning the cache was never hit for locally-computed embeddings. cache_key = "insightface" if entity_type == "face" else "siglip" # 1. Use Immich embedding if provided if immich_embedding is not None: if cache: cache.put(asset_id, immich_embedding, cache_key) return immich_embedding # 2. Check disk cache if cache: cached = cache.get(asset_id, cache_key) if cached is not None: return cached # 3. Compute locally if entity_type == "face": emb = get_face_embedding(img_pil) else: emb = get_object_embedding(img_pil) if emb is not None and cache: cache.put(asset_id, emb, cache_key) return emb def _is_module_available(module_name: str) -> bool: """Check if a Python module is importable without importing it fully.""" try: importlib.util.find_spec(module_name) return True except (ModuleNotFoundError, ValueError): return False def is_embedding_available(entity_type: str = "face", *, load: bool = False) -> bool: """Check if embedding model is available for the given entity type. By default this performs a lightweight import-check only (no model loading). Pass ``load=True`` to actually load the model (expensive, hundreds of MB). Args: entity_type: 'face' or 'object' load: If True, fully load the model to verify. If False (default), only check that the required packages are importable. """ if load: if entity_type == "face": return get_insightface_app() is not None model, _ = get_siglip_model() return model is not None # Lightweight check: just verify the packages are importable if entity_type == "face": return _is_module_available("insightface") and _is_module_available("onnxruntime") return _is_module_available("transformers") and _is_module_available("torch") def load_embedding_model(entity_type: str = "face") -> bool: """Explicitly load the embedding model for the given entity type. Returns True if the model loaded successfully. """ if entity_type == "face": return get_insightface_app() is not None model, _ = get_siglip_model() return model is not None