[RFC] Proposal to change API of `ConfigMixin` and `SchedulerMixin`.

[patrickvonplaten]

There are a couple of API problems currently related to the from_config method. This is a proposal to change the behavior of from_config and add a new load_config method instead.

1. Loading schedulers

Currently one has to load schedulers as follows:

from diffusers import DDIMScheduler

scheduler = DDIMScheduler.from_config("CompVis/stable-diffusion-v1-4", subfolder="scheduler")

We've received numerous feedback now (e.g. from @thomwolf, @keturn, @Narsil) that this is weird. Why do we pass a model_id to the from_config method? Both transformers and diffusers makes heavy usage of the from_pretrained(...) method, so it does indeed make much more sense to also better have the following API:

from diffusers import DDIMScheduler

scheduler = DDIMScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")

2. Switching between schedulers.

It's quite obvious that we need functionality to allow the user to easily switch between schedulers.
From the feedback of the set_scheduler PR here: #1247 as well as from feedback from @Narsil , we should probably re-visit the design of from_config a bit and instead of allowing to pass file names just make it accept a config dictionary - the one that is given to every object inheriting from ConfigMixin via the .config property. This would also directly solve the #1247 in a pretty clean way - we can just advertise users to do the following to switch between schedulers:

# 1. Switch to DDIM
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)

# 2. Switch to Euler
pipe.scheduler = EulerScheduler.from_config(pipe.scheduler.config)

...

Changing the API of from_config would be a pretty clean and also intuitive solution for the set_scheduler problem

3. Loading just the config of models and schedulers

Instead of instantiating the whole model or whole scheduler, users might want to just get the config of the components. This is currently a niche case, but might become more important going forward (e.g. if people want to train diffusion systems from scratch).

Right now people can to the following for models:

random_unet = Unet2DModel.from_config("CompVis/stable-diffusion-v1-4", subfolder="unet")

If we follow the API changes above long-term this would not be possible anymore if we make from_config(...) accept only a configuration dictionary. I would instead propose to add a new static method called load_config to ConfigMixin that would allow all classes inheriting from ConfigMixin to just load the config that can then be passed to from_config e.g. we could instead allow the following API:

config = UNet2DModel.load_config("CompVis/stable-diffusion-v1-4", subfolder="unet")

model = UNet2DModel.from_config(config)

This API has the following advantage - it clearly seperates the UNet object (with all it's random weights) from the configuration dictionary. Doing load_config would only load the config.json of the model and allow the user to quickly inspect a given config without having to init any random weights. The config can then be passed to the from_config(...) method to randomely init the model.

This RFC is heavily influenced by the very nice feedback from @keturn and also copies @keturn's suggestions here: #1247 (comment)

The more I think about, it, I'm very much in favor of changing the from_config(...) API and adding a new load_config(...) method to ConfigMixin as well as a from_pretrained(...) method to SchedulerMixin to nicely cover all three use cases.

What are your thoughts on this ? @patil-suraj @anton-l @pcuenca @Narsil @keturn ?

[patil-suraj]

Agree with changing the behaviour of from_config to address all these 3 use-cases!

QQ: Do we want to change the behaviour of from_config completely or deprecate the current behaviour and add support for this new logic ?

Thanks a lot for the really great feedback @keturn

[keturn]

Makes sense to me! Decoupling load_config from instantiation (and loading all the weights) sounds like a good move.

The phrase from_pretrained is an awkward fit for the schedulers—they have a config, but they are not trained, and do not have weights from training—but I don't mind overlooking that for the goal of giving all these ConfigMixin-children-of-Pipeline the same interface.

…

This just came to mind, perhaps food for thought for a future iteration: Schedulers are not sub-models themselves. Scheduler settings are properties of a sub-model. e.g. for Stable Diffusion, scheduler specifically applies to unet. If ensemble models like eDiff-I have different schedulers for different submodels, then maybe each submodel has its own scheduler settings? (This is speculation, as I haven't seen an eDiff-I model released yet.)

[patil-suraj]

This just came to mind, perhaps food for thought for a future iteration: Schedulers are not sub-models themselves. Scheduler settings are properties of a sub-model. e.g. for Stable Diffusion, scheduler specifically applies to unet. If ensemble models like eDiff-I have different schedulers for different submodels, then maybe each submodel has its own scheduler settings? (This is speculation, as I haven't seen an eDiff-I model released yet.)

This is totally possible, for example in imagen they use cosine schedule for the first two model (base 64x64 and SR256x256) and linear schedule for the SR1024 model. And in this case the pipeline can just have multiple schedulers, each with it's own setting.

[pcuenca]

The phrase from_pretrained is an awkward fit for the schedulers—they have a config, but they are not trained, and do not have weights from training—but I don't mind overlooking that for the goal of giving all these ConfigMixin-children-of-Pipeline the same interface.

I agree it's useful for consistency. I've convinced myself that from_pretrained actually means "from this particular trained model repo".

---

I like the general approach a lot, it's more consistent and allows us to access configurations independently of weight loading. A couple of minor points.

config = UNet2DModel.load_config("CompVis/stable-diffusion-v1-4", subfolder="unet")

Could we make it so the subfolder is optional? In this case we know we are loading a UNet configuration, so I'd rather skip the subfolder parameter. It's a bit silly, but then we don't require users to know the underlying tree structure of the repo.

On another note,

config = UNet2DModel.load_config("CompVis/stable-diffusion-v1-4", subfolder="unet")
model = UNet2DModel.from_config(config)

I think this could still be confused by some users as loading the pretrained weights instead of instantiating them with default/random values. The reason for that is that the configuration was retrieved from a pretrained repo, so I think it's not totally unreasonable to expect that creating the model from that particular config the weights would also be retrieved. I'd prefer to give it a different name, but it's better to keep consistency with transformers.

I'm mentioning this because there is currently no need to advertise from_config for any of the models (except for training). However, we do need to use from_config to switch schedulers for inference, and that instinctively feels a little off. What if the user wants to replace another component, like one of the fine-tuned VAEs? They'd still need to use from_pretrained, not from_config. I don't have a good solution for this at this point, maybe we want to overload from_pretrained instead?

[patil-suraj]

Could we make it so the subfolder is optional? In this case we know we are loading a UNet configuration, so I'd rather skip the subfolder parameter. It's a bit silly, but then we don't require users to know the underlying tree structure of the repo.

Not sure if it could be made completely optional, for example in case a pipeline has two unets (like in cascaded diffusion), we will need to provide the subfolder to specify which unet we want.