| |
| import time |
| import threading |
| import logging |
| import importlib |
| import gc |
| import os |
| from typing import Dict, Any, Optional |
| import json |
|
|
| |
| logging.basicConfig(level=logging.INFO) |
| logger = logging.getLogger(__name__) |
|
|
|
|
| class ModelManager: |
| """Manages loading, unloading, and lifecycle of multiple models""" |
|
|
| def __init__(self, unload_timeout: int = 1800): |
| self.models: Dict[str, Any] = {} |
| self.model_modules: Dict[str, Any] = {} |
| self.last_used: Dict[str, float] = {} |
| self.unload_timeout = unload_timeout |
| self.lock = threading.RLock() |
| self._cleanup_thread = None |
| self._start_cleanup_thread() |
|
|
| def _start_cleanup_thread(self): |
| """Start background thread for automatic model cleanup""" |
| if self._cleanup_thread is None or not self._cleanup_thread.is_alive(): |
| self._cleanup_thread = threading.Thread( |
| target=self._cleanup_models, daemon=True |
| ) |
| self._cleanup_thread.start() |
|
|
| def _cleanup_models(self): |
| """Background thread to unload unused models""" |
| while True: |
| try: |
| time.sleep(300) |
| current_time = time.time() |
| models_to_unload = [] |
|
|
| with self.lock: |
| for model_name, last_used_time in self.last_used.items(): |
| if current_time - last_used_time > self.unload_timeout: |
| models_to_unload.append(model_name) |
|
|
| for model_name in models_to_unload: |
| self.unload_model(model_name) |
| logger.info(f"Auto-unloaded model: {model_name}") |
|
|
| except Exception as e: |
| logger.error(f"Error in cleanup thread: {e}") |
|
|
| def register_model_module(self, name: str, module: Any): |
| """Register a model module""" |
| with self.lock: |
| if not hasattr(module, "load_model"): |
| raise ValueError(f"Model module {name} must have 'load_model' function") |
| if not hasattr(module, "predict"): |
| raise ValueError(f"Model module {name} must have 'predict' function") |
|
|
| self.model_modules[name] = module |
| logger.info(f"Registered model module: {name}") |
|
|
| def load_model(self, name: str) -> bool: |
| """Load a specific model""" |
| with self.lock: |
| try: |
| if name in self.models: |
| self.last_used[name] = time.time() |
| return True |
|
|
| if name not in self.model_modules: |
| logger.error(f"Model module not found: {name}") |
| return False |
|
|
| logger.info(f"Loading model: {name}") |
| module = self.model_modules[name] |
| model = module.load_model() |
| self.models[name] = model |
| self.last_used[name] = time.time() |
| logger.info(f"Successfully loaded model: {name}") |
| return True |
|
|
| except Exception as e: |
| logger.error(f"Failed to load model {name}: {e}") |
| return False |
|
|
| def unload_model(self, name: str) -> bool: |
| """Unload a specific model""" |
| with self.lock: |
| try: |
| if name in self.models: |
| if name in self.model_modules: |
| module = self.model_modules[name] |
| if hasattr(module, "unload_model"): |
| try: |
| module.unload_model(self.models[name]) |
| except Exception as e: |
| logger.warning( |
| f"Error in custom unload for {name}: {e}" |
| ) |
|
|
| del self.models[name] |
| if name in self.last_used: |
| del self.last_used[name] |
|
|
| gc.collect() |
| logger.info(f"Unloaded model: {name}") |
| return True |
| return False |
|
|
| except Exception as e: |
| logger.error(f"Failed to unload model {name}: {e}") |
| return False |
|
|
| def get_model(self, name: str) -> Optional[Any]: |
| """Get a loaded model, loading it if necessary""" |
| with self.lock: |
| if name not in self.models: |
| if not self.load_model(name): |
| return None |
|
|
| self.last_used[name] = time.time() |
| return self.models[name] |
|
|
| def predict(self, model_name: str, **kwargs) -> Any: |
| """Make prediction using specified model""" |
| try: |
| if model_name not in self.model_modules: |
| raise ValueError(f"Unknown model: {model_name}") |
|
|
| model = self.get_model(model_name) |
| if model is None: |
| raise RuntimeError(f"Failed to load model: {model_name}") |
|
|
| module = self.model_modules[model_name] |
|
|
| |
| import inspect |
|
|
| sig = inspect.signature(module.predict) |
| if len(sig.parameters) > 0: |
| first_param = list(sig.parameters.keys())[0] |
| if first_param in ["model", "m", "net", "network"]: |
| return module.predict(model, **kwargs) |
|
|
| return module.predict(**kwargs) |
|
|
| except Exception as e: |
| logger.error(f"Prediction failed for {model_name}: {e}") |
| raise |
|
|
|
|
| def discover_models_in_src(): |
| """Discover and register models from src directory""" |
| src_path = os.path.join(os.path.dirname(__file__), "src") |
| if not os.path.exists(src_path): |
| logger.warning("src directory not found") |
| return {} |
|
|
| discovered_models = {} |
|
|
| |
| for item in os.listdir(src_path): |
| if item.endswith(".py") and not item.startswith("__"): |
| model_name = item[:-3] |
| try: |
| module = importlib.import_module(f"src.{model_name}") |
| discovered_models[model_name] = module |
| logger.info(f"Discovered model: {model_name}") |
| except Exception as e: |
| logger.warning(f"Could not load model {model_name}: {e}") |
|
|
| |
| elif os.path.isdir(os.path.join(src_path, item)) and not item.startswith("__"): |
| model_dir = os.path.join(src_path, item) |
| possible_files = ["model.py", "main.py", "__init__.py", f"{item}.py"] |
|
|
| for filename in possible_files: |
| if os.path.exists(os.path.join(model_dir, filename)): |
| try: |
| if filename == "__init__.py": |
| module = importlib.import_module(f"src.{item}") |
| else: |
| module_name = filename[:-3] |
| module = importlib.import_module( |
| f"src.{item}.{module_name}" |
| ) |
|
|
| discovered_models[item] = module |
| logger.info(f"Discovered model: {item} (from {filename})") |
| break |
| except Exception as e: |
| logger.warning( |
| f"Could not load model {item} from {filename}: {e}" |
| ) |
| continue |
|
|
| return discovered_models |
|
|
|
|
| |
|
|
| |
| """ |
| import torch |
| import torchvision.models as models |
| |
| _resnet_model = None |
| |
| def load_model(): |
| global _resnet_model |
| if _resnet_model is None: |
| _resnet_model = models.resnet50(pretrained=True) |
| _resnet_model.eval() |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| _resnet_model.to(device) |
| return _resnet_model |
| |
| def predict(model, input_tensor, **kwargs): |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| with torch.no_grad(): |
| if isinstance(input_tensor, list): |
| input_tensor = torch.tensor(input_tensor) |
| input_tensor = input_tensor.to(device) |
| output = model(input_tensor) |
| return output.cpu().numpy().tolist() |
| |
| def unload_model(model): |
| global _resnet_model |
| _resnet_model = None |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
| """ |
|
|
| |
| """ |
| _birefnet_model = None |
| |
| def load_model(): |
| global _birefnet_model |
| if _birefnet_model is None: |
| # Your BiRefNet loading code here |
| _birefnet_model = "BiRefNet_Model_Placeholder" |
| return _birefnet_model |
| |
| def predict(input_data, **kwargs): |
| if _birefnet_model is None: |
| raise RuntimeError("Model not loaded") |
| |
| # Your prediction code here |
| return f"BiRefNet prediction for {input_data}" |
| |
| def unload_model(model): |
| global _birefnet_model |
| _birefnet_model = None |
| """ |
|
|
| |
| """ |
| torch>=1.9.0 |
| torchvision>=0.10.0 |
| transformers>=4.20.0 |
| numpy>=1.21.0 |
| Pillow>=8.3.0 |
| """ |
|
|
| |
| """ |
| # Input format 1 - Single model |
| { |
| "model": "resnet", |
| "inputs": { |
| "input_tensor": [[[[0.5, 0.3, ...]]]] |
| } |
| } |
| |
| # Input format 2 - Multiple models |
| { |
| "models": [ |
| {"name": "resnet", "inputs": {"input_tensor": [...]}}, |
| {"name": "birefnet", "inputs": {"input_data": "image.jpg"}} |
| ] |
| } |
| |
| # Input format 3 - Direct inputs (uses first available model) |
| { |
| "input_tensor": [[[[0.5, 0.3, ...]]]] |
| } |
| """ |
|
|
| |
| """ |
| PRELOAD_MODELS=true # Load all models at startup (default: true) |
| MODEL_UNLOAD_TIMEOUT=1800 # Unload models after 30 minutes (default: 1800) |
| LOG_LEVEL=INFO # Logging level |
| """ |
|
|
