
We intentionally minimized changes to the original model code. Further profiling and optimization can probably bring more improvements
At the moment compilation is applied only to the U-Net model inside the sampler. Since there is a lot happening outside of U-Net (e.g. operations directly in the sampling loop), it would be beneficial to compile the whole sampler. However, this would require analysis of the compilation process to avoid recompilation at every sampling step
Current code only applies compilation within the PLMS sampler, but it should be trivial to extend it to other samplers
Besides text-to-image generation, diffusion models are also applied to other tasks - image-to-image and inpainting. It would be interesting to measure how their performance improves from PyTorch 2 optimizations 

See if you can increase performance of open source diffusion models using the methods we described, and share the results! 
Resources