PixRestore: Unified Image Restoration via Pixel Diffusion Transformer

Faithful, compact, and fast unified restoration diffusion model in pixel space.

Lingchen Sun★1,2Rongyuan Wu★1,2Xiangtao Kong1,2Jixin Zhao1Qiaosi Yi1,2Yujing Sun1,2Shuaizheng Liu1,2Zhengqiang Zhang1,2Lei Zhang†1,2
1 The Hong Kong Polytechnic University    2 OPPO Research Institute
★ Equal contribution.   † Corresponding author.

Abstract

Unified image restoration (UIR) aims to recover high-quality (HQ) content from low-quality (LQ) images with different degradations using a single model. Most recent methods adapt large pretrained text-to-image (T2I) latent diffusion models for their strong capacity and generative priors. However, the variational autoencoder (VAE) in latent T2I models may discard restoration-sensitive details, while the open-ended synthesis prior can introduce content-inconsistent artifacts. We present PixRestore, a VAE-free pixel-space Diffusion Transformer (DiT) for UIR, where the diffusion backbone is trained entirely from scratch, without relying on T2I pretraining. PixRestore performs flow matching directly on patchified pixels, preserving fine-grained details while keeping the token sequence tractable. To adapt to different degradations, PixRestore learns to predict the reliability of layer features using LQ–HQ DINO feature similarity. Features from more reliable layers are fused as dense conditioning, while less reliable layers receive stronger HQ-feature supervision to encourage degradation removal. We train PixRestore on a large-scale corpus of diverse scenes and degradations, and further finetune it into a one-step generator using DINO-based adversarial objectives for efficient inference. Experiments on public benchmarks and real-world test sets show that, with only about 50M parameters and single-step inference, PixRestore achieves the best overall fidelity, perceptual quality, and robustness to degradations among competing UIR models while being far more efficient. Larger PixRestore variants can further boost performance, demonstrating the scalability of our pixel-space design.

PixRestore quantitative comparison across restoration tasks, model efficiency, and inference latency
Overall restoration performance and efficiency comparison
53.7MParameters
44msInference latency
1Function evaluation
8Restoration tasks

Restoration Gallery

Method

PixRestore framework with a VAE-free Pixel DiT, frozen vision encoder, adaptive layer router, and fused hierarchical features
Overview of the PixRestore framework

PixRestore formulates unified image restoration as conditional flow matching directly in pixel space. The low-quality image and noisy state are patchified and processed by a VAE-free Pixel DiT, preserving pixel-aligned evidence without a latent bottleneck. A frozen DINOv2 encoder extracts complementary multi-layer features, while an adaptive layer router predicts per-image reliability weights to fuse the most trustworthy features for cross-attention. Layers that differ more strongly from the high-quality target receive stronger hierarchical feature supervision, enabling the same model to adapt to diverse degradations.

Training Loss

The multi-step model combines the pixel-space flow objective with two DINO-guided auxiliary losses:

L = Lflow + λwpredLwpred + λfeatLfeat

Flow matching

Lflow matches the predicted and target velocities along the linear path from the high-quality image to Gaussian noise.

Layer-weight prediction

Lwpred uses cross-entropy to train the router against the reliability distribution derived from paired LQ-HQ DINO feature similarity.

Feature supervision

Lfeat applies cosine feature loss more strongly to DINO layers whose LQ features deviate most from the HQ reference.

Single-step finetuning. A DINO-based adversarial loss Ladv is added to retain fine textures when the flow time is fixed to t = 1. The auxiliary weights λwpred, λfeat, and λadv are all set to 0.5.

Degradation Removal Score

Full-reference metrics require a ground-truth image, while common no-reference quality metrics may reward rich textures even when they come from artifacts or residual degradation. DR-Score is introduced as an auxiliary diagnostic metric that directly asks a vision-language model whether the target degradation has been removed. Given the LQ input, restored result, and task description, the evaluator assigns a score from 0 to 100 and provides a short reason; each image is evaluated five times and the mean score is reported for more stable method-level comparison.

Example showing why degradation removal score is needed when conventional no-reference metrics favor residual degradation
DR-Score evaluation protocol using the low-quality image, restored output, task description, and a vision-language model

Public Benchmark Comparison

The table below compares representative unified restoration methods on eight public tasks using fidelity, perceptual, and degradation-removal metrics. Methods marked with * are retrained on the same training corpus as PixRestore; the best and second-best results are highlighted in red and blue.

Quantitative comparison on public benchmarks across eight restoration tasks
Quantitative comparison on public benchmarks
Qualitative comparisons for desnowing, low-light enhancement, derainstreak, and denoising
Qualitative comparisons for dehazing, deraindrop, deblurring, and super-resolution

PixRestore variants rank first or second on most metrics across the eight degradation types. The visual comparisons show more complete degradation removal and sharper, content-consistent structures.

Generalization to Real-world Testset

Real-world test images are evaluated without paired ground truth and contain out-of-domain degradation patterns. Across desnowing, dehazing, deblurring, raindrop removal, low-light enhancement, and rain-streak removal, PixRestore produces more balanced restorations with cleaner structures, more natural colors, and fewer residual artifacts than competing unified restoration methods.

Real-world qualitative comparisons for desnowing, dehazing, and deblurring
Real-world qualitative comparisons for raindrop removal, low-light enhancement, and rain-streak removal

Model Complexity Comparisons

Under the same 512 × 512 input and hardware setting, PixRestore uses one function evaluation, 53.70M parameters, 658G FLOPs, and only 44 ms per image. It is approximately 7–23× faster than PromptIR, DiffUIR, FoundIR, and FAPE-IR, while PixRestore-B remains compact at 210.89M parameters and 79 ms latency.

Model complexity comparison showing NFE, parameter count, FLOPs, and latency

Scalability

PixRestore scales smoothly along two axes: decreasing the pixel patch size preserves more local evidence, while increasing the DiT backbone size strengthens global modeling. Across the evaluated configurations, LPIPS consistently decreases as computation increases, with a correlation of −0.96, indicating predictable quality gains from additional model capacity and finer pixel tokenization.

PixRestore scalability plot of GFLOPs against LPIPS for different model and patch sizes

Citation

@article{sun2026pixrestore,
  title={PixRestore: Unified Image Restoration via Pixel Diffusion Transformer},
  author={Sun, Lingchen and Wu, Rongyuan and Kong, Xiangtao and Zhao, Jixin and Yi, Qiaosi and Sun, Yujing and Liu, Shuaizheng and Zhang, Zhengqiang and Zhang, Lei},
  journal={arXiv preprint arXiv: https://arxiv.org/pdf/2608.16793},
  year={2026}
}