# handler.py - Hugging Face Inference Endpoints Compatible Version import time import threading import logging import importlib import gc import os from typing import Dict, Any, Optional import json # Configure logging 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): # 30 minutes default 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) # Check every 5 minutes 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] # Try to determine if predict function expects model parameter 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 = {} # Look for Python files in src/ 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}") # Look for directories 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 # Example model files for reference: # src/resnet.py """ 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() """ # src/birefnet.py """ _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 """ # requirements.txt """ torch>=1.9.0 torchvision>=0.10.0 transformers>=4.20.0 numpy>=1.21.0 Pillow>=8.3.0 """ # Example usage in HF Inference Endpoints: """ # 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, ...]]]] } """ # Environment variables you can set in HF Inference Endpoints: """ 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 """