zjyao-PKU/Pref-Restore-PhaseA-Fidelity

Pref-Restore — PhaseA (Fidelity-leaning) checkpoint

Stage-A SFT checkpoint for Pref-Restore — released so the community can skip the most expensive stage and jump straight to PhaseB preference-RL training.

📄 Paper: Bridging Information Asymmetry: A Hierarchical Framework for Blind Face Restoration with Reduced Uncertainty (TPAMI 2026) 🧑‍💻 Code: https://github.com/zjYao36/Pref_Restore

Why this checkpoint exists

PhaseA (hierarchical SFT) is by far the most compute-heavy step of the full Pref-Restore pipeline. In our experiments, restoration quality keeps improving as PhaseA training continues, but with diminishing marginal returns — most of the easy gains land early; later iterations cost a lot of GPU-hours for a small numerical bump.

To make PhaseB training accessible without re-running our entire SFT, we publish this PhaseA snapshot. It is tuned to lean toward restoration fidelity and image realism, at the cost of slightly weaker aesthetic quality — exactly the trade-off you want as a base model that PhaseB's preference-RL will then push toward perceptual preference.

How to use it (as the base for PhaseB RL training)

# 1. Download with huggingface_hub
from huggingface_hub import snapshot_download
ckpt_dir = snapshot_download(repo_id="zjyao-PKU/Pref-Restore-PhaseA-Fidelity")

Then in DiffusionNFT/config/pref_restore_gt.py set:

config.pretrained.model = ckpt_dir   # the path returned above

and launch PhaseB as documented in the repo README.

You can also use it directly for inference with the inference_batch_noPrompt_fixLQ_vae.py script in the code repo — pass this folder as --model_path.

Architecture

blip3oQwenForCausalLMVAE — a BLIP-3o-NEXT-style multimodal AR backbone (~3B) coupled with a SANA 1.5 diffusion decoder and a VAE-encoded HQ branch (see config.json).

Component	Source
Backbone	BLIP3o-NEXT-SFT-3B (modified)
Vision tower	TA-Tok
Diffusion decoder	SANA 1.5 1.6B 1024px
Training	hierarchical SFT on FFHQ + CelebA-HQ with on-the-fly degradations

The checkpoint contains only the inference weights (model-0000{1,2,3}-of-00003.safetensors + config + tokenizer); the DeepSpeed optimizer state has been stripped.

License & citation

Apache-2.0. If you use this checkpoint, please cite the paper:

@article{yao2026prefrestore,
  title   = {Bridging Information Asymmetry: A Hierarchical Framework for Deterministic Blind Face Restoration},
  author  = {Yao, Zhengjian and Hu, Jiakui and Li, Kaiwen and He, Hangzhou and
             Zhang, Xinliang and Zeng, Shuang and Zhu, Lei and Lu, Yanye},
  journal = {IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI)},
  year    = {2026}
}