""" Model format handlers for different file formats (.pth, .pt, .onnx). This module provides handlers for saving and loading models in various formats, including password protection and metadata integration as specified in the project requirements. """ import os import json import zipfile import tempfile import hashlib import io from abc import ABC, abstractmethod from typing import Dict, Any, Optional, Union, BinaryIO from pathlib import Path import torch import torch.nn as nn from cryptography.fernet import Fernet from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC import base64 import logging try: import onnx import onnxruntime as ort ONNX_AVAILABLE = True except ImportError: ONNX_AVAILABLE = False onnx = None ort = None class ModelFormatError(Exception): """Exception raised for model format related errors.""" pass class EncryptionManager: """Handles encryption and decryption for password-protected models.""" @staticmethod def generate_key_from_password(password: str, salt: bytes = None) -> bytes: """Generate encryption key from password using PBKDF2.""" if salt is None: salt = os.urandom(16) kdf = PBKDF2HMAC( algorithm=hashes.SHA256(), length=32, salt=salt, iterations=100000, ) key = base64.urlsafe_b64encode(kdf.derive(password.encode())) return key, salt @staticmethod def encrypt_data(data: bytes, password: str) -> bytes: """Encrypt data with password.""" key, salt = EncryptionManager.generate_key_from_password(password) f = Fernet(key) encrypted_data = f.encrypt(data) # Prepend salt to encrypted data return salt + encrypted_data @staticmethod def decrypt_data(encrypted_data: bytes, password: str) -> bytes: """Decrypt data with password.""" # Extract salt from the beginning salt = encrypted_data[:16] encrypted_content = encrypted_data[16:] key, _ = EncryptionManager.generate_key_from_password(password, salt) f = Fernet(key) return f.decrypt(encrypted_content) class ModelFormatHandler(ABC): """Abstract base class for model format handlers.""" def __init__(self, logger: Optional[logging.Logger] = None): self.logger = logger or logging.getLogger(__name__) @abstractmethod def save_model(self, model: nn.Module, filepath: str, metadata: Dict[str, Any], password: Optional[str] = None, **kwargs) -> bool: """Save model with metadata and optional password protection.""" pass @abstractmethod def load_model(self, filepath: str, password: Optional[str] = None, map_location: Optional[str] = None) -> tuple: """Load model and return (model_state_dict, metadata).""" pass @abstractmethod def extract_metadata(self, filepath: str, password: Optional[str] = None) -> Dict[str, Any]: """Extract metadata from model file without loading the full model.""" pass @abstractmethod def update_metadata(self, filepath: str, new_metadata: Dict[str, Any], password: Optional[str] = None) -> bool: """Update metadata in existing model file.""" pass def _generate_file_hash(self, filepath: str) -> str: """Generate SHA256 hash of the file.""" hash_sha256 = hashlib.sha256() with open(filepath, "rb") as f: for chunk in iter(lambda: f.read(4096), b""): hash_sha256.update(chunk) return hash_sha256.hexdigest() class PTHHandler(ModelFormatHandler): """Handler for .pth files (ZIP archives with JSON metadata).""" def save_model(self, model: nn.Module, filepath: str, metadata: Dict[str, Any], password: Optional[str] = None, **kwargs) -> bool: """ Save model as .pth file (ZIP archive) with metadata.json. Args: model: PyTorch model to save filepath: Path to save the model metadata: Metadata dictionary to include password: Optional password for encryption **kwargs: Additional arguments Returns: bool: True if successful """ try: # Ensure directory exists os.makedirs(os.path.dirname(filepath), exist_ok=True) # Create temporary directory for building the archive with tempfile.TemporaryDirectory() as temp_dir: # Save model state dict model_path = os.path.join(temp_dir, "model.pth") torch.save(model.state_dict(), model_path) # Save metadata metadata_path = os.path.join(temp_dir, "metafile.json") metadata['file_format'] = 'pth' metadata['pytorch_version'] = torch.__version__ metadata['file_hash'] = None # Will be calculated after creation with open(metadata_path, 'w', encoding='utf-8') as f: json.dump(metadata, f, indent=2, ensure_ascii=False) # Create ZIP archive with zipfile.ZipFile(filepath, 'w', zipfile.ZIP_DEFLATED) as zipf: zipf.write(model_path, "model.pth") zipf.write(metadata_path, "metafile.json") # Add any additional files from kwargs for key, value in kwargs.items(): if key.startswith('add_file_'): file_path = value.get('path') archive_name = value.get('name', os.path.basename(file_path)) if os.path.exists(file_path): zipf.write(file_path, archive_name) # Apply password protection if requested if password: self._encrypt_file(filepath, password) # Update file hash in metadata file_hash = self._generate_file_hash(filepath) self._update_file_hash(filepath, file_hash, password) self.logger.info(f"Model saved successfully: {filepath}") return True except Exception as e: self.logger.error(f"Failed to save model: {str(e)}") return False def load_model(self, filepath: str, password: Optional[str] = None, map_location: Optional[str] = None) -> tuple: """ Load model from .pth file. Returns: tuple: (model_state_dict, metadata) """ try: # Handle password protection if password: filepath = self._decrypt_file(filepath, password) with zipfile.ZipFile(filepath, 'r') as zipf: # Load metadata with zipf.open("metafile.json") as f: metadata = json.load(f) # Load model with zipf.open("model.pth") as f: model_data = f.read() model_state_dict = torch.load( io.BytesIO(model_data), map_location=map_location ) self.logger.info(f"Model loaded successfully: {filepath}") return model_state_dict, metadata except Exception as e: self.logger.error(f"Failed to load model: {str(e)}") raise ModelFormatError(f"Failed to load model: {str(e)}") def extract_metadata(self, filepath: str, password: Optional[str] = None) -> Dict[str, Any]: """Extract metadata from .pth file without loading the model.""" try: # Handle password protection if password: filepath = self._decrypt_file(filepath, password) with zipfile.ZipFile(filepath, 'r') as zipf: with zipf.open("metafile.json") as f: metadata = json.load(f) return metadata except Exception as e: self.logger.error(f"Failed to extract metadata: {str(e)}") return {} def update_metadata(self, filepath: str, new_metadata: Dict[str, Any], password: Optional[str] = None) -> bool: """Update metadata in existing .pth file.""" try: # Load existing data model_state_dict, existing_metadata = self.load_model(filepath, password) # Merge metadata existing_metadata.update(new_metadata) # Create temporary model for saving temp_model = nn.Module() temp_model.load_state_dict = lambda: None temp_model.state_dict = lambda: model_state_dict # Save with updated metadata return self.save_model(temp_model, filepath, existing_metadata, password) except Exception as e: self.logger.error(f"Failed to update metadata: {str(e)}") return False def _encrypt_file(self, filepath: str, password: str): """Encrypt file in place.""" with open(filepath, 'rb') as f: data = f.read() encrypted_data = EncryptionManager.encrypt_data(data, password) with open(filepath, 'wb') as f: f.write(encrypted_data) def _decrypt_file(self, filepath: str, password: str) -> str: """Decrypt file to temporary location and return temp path.""" with open(filepath, 'rb') as f: encrypted_data = f.read() try: decrypted_data = EncryptionManager.decrypt_data(encrypted_data, password) except Exception: raise ModelFormatError("Invalid password or corrupted file") # Save to temporary file temp_fd, temp_path = tempfile.mkstemp(suffix='.pth') with os.fdopen(temp_fd, 'wb') as f: f.write(decrypted_data) return temp_path def _update_file_hash(self, filepath: str, file_hash: str, password: Optional[str] = None): """Update file hash in metadata.""" try: # Load and update metadata metadata = self.extract_metadata(filepath, password) metadata['file_hash'] = file_hash self.update_metadata(filepath, metadata, password) except Exception as e: self.logger.warning(f"Failed to update file hash: {str(e)}") class PTHandler(ModelFormatHandler): """Handler for .pt files using torch.package.""" def save_model(self, model: nn.Module, filepath: str, metadata: Dict[str, Any], password: Optional[str] = None, **kwargs) -> bool: """ Save model as .pt file using torch.package. Args: model: PyTorch model to save filepath: Path to save the model metadata: Metadata dictionary to include password: Optional password for encryption **kwargs: Additional arguments Returns: bool: True if successful """ try: # Ensure directory exists os.makedirs(os.path.dirname(filepath), exist_ok=True) # Create package with torch.package.PackageExporter(filepath) as exp: # Add model exp.intern('**') exp.save_pickle('model', 'model.pkl', model.state_dict()) # Add metadata metadata['file_format'] = 'pt' metadata['pytorch_version'] = torch.__version__ metadata['package_version'] = '1.0' # Save metadata as JSON string in package exp.save_text('metadata', 'metafile.json', json.dumps(metadata, indent=2, ensure_ascii=False)) # Add custom resources if specified for key, value in kwargs.items(): if key.startswith('add_resource_'): resource_name = key.replace('add_resource_', '') if isinstance(value, str): exp.save_text('resources', f'{resource_name}.txt', value) elif isinstance(value, (dict, list)): exp.save_text('resources', f'{resource_name}.json', json.dumps(value, indent=2)) # Apply password protection if requested if password: self._encrypt_file(filepath, password) # Update file hash file_hash = self._generate_file_hash(filepath) self._update_file_hash(filepath, file_hash, password) self.logger.info(f"Model saved successfully: {filepath}") return True except Exception as e: self.logger.error(f"Failed to save model: {str(e)}") return False def load_model(self, filepath: str, password: Optional[str] = None, map_location: Optional[str] = None) -> tuple: """ Load model from .pt file. Returns: tuple: (model_state_dict, metadata) """ try: # Handle password protection if password: filepath = self._decrypt_file(filepath, password) # Load package imp = torch.package.PackageImporter(filepath) # Load model state dict model_state_dict = imp.load_pickle('model', 'model.pkl', map_location=map_location) # Load metadata metadata_json = imp.load_text('metadata', 'metafile.json') metadata = json.loads(metadata_json) self.logger.info(f"Model loaded successfully: {filepath}") return model_state_dict, metadata except Exception as e: self.logger.error(f"Failed to load model: {str(e)}") raise ModelFormatError(f"Failed to load model: {str(e)}") def extract_metadata(self, filepath: str, password: Optional[str] = None) -> Dict[str, Any]: """Extract metadata from .pt file without loading the model.""" try: # Handle password protection if password: filepath = self._decrypt_file(filepath, password) imp = torch.package.PackageImporter(filepath) metadata_json = imp.load_text('metadata', 'metafile.json') metadata = json.loads(metadata_json) return metadata except Exception as e: self.logger.error(f"Failed to extract metadata: {str(e)}") return {} def update_metadata(self, filepath: str, new_metadata: Dict[str, Any], password: Optional[str] = None) -> bool: """Update metadata in existing .pt file.""" try: # Load existing data model_state_dict, existing_metadata = self.load_model(filepath, password) # Merge metadata existing_metadata.update(new_metadata) # Create temporary model for saving temp_model = nn.Module() temp_model.load_state_dict = lambda: None temp_model.state_dict = lambda: model_state_dict # Save with updated metadata return self.save_model(temp_model, filepath, existing_metadata, password) except Exception as e: self.logger.error(f"Failed to update metadata: {str(e)}") return False def _encrypt_file(self, filepath: str, password: str): """Encrypt file in place.""" with open(filepath, 'rb') as f: data = f.read() encrypted_data = EncryptionManager.encrypt_data(data, password) with open(filepath, 'wb') as f: f.write(encrypted_data) def _decrypt_file(self, filepath: str, password: str) -> str: """Decrypt file to temporary location and return temp path.""" with open(filepath, 'rb') as f: encrypted_data = f.read() try: decrypted_data = EncryptionManager.decrypt_data(encrypted_data, password) except Exception: raise ModelFormatError("Invalid password or corrupted file") # Save to temporary file temp_fd, temp_path = tempfile.mkstemp(suffix='.pt') with os.fdopen(temp_fd, 'wb') as f: f.write(decrypted_data) return temp_path def _update_file_hash(self, filepath: str, file_hash: str, password: Optional[str] = None): """Update file hash in metadata.""" try: metadata = self.extract_metadata(filepath, password) metadata['file_hash'] = file_hash self.update_metadata(filepath, metadata, password) except Exception as e: self.logger.warning(f"Failed to update file hash: {str(e)}") class ONNXHandler(ModelFormatHandler): """Handler for .onnx files with metadata_props.""" def __init__(self, logger: Optional[logging.Logger] = None): super().__init__(logger) if not ONNX_AVAILABLE: raise ImportError("ONNX not available. Install with: pip install onnx onnxruntime") def save_model(self, model: Union[nn.Module, str], filepath: str, metadata: Dict[str, Any], password: Optional[str] = None, **kwargs) -> bool: """ Save model as .onnx file with metadata_props. Args: model: PyTorch model or path to existing ONNX model filepath: Path to save the model metadata: Metadata dictionary to include password: Optional password for encryption **kwargs: Additional arguments for ONNX export Returns: bool: True if successful """ try: # Ensure directory exists os.makedirs(os.path.dirname(filepath), exist_ok=True) if isinstance(model, str): # Load existing ONNX model onnx_model = onnx.load(model) else: # Export PyTorch model to ONNX dummy_input = kwargs.get('dummy_input') if dummy_input is None: raise ValueError("dummy_input required for PyTorch to ONNX conversion") # Create temporary ONNX file temp_fd, temp_onnx_path = tempfile.mkstemp(suffix='.onnx') os.close(temp_fd) torch.onnx.export( model, dummy_input, temp_onnx_path, export_params=True, opset_version=kwargs.get('opset_version', 11), do_constant_folding=True, input_names=kwargs.get('input_names', ['input']), output_names=kwargs.get('output_names', ['output']), dynamic_axes=kwargs.get('dynamic_axes', {}), verbose=False ) onnx_model = onnx.load(temp_onnx_path) os.unlink(temp_onnx_path) # Add metadata to model metadata['file_format'] = 'onnx' metadata['onnx_version'] = onnx.__version__ # Clear existing metadata_props del onnx_model.metadata_props[:] # Add new metadata as key-value pairs for key, value in metadata.items(): prop = onnx_model.metadata_props.add() prop.key = key if isinstance(value, (dict, list)): prop.value = json.dumps(value, ensure_ascii=False) else: prop.value = str(value) # Save model temp_filepath = filepath if password: temp_fd, temp_filepath = tempfile.mkstemp(suffix='.onnx') os.close(temp_fd) onnx.save(onnx_model, temp_filepath) # Apply password protection if requested if password: self._encrypt_file(temp_filepath, password) # Move encrypted file to final location with open(temp_filepath, 'rb') as src, open(filepath, 'wb') as dst: dst.write(src.read()) os.unlink(temp_filepath) # Update file hash file_hash = self._generate_file_hash(filepath) self._update_file_hash(filepath, file_hash, password) self.logger.info(f"Model saved successfully: {filepath}") return True except Exception as e: self.logger.error(f"Failed to save model: {str(e)}") return False def load_model(self, filepath: str, password: Optional[str] = None, map_location: Optional[str] = None) -> tuple: """ Load model from .onnx file. Returns: tuple: (onnx_model, metadata) """ try: # Handle password protection if password: filepath = self._decrypt_file(filepath, password) # Load ONNX model onnx_model = onnx.load(filepath) # Extract metadata from metadata_props metadata = {} for prop in onnx_model.metadata_props: try: # Try to parse as JSON first metadata[prop.key] = json.loads(prop.value) except json.JSONDecodeError: # If not JSON, store as string metadata[prop.key] = prop.value self.logger.info(f"Model loaded successfully: {filepath}") return onnx_model, metadata except Exception as e: self.logger.error(f"Failed to load model: {str(e)}") raise ModelFormatError(f"Failed to load model: {str(e)}") def extract_metadata(self, filepath: str, password: Optional[str] = None) -> Dict[str, Any]: """Extract metadata from .onnx file without loading the full model.""" try: # Handle password protection if password: filepath = self._decrypt_file(filepath, password) # Load model to extract metadata onnx_model = onnx.load(filepath) # Extract metadata metadata = {} for prop in onnx_model.metadata_props: try: metadata[prop.key] = json.loads(prop.value) except json.JSONDecodeError: metadata[prop.key] = prop.value return metadata except Exception as e: self.logger.error(f"Failed to extract metadata: {str(e)}") return {} def update_metadata(self, filepath: str, new_metadata: Dict[str, Any], password: Optional[str] = None) -> bool: """Update metadata in existing .onnx file.""" try: # Load existing model and metadata onnx_model, existing_metadata = self.load_model(filepath, password) # Merge metadata existing_metadata.update(new_metadata) # Save with updated metadata return self.save_model(onnx_model, filepath, existing_metadata, password) except Exception as e: self.logger.error(f"Failed to update metadata: {str(e)}") return False def _encrypt_file(self, filepath: str, password: str): """Encrypt file in place.""" with open(filepath, 'rb') as f: data = f.read() encrypted_data = EncryptionManager.encrypt_data(data, password) with open(filepath, 'wb') as f: f.write(encrypted_data) def _decrypt_file(self, filepath: str, password: str) -> str: """Decrypt file to temporary location and return temp path.""" with open(filepath, 'rb') as f: encrypted_data = f.read() try: decrypted_data = EncryptionManager.decrypt_data(encrypted_data, password) except Exception: raise ModelFormatError("Invalid password or corrupted file") # Save to temporary file temp_fd, temp_path = tempfile.mkstemp(suffix='.onnx') with os.fdopen(temp_fd, 'wb') as f: f.write(decrypted_data) return temp_path def _update_file_hash(self, filepath: str, file_hash: str, password: Optional[str] = None): """Update file hash in metadata.""" try: metadata = self.extract_metadata(filepath, password) metadata['file_hash'] = file_hash self.update_metadata(filepath, metadata, password) except Exception as e: self.logger.warning(f"Failed to update file hash: {str(e)}") class ModelFormatManager: """Manager class for handling different model formats.""" def __init__(self, logger: Optional[logging.Logger] = None): self.logger = logger or logging.getLogger(__name__) self.handlers = { '.pth': PTHHandler(logger), '.pt': PTHandler(logger), } # Add ONNX handler if available if ONNX_AVAILABLE: self.handlers['.onnx'] = ONNXHandler(logger) else: self.logger.warning("ONNX support not available") def get_handler(self, filepath: str) -> ModelFormatHandler: """Get appropriate handler for file format.""" ext = Path(filepath).suffix.lower() if ext not in self.handlers: raise ModelFormatError(f"Unsupported format: {ext}") return self.handlers[ext] def save_model(self, model: Union[nn.Module, str], filepath: str, metadata: Dict[str, Any], password: Optional[str] = None, **kwargs) -> bool: """Save model using appropriate handler.""" handler = self.get_handler(filepath) return handler.save_model(model, filepath, metadata, password, **kwargs) def load_model(self, filepath: str, password: Optional[str] = None, map_location: Optional[str] = None) -> tuple: """Load model using appropriate handler.""" handler = self.get_handler(filepath) return handler.load_model(filepath, password, map_location) def extract_metadata(self, filepath: str, password: Optional[str] = None) -> Dict[str, Any]: """Extract metadata using appropriate handler.""" handler = self.get_handler(filepath) return handler.extract_metadata(filepath, password) def update_metadata(self, filepath: str, new_metadata: Dict[str, Any], password: Optional[str] = None) -> bool: """Update metadata using appropriate handler.""" handler = self.get_handler(filepath) return handler.update_metadata(filepath, new_metadata, password) def get_supported_formats(self) -> List[str]: """Get list of supported file formats.""" return list(self.handlers.keys())