Skip to content

Commit e931bef

Browse files
DN6sayakpaulyiyixuxu
committed
[Refactor] Better align from_single_file logic with from_pretrained (#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 <[email protected]> * Update docs/source/en/api/single_file.md Co-authored-by: Sayak Paul <[email protected]> * 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 <[email protected]> * Update src/diffusers/loaders/single_file.py Co-authored-by: YiYi Xu <[email protected]> * Update docs/source/en/api/loaders/single_file.md Co-authored-by: Sayak Paul <[email protected]> * Update docs/source/en/api/loaders/single_file.md Co-authored-by: Sayak Paul <[email protected]> * Update docs/source/en/api/loaders/single_file.md Co-authored-by: Sayak Paul <[email protected]> * Update docs/source/en/api/loaders/single_file.md Co-authored-by: Sayak Paul <[email protected]> * 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 <[email protected]> Co-authored-by: YiYi Xu <[email protected]>
1 parent 8bc357b commit e931bef

File tree

45 files changed

+3790
-1747
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

45 files changed

+3790
-1747
lines changed

.github/workflows/push_tests.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ jobs:
124124
shell: bash
125125
strategy:
126126
matrix:
127-
module: [models, schedulers, lora, others]
127+
module: [models, schedulers, lora, others, single_file]
128128
steps:
129129
- name: Checkout diffusers
130130
uses: actions/checkout@v3

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

Lines changed: 222 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -10,28 +10,239 @@ an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express o
1010
specific language governing permissions and limitations under the License.
1111
-->
1212

13-
# Single files
13+
# Loading Pipelines and Models via `from_single_file`
1414

15-
Diffusers supports loading pretrained pipeline (or model) weights stored in a single file, such as a `ckpt` or `safetensors` file. These single file types are typically produced from community trained models. There are three classes for loading single file weights:
15+
The `from_single_file` method allows you to load supported pipelines using a single checkpoint file as opposed to the folder format used by Diffusers. This is useful if you are working with many of the Stable Diffusion Web UI's (such as A1111) that extensively rely on a single file to distribute all the components of a diffusion model.
1616

17-
- [`FromSingleFileMixin`] supports loading pretrained pipeline weights stored in a single file, which can either be a `ckpt` or `safetensors` file.
18-
- [`FromOriginalVAEMixin`] supports loading a pretrained [`AutoencoderKL`] from pretrained ControlNet weights stored in a single file, which can either be a `ckpt` or `safetensors` file.
19-
- [`FromOriginalControlnetMixin`] supports loading pretrained ControlNet weights stored in a single file, which can either be a `ckpt` or `safetensors` file.
17+
The `from_single_file` method also supports loading models in their originally distributed format. This means that supported models that have been finetuned with other services can be loaded directly into supported Diffusers model objects and pipelines.
18+
19+
## Pipelines that currently support `from_single_file` loading
20+
21+
- [`StableDiffusionPipeline`]
22+
- [`StableDiffusionImg2ImgPipeline`]
23+
- [`StableDiffusionInpaintPipeline`]
24+
- [`StableDiffusionControlNetPipeline`]
25+
- [`StableDiffusionControlNetImg2ImgPipeline`]
26+
- [`StableDiffusionControlNetInpaintPipeline`]
27+
- [`StableDiffusionUpscalePipeline`]
28+
- [`StableDiffusionXLPipeline`]
29+
- [`StableDiffusionXLImg2ImgPipeline`]
30+
- [`StableDiffusionXLInpaintPipeline`]
31+
- [`StableDiffusionXLInstructPix2PixPipeline`]
32+
- [`StableDiffusionXLControlNetPipeline`]
33+
- [`StableDiffusionXLKDiffusionPipeline`]
34+
- [`LatentConsistencyModelPipeline`]
35+
- [`LatentConsistencyModelImg2ImgPipeline`]
36+
- [`StableDiffusionControlNetXSPipeline`]
37+
- [`StableDiffusionXLControlNetXSPipeline`]
38+
- [`LEditsPPPipelineStableDiffusion`]
39+
- [`LEditsPPPipelineStableDiffusionXL`]
40+
- [`PIAPipeline`]
41+
42+
## Models that currently support `from_single_file` loading
43+
44+
- [`UNet2DConditionModel`]
45+
- [`StableCascadeUNet`]
46+
- [`AutoencoderKL`]
47+
- [`ControlNetModel`]
48+
49+
## Usage Examples
50+
51+
## Loading a Pipeline using `from_single_file`
52+
53+
```python
54+
from diffusers import StableDiffusionXLPipeline
55+
56+
ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors"
57+
pipe = StableDiffusionXLPipeline.from_single_file(ckpt_path)
58+
```
59+
60+
## Setting components in a Pipeline using `from_single_file`
61+
62+
Swap components of the pipeline by passing them directly to the `from_single_file` method. e.g If you would like use a different scheduler than the pipeline default.
63+
64+
```python
65+
from diffusers import StableDiffusionXLPipeline, DDIMScheduler
66+
67+
ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors"
68+
69+
scheduler = DDIMScheduler()
70+
pipe = StableDiffusionXLPipeline.from_single_file(ckpt_path, scheduler=scheduler)
71+
72+
```
73+
74+
```python
75+
from diffusers import StableDiffusionPipeline, ControlNetModel
76+
77+
ckpt_path = "https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/v1-5-pruned-emaonly.safetensors"
78+
79+
controlnet = ControlNetModel.from_pretrained("https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/v1-5-pruned-emaonly.safetensors")
80+
pipe = StableDiffusionPipeline.from_single_file(ckpt_path, controlnet=controlnet)
81+
82+
```
83+
84+
## Loading a Model using `from_single_file`
85+
86+
```python
87+
from diffusers import StableCascadeUNet
88+
89+
ckpt_path = "https://huggingface.co/stabilityai/stable-cascade/blob/main/stage_b_lite.safetensors"
90+
model = StableCascadeUNet.from_single_file(ckpt_path)
91+
92+
```
93+
94+
## Using a Diffusers model repository to configure single file loading
95+
96+
Under the hood, `from_single_file` will try to determine a model repository to use to configure the components of the pipeline. You can also pass in a repository id to the `config` argument of the `from_single_file` method to explicitly set the repository to use.
97+
98+
```python
99+
from diffusers import StableDiffusionXLPipeline
100+
101+
ckpt_path = "https://huggingface.co/segmind/SSD-1B/blob/main/SSD-1B.safetensors"
102+
repo_id = "segmind/SSD-1B"
103+
104+
pipe = StableDiffusionXLPipeline.from_single_file(ckpt_path, config=repo_id)
105+
106+
```
107+
108+
## Override configuration options when using single file loading
109+
110+
Override the default model or pipeline configuration options when using `from_single_file` by passing in the relevant arguments directly to the `from_single_file` method. Any argument that is supported by the model or pipeline class can be configured in this way:
111+
112+
```python
113+
from diffusers import StableDiffusionXLInstructPix2PixPipeline
114+
115+
ckpt_path = "https://huggingface.co/stabilityai/cosxl/blob/main/cosxl_edit.safetensors"
116+
pipe = StableDiffusionXLInstructPix2PixPipeline.from_single_file(ckpt_path, config="diffusers/sdxl-instructpix2pix-768", is_cosxl_edit=True)
117+
118+
```
119+
120+
```python
121+
from diffusers import UNet2DConditionModel
122+
123+
ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors"
124+
model = UNet2DConditionModel.from_single_file(ckpt_path, upcast_attention=True)
125+
126+
```
127+
128+
In the example above, since we explicitly passed `repo_id="segmind/SSD-1B"`, it will use this [configuration file](https://huggingface.co/segmind/SSD-1B/blob/main/unet/config.json) from the "unet" subfolder in `"segmind/SSD-1B"` to configure the unet component included in the checkpoint; Similarly, it will use the `config.json` file from `"vae"` subfolder to configure the vae model, `config.json` file from text_encoder folder to configure text_encoder and so on.
129+
130+
Note that most of the time you do not need to explicitly a `config` argument, `from_single_file` will automatically map the checkpoint to a repo id (we will discuss this in more details in next section). However, this can be useful in cases where model components might have been changed from what was originally distributed or in cases where a checkpoint file might not have the necessary metadata to correctly determine the configuration to use for the pipeline.
20131

21132
<Tip>
22133

23134
To learn more about how to load single file weights, see the [Load different Stable Diffusion formats](../../using-diffusers/other-formats) loading guide.
24135

25136
</Tip>
26137

27-
## FromSingleFileMixin
138+
## Working with local files
28139

29-
[[autodoc]] loaders.single_file.FromSingleFileMixin
140+
As of `diffusers>=0.28.0` the `from_single_file` method will attempt to configure a pipeline or model by first inferring the model type from the checkpoint file and then using the model type to determine the appropriate model repo configuration to use from the Hugging Face Hub. For example, any single file checkpoint based on the Stable Diffusion XL base model will use the [`stabilityai/stable-diffusion-xl-base-1.0`](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) model repo to configure the pipeline.
30141

31-
## FromOriginalVAEMixin
142+
If you are working in an environment with restricted internet access, it is recommended to download the config files and checkpoints for the model to your preferred directory and pass the local paths to the `pretrained_model_link_or_path` and `config` arguments of the `from_single_file` method.
32143

33-
[[autodoc]] loaders.autoencoder.FromOriginalVAEMixin
144+
```python
145+
from huggingface_hub import hf_hub_download, snapshot_download
146+
147+
my_local_checkpoint_path = hf_hub_download(
148+
repo_id="segmind/SSD-1B",
149+
filename="SSD-1B.safetensors"
150+
)
151+
152+
my_local_config_path = snapshot_download(
153+
repo_id="segmind/SSD-1B",
154+
allowed_patterns=["*.json", "**/*.json", "*.txt", "**/*.txt"]
155+
)
156+
157+
pipe = StableDiffusionXLPipeline.from_single_file(my_local_checkpoint_path, config=my_local_config_path, local_files_only=True)
158+
159+
```
160+
161+
By default this will download the checkpoints and config files to the [Hugging Face Hub cache directory](https://huggingface.co/docs/huggingface_hub/en/guides/manage-cache). You can also specify a local directory to download the files to by passing the `local_dir` argument to the `hf_hub_download` and `snapshot_download` functions.
162+
163+
```python
164+
from huggingface_hub import hf_hub_download, snapshot_download
165+
166+
my_local_checkpoint_path = hf_hub_download(
167+
repo_id="segmind/SSD-1B",
168+
filename="SSD-1B.safetensors"
169+
local_dir="my_local_checkpoints"
170+
)
171+
172+
my_local_config_path = snapshot_download(
173+
repo_id="segmind/SSD-1B",
174+
allowed_patterns=["*.json", "**/*.json", "*.txt", "**/*.txt"]
175+
local_dir="my_local_config"
176+
)
177+
178+
pipe = StableDiffusionXLPipeline.from_single_file(my_local_checkpoint_path, config=my_local_config_path, local_files_only=True)
179+
180+
```
181+
182+
## Working with local files on file systems that do not support symlinking
183+
184+
By default the `from_single_file` method relies on the `huggingface_hub` caching mechanism to fetch and store checkpoints and config files for models and pipelines. If you are working with a file system that does not support symlinking, it is recommended that you first download the checkpoint file to a local directory and disable symlinking by passing the `local_dir_use_symlink=False` argument to the `hf_hub_download` and `snapshot_download` functions.
185+
186+
```python
187+
from huggingface_hub import hf_hub_download, snapshot_download
188+
189+
my_local_checkpoint_path = hf_hub_download(
190+
repo_id="segmind/SSD-1B",
191+
filename="SSD-1B.safetensors"
192+
local_dir="my_local_checkpoints",
193+
local_dir_use_symlinks=False
194+
)
195+
print("My local checkpoint: ", my_local_checkpoint_path)
196+
197+
my_local_config_path = snapshot_download(
198+
repo_id="segmind/SSD-1B",
199+
allowed_patterns=["*.json", "**/*.json", "*.txt", "**/*.txt"]
200+
local_dir_use_symlinks=False,
201+
)
202+
print("My local config: ", my_local_config_path)
203+
204+
```
205+
206+
Then pass the local paths to the `pretrained_model_link_or_path` and `config` arguments of the `from_single_file` method.
207+
208+
```python
209+
pipe = StableDiffusionXLPipeline.from_single_file(my_local_checkpoint_path, config=my_local_config_path, local_files_only=True)
210+
211+
```
212+
213+
<Tip>
214+
Disabling symlinking means that the `huggingface_hub` caching mechanism has no way to determine whether a file has already been downloaded to the local directory. This means that the `hf_hub_download` and `snapshot_download` functions will download files to the local directory each time they are executed. If you are disabling symlinking, it is recommended that you separate the model download and loading steps to avoid downloading the same file multiple times.
215+
216+
</Tip>
217+
218+
## Using the original configuration file of a model
219+
220+
If you would like to configure the parameters of the model components in the pipeline using the orignal YAML configuration file, you can pass a local path or url to the original configuration file to the `original_config` argument of the `from_single_file` method.
221+
222+
```python
223+
from diffusers import StableDiffusionXLPipeline
224+
225+
ckpt_path = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors"
226+
repo_id = "stabilityai/stable-diffusion-xl-base-1.0"
227+
original_config = "https://raw.githubusercontent.com/Stability-AI/generative-models/main/configs/inference/sd_xl_base.yaml"
228+
229+
pipe = StableDiffusionXLPipeline.from_single_file(ckpt_path, original_config=original_config)
230+
```
231+
232+
In the example above, the `original_config` file is only used to configure the parameters of the individual model components of the pipeline. For example it will be used to configure parameters such as the `in_channels` of the `vae` model and `unet` model. It is not used to determine the type of component objects in the pipeline.
233+
234+
235+
<Tip>
236+
When using `original_config` with local_files_only=True`, Diffusers will attempt to infer the components based on the type signatures of pipeline class, rather than attempting to fetch the pipeline config from the Hugging Face Hub. This is to prevent backwards breaking changes in existing code that might not be able to connect to the internet to fetch the necessary pipeline config files.
237+
238+
This is not as reliable as providing a path to a local config repo and might lead to errors when configuring the pipeline. To avoid this, please run the pipeline with `local_files_only=False` once to download the appropriate pipeline config files to the local cache.
239+
</Tip>
240+
241+
242+
## FromSingleFileMixin
243+
244+
[[autodoc]] loaders.single_file.FromSingleFileMixin
34245

35-
## FromOriginalControlnetMixin
246+
## FromOriginalModelMixin
36247

37-
[[autodoc]] loaders.controlnet.FromOriginalControlNetMixin
248+
[[autodoc]] loaders.single_file_model.FromOriginalModelMixin

src/diffusers/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727

2828
_import_structure = {
2929
"configuration_utils": ["ConfigMixin"],
30+
"loaders": ["FromOriginalModelMixin"],
3031
"models": [],
3132
"pipelines": [],
3233
"schedulers": [],

src/diffusers/configuration_utils.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -340,6 +340,8 @@ def load_config(
340340
341341
"""
342342
cache_dir = kwargs.pop("cache_dir", None)
343+
local_dir = kwargs.pop("local_dir", None)
344+
local_dir_use_symlinks = kwargs.pop("local_dir_use_symlinks", "auto")
343345
force_download = kwargs.pop("force_download", False)
344346
resume_download = kwargs.pop("resume_download", None)
345347
proxies = kwargs.pop("proxies", None)
@@ -364,13 +366,13 @@ def load_config(
364366
if os.path.isfile(pretrained_model_name_or_path):
365367
config_file = pretrained_model_name_or_path
366368
elif os.path.isdir(pretrained_model_name_or_path):
367-
if os.path.isfile(os.path.join(pretrained_model_name_or_path, cls.config_name)):
368-
# Load from a PyTorch checkpoint
369-
config_file = os.path.join(pretrained_model_name_or_path, cls.config_name)
370-
elif subfolder is not None and os.path.isfile(
369+
if subfolder is not None and os.path.isfile(
371370
os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
372371
):
373372
config_file = os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
373+
elif os.path.isfile(os.path.join(pretrained_model_name_or_path, cls.config_name)):
374+
# Load from a PyTorch checkpoint
375+
config_file = os.path.join(pretrained_model_name_or_path, cls.config_name)
374376
else:
375377
raise EnvironmentError(
376378
f"Error no file named {cls.config_name} found in directory {pretrained_model_name_or_path}."
@@ -390,6 +392,8 @@ def load_config(
390392
user_agent=user_agent,
391393
subfolder=subfolder,
392394
revision=revision,
395+
local_dir=local_dir,
396+
local_dir_use_symlinks=local_dir_use_symlinks,
393397
)
394398
except RepositoryNotFoundError:
395399
raise EnvironmentError(

src/diffusers/loaders/__init__.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -54,9 +54,7 @@ def text_encoder_attn_modules(text_encoder):
5454
_import_structure = {}
5555

5656
if is_torch_available():
57-
_import_structure["autoencoder"] = ["FromOriginalVAEMixin"]
58-
59-
_import_structure["controlnet"] = ["FromOriginalControlNetMixin"]
57+
_import_structure["single_file_model"] = ["FromOriginalModelMixin"]
6058
_import_structure["unet"] = ["UNet2DConditionLoadersMixin"]
6159
_import_structure["utils"] = ["AttnProcsLayers"]
6260
if is_transformers_available():
@@ -70,8 +68,7 @@ def text_encoder_attn_modules(text_encoder):
7068

7169
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
7270
if is_torch_available():
73-
from .autoencoder import FromOriginalVAEMixin
74-
from .controlnet import FromOriginalControlNetMixin
71+
from .single_file_model import FromOriginalModelMixin
7572
from .unet import UNet2DConditionLoadersMixin
7673
from .utils import AttnProcsLayers
7774

0 commit comments

Comments
 (0)