Create models/loader.py

models/models/loader.py

@@ -0,0 +1,515 @@
+"""
+Model loader for BackgroundFX Pro.
+Handles loading, initialization, and management of ML models.
+"""
+
+import torch
+import torch.nn as nn
+import onnxruntime as ort
+import numpy as np
+from pathlib import Path
+from typing import Dict, Optional, Any, Union, List, Tuple
+from dataclasses import dataclass
+import logging
+import gc
+import psutil
+from functools import lru_cache
+
+from .registry import ModelInfo, ModelFramework, ModelTask, ModelRegistry
+from .downloader import ModelDownloader
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class LoadedModel:
+    """Container for loaded model."""
+    model_id: str
+    model: Any  # Actual model object
+    framework: ModelFramework
+    device: str
+    memory_usage: int  # In bytes
+    load_time: float  # In seconds
+    metadata: Dict[str, Any]
+
+
+class ModelLoader:
+    """
+    Load and manage ML models with automatic memory management.
+    """
+    
+    def __init__(self,
+                 registry: ModelRegistry,
+                 device: Optional[str] = None,
+                 max_memory_gb: float = 4.0,
+                 enable_cache: bool = True):
+        """
+        Initialize model loader.
+        
+        Args:
+            registry: Model registry instance
+            device: Device to load models on ('cuda', 'cpu', 'auto')
+            max_memory_gb: Maximum memory usage in GB
+            enable_cache: Enable model caching
+        """
+        self.registry = registry
+        self.downloader = ModelDownloader(registry)
+        self.max_memory_bytes = int(max_memory_gb * 1024 * 1024 * 1024)
+        self.enable_cache = enable_cache
+        
+        # Device management
+        self.device = self._setup_device(device)
+        self.providers = self._setup_providers()
+        
+        # Model cache
+        self.loaded_models: Dict[str, LoadedModel] = {}
+        self.current_memory_usage = 0
+        
+        logger.info(f"ModelLoader initialized with device: {self.device}")
+    
+    def _setup_device(self, device: Optional[str]) -> str:
+        """Setup computation device."""
+        if device == 'auto' or device is None:
+            if torch.cuda.is_available():
+                return 'cuda'
+            elif torch.backends.mps.is_available():
+                return 'mps'
+            else:
+                return 'cpu'
+        return device
+    
+    def _setup_providers(self) -> List[str]:
+        """Setup ONNX Runtime providers."""
+        providers = []
+        
+        if self.device == 'cuda':
+            providers.extend([
+                'CUDAExecutionProvider',
+                'TensorrtExecutionProvider'
+            ])
+        elif self.device == 'mps':
+            providers.append('CoreMLExecutionProvider')
+        
+        providers.append('CPUExecutionProvider')
+        
+        return providers
+    
+    def load_model(self,
+                  model_id: str,
+                  force_reload: bool = False,
+                  device_override: Optional[str] = None) -> Optional[LoadedModel]:
+        """
+        Load a model by ID.
+        
+        Args:
+            model_id: Model ID to load
+            force_reload: Force reload even if cached
+            device_override: Override default device
+            
+        Returns:
+            Loaded model or None if failed
+        """
+        # Check cache
+        if not force_reload and model_id in self.loaded_models:
+            logger.info(f"Using cached model: {model_id}")
+            self.registry.update_model_usage(model_id)
+            return self.loaded_models[model_id]
+        
+        # Get model info
+        model_info = self.registry.get_model(model_id)
+        if not model_info:
+            logger.error(f"Model not found: {model_id}")
+            return None
+        
+        # Download if needed
+        if model_info.status != "available":
+            logger.info(f"Downloading model: {model_id}")
+            if not self.downloader.download_model(model_id):
+                logger.error(f"Failed to download model: {model_id}")
+                return None
+        
+        # Check memory
+        if not self._check_memory_available(model_info):
+            logger.warning(f"Insufficient memory for model: {model_id}")
+            self._free_memory(model_info.memory_mb * 1024 * 1024 if model_info.memory_mb else 0)
+        
+        # Load model
+        device = device_override or self.device
+        loaded = self._load_model_impl(model_info, device)
+        
+        if loaded:
+            # Cache model
+            if self.enable_cache:
+                self.loaded_models[model_id] = loaded
+                self.current_memory_usage += loaded.memory_usage
+            
+            # Update registry
+            self.registry.update_model_usage(model_id)
+            
+            logger.info(f"Successfully loaded model: {model_id}")
+            return loaded
+        
+        return None
+    
+    def _load_model_impl(self, model_info: ModelInfo, device: str) -> Optional[LoadedModel]:
+        """
+        Implementation of model loading based on framework.
+        
+        Args:
+            model_info: Model information
+            device: Device to load on
+            
+        Returns:
+            Loaded model or None
+        """
+        import time
+        start_time = time.time()
+        
+        try:
+            if model_info.framework == ModelFramework.PYTORCH:
+                model = self._load_pytorch_model(model_info, device)
+            elif model_info.framework == ModelFramework.ONNX:
+                model = self._load_onnx_model(model_info)
+            elif model_info.framework == ModelFramework.TFLITE:
+                model = self._load_tflite_model(model_info)
+            elif model_info.framework == ModelFramework.TENSORRT:
+                model = self._load_tensorrt_model(model_info)
+            else:
+                logger.error(f"Unsupported framework: {model_info.framework}")
+                return None
+            
+            if model is None:
+                return None
+            
+            # Estimate memory usage
+            memory_usage = self._estimate_model_memory(model, model_info)
+            
+            loaded = LoadedModel(
+                model_id=model_info.model_id,
+                model=model,
+                framework=model_info.framework,
+                device=device,
+                memory_usage=memory_usage,
+                load_time=time.time() - start_time,
+                metadata=model_info.config
+            )
+            
+            return loaded
+            
+        except Exception as e:
+            logger.error(f"Failed to load model {model_info.model_id}: {e}")
+            return None
+    
+    def _load_pytorch_model(self, model_info: ModelInfo, device: str) -> Optional[Any]:
+        """Load PyTorch model."""
+        try:
+            model_path = Path(model_info.local_path)
+            
+            # Load model
+            if model_path.suffix == '.pth':
+                # Load state dict
+                state_dict = torch.load(model_path, map_location=device)
+                
+                # Create model architecture (model-specific)
+                model = self._create_model_architecture(model_info)
+                if model:
+                    model.load_state_dict(state_dict)
+                else:
+                    # Try loading as complete model
+                    model = torch.load(model_path, map_location=device)
+            else:
+                # Load complete model
+                model = torch.load(model_path, map_location=device)
+            
+            # Move to device
+            if isinstance(model, nn.Module):
+                model = model.to(device)
+                model.eval()
+            
+            return model
+            
+        except Exception as e:
+            logger.error(f"PyTorch model loading failed: {e}")
+            return None
+    
+    def _load_onnx_model(self, model_info: ModelInfo) -> Optional[Any]:
+        """Load ONNX model."""
+        try:
+            model_path = str(model_info.local_path)
+            
+            # Create session options
+            sess_options = ort.SessionOptions()
+            sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+            
+            # Add providers based on device
+            providers = self.providers
+            
+            # Create inference session
+            session = ort.InferenceSession(
+                model_path,
+                sess_options=sess_options,
+                providers=providers
+            )
+            
+            return session
+            
+        except Exception as e:
+            logger.error(f"ONNX model loading failed: {e}")
+            return None
+    
+    def _load_tflite_model(self, model_info: ModelInfo) -> Optional[Any]:
+        """Load TFLite model."""
+        try:
+            import tensorflow as tf
+            
+            model_path = str(model_info.local_path)
+            
+            # Load TFLite model
+            interpreter = tf.lite.Interpreter(model_path=model_path)
+            interpreter.allocate_tensors()
+            
+            return interpreter
+            
+        except Exception as e:
+            logger.error(f"TFLite model loading failed: {e}")
+            return None
+    
+    def _load_tensorrt_model(self, model_info: ModelInfo) -> Optional[Any]:
+        """Load TensorRT model."""
+        try:
+            import tensorrt as trt
+            import pycuda.driver as cuda
+            import pycuda.autoinit
+            
+            model_path = str(model_info.local_path)
+            
+            # Load TensorRT engine
+            with open(model_path, 'rb') as f:
+                engine_data = f.read()
+            
+            runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
+            engine = runtime.deserialize_cuda_engine(engine_data)
+            context = engine.create_execution_context()
+            
+            return {'engine': engine, 'context': context}
+            
+        except Exception as e:
+            logger.error(f"TensorRT model loading failed: {e}")
+            return None
+    
+    def _create_model_architecture(self, model_info: ModelInfo) -> Optional[nn.Module]:
+        """
+        Create model architecture for specific models.
+        This would need to be implemented for each model type.
+        """
+        # Model-specific architecture creation
+        # This is where you'd define the architecture for models
+        # that are loaded as state_dicts
+        
+        if model_info.model_id == "u2net":
+            # Example: Create U2Net architecture
+            try:
+                from ..core.models import U2NET
+                return U2NET()
+            except:
+                pass
+        
+        return None
+    
+    def _estimate_model_memory(self, model: Any, model_info: ModelInfo) -> int:
+        """Estimate model memory usage in bytes."""
+        if model_info.memory_mb:
+            return model_info.memory_mb * 1024 * 1024
+        
+        # Estimate based on model type
+        if isinstance(model, nn.Module):
+            # PyTorch model
+            param_size = sum(p.numel() * p.element_size() for p in model.parameters())
+            buffer_size = sum(b.numel() * b.element_size() for b in model.buffers())
+            return param_size + buffer_size
+        
+        elif hasattr(model, 'get_inputs'):
+            # ONNX model
+            # Rough estimate based on file size
+            file_size = Path(model_info.local_path).stat().st_size
+            return int(file_size * 2)  # Account for runtime overhead
+        
+        else:
+            # Default estimate
+            return 500 * 1024 * 1024  # 500MB default
+    
+    def _check_memory_available(self, model_info: ModelInfo) -> bool:
+        """Check if enough memory is available."""
+        required = model_info.memory_mb * 1024 * 1024 if model_info.memory_mb else 500 * 1024 * 1024
+        
+        if self.device == 'cuda':
+            # Check GPU memory
+            try:
+                import torch
+                free_memory = torch.cuda.mem_get_info()[0]
+                return free_memory > required
+            except:
+                pass
+        
+        # Check system memory
+        available = psutil.virtual_memory().available
+        return available > required
+    
+    def _free_memory(self, required_bytes: int):
+        """Free memory by unloading models."""
+        if not self.enable_cache:
+            return
+        
+        # Sort models by last used time
+        models_by_usage = sorted(
+            self.loaded_models.items(),
+            key=lambda x: self.registry.models[x[0]].last_used or 0
+        )
+        
+        freed = 0
+        for model_id, loaded_model in models_by_usage:
+            if freed >= required_bytes:
+                break
+            
+            # Unload model
+            self.unload_model(model_id)
+            freed += loaded_model.memory_usage
+            
+            logger.info(f"Freed memory by unloading: {model_id}")
+    
+    def unload_model(self, model_id: str) -> bool:
+        """
+        Unload a model from memory.
+        
+        Args:
+            model_id: Model ID to unload
+            
+        Returns:
+            True if unloaded
+        """
+        if model_id in self.loaded_models:
+            loaded = self.loaded_models[model_id]
+            
+            # Clean up model
+            if isinstance(loaded.model, nn.Module):
+                del loaded.model
+                if self.device == 'cuda':
+                    torch.cuda.empty_cache()
+            else:
+                del loaded.model
+            
+            # Update tracking
+            self.current_memory_usage -= loaded.memory_usage
+            del self.loaded_models[model_id]
+            
+            # Force garbage collection
+            gc.collect()
+            
+            logger.info(f"Unloaded model: {model_id}")
+            return True
+        
+        return False
+    
+    def unload_all(self):
+        """Unload all models."""
+        model_ids = list(self.loaded_models.keys())
+        for model_id in model_ids:
+            self.unload_model(model_id)
+    
+    def get_loaded_models(self) -> List[str]:
+        """Get list of loaded model IDs."""
+        return list(self.loaded_models.keys())
+    
+    def get_memory_usage(self) -> Dict[str, Any]:
+        """Get memory usage statistics."""
+        return {
+            'current_usage_mb': self.current_memory_usage / (1024 * 1024),
+            'max_usage_mb': self.max_memory_bytes / (1024 * 1024),
+            'loaded_models': len(self.loaded_models),
+            'models': {
+                model_id: loaded.memory_usage / (1024 * 1024)
+                for model_id, loaded in self.loaded_models.items()
+            }
+        }
+    
+    def predict(self,
+               model_id: str,
+               input_data: Union[np.ndarray, torch.Tensor],
+               **kwargs) -> Optional[Any]:
+        """
+        Run prediction with a model.
+        
+        Args:
+            model_id: Model ID
+            input_data: Input data
+            **kwargs: Additional arguments
+            
+        Returns:
+            Prediction result
+        """
+        # Load model if needed
+        loaded = self.load_model(model_id)
+        if not loaded:
+            return None
+        
+        try:
+            if loaded.framework == ModelFramework.PYTORCH:
+                return self._predict_pytorch(loaded.model, input_data, **kwargs)
+            elif loaded.framework == ModelFramework.ONNX:
+                return self._predict_onnx(loaded.model, input_data, **kwargs)
+            elif loaded.framework == ModelFramework.TFLITE:
+                return self._predict_tflite(loaded.model, input_data, **kwargs)
+            else:
+                logger.error(f"Prediction not implemented for: {loaded.framework}")
+                return None
+                
+        except Exception as e:
+            logger.error(f"Prediction failed: {e}")
+            return None
+    
+    def _predict_pytorch(self, model: nn.Module, input_data: Any, **kwargs) -> Any:
+        """Run PyTorch prediction."""
+        with torch.no_grad():
+            if not isinstance(input_data, torch.Tensor):
+                input_data = torch.from_numpy(input_data)
+            
+            input_data = input_data.to(self.device)
+            output = model(input_data)
+            
+            if isinstance(output, torch.Tensor):
+                output = output.cpu().numpy()
+            
+            return output
+    
+    def _predict_onnx(self, session: ort.InferenceSession, input_data: Any, **kwargs) -> Any:
+        """Run ONNX prediction."""
+        if isinstance(input_data, torch.Tensor):
+            input_data = input_data.numpy()
+        
+        # Get input name
+        input_name = session.get_inputs()[0].name
+        
+        # Run inference
+        outputs = session.run(None, {input_name: input_data})
+        
+        return outputs[0] if len(outputs) == 1 else outputs
+    
+    def _predict_tflite(self, interpreter: Any, input_data: Any, **kwargs) -> Any:
+        """Run TFLite prediction."""
+        if isinstance(input_data, torch.Tensor):
+            input_data = input_data.numpy()
+        
+        # Get input/output details
+        input_details = interpreter.get_input_details()
+        output_details = interpreter.get_output_details()
+        
+        # Set input
+        interpreter.set_tensor(input_details[0]['index'], input_data)
+        
+        # Run inference
+        interpreter.invoke()
+        
+        # Get output
+        output = interpreter.get_tensor(output_details[0]['index'])
+        
+        return output