level#

models/sam2_loader.py

@@ -1,413 +1,262 @@
 #!/usr/bin/env python3
 """
-SAM2 Loader — Robust loading and mask generation for SAM2
-========================================================
-- Loads SAM2 model with Hydra config resolution
-- Auto-downloads missing checkpoint files
-- Generates seed masks for MatAnyone
-- Aligned with torch==2.3.1+cu121 and SAM2 commit 3c76f73c1a7e7b4a2e8a0a9a3e5b92f7e6e3f2f5
-
-Changes (2025-09-17):
-- Added automatic checkpoint download functionality
-- Enhanced error handling and logging
-- Fixed missing checkpoint issue that was causing fallback mask generation
+SAM2 Loader with T4-optimized predictor wrapper
+Provides SAM2Predictor class with memory management and optimization features
 """
 
-from __future__ import annotations
-
 import os
-import sys
-import inspect
+import gc
+import torch
 import logging
-import importlib.metadata
-import urllib.request
-import urllib.error
-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)
+from pathlib import Path
+from typing import Optional, Any, Dict, List, Tuple
 
-# --------------------------------------------------------------------------------------
-# Checkpoint Download Functionality
-# --------------------------------------------------------------------------------------
-SAM2_CHECKPOINT_URLS = {
-    "sam2_hiera_large.pt": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt",
-    "sam2_hiera_base_plus.pt": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt", 
-    "sam2_hiera_small.pt": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt",
-    "sam2_hiera_tiny.pt": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt"
-}
+logger = logging.getLogger(__name__)
 
-def _download_checkpoint(checkpoint_path: str, checkpoint_name: str) -> bool:
-    """Download SAM2 checkpoint if it doesn't exist."""
-    if os.path.exists(checkpoint_path):
-        logger.info(f"Checkpoint already exists: {checkpoint_path}")
-        return True
-    
-    if checkpoint_name not in SAM2_CHECKPOINT_URLS:
-        logger.error(f"Unknown checkpoint: {checkpoint_name}. Available: {list(SAM2_CHECKPOINT_URLS.keys())}")
-        return False
-    
-    url = SAM2_CHECKPOINT_URLS[checkpoint_name]
-    logger.info(f"Downloading SAM2 checkpoint: {checkpoint_name}")
-    logger.info(f"URL: {url}")
-    logger.info(f"Destination: {checkpoint_path}")
+class SAM2Predictor:
+    """
+    T4-optimized SAM2 video predictor wrapper with memory management
+    """
     
-    try:
-        # Create directory if it doesn't exist
-        os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)
+    def __init__(self, device: torch.device, model_size: str = "small"):
+        self.device = device
+        self.model_size = model_size
+        self.predictor = None
+        self.model = None
+        self._load_predictor()
         
-        # Download with progress
-        def _progress_hook(block_num, block_size, total_size):
-            if total_size > 0:
-                percent = min(100, block_num * block_size * 100 // total_size)
-                if percent % 10 == 0:  # Log every 10%
-                    logger.info(f"Download progress: {percent}%")
+    def _load_predictor(self):
+        """Load SAM2 predictor with optimizations"""
+        try:
+            from sam2.build_sam import build_sam2_video_predictor
+            
+            # Download checkpoint if needed
+            checkpoint_path = f"./checkpoints/sam2_hiera_{self.model_size}.pt"
+            if not self._ensure_checkpoint(checkpoint_path):
+                raise RuntimeError(f"Failed to get SAM2 {self.model_size} checkpoint")
+            
+            # Build predictor
+            model_cfg = f"sam2_hiera_{self.model_size[0]}.yaml"  # small -> s, base -> b, large -> l
+            self.predictor = build_sam2_video_predictor(model_cfg, checkpoint_path, device=self.device)
+            
+            # Apply T4 optimizations
+            self._optimize_for_t4()
+            
+            logger.info(f"SAM2 {self.model_size} predictor loaded successfully")
+            
+        except ImportError as e:
+            logger.error(f"SAM2 import failed: {e}")
+            raise RuntimeError("SAM2 not available - check third_party/sam2 installation")
+        except Exception as e:
+            logger.error(f"SAM2 loading failed: {e}")
+            raise
+    
+    def _ensure_checkpoint(self, checkpoint_path: str) -> bool:
+        """Ensure checkpoint exists, download if needed"""
+        checkpoint_file = Path(checkpoint_path)
         
-        urllib.request.urlretrieve(url, checkpoint_path, reporthook=_progress_hook)
+        if checkpoint_file.exists():
+            file_size = checkpoint_file.stat().st_size / (1024**2)
+            if file_size > 50:  # At least 50MB
+                logger.info(f"SAM2 checkpoint exists: {file_size:.1f}MB")
+                return True
+            else:
+                logger.warning(f"Checkpoint too small ({file_size:.1f}MB), re-downloading")
+                checkpoint_file.unlink()
         
-        # Verify the file was downloaded
-        if os.path.exists(checkpoint_path) and os.path.getsize(checkpoint_path) > 0:
-            size_mb = os.path.getsize(checkpoint_path) / (1024 * 1024)
-            logger.info(f"Successfully downloaded {checkpoint_name} ({size_mb:.1f} MB)")
+        return self._download_checkpoint(checkpoint_path)
+    
+    def _download_checkpoint(self, checkpoint_path: str, timeout_seconds: int = 600) -> bool:
+        """Download SAM2 checkpoint"""
+        try:
+            logger.info(f"Downloading SAM2 {self.model_size} checkpoint...")
+            
+            checkpoint_file = Path(checkpoint_path)
+            checkpoint_file.parent.mkdir(parents=True, exist_ok=True)
+            
+            import requests
+            
+            # Checkpoint URLs
+            urls = {
+                "small": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt",
+                "base": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt", 
+                "large": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt"
+            }
+            
+            if self.model_size not in urls:
+                raise ValueError(f"Unknown model size: {self.model_size}")
+            
+            checkpoint_url = urls[self.model_size]
+            
+            import time
+            start_time = time.time()
+            response = requests.get(checkpoint_url, stream=True, timeout=30)
+            response.raise_for_status()
+            
+            total_size = int(response.headers.get('content-length', 0))
+            
+            temp_path = checkpoint_file.with_suffix('.download')
+            downloaded = 0
+            last_log = start_time
+            
+            with open(temp_path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=1024*1024):
+                    if chunk:
+                        f.write(chunk)
+                        downloaded += len(chunk)
+                        
+                        current_time = time.time()
+                        if current_time - start_time > timeout_seconds:
+                            raise TimeoutError(f"Download timeout after {timeout_seconds}s")
+                        
+                        # Progress logging every 15 seconds
+                        if current_time - last_log > 15:
+                            progress = (downloaded / total_size * 100) if total_size > 0 else 0
+                            speed = downloaded / (current_time - start_time) / (1024**2)
+                            logger.info(f"Download: {progress:.1f}% ({speed:.1f}MB/s)")
+                            last_log = current_time
+            
+            temp_path.rename(checkpoint_file)
+            
+            download_time = time.time() - start_time
+            speed = downloaded / download_time / (1024**2)
+            logger.info(f"Download complete: {downloaded/(1024**2):.1f}MB in {download_time:.1f}s ({speed:.1f}MB/s)")
+            
             return True
-        else:
-            logger.error(f"Download failed: {checkpoint_path} does not exist or is empty")
-            return False
             
-    except urllib.error.URLError as e:
-        logger.error(f"URL error downloading checkpoint: {e}")
-        return False
-    except Exception as e:
-        logger.error(f"Error downloading checkpoint: {e}")
-        return False
-
-# --------------------------------------------------------------------------------------
-# 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: {dict((k, v) for k, v in os.environ.items() if 'CUDA' in k or 'SAM2' in k)}")
-    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}")
-    logger.info(f"TP_SAM2 = {TP_SAM2}")
-    
-    candidate = TP_SAM2 / cfg_str
-    logger.info(f"Candidate path: {candidate}")
-    logger.info(f"Candidate exists: {candidate.exists()}")
-    
-    if candidate.exists():
-        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 = [
-        TP_SAM2 / "sam2" / cfg_str,
-        TP_SAM2 / "configs" / cfg_str,
-    ]
-    
-    for fallback in fallbacks:
-        logger.info(f"Trying fallback: {fallback}")
-        if fallback.exists():
-            if "configs" in str(fallback):
-                relative_path = "configs/" + str(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 checkpoint auto-download."""
-    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.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)
-    
-    # Handle checkpoint with auto-download
-    ckpt = os.environ.get("SAM2_CHECKPOINT", "/home/user/app/checkpoints/sam2_hiera_large.pt")
-    checkpoint_name = os.path.basename(ckpt)
+        except Exception as e:
+            logger.error(f"Checkpoint download failed: {e}")
+            if Path(checkpoint_path).exists():
+                Path(checkpoint_path).unlink()
+            return False
     
-    # Auto-download checkpoint if missing
-    if not os.path.exists(ckpt):
-        logger.info(f"SAM2 checkpoint not found: {ckpt}")
-        if not _download_checkpoint(ckpt, checkpoint_name):
-            logger.error(f"Failed to download SAM2 checkpoint: {checkpoint_name}")
-            return None, False, meta
-    else:
-        logger.info(f"Using existing SAM2 checkpoint: {ckpt}")
-
-    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:
+    def _optimize_for_t4(self):
+        """Apply T4-specific optimizations"""
         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 hasattr(self.predictor, "model") and self.predictor.model is not None:
+                self.model = self.predictor.model
+                
+                # Apply fp16 and channels_last for T4 efficiency
+                self.model = self.model.half().to(self.device)
+                self.model = self.model.to(memory_format=torch.channels_last)
+                
+                logger.info("SAM2: fp16 + channels_last applied for T4 optimization")
+                
+        except Exception as e:
+            logger.warning(f"SAM2 T4 optimization warning: {e}")
+    
+    def init_state(self, video_path: str):
+        """Initialize video processing state"""
+        if self.predictor is None:
+            raise RuntimeError("Predictor not loaded")
         
-        if sam is not None:
-            predictor = SAM2ImagePredictor(sam)
-            meta["sam2_init_ok"] = True
-            meta["sam2_device"] = device
-            logger.info("✅ SAM2 loaded successfully with auto-downloaded checkpoint")
-            return predictor, True, meta
-        else:
-            logger.error("❌ SAM2 initialization returned None")
-            return None, False, meta
-            
-    except Exception as e:
-        logger.error(f"❌ SAM2 loading failed: {e}")
-        import traceback
-        logger.error(f"SAM2 loading traceback: {traceback.format_exc()}")
-        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
-
-# --------------------------------------------------------------------------------------
-# SAM2Model Wrapper Class for app_hf.py compatibility
-# --------------------------------------------------------------------------------------
-class SAM2Model:
-    """Wrapper class for SAM2 model to match app_hf.py interface"""
+        try:
+            return self.predictor.init_state(video_path=video_path)
+        except Exception as e:
+            logger.error(f"Failed to initialize video state: {e}")
+            raise
     
-    def __init__(self, device="cuda"):
-        self.device = device
-        self.predictor = None
-        self.loaded = False
-        logger.info(f"Initializing SAM2Model on device: {device}")
+    def add_new_points(self, inference_state, frame_idx: int, obj_id: int, 
+                      points: np.ndarray, labels: np.ndarray):
+        """Add new points for tracking"""
+        if self.predictor is None:
+            raise RuntimeError("Predictor not loaded")
         
-        # Load the model immediately
-        self._load_model()
+        try:
+            return self.predictor.add_new_points(
+                inference_state=inference_state,
+                frame_idx=frame_idx,
+                obj_id=obj_id,
+                points=points,
+                labels=labels
+            )
+        except Exception as e:
+            logger.error(f"Failed to add new points: {e}")
+            raise
     
-    def _load_model(self):
-        """Load the SAM2 model"""
+    def propagate_in_video(self, inference_state, scale: float = 1.0, **kwargs):
+        """Propagate through video with optional scaling"""
+        if self.predictor is None:
+            raise RuntimeError("Predictor not loaded")
+        
         try:
-            self.predictor, self.loaded, meta = load_sam2()
-            if self.loaded:
-                logger.info("SAM2Model loaded successfully")
-            else:
-                logger.error("Failed to load SAM2Model")
+            # Use the predictor's propagate_in_video method
+            return self.predictor.propagate_in_video(inference_state, **kwargs)
         except Exception as e:
-            logger.error(f"Error loading SAM2Model: {e}")
-            self.loaded = False
+            logger.error(f"Failed to propagate in video: {e}")
+            raise
     
-    def predict(self, video_path):
-        """Generate masks for video frames"""
-        if not self.loaded:
-            logger.error("SAM2Model not loaded")
-            return None
-            
+    def prune_state(self, inference_state, keep: int):
+        """Prune SAM2 state to keep only recent frames in memory"""
         try:
-            import cv2
-            
-            # Read first frame of video to generate initial mask
-            cap = cv2.VideoCapture(video_path)
-            ret, frame = cap.read()
-            cap.release()
+            # Try to access and prune internal caches
+            # This is model-specific and may need adjustment based on SAM2 internals
+            if hasattr(inference_state, 'cached_features'):
+                # Keep only the most recent 'keep' frames
+                cached_keys = list(inference_state.cached_features.keys())
+                if len(cached_keys) > keep:
+                    keys_to_remove = cached_keys[:-keep]
+                    for key in keys_to_remove:
+                        if key in inference_state.cached_features:
+                            del inference_state.cached_features[key]
+                    logger.debug(f"Pruned {len(keys_to_remove)} old cached features")
             
-            if not ret:
-                logger.error(f"Could not read video: {video_path}")
-                return None
+            # Clear other potential caches
+            if hasattr(inference_state, 'point_inputs_per_obj'):
+                # Keep recent point inputs only
+                for obj_id in list(inference_state.point_inputs_per_obj.keys()):
+                    obj_inputs = inference_state.point_inputs_per_obj[obj_id]
+                    if len(obj_inputs) > keep:
+                        # Keep only recent entries
+                        recent_keys = sorted(obj_inputs.keys())[-keep:]
+                        new_inputs = {k: obj_inputs[k] for k in recent_keys}
+                        inference_state.point_inputs_per_obj[obj_id] = new_inputs
             
-            # Generate mask for the frame
-            mask, success = run_sam2_mask(self.predictor, frame, auto=True)
+            # Force garbage collection
+            torch.cuda.empty_cache() if self.device.type == 'cuda' else None
             
-            if success:
-                logger.info("Successfully generated mask from video")
-                return mask
-            else:
-                logger.error("Failed to generate mask from video")
-                return None
-                
         except Exception as e:
-            logger.error(f"Error predicting masks: {e}")
-            return None
+            logger.debug(f"State pruning warning: {e}")
+    
+    def clear_memory(self):
+        """Clear GPU memory aggressively"""
+        try:
+            if self.device.type == 'cuda':
+                torch.cuda.empty_cache()
+                torch.cuda.synchronize()
+                torch.cuda.ipc_collect()
+            gc.collect()
+        except Exception as e:
+            logger.warning(f"Memory clearing warning: {e}")
+    
+    def get_memory_usage(self) -> Dict[str, float]:
+        """Get current memory usage statistics"""
+        if self.device.type != 'cuda':
+            return {"allocated_gb": 0.0, "reserved_gb": 0.0, "free_gb": 0.0}
+        
+        try:
+            allocated = torch.cuda.memory_allocated(self.device) / (1024**3)
+            reserved = torch.cuda.memory_reserved(self.device) / (1024**3) 
+            free, total = torch.cuda.mem_get_info(self.device)
+            free_gb = free / (1024**3)
+            
+            return {
+                "allocated_gb": allocated,
+                "reserved_gb": reserved, 
+                "free_gb": free_gb,
+                "total_gb": total / (1024**3)
+            }
+        except Exception:
+            return {"allocated_gb": 0.0, "reserved_gb": 0.0, "free_gb": 0.0}
+    
+    def __del__(self):
+        """Cleanup on deletion"""
+        try:
+            if hasattr(self, 'predictor') and self.predictor is not None:
+                del self.predictor
+            if hasattr(self, 'model') and self.model is not None:
+                del self.model
+            self.clear_memory()
+        except Exception:
+            pass