Choosing the Appropriate add_noise() Method in diffusers.schedulers and Compatibility with Different Schedulers for Data Generation

[llmydota]

Hello everyone,

I’m training my own Diffusion model to generate data, and I’ve encountered some questions regarding the add_noise methods. In the diffusers library, I’ve noticed that the scheduler module provides three different methods for adding noise:

EulerDiscreteScheduler.add_noise

DDPMScheduler.add_noise

DPMSolverMultistepScheduler.add_noise

Each of these methods seems to handle noise addition in a different way, and I’m wondering which one I should use during training. Additionally, if I train the model with one scheduler, am I limited to using the corresponding scheduler for data generation? If I wish to generate data using a different scheduler, do I need to retrain the model with the corresponding add_noise method?

I would appreciate any insights or advice on this matter. Thank you!

[llmydota]

I found that all schedulers can be interchanged, but during training, it's best to use the add_noise method from the DDPMScheduler:

# Add noise to the original samples
noisy_sample = scheduler.add_noise(original_samples, noise, timesteps)

# Train the model to predict noise
pred_noise = denoise_model(noisy_sample, timesteps)

During generation, you can choose any scheduler. However, when using a scheduler that implements the EulerDiscreteScheduler's add_noise method, you need to be cautious:

# Initialize model input as a zero tensor, similar to target data
model_input = torch.zeros(datashape)

# Add the maximum noise to the model input
model_input = scheduler.add_noise(model_input, torch.randn_like(model_input), timesteps=scheduler.timesteps[0])

# Denoise the input step by step
for t in scheduler.timesteps:
    scaled_input = scheduler.scale_model_input(model_input, t)  # Scale input before predicting noise
    pred_noise = denoise_model(scaled_input, t)  # Predict noise
    model_input = scheduler.step(pred_noise, t, model_input).prev_sample # Denoise on the un-scaled input

In summary, you can mix and match schedulers during inference, but when using a scheduler like EulerDiscreteScheduler, it's important to scale the input at each step.