Eliminating VAE for Fast and High-Resolution Generative Detail Restoration

TL;DR

The paper tackles the bottleneck of variational auto-encoders in one-step diffusion super-resolution by transferring the diffusion process from latent space to pixel space through eightfold pixel-unshuffle/shuffle, effectively removing the VAE. It introduces a two-stage adversarial distillation to replace the encoder and decoder, plus RandPad and a masked Fourier space loss to mitigate artifacts, and a PadCFG scheme for stable, guided inference. GenDR-Pix achieves about 2.8× speedup and over 60% memory savings over GenDR, while preserving fidelity enough to restore 4K images in roughly 1 s on modest hardware. The approach delivers strong runtime efficiency and competitive image quality against state-of-the-art one-step diffusion SR methods, demonstrating practical potential for real-world, high-resolution restoration.

Abstract

Diffusion models have attained remarkable breakthroughs in the real-world super-resolution (SR) task, albeit at slow inference and high demand on devices. To accelerate inference, recent works like GenDR adopt step distillation to minimize the step number to one. However, the memory boundary still restricts the maximum processing size, necessitating tile-by-tile restoration of high-resolution images. Through profiling the pipeline, we pinpoint that the variational auto-encoder (VAE) is the bottleneck of latency and memory. To completely solve the problem, we leverage pixel-(un)shuffle operations to eliminate the VAE, reversing the latent-based GenDR to pixel-space GenDR-Pix. However, upscale with x8 pixelshuffle may induce artifacts of repeated patterns. To alleviate the distortion, we propose a multi-stage adversarial distillation to progressively remove the encoder and decoder. Specifically, we utilize generative features from the previous stage models to guide adversarial discrimination. Moreover, we propose random padding to augment generative features and avoid discriminator collapse. We also introduce a masked Fourier space loss to penalize the outliers of amplitude. To improve inference performance, we empirically integrate a padding-based self-ensemble with classifier-free guidance to improve inference scaling. Experimental results show that GenDR-Pix performs 2.8x acceleration and 60% memory-saving compared to GenDR with negligible visual degradation, surpassing other one-step diffusion SR. Against all odds, GenDR-Pix can restore 4K image in only 1 second and 6GB.

Figures (18)

• Figure 1: Profiling for GenDR with input size of 1024$^2$px.
• Figure 2: We remove the entire VAE for GenDR for higher efficiency. The bars are model MUSIQ score on RealLQ250. The hatched bar means the modification has not been adopted. The dotted curve shows runtime with 2560$\times$1440 output.
• Figure 3: We visualize the effect of our modifications step by step. The left-top corner shows the operation types corresponding to \ref{['fig:performance-curve']}, and the left-down numbers are the operation order. The proposed strategies can effectively remove most artifacts.
• Figure 4: Illustration of training processes of GenDR-Pix. In Stage I, we replace the encoder with $\times$8 pixel-unshuffle and utilize half of the teacher (gray) to generate discriminative features, and blue blocks are trainable. In Stage II, we further remove the decoder with pixel-shuffle and utilize the Stage I model (blue), and the purple blocks are trainable.
• Figure 4: Ablation study on discriminator.