e4cfa9d1e9
Fork of dgrauet/Hunyuan3D-2.1-mlx + our generate_e2e.py CLI, env-tunable remesh (HY3D_REMESH_FACES), and HARDWARE.md. Upstream training data (mini_trainset) and demo images stripped — inference needs none of it. Full upstream history: github.com/dgrauet/Hunyuan3D-2.1-mlx
79 lines
2.6 KiB
Python
79 lines
2.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Convert RealESRGAN (RRDBNet) PyTorch weights to MLX format.
|
|
|
|
Usage:
|
|
python convert_realesrgan.py --input path/to/RealESRGAN_x4plus.pth --output path/to/realesrgan_mlx.safetensors
|
|
|
|
The script:
|
|
1. Loads the PyTorch .pth checkpoint.
|
|
2. Remaps the keys from PyTorch's Sequential body (body.0, body.1, ...) to MLX list
|
|
indexing (body.0, body.1, ...) -- these happen to match.
|
|
3. Transposes Conv2d weights from PyTorch layout (O, I, H, W) to MLX layout (O, H, W, I).
|
|
4. Saves the result as .safetensors or .npz.
|
|
"""
|
|
|
|
import argparse
|
|
import re
|
|
from pathlib import Path
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
|
|
|
|
def convert_weights(input_path: str, output_path: str) -> None:
|
|
"""Convert PyTorch RRDBNet weights to MLX format.
|
|
|
|
Args:
|
|
input_path: Path to PyTorch .pth file.
|
|
output_path: Path to output .safetensors or .npz file.
|
|
"""
|
|
import torch
|
|
|
|
# Load PyTorch checkpoint
|
|
state_dict = torch.load(input_path, map_location="cpu", weights_only=True)
|
|
|
|
# Some checkpoints wrap state_dict under a key
|
|
if "params_ema" in state_dict:
|
|
state_dict = state_dict["params_ema"]
|
|
elif "params" in state_dict:
|
|
state_dict = state_dict["params"]
|
|
|
|
mlx_weights = {}
|
|
for key, value in state_dict.items():
|
|
arr = value.numpy()
|
|
|
|
# Conv2d weights: PyTorch (O, I, H, W) -> MLX (O, H, W, I)
|
|
if "weight" in key and arr.ndim == 4:
|
|
arr = np.transpose(arr, (0, 2, 3, 1))
|
|
|
|
mlx_weights[key] = mx.array(arr)
|
|
|
|
# Save
|
|
if output_path.endswith(".npz"):
|
|
mx.savez(output_path, **mlx_weights)
|
|
elif output_path.endswith(".safetensors"):
|
|
mx.save_safetensors(output_path, mlx_weights)
|
|
else:
|
|
raise ValueError(f"Unsupported output format: {output_path}. Use .safetensors or .npz")
|
|
|
|
print(f"Converted {len(mlx_weights)} tensors from {input_path} -> {output_path}")
|
|
|
|
# Print some stats
|
|
conv_count = sum(1 for k in mlx_weights if "weight" in k and mlx_weights[k].ndim == 4)
|
|
print(f" Conv2d weight tensors transposed: {conv_count}")
|
|
total_params = sum(v.size for v in mlx_weights.values())
|
|
print(f" Total parameters: {total_params:,}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Convert RealESRGAN PyTorch weights to MLX format")
|
|
parser.add_argument("--input", "-i", required=True, help="Path to PyTorch .pth checkpoint")
|
|
parser.add_argument("--output", "-o", required=True, help="Output path (.safetensors or .npz)")
|
|
args = parser.parse_args()
|
|
|
|
convert_weights(args.input, args.output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|