[Refactor] Better align `from_single_file` logic with `from_pretrained` #7496

DN6 · 2024-03-27T16:46:05Z

What does this PR do?

Single file loading is still far from ideal. We should aim to align the behaviour as close to from_pretrained as possible (with the goal of converging to a single loading method)

Some of the issues that need to be addressed

Different ways of setting model/pipeline configurations. We currently support arguments such as num_in_channels, scheduler_type , load_safety_checker which are not supported in the diffusers model configs or in from_pretrained. We should deprecate these in favour of using the same configuration override methods that we use in from_pretrained
Configuring model parameters with heuristics based on the pipeline name, e.g. Setting the number of In Channels for a pipeline based on the invoking class. This isn't great because we should be able to configure a model based on the information within the checkpoint alone.
Loading pipeline components is still quite rigid and relies a lot of heuristics to fetch/load the components. We should be able to reuse model_index.json files to determine the correct classes for each component in the pipeline
Not respecting the configured scheduler type in the model_index.json file on the hub. Single file currently defaults to always using the DDIM scheduler

This PR attempts to get single file loading behaviour much closer to the logic used in from_pretrained

For Models

It pushes the model loading logic into a FromOriginalModelMixin that fetches the appropriate model config.json file from the hub based on the keys provided in the checkpoint
Defines model specific mapping functions that converts the original state dict to a diffusers state dict.
Apply this OriginalModelMixin to all Diffusers models that are meant to support from_single_file loading

This allows us to rely on a lot of the functionality already defined in ModelMixin and DiffusionPipeline to create the correct model

For Pipelines

Rather than relying on the original yaml files to configure the Pipeline and Models, we should opt to identify the appropriate model repo related to the single file checkpoint based on the checkpoint keys. This allows us to fetch the model_index.json file for this checkpoint and load the components using similar logic to from_pretrained. Note loading a Pipeline/Model via YAML is still supported, it is just not the default anymore.
Allow passing in local_dir and local_dir_use_symlinks arguments to control where checkpoints are downloaded and to disable symlinking if users request it.

TODO:

General clean up
Some clean up in single_file_utils
Move the Cascade loading logic to follow this system
Add tests to ensure this isn't backwards breaking
Improve the docs a bit more to demonstrate single file functionality.
Look into supporting connected pipelines via this approach.

This should make it a bit easier to add more Pipelines/Models with single file support. The process would become

Define a mapping function for the model from the original dict to diffusers dict
Create a model repo with the config files
Find a way to infer the model type from the checkpoint and fetch the appropriate config

We already have to take care of steps 1 and 2 when adding a model to diffusers. So it's just a matter inferring the proper config from the checkpoint (not easy but very possible).

Fixes # (issue)

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline?
Did you read our philosophy doc (important for complex PRs)?
Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

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.

sayakpaul · 2024-03-28T02:09:11Z

docs/source/en/api/single_file.md

+- [`StableDiffusionPipeline`]
+- [`StableDiffusionImg2ImgPipeline`]
+- [`StableDiffusionInpaintPipeline`]
+- [`StableDiffusionControlNetPipeline`]
+- [`StableDiffusionControlNetImg2ImgPipeline`]
+- [`StableDiffusionControlNetInpaintPipeline`]
+- [`StableDiffusionUpscalePipeline`]
+- [`StableDiffusionXLPipeline`]
+- [`StableDiffusionXLImg2ImgPipeline`]
+- [`StableDiffusionXLInpaintPipeline`]
+- [`StableDiffusionXLControlNetPipeline`]


We should consider having a utility method to find this programmatically as well.

sayakpaul · 2024-03-28T02:09:26Z

docs/source/en/api/single_file.md

+## Models that currently support `from_single_file` loading
+
+- [`UNet2DConditionModel`]
+- [`StableCascadeUNet`]
+- [`AutoencoderKL`]
+- [`ControlNetModel`]


Same here as well.

docs/source/en/api/single_file.md

sayakpaul · 2024-03-28T02:10:37Z

docs/source/en/api/single_file.md

+scheduler = DDIMScheduler()
+pipe = StableDiffusionXLPipeline.from_single_file(ckpt_path, scheduler=scheduler)
+```
+


Add a line here?

docs/source/en/api/single_file.md

src/diffusers/loaders/single_file.py

sayakpaul · 2024-03-28T02:16:13Z

src/diffusers/loaders/single_file.py

+        config_file = hf_hub_download(
+            default_pipeline_config["pretrained_model_name_or_path"],
+            filename=cls.config_name,
+            cache_dir=cache_dir,
+            revision=revision,
+            proxies=proxies,
+            force_download=force_download,
+            resume_download=resume_download,
+            token=token,
+            local_files_only=local_files_only,
+        )


There are many single file checkpoints for which we don't have equivalent diffusers checkpoints, but they can be loaded into a pipeline. How does this stand in that case? We fetch the corresponding config file examining the checkpoint and do a best effort to map accordingly? If so, I think there should be a comment about it.

do you have examples of these?

There are many single file checkpoints for which we don't have equivalent diffusers checkpoints

Have seen many but cannot recall any exact names of the top of my head. You can find a couple on CivitAI too (it hosts full finetunes of SD and SDXL too not just LoRAs). Cc: @vladmandic here.

HuggingFaceDocBuilderDev · 2024-04-29T08:03:53Z

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.

DN6 · 2024-05-01T12:28:51Z

@vladmandic I think this is almost ready to merge, if you have the time would you be able to test the branch to see if any breaking changes occur. I tried to make sure we have full backwards compatibility, but given the number of changes here, I just want to be on the safe side.

sayakpaul

Thanks for working on this. Some tests seem to be failing. Here's you can find the detailed logs:

https://huggingface.co/datasets/sayakpaul/sample-datasets/blob/main/text_outputs_single_file/

sayakpaul · 2024-05-02T00:39:22Z

src/diffusers/configuration_utils.py

                os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
            ):
                config_file = os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
+            elif os.path.isfile(os.path.join(pretrained_model_name_or_path, cls.config_name)):
+                # Load from a PyTorch checkpoint


Suggested change

# Load from a PyTorch checkpoint

# Load from a PyTorch checkpoint (SD checkpoints usually have some configuration details in them)

sayakpaul · 2024-05-02T00:44:11Z

src/diffusers/loaders/single_file.py

+    local_files_only=None,
+    token=None,
+):
+    allow_patterns = ["**/*.json", "*.json", "*.txt", "**/*.txt"]


Why do we allow txt files here?

Also, nit: ["**/*.json", "*.json", "*.txt", "**/*.txt"] can be reduced to just ["*.json", "*.txt"], I think.

Tokenizer subfolders have text files as part of their configs sometimes. e.g.
https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/tree/main/tokenizer_2

sayakpaul · 2024-05-02T00:44:44Z

src/diffusers/loaders/single_file.py

+        allow_patterns=allow_patterns,
+    )
+
+    return cached_model_path


Prefer cached_model_config_path.

sayakpaul · 2024-05-02T00:45:43Z

src/diffusers/loaders/single_file.py

+        # We shouldn't allow configuring individual models components through a Pipeline creation method
+        # These model kwargs should be deprecated
+        scaling_factor = kwargs.get("scaling_factor", None)
+        if scaling_factor is not None:
+            deprecation_message = (
+                "Passing the `scaling_factor` argument to `from_single_file is deprecated "
+                "and will be ignored in future versions."
+            )
+            deprecate("scaling_factor", "1.0.0", deprecation_message)


Can we include instructions on what the users should use instead in this message?

sayakpaul · 2024-05-02T00:54:02Z

src/diffusers/loaders/single_file.py

+                        "Detected legacy `from_single_file` loading behavior. Attempting to create the pipeline based on inferred components.\n"
+                        "This may lead to errors if the model components are not correctly inferred. "
+                        "To avoid this warning, please explicity pass the `config` argument to `from_single_file` with a path to a local diffusers model repo "
+                        "or run `from_single_file` with `local_files_only=False` first to update the local cache directory with "
+                        "the necessary config files.\n"


An example config="..." in the message would be helpful for the users.

sayakpaul · 2024-05-02T01:16:51Z

.github/workflows/push_tests.yml

@@ -123,7 +123,7 @@ jobs:
        shell: bash
    strategy:
      matrix:
-        module: [models, schedulers, lora, others]
+        module: [models, schedulers, lora, others, single_file]


Shouldn't we also include this in the nightly tests?

sayakpaul · 2024-05-02T01:25:03Z

tests/single_file/test_model_controlnet_single_file.py

@@ -0,0 +1,78 @@
+# coding=utf-8


WDYT about further grouping this so that it's easier for the contributors to pique through a certain test?

Grouping like so:

single_file controlnet vae stable_diffusion stable_cascade ...

sayakpaul · 2024-05-02T01:33:30Z

tests/single_file/test_model_vae_single_file.py

+        image = torch.from_numpy(load_hf_numpy(self.get_file_format(seed, shape))).to(torch_device).to(dtype)
+        return image
+
+    def test_single_file_inference_same_as_pretrained(self):


This is failing on a T4:

FAILED tests/single_file/test_model_vae_single_file.py::AutoencoderKLSingleFileTests::test_single_file_inference_same_as_pretrained - AssertionError: Max diff is absolute 0.9882212281227112. Diff tensor is tensor([0.0288, 0.9882, 0.0790, 0.1767, 0.1501, 0.9556, 0.1592, 0.0806]).

Could you please check?

sayakpaul · 2024-05-02T01:36:32Z

tests/single_file/single_file_testing_utils.py

+        self._compare_component_configs(pipe, single_file_pipe)
+
+
+class SDXLSingleFileTesterMixin:


Is there a major difference between SDXLSingleFileTesterMixin and SDSingleFileTesterMixin?

SDXLSingleFileTesterMixin accounts for the multiple text encoder components that may or may not be present in the pipeline depending on the checkpoint.

sayakpaul · 2024-05-02T01:40:39Z

tests/single_file/test_stable_diffusion_controlnet_inpaint_single_file.py

@@ -0,0 +1,183 @@
+import gc


Failing tests:

FAILED tests/single_file/test_stable_diffusion_controlnet_inpaint_single_file.py::StableDiffusionControlNetInpaintPipelineSingleFileSlowTests::test_single_file_components_with_original_config - AssertionError: single file sample_size: 512 differs from pretrained 256 FAILED tests/single_file/test_stable_diffusion_controlnet_inpaint_single_file.py::StableDiffusionControlNetInpaintPipelineSingleFileSlowTests::test_single_file_components_with_original_config_local_files_only - AssertionError: single file sample_size: 512 differs from pretrained 256

Issue is that the config on the hub is incorrect.

yiyixuxu

thanks!!!!!
i left only some nits! let's merge this in soon!!!

yiyixuxu · 2024-05-02T20:06:36Z

src/diffusers/loaders/single_file.py

@@ -252,69 +362,190 @@ def from_single_file(cls, pretrained_model_link_or_path, **kwargs):
        local_files_only = kwargs.pop("local_files_only", False)
        revision = kwargs.pop("revision", None)
        torch_dtype = kwargs.pop("torch_dtype", None)
+        use_safetensors = kwargs.pop("use_safetensors", None)
+
+        is_legacy_loading = False


docstring needed here

yiyixuxu · 2024-05-02T20:06:49Z

src/diffusers/loaders/single_file.py

@@ -252,69 +362,190 @@ def from_single_file(cls, pretrained_model_link_or_path, **kwargs):
        local_files_only = kwargs.pop("local_files_only", False)
        revision = kwargs.pop("revision", None)
        torch_dtype = kwargs.pop("torch_dtype", None)
+        use_safetensors = kwargs.pop("use_safetensors", None)


is this relevant now because we will run from_pretrained() for missing components? need a doc string here:)

yiyixuxu · 2024-05-02T20:09:07Z

docs/source/en/api/loaders/single_file.md

+
+```
+
+## Override configuration options when using single file loading


I think we should swap this section with the next one, because in the example we used the config argument that hasn't been explained yet (but it will be in the next section)

yiyixuxu · 2024-05-02T20:26:34Z

src/diffusers/loaders/single_file.py

+        use_safetensors = kwargs.pop("use_safetensors", None)
+
+        is_legacy_loading = False
+


should we also accept a local_dir_use_symlinks here and consistently pass it down to any place that downloads things?
currently, the from_single_file for models accepts local_dir_use_symlinks but not pipeline

Hmm actually might be better to not include local_dir support in this version. I think it's better for the user to predownload stuff to local paths beforehand. I'll remove this.

src/diffusers/loaders/single_file.py

yiyixuxu · 2024-05-02T21:24:24Z

src/diffusers/loaders/single_file_model.py

+
+SINGLE_FILE_LOADABLE_CLASSES = {
+    "StableCascadeUNet": {
+        "checkpoint_mapping_fn": convert_stable_cascade_unet_single_file_to_diffusers,


the single checkpoints are not in subfolders
https://huggingface.co/stabilityai/stable-cascade/tree/main

yiyixuxu · 2024-05-02T21:31:09Z

src/diffusers/loaders/single_file_model.py

+                f"FromOriginalModelMixin is currently only compatible with {', '.join(SINGLE_FILE_LOADABLE_CLASSES.keys())}"
+            )
+
+        checkpoint = kwargs.pop("checkpoint", None)


oh maybe rename it to pretrained_model_name_or_path_or_dict and consolidate these two arguments?
it is done like this for load_lora_weights and load_ip_adapter too

diffusers/src/diffusers/loaders/lora.py

Line 82 in 0d7c479

def load_lora_weights(

yiyixuxu · 2024-05-02T21:56:27Z

src/diffusers/loaders/single_file_model.py

+        mapping_functions = SINGLE_FILE_LOADABLE_CLASSES[class_name]
+
+        checkpoint_mapping_fn = mapping_functions["checkpoint_mapping_fn"]
+        if "config_mapping_fn" in mapping_functions:


nit: I would move this inside the if original_config: ... block because it is only used there (less code to read when it is not applicable)

yiyixuxu · 2024-05-02T22:31:06Z

src/diffusers/loaders/unet.py

@@ -1002,156 +997,3 @@ def _load_ip_adapter_weights(self, state_dicts, low_cpu_mem_usage=False):
        self.config.encoder_hid_dim_type = "ip_image_proj"

        self.to(dtype=self.dtype, device=self.device)
-
-    def _load_ip_adapter_loras(self, state_dicts):


+1
should not go away I think

…d` (#7496) * refactor unet single file loading a bit. * retrieve the unet from create_diffusers_unet_model_from_ldm * update * update * updae * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * tests * update * update * update * Update docs/source/en/api/single_file.md Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Update docs/source/en/api/single_file.md Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * update * update * update * update * update * update * update * update * update * update * update * update * update * Update docs/source/en/api/loaders/single_file.md Co-authored-by: YiYi Xu <yixu310@gmail.com> * Update src/diffusers/loaders/single_file.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * Update docs/source/en/api/loaders/single_file.md Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Update docs/source/en/api/loaders/single_file.md Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Update docs/source/en/api/loaders/single_file.md Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * Update docs/source/en/api/loaders/single_file.md Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update --------- Co-authored-by: sayakpaul <spsayakpaul@gmail.com> Co-authored-by: YiYi Xu <yixu310@gmail.com>

sayakpaul and others added 19 commits March 14, 2024 15:33

refactor unet single file loading a bit.

f03ea10

retrieve the unet from create_diffusers_unet_model_from_ldm

bfaa0d8

update

bc32a9d

update

5bb7d56

updae

56863a2

update

2cd8175

update

92cf552

update

57aa8be

update

8c9a890

update

5cb4f12

update

0dd26eb

update

8e0cdd2

update

5203eb6

update

17f3cbd

update

3d0bc40

update

88389a2

update

3850ef8

update

7997372

update

64bdee0

DN6 changed the title ~~[Refactor] Align from_single_file with from_pretrained~~ [Refactor] Better align from_single_file logic with from_pretrained Mar 27, 2024

sayakpaul reviewed Mar 28, 2024

View reviewed changes

docs/source/en/api/single_file.md Outdated Show resolved Hide resolved

sayakpaul reviewed Mar 28, 2024

View reviewed changes

docs/source/en/api/single_file.md Outdated Show resolved Hide resolved

sayakpaul reviewed Mar 28, 2024

View reviewed changes

docs/source/en/api/single_file.md Outdated Show resolved Hide resolved

sayakpaul reviewed Mar 28, 2024

View reviewed changes

docs/source/en/api/single_file.md Outdated Show resolved Hide resolved

sayakpaul reviewed Mar 28, 2024

View reviewed changes

src/diffusers/loaders/single_file.py Outdated Show resolved Hide resolved

sayakpaul reviewed Mar 28, 2024

View reviewed changes

src/diffusers/loaders/single_file.py Outdated Show resolved Hide resolved

sayakpaul reviewed Mar 28, 2024

View reviewed changes

DN6 added 3 commits April 26, 2024 11:19

update

03a2ed8

update

0051843

update

a5c78c2

DN6 added 2 commits April 30, 2024 11:28

Merge branch 'main' into single-file-updates

96f1b2e

update

47f825d

sayakpaul mentioned this pull request May 2, 2024

[Tests] adds fast tests for single-file checkpoint loaders #6961

Closed

sayakpaul reviewed May 2, 2024

View reviewed changes

yiyixuxu approved these changes May 2, 2024

View reviewed changes

DN6 added 15 commits May 7, 2024 12:14

Merge branch 'main' into single-file-updates

8e41325

update

bd2e73f

update

f5e4017

update

954c20a

update

a04562f

update

4a8f072

update

cc16cc8

update

28bf5ad

update

8387950

update

fff5297

update

696b258

update

d364604

update

6a22444

update

c61779d

Merge branch 'main' into single-file-updates

f211c04

DN6 merged commit cb0f3b4 into main May 9, 2024
17 checks passed

yiyixuxu mentioned this pull request May 9, 2024

single_file_utils.py#load_single_file_model_checkpoint bugs #7889

Open

yiyixuxu deleted the single-file-updates branch May 9, 2024 16:59

suzukimain mentioned this pull request Dec 14, 2024

Integration of from_pretrained and from_single_file #10208

Open

6 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Refactor] Better align `from_single_file` logic with `from_pretrained` #7496

[Refactor] Better align `from_single_file` logic with `from_pretrained` #7496

DN6 commented Mar 27, 2024

sayakpaul Mar 28, 2024

sayakpaul Mar 28, 2024

sayakpaul Mar 28, 2024

sayakpaul Mar 28, 2024

yiyixuxu Mar 28, 2024

sayakpaul Mar 29, 2024

HuggingFaceDocBuilderDev commented Apr 29, 2024

DN6 commented May 1, 2024

sayakpaul left a comment

sayakpaul May 2, 2024

sayakpaul May 2, 2024

DN6 May 8, 2024

sayakpaul May 2, 2024

sayakpaul May 2, 2024

sayakpaul May 2, 2024

sayakpaul May 2, 2024

sayakpaul May 2, 2024

sayakpaul May 2, 2024

sayakpaul May 2, 2024

DN6 May 7, 2024

sayakpaul May 2, 2024

DN6 May 8, 2024

yiyixuxu left a comment

yiyixuxu May 2, 2024

yiyixuxu May 2, 2024

yiyixuxu May 2, 2024

yiyixuxu May 2, 2024

DN6 May 7, 2024 •

edited

Loading

yiyixuxu May 2, 2024

yiyixuxu May 2, 2024

yiyixuxu May 2, 2024

yiyixuxu May 2, 2024

	# Load from a PyTorch checkpoint
	# Load from a PyTorch checkpoint (SD checkpoints usually have some configuration details in them)

		self._compare_component_configs(pipe, single_file_pipe)


		class SDXLSingleFileTesterMixin:


		```

		## Override configuration options when using single file loading

		use_safetensors = kwargs.pop("use_safetensors", None)

		is_legacy_loading = False

[Refactor] Better align from_single_file logic with from_pretrained #7496

[Refactor] Better align from_single_file logic with from_pretrained #7496

Conversation

DN6 commented Mar 27, 2024

What does this PR do?

For Models

For Pipelines

Before submitting

Who can review?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

HuggingFaceDocBuilderDev commented Apr 29, 2024

DN6 commented May 1, 2024

sayakpaul left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yiyixuxu left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

DN6 May 7, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

[Refactor] Better align `from_single_file` logic with `from_pretrained` #7496

[Refactor] Better align `from_single_file` logic with `from_pretrained` #7496

DN6 May 7, 2024 •

edited

Loading