winnow/winnow/embeddings.py

"""
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:
        os.dup2(saved_out, 1)
        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

    # Preload CUDA/cuDNN DLLs BEFORE any ORT InferenceSession is created
    _preload_cuda_libs()

    ctx_id = -1
    insightface_home = os.environ.get("INSIGHTFACE_HOME", os.path.expanduser("~/.insightface"))
    try:
        import onnxruntime as ort
        from insightface.app import FaceAnalysis

        # 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",
        }
        has_gpu_provider = bool(gpu_providers & set(providers))
        ctx_id = -1 if _is_force_cpu() else (0 if has_gpu_provider else -1)

        if not has_gpu_provider and not _is_force_cpu():
            logger.warning(
                "No GPU execution provider found — running InsightFace on CPU. "
                "If you have an NVIDIA GPU, ensure the NVIDIA Container Toolkit is "
                "installed and the container has GPU access (deploy.resources in compose)."
            )

        device_str = "GPU" if ctx_id >= 0 else "CPU"
        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
        if not _is_force_cpu():
            if torch.cuda.is_available():
                _siglip_model = _siglip_model.cuda()
                device_name = "CUDA GPU"
            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