6135 lines
282 KiB
Python
6135 lines
282 KiB
Python
# Copyright 2025 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 os
|
|
from typing import Callable, Dict, List, Optional, Union
|
|
|
|
import torch
|
|
from huggingface_hub.utils import validate_hf_hub_args
|
|
|
|
from ..utils import (
|
|
USE_PEFT_BACKEND,
|
|
deprecate,
|
|
get_submodule_by_name,
|
|
is_bitsandbytes_available,
|
|
is_gguf_available,
|
|
is_peft_available,
|
|
is_peft_version,
|
|
is_torch_version,
|
|
is_transformers_available,
|
|
is_transformers_version,
|
|
logging,
|
|
)
|
|
from .lora_base import ( # noqa
|
|
LORA_WEIGHT_NAME,
|
|
LORA_WEIGHT_NAME_SAFE,
|
|
LoraBaseMixin,
|
|
_fetch_state_dict,
|
|
_load_lora_into_text_encoder,
|
|
_pack_dict_with_prefix,
|
|
)
|
|
from .lora_conversion_utils import (
|
|
_convert_bfl_flux_control_lora_to_diffusers,
|
|
_convert_hunyuan_video_lora_to_diffusers,
|
|
_convert_kohya_flux_lora_to_diffusers,
|
|
_convert_musubi_wan_lora_to_diffusers,
|
|
_convert_non_diffusers_hidream_lora_to_diffusers,
|
|
_convert_non_diffusers_lora_to_diffusers,
|
|
_convert_non_diffusers_ltxv_lora_to_diffusers,
|
|
_convert_non_diffusers_lumina2_lora_to_diffusers,
|
|
_convert_non_diffusers_wan_lora_to_diffusers,
|
|
_convert_xlabs_flux_lora_to_diffusers,
|
|
_maybe_map_sgm_blocks_to_diffusers,
|
|
)
|
|
|
|
|
|
_LOW_CPU_MEM_USAGE_DEFAULT_LORA = False
|
|
if is_torch_version(">=", "1.9.0"):
|
|
if (
|
|
is_peft_available()
|
|
and is_peft_version(">=", "0.13.1")
|
|
and is_transformers_available()
|
|
and is_transformers_version(">", "4.45.2")
|
|
):
|
|
_LOW_CPU_MEM_USAGE_DEFAULT_LORA = True
|
|
|
|
|
|
logger = logging.get_logger(__name__)
|
|
|
|
TEXT_ENCODER_NAME = "text_encoder"
|
|
UNET_NAME = "unet"
|
|
TRANSFORMER_NAME = "transformer"
|
|
|
|
_MODULE_NAME_TO_ATTRIBUTE_MAP_FLUX = {"x_embedder": "in_channels"}
|
|
|
|
|
|
def _maybe_dequantize_weight_for_expanded_lora(model, module):
|
|
if is_bitsandbytes_available():
|
|
from ..quantizers.bitsandbytes import dequantize_bnb_weight
|
|
|
|
if is_gguf_available():
|
|
from ..quantizers.gguf.utils import dequantize_gguf_tensor
|
|
|
|
is_bnb_4bit_quantized = module.weight.__class__.__name__ == "Params4bit"
|
|
is_bnb_8bit_quantized = module.weight.__class__.__name__ == "Int8Params"
|
|
is_gguf_quantized = module.weight.__class__.__name__ == "GGUFParameter"
|
|
|
|
if is_bnb_4bit_quantized and not is_bitsandbytes_available():
|
|
raise ValueError(
|
|
"The checkpoint seems to have been quantized with `bitsandbytes` (4bits). Install `bitsandbytes` to load quantized checkpoints."
|
|
)
|
|
if is_bnb_8bit_quantized and not is_bitsandbytes_available():
|
|
raise ValueError(
|
|
"The checkpoint seems to have been quantized with `bitsandbytes` (8bits). Install `bitsandbytes` to load quantized checkpoints."
|
|
)
|
|
if is_gguf_quantized and not is_gguf_available():
|
|
raise ValueError(
|
|
"The checkpoint seems to have been quantized with `gguf`. Install `gguf` to load quantized checkpoints."
|
|
)
|
|
|
|
weight_on_cpu = False
|
|
if module.weight.device.type == "cpu":
|
|
weight_on_cpu = True
|
|
|
|
device = torch.accelerator.current_accelerator().type if hasattr(torch, "accelerator") else "cuda"
|
|
if is_bnb_4bit_quantized or is_bnb_8bit_quantized:
|
|
module_weight = dequantize_bnb_weight(
|
|
module.weight.to(device) if weight_on_cpu else module.weight,
|
|
state=module.weight.quant_state if is_bnb_4bit_quantized else module.state,
|
|
dtype=model.dtype,
|
|
).data
|
|
elif is_gguf_quantized:
|
|
module_weight = dequantize_gguf_tensor(
|
|
module.weight.to(device) if weight_on_cpu else module.weight,
|
|
)
|
|
module_weight = module_weight.to(model.dtype)
|
|
else:
|
|
module_weight = module.weight.data
|
|
|
|
if weight_on_cpu:
|
|
module_weight = module_weight.cpu()
|
|
|
|
return module_weight
|
|
|
|
|
|
class StableDiffusionLoraLoaderMixin(LoraBaseMixin):
|
|
r"""
|
|
Load LoRA layers into Stable Diffusion [`UNet2DConditionModel`] and
|
|
[`CLIPTextModel`](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel).
|
|
"""
|
|
|
|
_lora_loadable_modules = ["unet", "text_encoder"]
|
|
unet_name = UNET_NAME
|
|
text_encoder_name = TEXT_ENCODER_NAME
|
|
|
|
def load_lora_weights(
|
|
self,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name: Optional[str] = None,
|
|
hotswap: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.unet` and
|
|
`self.text_encoder`.
|
|
|
|
All kwargs are forwarded to `self.lora_state_dict`.
|
|
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is
|
|
loaded.
|
|
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details on how the state dict is
|
|
loaded into `self.unet`.
|
|
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder`] for more details on how the state
|
|
dict is loaded into `self.text_encoder`.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
Defaults to `False`. Whether to substitute an existing (LoRA) adapter with the newly loaded adapter
|
|
in-place. This means that, instead of loading an additional adapter, this will take the existing
|
|
adapter weights and replace them with the weights of the new adapter. This can be faster and more
|
|
memory efficient. However, the main advantage of hotswapping is that when the model is compiled with
|
|
torch.compile, loading the new adapter does not require recompilation of the model. When using
|
|
hotswapping, the passed `adapter_name` should be the name of an already loaded adapter.
|
|
|
|
If the new adapter and the old adapter have different ranks and/or LoRA alphas (i.e. scaling), you need
|
|
to call an additional method before loading the adapter:
|
|
|
|
```py
|
|
pipeline = ... # load diffusers pipeline
|
|
max_rank = ... # the highest rank among all LoRAs that you want to load
|
|
# call *before* compiling and loading the LoRA adapter
|
|
pipeline.enable_lora_hotswap(target_rank=max_rank)
|
|
pipeline.load_lora_weights(file_name)
|
|
# optionally compile the model now
|
|
```
|
|
|
|
Note that hotswapping adapters of the text encoder is not yet supported. There are some further
|
|
limitations to this technique, which are documented here:
|
|
https://huggingface.co/docs/peft/main/en/package_reference/hotswap
|
|
kwargs (`dict`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
kwargs["return_lora_metadata"] = True
|
|
state_dict, network_alphas, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
|
|
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
self.load_lora_into_unet(
|
|
state_dict,
|
|
network_alphas=network_alphas,
|
|
unet=getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
self.load_lora_into_text_encoder(
|
|
state_dict,
|
|
network_alphas=network_alphas,
|
|
text_encoder=getattr(self, self.text_encoder_name)
|
|
if not hasattr(self, "text_encoder")
|
|
else self.text_encoder,
|
|
lora_scale=self.lora_scale,
|
|
adapter_name=adapter_name,
|
|
_pipeline=self,
|
|
metadata=metadata,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
@validate_hf_hub_args
|
|
def lora_state_dict(
|
|
cls,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Return state dict for lora weights and the network alphas.
|
|
|
|
<Tip warning={true}>
|
|
|
|
We support loading A1111 formatted LoRA checkpoints in a limited capacity.
|
|
|
|
This function is experimental and might change in the future.
|
|
|
|
</Tip>
|
|
|
|
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.
|
|
|
|
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.
|
|
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.
|
|
weight_name (`str`, *optional*, defaults to None):
|
|
Name of the serialized state dict file.
|
|
return_lora_metadata (`bool`, *optional*, defaults to False):
|
|
When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
|
|
"""
|
|
# Load the main state dict first which has the LoRA layers for either of
|
|
# UNet and text encoder or both.
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
force_download = kwargs.pop("force_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)
|
|
unet_config = kwargs.pop("unet_config", None)
|
|
use_safetensors = kwargs.pop("use_safetensors", None)
|
|
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = True
|
|
allow_pickle = True
|
|
|
|
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
|
|
|
|
state_dict, metadata = _fetch_state_dict(
|
|
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
|
|
weight_name=weight_name,
|
|
use_safetensors=use_safetensors,
|
|
local_files_only=local_files_only,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
token=token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
allow_pickle=allow_pickle,
|
|
)
|
|
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
|
|
if is_dora_scale_present:
|
|
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
|
|
logger.warning(warn_msg)
|
|
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
|
|
|
|
network_alphas = None
|
|
# TODO: replace it with a method from `state_dict_utils`
|
|
if all(
|
|
(
|
|
k.startswith("lora_te_")
|
|
or k.startswith("lora_unet_")
|
|
or k.startswith("lora_te1_")
|
|
or k.startswith("lora_te2_")
|
|
)
|
|
for k in state_dict.keys()
|
|
):
|
|
# Map SDXL blocks correctly.
|
|
if unet_config is not None:
|
|
# use unet config to remap block numbers
|
|
state_dict = _maybe_map_sgm_blocks_to_diffusers(state_dict, unet_config)
|
|
state_dict, network_alphas = _convert_non_diffusers_lora_to_diffusers(state_dict)
|
|
|
|
out = (state_dict, network_alphas, metadata) if return_lora_metadata else (state_dict, network_alphas)
|
|
return out
|
|
|
|
@classmethod
|
|
def load_lora_into_unet(
|
|
cls,
|
|
state_dict,
|
|
network_alphas,
|
|
unet,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `unet`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
network_alphas (`Dict[str, float]`):
|
|
The value of the network alpha used for stable learning and preventing underflow. This value has the
|
|
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
|
|
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
|
|
unet (`UNet2DConditionModel`):
|
|
The UNet model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# If the serialization format is new (introduced in https://github.com/huggingface/diffusers/pull/2918),
|
|
# then the `state_dict` keys should have `cls.unet_name` and/or `cls.text_encoder_name` as
|
|
# their prefixes.
|
|
logger.info(f"Loading {cls.unet_name}.")
|
|
unet.load_lora_adapter(
|
|
state_dict,
|
|
prefix=cls.unet_name,
|
|
network_alphas=network_alphas,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
def load_lora_into_text_encoder(
|
|
cls,
|
|
state_dict,
|
|
network_alphas,
|
|
text_encoder,
|
|
prefix=None,
|
|
lora_scale=1.0,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `text_encoder`
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The key should be prefixed with an
|
|
additional `text_encoder` to distinguish between unet lora layers.
|
|
network_alphas (`Dict[str, float]`):
|
|
The value of the network alpha used for stable learning and preventing underflow. This value has the
|
|
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
|
|
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
|
|
text_encoder (`CLIPTextModel`):
|
|
The text encoder model to load the LoRA layers into.
|
|
prefix (`str`):
|
|
Expected prefix of the `text_encoder` in the `state_dict`.
|
|
lora_scale (`float`):
|
|
How much to scale the output of the lora linear layer before it is added with the output of the regular
|
|
lora layer.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
_load_lora_into_text_encoder(
|
|
state_dict=state_dict,
|
|
network_alphas=network_alphas,
|
|
lora_scale=lora_scale,
|
|
text_encoder=text_encoder,
|
|
prefix=prefix,
|
|
text_encoder_name=cls.text_encoder_name,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
text_encoder_lora_layers: Dict[str, torch.nn.Module] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
unet_lora_adapter_metadata=None,
|
|
text_encoder_lora_adapter_metadata=None,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the UNet and text encoder.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
unet_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `unet`.
|
|
text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text
|
|
encoder LoRA state dict because it comes from 🤗 Transformers.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
unet_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the unet to be serialized with the state dict.
|
|
text_encoder_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the text encoder to be serialized with the state dict.
|
|
"""
|
|
state_dict = {}
|
|
lora_adapter_metadata = {}
|
|
|
|
if not (unet_lora_layers or text_encoder_lora_layers):
|
|
raise ValueError("You must pass at least one of `unet_lora_layers` and `text_encoder_lora_layers`.")
|
|
|
|
if unet_lora_layers:
|
|
state_dict.update(cls.pack_weights(unet_lora_layers, cls.unet_name))
|
|
|
|
if text_encoder_lora_layers:
|
|
state_dict.update(cls.pack_weights(text_encoder_lora_layers, cls.text_encoder_name))
|
|
|
|
if unet_lora_adapter_metadata:
|
|
lora_adapter_metadata.update(_pack_dict_with_prefix(unet_lora_adapter_metadata, cls.unet_name))
|
|
|
|
if text_encoder_lora_adapter_metadata:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name)
|
|
)
|
|
|
|
# Save the model
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
lora_adapter_metadata=lora_adapter_metadata,
|
|
)
|
|
|
|
def fuse_lora(
|
|
self,
|
|
components: List[str] = ["unet", "text_encoder"],
|
|
lora_scale: float = 1.0,
|
|
safe_fusing: bool = False,
|
|
adapter_names: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Fuses the LoRA parameters into the original parameters of the corresponding blocks.
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
|
|
lora_scale (`float`, defaults to 1.0):
|
|
Controls how much to influence the outputs with the LoRA parameters.
|
|
safe_fusing (`bool`, defaults to `False`):
|
|
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
|
|
adapter_names (`List[str]`, *optional*):
|
|
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
|
|
|
|
Example:
|
|
|
|
```py
|
|
from diffusers import DiffusionPipeline
|
|
import torch
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
|
).to("cuda")
|
|
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
|
|
pipeline.fuse_lora(lora_scale=0.7)
|
|
```
|
|
"""
|
|
super().fuse_lora(
|
|
components=components,
|
|
lora_scale=lora_scale,
|
|
safe_fusing=safe_fusing,
|
|
adapter_names=adapter_names,
|
|
**kwargs,
|
|
)
|
|
|
|
def unfuse_lora(self, components: List[str] = ["unet", "text_encoder"], **kwargs):
|
|
r"""
|
|
Reverses the effect of
|
|
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
|
|
unfuse_unet (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
|
|
unfuse_text_encoder (`bool`, defaults to `True`):
|
|
Whether to unfuse the text encoder LoRA parameters. If the text encoder wasn't monkey-patched with the
|
|
LoRA parameters then it won't have any effect.
|
|
"""
|
|
super().unfuse_lora(components=components, **kwargs)
|
|
|
|
|
|
class StableDiffusionXLLoraLoaderMixin(LoraBaseMixin):
|
|
r"""
|
|
Load LoRA layers into Stable Diffusion XL [`UNet2DConditionModel`],
|
|
[`CLIPTextModel`](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), and
|
|
[`CLIPTextModelWithProjection`](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection).
|
|
"""
|
|
|
|
_lora_loadable_modules = ["unet", "text_encoder", "text_encoder_2"]
|
|
unet_name = UNET_NAME
|
|
text_encoder_name = TEXT_ENCODER_NAME
|
|
|
|
def load_lora_weights(
|
|
self,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name: Optional[str] = None,
|
|
hotswap: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.unet` and
|
|
`self.text_encoder`.
|
|
|
|
All kwargs are forwarded to `self.lora_state_dict`.
|
|
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is
|
|
loaded.
|
|
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details on how the state dict is
|
|
loaded into `self.unet`.
|
|
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder`] for more details on how the state
|
|
dict is loaded into `self.text_encoder`.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
kwargs (`dict`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# We could have accessed the unet config from `lora_state_dict()` too. We pass
|
|
# it here explicitly to be able to tell that it's coming from an SDXL
|
|
# pipeline.
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
kwargs["return_lora_metadata"] = True
|
|
state_dict, network_alphas, metadata = self.lora_state_dict(
|
|
pretrained_model_name_or_path_or_dict,
|
|
unet_config=self.unet.config,
|
|
**kwargs,
|
|
)
|
|
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
self.load_lora_into_unet(
|
|
state_dict,
|
|
network_alphas=network_alphas,
|
|
unet=self.unet,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
self.load_lora_into_text_encoder(
|
|
state_dict,
|
|
network_alphas=network_alphas,
|
|
text_encoder=self.text_encoder,
|
|
prefix=self.text_encoder_name,
|
|
lora_scale=self.lora_scale,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
self.load_lora_into_text_encoder(
|
|
state_dict,
|
|
network_alphas=network_alphas,
|
|
text_encoder=self.text_encoder_2,
|
|
prefix=f"{self.text_encoder_name}_2",
|
|
lora_scale=self.lora_scale,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
@validate_hf_hub_args
|
|
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.lora_state_dict
|
|
def lora_state_dict(
|
|
cls,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Return state dict for lora weights and the network alphas.
|
|
|
|
<Tip warning={true}>
|
|
|
|
We support loading A1111 formatted LoRA checkpoints in a limited capacity.
|
|
|
|
This function is experimental and might change in the future.
|
|
|
|
</Tip>
|
|
|
|
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.
|
|
|
|
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.
|
|
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.
|
|
weight_name (`str`, *optional*, defaults to None):
|
|
Name of the serialized state dict file.
|
|
return_lora_metadata (`bool`, *optional*, defaults to False):
|
|
When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
|
|
"""
|
|
# Load the main state dict first which has the LoRA layers for either of
|
|
# UNet and text encoder or both.
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
force_download = kwargs.pop("force_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)
|
|
unet_config = kwargs.pop("unet_config", None)
|
|
use_safetensors = kwargs.pop("use_safetensors", None)
|
|
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = True
|
|
allow_pickle = True
|
|
|
|
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
|
|
|
|
state_dict, metadata = _fetch_state_dict(
|
|
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
|
|
weight_name=weight_name,
|
|
use_safetensors=use_safetensors,
|
|
local_files_only=local_files_only,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
token=token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
allow_pickle=allow_pickle,
|
|
)
|
|
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
|
|
if is_dora_scale_present:
|
|
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
|
|
logger.warning(warn_msg)
|
|
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
|
|
|
|
network_alphas = None
|
|
# TODO: replace it with a method from `state_dict_utils`
|
|
if all(
|
|
(
|
|
k.startswith("lora_te_")
|
|
or k.startswith("lora_unet_")
|
|
or k.startswith("lora_te1_")
|
|
or k.startswith("lora_te2_")
|
|
)
|
|
for k in state_dict.keys()
|
|
):
|
|
# Map SDXL blocks correctly.
|
|
if unet_config is not None:
|
|
# use unet config to remap block numbers
|
|
state_dict = _maybe_map_sgm_blocks_to_diffusers(state_dict, unet_config)
|
|
state_dict, network_alphas = _convert_non_diffusers_lora_to_diffusers(state_dict)
|
|
|
|
out = (state_dict, network_alphas, metadata) if return_lora_metadata else (state_dict, network_alphas)
|
|
return out
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.load_lora_into_unet
|
|
def load_lora_into_unet(
|
|
cls,
|
|
state_dict,
|
|
network_alphas,
|
|
unet,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `unet`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
network_alphas (`Dict[str, float]`):
|
|
The value of the network alpha used for stable learning and preventing underflow. This value has the
|
|
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
|
|
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
|
|
unet (`UNet2DConditionModel`):
|
|
The UNet model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# If the serialization format is new (introduced in https://github.com/huggingface/diffusers/pull/2918),
|
|
# then the `state_dict` keys should have `cls.unet_name` and/or `cls.text_encoder_name` as
|
|
# their prefixes.
|
|
logger.info(f"Loading {cls.unet_name}.")
|
|
unet.load_lora_adapter(
|
|
state_dict,
|
|
prefix=cls.unet_name,
|
|
network_alphas=network_alphas,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder
|
|
def load_lora_into_text_encoder(
|
|
cls,
|
|
state_dict,
|
|
network_alphas,
|
|
text_encoder,
|
|
prefix=None,
|
|
lora_scale=1.0,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `text_encoder`
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The key should be prefixed with an
|
|
additional `text_encoder` to distinguish between unet lora layers.
|
|
network_alphas (`Dict[str, float]`):
|
|
The value of the network alpha used for stable learning and preventing underflow. This value has the
|
|
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
|
|
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
|
|
text_encoder (`CLIPTextModel`):
|
|
The text encoder model to load the LoRA layers into.
|
|
prefix (`str`):
|
|
Expected prefix of the `text_encoder` in the `state_dict`.
|
|
lora_scale (`float`):
|
|
How much to scale the output of the lora linear layer before it is added with the output of the regular
|
|
lora layer.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
_load_lora_into_text_encoder(
|
|
state_dict=state_dict,
|
|
network_alphas=network_alphas,
|
|
lora_scale=lora_scale,
|
|
text_encoder=text_encoder,
|
|
prefix=prefix,
|
|
text_encoder_name=cls.text_encoder_name,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
unet_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
unet_lora_adapter_metadata=None,
|
|
text_encoder_lora_adapter_metadata=None,
|
|
text_encoder_2_lora_adapter_metadata=None,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the UNet and text encoder.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
unet_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `unet`.
|
|
text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text
|
|
encoder LoRA state dict because it comes from 🤗 Transformers.
|
|
text_encoder_2_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `text_encoder_2`. Must explicitly pass the text
|
|
encoder LoRA state dict because it comes from 🤗 Transformers.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
unet_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the unet to be serialized with the state dict.
|
|
text_encoder_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the text encoder to be serialized with the state dict.
|
|
text_encoder_2_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the second text encoder to be serialized with the state dict.
|
|
"""
|
|
state_dict = {}
|
|
lora_adapter_metadata = {}
|
|
|
|
if not (unet_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers):
|
|
raise ValueError(
|
|
"You must pass at least one of `unet_lora_layers`, `text_encoder_lora_layers`, `text_encoder_2_lora_layers`."
|
|
)
|
|
|
|
if unet_lora_layers:
|
|
state_dict.update(cls.pack_weights(unet_lora_layers, cls.unet_name))
|
|
|
|
if text_encoder_lora_layers:
|
|
state_dict.update(cls.pack_weights(text_encoder_lora_layers, "text_encoder"))
|
|
|
|
if text_encoder_2_lora_layers:
|
|
state_dict.update(cls.pack_weights(text_encoder_2_lora_layers, "text_encoder_2"))
|
|
|
|
if unet_lora_adapter_metadata is not None:
|
|
lora_adapter_metadata.update(_pack_dict_with_prefix(unet_lora_adapter_metadata, cls.unet_name))
|
|
|
|
if text_encoder_lora_adapter_metadata:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name)
|
|
)
|
|
|
|
if text_encoder_2_lora_adapter_metadata:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(text_encoder_2_lora_adapter_metadata, "text_encoder_2")
|
|
)
|
|
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
lora_adapter_metadata=lora_adapter_metadata,
|
|
)
|
|
|
|
def fuse_lora(
|
|
self,
|
|
components: List[str] = ["unet", "text_encoder", "text_encoder_2"],
|
|
lora_scale: float = 1.0,
|
|
safe_fusing: bool = False,
|
|
adapter_names: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Fuses the LoRA parameters into the original parameters of the corresponding blocks.
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
|
|
lora_scale (`float`, defaults to 1.0):
|
|
Controls how much to influence the outputs with the LoRA parameters.
|
|
safe_fusing (`bool`, defaults to `False`):
|
|
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
|
|
adapter_names (`List[str]`, *optional*):
|
|
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
|
|
|
|
Example:
|
|
|
|
```py
|
|
from diffusers import DiffusionPipeline
|
|
import torch
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
|
).to("cuda")
|
|
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
|
|
pipeline.fuse_lora(lora_scale=0.7)
|
|
```
|
|
"""
|
|
super().fuse_lora(
|
|
components=components,
|
|
lora_scale=lora_scale,
|
|
safe_fusing=safe_fusing,
|
|
adapter_names=adapter_names,
|
|
**kwargs,
|
|
)
|
|
|
|
def unfuse_lora(self, components: List[str] = ["unet", "text_encoder", "text_encoder_2"], **kwargs):
|
|
r"""
|
|
Reverses the effect of
|
|
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
|
|
unfuse_unet (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
|
|
unfuse_text_encoder (`bool`, defaults to `True`):
|
|
Whether to unfuse the text encoder LoRA parameters. If the text encoder wasn't monkey-patched with the
|
|
LoRA parameters then it won't have any effect.
|
|
"""
|
|
super().unfuse_lora(components=components, **kwargs)
|
|
|
|
|
|
class SD3LoraLoaderMixin(LoraBaseMixin):
|
|
r"""
|
|
Load LoRA layers into [`SD3Transformer2DModel`],
|
|
[`CLIPTextModel`](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), and
|
|
[`CLIPTextModelWithProjection`](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection).
|
|
|
|
Specific to [`StableDiffusion3Pipeline`].
|
|
"""
|
|
|
|
_lora_loadable_modules = ["transformer", "text_encoder", "text_encoder_2"]
|
|
transformer_name = TRANSFORMER_NAME
|
|
text_encoder_name = TEXT_ENCODER_NAME
|
|
|
|
@classmethod
|
|
@validate_hf_hub_args
|
|
def lora_state_dict(
|
|
cls,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Return state dict for lora weights and the network alphas.
|
|
|
|
<Tip warning={true}>
|
|
|
|
We support loading A1111 formatted LoRA checkpoints in a limited capacity.
|
|
|
|
This function is experimental and might change in the future.
|
|
|
|
</Tip>
|
|
|
|
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.
|
|
|
|
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.
|
|
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.
|
|
return_lora_metadata (`bool`, *optional*, defaults to False):
|
|
When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
|
|
|
|
"""
|
|
# Load the main state dict first which has the LoRA layers for either of
|
|
# transformer and text encoder or both.
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
force_download = kwargs.pop("force_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)
|
|
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = True
|
|
allow_pickle = True
|
|
|
|
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
|
|
|
|
state_dict, metadata = _fetch_state_dict(
|
|
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
|
|
weight_name=weight_name,
|
|
use_safetensors=use_safetensors,
|
|
local_files_only=local_files_only,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
token=token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
allow_pickle=allow_pickle,
|
|
)
|
|
|
|
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
|
|
if is_dora_scale_present:
|
|
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
|
|
logger.warning(warn_msg)
|
|
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
|
|
|
|
out = (state_dict, metadata) if return_lora_metadata else state_dict
|
|
return out
|
|
|
|
def load_lora_weights(
|
|
self,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name=None,
|
|
hotswap: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.unet` and
|
|
`self.text_encoder`.
|
|
|
|
All kwargs are forwarded to `self.lora_state_dict`.
|
|
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is
|
|
loaded.
|
|
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state
|
|
dict is loaded into `self.transformer`.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
kwargs (`dict`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
kwargs["return_lora_metadata"] = True
|
|
state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
|
|
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
self.load_lora_into_transformer(
|
|
state_dict,
|
|
transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
self.load_lora_into_text_encoder(
|
|
state_dict,
|
|
network_alphas=None,
|
|
text_encoder=self.text_encoder,
|
|
prefix=self.text_encoder_name,
|
|
lora_scale=self.lora_scale,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
self.load_lora_into_text_encoder(
|
|
state_dict,
|
|
network_alphas=None,
|
|
text_encoder=self.text_encoder_2,
|
|
prefix=f"{self.text_encoder_name}_2",
|
|
lora_scale=self.lora_scale,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
def load_lora_into_transformer(
|
|
cls,
|
|
state_dict,
|
|
transformer,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `transformer`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
transformer (`SD3Transformer2DModel`):
|
|
The Transformer model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# Load the layers corresponding to transformer.
|
|
logger.info(f"Loading {cls.transformer_name}.")
|
|
transformer.load_lora_adapter(
|
|
state_dict,
|
|
network_alphas=None,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder
|
|
def load_lora_into_text_encoder(
|
|
cls,
|
|
state_dict,
|
|
network_alphas,
|
|
text_encoder,
|
|
prefix=None,
|
|
lora_scale=1.0,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `text_encoder`
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The key should be prefixed with an
|
|
additional `text_encoder` to distinguish between unet lora layers.
|
|
network_alphas (`Dict[str, float]`):
|
|
The value of the network alpha used for stable learning and preventing underflow. This value has the
|
|
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
|
|
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
|
|
text_encoder (`CLIPTextModel`):
|
|
The text encoder model to load the LoRA layers into.
|
|
prefix (`str`):
|
|
Expected prefix of the `text_encoder` in the `state_dict`.
|
|
lora_scale (`float`):
|
|
How much to scale the output of the lora linear layer before it is added with the output of the regular
|
|
lora layer.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
_load_lora_into_text_encoder(
|
|
state_dict=state_dict,
|
|
network_alphas=network_alphas,
|
|
lora_scale=lora_scale,
|
|
text_encoder=text_encoder,
|
|
prefix=prefix,
|
|
text_encoder_name=cls.text_encoder_name,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionXLLoraLoaderMixin.save_lora_weights with unet->transformer
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
text_encoder_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
text_encoder_2_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
transformer_lora_adapter_metadata=None,
|
|
text_encoder_lora_adapter_metadata=None,
|
|
text_encoder_2_lora_adapter_metadata=None,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the UNet and text encoder.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `transformer`.
|
|
text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text
|
|
encoder LoRA state dict because it comes from 🤗 Transformers.
|
|
text_encoder_2_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `text_encoder_2`. Must explicitly pass the text
|
|
encoder LoRA state dict because it comes from 🤗 Transformers.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
transformer_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the transformer to be serialized with the state dict.
|
|
text_encoder_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the text encoder to be serialized with the state dict.
|
|
text_encoder_2_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the second text encoder to be serialized with the state dict.
|
|
"""
|
|
state_dict = {}
|
|
lora_adapter_metadata = {}
|
|
|
|
if not (transformer_lora_layers or text_encoder_lora_layers or text_encoder_2_lora_layers):
|
|
raise ValueError(
|
|
"You must pass at least one of `transformer_lora_layers`, `text_encoder_lora_layers`, `text_encoder_2_lora_layers`."
|
|
)
|
|
|
|
if transformer_lora_layers:
|
|
state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
|
|
|
|
if text_encoder_lora_layers:
|
|
state_dict.update(cls.pack_weights(text_encoder_lora_layers, "text_encoder"))
|
|
|
|
if text_encoder_2_lora_layers:
|
|
state_dict.update(cls.pack_weights(text_encoder_2_lora_layers, "text_encoder_2"))
|
|
|
|
if transformer_lora_adapter_metadata is not None:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
|
|
)
|
|
|
|
if text_encoder_lora_adapter_metadata:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name)
|
|
)
|
|
|
|
if text_encoder_2_lora_adapter_metadata:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(text_encoder_2_lora_adapter_metadata, "text_encoder_2")
|
|
)
|
|
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
lora_adapter_metadata=lora_adapter_metadata,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionXLLoraLoaderMixin.fuse_lora with unet->transformer
|
|
def fuse_lora(
|
|
self,
|
|
components: List[str] = ["transformer", "text_encoder", "text_encoder_2"],
|
|
lora_scale: float = 1.0,
|
|
safe_fusing: bool = False,
|
|
adapter_names: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Fuses the LoRA parameters into the original parameters of the corresponding blocks.
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
|
|
lora_scale (`float`, defaults to 1.0):
|
|
Controls how much to influence the outputs with the LoRA parameters.
|
|
safe_fusing (`bool`, defaults to `False`):
|
|
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
|
|
adapter_names (`List[str]`, *optional*):
|
|
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
|
|
|
|
Example:
|
|
|
|
```py
|
|
from diffusers import DiffusionPipeline
|
|
import torch
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
|
).to("cuda")
|
|
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
|
|
pipeline.fuse_lora(lora_scale=0.7)
|
|
```
|
|
"""
|
|
super().fuse_lora(
|
|
components=components,
|
|
lora_scale=lora_scale,
|
|
safe_fusing=safe_fusing,
|
|
adapter_names=adapter_names,
|
|
**kwargs,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionXLLoraLoaderMixin.unfuse_lora with unet->transformer
|
|
def unfuse_lora(self, components: List[str] = ["transformer", "text_encoder", "text_encoder_2"], **kwargs):
|
|
r"""
|
|
Reverses the effect of
|
|
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
|
|
unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
|
|
unfuse_text_encoder (`bool`, defaults to `True`):
|
|
Whether to unfuse the text encoder LoRA parameters. If the text encoder wasn't monkey-patched with the
|
|
LoRA parameters then it won't have any effect.
|
|
"""
|
|
super().unfuse_lora(components=components, **kwargs)
|
|
|
|
|
|
class AuraFlowLoraLoaderMixin(LoraBaseMixin):
|
|
r"""
|
|
Load LoRA layers into [`AuraFlowTransformer2DModel`] Specific to [`AuraFlowPipeline`].
|
|
"""
|
|
|
|
_lora_loadable_modules = ["transformer"]
|
|
transformer_name = TRANSFORMER_NAME
|
|
|
|
@classmethod
|
|
@validate_hf_hub_args
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.lora_state_dict
|
|
def lora_state_dict(
|
|
cls,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Return state dict for lora weights and the network alphas.
|
|
|
|
<Tip warning={true}>
|
|
|
|
We support loading A1111 formatted LoRA checkpoints in a limited capacity.
|
|
|
|
This function is experimental and might change in the future.
|
|
|
|
</Tip>
|
|
|
|
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.
|
|
|
|
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.
|
|
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.
|
|
return_lora_metadata (`bool`, *optional*, defaults to False):
|
|
When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
|
|
|
|
"""
|
|
# Load the main state dict first which has the LoRA layers for either of
|
|
# transformer and text encoder or both.
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
force_download = kwargs.pop("force_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)
|
|
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = True
|
|
allow_pickle = True
|
|
|
|
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
|
|
|
|
state_dict, metadata = _fetch_state_dict(
|
|
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
|
|
weight_name=weight_name,
|
|
use_safetensors=use_safetensors,
|
|
local_files_only=local_files_only,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
token=token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
allow_pickle=allow_pickle,
|
|
)
|
|
|
|
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
|
|
if is_dora_scale_present:
|
|
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
|
|
logger.warning(warn_msg)
|
|
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
|
|
|
|
out = (state_dict, metadata) if return_lora_metadata else state_dict
|
|
return out
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
|
|
def load_lora_weights(
|
|
self,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name: Optional[str] = None,
|
|
hotswap: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and
|
|
`self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See
|
|
[`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded.
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state
|
|
dict is loaded into `self.transformer`.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
kwargs (`dict`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
kwargs["return_lora_metadata"] = True
|
|
state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
|
|
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
self.load_lora_into_transformer(
|
|
state_dict,
|
|
transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer with SD3Transformer2DModel->AuraFlowTransformer2DModel
|
|
def load_lora_into_transformer(
|
|
cls,
|
|
state_dict,
|
|
transformer,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `transformer`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
transformer (`AuraFlowTransformer2DModel`):
|
|
The Transformer model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# Load the layers corresponding to transformer.
|
|
logger.info(f"Loading {cls.transformer_name}.")
|
|
transformer.load_lora_adapter(
|
|
state_dict,
|
|
network_alphas=None,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
transformer_lora_adapter_metadata: Optional[dict] = None,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the transformer.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `transformer`.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
transformer_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the transformer to be serialized with the state dict.
|
|
"""
|
|
state_dict = {}
|
|
lora_adapter_metadata = {}
|
|
|
|
if not transformer_lora_layers:
|
|
raise ValueError("You must pass `transformer_lora_layers`.")
|
|
|
|
state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
|
|
|
|
if transformer_lora_adapter_metadata is not None:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
|
|
)
|
|
|
|
# Save the model
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
lora_adapter_metadata=lora_adapter_metadata,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora
|
|
def fuse_lora(
|
|
self,
|
|
components: List[str] = ["transformer"],
|
|
lora_scale: float = 1.0,
|
|
safe_fusing: bool = False,
|
|
adapter_names: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Fuses the LoRA parameters into the original parameters of the corresponding blocks.
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
|
|
lora_scale (`float`, defaults to 1.0):
|
|
Controls how much to influence the outputs with the LoRA parameters.
|
|
safe_fusing (`bool`, defaults to `False`):
|
|
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
|
|
adapter_names (`List[str]`, *optional*):
|
|
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
|
|
|
|
Example:
|
|
|
|
```py
|
|
from diffusers import DiffusionPipeline
|
|
import torch
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
|
).to("cuda")
|
|
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
|
|
pipeline.fuse_lora(lora_scale=0.7)
|
|
```
|
|
"""
|
|
super().fuse_lora(
|
|
components=components,
|
|
lora_scale=lora_scale,
|
|
safe_fusing=safe_fusing,
|
|
adapter_names=adapter_names,
|
|
**kwargs,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora
|
|
def unfuse_lora(self, components: List[str] = ["transformer", "text_encoder"], **kwargs):
|
|
r"""
|
|
Reverses the effect of
|
|
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
|
|
unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
|
|
"""
|
|
super().unfuse_lora(components=components, **kwargs)
|
|
|
|
|
|
class FluxLoraLoaderMixin(LoraBaseMixin):
|
|
r"""
|
|
Load LoRA layers into [`FluxTransformer2DModel`],
|
|
[`CLIPTextModel`](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel).
|
|
|
|
Specific to [`StableDiffusion3Pipeline`].
|
|
"""
|
|
|
|
_lora_loadable_modules = ["transformer", "text_encoder"]
|
|
transformer_name = TRANSFORMER_NAME
|
|
text_encoder_name = TEXT_ENCODER_NAME
|
|
_control_lora_supported_norm_keys = ["norm_q", "norm_k", "norm_added_q", "norm_added_k"]
|
|
|
|
@classmethod
|
|
@validate_hf_hub_args
|
|
def lora_state_dict(
|
|
cls,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
return_alphas: bool = False,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Return state dict for lora weights and the network alphas.
|
|
|
|
<Tip warning={true}>
|
|
|
|
We support loading A1111 formatted LoRA checkpoints in a limited capacity.
|
|
|
|
This function is experimental and might change in the future.
|
|
|
|
</Tip>
|
|
|
|
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.
|
|
|
|
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.
|
|
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.
|
|
return_lora_metadata (`bool`, *optional*, defaults to False):
|
|
When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
|
|
"""
|
|
# Load the main state dict first which has the LoRA layers for either of
|
|
# transformer and text encoder or both.
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
force_download = kwargs.pop("force_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)
|
|
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = True
|
|
allow_pickle = True
|
|
|
|
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
|
|
|
|
state_dict, metadata = _fetch_state_dict(
|
|
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
|
|
weight_name=weight_name,
|
|
use_safetensors=use_safetensors,
|
|
local_files_only=local_files_only,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
token=token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
allow_pickle=allow_pickle,
|
|
)
|
|
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
|
|
if is_dora_scale_present:
|
|
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
|
|
logger.warning(warn_msg)
|
|
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
|
|
|
|
# TODO (sayakpaul): to a follow-up to clean and try to unify the conditions.
|
|
is_kohya = any(".lora_down.weight" in k for k in state_dict)
|
|
if is_kohya:
|
|
state_dict = _convert_kohya_flux_lora_to_diffusers(state_dict)
|
|
# Kohya already takes care of scaling the LoRA parameters with alpha.
|
|
return cls._prepare_outputs(
|
|
state_dict,
|
|
metadata=metadata,
|
|
alphas=None,
|
|
return_alphas=return_alphas,
|
|
return_metadata=return_lora_metadata,
|
|
)
|
|
|
|
is_xlabs = any("processor" in k for k in state_dict)
|
|
if is_xlabs:
|
|
state_dict = _convert_xlabs_flux_lora_to_diffusers(state_dict)
|
|
# xlabs doesn't use `alpha`.
|
|
return cls._prepare_outputs(
|
|
state_dict,
|
|
metadata=metadata,
|
|
alphas=None,
|
|
return_alphas=return_alphas,
|
|
return_metadata=return_lora_metadata,
|
|
)
|
|
|
|
is_bfl_control = any("query_norm.scale" in k for k in state_dict)
|
|
if is_bfl_control:
|
|
state_dict = _convert_bfl_flux_control_lora_to_diffusers(state_dict)
|
|
return cls._prepare_outputs(
|
|
state_dict,
|
|
metadata=metadata,
|
|
alphas=None,
|
|
return_alphas=return_alphas,
|
|
return_metadata=return_lora_metadata,
|
|
)
|
|
|
|
# For state dicts like
|
|
# https://huggingface.co/TheLastBen/Jon_Snow_Flux_LoRA
|
|
keys = list(state_dict.keys())
|
|
network_alphas = {}
|
|
for k in keys:
|
|
if "alpha" in k:
|
|
alpha_value = state_dict.get(k)
|
|
if (torch.is_tensor(alpha_value) and torch.is_floating_point(alpha_value)) or isinstance(
|
|
alpha_value, float
|
|
):
|
|
network_alphas[k] = state_dict.pop(k)
|
|
else:
|
|
raise ValueError(
|
|
f"The alpha key ({k}) seems to be incorrect. If you think this error is unexpected, please open as issue."
|
|
)
|
|
|
|
if return_alphas or return_lora_metadata:
|
|
return cls._prepare_outputs(
|
|
state_dict,
|
|
metadata=metadata,
|
|
alphas=network_alphas,
|
|
return_alphas=return_alphas,
|
|
return_metadata=return_lora_metadata,
|
|
)
|
|
else:
|
|
return state_dict
|
|
|
|
def load_lora_weights(
|
|
self,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name: Optional[str] = None,
|
|
hotswap: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and
|
|
`self.text_encoder`.
|
|
|
|
All kwargs are forwarded to `self.lora_state_dict`.
|
|
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is
|
|
loaded.
|
|
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state
|
|
dict is loaded into `self.transformer`.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
`Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
kwargs (`dict`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
kwargs["return_lora_metadata"] = True
|
|
state_dict, network_alphas, metadata = self.lora_state_dict(
|
|
pretrained_model_name_or_path_or_dict, return_alphas=True, **kwargs
|
|
)
|
|
|
|
has_lora_keys = any("lora" in key for key in state_dict.keys())
|
|
|
|
# Flux Control LoRAs also have norm keys
|
|
has_norm_keys = any(
|
|
norm_key in key for key in state_dict.keys() for norm_key in self._control_lora_supported_norm_keys
|
|
)
|
|
|
|
if not (has_lora_keys or has_norm_keys):
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
transformer_lora_state_dict = {
|
|
k: state_dict.get(k)
|
|
for k in list(state_dict.keys())
|
|
if k.startswith(f"{self.transformer_name}.") and "lora" in k
|
|
}
|
|
transformer_norm_state_dict = {
|
|
k: state_dict.pop(k)
|
|
for k in list(state_dict.keys())
|
|
if k.startswith(f"{self.transformer_name}.")
|
|
and any(norm_key in k for norm_key in self._control_lora_supported_norm_keys)
|
|
}
|
|
|
|
transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer
|
|
has_param_with_expanded_shape = False
|
|
if len(transformer_lora_state_dict) > 0:
|
|
has_param_with_expanded_shape = self._maybe_expand_transformer_param_shape_or_error_(
|
|
transformer, transformer_lora_state_dict, transformer_norm_state_dict
|
|
)
|
|
|
|
if has_param_with_expanded_shape:
|
|
logger.info(
|
|
"The LoRA weights contain parameters that have different shapes that expected by the transformer. "
|
|
"As a result, the state_dict of the transformer has been expanded to match the LoRA parameter shapes. "
|
|
"To get a comprehensive list of parameter names that were modified, enable debug logging."
|
|
)
|
|
if len(transformer_lora_state_dict) > 0:
|
|
transformer_lora_state_dict = self._maybe_expand_lora_state_dict(
|
|
transformer=transformer, lora_state_dict=transformer_lora_state_dict
|
|
)
|
|
for k in transformer_lora_state_dict:
|
|
state_dict.update({k: transformer_lora_state_dict[k]})
|
|
|
|
self.load_lora_into_transformer(
|
|
state_dict,
|
|
network_alphas=network_alphas,
|
|
transformer=transformer,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
if len(transformer_norm_state_dict) > 0:
|
|
transformer._transformer_norm_layers = self._load_norm_into_transformer(
|
|
transformer_norm_state_dict,
|
|
transformer=transformer,
|
|
discard_original_layers=False,
|
|
)
|
|
|
|
self.load_lora_into_text_encoder(
|
|
state_dict,
|
|
network_alphas=network_alphas,
|
|
text_encoder=self.text_encoder,
|
|
prefix=self.text_encoder_name,
|
|
lora_scale=self.lora_scale,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
def load_lora_into_transformer(
|
|
cls,
|
|
state_dict,
|
|
network_alphas,
|
|
transformer,
|
|
adapter_name=None,
|
|
metadata=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `transformer`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
network_alphas (`Dict[str, float]`):
|
|
The value of the network alpha used for stable learning and preventing underflow. This value has the
|
|
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
|
|
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
|
|
transformer (`FluxTransformer2DModel`):
|
|
The Transformer model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# Load the layers corresponding to transformer.
|
|
logger.info(f"Loading {cls.transformer_name}.")
|
|
transformer.load_lora_adapter(
|
|
state_dict,
|
|
network_alphas=network_alphas,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
def _load_norm_into_transformer(
|
|
cls,
|
|
state_dict,
|
|
transformer,
|
|
prefix=None,
|
|
discard_original_layers=False,
|
|
) -> Dict[str, torch.Tensor]:
|
|
# Remove prefix if present
|
|
prefix = prefix or cls.transformer_name
|
|
for key in list(state_dict.keys()):
|
|
if key.split(".")[0] == prefix:
|
|
state_dict[key.removeprefix(f"{prefix}.")] = state_dict.pop(key)
|
|
|
|
# Find invalid keys
|
|
transformer_state_dict = transformer.state_dict()
|
|
transformer_keys = set(transformer_state_dict.keys())
|
|
state_dict_keys = set(state_dict.keys())
|
|
extra_keys = list(state_dict_keys - transformer_keys)
|
|
|
|
if extra_keys:
|
|
logger.warning(
|
|
f"Unsupported keys found in state dict when trying to load normalization layers into the transformer. The following keys will be ignored:\n{extra_keys}."
|
|
)
|
|
|
|
for key in extra_keys:
|
|
state_dict.pop(key)
|
|
|
|
# Save the layers that are going to be overwritten so that unload_lora_weights can work as expected
|
|
overwritten_layers_state_dict = {}
|
|
if not discard_original_layers:
|
|
for key in state_dict.keys():
|
|
overwritten_layers_state_dict[key] = transformer_state_dict[key].clone()
|
|
|
|
logger.info(
|
|
"The provided state dict contains normalization layers in addition to LoRA layers. The normalization layers will directly update the state_dict of the transformer "
|
|
'as opposed to the LoRA layers that will co-exist separately until the "fuse_lora()" method is called. That is to say, the normalization layers will always be directly '
|
|
"fused into the transformer and can only be unfused if `discard_original_layers=True` is passed. This might also have implications when dealing with multiple LoRAs. "
|
|
"If you notice something unexpected, please open an issue: https://github.com/huggingface/diffusers/issues."
|
|
)
|
|
|
|
# We can't load with strict=True because the current state_dict does not contain all the transformer keys
|
|
incompatible_keys = transformer.load_state_dict(state_dict, strict=False)
|
|
unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
|
|
|
|
# We shouldn't expect to see the supported norm keys here being present in the unexpected keys.
|
|
if unexpected_keys:
|
|
if any(norm_key in k for k in unexpected_keys for norm_key in cls._control_lora_supported_norm_keys):
|
|
raise ValueError(
|
|
f"Found {unexpected_keys} as unexpected keys while trying to load norm layers into the transformer."
|
|
)
|
|
|
|
return overwritten_layers_state_dict
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder
|
|
def load_lora_into_text_encoder(
|
|
cls,
|
|
state_dict,
|
|
network_alphas,
|
|
text_encoder,
|
|
prefix=None,
|
|
lora_scale=1.0,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `text_encoder`
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The key should be prefixed with an
|
|
additional `text_encoder` to distinguish between unet lora layers.
|
|
network_alphas (`Dict[str, float]`):
|
|
The value of the network alpha used for stable learning and preventing underflow. This value has the
|
|
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
|
|
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
|
|
text_encoder (`CLIPTextModel`):
|
|
The text encoder model to load the LoRA layers into.
|
|
prefix (`str`):
|
|
Expected prefix of the `text_encoder` in the `state_dict`.
|
|
lora_scale (`float`):
|
|
How much to scale the output of the lora linear layer before it is added with the output of the regular
|
|
lora layer.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
_load_lora_into_text_encoder(
|
|
state_dict=state_dict,
|
|
network_alphas=network_alphas,
|
|
lora_scale=lora_scale,
|
|
text_encoder=text_encoder,
|
|
prefix=prefix,
|
|
text_encoder_name=cls.text_encoder_name,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.save_lora_weights with unet->transformer
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
text_encoder_lora_layers: Dict[str, torch.nn.Module] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
transformer_lora_adapter_metadata=None,
|
|
text_encoder_lora_adapter_metadata=None,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the UNet and text encoder.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `transformer`.
|
|
text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text
|
|
encoder LoRA state dict because it comes from 🤗 Transformers.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
transformer_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the transformer to be serialized with the state dict.
|
|
text_encoder_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the text encoder to be serialized with the state dict.
|
|
"""
|
|
state_dict = {}
|
|
lora_adapter_metadata = {}
|
|
|
|
if not (transformer_lora_layers or text_encoder_lora_layers):
|
|
raise ValueError("You must pass at least one of `transformer_lora_layers` and `text_encoder_lora_layers`.")
|
|
|
|
if transformer_lora_layers:
|
|
state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
|
|
|
|
if text_encoder_lora_layers:
|
|
state_dict.update(cls.pack_weights(text_encoder_lora_layers, cls.text_encoder_name))
|
|
|
|
if transformer_lora_adapter_metadata:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
|
|
)
|
|
|
|
if text_encoder_lora_adapter_metadata:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(text_encoder_lora_adapter_metadata, cls.text_encoder_name)
|
|
)
|
|
|
|
# Save the model
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
lora_adapter_metadata=lora_adapter_metadata,
|
|
)
|
|
|
|
def fuse_lora(
|
|
self,
|
|
components: List[str] = ["transformer"],
|
|
lora_scale: float = 1.0,
|
|
safe_fusing: bool = False,
|
|
adapter_names: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Fuses the LoRA parameters into the original parameters of the corresponding blocks.
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
|
|
lora_scale (`float`, defaults to 1.0):
|
|
Controls how much to influence the outputs with the LoRA parameters.
|
|
safe_fusing (`bool`, defaults to `False`):
|
|
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
|
|
adapter_names (`List[str]`, *optional*):
|
|
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
|
|
|
|
Example:
|
|
|
|
```py
|
|
from diffusers import DiffusionPipeline
|
|
import torch
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
|
).to("cuda")
|
|
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
|
|
pipeline.fuse_lora(lora_scale=0.7)
|
|
```
|
|
"""
|
|
|
|
transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer
|
|
if (
|
|
hasattr(transformer, "_transformer_norm_layers")
|
|
and isinstance(transformer._transformer_norm_layers, dict)
|
|
and len(transformer._transformer_norm_layers.keys()) > 0
|
|
):
|
|
logger.info(
|
|
"The provided state dict contains normalization layers in addition to LoRA layers. The normalization layers will be directly updated the state_dict of the transformer "
|
|
"as opposed to the LoRA layers that will co-exist separately until the 'fuse_lora()' method is called. That is to say, the normalization layers will always be directly "
|
|
"fused into the transformer and can only be unfused if `discard_original_layers=True` is passed."
|
|
)
|
|
|
|
super().fuse_lora(
|
|
components=components,
|
|
lora_scale=lora_scale,
|
|
safe_fusing=safe_fusing,
|
|
adapter_names=adapter_names,
|
|
**kwargs,
|
|
)
|
|
|
|
def unfuse_lora(self, components: List[str] = ["transformer", "text_encoder"], **kwargs):
|
|
r"""
|
|
Reverses the effect of
|
|
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
|
|
"""
|
|
transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer
|
|
if hasattr(transformer, "_transformer_norm_layers") and transformer._transformer_norm_layers:
|
|
transformer.load_state_dict(transformer._transformer_norm_layers, strict=False)
|
|
|
|
super().unfuse_lora(components=components, **kwargs)
|
|
|
|
# We override this here account for `_transformer_norm_layers` and `_overwritten_params`.
|
|
def unload_lora_weights(self, reset_to_overwritten_params=False):
|
|
"""
|
|
Unloads the LoRA parameters.
|
|
|
|
Args:
|
|
reset_to_overwritten_params (`bool`, defaults to `False`): Whether to reset the LoRA-loaded modules
|
|
to their original params. Refer to the [Flux
|
|
documentation](https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux) to learn more.
|
|
|
|
Examples:
|
|
|
|
```python
|
|
>>> # Assuming `pipeline` is already loaded with the LoRA parameters.
|
|
>>> pipeline.unload_lora_weights()
|
|
>>> ...
|
|
```
|
|
"""
|
|
super().unload_lora_weights()
|
|
|
|
transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer
|
|
if hasattr(transformer, "_transformer_norm_layers") and transformer._transformer_norm_layers:
|
|
transformer.load_state_dict(transformer._transformer_norm_layers, strict=False)
|
|
transformer._transformer_norm_layers = None
|
|
|
|
if reset_to_overwritten_params and getattr(transformer, "_overwritten_params", None) is not None:
|
|
overwritten_params = transformer._overwritten_params
|
|
module_names = set()
|
|
|
|
for param_name in overwritten_params:
|
|
if param_name.endswith(".weight"):
|
|
module_names.add(param_name.replace(".weight", ""))
|
|
|
|
for name, module in transformer.named_modules():
|
|
if isinstance(module, torch.nn.Linear) and name in module_names:
|
|
module_weight = module.weight.data
|
|
module_bias = module.bias.data if module.bias is not None else None
|
|
bias = module_bias is not None
|
|
|
|
parent_module_name, _, current_module_name = name.rpartition(".")
|
|
parent_module = transformer.get_submodule(parent_module_name)
|
|
|
|
current_param_weight = overwritten_params[f"{name}.weight"]
|
|
in_features, out_features = current_param_weight.shape[1], current_param_weight.shape[0]
|
|
with torch.device("meta"):
|
|
original_module = torch.nn.Linear(
|
|
in_features,
|
|
out_features,
|
|
bias=bias,
|
|
dtype=module_weight.dtype,
|
|
)
|
|
|
|
tmp_state_dict = {"weight": current_param_weight}
|
|
if module_bias is not None:
|
|
tmp_state_dict.update({"bias": overwritten_params[f"{name}.bias"]})
|
|
original_module.load_state_dict(tmp_state_dict, assign=True, strict=True)
|
|
setattr(parent_module, current_module_name, original_module)
|
|
|
|
del tmp_state_dict
|
|
|
|
if current_module_name in _MODULE_NAME_TO_ATTRIBUTE_MAP_FLUX:
|
|
attribute_name = _MODULE_NAME_TO_ATTRIBUTE_MAP_FLUX[current_module_name]
|
|
new_value = int(current_param_weight.shape[1])
|
|
old_value = getattr(transformer.config, attribute_name)
|
|
setattr(transformer.config, attribute_name, new_value)
|
|
logger.info(
|
|
f"Set the {attribute_name} attribute of the model to {new_value} from {old_value}."
|
|
)
|
|
|
|
@classmethod
|
|
def _maybe_expand_transformer_param_shape_or_error_(
|
|
cls,
|
|
transformer: torch.nn.Module,
|
|
lora_state_dict=None,
|
|
norm_state_dict=None,
|
|
prefix=None,
|
|
) -> bool:
|
|
"""
|
|
Control LoRA expands the shape of the input layer from (3072, 64) to (3072, 128). This method handles that and
|
|
generalizes things a bit so that any parameter that needs expansion receives appropriate treatment.
|
|
"""
|
|
state_dict = {}
|
|
if lora_state_dict is not None:
|
|
state_dict.update(lora_state_dict)
|
|
if norm_state_dict is not None:
|
|
state_dict.update(norm_state_dict)
|
|
|
|
# Remove prefix if present
|
|
prefix = prefix or cls.transformer_name
|
|
for key in list(state_dict.keys()):
|
|
if key.split(".")[0] == prefix:
|
|
state_dict[key.removeprefix(f"{prefix}.")] = state_dict.pop(key)
|
|
|
|
# Expand transformer parameter shapes if they don't match lora
|
|
has_param_with_shape_update = False
|
|
overwritten_params = {}
|
|
|
|
is_peft_loaded = getattr(transformer, "peft_config", None) is not None
|
|
is_quantized = hasattr(transformer, "hf_quantizer")
|
|
for name, module in transformer.named_modules():
|
|
if isinstance(module, torch.nn.Linear):
|
|
module_weight = module.weight.data
|
|
module_bias = module.bias.data if module.bias is not None else None
|
|
bias = module_bias is not None
|
|
|
|
lora_base_name = name.replace(".base_layer", "") if is_peft_loaded else name
|
|
lora_A_weight_name = f"{lora_base_name}.lora_A.weight"
|
|
lora_B_weight_name = f"{lora_base_name}.lora_B.weight"
|
|
if lora_A_weight_name not in state_dict:
|
|
continue
|
|
|
|
in_features = state_dict[lora_A_weight_name].shape[1]
|
|
out_features = state_dict[lora_B_weight_name].shape[0]
|
|
|
|
# Model maybe loaded with different quantization schemes which may flatten the params.
|
|
# `bitsandbytes`, for example, flatten the weights when using 4bit. 8bit bnb models
|
|
# preserve weight shape.
|
|
module_weight_shape = cls._calculate_module_shape(model=transformer, base_module=module)
|
|
|
|
# This means there's no need for an expansion in the params, so we simply skip.
|
|
if tuple(module_weight_shape) == (out_features, in_features):
|
|
continue
|
|
|
|
module_out_features, module_in_features = module_weight_shape
|
|
debug_message = ""
|
|
if in_features > module_in_features:
|
|
debug_message += (
|
|
f'Expanding the nn.Linear input/output features for module="{name}" because the provided LoRA '
|
|
f"checkpoint contains higher number of features than expected. The number of input_features will be "
|
|
f"expanded from {module_in_features} to {in_features}"
|
|
)
|
|
if out_features > module_out_features:
|
|
debug_message += (
|
|
", and the number of output features will be "
|
|
f"expanded from {module_out_features} to {out_features}."
|
|
)
|
|
else:
|
|
debug_message += "."
|
|
if debug_message:
|
|
logger.debug(debug_message)
|
|
|
|
if out_features > module_out_features or in_features > module_in_features:
|
|
has_param_with_shape_update = True
|
|
parent_module_name, _, current_module_name = name.rpartition(".")
|
|
parent_module = transformer.get_submodule(parent_module_name)
|
|
|
|
if is_quantized:
|
|
module_weight = _maybe_dequantize_weight_for_expanded_lora(transformer, module)
|
|
|
|
# TODO: consider if this layer needs to be a quantized layer as well if `is_quantized` is True.
|
|
with torch.device("meta"):
|
|
expanded_module = torch.nn.Linear(
|
|
in_features, out_features, bias=bias, dtype=module_weight.dtype
|
|
)
|
|
# Only weights are expanded and biases are not. This is because only the input dimensions
|
|
# are changed while the output dimensions remain the same. The shape of the weight tensor
|
|
# is (out_features, in_features), while the shape of bias tensor is (out_features,), which
|
|
# explains the reason why only weights are expanded.
|
|
new_weight = torch.zeros_like(
|
|
expanded_module.weight.data, device=module_weight.device, dtype=module_weight.dtype
|
|
)
|
|
slices = tuple(slice(0, dim) for dim in module_weight_shape)
|
|
new_weight[slices] = module_weight
|
|
tmp_state_dict = {"weight": new_weight}
|
|
if module_bias is not None:
|
|
tmp_state_dict["bias"] = module_bias
|
|
expanded_module.load_state_dict(tmp_state_dict, strict=True, assign=True)
|
|
|
|
setattr(parent_module, current_module_name, expanded_module)
|
|
|
|
del tmp_state_dict
|
|
|
|
if current_module_name in _MODULE_NAME_TO_ATTRIBUTE_MAP_FLUX:
|
|
attribute_name = _MODULE_NAME_TO_ATTRIBUTE_MAP_FLUX[current_module_name]
|
|
new_value = int(expanded_module.weight.data.shape[1])
|
|
old_value = getattr(transformer.config, attribute_name)
|
|
setattr(transformer.config, attribute_name, new_value)
|
|
logger.info(
|
|
f"Set the {attribute_name} attribute of the model to {new_value} from {old_value}."
|
|
)
|
|
|
|
# For `unload_lora_weights()`.
|
|
# TODO: this could lead to more memory overhead if the number of overwritten params
|
|
# are large. Should be revisited later and tackled through a `discard_original_layers` arg.
|
|
overwritten_params[f"{current_module_name}.weight"] = module_weight
|
|
if module_bias is not None:
|
|
overwritten_params[f"{current_module_name}.bias"] = module_bias
|
|
|
|
if len(overwritten_params) > 0:
|
|
transformer._overwritten_params = overwritten_params
|
|
|
|
return has_param_with_shape_update
|
|
|
|
@classmethod
|
|
def _maybe_expand_lora_state_dict(cls, transformer, lora_state_dict):
|
|
expanded_module_names = set()
|
|
transformer_state_dict = transformer.state_dict()
|
|
prefix = f"{cls.transformer_name}."
|
|
|
|
lora_module_names = [
|
|
key[: -len(".lora_A.weight")] for key in lora_state_dict if key.endswith(".lora_A.weight")
|
|
]
|
|
lora_module_names = [name[len(prefix) :] for name in lora_module_names if name.startswith(prefix)]
|
|
lora_module_names = sorted(set(lora_module_names))
|
|
transformer_module_names = sorted({name for name, _ in transformer.named_modules()})
|
|
unexpected_modules = set(lora_module_names) - set(transformer_module_names)
|
|
if unexpected_modules:
|
|
logger.debug(f"Found unexpected modules: {unexpected_modules}. These will be ignored.")
|
|
|
|
for k in lora_module_names:
|
|
if k in unexpected_modules:
|
|
continue
|
|
|
|
base_param_name = (
|
|
f"{k.replace(prefix, '')}.base_layer.weight"
|
|
if f"{k.replace(prefix, '')}.base_layer.weight" in transformer_state_dict
|
|
else f"{k.replace(prefix, '')}.weight"
|
|
)
|
|
base_weight_param = transformer_state_dict[base_param_name]
|
|
lora_A_param = lora_state_dict[f"{prefix}{k}.lora_A.weight"]
|
|
|
|
# TODO (sayakpaul): Handle the cases when we actually need to expand when using quantization.
|
|
base_module_shape = cls._calculate_module_shape(model=transformer, base_weight_param_name=base_param_name)
|
|
|
|
if base_module_shape[1] > lora_A_param.shape[1]:
|
|
shape = (lora_A_param.shape[0], base_weight_param.shape[1])
|
|
expanded_state_dict_weight = torch.zeros(shape, device=base_weight_param.device)
|
|
expanded_state_dict_weight[:, : lora_A_param.shape[1]].copy_(lora_A_param)
|
|
lora_state_dict[f"{prefix}{k}.lora_A.weight"] = expanded_state_dict_weight
|
|
expanded_module_names.add(k)
|
|
elif base_module_shape[1] < lora_A_param.shape[1]:
|
|
raise NotImplementedError(
|
|
f"This LoRA param ({k}.lora_A.weight) has an incompatible shape {lora_A_param.shape}. Please open an issue to file for a feature request - https://github.com/huggingface/diffusers/issues/new."
|
|
)
|
|
|
|
if expanded_module_names:
|
|
logger.info(
|
|
f"The following LoRA modules were zero padded to match the state dict of {cls.transformer_name}: {expanded_module_names}. Please open an issue if you think this was unexpected - https://github.com/huggingface/diffusers/issues/new."
|
|
)
|
|
|
|
return lora_state_dict
|
|
|
|
@staticmethod
|
|
def _calculate_module_shape(
|
|
model: "torch.nn.Module",
|
|
base_module: "torch.nn.Linear" = None,
|
|
base_weight_param_name: str = None,
|
|
) -> "torch.Size":
|
|
def _get_weight_shape(weight: torch.Tensor):
|
|
if weight.__class__.__name__ == "Params4bit":
|
|
return weight.quant_state.shape
|
|
elif weight.__class__.__name__ == "GGUFParameter":
|
|
return weight.quant_shape
|
|
else:
|
|
return weight.shape
|
|
|
|
if base_module is not None:
|
|
return _get_weight_shape(base_module.weight)
|
|
elif base_weight_param_name is not None:
|
|
if not base_weight_param_name.endswith(".weight"):
|
|
raise ValueError(
|
|
f"Invalid `base_weight_param_name` passed as it does not end with '.weight' {base_weight_param_name=}."
|
|
)
|
|
module_path = base_weight_param_name.rsplit(".weight", 1)[0]
|
|
submodule = get_submodule_by_name(model, module_path)
|
|
return _get_weight_shape(submodule.weight)
|
|
|
|
raise ValueError("Either `base_module` or `base_weight_param_name` must be provided.")
|
|
|
|
@staticmethod
|
|
def _prepare_outputs(state_dict, metadata, alphas=None, return_alphas=False, return_metadata=False):
|
|
outputs = [state_dict]
|
|
if return_alphas:
|
|
outputs.append(alphas)
|
|
if return_metadata:
|
|
outputs.append(metadata)
|
|
return tuple(outputs) if (return_alphas or return_metadata) else state_dict
|
|
|
|
|
|
# The reason why we subclass from `StableDiffusionLoraLoaderMixin` here is because Amused initially
|
|
# relied on `StableDiffusionLoraLoaderMixin` for its LoRA support.
|
|
class AmusedLoraLoaderMixin(StableDiffusionLoraLoaderMixin):
|
|
_lora_loadable_modules = ["transformer", "text_encoder"]
|
|
transformer_name = TRANSFORMER_NAME
|
|
text_encoder_name = TEXT_ENCODER_NAME
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.FluxLoraLoaderMixin.load_lora_into_transformer with FluxTransformer2DModel->UVit2DModel
|
|
def load_lora_into_transformer(
|
|
cls,
|
|
state_dict,
|
|
network_alphas,
|
|
transformer,
|
|
adapter_name=None,
|
|
metadata=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `transformer`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
network_alphas (`Dict[str, float]`):
|
|
The value of the network alpha used for stable learning and preventing underflow. This value has the
|
|
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
|
|
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
|
|
transformer (`UVit2DModel`):
|
|
The Transformer model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if low_cpu_mem_usage and not is_peft_version(">=", "0.13.1"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# Load the layers corresponding to transformer.
|
|
logger.info(f"Loading {cls.transformer_name}.")
|
|
transformer.load_lora_adapter(
|
|
state_dict,
|
|
network_alphas=network_alphas,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.load_lora_into_text_encoder
|
|
def load_lora_into_text_encoder(
|
|
cls,
|
|
state_dict,
|
|
network_alphas,
|
|
text_encoder,
|
|
prefix=None,
|
|
lora_scale=1.0,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `text_encoder`
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The key should be prefixed with an
|
|
additional `text_encoder` to distinguish between unet lora layers.
|
|
network_alphas (`Dict[str, float]`):
|
|
The value of the network alpha used for stable learning and preventing underflow. This value has the
|
|
same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this
|
|
link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning).
|
|
text_encoder (`CLIPTextModel`):
|
|
The text encoder model to load the LoRA layers into.
|
|
prefix (`str`):
|
|
Expected prefix of the `text_encoder` in the `state_dict`.
|
|
lora_scale (`float`):
|
|
How much to scale the output of the lora linear layer before it is added with the output of the regular
|
|
lora layer.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
_load_lora_into_text_encoder(
|
|
state_dict=state_dict,
|
|
network_alphas=network_alphas,
|
|
lora_scale=lora_scale,
|
|
text_encoder=text_encoder,
|
|
prefix=prefix,
|
|
text_encoder_name=cls.text_encoder_name,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
text_encoder_lora_layers: Dict[str, torch.nn.Module] = None,
|
|
transformer_lora_layers: Dict[str, torch.nn.Module] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the UNet and text encoder.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
unet_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `unet`.
|
|
text_encoder_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `text_encoder`. Must explicitly pass the text
|
|
encoder LoRA state dict because it comes from 🤗 Transformers.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
"""
|
|
state_dict = {}
|
|
|
|
if not (transformer_lora_layers or text_encoder_lora_layers):
|
|
raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.")
|
|
|
|
if transformer_lora_layers:
|
|
state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
|
|
|
|
if text_encoder_lora_layers:
|
|
state_dict.update(cls.pack_weights(text_encoder_lora_layers, cls.text_encoder_name))
|
|
|
|
# Save the model
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
)
|
|
|
|
|
|
class CogVideoXLoraLoaderMixin(LoraBaseMixin):
|
|
r"""
|
|
Load LoRA layers into [`CogVideoXTransformer3DModel`]. Specific to [`CogVideoXPipeline`].
|
|
"""
|
|
|
|
_lora_loadable_modules = ["transformer"]
|
|
transformer_name = TRANSFORMER_NAME
|
|
|
|
@classmethod
|
|
@validate_hf_hub_args
|
|
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict
|
|
def lora_state_dict(
|
|
cls,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Return state dict for lora weights and the network alphas.
|
|
|
|
<Tip warning={true}>
|
|
|
|
We support loading A1111 formatted LoRA checkpoints in a limited capacity.
|
|
|
|
This function is experimental and might change in the future.
|
|
|
|
</Tip>
|
|
|
|
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.
|
|
|
|
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.
|
|
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.
|
|
return_lora_metadata (`bool`, *optional*, defaults to False):
|
|
When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
|
|
|
|
"""
|
|
# Load the main state dict first which has the LoRA layers for either of
|
|
# transformer and text encoder or both.
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
force_download = kwargs.pop("force_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)
|
|
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = True
|
|
allow_pickle = True
|
|
|
|
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
|
|
|
|
state_dict, metadata = _fetch_state_dict(
|
|
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
|
|
weight_name=weight_name,
|
|
use_safetensors=use_safetensors,
|
|
local_files_only=local_files_only,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
token=token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
allow_pickle=allow_pickle,
|
|
)
|
|
|
|
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
|
|
if is_dora_scale_present:
|
|
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
|
|
logger.warning(warn_msg)
|
|
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
|
|
|
|
out = (state_dict, metadata) if return_lora_metadata else state_dict
|
|
return out
|
|
|
|
def load_lora_weights(
|
|
self,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name: Optional[str] = None,
|
|
hotswap: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and
|
|
`self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See
|
|
[`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded.
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state
|
|
dict is loaded into `self.transformer`.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
kwargs (`dict`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
kwargs["return_lora_metadata"] = True
|
|
state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
|
|
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
self.load_lora_into_transformer(
|
|
state_dict,
|
|
transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer with SD3Transformer2DModel->CogVideoXTransformer3DModel
|
|
def load_lora_into_transformer(
|
|
cls,
|
|
state_dict,
|
|
transformer,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `transformer`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
transformer (`CogVideoXTransformer3DModel`):
|
|
The Transformer model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# Load the layers corresponding to transformer.
|
|
logger.info(f"Loading {cls.transformer_name}.")
|
|
transformer.load_lora_adapter(
|
|
state_dict,
|
|
network_alphas=None,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Adapted from diffusers.loaders.lora_pipeline.StableDiffusionLoraLoaderMixin.save_lora_weights without support for text encoder
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
transformer_lora_adapter_metadata: Optional[dict] = None,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the transformer.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `transformer`.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
transformer_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the transformer to be serialized with the state dict.
|
|
"""
|
|
state_dict = {}
|
|
lora_adapter_metadata = {}
|
|
|
|
if not transformer_lora_layers:
|
|
raise ValueError("You must pass `transformer_lora_layers`.")
|
|
|
|
state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
|
|
|
|
if transformer_lora_adapter_metadata is not None:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
|
|
)
|
|
|
|
# Save the model
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
lora_adapter_metadata=lora_adapter_metadata,
|
|
)
|
|
|
|
def fuse_lora(
|
|
self,
|
|
components: List[str] = ["transformer"],
|
|
lora_scale: float = 1.0,
|
|
safe_fusing: bool = False,
|
|
adapter_names: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Fuses the LoRA parameters into the original parameters of the corresponding blocks.
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
|
|
lora_scale (`float`, defaults to 1.0):
|
|
Controls how much to influence the outputs with the LoRA parameters.
|
|
safe_fusing (`bool`, defaults to `False`):
|
|
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
|
|
adapter_names (`List[str]`, *optional*):
|
|
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
|
|
|
|
Example:
|
|
|
|
```py
|
|
from diffusers import DiffusionPipeline
|
|
import torch
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
|
).to("cuda")
|
|
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
|
|
pipeline.fuse_lora(lora_scale=0.7)
|
|
```
|
|
"""
|
|
super().fuse_lora(
|
|
components=components,
|
|
lora_scale=lora_scale,
|
|
safe_fusing=safe_fusing,
|
|
adapter_names=adapter_names,
|
|
**kwargs,
|
|
)
|
|
|
|
def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
|
|
r"""
|
|
Reverses the effect of
|
|
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
|
|
unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
|
|
"""
|
|
super().unfuse_lora(components=components, **kwargs)
|
|
|
|
|
|
class Mochi1LoraLoaderMixin(LoraBaseMixin):
|
|
r"""
|
|
Load LoRA layers into [`MochiTransformer3DModel`]. Specific to [`MochiPipeline`].
|
|
"""
|
|
|
|
_lora_loadable_modules = ["transformer"]
|
|
transformer_name = TRANSFORMER_NAME
|
|
|
|
@classmethod
|
|
@validate_hf_hub_args
|
|
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict
|
|
def lora_state_dict(
|
|
cls,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Return state dict for lora weights and the network alphas.
|
|
|
|
<Tip warning={true}>
|
|
|
|
We support loading A1111 formatted LoRA checkpoints in a limited capacity.
|
|
|
|
This function is experimental and might change in the future.
|
|
|
|
</Tip>
|
|
|
|
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.
|
|
|
|
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.
|
|
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.
|
|
return_lora_metadata (`bool`, *optional*, defaults to False):
|
|
When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
|
|
|
|
"""
|
|
# Load the main state dict first which has the LoRA layers for either of
|
|
# transformer and text encoder or both.
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
force_download = kwargs.pop("force_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)
|
|
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = True
|
|
allow_pickle = True
|
|
|
|
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
|
|
|
|
state_dict, metadata = _fetch_state_dict(
|
|
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
|
|
weight_name=weight_name,
|
|
use_safetensors=use_safetensors,
|
|
local_files_only=local_files_only,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
token=token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
allow_pickle=allow_pickle,
|
|
)
|
|
|
|
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
|
|
if is_dora_scale_present:
|
|
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
|
|
logger.warning(warn_msg)
|
|
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
|
|
|
|
out = (state_dict, metadata) if return_lora_metadata else state_dict
|
|
return out
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
|
|
def load_lora_weights(
|
|
self,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name: Optional[str] = None,
|
|
hotswap: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and
|
|
`self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See
|
|
[`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded.
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state
|
|
dict is loaded into `self.transformer`.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
kwargs (`dict`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
kwargs["return_lora_metadata"] = True
|
|
state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
|
|
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
self.load_lora_into_transformer(
|
|
state_dict,
|
|
transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer with SD3Transformer2DModel->MochiTransformer3DModel
|
|
def load_lora_into_transformer(
|
|
cls,
|
|
state_dict,
|
|
transformer,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `transformer`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
transformer (`MochiTransformer3DModel`):
|
|
The Transformer model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# Load the layers corresponding to transformer.
|
|
logger.info(f"Loading {cls.transformer_name}.")
|
|
transformer.load_lora_adapter(
|
|
state_dict,
|
|
network_alphas=None,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
transformer_lora_adapter_metadata: Optional[dict] = None,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the transformer.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `transformer`.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
transformer_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the transformer to be serialized with the state dict.
|
|
"""
|
|
state_dict = {}
|
|
lora_adapter_metadata = {}
|
|
|
|
if not transformer_lora_layers:
|
|
raise ValueError("You must pass `transformer_lora_layers`.")
|
|
|
|
state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
|
|
|
|
if transformer_lora_adapter_metadata is not None:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
|
|
)
|
|
|
|
# Save the model
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
lora_adapter_metadata=lora_adapter_metadata,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
|
|
def fuse_lora(
|
|
self,
|
|
components: List[str] = ["transformer"],
|
|
lora_scale: float = 1.0,
|
|
safe_fusing: bool = False,
|
|
adapter_names: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Fuses the LoRA parameters into the original parameters of the corresponding blocks.
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
|
|
lora_scale (`float`, defaults to 1.0):
|
|
Controls how much to influence the outputs with the LoRA parameters.
|
|
safe_fusing (`bool`, defaults to `False`):
|
|
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
|
|
adapter_names (`List[str]`, *optional*):
|
|
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
|
|
|
|
Example:
|
|
|
|
```py
|
|
from diffusers import DiffusionPipeline
|
|
import torch
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
|
).to("cuda")
|
|
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
|
|
pipeline.fuse_lora(lora_scale=0.7)
|
|
```
|
|
"""
|
|
super().fuse_lora(
|
|
components=components,
|
|
lora_scale=lora_scale,
|
|
safe_fusing=safe_fusing,
|
|
adapter_names=adapter_names,
|
|
**kwargs,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
|
|
def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
|
|
r"""
|
|
Reverses the effect of
|
|
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
|
|
unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
|
|
"""
|
|
super().unfuse_lora(components=components, **kwargs)
|
|
|
|
|
|
class LTXVideoLoraLoaderMixin(LoraBaseMixin):
|
|
r"""
|
|
Load LoRA layers into [`LTXVideoTransformer3DModel`]. Specific to [`LTXPipeline`].
|
|
"""
|
|
|
|
_lora_loadable_modules = ["transformer"]
|
|
transformer_name = TRANSFORMER_NAME
|
|
|
|
@classmethod
|
|
@validate_hf_hub_args
|
|
def lora_state_dict(
|
|
cls,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Return state dict for lora weights and the network alphas.
|
|
|
|
<Tip warning={true}>
|
|
|
|
We support loading A1111 formatted LoRA checkpoints in a limited capacity.
|
|
|
|
This function is experimental and might change in the future.
|
|
|
|
</Tip>
|
|
|
|
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.
|
|
|
|
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.
|
|
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.
|
|
return_lora_metadata (`bool`, *optional*, defaults to False):
|
|
When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
|
|
"""
|
|
# Load the main state dict first which has the LoRA layers for either of
|
|
# transformer and text encoder or both.
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
force_download = kwargs.pop("force_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)
|
|
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = True
|
|
allow_pickle = True
|
|
|
|
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
|
|
|
|
state_dict, metadata = _fetch_state_dict(
|
|
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
|
|
weight_name=weight_name,
|
|
use_safetensors=use_safetensors,
|
|
local_files_only=local_files_only,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
token=token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
allow_pickle=allow_pickle,
|
|
)
|
|
|
|
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
|
|
if is_dora_scale_present:
|
|
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
|
|
logger.warning(warn_msg)
|
|
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
|
|
|
|
is_non_diffusers_format = any(k.startswith("diffusion_model.") for k in state_dict)
|
|
if is_non_diffusers_format:
|
|
state_dict = _convert_non_diffusers_ltxv_lora_to_diffusers(state_dict)
|
|
|
|
out = (state_dict, metadata) if return_lora_metadata else state_dict
|
|
return out
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
|
|
def load_lora_weights(
|
|
self,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name: Optional[str] = None,
|
|
hotswap: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and
|
|
`self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See
|
|
[`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded.
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state
|
|
dict is loaded into `self.transformer`.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
kwargs (`dict`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
kwargs["return_lora_metadata"] = True
|
|
state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
|
|
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
self.load_lora_into_transformer(
|
|
state_dict,
|
|
transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer with SD3Transformer2DModel->LTXVideoTransformer3DModel
|
|
def load_lora_into_transformer(
|
|
cls,
|
|
state_dict,
|
|
transformer,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `transformer`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
transformer (`LTXVideoTransformer3DModel`):
|
|
The Transformer model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# Load the layers corresponding to transformer.
|
|
logger.info(f"Loading {cls.transformer_name}.")
|
|
transformer.load_lora_adapter(
|
|
state_dict,
|
|
network_alphas=None,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
transformer_lora_adapter_metadata: Optional[dict] = None,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the transformer.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `transformer`.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
transformer_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the transformer to be serialized with the state dict.
|
|
"""
|
|
state_dict = {}
|
|
lora_adapter_metadata = {}
|
|
|
|
if not transformer_lora_layers:
|
|
raise ValueError("You must pass `transformer_lora_layers`.")
|
|
|
|
state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
|
|
|
|
if transformer_lora_adapter_metadata is not None:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
|
|
)
|
|
|
|
# Save the model
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
lora_adapter_metadata=lora_adapter_metadata,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
|
|
def fuse_lora(
|
|
self,
|
|
components: List[str] = ["transformer"],
|
|
lora_scale: float = 1.0,
|
|
safe_fusing: bool = False,
|
|
adapter_names: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Fuses the LoRA parameters into the original parameters of the corresponding blocks.
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
|
|
lora_scale (`float`, defaults to 1.0):
|
|
Controls how much to influence the outputs with the LoRA parameters.
|
|
safe_fusing (`bool`, defaults to `False`):
|
|
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
|
|
adapter_names (`List[str]`, *optional*):
|
|
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
|
|
|
|
Example:
|
|
|
|
```py
|
|
from diffusers import DiffusionPipeline
|
|
import torch
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
|
).to("cuda")
|
|
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
|
|
pipeline.fuse_lora(lora_scale=0.7)
|
|
```
|
|
"""
|
|
super().fuse_lora(
|
|
components=components,
|
|
lora_scale=lora_scale,
|
|
safe_fusing=safe_fusing,
|
|
adapter_names=adapter_names,
|
|
**kwargs,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
|
|
def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
|
|
r"""
|
|
Reverses the effect of
|
|
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
|
|
unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
|
|
"""
|
|
super().unfuse_lora(components=components, **kwargs)
|
|
|
|
|
|
class SanaLoraLoaderMixin(LoraBaseMixin):
|
|
r"""
|
|
Load LoRA layers into [`SanaTransformer2DModel`]. Specific to [`SanaPipeline`].
|
|
"""
|
|
|
|
_lora_loadable_modules = ["transformer"]
|
|
transformer_name = TRANSFORMER_NAME
|
|
|
|
@classmethod
|
|
@validate_hf_hub_args
|
|
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.lora_state_dict
|
|
def lora_state_dict(
|
|
cls,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Return state dict for lora weights and the network alphas.
|
|
|
|
<Tip warning={true}>
|
|
|
|
We support loading A1111 formatted LoRA checkpoints in a limited capacity.
|
|
|
|
This function is experimental and might change in the future.
|
|
|
|
</Tip>
|
|
|
|
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.
|
|
|
|
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.
|
|
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.
|
|
return_lora_metadata (`bool`, *optional*, defaults to False):
|
|
When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
|
|
|
|
"""
|
|
# Load the main state dict first which has the LoRA layers for either of
|
|
# transformer and text encoder or both.
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
force_download = kwargs.pop("force_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)
|
|
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = True
|
|
allow_pickle = True
|
|
|
|
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
|
|
|
|
state_dict, metadata = _fetch_state_dict(
|
|
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
|
|
weight_name=weight_name,
|
|
use_safetensors=use_safetensors,
|
|
local_files_only=local_files_only,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
token=token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
allow_pickle=allow_pickle,
|
|
)
|
|
|
|
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
|
|
if is_dora_scale_present:
|
|
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
|
|
logger.warning(warn_msg)
|
|
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
|
|
|
|
out = (state_dict, metadata) if return_lora_metadata else state_dict
|
|
return out
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
|
|
def load_lora_weights(
|
|
self,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name: Optional[str] = None,
|
|
hotswap: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and
|
|
`self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See
|
|
[`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded.
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state
|
|
dict is loaded into `self.transformer`.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
kwargs (`dict`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
kwargs["return_lora_metadata"] = True
|
|
state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
|
|
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
self.load_lora_into_transformer(
|
|
state_dict,
|
|
transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer with SD3Transformer2DModel->SanaTransformer2DModel
|
|
def load_lora_into_transformer(
|
|
cls,
|
|
state_dict,
|
|
transformer,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `transformer`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
transformer (`SanaTransformer2DModel`):
|
|
The Transformer model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# Load the layers corresponding to transformer.
|
|
logger.info(f"Loading {cls.transformer_name}.")
|
|
transformer.load_lora_adapter(
|
|
state_dict,
|
|
network_alphas=None,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
transformer_lora_adapter_metadata: Optional[dict] = None,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the transformer.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `transformer`.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
transformer_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the transformer to be serialized with the state dict.
|
|
"""
|
|
state_dict = {}
|
|
lora_adapter_metadata = {}
|
|
|
|
if not transformer_lora_layers:
|
|
raise ValueError("You must pass `transformer_lora_layers`.")
|
|
|
|
state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
|
|
|
|
if transformer_lora_adapter_metadata is not None:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
|
|
)
|
|
|
|
# Save the model
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
lora_adapter_metadata=lora_adapter_metadata,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
|
|
def fuse_lora(
|
|
self,
|
|
components: List[str] = ["transformer"],
|
|
lora_scale: float = 1.0,
|
|
safe_fusing: bool = False,
|
|
adapter_names: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Fuses the LoRA parameters into the original parameters of the corresponding blocks.
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
|
|
lora_scale (`float`, defaults to 1.0):
|
|
Controls how much to influence the outputs with the LoRA parameters.
|
|
safe_fusing (`bool`, defaults to `False`):
|
|
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
|
|
adapter_names (`List[str]`, *optional*):
|
|
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
|
|
|
|
Example:
|
|
|
|
```py
|
|
from diffusers import DiffusionPipeline
|
|
import torch
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
|
).to("cuda")
|
|
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
|
|
pipeline.fuse_lora(lora_scale=0.7)
|
|
```
|
|
"""
|
|
super().fuse_lora(
|
|
components=components,
|
|
lora_scale=lora_scale,
|
|
safe_fusing=safe_fusing,
|
|
adapter_names=adapter_names,
|
|
**kwargs,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
|
|
def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
|
|
r"""
|
|
Reverses the effect of
|
|
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
|
|
unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
|
|
"""
|
|
super().unfuse_lora(components=components, **kwargs)
|
|
|
|
|
|
class HunyuanVideoLoraLoaderMixin(LoraBaseMixin):
|
|
r"""
|
|
Load LoRA layers into [`HunyuanVideoTransformer3DModel`]. Specific to [`HunyuanVideoPipeline`].
|
|
"""
|
|
|
|
_lora_loadable_modules = ["transformer"]
|
|
transformer_name = TRANSFORMER_NAME
|
|
|
|
@classmethod
|
|
@validate_hf_hub_args
|
|
def lora_state_dict(
|
|
cls,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Return state dict for lora weights and the network alphas.
|
|
|
|
<Tip warning={true}>
|
|
|
|
We support loading original format HunyuanVideo LoRA checkpoints.
|
|
|
|
This function is experimental and might change in the future.
|
|
|
|
</Tip>
|
|
|
|
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.
|
|
|
|
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.
|
|
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.
|
|
return_lora_metadata (`bool`, *optional*, defaults to False):
|
|
When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
|
|
"""
|
|
# Load the main state dict first which has the LoRA layers for either of
|
|
# transformer and text encoder or both.
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
force_download = kwargs.pop("force_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)
|
|
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = True
|
|
allow_pickle = True
|
|
|
|
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
|
|
|
|
state_dict, metadata = _fetch_state_dict(
|
|
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
|
|
weight_name=weight_name,
|
|
use_safetensors=use_safetensors,
|
|
local_files_only=local_files_only,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
token=token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
allow_pickle=allow_pickle,
|
|
)
|
|
|
|
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
|
|
if is_dora_scale_present:
|
|
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
|
|
logger.warning(warn_msg)
|
|
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
|
|
|
|
is_original_hunyuan_video = any("img_attn_qkv" in k for k in state_dict)
|
|
if is_original_hunyuan_video:
|
|
state_dict = _convert_hunyuan_video_lora_to_diffusers(state_dict)
|
|
|
|
out = (state_dict, metadata) if return_lora_metadata else state_dict
|
|
return out
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
|
|
def load_lora_weights(
|
|
self,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name: Optional[str] = None,
|
|
hotswap: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and
|
|
`self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See
|
|
[`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded.
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state
|
|
dict is loaded into `self.transformer`.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
kwargs (`dict`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
kwargs["return_lora_metadata"] = True
|
|
state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
|
|
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
self.load_lora_into_transformer(
|
|
state_dict,
|
|
transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer with SD3Transformer2DModel->HunyuanVideoTransformer3DModel
|
|
def load_lora_into_transformer(
|
|
cls,
|
|
state_dict,
|
|
transformer,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `transformer`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
transformer (`HunyuanVideoTransformer3DModel`):
|
|
The Transformer model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# Load the layers corresponding to transformer.
|
|
logger.info(f"Loading {cls.transformer_name}.")
|
|
transformer.load_lora_adapter(
|
|
state_dict,
|
|
network_alphas=None,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
transformer_lora_adapter_metadata: Optional[dict] = None,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the transformer.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `transformer`.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
transformer_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the transformer to be serialized with the state dict.
|
|
"""
|
|
state_dict = {}
|
|
lora_adapter_metadata = {}
|
|
|
|
if not transformer_lora_layers:
|
|
raise ValueError("You must pass `transformer_lora_layers`.")
|
|
|
|
state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
|
|
|
|
if transformer_lora_adapter_metadata is not None:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
|
|
)
|
|
|
|
# Save the model
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
lora_adapter_metadata=lora_adapter_metadata,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
|
|
def fuse_lora(
|
|
self,
|
|
components: List[str] = ["transformer"],
|
|
lora_scale: float = 1.0,
|
|
safe_fusing: bool = False,
|
|
adapter_names: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Fuses the LoRA parameters into the original parameters of the corresponding blocks.
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
|
|
lora_scale (`float`, defaults to 1.0):
|
|
Controls how much to influence the outputs with the LoRA parameters.
|
|
safe_fusing (`bool`, defaults to `False`):
|
|
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
|
|
adapter_names (`List[str]`, *optional*):
|
|
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
|
|
|
|
Example:
|
|
|
|
```py
|
|
from diffusers import DiffusionPipeline
|
|
import torch
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
|
).to("cuda")
|
|
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
|
|
pipeline.fuse_lora(lora_scale=0.7)
|
|
```
|
|
"""
|
|
super().fuse_lora(
|
|
components=components,
|
|
lora_scale=lora_scale,
|
|
safe_fusing=safe_fusing,
|
|
adapter_names=adapter_names,
|
|
**kwargs,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
|
|
def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
|
|
r"""
|
|
Reverses the effect of
|
|
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
|
|
unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
|
|
"""
|
|
super().unfuse_lora(components=components, **kwargs)
|
|
|
|
|
|
class Lumina2LoraLoaderMixin(LoraBaseMixin):
|
|
r"""
|
|
Load LoRA layers into [`Lumina2Transformer2DModel`]. Specific to [`Lumina2Text2ImgPipeline`].
|
|
"""
|
|
|
|
_lora_loadable_modules = ["transformer"]
|
|
transformer_name = TRANSFORMER_NAME
|
|
|
|
@classmethod
|
|
@validate_hf_hub_args
|
|
def lora_state_dict(
|
|
cls,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Return state dict for lora weights and the network alphas.
|
|
|
|
<Tip warning={true}>
|
|
|
|
We support loading A1111 formatted LoRA checkpoints in a limited capacity.
|
|
|
|
This function is experimental and might change in the future.
|
|
|
|
</Tip>
|
|
|
|
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.
|
|
|
|
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.
|
|
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.
|
|
return_lora_metadata (`bool`, *optional*, defaults to False):
|
|
When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
|
|
"""
|
|
# Load the main state dict first which has the LoRA layers for either of
|
|
# transformer and text encoder or both.
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
force_download = kwargs.pop("force_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)
|
|
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = True
|
|
allow_pickle = True
|
|
|
|
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
|
|
|
|
state_dict, metadata = _fetch_state_dict(
|
|
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
|
|
weight_name=weight_name,
|
|
use_safetensors=use_safetensors,
|
|
local_files_only=local_files_only,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
token=token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
allow_pickle=allow_pickle,
|
|
)
|
|
|
|
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
|
|
if is_dora_scale_present:
|
|
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
|
|
logger.warning(warn_msg)
|
|
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
|
|
|
|
# conversion.
|
|
non_diffusers = any(k.startswith("diffusion_model.") for k in state_dict)
|
|
if non_diffusers:
|
|
state_dict = _convert_non_diffusers_lumina2_lora_to_diffusers(state_dict)
|
|
|
|
out = (state_dict, metadata) if return_lora_metadata else state_dict
|
|
return out
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
|
|
def load_lora_weights(
|
|
self,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name: Optional[str] = None,
|
|
hotswap: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and
|
|
`self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See
|
|
[`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded.
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state
|
|
dict is loaded into `self.transformer`.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
kwargs (`dict`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
kwargs["return_lora_metadata"] = True
|
|
state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
|
|
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
self.load_lora_into_transformer(
|
|
state_dict,
|
|
transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer with SD3Transformer2DModel->Lumina2Transformer2DModel
|
|
def load_lora_into_transformer(
|
|
cls,
|
|
state_dict,
|
|
transformer,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `transformer`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
transformer (`Lumina2Transformer2DModel`):
|
|
The Transformer model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# Load the layers corresponding to transformer.
|
|
logger.info(f"Loading {cls.transformer_name}.")
|
|
transformer.load_lora_adapter(
|
|
state_dict,
|
|
network_alphas=None,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
transformer_lora_adapter_metadata: Optional[dict] = None,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the transformer.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `transformer`.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
transformer_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the transformer to be serialized with the state dict.
|
|
"""
|
|
state_dict = {}
|
|
lora_adapter_metadata = {}
|
|
|
|
if not transformer_lora_layers:
|
|
raise ValueError("You must pass `transformer_lora_layers`.")
|
|
|
|
state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
|
|
|
|
if transformer_lora_adapter_metadata is not None:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
|
|
)
|
|
|
|
# Save the model
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
lora_adapter_metadata=lora_adapter_metadata,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora
|
|
def fuse_lora(
|
|
self,
|
|
components: List[str] = ["transformer"],
|
|
lora_scale: float = 1.0,
|
|
safe_fusing: bool = False,
|
|
adapter_names: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Fuses the LoRA parameters into the original parameters of the corresponding blocks.
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
|
|
lora_scale (`float`, defaults to 1.0):
|
|
Controls how much to influence the outputs with the LoRA parameters.
|
|
safe_fusing (`bool`, defaults to `False`):
|
|
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
|
|
adapter_names (`List[str]`, *optional*):
|
|
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
|
|
|
|
Example:
|
|
|
|
```py
|
|
from diffusers import DiffusionPipeline
|
|
import torch
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
|
).to("cuda")
|
|
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
|
|
pipeline.fuse_lora(lora_scale=0.7)
|
|
```
|
|
"""
|
|
super().fuse_lora(
|
|
components=components,
|
|
lora_scale=lora_scale,
|
|
safe_fusing=safe_fusing,
|
|
adapter_names=adapter_names,
|
|
**kwargs,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora
|
|
def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
|
|
r"""
|
|
Reverses the effect of
|
|
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
|
|
unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
|
|
"""
|
|
super().unfuse_lora(components=components, **kwargs)
|
|
|
|
|
|
class WanLoraLoaderMixin(LoraBaseMixin):
|
|
r"""
|
|
Load LoRA layers into [`WanTransformer3DModel`]. Specific to [`WanPipeline`] and `[WanImageToVideoPipeline`].
|
|
"""
|
|
|
|
_lora_loadable_modules = ["transformer"]
|
|
transformer_name = TRANSFORMER_NAME
|
|
|
|
@classmethod
|
|
@validate_hf_hub_args
|
|
def lora_state_dict(
|
|
cls,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Return state dict for lora weights and the network alphas.
|
|
|
|
<Tip warning={true}>
|
|
|
|
We support loading A1111 formatted LoRA checkpoints in a limited capacity.
|
|
|
|
This function is experimental and might change in the future.
|
|
|
|
</Tip>
|
|
|
|
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.
|
|
|
|
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.
|
|
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.
|
|
return_lora_metadata (`bool`, *optional*, defaults to False):
|
|
When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
|
|
"""
|
|
# Load the main state dict first which has the LoRA layers for either of
|
|
# transformer and text encoder or both.
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
force_download = kwargs.pop("force_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)
|
|
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = True
|
|
allow_pickle = True
|
|
|
|
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
|
|
|
|
state_dict, metadata = _fetch_state_dict(
|
|
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
|
|
weight_name=weight_name,
|
|
use_safetensors=use_safetensors,
|
|
local_files_only=local_files_only,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
token=token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
allow_pickle=allow_pickle,
|
|
)
|
|
if any(k.startswith("diffusion_model.") for k in state_dict):
|
|
state_dict = _convert_non_diffusers_wan_lora_to_diffusers(state_dict)
|
|
elif any(k.startswith("lora_unet_") for k in state_dict):
|
|
state_dict = _convert_musubi_wan_lora_to_diffusers(state_dict)
|
|
|
|
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
|
|
if is_dora_scale_present:
|
|
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
|
|
logger.warning(warn_msg)
|
|
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
|
|
|
|
out = (state_dict, metadata) if return_lora_metadata else state_dict
|
|
return out
|
|
|
|
@classmethod
|
|
def _maybe_expand_t2v_lora_for_i2v(
|
|
cls,
|
|
transformer: torch.nn.Module,
|
|
state_dict,
|
|
):
|
|
if transformer.config.image_dim is None:
|
|
return state_dict
|
|
|
|
target_device = transformer.device
|
|
|
|
if any(k.startswith("transformer.blocks.") for k in state_dict):
|
|
num_blocks = len({k.split("blocks.")[1].split(".")[0] for k in state_dict if "blocks." in k})
|
|
is_i2v_lora = any("add_k_proj" in k for k in state_dict) and any("add_v_proj" in k for k in state_dict)
|
|
has_bias = any(".lora_B.bias" in k for k in state_dict)
|
|
|
|
if is_i2v_lora:
|
|
return state_dict
|
|
|
|
for i in range(num_blocks):
|
|
for o, c in zip(["k_img", "v_img"], ["add_k_proj", "add_v_proj"]):
|
|
# These keys should exist if the block `i` was part of the T2V LoRA.
|
|
ref_key_lora_A = f"transformer.blocks.{i}.attn2.to_k.lora_A.weight"
|
|
ref_key_lora_B = f"transformer.blocks.{i}.attn2.to_k.lora_B.weight"
|
|
|
|
if ref_key_lora_A not in state_dict or ref_key_lora_B not in state_dict:
|
|
continue
|
|
|
|
state_dict[f"transformer.blocks.{i}.attn2.{c}.lora_A.weight"] = torch.zeros_like(
|
|
state_dict[f"transformer.blocks.{i}.attn2.to_k.lora_A.weight"], device=target_device
|
|
)
|
|
state_dict[f"transformer.blocks.{i}.attn2.{c}.lora_B.weight"] = torch.zeros_like(
|
|
state_dict[f"transformer.blocks.{i}.attn2.to_k.lora_B.weight"], device=target_device
|
|
)
|
|
|
|
# If the original LoRA had biases (indicated by has_bias)
|
|
# AND the specific reference bias key exists for this block.
|
|
|
|
ref_key_lora_B_bias = f"transformer.blocks.{i}.attn2.to_k.lora_B.bias"
|
|
if has_bias and ref_key_lora_B_bias in state_dict:
|
|
ref_lora_B_bias_tensor = state_dict[ref_key_lora_B_bias]
|
|
state_dict[f"transformer.blocks.{i}.attn2.{c}.lora_B.bias"] = torch.zeros_like(
|
|
ref_lora_B_bias_tensor,
|
|
device=target_device,
|
|
)
|
|
|
|
return state_dict
|
|
|
|
def load_lora_weights(
|
|
self,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name: Optional[str] = None,
|
|
hotswap: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and
|
|
`self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See
|
|
[`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded.
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state
|
|
dict is loaded into `self.transformer`.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
kwargs (`dict`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
kwargs["return_lora_metadata"] = True
|
|
state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
|
|
# convert T2V LoRA to I2V LoRA (when loaded to Wan I2V) by adding zeros for the additional (missing) _img layers
|
|
state_dict = self._maybe_expand_t2v_lora_for_i2v(
|
|
transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
|
|
state_dict=state_dict,
|
|
)
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
self.load_lora_into_transformer(
|
|
state_dict,
|
|
transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer with SD3Transformer2DModel->WanTransformer3DModel
|
|
def load_lora_into_transformer(
|
|
cls,
|
|
state_dict,
|
|
transformer,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `transformer`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
transformer (`WanTransformer3DModel`):
|
|
The Transformer model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# Load the layers corresponding to transformer.
|
|
logger.info(f"Loading {cls.transformer_name}.")
|
|
transformer.load_lora_adapter(
|
|
state_dict,
|
|
network_alphas=None,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
transformer_lora_adapter_metadata: Optional[dict] = None,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the transformer.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `transformer`.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
transformer_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the transformer to be serialized with the state dict.
|
|
"""
|
|
state_dict = {}
|
|
lora_adapter_metadata = {}
|
|
|
|
if not transformer_lora_layers:
|
|
raise ValueError("You must pass `transformer_lora_layers`.")
|
|
|
|
state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
|
|
|
|
if transformer_lora_adapter_metadata is not None:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
|
|
)
|
|
|
|
# Save the model
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
lora_adapter_metadata=lora_adapter_metadata,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
|
|
def fuse_lora(
|
|
self,
|
|
components: List[str] = ["transformer"],
|
|
lora_scale: float = 1.0,
|
|
safe_fusing: bool = False,
|
|
adapter_names: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Fuses the LoRA parameters into the original parameters of the corresponding blocks.
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
|
|
lora_scale (`float`, defaults to 1.0):
|
|
Controls how much to influence the outputs with the LoRA parameters.
|
|
safe_fusing (`bool`, defaults to `False`):
|
|
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
|
|
adapter_names (`List[str]`, *optional*):
|
|
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
|
|
|
|
Example:
|
|
|
|
```py
|
|
from diffusers import DiffusionPipeline
|
|
import torch
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
|
).to("cuda")
|
|
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
|
|
pipeline.fuse_lora(lora_scale=0.7)
|
|
```
|
|
"""
|
|
super().fuse_lora(
|
|
components=components,
|
|
lora_scale=lora_scale,
|
|
safe_fusing=safe_fusing,
|
|
adapter_names=adapter_names,
|
|
**kwargs,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
|
|
def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
|
|
r"""
|
|
Reverses the effect of
|
|
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
|
|
unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
|
|
"""
|
|
super().unfuse_lora(components=components, **kwargs)
|
|
|
|
|
|
class CogView4LoraLoaderMixin(LoraBaseMixin):
|
|
r"""
|
|
Load LoRA layers into [`WanTransformer3DModel`]. Specific to [`CogView4Pipeline`].
|
|
"""
|
|
|
|
_lora_loadable_modules = ["transformer"]
|
|
transformer_name = TRANSFORMER_NAME
|
|
|
|
@classmethod
|
|
@validate_hf_hub_args
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.lora_state_dict
|
|
def lora_state_dict(
|
|
cls,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Return state dict for lora weights and the network alphas.
|
|
|
|
<Tip warning={true}>
|
|
|
|
We support loading A1111 formatted LoRA checkpoints in a limited capacity.
|
|
|
|
This function is experimental and might change in the future.
|
|
|
|
</Tip>
|
|
|
|
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.
|
|
|
|
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.
|
|
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.
|
|
return_lora_metadata (`bool`, *optional*, defaults to False):
|
|
When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
|
|
|
|
"""
|
|
# Load the main state dict first which has the LoRA layers for either of
|
|
# transformer and text encoder or both.
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
force_download = kwargs.pop("force_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)
|
|
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = True
|
|
allow_pickle = True
|
|
|
|
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
|
|
|
|
state_dict, metadata = _fetch_state_dict(
|
|
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
|
|
weight_name=weight_name,
|
|
use_safetensors=use_safetensors,
|
|
local_files_only=local_files_only,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
token=token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
allow_pickle=allow_pickle,
|
|
)
|
|
|
|
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
|
|
if is_dora_scale_present:
|
|
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
|
|
logger.warning(warn_msg)
|
|
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
|
|
|
|
out = (state_dict, metadata) if return_lora_metadata else state_dict
|
|
return out
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
|
|
def load_lora_weights(
|
|
self,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name: Optional[str] = None,
|
|
hotswap: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and
|
|
`self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See
|
|
[`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded.
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state
|
|
dict is loaded into `self.transformer`.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
kwargs (`dict`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
kwargs["return_lora_metadata"] = True
|
|
state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
|
|
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
self.load_lora_into_transformer(
|
|
state_dict,
|
|
transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer with SD3Transformer2DModel->CogView4Transformer2DModel
|
|
def load_lora_into_transformer(
|
|
cls,
|
|
state_dict,
|
|
transformer,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `transformer`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
transformer (`CogView4Transformer2DModel`):
|
|
The Transformer model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# Load the layers corresponding to transformer.
|
|
logger.info(f"Loading {cls.transformer_name}.")
|
|
transformer.load_lora_adapter(
|
|
state_dict,
|
|
network_alphas=None,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
transformer_lora_adapter_metadata: Optional[dict] = None,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the transformer.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `transformer`.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
transformer_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the transformer to be serialized with the state dict.
|
|
"""
|
|
state_dict = {}
|
|
lora_adapter_metadata = {}
|
|
|
|
if not transformer_lora_layers:
|
|
raise ValueError("You must pass `transformer_lora_layers`.")
|
|
|
|
state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
|
|
|
|
if transformer_lora_adapter_metadata is not None:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
|
|
)
|
|
|
|
# Save the model
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
lora_adapter_metadata=lora_adapter_metadata,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora
|
|
def fuse_lora(
|
|
self,
|
|
components: List[str] = ["transformer"],
|
|
lora_scale: float = 1.0,
|
|
safe_fusing: bool = False,
|
|
adapter_names: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Fuses the LoRA parameters into the original parameters of the corresponding blocks.
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
|
|
lora_scale (`float`, defaults to 1.0):
|
|
Controls how much to influence the outputs with the LoRA parameters.
|
|
safe_fusing (`bool`, defaults to `False`):
|
|
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
|
|
adapter_names (`List[str]`, *optional*):
|
|
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
|
|
|
|
Example:
|
|
|
|
```py
|
|
from diffusers import DiffusionPipeline
|
|
import torch
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
|
).to("cuda")
|
|
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
|
|
pipeline.fuse_lora(lora_scale=0.7)
|
|
```
|
|
"""
|
|
super().fuse_lora(
|
|
components=components,
|
|
lora_scale=lora_scale,
|
|
safe_fusing=safe_fusing,
|
|
adapter_names=adapter_names,
|
|
**kwargs,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora
|
|
def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
|
|
r"""
|
|
Reverses the effect of
|
|
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
|
|
unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
|
|
"""
|
|
super().unfuse_lora(components=components, **kwargs)
|
|
|
|
|
|
class HiDreamImageLoraLoaderMixin(LoraBaseMixin):
|
|
r"""
|
|
Load LoRA layers into [`HiDreamImageTransformer2DModel`]. Specific to [`HiDreamImagePipeline`].
|
|
"""
|
|
|
|
_lora_loadable_modules = ["transformer"]
|
|
transformer_name = TRANSFORMER_NAME
|
|
|
|
@classmethod
|
|
@validate_hf_hub_args
|
|
def lora_state_dict(
|
|
cls,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Return state dict for lora weights and the network alphas.
|
|
|
|
<Tip warning={true}>
|
|
|
|
We support loading A1111 formatted LoRA checkpoints in a limited capacity.
|
|
|
|
This function is experimental and might change in the future.
|
|
|
|
</Tip>
|
|
|
|
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.
|
|
|
|
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.
|
|
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.
|
|
return_lora_metadata (`bool`, *optional*, defaults to False):
|
|
When enabled, additionally return the LoRA adapter metadata, typically found in the state dict.
|
|
"""
|
|
# Load the main state dict first which has the LoRA layers for either of
|
|
# transformer and text encoder or both.
|
|
cache_dir = kwargs.pop("cache_dir", None)
|
|
force_download = kwargs.pop("force_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)
|
|
return_lora_metadata = kwargs.pop("return_lora_metadata", False)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = True
|
|
allow_pickle = True
|
|
|
|
user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"}
|
|
|
|
state_dict, metadata = _fetch_state_dict(
|
|
pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict,
|
|
weight_name=weight_name,
|
|
use_safetensors=use_safetensors,
|
|
local_files_only=local_files_only,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
token=token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
allow_pickle=allow_pickle,
|
|
)
|
|
|
|
is_dora_scale_present = any("dora_scale" in k for k in state_dict)
|
|
if is_dora_scale_present:
|
|
warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new."
|
|
logger.warning(warn_msg)
|
|
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
|
|
|
|
is_non_diffusers_format = any("diffusion_model" in k for k in state_dict)
|
|
if is_non_diffusers_format:
|
|
state_dict = _convert_non_diffusers_hidream_lora_to_diffusers(state_dict)
|
|
|
|
out = (state_dict, metadata) if return_lora_metadata else state_dict
|
|
return out
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights
|
|
def load_lora_weights(
|
|
self,
|
|
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
|
|
adapter_name: Optional[str] = None,
|
|
hotswap: bool = False,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Load LoRA weights specified in `pretrained_model_name_or_path_or_dict` into `self.transformer` and
|
|
`self.text_encoder`. All kwargs are forwarded to `self.lora_state_dict`. See
|
|
[`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded.
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_transformer`] for more details on how the state
|
|
dict is loaded into `self.transformer`.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
kwargs (`dict`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`].
|
|
"""
|
|
if not USE_PEFT_BACKEND:
|
|
raise ValueError("PEFT backend is required for this method.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA)
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# if a dict is passed, copy it instead of modifying it inplace
|
|
if isinstance(pretrained_model_name_or_path_or_dict, dict):
|
|
pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy()
|
|
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
kwargs["return_lora_metadata"] = True
|
|
state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
|
|
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
self.load_lora_into_transformer(
|
|
state_dict,
|
|
transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=self,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer with SD3Transformer2DModel->HiDreamImageTransformer2DModel
|
|
def load_lora_into_transformer(
|
|
cls,
|
|
state_dict,
|
|
transformer,
|
|
adapter_name=None,
|
|
_pipeline=None,
|
|
low_cpu_mem_usage=False,
|
|
hotswap: bool = False,
|
|
metadata=None,
|
|
):
|
|
"""
|
|
This will load the LoRA layers specified in `state_dict` into `transformer`.
|
|
|
|
Parameters:
|
|
state_dict (`dict`):
|
|
A standard state dict containing the lora layer parameters. The keys can either be indexed directly
|
|
into the unet or prefixed with an additional `unet` which can be used to distinguish between text
|
|
encoder lora layers.
|
|
transformer (`HiDreamImageTransformer2DModel`):
|
|
The Transformer model to load the LoRA layers into.
|
|
adapter_name (`str`, *optional*):
|
|
Adapter name to be used for referencing the loaded adapter model. If not specified, it will use
|
|
`default_{i}` where i is the total number of adapters being loaded.
|
|
low_cpu_mem_usage (`bool`, *optional*):
|
|
Speed up model loading by only loading the pretrained LoRA weights and not initializing the random
|
|
weights.
|
|
hotswap (`bool`, *optional*):
|
|
See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`].
|
|
metadata (`dict`):
|
|
Optional LoRA adapter metadata. When supplied, the `LoraConfig` arguments of `peft` won't be derived
|
|
from the state dict.
|
|
"""
|
|
if low_cpu_mem_usage and is_peft_version("<", "0.13.0"):
|
|
raise ValueError(
|
|
"`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`."
|
|
)
|
|
|
|
# Load the layers corresponding to transformer.
|
|
logger.info(f"Loading {cls.transformer_name}.")
|
|
transformer.load_lora_adapter(
|
|
state_dict,
|
|
network_alphas=None,
|
|
adapter_name=adapter_name,
|
|
metadata=metadata,
|
|
_pipeline=_pipeline,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
hotswap=hotswap,
|
|
)
|
|
|
|
@classmethod
|
|
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights
|
|
def save_lora_weights(
|
|
cls,
|
|
save_directory: Union[str, os.PathLike],
|
|
transformer_lora_layers: Dict[str, Union[torch.nn.Module, torch.Tensor]] = None,
|
|
is_main_process: bool = True,
|
|
weight_name: str = None,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = True,
|
|
transformer_lora_adapter_metadata: Optional[dict] = None,
|
|
):
|
|
r"""
|
|
Save the LoRA parameters corresponding to the transformer.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to save LoRA parameters to. Will be created if it doesn't exist.
|
|
transformer_lora_layers (`Dict[str, torch.nn.Module]` or `Dict[str, torch.Tensor]`):
|
|
State dict of the LoRA layers corresponding to the `transformer`.
|
|
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 the traditional PyTorch way with `pickle`.
|
|
transformer_lora_adapter_metadata:
|
|
LoRA adapter metadata associated with the transformer to be serialized with the state dict.
|
|
"""
|
|
state_dict = {}
|
|
lora_adapter_metadata = {}
|
|
|
|
if not transformer_lora_layers:
|
|
raise ValueError("You must pass `transformer_lora_layers`.")
|
|
|
|
state_dict.update(cls.pack_weights(transformer_lora_layers, cls.transformer_name))
|
|
|
|
if transformer_lora_adapter_metadata is not None:
|
|
lora_adapter_metadata.update(
|
|
_pack_dict_with_prefix(transformer_lora_adapter_metadata, cls.transformer_name)
|
|
)
|
|
|
|
# Save the model
|
|
cls.write_lora_layers(
|
|
state_dict=state_dict,
|
|
save_directory=save_directory,
|
|
is_main_process=is_main_process,
|
|
weight_name=weight_name,
|
|
save_function=save_function,
|
|
safe_serialization=safe_serialization,
|
|
lora_adapter_metadata=lora_adapter_metadata,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.fuse_lora
|
|
def fuse_lora(
|
|
self,
|
|
components: List[str] = ["transformer"],
|
|
lora_scale: float = 1.0,
|
|
safe_fusing: bool = False,
|
|
adapter_names: Optional[List[str]] = None,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Fuses the LoRA parameters into the original parameters of the corresponding blocks.
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components: (`List[str]`): List of LoRA-injectable components to fuse the LoRAs into.
|
|
lora_scale (`float`, defaults to 1.0):
|
|
Controls how much to influence the outputs with the LoRA parameters.
|
|
safe_fusing (`bool`, defaults to `False`):
|
|
Whether to check fused weights for NaN values before fusing and if values are NaN not fusing them.
|
|
adapter_names (`List[str]`, *optional*):
|
|
Adapter names to be used for fusing. If nothing is passed, all active adapters will be fused.
|
|
|
|
Example:
|
|
|
|
```py
|
|
from diffusers import DiffusionPipeline
|
|
import torch
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
|
|
).to("cuda")
|
|
pipeline.load_lora_weights("nerijs/pixel-art-xl", weight_name="pixel-art-xl.safetensors", adapter_name="pixel")
|
|
pipeline.fuse_lora(lora_scale=0.7)
|
|
```
|
|
"""
|
|
super().fuse_lora(
|
|
components=components,
|
|
lora_scale=lora_scale,
|
|
safe_fusing=safe_fusing,
|
|
adapter_names=adapter_names,
|
|
**kwargs,
|
|
)
|
|
|
|
# Copied from diffusers.loaders.lora_pipeline.SanaLoraLoaderMixin.unfuse_lora
|
|
def unfuse_lora(self, components: List[str] = ["transformer"], **kwargs):
|
|
r"""
|
|
Reverses the effect of
|
|
[`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora).
|
|
|
|
<Tip warning={true}>
|
|
|
|
This is an experimental API.
|
|
|
|
</Tip>
|
|
|
|
Args:
|
|
components (`List[str]`): List of LoRA-injectable components to unfuse LoRA from.
|
|
unfuse_transformer (`bool`, defaults to `True`): Whether to unfuse the UNet LoRA parameters.
|
|
"""
|
|
super().unfuse_lora(components=components, **kwargs)
|
|
|
|
|
|
class LoraLoaderMixin(StableDiffusionLoraLoaderMixin):
|
|
def __init__(self, *args, **kwargs):
|
|
deprecation_message = "LoraLoaderMixin is deprecated and this will be removed in a future version. Please use `StableDiffusionLoraLoaderMixin`, instead."
|
|
deprecate("LoraLoaderMixin", "1.0.0", deprecation_message)
|
|
super().__init__(*args, **kwargs)
|