CoDi: Conditional Diffusion Distillation for Higher-Fidelity and Faster Image Generation

TL;DR

CoDi presents a single-stage method to distill an unconditional diffusion model into a conditional one by injecting a lightweight conditional encoder and enforcing conditional diffusion consistency along the probability-flow trajectory, enabling high-quality generation in as few as one to four steps. The approach preserves the diffusion prior, avoids reliance on the original training data, and supports parameter-efficient conditioning via adapters such as ControlNet. Empirical results across real-world tasks including super-resolution, inpainting, depth-to-image, and instructed editing show CoDi outperforms prior distillation methods in both quality and speed, with notable gains from the conditional consistency loss and the PF-ODE sampling framework. The work also introduces PE-CoDi, a parameter-efficient distillation variant that shares backbone parameters across tasks, significantly improving practicality for accelerating large-scale conditional diffusion models.

Abstract

Large generative diffusion models have revolutionized text-to-image generation and offer immense potential for conditional generation tasks such as image enhancement, restoration, editing, and compositing. However, their widespread adoption is hindered by the high computational cost, which limits their real-time application. To address this challenge, we introduce a novel method dubbed CoDi, that adapts a pre-trained latent diffusion model to accept additional image conditioning inputs while significantly reducing the sampling steps required to achieve high-quality results. Our method can leverage architectures such as ControlNet to incorporate conditioning inputs without compromising the model's prior knowledge gained during large scale pre-training. Additionally, a conditional consistency loss enforces consistent predictions across diffusion steps, effectively compelling the model to generate high-quality images with conditions in a few steps. Our conditional-task learning and distillation approach outperforms previous distillation methods, achieving a new state-of-the-art in producing high-quality images with very few steps (e.g., 1-4) across multiple tasks, including super-resolution, text-guided image editing, and depth-to-image generation.

Figures (11)

• Figure 1: Our proposed CoDi efficiently distills a conditional diffusion model from an unconditional one, enabling rapid generation of high-quality images under various conditional settings. We demonstrate CoDi's capabilities through generated results across various tasks.
• Figure 2: Sampled results between distilled models learned with alternative conditional guidance. Left curves shows the quantitative performance between the LPIPS and FID in $\{1,2,4,8\}$ steps. Right part show the visual results where each result comes from the 1 sampling step (top) or 4 sampling steps (bottom). The distance function from the left to right is $\| \mathbf{x} - \mathbb{E}(\mathbb{D}(\hat{\mathbf{x}}_\theta(\mathbf{z}_t, c))) \|^2_2$, $\| \mathbb{D}(\mathbf{x}) - \mathbb{D}(\hat{\mathbf{x}}_\theta(\mathbf{z}_t, c)) \|^2_2$, $F_\mathrm{lpips}(\mathbb{D}(\mathbf{x}), \mathbb{D}(\hat{\mathbf{x}}_\theta(\mathbf{z}_t, c))$, and our default $\| \mathbf{x} - \hat{\mathbf{x}}_\theta(\mathbf{z}_t) \|^2_2$, respectively.
• Figure 3: Network architecture illustration of our parameter-efficient conditional distillation framework.
• Figure 4: We show the results sampled in 4 steps by different models. Samples generated according to the low-resolution images (left) and masks (right) respectively. Please see our supplement for many more examples such as visual comparisons with the other methods.
• Figure 5: Samples generated according to the depth image (left) from ControlNet sampled in 4 steps (middle), and ours from the unconditional pretraining sampled in 4 steps (right). Please see our supplement for many more examples.

Theorems & Definitions (3)

• Remark 1
• Remark
• proof