File size: 4,964 Bytes
a42ddbb | 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 | # 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
from model_manager import discover_models_in_src, ModelManager
from download_models import download_models
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Global model manager
_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()
# Initialize model manager
_model_manager = ModelManager()
# Change to the model directory
original_cwd = os.getcwd()
if path != ".":
os.chdir(path)
try:
# Discover and register models
discovered_models = discover_models_in_src()
for name, module in discovered_models.items():
_model_manager.register_model_module(name, module)
# Load models at startup (optional - you can also load on-demand)
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:
# Restore original working directory
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:
# Handle different input formats
if isinstance(data, dict):
# Format 1: Single model specified
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}
# Format 2: Multiple models
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}
# Format 3: Direct inputs - try to infer model or use default
else:
# Try to use the first available model as default
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)}
# Standalone functions for compatibility (if needed)
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)
|