File size: 2,685 Bytes
267d7ad | 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 | """Load MLX weights into UnlimitedOCR model.
Handles the complete weight loading with proper name mapping and validation.
"""
from typing import Dict, List, Tuple
import mlx.core as mx
import mlx.nn as nn
from .model import UnlimitedOCRModel, SAMVisionEncoder, CLIPVisionTransformer
from .config import UnlimitedOCRConfig
def load_weights_from_safetensors(model: nn.Module, weights_path: str) -> nn.Module:
"""Load MLX-compatible weights from safetensors file.
Args:
model: MLX model instance
weights_path: Path to safetensors file
Returns:
Model with loaded weights
"""
import safetensors.torch
import numpy as np
print(f"Loading weights from {weights_path}...")
st_weights = safetensors.torch.load_file(weights_path, device="cpu")
# Convert to MLX arrays
mlx_weights = {}
for name, tensor in st_weights.items():
mlx_weights[name] = mx.array(tensor.float().numpy())
# Load into model
model.load_weights(list(mlx_weights.items()))
mx.eval(model.parameters())
total = sum(v.size for v in mlx_weights.values())
print(f"Loaded {len(mlx_weights)} tensors, {total:,} parameters")
return model
def create_model_from_dir(model_dir: str) -> Tuple[UnlimitedOCRModel, UnlimitedOCRConfig]:
"""Create model instance from model directory.
Args:
model_dir: Directory containing config.json and model.safetensors
Returns:
Tuple of (model, config)
"""
import json
config_path = f"{model_dir}/config.json"
weights_path = f"{model_dir}/model.safetensors"
with open(config_path) as f:
config_dict = json.load(f)
config = UnlimitedOCRConfig.from_original_config(config_dict)
model = UnlimitedOCRModel(config)
model = load_weights_from_safetensors(model, weights_path)
return model, config
def verify_weights(model: UnlimitedOCRModel) -> Dict[str, any]:
"""Verify that all model weights are properly loaded.
Returns:
Dict with verification statistics
"""
stats = {"total_params": 0, "num_layers": {}, "issues": []}
params = dict(model.parameters())
for name, param in params.items():
size = param.numpy().size if hasattr(param, 'numpy') else 1
stats["total_params"] += size
# Check for NaN values
val = param
if hasattr(param, 'numpy'):
arr = param.numpy()
if hasattr(arr, 'isnan'):
nans = arr.isnan().sum()
if nans > 0:
stats["issues"].append(f"NaN values in {name}: {nans}")
stats["total_params_formatted"] = f"{stats['total_params']:,}"
return stats
|