#

models/sam2_loader.py

@@ -1,279 +1,107 @@
-#!/usr/bin/env python3
-"""
-SAM2 Loader — Robust loading and mask generation for SAM2
-========================================================
-- Loads SAM2 model with Hydra config resolution
-- Generates seed masks for MatAnyone
-- Aligned with torch==2.3.1+cu121 and SAM2 commit 3c76f73c1a7e7b4a2e8a0a9a3e5b92f7e6e3f2f5
-
-Changes (2025-09-16):
-- Aligned with torch==2.3.1+cu121 and SAM2 commit
-- Added GPU memory logging for Tesla T4
-- Added SAM2 version logging via importlib.metadata
-- Simplified config resolution to match __init__.py
-"""
-
-from __future__ import annotations
-
-import os
-import logging
-import importlib.metadata
-from pathlib import Path
-from typing import Optional, Tuple, Dict, Any
-
-import numpy as np
-import yaml
-import torch
-
-# --------------------------------------------------------------------------------------
-# Logging
-# --------------------------------------------------------------------------------------
-logger = logging.getLogger("backgroundfx_pro")
-if not logger.handlers:
-    _h = logging.StreamHandler()
-    _h.setFormatter(logging.Formatter("[%(asctime)s] %(levelname)s: %(message)s"))
-    logger.addHandler(_h)
-logger.setLevel(logging.INFO)
-
-# --------------------------------------------------------------------------------------
-# Path setup for third_party repos
-# --------------------------------------------------------------------------------------
-ROOT = Path(__file__).resolve().parent.parent  # project root
-TP_SAM2 = Path(os.environ.get("THIRD_PARTY_SAM2_DIR", ROOT / "third_party" / "sam2")).resolve()
-
-def _add_sys_path(p: Path) -> None:
-    if p.exists():
-        p_str = str(p)
-        if p_str not in sys.path:
-            sys.path.insert(0, p_str)
-    else:
-        logger.warning(f"third_party path not found: {p}")
-
-_add_sys_path(TP_SAM2)
-
-# --------------------------------------------------------------------------------------
-# Safe Torch accessors
-# --------------------------------------------------------------------------------------
-def _torch():
-    try:
-        import torch
-        return torch
-    except Exception as e:
-        logger.warning(f"[sam2_loader.safe-torch] import failed: {e}")
-        return None
-
-def _has_cuda() -> bool:
-    t = _torch()
-    if t is None:
-        return False
-    try:
-        return bool(t.cuda.is_available())
-    except Exception as e:
-        logger.warning(f"[sam2_loader.safe-torch] cuda.is_available() failed: {e}")
-        return False
-
-def _pick_device(env_key: str) -> str:
-    requested = os.environ.get(env_key, "").strip().lower()
-    has_cuda = _has_cuda()
-    
-    logger.info(f"CUDA environment variables: {{'SAM2_DEVICE': '{os.environ.get('SAM2_DEVICE', '')}'}}")
-    logger.info(f"_pick_device({env_key}): requested='{requested}', has_cuda={has_cuda}")
-    
-    if has_cuda and requested not in {"cpu"}:
-        logger.info(f"FORCING CUDA device (GPU available, requested='{requested}')")
-        return "cuda"
-    elif requested in {"cuda", "cpu"}:
-        logger.info(f"Using explicitly requested device: {requested}")
-        return requested
-    
-    result = "cuda" if has_cuda else "cpu"
-    logger.info(f"Auto-selected device: {result}")
-    return result
-
-# --------------------------------------------------------------------------------------
-# SAM2 Loading and Mask Generation
-# --------------------------------------------------------------------------------------
-def _resolve_sam2_cfg(cfg_str: str) -> str:
-    """Resolve SAM2 config path - return relative path for Hydra compatibility."""
-    logger.info(f"_resolve_sam2_cfg called with cfg_str={cfg_str}")
-    
-    candidate = os.path.join(TP_SAM2, cfg_str)
-    logger.info(f"Candidate path: {candidate}")
-    logger.info(f"Candidate exists: {os.path.exists(candidate)}")
-    
-    if os.path.exists(candidate):
-        if cfg_str.startswith("sam2/configs/"):
-            relative_path = cfg_str.replace("sam2/configs/", "configs/")
-        else:
-            relative_path = cfg_str
-        logger.info(f"Returning Hydra-compatible relative path: {relative_path}")
-        return relative_path
-    
-    fallbacks = [
-        os.path.join(TP_SAM2, "sam2", cfg_str),
-        os.path.join(TP_SAM2, "configs", cfg_str),
-    ]
-    
-    for fallback in fallbacks:
-        logger.info(f"Trying fallback: {fallback}")
-        if os.path.exists(fallback):
-            if "configs/" in fallback:
-                relative_path = "configs/" + fallback.split("configs/")[-1]
-                logger.info(f"Returning fallback relative path: {relative_path}")
-                return relative_path
-    
-    logger.warning(f"Config not found, returning original: {cfg_str}")
-    return cfg_str
-
-def _find_hiera_config_if_hieradet(cfg_path: str) -> Optional[str]:
-    """If config references 'hieradet', try to find a 'hiera' config."""
-    try:
-        with open(cfg_path, "r") as f:
-            data = yaml.safe_load(f)
-        model = data.get("model", {}) or {}
-        enc = model.get("image_encoder") or {}
-        trunk = enc.get("trunk") or {}
-        target = trunk.get("_target_") or trunk.get("target")
-        if isinstance(target, str) and "hieradet" in target:
-            for y in TP_SAM2.rglob("*.yaml"):
-                try:
-                    with open(y, "r") as f2:
-                        d2 = yaml.safe_load(f2) or {}
-                    e2 = (d2.get("model", {}) or {}).get("image_encoder") or {}
-                    t2 = (e2.get("trunk") or {})
-                    tgt2 = t2.get("_target_") or t2.get("target")
-                    if isinstance(tgt2, str) and ".hiera." in tgt2:
-                        logger.info(f"SAM2: switching config from 'hieradet' → 'hiera': {y}")
-                        return str(y)
-                except Exception:
-                    continue
-    except Exception:
-        pass
-    return None
-
-def load_sam2() -> Tuple[Optional[object], bool, Dict[str, Any]]:
-    """Robust SAM2 loader with config resolution and error handling."""
-    meta = {"sam2_import_ok": False, "sam2_init_ok": False}
-    try:
-        from sam2.build_sam import build_sam2
-        from sam2.sam2_image_predictor import SAM2ImagePredictor
-        meta["sam2_import_ok"] = True
-    except Exception as e:
-        logger.warning(f"SAM2 import failed: {e}")
-        return None, False, meta
-
-    # Log SAM2 version
-    try:
-        version = importlib.metadata.version("segment-anything-2")
-        logger.info(f"[SAM2] SAM2 version: {version}")
-    except Exception:
-        logger.info("[SAM2] SAM2 version unknown")
-
-    # Check GPU memory before loading
-    if torch and torch.cuda.is_available():
-        mem_before = torch.cuda.memory_allocated() / 1024**3
-        logger.info(f"🔍 GPU memory before SAM2 load: {mem_before:.2f}GB")
-
-    device = _pick_device("SAM2_DEVICE")
-    cfg_env = os.environ.get("SAM2_MODEL_CFG", "sam2/configs/sam2/sam2_hiera_l.yaml")
-    cfg = _resolve_sam2_cfg(cfg_env)
-    ckpt = os.environ.get("SAM2_CHECKPOINT", "")
-
-    def _try_build(cfg_path: str):
-        logger.info(f"_try_build called with cfg_path: {cfg_path}")
-        params = set(inspect.signature(build_sam2).parameters.keys())
-        logger.info(f"build_sam2 parameters: {list(params)}")
-        kwargs = {}
-        if "config_file" in params:
-            kwargs["config_file"] = cfg_path
-            logger.info(f"Using config_file parameter: {cfg_path}")
-        elif "model_cfg" in params:
-            kwargs["model_cfg"] = cfg_path
-            logger.info(f"Using model_cfg parameter: {cfg_path}")
-        if ckpt:
-            if "checkpoint" in params:
-                kwargs["checkpoint"] = ckpt
-            elif "ckpt_path" in params:
-                kwargs["ckpt_path"] = ckpt
-            elif "weights" in params:
-                kwargs["weights"] = ckpt
-        if "device" in params:
-            kwargs["device"] = device
-        try:
-            logger.info(f"Calling build_sam2 with kwargs: {kwargs}")
-            result = build_sam2(**kwargs)
-            logger.info(f"build_sam2 succeeded with kwargs")
-            if hasattr(result, 'device'):
-                logger.info(f"SAM2 model device: {result.device}")
-            elif hasattr(result, 'image_encoder') and hasattr(result.image_encoder, 'device'):
-                logger.info(f"SAM2 model device: {result.image_encoder.device}")
-            return result
-        except TypeError as e:
-            logger.info(f"build_sam2 kwargs failed: {e}, trying positional args")
-            pos = [cfg_path]
-            if ckpt:
-                pos.append(ckpt)
-            if "device" not in kwargs:
-                pos.append(device)
-            logger.info(f"Calling build_sam2 with positional args: {pos}")
-            result = build_sam2(*pos)
-            logger.info(f"build_sam2 succeeded with positional args")
-            return result
-
-    try:
-        try:
-            sam = _try_build(cfg)
-        except Exception:
-            alt_cfg = _find_hiera_config_if_hieradet(cfg)
-            if alt_cfg:
-                sam = _try_build(alt_cfg)
-            else:
-                raise
-        
-        if sam is not None:
-            predictor = SAM2ImagePredictor(sam)
-            meta["sam2_init_ok"] = True
-            meta["sam2_device"] = device
-            return predictor, True, meta
-        else:
-            return None, False, meta
-            
-    except Exception as e:
-        logger.error(f"SAM2 loading failed: {e}")
-        return None, False, meta
-
-def run_sam2_mask(predictor: object,
-                  first_frame_bgr: np.ndarray,
-                  point: Optional[Tuple[int, int]] = None,
-                  auto: bool = False) -> Tuple[Optional[np.ndarray], bool]:
-    """Generate a seed mask for MatAnyone. Returns (mask_uint8_0_255, ok)."""
-    if predictor is None:
-        return None, False
-    try:
-        import cv2
-        rgb = cv2.cvtColor(first_frame_bgr, cv2.COLOR_BGR2RGB)
-        predictor.set_image(rgb)
-
-        if auto:
-            h, w = rgb.shape[:2]
-            box = np.array([int(0.05*w), int(0.05*h), int(0.95*w), int(0.95*h)])
-            masks, _, _ = predictor.predict(box=box)
-        elif point is not None:
-            x, y = int(point[0]), int(point[1])
-            pts = np.array([[x, y]], dtype=np.int32)
-            labels = np.array([1], dtype=np.int32)
-            masks, _, _ = predictor.predict(point_coords=pts, point_labels=labels)
-        else:
-            h, w = rgb.shape[:2]
-            box = np.array([int(0.1*w), int(0.1*h), int(0.9*w), int(0.9*h)])
-            masks, _, _ = predictor.predict(box=box)
-
-        if masks is None or len(masks) == 0:
-            return None, False
-
-        m = masks[0].astype(np.uint8) * 255
-        logger.info(f"[SAM2] Generated mask: shape={m.shape}, dtype={m.dtype}")
-        return m, True
-    except Exception as e:
-        logger.warning(f"SAM2 mask generation failed: {e}")
-        return None, False
+# ===============================
+# Optimized Dockerfile for Hugging Face Spaces
+# PyTorch 2.3.1 + CUDA 12.1
+# ===============================
+
+# Base image with CUDA 12.1.1
+FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
+
+# Environment variables
+ENV DEBIAN_FRONTEND=noninteractive \
+    PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PIP_NO_CACHE_DIR=1 \
+    PIP_DISABLE_PIP_VERSION_CHECK=1 \
+    TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6+PTX" \
+    FORCE_CUDA="1" \
+    CUDA_VISIBLE_DEVICES="0"
+
+# Create non-root user
+RUN useradd -m -u 1000 user
+ENV HOME=/home/user
+WORKDIR $HOME/app
+
+# Install system dependencies in a single layer
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git \
+    ffmpeg \
+    wget \
+    python3 \
+    python3-pip \
+    python3-venv \
+    python3-dev \
+    build-essential \
+    gcc \
+    g++ \
+    pkg-config \
+    libffi-dev \
+    libssl-dev \
+    libc6-dev \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    libsm6 \
+    libxext6 \
+    libxrender1 \
+    libgomp1 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set up Python environment
+RUN python3 -m pip install --no-cache-dir --upgrade pip setuptools wheel
+
+# Install PyTorch with CUDA 12.1 first (base for other dependencies)
+RUN python3 -m pip install --no-cache-dir \
+    --extra-index-url https://download.pytorch.org/whl/cu121 \
+    torch==2.3.1+cu121 \
+    torchvision==0.18.1+cu121 \
+    torchaudio==2.3.1+cu121 \
+    && python3 -c "import torch; print(f'PyTorch version: {torch.__version__}'); print(f'CUDA available: {torch.cuda.is_available()}'); print(f'CUDA version: {torch.version.cuda if torch.cuda.is_available() else \"N/A\"}'); print(f'cuDNN version: {torch.backends.cudnn.version() if torch.cuda.is_available() else \"N/A\"}')"
+
+# Copy requirements files first for better caching
+COPY --chown=user requirements.txt requirements-hf.txt ./
+
+# Install Python dependencies
+RUN python3 -m pip install --no-cache-dir -r requirements.txt
+
+# Install MatAnyone with retry logic and fallback dependencies
+RUN echo "Installing problematic dependencies first..." && \
+    python3 -m pip install --no-cache-dir chardet charset-normalizer && \
+    echo "Installing MatAnyone..." && \
+    (python3 -m pip install --no-cache-dir -v git+https://github.com/pq-yang/MatAnyone@main#egg=matanyone || \
+     (echo "Retrying MatAnyone installation..." && \
+      python3 -m pip install --no-cache-dir -v git+https://github.com/pq-yang/MatAnyone@main#egg=matanyone)) && \
+    python3 -c "import matanyone; print('MatAnyone import successful')"
+
+# Copy application code
+COPY --chown=user . .
+
+# Install SAM2
+RUN echo "Installing SAM2..." && \
+    git clone --depth=1 https://github.com/facebookresearch/segment-anything-2.git third_party/sam2 && \
+    cd third_party/sam2 && \
+    python3 -m pip install --no-cache-dir -e .
+
+# Set up environment variables
+ENV PYTHONPATH=/home/user/app:/home/user/app/third_party:/home/user/app/third_party/sam2 \
+    FFMPEG_BIN=ffmpeg \
+    THIRD_PARTY_SAM2_DIR=/home/user/app/third_party/sam2 \
+    ENABLE_MATANY=1 \
+    SAM2_DEVICE=cuda \
+    MATANY_DEVICE=cuda \
+    OMP_NUM_THREADS=2 \
+    TF_CPP_MIN_LOG_LEVEL=2 \
+    SAM2_CHECKPOINT=/home/user/app/checkpoints/sam2_hiera_large.pt
+
+# Create checkpoints directory
+RUN mkdir -p /home/user/app/checkpoints
+
+# Note: SAM2 model will be downloaded at runtime via lazy loading
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=5s --retries=3 CMD python3 -c "import torch; print(f'PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}')"
+
+# Run as non-root user
+USER user
+EXPOSE 7860
+
+# Start the application
+CMD ["python3", "-u", "app_hf.py"]