fejl 40

Dockerfile

@@ -89,16 +89,12 @@ ENV PYTHONPATH=/home/user/app:/home/user/app/third_party:/home/user/app/third_pa
     MATANY_DEVICE=cuda \
     OMP_NUM_THREADS=2 \
     TF_CPP_MIN_LOG_LEVEL=2 \
-    SAM2_MODEL_PATH=/home/user/app/checkpoints/sam2_hiera_large.pt
+    SAM2_CHECKPOINT=/home/user/app/checkpoints/sam2_hiera_large.pt
 
 # Create checkpoints directory
 RUN mkdir -p /home/user/app/checkpoints
 
-# Download SAM2 model
-RUN echo "Downloading SAM2 model..." && \
-    wget -q -O /home/user/app/checkpoints/sam2_hiera_large.pt \
-    https://dl.fbaipublicfiles.com/segment_anything/sam2_hiera_large.pt && \
-    chown -R user:user /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()}')"

models/sam2_loader.py

@@ -1,107 +1,279 @@
-# ===============================
-# 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"]
+#!/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