models/model_loaders.py

#!/usr/bin/env python3
"""
Model Loading and Memory Management
Handles lazy loading of SAM2 and MatAnyone models with caching.
(Enhanced logging, error handling, and memory safety)
"""

import os
import gc
import logging
import streamlit as st
import torch
import psutil
from contextlib import contextmanager

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@contextmanager
def torch_memory_manager():
    try:
        logger.info("[torch_memory_manager] Enter")
        yield
    finally:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
        logger.info("[torch_memory_manager] Exit, cleaned up")

def get_memory_usage():
    memory_info = {}
    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
    memory_info['ram_used'] = psutil.virtual_memory().used / 1e9
    memory_info['ram_available'] = psutil.virtual_memory().available / 1e9
    logger.info(f"[get_memory_usage] {memory_info}")
    return memory_info

def clear_model_cache():
    """Manual/debug only: Clear Streamlit resource cache and free memory."""
    logger.info("[clear_model_cache] Clearing all model caches...")
    if hasattr(st, 'cache_resource'):
        st.cache_resource.clear()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()
    logger.info("[clear_model_cache] Model cache cleared")

@st.cache_resource(show_spinner=False)
def load_sam2_predictor():
    """Load SAM2 image predictor, choosing model size based on available GPU memory."""
    try:
        logger.info("[load_sam2_predictor] Loading SAM2 image predictor...")
        from sam2.build_sam import build_sam2
        from sam2.sam2_image_predictor import SAM2ImagePredictor
        device = "cuda" if torch.cuda.is_available() else "cpu"
        logger.info(f"[load_sam2_predictor] Using device: {device}")
        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("[load_sam2_predictor] Local checkpoints not found, using Hugging Face.")
            predictor = SAM2ImagePredictor.from_pretrained(
                "facebook/sam2-hiera-large",
                device=device
            )
        else:
            memory_info = get_memory_usage()
            gpu_free = memory_info.get('gpu_free', 0)
            if device == "cuda" and gpu_free < 4.0:
                logger.warning(f"[load_sam2_predictor] Limited GPU memory ({gpu_free:.1f}GB), using smaller SAM2 model.")
                try:
                    predictor = SAM2ImagePredictor.from_pretrained(
                        "facebook/sam2-hiera-tiny",
                        device=device
                    )
                except Exception as e:
                    logger.warning(f"[load_sam2_predictor] Tiny model failed, trying small. {e}")
                    predictor = SAM2ImagePredictor.from_pretrained(
                        "facebook/sam2-hiera-small",
                        device=device
                    )
            else:
                logger.info("[load_sam2_predictor] Using local large model")
                sam2_model = build_sam2(model_cfg, checkpoint_path, device=device)
                predictor = SAM2ImagePredictor(sam2_model)
        if hasattr(predictor, 'model'):
            predictor.model.to(device)
            predictor.model.eval()
            logger.info(f"[load_sam2_predictor] SAM2 model moved to {device} and set to eval mode")
        logger.info(f"✅ SAM2 loaded successfully on {device}!")
        return predictor
    except Exception as e:
        logger.error(f"❌ Failed to load SAM2 predictor: {e}", exc_info=True)
        import traceback
        traceback.print_exc()
        return None

def load_sam2():
    """Convenience alias for legacy code: returns only the predictor object."""
    predictor = load_sam2_predictor()
    return predictor

@st.cache_resource(show_spinner=False)
def load_matanyone_processor():
    """Load MatAnyone processor (inference core) on the best available device."""
    try:
        logger.info("[load_matanyone_processor] Loading MatAnyone processor...")
        from matanyone import InferenceCore
        device = "cuda" if torch.cuda.is_available() else "cpu"
        logger.info(f"[load_matanyone_processor] MatAnyone using device: {device}")
        try:
            processor = InferenceCore("PeiqingYang/MatAnyone", device=device)
        except Exception as e:
            logger.warning(f"[load_matanyone_processor] Path warning caught: {e}")
            processor = InferenceCore("PeiqingYang/MatAnyone", device=device)  # Retry
        if hasattr(processor, 'model'):
            processor.model.to(device)
            processor.model.eval()
            logger.info(f"[load_matanyone_processor] MatAnyone model explicitly moved to {device}")
        if not hasattr(processor, 'device'):
            processor.device = device
            logger.info(f"[load_matanyone_processor] Set processor.device to {device}")
        logger.info(f"✅ MatAnyone loaded successfully on {device}!")
        return processor
    except Exception as e:
        logger.error(f"❌ Failed to load MatAnyone: {e}", exc_info=True)
        import traceback
        traceback.print_exc()
        return None

def load_matanyone():
    """Convenience alias for legacy code: returns only the processor object."""
    processor = load_matanyone_processor()
    return processor

def test_models():
    """For admin/diagnosis: attempts to load both models and returns status."""
    results = {
        'sam2': {'loaded': False, 'error': None},
        'matanyone': {'loaded': False, 'error': None}
    }
    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)
        logger.error(f"[test_models] SAM2 error: {e}", exc_info=True)
    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)
        logger.error(f"[test_models] MatAnyone error: {e}", exc_info=True)
    logger.info(f"[test_models] Results: {results}")
    return results

def log_memory_usage(stage=""):
    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"
    print(log_msg, flush=True)
    logger.info(log_msg)
    return memory_info

def check_memory_available(required_gb=2.0):
    if not torch.cuda.is_available():
        return False, 0.0
    memory_info = get_memory_usage()
    free_gb = memory_info.get('gpu_free', 0)
    logger.info(f"[check_memory_available] free_gb={free_gb}, required={required_gb}")
    return free_gb >= required_gb, free_gb

def free_memory_aggressive():
    """For emergency/manual use only! Do NOT call after every video or from UI!"""
    logger.info("[free_memory_aggressive] Performing aggressive memory cleanup...")
    print("Performing aggressive memory cleanup...", flush=True)
    clear_model_cache()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()
        try:
            torch.cuda.ipc_collect()
        except Exception:
            pass
    gc.collect()
    print("Memory cleanup complete", flush=True)
    logger.info("Memory cleanup complete")
    log_memory_usage("after cleanup")