Repository navigation
Support batches in noise scheduler #25
Description
Firstly, kudos to the HF team for building this awesome library!
I would like to propose some modifications to the DDPM noise scheduler (and potentially other schedulers) to support batching. Based on training_ddpm.py, it appears that the DDPM noise scheduler does not batch forward diffusion steps:
diffusers/examples/train_ddpm.py
Lines 78 to 81 in 694ad98
Supporting batches should be possible in theory by crudely modifying DDPMScheduler.get_alpha_prod to something like
def get_alpha_prod(self, time_steps):
mask = (time_steps < 0).float().numpy()
return torch.from_numpy(self.one * mask + self.alphas_cumprod[time_steps] * (1 - mask))The referenced line in the training example script would now be reduced to
# noises = torch.randn(clean_images.shape).to(clean_images.device)
# timesteps = torch.randint(0, noise_scheduler.timesteps, (16, 1, 1, 1)).long()
noise_images = noise_scheduler.forward_step(clean_images, noises, timesteps)Is there a particular reason why we should avoid batching the forward process? Curious to hear your thoughts!
Metadata
Metadata
Assignees
Labels
No labels
Activity
patrickvonplaten commentedon Jun 22, 2022
Very much agree! @anton-l is working on this right now I think :-)
anton-l commentedon Jun 22, 2022
Yep, much cleaner now :) 848c86c
Image to image (huggingface#25)
Eager mode with IREE (huggingface#25)
Merge pull request huggingface#25 from mattjamo/main
Merge pull request huggingface#25 from wdhongtw/dep-setup
Flux3 image (#15046)