Load model weights (using fp32 precision) Args: model_path: Model file path or directory path model_type: Model type ("safetensors" or "pytorch") Returns: Model weights dictionary (fp32 precision)
(model_path: str, model_type: str)
| 66 | |
| 67 | |
| 68 | def load_model_weights(model_path: str, model_type: str) -> Dict[str, torch.Tensor]: |
| 69 | """ |
| 70 | Load model weights (using fp32 precision) |
| 71 | |
| 72 | Args: |
| 73 | model_path: Model file path or directory path |
| 74 | model_type: Model type ("safetensors" or "pytorch") |
| 75 | |
| 76 | Returns: |
| 77 | Model weights dictionary (fp32 precision) |
| 78 | """ |
| 79 | print(f"Loading model: {model_path} (type: {model_type}, precision: fp32)") |
| 80 | |
| 81 | if not os.path.exists(model_path): |
| 82 | raise FileNotFoundError(f"Model path does not exist: {model_path}") |
| 83 | |
| 84 | weights = {} |
| 85 | |
| 86 | if model_type == "safetensors": |
| 87 | if os.path.isdir(model_path): |
| 88 | # If it's a directory, load all .safetensors files in the directory |
| 89 | safetensors_files = [] |
| 90 | for file in os.listdir(model_path): |
| 91 | if file.endswith(".safetensors"): |
| 92 | safetensors_files.append(os.path.join(model_path, file)) |
| 93 | |
| 94 | if not safetensors_files: |
| 95 | raise ValueError(f"No .safetensors files found in directory: {model_path}") |
| 96 | |
| 97 | print(f"Found {len(safetensors_files)} safetensors files") |
| 98 | |
| 99 | # Load all files and merge weights |
| 100 | for file_path in sorted(safetensors_files): |
| 101 | print(f" Loading file: {os.path.basename(file_path)}") |
| 102 | with safe_open(file_path, framework="pt", device="cpu") as f: |
| 103 | for key in f.keys(): |
| 104 | if key in weights: |
| 105 | print(f"Warning: weight key '{key}' is duplicated in multiple files, will be overwritten") |
| 106 | weights[key] = f.get_tensor(key) |
| 107 | |
| 108 | elif os.path.isfile(model_path): |
| 109 | # If it's a single file |
| 110 | if model_path.endswith(".safetensors"): |
| 111 | with safe_open(model_path, framework="pt", device="cpu") as f: |
| 112 | for key in f.keys(): |
| 113 | weights[key] = f.get_tensor(key) |
| 114 | else: |
| 115 | raise ValueError(f"safetensors type file should end with .safetensors: {model_path}") |
| 116 | else: |
| 117 | raise ValueError(f"Invalid path type: {model_path}") |
| 118 | |
| 119 | elif model_type == "pytorch": |
| 120 | # Load pytorch format (.pt, .pth) |
| 121 | if model_path.endswith((".pt", ".pth")): |
| 122 | checkpoint = torch.load(model_path, map_location="cpu") |
| 123 | |
| 124 | # Handle possible nested structure |
| 125 | if isinstance(checkpoint, dict): |