Create models.py

models.py

@@ -0,0 +1,212 @@
+#!/usr/bin/env python3
+"""
+Model Loading and Memory Management
+Handles lazy loading of SAM2 and MatAnyone models with caching
+"""
+
+import os
+import gc
+import logging
+import streamlit as st
+import torch
+import psutil
+
+logger = logging.getLogger(__name__)
+
+# Context manager for CUDA memory cleanup
+from contextlib import contextmanager
+
+@contextmanager
+def torch_memory_manager():
+    """Context manager for CUDA memory cleanup."""
+    try:
+        yield
+    finally:
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
+
+def get_memory_usage():
+    """Get current memory usage statistics."""
+    memory_info = {}
+    
+    # GPU memory if available
+    if torch.cuda.is_available():
+        memory_info['gpu_allocated'] = torch.cuda.memory_allocated() / 1e9
+        memory_info['gpu_reserved'] = torch.cuda.memory_reserved() / 1e9
+        memory_info['gpu_free'] = (torch.cuda.get_device_properties(0).total_memory - 
+                                  torch.cuda.memory_allocated()) / 1e9
+    
+    # RAM memory
+    memory_info['ram_used'] = psutil.virtual_memory().used / 1e9
+    memory_info['ram_available'] = psutil.virtual_memory().available / 1e9
+    
+    return memory_info
+
+def clear_model_cache():
+    """Clear all cached models and free memory."""
+    if hasattr(st, 'cache_resource'):
+        st.cache_resource.clear()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
+    logger.info("Model cache cleared")
+
+# ============================================================================
+# SAM2 Model Loading
+# ============================================================================
+
+@st.cache_resource(show_spinner=False)
+def load_sam2_predictor():
+    """
+    Lazy load SAM2 image predictor with fallback strategies.
+    Returns None if loading fails.
+    """
+    try:
+        logger.info("Loading SAM2 image predictor...")
+        from sam2.build_sam import build_sam2
+        from sam2.sam2_image_predictor import SAM2ImagePredictor
+        
+        # Try local checkpoints first
+        checkpoint_path = "/home/user/app/checkpoints/sam2.1_hiera_large.pt"
+        model_cfg = "/home/user/app/configs/sam2.1/sam2.1_hiera_l.yaml"
+        
+        if not os.path.exists(checkpoint_path) or not os.path.exists(model_cfg):
+            logger.warning("Local checkpoints not found, using Hugging Face...")
+            predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
+        else:
+            # Check available GPU memory
+            memory_info = get_memory_usage()
+            gpu_free = memory_info.get('gpu_free', 0)
+            
+            if gpu_free < 4.0:
+                logger.warning(f"Limited GPU memory ({gpu_free:.1f}GB), using smaller SAM2 model...")
+                try:
+                    predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-tiny")
+                except Exception:
+                    predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-small")
+            else:
+                # Use local large model
+                predictor = SAM2ImagePredictor(build_sam2(model_cfg, checkpoint_path))
+        
+        logger.info("✅ SAM2 image predictor loaded successfully!")
+        return predictor
+        
+    except Exception as e:
+        logger.error(f"Failed to load SAM2 predictor: {e}")
+        return None
+
+# ============================================================================
+# MatAnyone Model Loading
+# ============================================================================
+
+@st.cache_resource(show_spinner=False)
+def load_matanyone_processor():
+    """
+    Lazy load MatAnyone processor.
+    Returns None if loading fails.
+    """
+    try:
+        logger.info("Loading MatAnyone processor...")
+        from matanyone import InferenceCore
+        
+        processor = InferenceCore("PeiqingYang/MatAnyone")
+        
+        logger.info("✅ MatAnyone processor loaded successfully!")
+        return processor
+        
+    except Exception as e:
+        logger.error(f"Failed to load MatAnyone: {e}")
+        return None
+
+# ============================================================================
+# Model Health Check
+# ============================================================================
+
+def test_models():
+    """
+    Test if both models can load successfully.
+    Returns dict with test results.
+    """
+    results = {
+        'sam2': {'loaded': False, 'error': None},
+        'matanyone': {'loaded': False, 'error': None}
+    }
+    
+    # Test SAM2
+    try:
+        sam2_predictor = load_sam2_predictor()
+        if sam2_predictor is not None:
+            results['sam2']['loaded'] = True
+        else:
+            results['sam2']['error'] = "Predictor returned None"
+    except Exception as e:
+        results['sam2']['error'] = str(e)
+    
+    # Test MatAnyone
+    try:
+        matanyone_processor = load_matanyone_processor()
+        if matanyone_processor is not None:
+            results['matanyone']['loaded'] = True
+        else:
+            results['matanyone']['error'] = "Processor returned None"
+    except Exception as e:
+        results['matanyone']['error'] = str(e)
+    
+    return results
+
+# ============================================================================
+# Memory Monitoring
+# ============================================================================
+
+def log_memory_usage(stage=""):
+    """Log current memory usage with optional stage label."""
+    memory_info = get_memory_usage()
+    
+    log_msg = f"Memory usage"
+    if stage:
+        log_msg += f" ({stage})"
+    log_msg += ":"
+    
+    if 'gpu_allocated' in memory_info:
+        log_msg += f" GPU {memory_info['gpu_allocated']:.1f}GB allocated, {memory_info['gpu_free']:.1f}GB free"
+    
+    log_msg += f" | RAM {memory_info['ram_used']:.1f}GB used"
+    
+    logger.info(log_msg)
+    return memory_info
+
+def check_memory_available(required_gb=2.0):
+    """
+    Check if enough GPU memory is available.
+    Returns (bool, float) - (is_available, free_gb)
+    """
+    if not torch.cuda.is_available():
+        return False, 0.0
+    
+    memory_info = get_memory_usage()
+    free_gb = memory_info.get('gpu_free', 0)
+    
+    return free_gb >= required_gb, free_gb
+
+def free_memory_aggressive():
+    """Aggressively free GPU and system memory."""
+    logger.info("Performing aggressive memory cleanup...")
+    
+    # Clear model cache
+    clear_model_cache()
+    
+    # CUDA cleanup
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+        try:
+            torch.cuda.ipc_collect()
+        except Exception:
+            pass
+    
+    # System cleanup
+    gc.collect()
+    
+    logger.info("Memory cleanup complete")
+    log_memory_usage("after cleanup")