search

LaurentMazare / diffusers-rs

An implementation of the diffusers api in Rust
Apache License 2.0
537 stars 55 forks source link

Implement DDPM Scheduler #21

Closed mspronesti closed 1 year ago

mspronesti commented 1 year ago

Hello, this PR aims at integrating the DDPMScheduler into this beautiful rust implementation of the diffusers API.

I made it a draft because, despite the implementation seems (to me) to be sound, using the stable_diffusion pipeline as it is with this new scheduler, the final image is nothing but noise. I'm still working on it, but I wanted to show this "work-in-progress" result in case anybody wanted to point out anything / help .

LaurentMazare commented 1 year ago

Looks pretty interesting! What I usually did when getting noise back for the generated image was lining things up between the python and rust implementation, they are close enough so that it's easy to fix the seed and get almost identical numbers on both sides (and for the scheduler I guess there is even a lot less randomness involved in the process).

mspronesti commented 1 year ago

Hello @LaurentMazare, thanks for the feedback. This implementation is in fact sound, I had noise back because the number of inference steps was too small : I tried again with the official python implementation, specifically this notebook and the snippet below and nothing smaller than 1000 produced a "meaningful" image (I was "testing" the rust version with 30, 100, 200, 500). DDPMPipeline didn't even allow to set the number of inference steps in the latest release and __call__ was silently setting it to 1000. 

import torch
from diffusers import StableDiffusionPipeline, DDPMScheduler

pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)

pipe.scheduler = DDPMScheduler()

pipe = pipe.to("cuda")

prompt = "A very rusty robot holding a fire torch."
image = pipe(prompt, num_inference_steps=1000).images[0]

Please notice that the step ratio is given by num_train_timesteps / num_inference_steps, which means that, keeping num_train_timesteps=1000, any value of num_inference_steps greater than 500 will still make the scheduler work on 1000 timesteps, i.e. with a step ratio equal to 1.

With 1000 steps, I'm getting back a meaningful result with the code of this PR. I'm therefore marking it as "Ready for review" :)

LaurentMazare commented 1 year ago

Merged, thanks for the PR!

JohnAlcatraz commented 1 year ago

Can someone maybe explain what this DDPM schedulers is, what makes it cool, when it should be used? If it needs 1000 steps for a proper result, it sounds extremely slow?

mspronesti commented 1 year ago

@JohnAlcatraz Sorry, I must have missed your comment. The scheduler contained a bug when used with less inference steps than training timesteps (both the official python version from HF and the rust version ported from it). Should be fixed now.