nanguoyu/restoreformerpp-mlx

RestoreFormer++ — Apple MLX safetensors

MLX-compatible safetensors for RestoreFormer++ (TPAMI 2023), converted from the upstream PyTorch checkpoint for use with the mlx-restoreformerpp-swift Swift port and ultimately the OtterPix macOS app.

What's in this repo

File	Size	Contents
model.safetensors	~280 MB	RestoreFormer++ (head_size=4, ex_multi_scale_num=1) weights

The weights are the inference-only path — encoder + VectorQuantizer + decoder + the two 1×1 quant_conv / post_quant_conv projections. Training-only components (discriminator, optimizer state) and the orthogonal RestoreFormer (CVPR 2022; head_size=8) variant are not included.

What's different from upstream

Two transformations applied vs RestoreFormer++.ckpt:

1. vqvae. state-dict prefix stripped. The training-time wrapper saved the model under vqvae.<name>; the inference module tree omits that prefix.
2. Conv2d weights permuted from PyTorch OIHW → MLX OHWI. PyTorch stores conv kernels as (out, in, kH, kW); MLX expects (out, kH, kW, in). Detection rule used during conversion: any .weight tensor with 4D shape gets the (0, 2, 3, 1) axis transpose.

No other tensor manipulation. GroupNorm, Embedding, and Linear weights pass through unchanged.

Total: 441 tensors, of which 130 Conv2d weights were permuted and 311 are pass-through.

Loading from Swift

import MLXRestoreFormerPP

let model = RestoreFormerPP(variant: .plusPlus)
try model.loadWeights(from: weightsURL)
let restored = model(degradedFaceCrop)   // [1, 512, 512, 3] in [-1, 1]

See mlx-restoreformerpp-swift for the full Swift port + parity tests.

License

Apache 2.0 — matches the upstream RestoreFormerPlusPlus repo. No NVIDIA / StyleGAN2 dependencies; no S-Lab non-commercial clauses. Commercial use permitted.

Citation

If you use this in research or a product, please cite the upstream paper:

@article{wang2023restoreformer++,
  title={RestoreFormer++: Towards Real-World Blind Face Restoration from Undegraded Key-Value Pairs},
  author={Wang, Zhouxia and Zhang, Jiawei and Chen, Tianshui and Wang, Wenping and Luo, Ping},
  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
  year={2023}
}