Phase 1: Add SAM2/MatAnyone optimization infrastructure

Dockerfile

@@ -1,11 +1,57 @@
-# syntax=docker/dockerfile:1
-FROM python:3.11-slim
+# Production Dockerfile for BackgroundFX Pro with SAM2 + MatAnyone
+FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
 
-WORKDIR /code
+# System dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git ffmpeg libglib2.0-0 libgl1 libglib2.0-0 libsm6 libxrender1 libxext6 \
+    python3.10 python3.10-venv python3-pip \
+    && rm -rf /var/lib/apt/lists/*
 
-COPY requirements.txt .
+# Upgrade pip
+RUN python3 -m pip install --upgrade pip
+
+# Environment variables for caching and performance
+ENV HF_HOME=/home/user/.cache/huggingface \
+    TORCH_HOME=/home/user/.cache/torch \
+    TRANSFORMERS_CACHE=/home/user/.cache/transformers \
+    MPLCONFIGDIR=/home/user/.config/matplotlib
+
+# CUDA and memory optimizations for T4
+ENV PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128,expandable_segments:True \
+    CUDA_LAUNCH_BLOCKING=0 \
+    OMP_NUM_THREADS=2 \
+    MKL_NUM_THREADS=2 \
+    HF_HUB_ENABLE_HF_TRANSFER=1 \
+    TOKENIZERS_PARALLELISM=false
+
+# Create working directory
+WORKDIR /home/user/app
+
+# Copy and install Python dependencies
+COPY requirements.txt ./requirements.txt
 RUN pip install --no-cache-dir -r requirements.txt
 
+# Vendor SAM2 and MatAnyone at build time (more reliable than runtime git)
+# SAM2
+RUN git clone --depth=1 https://github.com/facebookresearch/segment-anything-2 /home/user/app/third_party/sam2
+ENV PYTHONPATH=/home/user/app/third_party/sam2:${PYTHONPATH}
+
+# MatAnyone (official repo)
+RUN git clone --depth=1 https://github.com/pq-yang/MatAnyone /home/user/app/third_party/matanyone
+ENV PYTHONPATH=/home/user/app/third_party/matanyone:${PYTHONPATH}
+
+# Copy application code
 COPY . .
 
-CMD ["python", "app.py"]
+# Create cache directories
+RUN mkdir -p /home/user/.cache/huggingface /home/user/.cache/torch /home/user/.cache/transformers
+
+# Expose Gradio port
+EXPOSE 7860
+
+# Environment for Gradio
+ENV GRADIO_SERVER_NAME=0.0.0.0 \
+    GRADIO_SERVER_PORT=7860
+
+# Run the application
+CMD ["python3", "app.py"]

models/matanyone_loader.py

@@ -0,0 +1,72 @@
+# models/matanyone_loader.py
+import os, logging, torch, gc
+import numpy as np
+from typing import Optional, Tuple
+
+log = logging.getLogger("matany_loader")
+
+def _import_inference_core():
+    try:
+        # Check the actual import path from pq-yang/MatAnyone repo
+        from matanyone.inference_core import InferenceCore
+        return InferenceCore
+    except Exception as e:
+        log.error("MatAnyone import failed (vendoring/repo path?): %s", e)
+        return None
+
+def _to_chw01(img):
+    # img: HWC uint8 or float01 -> CHW float01
+    if img.dtype != np.float32:
+        img = img.astype("float32")/255.0
+    return np.transpose(img, (2,0,1))
+
+def _to_1hw01(mask):
+    # mask: HxW [0,1]
+    m = mask.astype("float32")
+    return m[None, ...]
+
+class MatAnyoneSession:
+    def __init__(self, device: torch.device, precision: str = "fp16"):
+        self.device = device
+        self.precision = precision
+        self.core = None
+
+    def load(self, ckpt_path: Optional[str] = None, repo_id: Optional[str] = None, filename: Optional[str] = None):
+        InferenceCore = _import_inference_core()
+        if InferenceCore is None:
+            raise RuntimeError("MatAnyone not importable")
+
+        if ckpt_path is None and repo_id and filename:
+            from huggingface_hub import hf_hub_download
+            ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename, local_dir=os.environ.get("HF_HOME"))
+
+        # init model
+        self.core = InferenceCore(ckpt_path, device=str(self.device))
+        return self
+
+    @torch.inference_mode()
+    def step(self, image_rgb, seed_mask: Optional[np.ndarray]=None):
+        """
+        image_rgb: HxWx3 uint8/float01
+        seed_mask: HxW float01 for first frame, else None
+        returns alpha HxW float01
+        """
+        assert self.core is not None, "MatAnyone not loaded"
+        img = _to_chw01(image_rgb)    # CHW
+        if seed_mask is not None:
+            mask = _to_1hw01(seed_mask)  # 1HW
+            alpha = self.core.step(img, mask)
+        else:
+            alpha = self.core.step(img, None)
+        # ensure HxW
+        if isinstance(alpha, np.ndarray):
+            return alpha.astype("float32")
+        if torch.is_tensor(alpha):
+            return alpha.detach().float().cpu().numpy()
+        raise RuntimeError("MatAnyone returned unknown alpha type")
+
+    def reset(self):
+        if self.core and hasattr(self.core, "reset"):
+            self.core.reset()
+        torch.cuda.empty_cache()
+        gc.collect()

models/sam2_loader.py

@@ -0,0 +1,75 @@
+# models/sam2_loader.py
+import os, logging, torch
+from huggingface_hub import hf_hub_download
+from pathlib import Path
+import numpy as np
+
+log = logging.getLogger("sam2_loader")
+
+DEFAULT_MODEL_ID = os.environ.get("SAM2_MODEL_ID", "facebook/sam2")
+DEFAULT_VARIANT = os.environ.get("SAM2_VARIANT", "sam2_hiera_large")
+
+# Map variant -> filenames (SAM2 releases follow this pattern)
+VARIANT_FILES = {
+    "sam2_hiera_small": ("sam2_hiera_small.pt", "configs/sam2/sam2_hiera_s.yaml"),
+    "sam2_hiera_base": ("sam2_hiera_base.pt", "configs/sam2/sam2_hiera_b.yaml"),
+    "sam2_hiera_large": ("sam2_hiera_large.pt", "configs/sam2/sam2_hiera_l.yaml"),
+}
+
+def _download_checkpoint(model_id: str, ckpt_name: str) -> str:
+    return hf_hub_download(repo_id=model_id, filename=ckpt_name, local_dir=os.environ.get("HF_HOME"))
+
+def _find_sam2_build():
+    try:
+        from sam2.build_sam import build_sam2
+        return build_sam2
+    except Exception as e:
+        log.error("SAM2 not importable (check Dockerfile vendoring): %s", e)
+        return None
+
+class SAM2Predictor:
+    def __init__(self, device: torch.device):
+        self.device = device
+        self.model = None
+        self.predictor = None
+
+    def load(self, variant: str = DEFAULT_VARIANT, model_id: str = DEFAULT_MODEL_ID):
+        build_sam2 = _find_sam2_build()
+        if build_sam2 is None:
+            raise RuntimeError("SAM2 build function not available")
+
+        ckpt_name, cfg_path = VARIANT_FILES.get(variant, VARIANT_FILES["sam2_hiera_large"])
+        ckpt = _download_checkpoint(model_id, ckpt_name)
+
+        # Compose config via hydra-free path (using explicit path args)
+        model = build_sam2(config_file=cfg_path, ckpt_path=ckpt, device=str(self.device))
+        model.eval()
+        self.model = model
+
+        try:
+            from sam2.sam2_video_predictor import SAM2VideoPredictor
+            self.predictor = SAM2VideoPredictor(self.model)
+        except Exception:
+            # Fallback to image predictor if video predictor missing
+            from sam2.sam2_image_predictor import SAM2ImagePredictor
+            self.predictor = SAM2ImagePredictor(self.model)
+
+        return self
+
+    @torch.inference_mode()
+    def first_frame_mask(self, image_rgb01):
+        """
+        Returns an initial binary-ish mask for the foreground subject from first frame.
+        You can refine prompts here (points/boxes) if you add UI hooks later.
+        """
+        if hasattr(self.predictor, "set_image"):
+            self.predictor.set_image((image_rgb01*255).astype("uint8"))
+            # simple auto-box prompt (tight box)
+            h, w = image_rgb01.shape[:2]
+            box = np.array([1, 1, w-2, h-2])
+            masks, _, _ = self.predictor.predict(box=box, multimask_output=False)
+            mask = masks[0]  # HxW bool/float
+        else:
+            # video predictor path: run_single_frame if available
+            mask = (image_rgb01[...,0] > -1)  # dummy, should not happen
+        return mask.astype("float32")

requirements.txt

@@ -1,8 +1,43 @@
+# ===== Core runtime =====
+# Option A: Keep your current Torch stack (safe for existing builds)
 torch==2.2.2
 torchvision==0.17.2
+torchaudio==2.2.2
+
+# Option B: Faster CUDA 12.1 wheels for T4 (uncomment to use instead)
+# torch==2.3.1+cu121
+# torchvision==0.18.1+cu121  
+# torchaudio==2.3.1+cu121
+# --extra-index-url https://download.pytorch.org/whl/cu121
+
+# ===== Video / image IO =====
 opencv-python-headless==4.10.0.84
+imageio==2.35.1
+imageio-ffmpeg==0.5.1
+moviepy==1.0.3
+decord==0.6.0
+Pillow==10.4.0
 numpy==1.26.4
-pillow==10.4.0
+
+# ===== Gradio UI =====
 gradio==5.42.0
+
+# ===== SAM2 Dependencies =====
+hydra-core==1.3.2
+omegaconf==2.3.0
+einops==0.8.0
+timm==1.0.9
+pyyaml==6.0.2
+matplotlib==3.9.2
+
+# ===== MatAnyone Dependencies =====
+kornia==0.7.3
+scikit-image==0.24.0
+tqdm==4.66.5
+
+# ===== Helpers / caching =====
+huggingface_hub==0.24.6
+ffmpeg-python==0.2.0
+psutil==6.0.0
 requests==2.31.0
-huggingface-hub>=0.33.5
+scikit-learn==1.5.1

utils/accelerator.py

@@ -0,0 +1,34 @@
+# utils/accelerator.py
+import os, torch, logging, psutil, gc
+
+log = logging.getLogger("accelerator")
+
+def pick_device():
+    if torch.cuda.is_available():
+        dev = torch.device("cuda")
+        name = torch.cuda.get_device_name(0)
+        log.info(f"Using GPU: {name}")
+        return dev
+    log.warning("CUDA not available; falling back to CPU.")
+    return torch.device("cpu")
+
+def torch_global_tuning():
+    # better matmul perf without crazy memory
+    try:
+        torch.set_float32_matmul_precision("high")
+    except Exception:
+        pass
+
+def memory_checkpoint(tag=""):
+    try:
+        if torch.cuda.is_available():
+            mem = torch.cuda.memory_allocated() / (1024**2)
+            log.info(f"[CUDA mem] {tag}: {mem:.1f} MB")
+    except Exception:
+        pass
+
+def cleanup():
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
+    gc.collect()