| |
| """ |
| 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) |
| 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") |
|
|
