Add single file support for Stable Cascade #7274
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| @@ -257,6 +257,39 @@ def fetch_ldm_config_and_checkpoint( | |||
| return original_config, checkpoint | |||
|
|
|||
|
|
|||
| def load_single_file_model_checkpoint(pretrained_model_link_or_path, **kwargs): | |||
There was a problem hiding this comment.
I find it confusing with fetch_ldm_config_and_checkpoint():
Can't we repurpose fetch_ldm_config_and_checkpoint() here? Code looks almost the same to me.
There was a problem hiding this comment.
Can move load_single_file_model_checkpoint into the fetch_ldm_config_and_checkpoint` function.
| } | ||
|
|
||
|
|
||
| def convert_single_file_to_diffusers(original_state_dict): |
There was a problem hiding this comment.
We could leverage these functions in the conversion script to massively reduce the duplicated code, no?
(Future PR candidate)
There was a problem hiding this comment.
I copied these over from the conversion script. But yes, I can update the scripts to use this function once merged.
|
|
||
| assert unet.config[param_name] == param_value | ||
|
|
||
| def test_stable_cascade_config_loading(self): |
There was a problem hiding this comment.
(nit): We're testing the config for the UNet. The test, hence, should be named accordingly.
| unet = StableCascadeUNet.from_pretrained( | ||
| "stabilityai/stable-cascade-prior", | ||
| subfolder="prior", | ||
| revision="refs/pr/2", |
There was a problem hiding this comment.
Left a couple of nits.
My biggest concern is that with the current design to support single-file checkpoint loading, we are breaking the design of how we support it for ControlNet and VAE massively. Personally, I like this current approach, though!
I also think we should have a dedicated section in the Stable Cascade documentation page about how one should properly load a single-file checkpoint and run inference to obtain valid results.
I think what we can do is move towards a single |
I don't fully understand. Could you maybe elaborate this with pseudocode? |
| @@ -387,6 +475,52 @@ def get_block(block_type, in_channels, nhead, c_skip=0, dropout=0, self_attn=Tru | |||
|
|
|||
| self.gradient_checkpointing = False | |||
|
|
|||
| @classmethod | |||
| def from_single_file(cls, pretrained_model_link_or_path, **kwargs): | |||
There was a problem hiding this comment.
can we quickly create a FromOriginalUnetMixin, move everything there, and throw a warning if the cls.__name != StableCascadeUNet?
we can refactor the code later and apply to other unets
There was a problem hiding this comment.
Would prefer the other way around then. I.e., create a FromOriginalUnetMixin and then land this PR. Since with this PR, we are including a big anti-pattern in the code, which I would like to avoid at all costs.
|
great stuff. |
|
The docstring is. #7295 should clean it up. |
What does this PR do?
Adds single file support to the StableCascadeUNet, which allows users to load in checkpoints published in the original format.
The single file loading logic and mappings are defined within the UNet without a Mixin, since the mapping functions have be accessible to the
from_single_filemethod somehow. It relies on fetching the configs from a hosted repo of configs hosted on the diffusers org.This is a more practical approach, since a single file pipeline checkpoint would be quite large (~34GB) to load. Single File loading with combined pipelines is also not something we support at the moment and it is particularly challenging with Cascade due to the dtype restrictions in the Prior (Stage C) pipeline, which cannot work with float16, while the decoder (Stage B) can.
Fixes # (issue)
Before submitting
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.