[Community Discussion] Scheduler design

[patrickvonplaten]

This is a very nice PR by @pcuenca showing what changes need to be done to the PNDM/PMLS scheduler to make it work with JAX/XLA - it's actually more then anticipated and shows that the scheduler now substantially differs from the original implementation that we use for PyTorch.

In the coming days we will integrate these changes into main diffusers to make the library compatible with Flax/JAX. Now the big question is should we:
a) Make each scheduler very generic and continue the set_format("pt") logic? While this would make sense logically as the schedulers don't store any trainable weights really - this could potentially lead to quite some if - else statements and too much abstracted code, e.g. lots of self.where(...) functions in scheduler_utils.py. Also maybe we want schedulers to have trainable weights in the future? Also do we anticipate schedulers to be more or less complex in the future?
b) Make one scheduler file for each framework. Instead of trying to fit all frameworks into one scheduler file, we make one scheduler for one framework. The advantage is clearly readability. Also most people probably always only work in one framework so for them it might be nicer to have schedulers seperate. However: Some schedulers will probably be 1-to-1 the same (which also might not be a problem necessarily)

I'm starting to be lean more and more towards b) actually here.

Would love to discuss - cc @anton-l @patil-suraj @natolambert @pcuenca

[patil-suraj]

I agree with your points, creating a generic scheduler here, will complicate a the overall design. In JAX we need to take special care of state, keeping everything scheduler as jnp arrays, avoid device to host communication etc. So I'm also in favor of option b. B is good for both readability and maintainability.

[natolambert]

Having done a bunch of work (maybe the most out of us) on trying to numpy-ify schedulers and the related tests, I think option b) will also result in much more accessible code. If we lay out the directory correct, it'll be very obvious what the different files do, and fitting in with the style designs of the library: one-file one purpose.

When trying to make them work with numpy it generally is a lot of self.function and if-else statements. Then you need to add tests for each option anyways because the random number generation is handled differently in each.

I'll go look at the PR now.

[vishnu-anirudh]

Hello
I hope it is okay for me to pitch in my two cents.
I prefer option b as well. It is very difficult to achieve a generic scheduler now. Given that this is an evolving field, it is easier to dedicated for every framework.
However, from interface perspective, it would be better if each of the scheduler implements an Abstract Base Class. This gives a good template on how the scheduler can be accessed/used when it is imported in other code. Also, as the diffuser library evolves, it becomes easier to develop functions/classes that can take in multiple schedulers (the type-hint would be the abstract base class).
In this way, we won't have complicated one-scheduler-fits-all situation but at the same time we have a basic interface which will be consistent for all the implemented schedulers, this will improve the maintainability and the readability of the code base.

[patrickvonplaten]

Cool let's go for separate schedulers then!

[vishnu-anirudh]

Just out of curiosity, can I ask how you design the software architectures for the libraries? Do you do that internally or are there sprints where the community gets involved?

[patrickvonplaten]

What exactly do you mean by software architectures? Like the whole library here? Usually we start a library internally but try to include the community in design decisions as quickly as possible :-)

[vishnu-anirudh]

Yes meant the whole library. Okay makes sense. Thank you.

[patrickvonplaten]

Cool design taken -> we'll have framework-dependent schedulers :-)

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.