Congrats! Super great work!

I've noticed that you're currently using the original DDPM scheduler, which is rather slow. It would be much faster if we could apply DPM-Solver++ into this work to accelerate the sampling.

Note that the original DPM-Solver++ may have numerical issues when using the cosine beta schedule, and I've added a fix here: https://github.com/LuChengTHU/dpm-solver/blob/5c6ee9f1e6b60c8c54f955fbaab0a6717fc2b75b/dpm_solver_pytorch.py#L105

I'm happy to help to integrate DPM-Solver++ into IF when the model is released :)