879 lines
41 KiB
Python
879 lines
41 KiB
Python
# coding=utf-8
|
|
# Copyright 2022 The HuggingFace Inc. team.
|
|
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import importlib
|
|
import inspect
|
|
import os
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
import diffusers
|
|
import PIL
|
|
from huggingface_hub import model_info, snapshot_download
|
|
from packaging import version
|
|
from PIL import Image
|
|
from tqdm.auto import tqdm
|
|
|
|
from .configuration_utils import ConfigMixin
|
|
from .dynamic_modules_utils import get_class_from_dynamic_module
|
|
from .hub_utils import http_user_agent
|
|
from .modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT
|
|
from .schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME
|
|
from .utils import (
|
|
CONFIG_NAME,
|
|
DIFFUSERS_CACHE,
|
|
ONNX_WEIGHTS_NAME,
|
|
WEIGHTS_NAME,
|
|
BaseOutput,
|
|
deprecate,
|
|
is_accelerate_available,
|
|
is_safetensors_available,
|
|
is_torch_version,
|
|
is_transformers_available,
|
|
logging,
|
|
)
|
|
|
|
|
|
if is_transformers_available():
|
|
import transformers
|
|
from transformers import PreTrainedModel
|
|
|
|
|
|
INDEX_FILE = "diffusion_pytorch_model.bin"
|
|
CUSTOM_PIPELINE_FILE_NAME = "pipeline.py"
|
|
DUMMY_MODULES_FOLDER = "diffusers.utils"
|
|
TRANSFORMERS_DUMMY_MODULES_FOLDER = "transformers.utils"
|
|
|
|
|
|
logger = logging.get_logger(__name__)
|
|
|
|
|
|
LOADABLE_CLASSES = {
|
|
"diffusers": {
|
|
"ModelMixin": ["save_pretrained", "from_pretrained"],
|
|
"SchedulerMixin": ["save_pretrained", "from_pretrained"],
|
|
"DiffusionPipeline": ["save_pretrained", "from_pretrained"],
|
|
"OnnxRuntimeModel": ["save_pretrained", "from_pretrained"],
|
|
},
|
|
"transformers": {
|
|
"PreTrainedTokenizer": ["save_pretrained", "from_pretrained"],
|
|
"PreTrainedTokenizerFast": ["save_pretrained", "from_pretrained"],
|
|
"PreTrainedModel": ["save_pretrained", "from_pretrained"],
|
|
"FeatureExtractionMixin": ["save_pretrained", "from_pretrained"],
|
|
"ProcessorMixin": ["save_pretrained", "from_pretrained"],
|
|
"ImageProcessingMixin": ["save_pretrained", "from_pretrained"],
|
|
},
|
|
"onnxruntime.training": {
|
|
"ORTModule": ["save_pretrained", "from_pretrained"],
|
|
},
|
|
}
|
|
|
|
ALL_IMPORTABLE_CLASSES = {}
|
|
for library in LOADABLE_CLASSES:
|
|
ALL_IMPORTABLE_CLASSES.update(LOADABLE_CLASSES[library])
|
|
|
|
|
|
@dataclass
|
|
class ImagePipelineOutput(BaseOutput):
|
|
"""
|
|
Output class for image pipelines.
|
|
|
|
Args:
|
|
images (`List[PIL.Image.Image]` or `np.ndarray`)
|
|
List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
|
|
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
|
|
"""
|
|
|
|
images: Union[List[PIL.Image.Image], np.ndarray]
|
|
|
|
|
|
@dataclass
|
|
class AudioPipelineOutput(BaseOutput):
|
|
"""
|
|
Output class for audio pipelines.
|
|
|
|
Args:
|
|
audios (`np.ndarray`)
|
|
List of denoised samples of shape `(batch_size, num_channels, sample_rate)`. Numpy array present the
|
|
denoised audio samples of the diffusion pipeline.
|
|
"""
|
|
|
|
audios: np.ndarray
|
|
|
|
|
|
def is_safetensors_compatible(info) -> bool:
|
|
filenames = set(sibling.rfilename for sibling in info.siblings)
|
|
pt_filenames = set(filename for filename in filenames if filename.endswith(".bin"))
|
|
is_safetensors_compatible = any(file.endswith(".safetensors") for file in filenames)
|
|
for pt_filename in pt_filenames:
|
|
prefix, raw = os.path.split(pt_filename)
|
|
if raw == "pytorch_model.bin":
|
|
# transformers specific
|
|
sf_filename = os.path.join(prefix, "model.safetensors")
|
|
else:
|
|
sf_filename = pt_filename[: -len(".bin")] + ".safetensors"
|
|
if is_safetensors_compatible and sf_filename not in filenames:
|
|
logger.warning(f"{sf_filename} not found")
|
|
is_safetensors_compatible = False
|
|
return is_safetensors_compatible
|
|
|
|
|
|
class DiffusionPipeline(ConfigMixin):
|
|
r"""
|
|
Base class for all models.
|
|
|
|
[`DiffusionPipeline`] takes care of storing all components (models, schedulers, processors) for diffusion pipelines
|
|
and handles methods for loading, downloading and saving models as well as a few methods common to all pipelines to:
|
|
|
|
- move all PyTorch modules to the device of your choice
|
|
- enabling/disabling the progress bar for the denoising iteration
|
|
|
|
Class attributes:
|
|
|
|
- **config_name** (`str`) -- name of the config file that will store the class and module names of all
|
|
components of the diffusion pipeline.
|
|
- **_optional_components** (List[`str`]) -- list of all components that are optional so they don't have to be
|
|
passed for the pipeline to function (should be overridden by subclasses).
|
|
"""
|
|
config_name = "model_index.json"
|
|
_optional_components = []
|
|
|
|
def register_modules(self, **kwargs):
|
|
# import it here to avoid circular import
|
|
from diffusers import pipelines
|
|
|
|
for name, module in kwargs.items():
|
|
# retrieve library
|
|
if module is None:
|
|
register_dict = {name: (None, None)}
|
|
else:
|
|
library = module.__module__.split(".")[0]
|
|
|
|
# check if the module is a pipeline module
|
|
pipeline_dir = module.__module__.split(".")[-2] if len(module.__module__.split(".")) > 2 else None
|
|
path = module.__module__.split(".")
|
|
is_pipeline_module = pipeline_dir in path and hasattr(pipelines, pipeline_dir)
|
|
|
|
# if library is not in LOADABLE_CLASSES, then it is a custom module.
|
|
# Or if it's a pipeline module, then the module is inside the pipeline
|
|
# folder so we set the library to module name.
|
|
if library not in LOADABLE_CLASSES or is_pipeline_module:
|
|
library = pipeline_dir
|
|
|
|
# retrieve class_name
|
|
class_name = module.__class__.__name__
|
|
|
|
register_dict = {name: (library, class_name)}
|
|
|
|
# save model index config
|
|
self.register_to_config(**register_dict)
|
|
|
|
# set models
|
|
setattr(self, name, module)
|
|
|
|
def save_pretrained(
|
|
self,
|
|
save_directory: Union[str, os.PathLike],
|
|
safe_serialization: bool = False,
|
|
):
|
|
"""
|
|
Save all variables of the pipeline that can be saved and loaded as well as the pipelines configuration file to
|
|
a directory. A pipeline variable can be saved and loaded if its class implements both a save and loading
|
|
method. The pipeline can easily be re-loaded using the `[`~DiffusionPipeline.from_pretrained`]` class method.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to which to save. Will be created if it doesn't exist.
|
|
safe_serialization (`bool`, *optional*, defaults to `False`):
|
|
Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
|
|
"""
|
|
self.save_config(save_directory)
|
|
|
|
model_index_dict = dict(self.config)
|
|
model_index_dict.pop("_class_name")
|
|
model_index_dict.pop("_diffusers_version")
|
|
model_index_dict.pop("_module", None)
|
|
|
|
expected_modules, optional_kwargs = self._get_signature_keys(self)
|
|
|
|
def is_saveable_module(name, value):
|
|
if name not in expected_modules:
|
|
return False
|
|
if name in self._optional_components and value[0] is None:
|
|
return False
|
|
return True
|
|
|
|
model_index_dict = {k: v for k, v in model_index_dict.items() if is_saveable_module(k, v)}
|
|
|
|
for pipeline_component_name in model_index_dict.keys():
|
|
sub_model = getattr(self, pipeline_component_name)
|
|
model_cls = sub_model.__class__
|
|
|
|
save_method_name = None
|
|
# search for the model's base class in LOADABLE_CLASSES
|
|
for library_name, library_classes in LOADABLE_CLASSES.items():
|
|
library = importlib.import_module(library_name)
|
|
for base_class, save_load_methods in library_classes.items():
|
|
class_candidate = getattr(library, base_class, None)
|
|
if class_candidate is not None and issubclass(model_cls, class_candidate):
|
|
# if we found a suitable base class in LOADABLE_CLASSES then grab its save method
|
|
save_method_name = save_load_methods[0]
|
|
break
|
|
if save_method_name is not None:
|
|
break
|
|
|
|
save_method = getattr(sub_model, save_method_name)
|
|
|
|
# Call the save method with the argument safe_serialization only if it's supported
|
|
save_method_signature = inspect.signature(save_method)
|
|
save_method_accept_safe = "safe_serialization" in save_method_signature.parameters
|
|
if save_method_accept_safe:
|
|
save_method(
|
|
os.path.join(save_directory, pipeline_component_name), safe_serialization=safe_serialization
|
|
)
|
|
else:
|
|
save_method(os.path.join(save_directory, pipeline_component_name))
|
|
|
|
def to(self, torch_device: Optional[Union[str, torch.device]] = None):
|
|
if torch_device is None:
|
|
return self
|
|
|
|
module_names, _, _ = self.extract_init_dict(dict(self.config))
|
|
for name in module_names.keys():
|
|
module = getattr(self, name)
|
|
if isinstance(module, torch.nn.Module):
|
|
if module.dtype == torch.float16 and str(torch_device) in ["cpu"]:
|
|
logger.warning(
|
|
"Pipelines loaded with `torch_dtype=torch.float16` cannot run with `cpu` device. It"
|
|
" is not recommended to move them to `cpu` as running them will fail. Please make"
|
|
" sure to use an accelerator to run the pipeline in inference, due to the lack of"
|
|
" support for`float16` operations on this device in PyTorch. Please, remove the"
|
|
" `torch_dtype=torch.float16` argument, or use another device for inference."
|
|
)
|
|
module.to(torch_device)
|
|
return self
|
|
|
|
@property
|
|
def device(self) -> torch.device:
|
|
r"""
|
|
Returns:
|
|
`torch.device`: The torch device on which the pipeline is located.
|
|
"""
|
|
module_names, _, _ = self.extract_init_dict(dict(self.config))
|
|
for name in module_names.keys():
|
|
module = getattr(self, name)
|
|
if isinstance(module, torch.nn.Module):
|
|
return module.device
|
|
return torch.device("cpu")
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
|
|
r"""
|
|
Instantiate a PyTorch diffusion pipeline from pre-trained pipeline weights.
|
|
|
|
The pipeline is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated).
|
|
|
|
The warning *Weights from XXX not initialized from pretrained model* means that the weights of XXX do not come
|
|
pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning
|
|
task.
|
|
|
|
The warning *Weights from XXX not used in YYY* means that the layer XXX is not used by YYY, therefore those
|
|
weights are discarded.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
|
|
Can be either:
|
|
|
|
- A string, the *repo id* of a pretrained pipeline hosted inside a model repo on
|
|
https://huggingface.co/ Valid repo ids have to be located under a user or organization name, like
|
|
`CompVis/ldm-text2im-large-256`.
|
|
- A path to a *directory* containing pipeline weights saved using
|
|
[`~DiffusionPipeline.save_pretrained`], e.g., `./my_pipeline_directory/`.
|
|
torch_dtype (`str` or `torch.dtype`, *optional*):
|
|
Override the default `torch.dtype` and load the model under this dtype. If `"auto"` is passed the dtype
|
|
will be automatically derived from the model's weights.
|
|
custom_pipeline (`str`, *optional*):
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental feature and is likely to change in the future.
|
|
|
|
</Tip>
|
|
|
|
Can be either:
|
|
|
|
- A string, the *repo id* of a custom pipeline hosted inside a model repo on
|
|
https://huggingface.co/. Valid repo ids have to be located under a user or organization name,
|
|
like `hf-internal-testing/diffusers-dummy-pipeline`.
|
|
|
|
<Tip>
|
|
|
|
It is required that the model repo has a file, called `pipeline.py` that defines the custom
|
|
pipeline.
|
|
|
|
</Tip>
|
|
|
|
- A string, the *file name* of a community pipeline hosted on GitHub under
|
|
https://github.com/huggingface/diffusers/tree/main/examples/community. Valid file names have to
|
|
match exactly the file name without `.py` located under the above link, *e.g.*
|
|
`clip_guided_stable_diffusion`.
|
|
|
|
<Tip>
|
|
|
|
Community pipelines are always loaded from the current `main` branch of GitHub.
|
|
|
|
</Tip>
|
|
|
|
- A path to a *directory* containing a custom pipeline, e.g., `./my_pipeline_directory/`.
|
|
|
|
<Tip>
|
|
|
|
It is required that the directory has a file, called `pipeline.py` that defines the custom
|
|
pipeline.
|
|
|
|
</Tip>
|
|
|
|
For more information on how to load and create custom pipelines, please have a look at [Loading and
|
|
Adding Custom
|
|
Pipelines](https://huggingface.co/docs/diffusers/using-diffusers/custom_pipeline_overview)
|
|
|
|
torch_dtype (`str` or `torch.dtype`, *optional*):
|
|
force_download (`bool`, *optional*, defaults to `False`):
|
|
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
|
|
cached versions if they exist.
|
|
resume_download (`bool`, *optional*, defaults to `False`):
|
|
Whether or not to delete incompletely received files. Will attempt to resume the download if such a
|
|
file exists.
|
|
proxies (`Dict[str, str]`, *optional*):
|
|
A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
|
|
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
|
|
output_loading_info(`bool`, *optional*, defaults to `False`):
|
|
Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
|
|
local_files_only(`bool`, *optional*, defaults to `False`):
|
|
Whether or not to only look at local files (i.e., do not try to download the model).
|
|
use_auth_token (`str` or *bool*, *optional*):
|
|
The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
|
|
when running `huggingface-cli login` (stored in `~/.huggingface`).
|
|
revision (`str`, *optional*, defaults to `"main"`):
|
|
The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
|
|
git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
|
|
identifier allowed by git.
|
|
custom_revision (`str`, *optional*, defaults to `"main"` when loading from the Hub and to local version of `diffusers` when loading from GitHub):
|
|
The specific model version to use. It can be a branch name, a tag name, or a commit id similar to
|
|
`revision` when loading a custom pipeline from the Hub. It can be a diffusers version when loading a
|
|
custom pipeline from GitHub.
|
|
mirror (`str`, *optional*):
|
|
Mirror source to accelerate downloads in China. If you are from China and have an accessibility
|
|
problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety.
|
|
Please refer to the mirror site for more information. specify the folder name here.
|
|
device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
|
|
A map that specifies where each submodule should go. It doesn't need to be refined to each
|
|
parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the
|
|
same device.
|
|
|
|
To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For
|
|
more information about each option see [designing a device
|
|
map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map).
|
|
low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
|
|
Speed up model loading by not initializing the weights and only loading the pre-trained weights. This
|
|
also tries to not use more than 1x model size in CPU memory (including peak memory) while loading the
|
|
model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch,
|
|
setting this argument to `True` will raise an error.
|
|
return_cached_folder (`bool`, *optional*, defaults to `False`):
|
|
If set to `True`, path to downloaded cached folder will be returned in addition to loaded pipeline.
|
|
kwargs (remaining dictionary of keyword arguments, *optional*):
|
|
Can be used to overwrite load - and saveable variables - *i.e.* the pipeline components - of the
|
|
specific pipeline class. The overwritten components are then directly passed to the pipelines
|
|
`__init__` method. See example below for more information.
|
|
|
|
<Tip>
|
|
|
|
It is required to be logged in (`huggingface-cli login`) when you want to use private or [gated
|
|
models](https://huggingface.co/docs/hub/models-gated#gated-models), *e.g.* `"runwayml/stable-diffusion-v1-5"`
|
|
|
|
</Tip>
|
|
|
|
<Tip>
|
|
|
|
Activate the special ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use
|
|
this method in a firewalled environment.
|
|
|
|
</Tip>
|
|
|
|
Examples:
|
|
|
|
```py
|
|
>>> from diffusers import DiffusionPipeline
|
|
|
|
>>> # Download pipeline from huggingface.co and cache.
|
|
>>> pipeline = DiffusionPipeline.from_pretrained("CompVis/ldm-text2im-large-256")
|
|
|
|
>>> # Download pipeline that requires an authorization token
|
|
>>> # For more information on access tokens, please refer to this section
|
|
>>> # of the documentation](https://huggingface.co/docs/hub/security-tokens)
|
|
>>> pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
|
|
|
|
>>> # Use a different scheduler
|
|
>>> from diffusers import LMSDiscreteScheduler
|
|
|
|
>>> scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config)
|
|
>>> pipeline.scheduler = scheduler
|
|
```
|
|
"""
|
|
cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)
|
|
resume_download = kwargs.pop("resume_download", False)
|
|
force_download = kwargs.pop("force_download", False)
|
|
proxies = kwargs.pop("proxies", None)
|
|
local_files_only = kwargs.pop("local_files_only", False)
|
|
use_auth_token = kwargs.pop("use_auth_token", None)
|
|
revision = kwargs.pop("revision", None)
|
|
torch_dtype = kwargs.pop("torch_dtype", None)
|
|
custom_pipeline = kwargs.pop("custom_pipeline", None)
|
|
custom_revision = kwargs.pop("custom_revision", None)
|
|
provider = kwargs.pop("provider", None)
|
|
sess_options = kwargs.pop("sess_options", None)
|
|
device_map = kwargs.pop("device_map", None)
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT)
|
|
return_cached_folder = kwargs.pop("return_cached_folder", False)
|
|
|
|
# 1. Download the checkpoints and configs
|
|
# use snapshot download here to get it working from from_pretrained
|
|
if not os.path.isdir(pretrained_model_name_or_path):
|
|
config_dict = cls.load_config(
|
|
pretrained_model_name_or_path,
|
|
cache_dir=cache_dir,
|
|
resume_download=resume_download,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
local_files_only=local_files_only,
|
|
use_auth_token=use_auth_token,
|
|
revision=revision,
|
|
)
|
|
# make sure we only download sub-folders and `diffusers` filenames
|
|
folder_names = [k for k in config_dict.keys() if not k.startswith("_")]
|
|
allow_patterns = [os.path.join(k, "*") for k in folder_names]
|
|
allow_patterns += [WEIGHTS_NAME, SCHEDULER_CONFIG_NAME, CONFIG_NAME, ONNX_WEIGHTS_NAME, cls.config_name]
|
|
|
|
# make sure we don't download flax weights
|
|
ignore_patterns = ["*.msgpack"]
|
|
|
|
if custom_pipeline is not None:
|
|
allow_patterns += [CUSTOM_PIPELINE_FILE_NAME]
|
|
|
|
if cls != DiffusionPipeline:
|
|
requested_pipeline_class = cls.__name__
|
|
else:
|
|
requested_pipeline_class = config_dict.get("_class_name", cls.__name__)
|
|
user_agent = {"pipeline_class": requested_pipeline_class}
|
|
if custom_pipeline is not None and not custom_pipeline.endswith(".py"):
|
|
user_agent["custom_pipeline"] = custom_pipeline
|
|
|
|
user_agent = http_user_agent(user_agent)
|
|
|
|
if is_safetensors_available():
|
|
info = model_info(
|
|
pretrained_model_name_or_path,
|
|
use_auth_token=use_auth_token,
|
|
revision=revision,
|
|
)
|
|
if is_safetensors_compatible(info):
|
|
ignore_patterns.append("*.bin")
|
|
|
|
# download all allow_patterns
|
|
cached_folder = snapshot_download(
|
|
pretrained_model_name_or_path,
|
|
cache_dir=cache_dir,
|
|
resume_download=resume_download,
|
|
proxies=proxies,
|
|
local_files_only=local_files_only,
|
|
use_auth_token=use_auth_token,
|
|
revision=revision,
|
|
allow_patterns=allow_patterns,
|
|
ignore_patterns=ignore_patterns,
|
|
user_agent=user_agent,
|
|
)
|
|
else:
|
|
cached_folder = pretrained_model_name_or_path
|
|
|
|
config_dict = cls.load_config(cached_folder)
|
|
|
|
# 2. Load the pipeline class, if using custom module then load it from the hub
|
|
# if we load from explicit class, let's use it
|
|
if custom_pipeline is not None:
|
|
if custom_pipeline.endswith(".py"):
|
|
path = Path(custom_pipeline)
|
|
# decompose into folder & file
|
|
file_name = path.name
|
|
custom_pipeline = path.parent.absolute()
|
|
else:
|
|
file_name = CUSTOM_PIPELINE_FILE_NAME
|
|
|
|
pipeline_class = get_class_from_dynamic_module(
|
|
custom_pipeline, module_file=file_name, cache_dir=cache_dir, revision=custom_revision
|
|
)
|
|
elif cls != DiffusionPipeline:
|
|
pipeline_class = cls
|
|
else:
|
|
diffusers_module = importlib.import_module(cls.__module__.split(".")[0])
|
|
pipeline_class = getattr(diffusers_module, config_dict["_class_name"])
|
|
|
|
# To be removed in 1.0.0
|
|
if pipeline_class.__name__ == "StableDiffusionInpaintPipeline" and version.parse(
|
|
version.parse(config_dict["_diffusers_version"]).base_version
|
|
) <= version.parse("0.5.1"):
|
|
from diffusers import StableDiffusionInpaintPipeline, StableDiffusionInpaintPipelineLegacy
|
|
|
|
pipeline_class = StableDiffusionInpaintPipelineLegacy
|
|
|
|
deprecation_message = (
|
|
"You are using a legacy checkpoint for inpainting with Stable Diffusion, therefore we are loading the"
|
|
f" {StableDiffusionInpaintPipelineLegacy} class instead of {StableDiffusionInpaintPipeline}. For"
|
|
" better inpainting results, we strongly suggest using Stable Diffusion's official inpainting"
|
|
" checkpoint: https://huggingface.co/runwayml/stable-diffusion-inpainting instead or adapting your"
|
|
f" checkpoint {pretrained_model_name_or_path} to the format of"
|
|
" https://huggingface.co/runwayml/stable-diffusion-inpainting. Note that we do not actively maintain"
|
|
" the {StableDiffusionInpaintPipelineLegacy} class and will likely remove it in version 1.0.0."
|
|
)
|
|
deprecate("StableDiffusionInpaintPipelineLegacy", "1.0.0", deprecation_message, standard_warn=False)
|
|
|
|
# some modules can be passed directly to the init
|
|
# in this case they are already instantiated in `kwargs`
|
|
# extract them here
|
|
expected_modules, optional_kwargs = cls._get_signature_keys(pipeline_class)
|
|
passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs}
|
|
passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs}
|
|
|
|
init_dict, unused_kwargs, _ = pipeline_class.extract_init_dict(config_dict, **kwargs)
|
|
|
|
# define init kwargs
|
|
init_kwargs = {k: init_dict.pop(k) for k in optional_kwargs if k in init_dict}
|
|
init_kwargs = {**init_kwargs, **passed_pipe_kwargs}
|
|
|
|
# remove `null` components
|
|
def load_module(name, value):
|
|
if value[0] is None:
|
|
return False
|
|
if name in passed_class_obj and passed_class_obj[name] is None:
|
|
return False
|
|
return True
|
|
|
|
init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)}
|
|
|
|
if len(unused_kwargs) > 0:
|
|
logger.warning(
|
|
f"Keyword arguments {unused_kwargs} are not expected by {pipeline_class.__name__} and will be ignored."
|
|
)
|
|
|
|
if low_cpu_mem_usage and not is_accelerate_available():
|
|
low_cpu_mem_usage = False
|
|
logger.warning(
|
|
"Cannot initialize model with low cpu memory usage because `accelerate` was not found in the"
|
|
" environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install"
|
|
" `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip"
|
|
" install accelerate\n```\n."
|
|
)
|
|
|
|
if device_map is not None and not is_torch_version(">=", "1.9.0"):
|
|
raise NotImplementedError(
|
|
"Loading and dispatching requires torch >= 1.9.0. Please either update your PyTorch version or set"
|
|
" `device_map=None`."
|
|
)
|
|
|
|
if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"):
|
|
raise NotImplementedError(
|
|
"Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set"
|
|
" `low_cpu_mem_usage=False`."
|
|
)
|
|
|
|
if low_cpu_mem_usage is False and device_map is not None:
|
|
raise ValueError(
|
|
f"You cannot set `low_cpu_mem_usage` to False while using device_map={device_map} for loading and"
|
|
" dispatching. Please make sure to set `low_cpu_mem_usage=True`."
|
|
)
|
|
|
|
# import it here to avoid circular import
|
|
from diffusers import pipelines
|
|
|
|
# 3. Load each module in the pipeline
|
|
for name, (library_name, class_name) in init_dict.items():
|
|
# 3.1 - now that JAX/Flax is an official framework of the library, we might load from Flax names
|
|
if class_name.startswith("Flax"):
|
|
class_name = class_name[4:]
|
|
|
|
is_pipeline_module = hasattr(pipelines, library_name)
|
|
loaded_sub_model = None
|
|
|
|
# if the model is in a pipeline module, then we load it from the pipeline
|
|
if name in passed_class_obj:
|
|
# 1. check that passed_class_obj has correct parent class
|
|
if not is_pipeline_module:
|
|
library = importlib.import_module(library_name)
|
|
class_obj = getattr(library, class_name)
|
|
importable_classes = LOADABLE_CLASSES[library_name]
|
|
class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()}
|
|
|
|
expected_class_obj = None
|
|
for class_name, class_candidate in class_candidates.items():
|
|
if class_candidate is not None and issubclass(class_obj, class_candidate):
|
|
expected_class_obj = class_candidate
|
|
|
|
if not issubclass(passed_class_obj[name].__class__, expected_class_obj):
|
|
raise ValueError(
|
|
f"{passed_class_obj[name]} is of type: {type(passed_class_obj[name])}, but should be"
|
|
f" {expected_class_obj}"
|
|
)
|
|
else:
|
|
logger.warning(
|
|
f"You have passed a non-standard module {passed_class_obj[name]}. We cannot verify whether it"
|
|
" has the correct type"
|
|
)
|
|
|
|
# set passed class object
|
|
loaded_sub_model = passed_class_obj[name]
|
|
elif is_pipeline_module:
|
|
pipeline_module = getattr(pipelines, library_name)
|
|
class_obj = getattr(pipeline_module, class_name)
|
|
importable_classes = ALL_IMPORTABLE_CLASSES
|
|
class_candidates = {c: class_obj for c in importable_classes.keys()}
|
|
else:
|
|
# else we just import it from the library.
|
|
library = importlib.import_module(library_name)
|
|
|
|
class_obj = getattr(library, class_name)
|
|
importable_classes = LOADABLE_CLASSES[library_name]
|
|
class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()}
|
|
|
|
if loaded_sub_model is None:
|
|
load_method_name = None
|
|
for class_name, class_candidate in class_candidates.items():
|
|
if class_candidate is not None and issubclass(class_obj, class_candidate):
|
|
load_method_name = importable_classes[class_name][1]
|
|
|
|
if load_method_name is None:
|
|
none_module = class_obj.__module__
|
|
is_dummy_path = none_module.startswith(DUMMY_MODULES_FOLDER) or none_module.startswith(
|
|
TRANSFORMERS_DUMMY_MODULES_FOLDER
|
|
)
|
|
if is_dummy_path and "dummy" in none_module:
|
|
# call class_obj for nice error message of missing requirements
|
|
class_obj()
|
|
|
|
raise ValueError(
|
|
f"The component {class_obj} of {pipeline_class} cannot be loaded as it does not seem to have"
|
|
f" any of the loading methods defined in {ALL_IMPORTABLE_CLASSES}."
|
|
)
|
|
|
|
load_method = getattr(class_obj, load_method_name)
|
|
loading_kwargs = {}
|
|
|
|
if issubclass(class_obj, torch.nn.Module):
|
|
loading_kwargs["torch_dtype"] = torch_dtype
|
|
if issubclass(class_obj, diffusers.OnnxRuntimeModel):
|
|
loading_kwargs["provider"] = provider
|
|
loading_kwargs["sess_options"] = sess_options
|
|
|
|
is_diffusers_model = issubclass(class_obj, diffusers.ModelMixin)
|
|
is_transformers_model = (
|
|
is_transformers_available()
|
|
and issubclass(class_obj, PreTrainedModel)
|
|
and version.parse(version.parse(transformers.__version__).base_version) >= version.parse("4.20.0")
|
|
)
|
|
|
|
# When loading a transformers model, if the device_map is None, the weights will be initialized as opposed to diffusers.
|
|
# To make default loading faster we set the `low_cpu_mem_usage=low_cpu_mem_usage` flag which is `True` by default.
|
|
# This makes sure that the weights won't be initialized which significantly speeds up loading.
|
|
if is_diffusers_model or is_transformers_model:
|
|
loading_kwargs["device_map"] = device_map
|
|
loading_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage
|
|
|
|
# check if the module is in a subdirectory
|
|
if os.path.isdir(os.path.join(cached_folder, name)):
|
|
loaded_sub_model = load_method(os.path.join(cached_folder, name), **loading_kwargs)
|
|
else:
|
|
# else load from the root directory
|
|
loaded_sub_model = load_method(cached_folder, **loading_kwargs)
|
|
|
|
init_kwargs[name] = loaded_sub_model # UNet(...), # DiffusionSchedule(...)
|
|
|
|
# 4. Potentially add passed objects if expected
|
|
missing_modules = set(expected_modules) - set(init_kwargs.keys())
|
|
passed_modules = list(passed_class_obj.keys())
|
|
optional_modules = pipeline_class._optional_components
|
|
if len(missing_modules) > 0 and missing_modules <= set(passed_modules + optional_modules):
|
|
for module in missing_modules:
|
|
init_kwargs[module] = passed_class_obj.get(module, None)
|
|
elif len(missing_modules) > 0:
|
|
passed_modules = set(list(init_kwargs.keys()) + list(passed_class_obj.keys())) - optional_kwargs
|
|
raise ValueError(
|
|
f"Pipeline {pipeline_class} expected {expected_modules}, but only {passed_modules} were passed."
|
|
)
|
|
|
|
# 5. Instantiate the pipeline
|
|
model = pipeline_class(**init_kwargs)
|
|
|
|
if return_cached_folder:
|
|
return model, cached_folder
|
|
return model
|
|
|
|
@staticmethod
|
|
def _get_signature_keys(obj):
|
|
parameters = inspect.signature(obj.__init__).parameters
|
|
required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty}
|
|
optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty})
|
|
expected_modules = set(required_parameters.keys()) - set(["self"])
|
|
return expected_modules, optional_parameters
|
|
|
|
@property
|
|
def components(self) -> Dict[str, Any]:
|
|
r"""
|
|
|
|
The `self.components` property can be useful to run different pipelines with the same weights and
|
|
configurations to not have to re-allocate memory.
|
|
|
|
Examples:
|
|
|
|
```py
|
|
>>> from diffusers import (
|
|
... StableDiffusionPipeline,
|
|
... StableDiffusionImg2ImgPipeline,
|
|
... StableDiffusionInpaintPipeline,
|
|
... )
|
|
|
|
>>> text2img = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
|
|
>>> img2img = StableDiffusionImg2ImgPipeline(**text2img.components)
|
|
>>> inpaint = StableDiffusionInpaintPipeline(**text2img.components)
|
|
```
|
|
|
|
Returns:
|
|
A dictionaly containing all the modules needed to initialize the pipeline.
|
|
"""
|
|
expected_modules, optional_parameters = self._get_signature_keys(self)
|
|
components = {
|
|
k: getattr(self, k) for k in self.config.keys() if not k.startswith("_") and k not in optional_parameters
|
|
}
|
|
|
|
if set(components.keys()) != expected_modules:
|
|
raise ValueError(
|
|
f"{self} has been incorrectly initialized or {self.__class__} is incorrectly implemented. Expected"
|
|
f" {expected_modules} to be defined, but {components} are defined."
|
|
)
|
|
|
|
return components
|
|
|
|
@staticmethod
|
|
def numpy_to_pil(images):
|
|
"""
|
|
Convert a numpy image or a batch of images to a PIL image.
|
|
"""
|
|
if images.ndim == 3:
|
|
images = images[None, ...]
|
|
images = (images * 255).round().astype("uint8")
|
|
if images.shape[-1] == 1:
|
|
# special case for grayscale (single channel) images
|
|
pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
|
|
else:
|
|
pil_images = [Image.fromarray(image) for image in images]
|
|
|
|
return pil_images
|
|
|
|
def progress_bar(self, iterable=None, total=None):
|
|
if not hasattr(self, "_progress_bar_config"):
|
|
self._progress_bar_config = {}
|
|
elif not isinstance(self._progress_bar_config, dict):
|
|
raise ValueError(
|
|
f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}."
|
|
)
|
|
|
|
if iterable is not None:
|
|
return tqdm(iterable, **self._progress_bar_config)
|
|
elif total is not None:
|
|
return tqdm(total=total, **self._progress_bar_config)
|
|
else:
|
|
raise ValueError("Either `total` or `iterable` has to be defined.")
|
|
|
|
def set_progress_bar_config(self, **kwargs):
|
|
self._progress_bar_config = kwargs
|
|
|
|
def enable_xformers_memory_efficient_attention(self):
|
|
r"""
|
|
Enable memory efficient attention as implemented in xformers.
|
|
|
|
When this option is enabled, you should observe lower GPU memory usage and a potential speed up at inference
|
|
time. Speed up at training time is not guaranteed.
|
|
|
|
Warning: When Memory Efficient Attention and Sliced attention are both enabled, the Memory Efficient Attention
|
|
is used.
|
|
"""
|
|
self.set_use_memory_efficient_attention_xformers(True)
|
|
|
|
def disable_xformers_memory_efficient_attention(self):
|
|
r"""
|
|
Disable memory efficient attention as implemented in xformers.
|
|
"""
|
|
self.set_use_memory_efficient_attention_xformers(False)
|
|
|
|
def set_use_memory_efficient_attention_xformers(self, valid: bool) -> None:
|
|
# Recursively walk through all the children.
|
|
# Any children which exposes the set_use_memory_efficient_attention_xformers method
|
|
# gets the message
|
|
def fn_recursive_set_mem_eff(module: torch.nn.Module):
|
|
if hasattr(module, "set_use_memory_efficient_attention_xformers"):
|
|
module.set_use_memory_efficient_attention_xformers(valid)
|
|
|
|
for child in module.children():
|
|
fn_recursive_set_mem_eff(child)
|
|
|
|
module_names, _, _ = self.extract_init_dict(dict(self.config))
|
|
for module_name in module_names:
|
|
module = getattr(self, module_name)
|
|
if isinstance(module, torch.nn.Module):
|
|
fn_recursive_set_mem_eff(module)
|
|
|
|
def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"):
|
|
r"""
|
|
Enable sliced attention computation.
|
|
|
|
When this option is enabled, the attention module will split the input tensor in slices, to compute attention
|
|
in several steps. This is useful to save some memory in exchange for a small speed decrease.
|
|
|
|
Args:
|
|
slice_size (`str` or `int`, *optional*, defaults to `"auto"`):
|
|
When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If
|
|
`"max"`, maxium amount of memory will be saved by running only one slice at a time. If a number is
|
|
provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
|
|
must be a multiple of `slice_size`.
|
|
"""
|
|
self.set_attention_slice(slice_size)
|
|
|
|
def disable_attention_slicing(self):
|
|
r"""
|
|
Disable sliced attention computation. If `enable_attention_slicing` was previously invoked, this method will go
|
|
back to computing attention in one step.
|
|
"""
|
|
# set slice_size = `None` to disable `attention slicing`
|
|
self.enable_attention_slicing(None)
|
|
|
|
def set_attention_slice(self, slice_size: Optional[int]):
|
|
module_names, _, _ = self.extract_init_dict(dict(self.config))
|
|
for module_name in module_names:
|
|
module = getattr(self, module_name)
|
|
if isinstance(module, torch.nn.Module) and hasattr(module, "set_attention_slice"):
|
|
module.set_attention_slice(slice_size)
|