""" Model loading and management system with password protection. This module provides secure loading of PyTorch models from password-protected ZIP files (.pth) with metadata validation and model verification. """ import json import zipfile import tempfile import hashlib import io import warnings from pathlib import Path from typing import Dict, Any, Optional, Union, List, Tuple from dataclasses import dataclass, asdict from datetime import datetime import torch import torch.nn as nn @dataclass class ModelMetadata: """Metadata structure for wakeword models.""" # Model identification model_name: str model_version: str model_type: str # "wakeword", "voice_recognition" architecture: str # "RepCNN", "ImprovedRepCNN", etc. # Training information training_date: str author: str description: str computer_name: str # Model configuration num_classes: int input_shape: List[int] # [channels, height, width] class_labels: List[str] # ["custom", "system_command", "negative"] # Training results final_accuracy: float final_loss: float num_epochs: int training_duration_hours: float # Audio configuration sample_rate: int n_mels: int time_frames: int chunk_duration_ms: float # Model performance inference_time_ms: float model_size_mb: float flops_estimate: int # Security and validation model_hash: str created_timestamp: str @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'ModelMetadata': """Create ModelMetadata from dictionary.""" return cls(**data) def to_dict(self) -> Dict[str, Any]: """Convert ModelMetadata to dictionary.""" return asdict(self) def validate(self) -> bool: """Validate metadata completeness and consistency.""" try: # Check required fields required_fields = [ 'model_name', 'model_version', 'model_type', 'architecture', 'num_classes', 'input_shape', 'class_labels' ] for field in required_fields: value = getattr(self, field) if value is None or (isinstance(value, (str, list)) and len(value) == 0): return False # Validate consistency if len(self.class_labels) != self.num_classes: return False if len(self.input_shape) != 3: # Should be [C, H, W] return False # Check for wakeword-specific requirements if self.model_type == "wakeword": expected_labels = ["custom", "system_command", "negative"] if not all(label in self.class_labels for label in expected_labels): return False return True except Exception: return False class ModelLoadError(Exception): """Exception raised when model loading fails.""" pass class SecureModelLoader: """ Secure model loader for password-protected PyTorch models. Handles loading of .pth files that are actually ZIP archives containing: - model.pth: The actual PyTorch model state dict - metadata.json: Model metadata and configuration - Optional additional files (training logs, etc.) """ def __init__(self, device: Optional[torch.device] = None): """ Initialize the secure model loader. Args: device: PyTorch device for loading models """ self.device = device or torch.device("cpu") self._loaded_models: Dict[str, Dict[str, Any]] = {} def load_model(self, model_path: Union[str, Path], password: str, model_class: Optional[type] = None, verify_hash: bool = True) -> Tuple[nn.Module, ModelMetadata]: """ Load a password-protected PyTorch model. Args: model_path: Path to the .pth file (ZIP archive) password: Password for decrypting the model model_class: Model class to instantiate (if None, attempts auto-detection) verify_hash: Whether to verify model integrity using hash Returns: Tuple of (model, metadata) Raises: ModelLoadError: If loading fails """ model_path = Path(model_path) if not model_path.exists(): raise ModelLoadError(f"Model file not found: {model_path}") try: # Load from ZIP archive with zipfile.ZipFile(model_path, 'r') as zip_file: # Try to extract with password try: zip_file.setpassword(password.encode('utf-8')) # Load metadata metadata = self._load_metadata(zip_file) # Load model state dict model_state = self._load_model_state(zip_file) # Verify hash if requested if verify_hash and metadata.model_hash: if not self._verify_model_hash(model_state, metadata.model_hash): raise ModelLoadError("Model hash verification failed") # Create model instance model = self._create_model_instance(metadata, model_class) # Reparameterize model if it has the method (RepCNN models) # This converts the model structure to match the saved state dict if hasattr(model, 'reparameterize'): model.reparameterize() # Load state dict model.load_state_dict(model_state, strict=True) model.to(self.device) model.eval() # Cache the loaded model cache_key = str(model_path) self._loaded_models[cache_key] = { 'model': model, 'metadata': metadata, 'load_time': datetime.now() } return model, metadata except RuntimeError as e: if "Bad password" in str(e) or "Bad CRC" in str(e): raise ModelLoadError("Invalid password for model file") else: raise ModelLoadError(f"Failed to extract model: {e}") except zipfile.BadZipFile: raise ModelLoadError("Invalid ZIP file format") except Exception as e: raise ModelLoadError(f"Failed to load model: {e}") def _load_metadata(self, zip_file: zipfile.ZipFile) -> ModelMetadata: """Load and validate metadata from ZIP file.""" try: # Try metafile.json first (used by trainer), then metadata.json (legacy) try: with zip_file.open('metafile.json') as meta_file: raw_metadata = json.load(meta_file) # Convert trainer format to ModelMetadata format metadata_dict = self._convert_trainer_metadata(raw_metadata) except KeyError: with zip_file.open('metadata.json') as meta_file: metadata_dict = json.load(meta_file) metadata = ModelMetadata.from_dict(metadata_dict) if not metadata.validate(): raise ModelLoadError("Invalid metadata structure") return metadata except KeyError: raise ModelLoadError("metadata.json/metafile.json not found in model file") except json.JSONDecodeError as e: raise ModelLoadError(f"Invalid JSON in metadata: {e}") def _convert_trainer_metadata(self, trainer_meta: Dict[str, Any]) -> Dict[str, Any]: """Convert trainer metafile.json format to ModelMetadata format.""" # Extract nested data training_stats = trainer_meta.get('training_statistics', {}) audio_cfg = trainer_meta.get('audio_config', {}) arch_info = trainer_meta.get('architecture_info', {}) hardware = trainer_meta.get('hardware_info', {}) model_specific = trainer_meta.get('model_specific', {}) # Calculate time_frames from audio config audio_length_sec = audio_cfg.get('audio_length_seconds', 1.5) hop_length = audio_cfg.get('hop_length', 160) sample_rate = audio_cfg.get('sample_rate', 16000) time_frames = int((audio_length_sec * sample_rate) / hop_length) + 1 # Get class labels from model_specific.wakewords or fallback class_labels = model_specific.get('wakewords', ['custom', 'system_command', 'negative']) num_classes = len(class_labels) # Build input_shape n_mels = audio_cfg.get('n_mels', 40) input_shape = [1, n_mels, time_frames] # Map to ModelMetadata fields return { 'model_name': trainer_meta.get('model_name', 'unknown'), 'model_version': trainer_meta.get('version', '1.0.0'), 'model_type': trainer_meta.get('model_type', 'wakeword'), 'architecture': arch_info.get('architecture_name', 'RepCNN'), 'training_date': trainer_meta.get('creation_timestamp', ''), 'author': trainer_meta.get('author', 'Unknown'), 'description': trainer_meta.get('description', ''), 'computer_name': hardware.get('hostname', 'Unknown'), 'num_classes': num_classes, 'input_shape': input_shape, 'class_labels': class_labels, 'final_accuracy': training_stats.get('best_val_accuracy', 0.0), 'final_loss': training_stats.get('best_val_loss', 0.0), 'num_epochs': training_stats.get('total_epochs', 0), 'training_duration_hours': training_stats.get('training_duration_seconds', 0.0) / 3600.0, 'sample_rate': sample_rate, 'n_mels': n_mels, 'time_frames': time_frames, 'chunk_duration_ms': audio_length_sec * 1000, 'inference_time_ms': 0.0, 'model_size_mb': training_stats.get('total_parameters', 0) * 4 / (1024 * 1024), 'flops_estimate': 0, 'model_hash': '', # Not available in trainer format 'created_timestamp': trainer_meta.get('creation_timestamp', ''), } def _load_model_state(self, zip_file: zipfile.ZipFile) -> Dict[str, torch.Tensor]: """Load model state dict from ZIP file.""" try: with zip_file.open('model.pth') as model_file: # Load state dict directly from bytes model_bytes = model_file.read() state_dict = torch.load( io.BytesIO(model_bytes), map_location=self.device, weights_only=True ) return state_dict except KeyError: raise ModelLoadError("model.pth not found in model file") except Exception as e: raise ModelLoadError(f"Failed to load model state dict: {e}") def _verify_model_hash(self, model_state: Dict[str, torch.Tensor], expected_hash: str) -> bool: """Verify model integrity using hash.""" try: # Create hash of model parameters hasher = hashlib.sha256() # Sort keys for consistent hashing for key in sorted(model_state.keys()): tensor = model_state[key] hasher.update(key.encode('utf-8')) hasher.update(tensor.detach().cpu().numpy().tobytes()) calculated_hash = hasher.hexdigest() return calculated_hash == expected_hash except Exception: return False def _create_model_instance(self, metadata: ModelMetadata, model_class: Optional[type] = None) -> nn.Module: """Create model instance based on metadata.""" if model_class is not None: # Use provided model class return self._instantiate_model_class(model_class, metadata) # Attempt auto-detection based on metadata architecture = metadata.architecture.lower() if architecture in ["repcnn", "standard"]: from trainer.wakeword.models import RepCNN return self._instantiate_repcnn(RepCNN, metadata) elif architecture in ["improvedrepcnn", "improved"]: from trainer.wakeword.models import ImprovedRepCNN return self._instantiate_repcnn(ImprovedRepCNN, metadata) elif architecture in ["lightweightrepcnn", "lightweight"]: from trainer.wakeword.models import LightweightRepCNN return self._instantiate_lightweight_repcnn(metadata) else: raise ModelLoadError(f"Unknown architecture: {metadata.architecture}") def _instantiate_repcnn(self, model_class: type, metadata: ModelMetadata) -> nn.Module: """Instantiate RepCNN-based models.""" return model_class( num_classes=metadata.num_classes, input_channels=metadata.input_shape[0], num_mels=metadata.n_mels, time_frames=metadata.time_frames ) def _instantiate_lightweight_repcnn(self, metadata: ModelMetadata) -> nn.Module: """Instantiate LightweightRepCNN.""" from trainer.wakeword.models import LightweightRepCNN return LightweightRepCNN( num_classes=metadata.num_classes, input_channels=metadata.input_shape[0], num_mels=metadata.n_mels, time_frames=metadata.time_frames ) def _instantiate_model_class(self, model_class: type, metadata: ModelMetadata) -> nn.Module: """Instantiate arbitrary model class with metadata parameters.""" try: # Try different initialization patterns init_params = { 'num_classes': metadata.num_classes, 'input_channels': metadata.input_shape[0], 'num_mels': metadata.n_mels, 'time_frames': metadata.time_frames } return model_class(**init_params) except Exception as e: raise ModelLoadError(f"Failed to instantiate model class {model_class.__name__}: {e}") def get_loaded_models(self) -> List[str]: """Get list of currently loaded model paths.""" return list(self._loaded_models.keys()) def unload_model(self, model_path: str) -> bool: """ Unload a model from cache. Args: model_path: Path of model to unload Returns: True if model was unloaded, False if not found """ if model_path in self._loaded_models: del self._loaded_models[model_path] return True return False def clear_cache(self): """Clear all cached models.""" self._loaded_models.clear() def get_model_info(self, model_path: str) -> Optional[Dict[str, Any]]: """Get information about a loaded model.""" return self._loaded_models.get(model_path) class ModelLoader: """ Standard model loader for unprotected PyTorch models. This is a simpler loader for development and testing purposes. """ def __init__(self, device: Optional[torch.device] = None): """Initialize the model loader.""" self.device = device or torch.device("cpu") def load_model(self, model_path: Union[str, Path], model_class: type, **model_kwargs) -> nn.Module: """ Load a standard PyTorch model. Args: model_path: Path to the .pth file model_class: Model class to instantiate **model_kwargs: Arguments for model initialization Returns: Loaded PyTorch model """ model_path = Path(model_path) if not model_path.exists(): raise ModelLoadError(f"Model file not found: {model_path}") try: # Create model instance model = model_class(**model_kwargs) # Load state dict state_dict = torch.load(model_path, map_location=self.device, weights_only=True) model.load_state_dict(state_dict, strict=True) model.to(self.device) model.eval() return model except Exception as e: raise ModelLoadError(f"Failed to load model: {e}") def validate_model_file(model_path: Union[str, Path], password: Optional[str] = None) -> bool: """ Validate a model file without fully loading it. Args: model_path: Path to model file password: Password for encrypted models (optional) Returns: True if model file is valid, False otherwise """ model_path = Path(model_path) if not model_path.exists(): return False try: with zipfile.ZipFile(model_path, 'r') as zip_file: if password: zip_file.setpassword(password.encode('utf-8')) # Check for required files file_list = zip_file.namelist() # Check for model.pth if 'model.pth' not in file_list: return False # Check for metadata file (metafile.json or metadata.json) has_metadata = 'metafile.json' in file_list or 'metadata.json' in file_list if not has_metadata: return False # Try to load metadata try: meta_file_name = 'metafile.json' if 'metafile.json' in file_list else 'metadata.json' with zip_file.open(meta_file_name) as meta_file: metadata_dict = json.load(meta_file) metadata = ModelMetadata.from_dict(metadata_dict) return metadata.validate() except KeyError: return False except (zipfile.BadZipFile, json.JSONDecodeError, RuntimeError): return False except Exception: return False def create_model_metadata(model: nn.Module, training_config: Dict[str, Any], training_results: Dict[str, Any], audio_config: Dict[str, Any]) -> ModelMetadata: """ Create model metadata from training information. Args: model: Trained PyTorch model training_config: Training configuration training_results: Training results and metrics audio_config: Audio processing configuration Returns: ModelMetadata instance """ # Calculate model hash hasher = hashlib.sha256() for name, param in model.named_parameters(): hasher.update(name.encode('utf-8')) hasher.update(param.detach().cpu().numpy().tobytes()) model_hash = hasher.hexdigest() # Estimate model size param_count = sum(p.numel() for p in model.parameters()) model_size_mb = param_count * 4 / (1024 * 1024) # Assuming float32 return ModelMetadata( model_name=training_config.get('model_name', 'wakeword_model'), model_version=training_config.get('model_version', '1.0.0'), model_type="wakeword", architecture=model.__class__.__name__, training_date=datetime.now().isoformat(), author=training_config.get('author', 'Unknown'), description=training_config.get('description', ''), computer_name=training_config.get('computer_name', 'Unknown'), num_classes=training_config.get('num_classes', 3), input_shape=[1, audio_config.get('n_mels', 40), audio_config.get('time_frames', 151)], class_labels=training_config.get('class_labels', ['custom', 'system_command', 'negative']), final_accuracy=training_results.get('accuracy', 0.0), final_loss=training_results.get('loss', 0.0), num_epochs=training_results.get('epochs', 0), training_duration_hours=training_results.get('duration_hours', 0.0), sample_rate=audio_config.get('sample_rate', 16000), n_mels=audio_config.get('n_mels', 40), time_frames=audio_config.get('time_frames', 151), chunk_duration_ms=audio_config.get('chunk_duration_ms', 1000), inference_time_ms=0.0, # To be measured during deployment model_size_mb=model_size_mb, flops_estimate=0, # To be calculated if needed model_hash=model_hash, created_timestamp=datetime.now().isoformat() )