From 1bec8113deb40b902563636584aca156c8b22f80 Mon Sep 17 00:00:00 2001
From: Scott Register <sreg@sreg.io>
Date: Sat, 26 Jul 2025 07:57:01 -0700
Subject: [PATCH] fix sam2 models

---
 config_example.yaml              |  2 ++
 config_runpod.yaml               |  2 ++
 vr180_matting/config.py          |  6 +++++-
 vr180_matting/sam2_wrapper.py    | 14 ++++++++++++++
 vr180_matting/video_processor.py |  2 ++
 5 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/config_example.yaml b/config_example.yaml
index 8c51a84..18dea93 100644
--- a/config_example.yaml
+++ b/config_example.yaml
@@ -14,6 +14,8 @@ matting:
   use_disparity_mapping: true
   memory_offload: true
   fp16: true
+  sam2_model_cfg: "sam2_hiera_l.yaml"
+  sam2_checkpoint: "models/sam2_hiera_large.pt"
   
 output:
   path: "path/to/output/"
diff --git a/config_runpod.yaml b/config_runpod.yaml
index ae13aa5..c803a89 100644
--- a/config_runpod.yaml
+++ b/config_runpod.yaml
@@ -14,6 +14,8 @@ matting:
   use_disparity_mapping: true
   memory_offload: false  # A40 has enough VRAM
   fp16: true
+  sam2_model_cfg: "sam2_hiera_l.yaml"
+  sam2_checkpoint: "models/sam2_hiera_large.pt"
   
 output:
   path: "/workspace/output/matted_video.mp4"
diff --git a/vr180_matting/config.py b/vr180_matting/config.py
index 5d621b8..e6eed49 100644
--- a/vr180_matting/config.py
+++ b/vr180_matting/config.py
@@ -27,6 +27,8 @@ class MattingConfig:
     use_disparity_mapping: bool = True
     memory_offload: bool = True
     fp16: bool = True
+    sam2_model_cfg: str = "sam2_hiera_l.yaml"
+    sam2_checkpoint: str = "sam2_hiera_large.pt"
 
 
 @dataclass
@@ -89,7 +91,9 @@ class VR180Config:
             'matting': {
                 'use_disparity_mapping': self.matting.use_disparity_mapping,
                 'memory_offload': self.matting.memory_offload,
-                'fp16': self.matting.fp16
+                'fp16': self.matting.fp16,
+                'sam2_model_cfg': self.matting.sam2_model_cfg,
+                'sam2_checkpoint': self.matting.sam2_checkpoint
             },
             'output': {
                 'path': self.output.path,
diff --git a/vr180_matting/sam2_wrapper.py b/vr180_matting/sam2_wrapper.py
index 4197c21..2d48cda 100644
--- a/vr180_matting/sam2_wrapper.py
+++ b/vr180_matting/sam2_wrapper.py
@@ -4,6 +4,7 @@ from typing import List, Dict, Any, Optional, Tuple
 import cv2
 from pathlib import Path
 import warnings
+import os
 
 try:
     from sam2.build_sam import build_sam2_video_predictor
@@ -38,6 +39,19 @@ class SAM2VideoMatting:
     def _load_model(self, model_cfg: str, checkpoint_path: str):
         """Load SAM2 video predictor with optimizations"""
         try:
+            # Check for checkpoint in models directory if not found
+            if not Path(checkpoint_path).exists():
+                models_path = Path("models") / checkpoint_path
+                if models_path.exists():
+                    checkpoint_path = str(models_path)
+                else:
+                    # Try relative to package
+                    import os
+                    package_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+                    models_path = Path(package_dir) / "models" / checkpoint_path
+                    if models_path.exists():
+                        checkpoint_path = str(models_path)
+            
             self.predictor = build_sam2_video_predictor(
                 model_cfg, 
                 checkpoint_path,
diff --git a/vr180_matting/video_processor.py b/vr180_matting/video_processor.py
index 4397fdd..2c241f4 100644
--- a/vr180_matting/video_processor.py
+++ b/vr180_matting/video_processor.py
@@ -51,6 +51,8 @@ class VideoProcessor:
             
             # Initialize SAM2 model
             self.sam2_model = SAM2VideoMatting(
+                model_cfg=self.config.matting.sam2_model_cfg,
+                checkpoint_path=self.config.matting.sam2_checkpoint,
                 device=self.config.hardware.device,
                 memory_offload=self.config.matting.memory_offload,
                 fp16=self.config.matting.fp16