Convert main model weights with key remapping.
(model_path: Path, output_path: Path)
| 84 | |
| 85 | |
| 86 | def convert_model(model_path: Path, output_path: Path): |
| 87 | """Convert main model weights with key remapping.""" |
| 88 | print(f"Loading model from {model_path}") |
| 89 | state_dict = torch.load(model_path, map_location='cpu', weights_only=False) |
| 90 | |
| 91 | if 'model' in state_dict: |
| 92 | state_dict = state_dict['model'] |
| 93 | |
| 94 | remapped = {} |
| 95 | remap_count = 0 |
| 96 | for key, tensor in sorted(state_dict.items()): |
| 97 | new_key = flatten_fm_key(key) |
| 98 | |
| 99 | # Scalar tensors need at least 1 dim for safetensors |
| 100 | if tensor.dim() == 0: |
| 101 | tensor = tensor.unsqueeze(0) |
| 102 | |
| 103 | # Convert to float16 |
| 104 | if tensor.dtype == torch.float32: |
| 105 | tensor = tensor.half() |
| 106 | |
| 107 | remapped[new_key] = tensor |
| 108 | if new_key != key: |
| 109 | remap_count += 1 |
| 110 | print(f" {key} -> {new_key} {list(tensor.shape)}") |
| 111 | |
| 112 | print(f"\nRemapped {remap_count}/{len(remapped)} keys") |
| 113 | print(f"Saving {len(remapped)} tensors to {output_path}") |
| 114 | save_file(remapped, str(output_path)) |
| 115 | print(f" Size: {output_path.stat().st_size / 1e6:.1f} MB") |
| 116 | |
| 117 | |
| 118 | def convert_vocoder(vocoder_path: Path, output_path: Path): |
no test coverage detected