[Flax] Stateless schedulers, fixes and refactors#1661

patrickvonplaten merged 3 commits intohuggingface:mainhuggingface/diffusers:mainfrom

skirsten:flax/stateless-schedulers-and-improvementsskirsten/diffusers:flax/stateless-schedulers-and-improvementsCopy head branch name to clipboard