# 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)
```bash
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](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
bash scripts/stage1_1.sh
```

**Step 2: Texture-aware Feature Refinement**
```bash
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.*

```bash
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.

```bash
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.

```bash
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.

```bash
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>
```

---