Pref-Restoration

Official implementation of the paper "Bridging Information Asymmetry: A Hierarchical Framework for Deterministic Blind Face Restoration". Pref-Restore is a hierarchical framework that synergizes Auto-Regressive (AR) modeling with Diffusion processes, further aligned with human preference via On-Policy Reinforcement Learning.

🛠️ Installation

To facilitate the two-stage training paradigm, we recommend creating separate environments for Stage 1 (AR+Diffusion Pre-training) and Stage 2 (RL Fine-tuning).

1. Environment for Stage 1 (Base Model)

conda create -n pref-restore python=3.11 -y
conda activate pref-restore
pip install --upgrade pip setuptools
pip install -r requirements.txt
pip install -e .

2. Environment for Stage 2 (RL Fine-tuning)

Please refer to the official installation guide of DiffusionNFT for the RL environment setup: 👉 https://github.com/NVlabs/DiffusionNFT

🚀 Training

Stage 1: Knowledge-Guided Pre-training

This stage aligns the AR semantic integrator with the diffusion backbone.

Step 1: Semantic-to-Diffusion Alignment

bash scripts/stage1_1.sh

Step 2: Texture-aware Feature Refinement

bash scripts/stage1_2.sh

Stage 2: Preference-Aware Fine-tuning (On-Policy RL)

In this stage, we employ On-Policy RL to align the restoration model with human preferences. > Note: Please ensure all model checkpoints and data paths are correctly configured in run.sh before running.

cd DiffusionNFT
bash run.sh

⚡ Inference

We provide scripts for three different inference modes: RL-aligned Model (LoRA), Text-Guided, and Unconditional (w.o. text).

1. Inference with RL-aligned Model (LoRA)

Use the LoRA checkpoint trained in Stage 2 for preference-aligned restoration.

CUDA_VISIBLE_DEVICES=0 \
python inference_batch_noPrompt_fixLQ_vae_lora.py \
    --model_path <PATH_TO_BASE_MODEL> \
    --json_path /Pref-Restore/FFHQ_val/CelebA_HQ/captions_lq.json \
    --output_dir <PATH_TO_OUTPUT_DIR> \
    --lora_path /DiffusionNFT/logs/nft/prefRestore/multi_reward/checkpoints/checkpoint-30 \
    --use_lora

2. Inference with Text Guidance

Use explicit text prompts (captions) to guide the restoration.

CUDA_VISIBLE_DEVICES=1 \
python inference_batch_Prompt_fixLQ_vae.py \
    --model_path <PATH_TO_BASE_MODEL> \
    --json_path /Pref-Restore/FFHQ_val/CelebA_HQ/long_captions_lq_part.json \
    --output_dir <PATH_TO_OUTPUT_DIR>

3. Inference without Text Guidance (Base)

Perform blind restoration without explicit text prompts.

CUDA_VISIBLE_DEVICES=1 \
python inference_batch_noPrompt_fixLQ_vae.py \
    --model_path <PATH_TO_BASE_MODEL> \
    --json_path /Pref-Restore/FFHQ_val/CelebA_HQ/long_captions_lq_part.json \
    --output_dir <PATH_TO_OUTPUT_DIR>