[Python] How do I implement a different type of scheduler after training my DDPM on a linear...

I trained a Diffusion model (DDPM) using a linear scheduler defined by the following code:

scheduler = DDPMScheduler(
num_train_timesteps=1000,
beta_start=0.0001,
beta_end=0.02,
beta_schedule="linear"
)


However, after training this model and loading it using the code:

def load_checkpoint(model, optimizer, checkpoint_path):
"""Load the model and optimizer state dictionaries from a checkpoint."""
checkpoint = torch.load(checkpoint_path, map_location=device)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch'] + 1 # Resume from the next epoch.
print(f"Resuming from epoch {start_epoch}")
return start_epoch


When I try to switch over to a cosine scheduler for inference using the following code:

scheduler = DDPMScheduler(
num_train_timesteps=1000,
beta_schedule="cosine"
)


I get the following error:

NotImplementedError: cosine is not implemented for <class 'diffusers.schedulers.scheduling_ddpm.DDPMScheduler'>


I tried changing my DDPM scheduler to a DDIM scheduler, changing the code to:

scheduler = DDIMScheduler(
num_train_timesteps=1000,
beta_schedule="cosine"
)


But I still get a very similar error:

NotImplementedError: cosine is not implemented for <class 'diffusers.schedulers.scheduling_ddim.DDIMScheduler'>


I don't understand what I am doing wrong, as I have read from multiple sources that it is normal to train a diffusion model on one scheduler (in my case, linear), and change to another scheduler for inference (in my case, cosine). How do I solve this problem?

Continue reading...