Spaces:
Running
on
Zero
Running
on
Zero
# Copyright 2023 The HuggingFace Team. 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 inspect | |
import os | |
from collections import defaultdict | |
from contextlib import nullcontext | |
from functools import partial | |
from typing import Callable, Dict, List, Optional, Union | |
import safetensors | |
import torch | |
import torch.nn.functional as F | |
from huggingface_hub.utils import validate_hf_hub_args | |
from torch import nn | |
from ..models.embeddings import ImageProjection, IPAdapterFullImageProjection, IPAdapterPlusImageProjection | |
from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, load_model_dict_into_meta | |
from ..utils import ( | |
USE_PEFT_BACKEND, | |
_get_model_file, | |
delete_adapter_layers, | |
is_accelerate_available, | |
logging, | |
set_adapter_layers, | |
set_weights_and_activate_adapters, | |
) | |
from .utils import AttnProcsLayers | |
if is_accelerate_available(): | |
from accelerate import init_empty_weights | |
from accelerate.hooks import AlignDevicesHook, CpuOffload, remove_hook_from_module | |
logger = logging.get_logger(__name__) | |
TEXT_ENCODER_NAME = "text_encoder" | |
UNET_NAME = "unet" | |
LORA_WEIGHT_NAME = "pytorch_lora_weights.bin" | |
LORA_WEIGHT_NAME_SAFE = "pytorch_lora_weights.safetensors" | |
CUSTOM_DIFFUSION_WEIGHT_NAME = "pytorch_custom_diffusion_weights.bin" | |
CUSTOM_DIFFUSION_WEIGHT_NAME_SAFE = "pytorch_custom_diffusion_weights.safetensors" | |
class UNet2DConditionLoadersMixin: | |
""" | |
Load LoRA layers into a [`UNet2DCondtionModel`]. | |
""" | |
text_encoder_name = TEXT_ENCODER_NAME | |
unet_name = UNET_NAME | |
def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs): | |
r""" | |
Load pretrained attention processor layers into [`UNet2DConditionModel`]. Attention processor layers have to be | |
defined in | |
[`attention_processor.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py) | |
and be a `torch.nn.Module` class. | |
Parameters: | |
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): | |
Can be either: | |
- A string, the model id (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on | |
the Hub. | |
- A path to a directory (for example `./my_model_directory`) containing the model weights saved | |
with [`ModelMixin.save_pretrained`]. | |
- A [torch state | |
dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict). | |
cache_dir (`Union[str, os.PathLike]`, *optional*): | |
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache | |
is not used. | |
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 resume downloading the model weights and configuration files. If set to `False`, any | |
incompletely downloaded files are deleted. | |
proxies (`Dict[str, str]`, *optional*): | |
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128', | |
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. | |
local_files_only (`bool`, *optional*, defaults to `False`): | |
Whether to only load local model weights and configuration files or not. If set to `True`, the model | |
won't be downloaded from the Hub. | |
token (`str` or *bool*, *optional*): | |
The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from | |
`diffusers-cli login` (stored in `~/.huggingface`) is used. | |
low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`): | |
Speed up model loading only loading the pretrained weights and not initializing the weights. This also | |
tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. | |
Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this | |
argument to `True` will raise an error. | |
revision (`str`, *optional*, defaults to `"main"`): | |
The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier | |
allowed by Git. | |
subfolder (`str`, *optional*, defaults to `""`): | |
The subfolder location of a model file within a larger model repository on the Hub or locally. | |
mirror (`str`, *optional*): | |
Mirror source to resolve accessibility issues if you’re downloading a model in China. We do not | |
guarantee the timeliness or safety of the source, and you should refer to the mirror site for more | |
information. | |
Example: | |
```py | |
from diffusers import AutoPipelineForText2Image | |
import torch | |
pipeline = AutoPipelineForText2Image.from_pretrained( | |
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 | |
).to("cuda") | |
pipeline.unet.load_attn_procs( | |
"jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_name="cinematic" | |
) | |
``` | |
""" | |
from ..models.attention_processor import CustomDiffusionAttnProcessor | |
from ..models.lora import LoRACompatibleConv, LoRACompatibleLinear, LoRAConv2dLayer, LoRALinearLayer | |
cache_dir = kwargs.pop("cache_dir", None) | |
force_download = kwargs.pop("force_download", False) | |
resume_download = kwargs.pop("resume_download", False) | |
proxies = kwargs.pop("proxies", None) | |
local_files_only = kwargs.pop("local_files_only", None) | |
token = kwargs.pop("token", None) | |
revision = kwargs.pop("revision", None) | |
subfolder = kwargs.pop("subfolder", None) | |
weight_name = kwargs.pop("weight_name", None) | |
use_safetensors = kwargs.pop("use_safetensors", None) | |
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) | |
# This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script. | |
# See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning | |
network_alphas = kwargs.pop("network_alphas", None) | |
_pipeline = kwargs.pop("_pipeline", None) | |
is_network_alphas_none = network_alphas is None | |
allow_pickle = False | |
if use_safetensors is None: | |
use_safetensors = True | |
allow_pickle = True | |
user_agent = { | |
"file_type": "attn_procs_weights", | |
"framework": "pytorch", | |
} | |
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." | |
) | |
model_file = None | |
if not isinstance(pretrained_model_name_or_path_or_dict, dict): | |
# Let's first try to load .safetensors weights | |
if (use_safetensors and weight_name is None) or ( | |
weight_name is not None and weight_name.endswith(".safetensors") | |
): | |
try: | |
model_file = _get_model_file( | |
pretrained_model_name_or_path_or_dict, | |
weights_name=weight_name or LORA_WEIGHT_NAME_SAFE, | |
cache_dir=cache_dir, | |
force_download=force_download, | |
resume_download=resume_download, | |
proxies=proxies, | |
local_files_only=local_files_only, | |
token=token, | |
revision=revision, | |
subfolder=subfolder, | |
user_agent=user_agent, | |
) | |
state_dict = safetensors.torch.load_file(model_file, device="cpu") | |
except IOError as e: | |
if not allow_pickle: | |
raise e | |
# try loading non-safetensors weights | |
pass | |
if model_file is None: | |
model_file = _get_model_file( | |
pretrained_model_name_or_path_or_dict, | |
weights_name=weight_name or LORA_WEIGHT_NAME, | |
cache_dir=cache_dir, | |
force_download=force_download, | |
resume_download=resume_download, | |
proxies=proxies, | |
local_files_only=local_files_only, | |
token=token, | |
revision=revision, | |
subfolder=subfolder, | |
user_agent=user_agent, | |
) | |
state_dict = torch.load(model_file, map_location="cpu") | |
else: | |
state_dict = pretrained_model_name_or_path_or_dict | |
# fill attn processors | |
lora_layers_list = [] | |
is_lora = all(("lora" in k or k.endswith(".alpha")) for k in state_dict.keys()) and not USE_PEFT_BACKEND | |
is_custom_diffusion = any("custom_diffusion" in k for k in state_dict.keys()) | |
if is_lora: | |
# correct keys | |
state_dict, network_alphas = self.convert_state_dict_legacy_attn_format(state_dict, network_alphas) | |
if network_alphas is not None: | |
network_alphas_keys = list(network_alphas.keys()) | |
used_network_alphas_keys = set() | |
lora_grouped_dict = defaultdict(dict) | |
mapped_network_alphas = {} | |
all_keys = list(state_dict.keys()) | |
for key in all_keys: | |
value = state_dict.pop(key) | |
attn_processor_key, sub_key = ".".join(key.split(".")[:-3]), ".".join(key.split(".")[-3:]) | |
lora_grouped_dict[attn_processor_key][sub_key] = value | |
# Create another `mapped_network_alphas` dictionary so that we can properly map them. | |
if network_alphas is not None: | |
for k in network_alphas_keys: | |
if k.replace(".alpha", "") in key: | |
mapped_network_alphas.update({attn_processor_key: network_alphas.get(k)}) | |
used_network_alphas_keys.add(k) | |
if not is_network_alphas_none: | |
if len(set(network_alphas_keys) - used_network_alphas_keys) > 0: | |
raise ValueError( | |
f"The `network_alphas` has to be empty at this point but has the following keys \n\n {', '.join(network_alphas.keys())}" | |
) | |
if len(state_dict) > 0: | |
raise ValueError( | |
f"The `state_dict` has to be empty at this point but has the following keys \n\n {', '.join(state_dict.keys())}" | |
) | |
for key, value_dict in lora_grouped_dict.items(): | |
attn_processor = self | |
for sub_key in key.split("."): | |
attn_processor = getattr(attn_processor, sub_key) | |
# Process non-attention layers, which don't have to_{k,v,q,out_proj}_lora layers | |
# or add_{k,v,q,out_proj}_proj_lora layers. | |
rank = value_dict["lora.down.weight"].shape[0] | |
if isinstance(attn_processor, LoRACompatibleConv): | |
in_features = attn_processor.in_channels | |
out_features = attn_processor.out_channels | |
kernel_size = attn_processor.kernel_size | |
ctx = init_empty_weights if low_cpu_mem_usage else nullcontext | |
with ctx(): | |
lora = LoRAConv2dLayer( | |
in_features=in_features, | |
out_features=out_features, | |
rank=rank, | |
kernel_size=kernel_size, | |
stride=attn_processor.stride, | |
padding=attn_processor.padding, | |
network_alpha=mapped_network_alphas.get(key), | |
) | |
elif isinstance(attn_processor, LoRACompatibleLinear): | |
ctx = init_empty_weights if low_cpu_mem_usage else nullcontext | |
with ctx(): | |
lora = LoRALinearLayer( | |
attn_processor.in_features, | |
attn_processor.out_features, | |
rank, | |
mapped_network_alphas.get(key), | |
) | |
else: | |
raise ValueError(f"Module {key} is not a LoRACompatibleConv or LoRACompatibleLinear module.") | |
value_dict = {k.replace("lora.", ""): v for k, v in value_dict.items()} | |
lora_layers_list.append((attn_processor, lora)) | |
if low_cpu_mem_usage: | |
device = next(iter(value_dict.values())).device | |
dtype = next(iter(value_dict.values())).dtype | |
load_model_dict_into_meta(lora, value_dict, device=device, dtype=dtype) | |
else: | |
lora.load_state_dict(value_dict) | |
elif is_custom_diffusion: | |
attn_processors = {} | |
custom_diffusion_grouped_dict = defaultdict(dict) | |
for key, value in state_dict.items(): | |
if len(value) == 0: | |
custom_diffusion_grouped_dict[key] = {} | |
else: | |
if "to_out" in key: | |
attn_processor_key, sub_key = ".".join(key.split(".")[:-3]), ".".join(key.split(".")[-3:]) | |
else: | |
attn_processor_key, sub_key = ".".join(key.split(".")[:-2]), ".".join(key.split(".")[-2:]) | |
custom_diffusion_grouped_dict[attn_processor_key][sub_key] = value | |
for key, value_dict in custom_diffusion_grouped_dict.items(): | |
if len(value_dict) == 0: | |
attn_processors[key] = CustomDiffusionAttnProcessor( | |
train_kv=False, train_q_out=False, hidden_size=None, cross_attention_dim=None | |
) | |
else: | |
cross_attention_dim = value_dict["to_k_custom_diffusion.weight"].shape[1] | |
hidden_size = value_dict["to_k_custom_diffusion.weight"].shape[0] | |
train_q_out = True if "to_q_custom_diffusion.weight" in value_dict else False | |
attn_processors[key] = CustomDiffusionAttnProcessor( | |
train_kv=True, | |
train_q_out=train_q_out, | |
hidden_size=hidden_size, | |
cross_attention_dim=cross_attention_dim, | |
) | |
attn_processors[key].load_state_dict(value_dict) | |
elif USE_PEFT_BACKEND: | |
# In that case we have nothing to do as loading the adapter weights is already handled above by `set_peft_model_state_dict` | |
# on the Unet | |
pass | |
else: | |
raise ValueError( | |
f"{model_file} does not seem to be in the correct format expected by LoRA or Custom Diffusion training." | |
) | |
# <Unsafe code | |
# We can be sure that the following works as it just sets attention processors, lora layers and puts all in the same dtype | |
# Now we remove any existing hooks to | |
is_model_cpu_offload = False | |
is_sequential_cpu_offload = False | |
# For PEFT backend the Unet is already offloaded at this stage as it is handled inside `lora_lora_weights_into_unet` | |
if not USE_PEFT_BACKEND: | |
if _pipeline is not None: | |
for _, component in _pipeline.components.items(): | |
if isinstance(component, nn.Module) and hasattr(component, "_hf_hook"): | |
is_model_cpu_offload = isinstance(getattr(component, "_hf_hook"), CpuOffload) | |
is_sequential_cpu_offload = isinstance(getattr(component, "_hf_hook"), AlignDevicesHook) | |
logger.info( | |
"Accelerate hooks detected. Since you have called `load_lora_weights()`, the previous hooks will be first removed. Then the LoRA parameters will be loaded and the hooks will be applied again." | |
) | |
remove_hook_from_module(component, recurse=is_sequential_cpu_offload) | |
# only custom diffusion needs to set attn processors | |
if is_custom_diffusion: | |
self.set_attn_processor(attn_processors) | |
# set lora layers | |
for target_module, lora_layer in lora_layers_list: | |
target_module.set_lora_layer(lora_layer) | |
self.to(dtype=self.dtype, device=self.device) | |
# Offload back. | |
if is_model_cpu_offload: | |
_pipeline.enable_model_cpu_offload() | |
elif is_sequential_cpu_offload: | |
_pipeline.enable_sequential_cpu_offload() | |
# Unsafe code /> | |
def convert_state_dict_legacy_attn_format(self, state_dict, network_alphas): | |
is_new_lora_format = all( | |
key.startswith(self.unet_name) or key.startswith(self.text_encoder_name) for key in state_dict.keys() | |
) | |
if is_new_lora_format: | |
# Strip the `"unet"` prefix. | |
is_text_encoder_present = any(key.startswith(self.text_encoder_name) for key in state_dict.keys()) | |
if is_text_encoder_present: | |
warn_message = "The state_dict contains LoRA params corresponding to the text encoder which are not being used here. To use both UNet and text encoder related LoRA params, use [`pipe.load_lora_weights()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraLoaderMixin.load_lora_weights)." | |
logger.warn(warn_message) | |
unet_keys = [k for k in state_dict.keys() if k.startswith(self.unet_name)] | |
state_dict = {k.replace(f"{self.unet_name}.", ""): v for k, v in state_dict.items() if k in unet_keys} | |
# change processor format to 'pure' LoRACompatibleLinear format | |
if any("processor" in k.split(".") for k in state_dict.keys()): | |
def format_to_lora_compatible(key): | |
if "processor" not in key.split("."): | |
return key | |
return key.replace(".processor", "").replace("to_out_lora", "to_out.0.lora").replace("_lora", ".lora") | |
state_dict = {format_to_lora_compatible(k): v for k, v in state_dict.items()} | |
if network_alphas is not None: | |
network_alphas = {format_to_lora_compatible(k): v for k, v in network_alphas.items()} | |
return state_dict, network_alphas | |
def save_attn_procs( | |
self, | |
save_directory: Union[str, os.PathLike], | |
is_main_process: bool = True, | |
weight_name: str = None, | |
save_function: Callable = None, | |
safe_serialization: bool = True, | |
**kwargs, | |
): | |
r""" | |
Save attention processor layers to a directory so that it can be reloaded with the | |
[`~loaders.UNet2DConditionLoadersMixin.load_attn_procs`] method. | |
Arguments: | |
save_directory (`str` or `os.PathLike`): | |
Directory to save an attention processor to (will be created if it doesn't exist). | |
is_main_process (`bool`, *optional*, defaults to `True`): | |
Whether the process calling this is the main process or not. Useful during distributed training and you | |
need to call this function on all processes. In this case, set `is_main_process=True` only on the main | |
process to avoid race conditions. | |
save_function (`Callable`): | |
The function to use to save the state dictionary. Useful during distributed training when you need to | |
replace `torch.save` with another method. Can be configured with the environment variable | |
`DIFFUSERS_SAVE_MODE`. | |
safe_serialization (`bool`, *optional*, defaults to `True`): | |
Whether to save the model using `safetensors` or with `pickle`. | |
Example: | |
```py | |
import torch | |
from diffusers import DiffusionPipeline | |
pipeline = DiffusionPipeline.from_pretrained( | |
"CompVis/stable-diffusion-v1-4", | |
torch_dtype=torch.float16, | |
).to("cuda") | |
pipeline.unet.load_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin") | |
pipeline.unet.save_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin") | |
``` | |
""" | |
from ..models.attention_processor import ( | |
CustomDiffusionAttnProcessor, | |
CustomDiffusionAttnProcessor2_0, | |
CustomDiffusionXFormersAttnProcessor, | |
) | |
if os.path.isfile(save_directory): | |
logger.error(f"Provided path ({save_directory}) should be a directory, not a file") | |
return | |
if save_function is None: | |
if safe_serialization: | |
def save_function(weights, filename): | |
return safetensors.torch.save_file(weights, filename, metadata={"format": "pt"}) | |
else: | |
save_function = torch.save | |
os.makedirs(save_directory, exist_ok=True) | |
is_custom_diffusion = any( | |
isinstance( | |
x, | |
(CustomDiffusionAttnProcessor, CustomDiffusionAttnProcessor2_0, CustomDiffusionXFormersAttnProcessor), | |
) | |
for (_, x) in self.attn_processors.items() | |
) | |
if is_custom_diffusion: | |
model_to_save = AttnProcsLayers( | |
{ | |
y: x | |
for (y, x) in self.attn_processors.items() | |
if isinstance( | |
x, | |
( | |
CustomDiffusionAttnProcessor, | |
CustomDiffusionAttnProcessor2_0, | |
CustomDiffusionXFormersAttnProcessor, | |
), | |
) | |
} | |
) | |
state_dict = model_to_save.state_dict() | |
for name, attn in self.attn_processors.items(): | |
if len(attn.state_dict()) == 0: | |
state_dict[name] = {} | |
else: | |
model_to_save = AttnProcsLayers(self.attn_processors) | |
state_dict = model_to_save.state_dict() | |
if weight_name is None: | |
if safe_serialization: | |
weight_name = CUSTOM_DIFFUSION_WEIGHT_NAME_SAFE if is_custom_diffusion else LORA_WEIGHT_NAME_SAFE | |
else: | |
weight_name = CUSTOM_DIFFUSION_WEIGHT_NAME if is_custom_diffusion else LORA_WEIGHT_NAME | |
# Save the model | |
save_function(state_dict, os.path.join(save_directory, weight_name)) | |
logger.info(f"Model weights saved in {os.path.join(save_directory, weight_name)}") | |
def fuse_lora(self, lora_scale=1.0, safe_fusing=False, adapter_names=None): | |
self.lora_scale = lora_scale | |
self._safe_fusing = safe_fusing | |
self.apply(partial(self._fuse_lora_apply, adapter_names=adapter_names)) | |
def _fuse_lora_apply(self, module, adapter_names=None): | |
if not USE_PEFT_BACKEND: | |
if hasattr(module, "_fuse_lora"): | |
module._fuse_lora(self.lora_scale, self._safe_fusing) | |
if adapter_names is not None: | |
raise ValueError( | |
"The `adapter_names` argument is not supported in your environment. Please switch" | |
" to PEFT backend to use this argument by installing latest PEFT and transformers." | |
" `pip install -U peft transformers`" | |
) | |
else: | |
from peft.tuners.tuners_utils import BaseTunerLayer | |
merge_kwargs = {"safe_merge": self._safe_fusing} | |
if isinstance(module, BaseTunerLayer): | |
if self.lora_scale != 1.0: | |
module.scale_layer(self.lora_scale) | |
# For BC with prevous PEFT versions, we need to check the signature | |
# of the `merge` method to see if it supports the `adapter_names` argument. | |
supported_merge_kwargs = list(inspect.signature(module.merge).parameters) | |
if "adapter_names" in supported_merge_kwargs: | |
merge_kwargs["adapter_names"] = adapter_names | |
elif "adapter_names" not in supported_merge_kwargs and adapter_names is not None: | |
raise ValueError( | |
"The `adapter_names` argument is not supported with your PEFT version. Please upgrade" | |
" to the latest version of PEFT. `pip install -U peft`" | |
) | |
module.merge(**merge_kwargs) | |
def unfuse_lora(self): | |
self.apply(self._unfuse_lora_apply) | |
def _unfuse_lora_apply(self, module): | |
if not USE_PEFT_BACKEND: | |
if hasattr(module, "_unfuse_lora"): | |
module._unfuse_lora() | |
else: | |
from peft.tuners.tuners_utils import BaseTunerLayer | |
if isinstance(module, BaseTunerLayer): | |
module.unmerge() | |
def set_adapters( | |
self, | |
adapter_names: Union[List[str], str], | |
weights: Optional[Union[List[float], float]] = None, | |
): | |
""" | |
Set the currently active adapters for use in the UNet. | |
Args: | |
adapter_names (`List[str]` or `str`): | |
The names of the adapters to use. | |
adapter_weights (`Union[List[float], float]`, *optional*): | |
The adapter(s) weights to use with the UNet. If `None`, the weights are set to `1.0` for all the | |
adapters. | |
Example: | |
```py | |
from diffusers import AutoPipelineForText2Image | |
import torch | |
pipeline = AutoPipelineForText2Image.from_pretrained( | |
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 | |
).to("cuda") | |
pipeline.load_lora_weights( | |
"jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_name="cinematic" | |
) | |
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel") | |
pipeline.set_adapters(["cinematic", "pixel"], adapter_weights=[0.5, 0.5]) | |
``` | |
""" | |
if not USE_PEFT_BACKEND: | |
raise ValueError("PEFT backend is required for `set_adapters()`.") | |
adapter_names = [adapter_names] if isinstance(adapter_names, str) else adapter_names | |
if weights is None: | |
weights = [1.0] * len(adapter_names) | |
elif isinstance(weights, float): | |
weights = [weights] * len(adapter_names) | |
if len(adapter_names) != len(weights): | |
raise ValueError( | |
f"Length of adapter names {len(adapter_names)} is not equal to the length of their weights {len(weights)}." | |
) | |
set_weights_and_activate_adapters(self, adapter_names, weights) | |
def disable_lora(self): | |
""" | |
Disable the UNet's active LoRA layers. | |
Example: | |
```py | |
from diffusers import AutoPipelineForText2Image | |
import torch | |
pipeline = AutoPipelineForText2Image.from_pretrained( | |
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 | |
).to("cuda") | |
pipeline.load_lora_weights( | |
"jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_name="cinematic" | |
) | |
pipeline.disable_lora() | |
``` | |
""" | |
if not USE_PEFT_BACKEND: | |
raise ValueError("PEFT backend is required for this method.") | |
set_adapter_layers(self, enabled=False) | |
def enable_lora(self): | |
""" | |
Enable the UNet's active LoRA layers. | |
Example: | |
```py | |
from diffusers import AutoPipelineForText2Image | |
import torch | |
pipeline = AutoPipelineForText2Image.from_pretrained( | |
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 | |
).to("cuda") | |
pipeline.load_lora_weights( | |
"jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_name="cinematic" | |
) | |
pipeline.enable_lora() | |
``` | |
""" | |
if not USE_PEFT_BACKEND: | |
raise ValueError("PEFT backend is required for this method.") | |
set_adapter_layers(self, enabled=True) | |
def delete_adapters(self, adapter_names: Union[List[str], str]): | |
""" | |
Delete an adapter's LoRA layers from the UNet. | |
Args: | |
adapter_names (`Union[List[str], str]`): | |
The names (single string or list of strings) of the adapter to delete. | |
Example: | |
```py | |
from diffusers import AutoPipelineForText2Image | |
import torch | |
pipeline = AutoPipelineForText2Image.from_pretrained( | |
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 | |
).to("cuda") | |
pipeline.load_lora_weights( | |
"jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_names="cinematic" | |
) | |
pipeline.delete_adapters("cinematic") | |
``` | |
""" | |
if not USE_PEFT_BACKEND: | |
raise ValueError("PEFT backend is required for this method.") | |
if isinstance(adapter_names, str): | |
adapter_names = [adapter_names] | |
for adapter_name in adapter_names: | |
delete_adapter_layers(self, adapter_name) | |
# Pop also the corresponding adapter from the config | |
if hasattr(self, "peft_config"): | |
self.peft_config.pop(adapter_name, None) | |
def _convert_ip_adapter_image_proj_to_diffusers(self, state_dict): | |
updated_state_dict = {} | |
image_projection = None | |
if "proj.weight" in state_dict: | |
# IP-Adapter | |
num_image_text_embeds = 4 | |
clip_embeddings_dim = state_dict["proj.weight"].shape[-1] | |
cross_attention_dim = state_dict["proj.weight"].shape[0] // 4 | |
image_projection = ImageProjection( | |
cross_attention_dim=cross_attention_dim, | |
image_embed_dim=clip_embeddings_dim, | |
num_image_text_embeds=num_image_text_embeds, | |
) | |
for key, value in state_dict.items(): | |
diffusers_name = key.replace("proj", "image_embeds") | |
updated_state_dict[diffusers_name] = value | |
elif "proj.3.weight" in state_dict: | |
# IP-Adapter Full | |
clip_embeddings_dim = state_dict["proj.0.weight"].shape[0] | |
cross_attention_dim = state_dict["proj.3.weight"].shape[0] | |
image_projection = IPAdapterFullImageProjection( | |
cross_attention_dim=cross_attention_dim, image_embed_dim=clip_embeddings_dim | |
) | |
for key, value in state_dict.items(): | |
diffusers_name = key.replace("proj.0", "ff.net.0.proj") | |
diffusers_name = diffusers_name.replace("proj.2", "ff.net.2") | |
diffusers_name = diffusers_name.replace("proj.3", "norm") | |
updated_state_dict[diffusers_name] = value | |
else: | |
# IP-Adapter Plus | |
num_image_text_embeds = state_dict["latents"].shape[1] | |
embed_dims = state_dict["proj_in.weight"].shape[1] | |
output_dims = state_dict["proj_out.weight"].shape[0] | |
hidden_dims = state_dict["latents"].shape[2] | |
heads = state_dict["layers.0.0.to_q.weight"].shape[0] // 64 | |
image_projection = IPAdapterPlusImageProjection( | |
embed_dims=embed_dims, | |
output_dims=output_dims, | |
hidden_dims=hidden_dims, | |
heads=heads, | |
num_queries=num_image_text_embeds, | |
) | |
for key, value in state_dict.items(): | |
diffusers_name = key.replace("0.to", "2.to") | |
diffusers_name = diffusers_name.replace("1.0.weight", "3.0.weight") | |
diffusers_name = diffusers_name.replace("1.0.bias", "3.0.bias") | |
diffusers_name = diffusers_name.replace("1.1.weight", "3.1.net.0.proj.weight") | |
diffusers_name = diffusers_name.replace("1.3.weight", "3.1.net.2.weight") | |
if "norm1" in diffusers_name: | |
updated_state_dict[diffusers_name.replace("0.norm1", "0")] = value | |
elif "norm2" in diffusers_name: | |
updated_state_dict[diffusers_name.replace("0.norm2", "1")] = value | |
elif "to_kv" in diffusers_name: | |
v_chunk = value.chunk(2, dim=0) | |
updated_state_dict[diffusers_name.replace("to_kv", "to_k")] = v_chunk[0] | |
updated_state_dict[diffusers_name.replace("to_kv", "to_v")] = v_chunk[1] | |
elif "to_out" in diffusers_name: | |
updated_state_dict[diffusers_name.replace("to_out", "to_out.0")] = value | |
else: | |
updated_state_dict[diffusers_name] = value | |
image_projection.load_state_dict(updated_state_dict) | |
return image_projection | |
def _load_ip_adapter_weights(self, state_dict): | |
from ..models.attention_processor import ( | |
AttnProcessor, | |
AttnProcessor2_0, | |
IPAdapterAttnProcessor, | |
IPAdapterAttnProcessor2_0, | |
) | |
if "proj.weight" in state_dict["image_proj"]: | |
# IP-Adapter | |
num_image_text_embeds = 4 | |
elif "proj.3.weight" in state_dict["image_proj"]: | |
# IP-Adapter Full Face | |
num_image_text_embeds = 257 # 256 CLIP tokens + 1 CLS token | |
else: | |
# IP-Adapter Plus | |
num_image_text_embeds = state_dict["image_proj"]["latents"].shape[1] | |
# Set encoder_hid_proj after loading ip_adapter weights, | |
# because `IPAdapterPlusImageProjection` also has `attn_processors`. | |
self.encoder_hid_proj = None | |
# set ip-adapter cross-attention processors & load state_dict | |
attn_procs = {} | |
key_id = 1 | |
for name in self.attn_processors.keys(): | |
cross_attention_dim = None if name.endswith("attn1.processor") else self.config.cross_attention_dim | |
if name.startswith("mid_block"): | |
hidden_size = self.config.block_out_channels[-1] | |
elif name.startswith("up_blocks"): | |
block_id = int(name[len("up_blocks.")]) | |
hidden_size = list(reversed(self.config.block_out_channels))[block_id] | |
elif name.startswith("down_blocks"): | |
block_id = int(name[len("down_blocks.")]) | |
hidden_size = self.config.block_out_channels[block_id] | |
if cross_attention_dim is None or "motion_modules" in name: | |
attn_processor_class = ( | |
AttnProcessor2_0 if hasattr(F, "scaled_dot_product_attention") else AttnProcessor | |
) | |
attn_procs[name] = attn_processor_class() | |
else: | |
attn_processor_class = ( | |
IPAdapterAttnProcessor2_0 if hasattr(F, "scaled_dot_product_attention") else IPAdapterAttnProcessor | |
) | |
attn_procs[name] = attn_processor_class( | |
hidden_size=hidden_size, | |
cross_attention_dim=cross_attention_dim, | |
scale=1.0, | |
num_tokens=num_image_text_embeds, | |
).to(dtype=self.dtype, device=self.device) | |
value_dict = {} | |
for k, w in attn_procs[name].state_dict().items(): | |
value_dict.update({f"{k}": state_dict["ip_adapter"][f"{key_id}.{k}"]}) | |
attn_procs[name].load_state_dict(value_dict) | |
key_id += 2 | |
self.set_attn_processor(attn_procs) | |
# convert IP-Adapter Image Projection layers to diffusers | |
image_projection = self._convert_ip_adapter_image_proj_to_diffusers(state_dict["image_proj"]) | |
self.encoder_hid_proj = image_projection.to(device=self.device, dtype=self.dtype) | |
self.config.encoder_hid_dim_type = "ip_image_proj" | |