Timestep-Aware Diffusion Model for Extreme Image Rescaling

TL;DR

This work tackles extreme image rescaling, where downsampling factors of $16\times$ or $32\times$ create highly ill-posed upscaling. It introduces TADM, a framework that performs rescaling in a pre-trained VAE latent space via a Decoupled Feature Rescaling Module and enhances details with a one-step diffusion update guided by a Time-Step Prediction Module, plus a hybrid time scheduler. The key contributions include latent-space bidirectional mapping for robust downscaling, adaptive diffusion utilization through time-step alignment, and a tiled inference strategy for ultra-high-resolution inputs, yielding superior quantitative and qualitative performance across DIV2K, Urban100, DIV8K, and CLIC2020 comparisons. This approach demonstrates strong semantic fidelity and texture realism while maintaining practical inference efficiency, with potential to inform image compression and large-scale deployment scenarios.

Abstract

Image rescaling aims to learn the optimal low-resolution (LR) image that can be accurately reconstructed to its original high-resolution (HR) counterpart, providing an efficient image processing and storage method for ultra-high definition media. However, extreme downscaling factors pose significant challenges to the upscaling process due to its highly ill-posed nature, causing existing image rescaling methods to struggle in generating semantically correct structures and perceptual friendly textures. In this work, we propose a novel framework called Timestep-Aware Diffusion Model (TADM) for extreme image rescaling, which performs rescaling operations in the latent space of a pre-trained autoencoder and effectively leverages powerful natural image priors learned by a pre-trained text-to-image diffusion model. Specifically, TADM adopts a pseudo-invertible module to establish the bidirectional mapping between the latent features of the HR image and the target-sized LR image. Then, the rescaled latent features are enhanced by a pre-trained diffusion model to generate more faithful details. Considering the spatially non-uniform degradation caused by the rescaling operation, we propose a novel time-step alignment strategy, which can adaptively allocate the generative capacity of the diffusion model based on the quality of the reconstructed latent features. Extensive experiments demonstrate the superiority of TADM over previous methods in both quantitative and qualitative evaluations.

Figures (18)

• Figure 1: Overview of the proposed Timestep-Aware Diffusion Model (TADM). First, the HR images $x$ are encoded to obtain the latent features $z$. The latent features are then rescaled to the target size using the decoupled feature rescaling module (DFRM), producing the LR images $y$ and outputting the rescaled latent features $\hat{z}$. Next, the rescaled latent features $\hat{z}$ are passed into a pre-trained diffusion model to perform a single denoising step for perceptual enhancement to obtain $\hat{z}_{0}$. In this process, a time prediction module (TPM) is used to estimate the time step $t$ based on $\hat{z}$ and the predicted $t$ is then fed into both the U-Net and the time scheduler. Finally, the perceptually enhanced latent features $\hat{z}_{0}$ are decoded to obtain the rescaled image $\hat{x}$.
• Figure 2: Detailed architecture of Decoupled Feature Rescaling Module (DFRM).
• Figure 3: We compare the MSE introduced by the rescaling operation with that caused by the forward diffusion process. The rescaling degradation severity is correlated with both the image content and the rescaling scale. Therefore, the model requires adaptive time-step to align the rescaling MSE with the diffusion MSE.
• Figure 4: The structure of the time scheduler module. It combines both the fixed scheduler and learnable scheduler to perform the denoising process. The two schedulers are connected by a zero-convolution layer to stabilize the training.
• Figure 5: Qualitative comparisons of 16$\times$ (the 1st and 2nd rows) and 32$\times$ (the 3rd and 4th rows) rescaling results. Our method achieves higher semantic accuracy, such as more clearer text, more recognizable faces, and more realistic structures.