| |
| import time |
| import threading |
| import logging |
| import importlib |
| import gc |
| import os |
| from typing import Dict, Any, Optional |
| import json |
| from model_manager import discover_models_in_src, ModelManager |
| from download_models import download_models |
|
|
| |
| logging.basicConfig(level=logging.INFO) |
| logger = logging.getLogger(__name__) |
|
|
|
|
| |
| _model_manager = None |
|
|
|
|
| class EndpointHandler: |
| """ |
| Hugging Face Inference Endpoints compatible handler for multiple models |
| |
| This class implements the required interface for HF Inference Endpoints: |
| - __init__(self, path=".") |
| - __call__(self, data) |
| """ |
|
|
| def __init__(self, path="."): |
| """Initialize the endpoint handler""" |
| global _model_manager |
|
|
| logger.info(f"Initializing EndpointHandler with path: {path}") |
| download_models() |
|
|
| |
| _model_manager = ModelManager() |
|
|
| |
| original_cwd = os.getcwd() |
| if path != ".": |
| os.chdir(path) |
|
|
| try: |
| |
| discovered_models = discover_models_in_src() |
|
|
| for name, module in discovered_models.items(): |
| _model_manager.register_model_module(name, module) |
|
|
| |
| self.preload_models = os.getenv("PRELOAD_MODELS", "true").lower() == "true" |
| if self.preload_models and discovered_models: |
| logger.info("Preloading all models...") |
| for model_name in discovered_models.keys(): |
| _model_manager.load_model(model_name) |
|
|
| logger.info(f"Handler initialized with {len(discovered_models)} models") |
|
|
| finally: |
| |
| os.chdir(original_cwd) |
|
|
| def __call__(self, data): |
| """ |
| Handle inference requests |
| |
| Expected input formats: |
| 1. Single model: {"model": "resnet", "inputs": {...}} |
| 2. Multiple models: {"models": [{"name": "resnet", "inputs": {...}}, {"name": "yolo", "inputs": {...}}]} |
| 3. Direct inputs (uses default model): {"input_tensor": [...]} |
| """ |
| try: |
| |
| if isinstance(data, dict): |
|
|
| |
| if "model" in data and "inputs" in data: |
| model_name = data["model"] |
| inputs = data["inputs"] |
| result = _model_manager.predict(model_name, **inputs) |
| return {"model": model_name, "prediction": result} |
|
|
| |
| elif "models" in data: |
| results = [] |
| for model_request in data["models"]: |
| model_name = model_request["name"] |
| inputs = model_request.get("inputs", {}) |
| result = _model_manager.predict(model_name, **inputs) |
| results.append({"model": model_name, "prediction": result}) |
| return {"results": results} |
|
|
| |
| else: |
| |
| available_models = list(_model_manager.model_modules.keys()) |
| if not available_models: |
| raise ValueError("No models available") |
|
|
| default_model = available_models[0] |
| logger.info(f"Using default model: {default_model}") |
|
|
| result = _model_manager.predict(default_model, **data) |
| return {"model": default_model, "prediction": result} |
|
|
| else: |
| raise ValueError("Invalid input format. Expected dictionary.") |
|
|
| except Exception as e: |
| logger.error(f"Inference error: {e}") |
| return {"error": str(e)} |
|
|
|
|
| |
| def load_models(model_names: Optional[list] = None) -> Dict[str, bool]: |
| """Load specified models or all models if none specified""" |
| if _model_manager is None: |
| raise RuntimeError("Handler not initialized") |
|
|
| if model_names is None: |
| results = {} |
| for model_name in _model_manager.model_modules.keys(): |
| results[model_name] = _model_manager.load_model(model_name) |
| return results |
|
|
| results = {} |
| for name in model_names: |
| results[name] = _model_manager.load_model(name) |
| return results |
|
|
|
|
| def predict(model_name: str, **kwargs) -> Any: |
| """Make prediction using specified model""" |
| if _model_manager is None: |
| raise RuntimeError("Handler not initialized") |
| return _model_manager.predict(model_name, **kwargs) |
|
|
