Spaces:
Sleeping
Sleeping
File size: 6,167 Bytes
2978bba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 |
import os
import gc
import time
import torch
import logging
import threading
from enum import Enum
from typing import Dict, Any, Optional
from dataclasses import dataclass
# Setup logging
logger = logging.getLogger("ModelManager")
class ModelType(Enum):
DETECTOR = "detector"
DEMORPH_TRANSFORMER = "demorph_transformer"
DEMORPH_GAN = "demorph_gan"
DEMORPH_SD = "demorph_sd"
DEMORPH_LDM = "demorph_ldm"
LIVENESS = "liveness"
FACE_MATCHER = "face_matcher"
@dataclass
class ModelInfo:
model_type: ModelType
instance: Any
last_used: float
size_mb: float = 0
device: str = "cpu"
class ModelManager:
"""
Singleton for managing AI model lifecycles.
Prevents CUDA OOM by unloading unused models.
"""
_instance = None
_lock = threading.RLock()
def __new__(cls):
if cls._instance is None:
with cls._lock:
if cls._instance is None:
cls._instance = super(ModelManager, cls).__new__(cls)
cls._instance._initialized = False
return cls._instance
def __init__(self):
if self._initialized:
return
self._initialized = True
self.models: Dict[ModelType, ModelInfo] = {}
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Max VRAM config (approximate safety margin)
self.max_vram_usage = 0.9 # 90%
logger.info(f"ModelManager initialized on {self.device}")
def _unload_least_used(self, required_mb: float = 0):
"""Unload the least recently used model to free space."""
if not self.models:
return
# Sort by last used (oldest first)
sorted_models = sorted(self.models.items(), key=lambda x: x[1].last_used)
# Don't unload the detector if possible, it's critical
for model_type, info in sorted_models:
if model_type == ModelType.DETECTOR and len(self.models) > 1:
continue
logger.info(f"Unloading model: {model_type.value} (Last used: {time.time() - info.last_used:.1f}s ago)")
# Explicitly delete and empty cache
del info.instance
del self.models[model_type]
gc.collect()
if self.device.type == 'cuda':
torch.cuda.empty_cache()
logger.info(f"Model unloaded: {model_type.value}")
return
def _check_memory(self):
"""Check VRAM usage and unload if critical."""
if self.device.type != 'cuda':
return
try:
total_memory = torch.cuda.get_device_properties(0).total_memory
allocated = torch.cuda.memory_allocated(0)
reserved = torch.cuda.memory_reserved(0)
usage = reserved / total_memory
if usage > self.max_vram_usage:
logger.warning(f"High VRAM usage: {usage:.1%} - Unloading implicit models...")
self._unload_least_used()
except Exception as e:
logger.warning(f"Failed to check memory: {e}")
def get_model(self, model_type: ModelType, loader_func: callable, **kwargs) -> Any:
"""
Get a model instance, loading it if necessary.
Args:
model_type: Type of model
loader_func: Function that returns the model instance
**kwargs: Arguments for the loader function
"""
with self._lock:
# Update last used if exists
if model_type in self.models:
self.models[model_type].last_used = time.time()
return self.models[model_type].instance
# Check memory before loading
self._check_memory()
# If we need to load a heavy model (like SD), unload others first aggressively
if model_type in [ModelType.DEMORPH_SD, ModelType.DEMORPH_GAN]:
# these are heavy, unload everything else except maybe detector
for m_type in list(self.models.keys()):
if m_type != ModelType.DETECTOR:
self.unload_model(m_type)
logger.info(f"Loading model: {model_type.value}...")
start_time = time.time()
try:
# Load the model
instance = loader_func(**kwargs)
# Register
self.models[model_type] = ModelInfo(
model_type=model_type,
instance=instance,
last_used=time.time(),
device=str(self.device)
)
logger.info(f"Model loaded: {model_type.value} in {time.time() - start_time:.2f}s")
return instance
except RuntimeError as e:
if "out of memory" in str(e):
logger.error("CUDA OOM during load! Attempting emergency cleanup...")
self.unload_all()
torch.cuda.empty_cache()
# Try once more? Or just fail
raise e
raise e
def unload_model(self, model_type: ModelType):
"""Explicitly unload a model."""
with self._lock:
if model_type in self.models:
logger.info(f"Unloading model: {model_type.value}")
del self.models[model_type].instance
del self.models[model_type]
gc.collect()
if self.device.type == 'cuda':
torch.cuda.empty_cache()
def unload_all(self):
"""Unload all models."""
with self._lock:
logger.info("Unloading ALL models")
self.models.clear()
gc.collect()
if self.device.type == 'cuda':
torch.cuda.empty_cache()
# Global accessor
def get_model_manager():
return ModelManager()
|