Enable Conditioning in DDPM/DDIM Diffusion code and Add Sample Representation

Author: Sane-violaCreated Feb 9, 2026Updated Feb 9, 2026

Hi,

Currently, the DDPM/DDIM models support conditioning on both the standard UNet and Karras UNet. I propose the following improvements for future implementation: 1. Conditioning Support: Extend the diffusion code to fully support conditioning for both UNet architectures. 2. Sample Representation: Add a function to visualize sample representations according to the class balance in the dataset, which can help better understand the model outputs and dataset distribution.

I precise that I works on DDPM 1D model, maybe that my propose is already requests on 2D/3D diffusion code.

This enhancement will improve flexibility and make it easier to work with conditioned diffusion models on imbalanced datasets.

thanks,

Sane

Source: lucidrains/denoising-diffusion-pytorch