diff --git a/vr180_matting/sam2_wrapper.py b/vr180_matting/sam2_wrapper.py index ac8774b..abc5253 100644 --- a/vr180_matting/sam2_wrapper.py +++ b/vr180_matting/sam2_wrapper.py @@ -7,6 +7,7 @@ import warnings import os import tempfile import shutil +import gc try: from sam2.build_sam import build_sam2_video_predictor @@ -257,11 +258,11 @@ class SAM2VideoMatting: if self.inference_state is not None: try: # Reset SAM2 state first (critical for memory cleanup) - if hasattr(self.predictor, 'reset_state'): + if self.predictor is not None and hasattr(self.predictor, 'reset_state'): self.predictor.reset_state(self.inference_state) # Fallback to cleanup_state if available - elif hasattr(self.predictor, 'cleanup_state'): + elif self.predictor is not None and hasattr(self.predictor, 'cleanup_state'): self.predictor.cleanup_state(self.inference_state) # Explicitly delete inference state and video segments @@ -275,15 +276,6 @@ class SAM2VideoMatting: finally: self.inference_state = None - # Explicitly delete predictor - if self.predictor is not None: - try: - del self.predictor - except Exception as e: - warnings.warn(f"Failed to delete predictor: {e}") - finally: - self.predictor = None - # Clean up temporary video file if self.temp_video_path is not None: try: @@ -299,9 +291,15 @@ class SAM2VideoMatting: torch.cuda.empty_cache() # Force garbage collection (critical for memory leak prevention) - import gc gc.collect() + + # Clear predictor reference (but don't delete the object itself) + self.predictor = None def __del__(self): """Destructor to ensure cleanup""" - self.cleanup() \ No newline at end of file + try: + self.cleanup() + except Exception: + # Ignore errors during Python shutdown + pass \ No newline at end of file