# Copyright 2023 The HuggingFace Team. All rights reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import copy

import os

import re

import warnings

from collections import defaultdict

from contextlib import nullcontext

from io import BytesIO

from pathlib import Path

from typing import Callable, Dict, List, Optional, Union

import requests

import safetensors

import torch

import torch.nn.functional as F

from huggingface_hub import hf_hub_download

from torch import nn

from .utils import (

DIFFUSERS_CACHE,

HF_HUB_OFFLINE,

_get_model_file,

deprecate,

is_accelerate_available,

is_omegaconf_available,

is_transformers_available,

logging,

)

from .utils.import_utils import BACKENDS_MAPPING

if is_transformers_available():

from transformers import CLIPTextModel, CLIPTextModelWithProjection, PreTrainedModel, PreTrainedTokenizer

if is_accelerate_available():

from accelerate import init_empty_weights

from accelerate.utils import set_module_tensor_to_device

logger = logging.get_logger(__name__)

TEXT_ENCODER_NAME = "text_encoder"

UNET_NAME = "unet"

LORA_WEIGHT_NAME = "pytorch_lora_weights.bin"

LORA_WEIGHT_NAME_SAFE = "pytorch_lora_weights.safetensors"

TEXT_INVERSION_NAME = "learned_embeds.bin"

TEXT_INVERSION_NAME_SAFE = "learned_embeds.safetensors"

CUSTOM_DIFFUSION_WEIGHT_NAME = "pytorch_custom_diffusion_weights.bin"

CUSTOM_DIFFUSION_WEIGHT_NAME_SAFE = "pytorch_custom_diffusion_weights.safetensors"

class PatchedLoraProjection(nn.Module):

def __init__(self, regular_linear_layer, lora_scale=1, network_alpha=None, rank=4, dtype=None):

super().__init__()

from .models.lora import LoRALinearLayer

self.regular_linear_layer = regular_linear_layer

device = self.regular_linear_layer.weight.device

if dtype is None:

dtype = self.regular_linear_layer.weight.dtype

self.lora_linear_layer = LoRALinearLayer(

self.regular_linear_layer.in_features,

self.regular_linear_layer.out_features,

network_alpha=network_alpha,

device=device,

dtype=dtype,

rank=rank,

)

self.lora_scale = lora_scale

def forward(self, input):

return self.regular_linear_layer(input) + self.lora_scale * self.lora_linear_layer(input)

def text_encoder_attn_modules(text_encoder):

attn_modules = []

if isinstance(text_encoder, (CLIPTextModel, CLIPTextModelWithProjection)):

for i, layer in enumerate(text_encoder.text_model.encoder.layers):

name = f"text_model.encoder.layers.{i}.self_attn"

mod = layer.self_attn

attn_modules.append((name, mod))

else:

raise ValueError(f"do not know how to get attention modules for: {text_encoder.__class__.__name__}")

return attn_modules

def text_encoder_mlp_modules(text_encoder):

mlp_modules = []

if isinstance(text_encoder, (CLIPTextModel, CLIPTextModelWithProjection)):

for i, layer in enumerate(text_encoder.text_model.encoder.layers):

mlp_mod = layer.mlp

name = f"text_model.encoder.layers.{i}.mlp"

mlp_modules.append((name, mlp_mod))

else:

raise ValueError(f"do not know how to get mlp modules for: {text_encoder.__class__.__name__}")

return mlp_modules

def text_encoder_lora_state_dict(text_encoder):

state_dict = {}

for name, module in text_encoder_attn_modules(text_encoder):

for k, v in module.q_proj.lora_linear_layer.state_dict().items():

state_dict[f"{name}.q_proj.lora_linear_layer.{k}"] = v

for k, v in module.k_proj.lora_linear_layer.state_dict().items():

state_dict[f"{name}.k_proj.lora_linear_layer.{k}"] = v

for k, v in module.v_proj.lora_linear_layer.state_dict().items():

state_dict[f"{name}.v_proj.lora_linear_layer.{k}"] = v

for k, v in module.out_proj.lora_linear_layer.state_dict().items():

state_dict[f"{name}.out_proj.lora_linear_layer.{k}"] = v

return state_dict

class AttnProcsLayers(torch.nn.Module):

def __init__(self, state_dict: Dict[str, torch.Tensor]):

super().__init__()

self.layers = torch.nn.ModuleList(state_dict.values())

self.mapping = dict(enumerate(state_dict.keys()))

self.rev_mapping = {v: k for k, v in enumerate(state_dict.keys())}

# .processor for unet, .self_attn for text encoder

self.split_keys = [".processor", ".self_attn"]

# we add a hook to state_dict() and load_state_dict() so that the

# naming fits with `unet.attn_processors`

def map_to(module, state_dict, *args, **kwargs):

new_state_dict = {}

for key, value in state_dict.items():

num = int(key.split(".")[1]) # 0 is always "layers"

new_key = key.replace(f"layers.{num}", module.mapping[num])

new_state_dict[new_key] = value

return new_state_dict

def remap_key(key, state_dict):

for k in self.split_keys:

if k in key:

return key.split(k)[0] + k

raise ValueError(

f"There seems to be a problem with the state_dict: {set(state_dict.keys())}. {key} has to have one of {self.split_keys}."

)

def map_from(module, state_dict, *args, **kwargs):

all_keys = list(state_dict.keys())

for key in all_keys:

replace_key = remap_key(key, state_dict)

new_key = key.replace(replace_key, f"layers.{module.rev_mapping[replace_key]}")

state_dict[new_key] = state_dict[key]

del state_dict[key]

self._register_state_dict_hook(map_to)

self._register_load_state_dict_pre_hook(map_from, with_module=True)

class UNet2DConditionLoadersMixin:

text_encoder_name = TEXT_ENCODER_NAME

unet_name = UNET_NAME

def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs):

r"""

Load pretrained attention processor layers into [`UNet2DConditionModel`]. Attention processor layers have to be

defined in

[`attention_processor.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py)

and be a `torch.nn.Module` class.

Parameters:

pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):

Can be either:

- A string, the model id (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on

the Hub.

- A path to a directory (for example `./my_model_directory`) containing the model weights saved

with [`ModelMixin.save_pretrained`].

- A [torch state

dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).

cache_dir (`Union[str, os.PathLike]`, *optional*):

Path to a directory where a downloaded pretrained model configuration is cached if the standard cache

is not used.

force_download (`bool`, *optional*, defaults to `False`):

Whether or not to force the (re-)download of the model weights and configuration files, overriding the

cached versions if they exist.

resume_download (`bool`, *optional*, defaults to `False`):

Whether or not to resume downloading the model weights and configuration files. If set to `False`, any

incompletely downloaded files are deleted.

proxies (`Dict[str, str]`, *optional*):

A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',

'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.

local_files_only (`bool`, *optional*, defaults to `False`):

Whether to only load local model weights and configuration files or not. If set to `True`, the model

won't be downloaded from the Hub.

use_auth_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.

mirror (`str`, *optional*):

Mirror source to resolve accessibility issues if you’re downloading a model in China. We do not

guarantee the timeliness or safety of the source, and you should refer to the mirror site for more

information.

"""

from .models.attention_processor import (

AttnAddedKVProcessor,

AttnAddedKVProcessor2_0,

CustomDiffusionAttnProcessor,

LoRAAttnAddedKVProcessor,

LoRAAttnProcessor,

LoRAAttnProcessor2_0,

LoRAXFormersAttnProcessor,

SlicedAttnAddedKVProcessor,

XFormersAttnProcessor,

)

from .models.lora import LoRACompatibleConv, LoRACompatibleLinear, LoRAConv2dLayer, LoRALinearLayer

cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)

force_download = kwargs.pop("force_download", False)

resume_download = kwargs.pop("resume_download", False)

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

local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE)

use_auth_token = kwargs.pop("use_auth_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)

# This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script.

# See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning

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

is_network_alphas_none = network_alphas is None

allow_pickle = False

if use_safetensors is None:

use_safetensors = True

allow_pickle = True

user_agent = {

"file_type": "attn_procs_weights",

"framework": "pytorch",

}

model_file = None

if not isinstance(pretrained_model_name_or_path_or_dict, dict):

# Let's first try to load .safetensors weights

if (use_safetensors and weight_name is None) or (

weight_name is not None and weight_name.endswith(".safetensors")

):

try:

model_file = _get_model_file(

pretrained_model_name_or_path_or_dict,

weights_name=weight_name or LORA_WEIGHT_NAME_SAFE,

cache_dir=cache_dir,

force_download=force_download,

resume_download=resume_download,

proxies=proxies,

local_files_only=local_files_only,

use_auth_token=use_auth_token,

revision=revision,

subfolder=subfolder,

user_agent=user_agent,

)

state_dict = safetensors.torch.load_file(model_file, device="cpu")

except IOError as e:

if not allow_pickle:

raise e

# try loading non-safetensors weights

pass

if model_file is None:

model_file = _get_model_file(

pretrained_model_name_or_path_or_dict,

weights_name=weight_name or LORA_WEIGHT_NAME,

cache_dir=cache_dir,

force_download=force_download,

resume_download=resume_download,

proxies=proxies,

local_files_only=local_files_only,

use_auth_token=use_auth_token,

revision=revision,

subfolder=subfolder,

user_agent=user_agent,

)

state_dict = torch.load(model_file, map_location="cpu")

else:

state_dict = pretrained_model_name_or_path_or_dict

# fill attn processors

attn_processors = {}

non_attn_lora_layers = []

is_lora = all(("lora" in k or k.endswith(".alpha")) for k in state_dict.keys())

is_custom_diffusion = any("custom_diffusion" in k for k in state_dict.keys())

if is_lora:

is_new_lora_format = all(

key.startswith(self.unet_name) or key.startswith(self.text_encoder_name) for key in state_dict.keys()

)

if is_new_lora_format:

# Strip the `"unet"` prefix.

is_text_encoder_present = any(key.startswith(self.text_encoder_name) for key in state_dict.keys())

if is_text_encoder_present:

warn_message = "The state_dict contains LoRA params corresponding to the text encoder which are not being used here. To use both UNet and text encoder related LoRA params, use [`pipe.load_lora_weights()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraLoaderMixin.load_lora_weights)."

warnings.warn(warn_message)

unet_keys = [k for k in state_dict.keys() if k.startswith(self.unet_name)]

state_dict = {k.replace(f"{self.unet_name}.", ""): v for k, v in state_dict.items() if k in unet_keys}

lora_grouped_dict = defaultdict(dict)

mapped_network_alphas = {}

all_keys = list(state_dict.keys())

for key in all_keys:

value = state_dict.pop(key)

attn_processor_key, sub_key = ".".join(key.split(".")[:-3]), ".".join(key.split(".")[-3:])

lora_grouped_dict[attn_processor_key][sub_key] = value

# Create another `mapped_network_alphas` dictionary so that we can properly map them.

if network_alphas is not None:

network_alphas_ = copy.deepcopy(network_alphas)

for k in network_alphas_:

if k.replace(".alpha", "") in key:

mapped_network_alphas.update({attn_processor_key: network_alphas.pop(k)})

if not is_network_alphas_none:

if len(network_alphas) > 0:

raise ValueError(

f"The `network_alphas` has to be empty at this point but has the following keys \n\n {', '.join(network_alphas.keys())}"

)

if len(state_dict) > 0:

raise ValueError(

f"The `state_dict` has to be empty at this point but has the following keys \n\n {', '.join(state_dict.keys())}"

)

for key, value_dict in lora_grouped_dict.items():

attn_processor = self

for sub_key in key.split("."):

attn_processor = getattr(attn_processor, sub_key)

# Process non-attention layers, which don't have to_{k,v,q,out_proj}_lora layers

# or add_{k,v,q,out_proj}_proj_lora layers.

if "lora.down.weight" in value_dict:

rank = value_dict["lora.down.weight"].shape[0]

if isinstance(attn_processor, LoRACompatibleConv):

in_features = attn_processor.in_channels

out_features = attn_processor.out_channels

kernel_size = attn_processor.kernel_size

lora = LoRAConv2dLayer(

in_features=in_features,

out_features=out_features,

rank=rank,

kernel_size=kernel_size,

stride=attn_processor.stride,

padding=attn_processor.padding,

network_alpha=mapped_network_alphas.get(key),

)

elif isinstance(attn_processor, LoRACompatibleLinear):

lora = LoRALinearLayer(

attn_processor.in_features,

attn_processor.out_features,

rank,

mapped_network_alphas.get(key),

)

else:

raise ValueError(f"Module {key} is not a LoRACompatibleConv or LoRACompatibleLinear module.")

value_dict = {k.replace("lora.", ""): v for k, v in value_dict.items()}

lora.load_state_dict(value_dict)

non_attn_lora_layers.append((attn_processor, lora))

else:

# To handle SDXL.

rank_mapping = {}

hidden_size_mapping = {}

for projection_id in ["to_k", "to_q", "to_v", "to_out"]:

rank = value_dict[f"{projection_id}_lora.down.weight"].shape[0]

hidden_size = value_dict[f"{projection_id}_lora.up.weight"].shape[0]

rank_mapping.update({f"{projection_id}_lora.down.weight": rank})

hidden_size_mapping.update({f"{projection_id}_lora.up.weight": hidden_size})

if isinstance(

attn_processor, (AttnAddedKVProcessor, SlicedAttnAddedKVProcessor, AttnAddedKVProcessor2_0)

):

cross_attention_dim = value_dict["add_k_proj_lora.down.weight"].shape[1]

attn_processor_class = LoRAAttnAddedKVProcessor

else:

cross_attention_dim = value_dict["to_k_lora.down.weight"].shape[1]

if isinstance(attn_processor, (XFormersAttnProcessor, LoRAXFormersAttnProcessor)):

attn_processor_class = LoRAXFormersAttnProcessor

else:

attn_processor_class = (

LoRAAttnProcessor2_0

if hasattr(F, "scaled_dot_product_attention")

else LoRAAttnProcessor

)

if attn_processor_class is not LoRAAttnAddedKVProcessor:

attn_processors[key] = attn_processor_class(

rank=rank_mapping.get("to_k_lora.down.weight"),

hidden_size=hidden_size_mapping.get("to_k_lora.up.weight"),

cross_attention_dim=cross_attention_dim,

network_alpha=mapped_network_alphas.get(key),

q_rank=rank_mapping.get("to_q_lora.down.weight"),

q_hidden_size=hidden_size_mapping.get("to_q_lora.up.weight"),

v_rank=rank_mapping.get("to_v_lora.down.weight"),

v_hidden_size=hidden_size_mapping.get("to_v_lora.up.weight"),

out_rank=rank_mapping.get("to_out_lora.down.weight"),

out_hidden_size=hidden_size_mapping.get("to_out_lora.up.weight"),

)

else:

attn_processors[key] = attn_processor_class(

rank=rank_mapping.get("to_k_lora.down.weight", None),

hidden_size=hidden_size_mapping.get("to_k_lora.up.weight", None),

cross_attention_dim=cross_attention_dim,

network_alpha=mapped_network_alphas.get(key),

)

attn_processors[key].load_state_dict(value_dict)

elif is_custom_diffusion:

custom_diffusion_grouped_dict = defaultdict(dict)

for key, value in state_dict.items():

if len(value) == 0:

custom_diffusion_grouped_dict[key] = {}

else:

if "to_out" in key:

attn_processor_key, sub_key = ".".join(key.split(".")[:-3]), ".".join(key.split(".")[-3:])

else:

attn_processor_key, sub_key = ".".join(key.split(".")[:-2]), ".".join(key.split(".")[-2:])

custom_diffusion_grouped_dict[attn_processor_key][sub_key] = value

for key, value_dict in custom_diffusion_grouped_dict.items():

if len(value_dict) == 0:

attn_processors[key] = CustomDiffusionAttnProcessor(

train_kv=False, train_q_out=False, hidden_size=None, cross_attention_dim=None

)

else:

cross_attention_dim = value_dict["to_k_custom_diffusion.weight"].shape[1]

hidden_size = value_dict["to_k_custom_diffusion.weight"].shape[0]

train_q_out = True if "to_q_custom_diffusion.weight" in value_dict else False

attn_processors[key] = CustomDiffusionAttnProcessor(

train_kv=True,

train_q_out=train_q_out,

hidden_size=hidden_size,

cross_attention_dim=cross_attention_dim,

)

attn_processors[key].load_state_dict(value_dict)

else:

raise ValueError(

f"{model_file} does not seem to be in the correct format expected by LoRA or Custom Diffusion training."

)

# set correct dtype & device

attn_processors = {k: v.to(device=self.device, dtype=self.dtype) for k, v in attn_processors.items()}

non_attn_lora_layers = [(t, l.to(device=self.device, dtype=self.dtype)) for t, l in non_attn_lora_layers]

# set layers

self.set_attn_processor(attn_processors)

# set ff layers

for target_module, lora_layer in non_attn_lora_layers:

target_module.set_lora_layer(lora_layer)

def save_attn_procs(

self,

save_directory: Union[str, os.PathLike],

is_main_process: bool = True,

weight_name: str = None,

save_function: Callable = None,

safe_serialization: bool = True,

**kwargs,

):

r"""

Save an attention processor to a directory so that it can be reloaded using the

[`~loaders.UNet2DConditionLoadersMixin.load_attn_procs`] method.

Arguments:

save_directory (`str` or `os.PathLike`):

Directory to save an attention processor to. Will be created if it doesn't exist.

is_main_process (`bool`, *optional*, defaults to `True`):

Whether the process calling this is the main process or not. Useful during distributed training and you

need to call this function on all processes. In this case, set `is_main_process=True` only on the main

process to avoid race conditions.

save_function (`Callable`):

The function to use to save the state dictionary. Useful during distributed training when you need to

replace `torch.save` with another method. Can be configured with the environment variable

`DIFFUSERS_SAVE_MODE`.

safe_serialization (`bool`, *optional*, defaults to `True`):

Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`.

"""

from .models.attention_processor import (

CustomDiffusionAttnProcessor,

CustomDiffusionXFormersAttnProcessor,

)

if os.path.isfile(save_directory):

logger.error(f"Provided path ({save_directory}) should be a directory, not a file")

return

if save_function is None:

if safe_serialization:

def save_function(weights, filename):

return safetensors.torch.save_file(weights, filename, metadata={"format": "pt"})

else:

save_function = torch.save

os.makedirs(save_directory, exist_ok=True)

is_custom_diffusion = any(

isinstance(x, (CustomDiffusionAttnProcessor, CustomDiffusionXFormersAttnProcessor))

for (_, x) in self.attn_processors.items()

)

if is_custom_diffusion:

model_to_save = AttnProcsLayers(

{

y: x

for (y, x) in self.attn_processors.items()

if isinstance(x, (CustomDiffusionAttnProcessor, CustomDiffusionXFormersAttnProcessor))

}

)

state_dict = model_to_save.state_dict()

for name, attn in self.attn_processors.items():

if len(attn.state_dict()) == 0:

state_dict[name] = {}

else:

model_to_save = AttnProcsLayers(self.attn_processors)

state_dict = model_to_save.state_dict()

if weight_name is None:

if safe_serialization:

weight_name = CUSTOM_DIFFUSION_WEIGHT_NAME_SAFE if is_custom_diffusion else LORA_WEIGHT_NAME_SAFE

else:

weight_name = CUSTOM_DIFFUSION_WEIGHT_NAME if is_custom_diffusion else LORA_WEIGHT_NAME

# Save the model

save_function(state_dict, os.path.join(save_directory, weight_name))

logger.info(f"Model weights saved in {os.path.join(save_directory, weight_name)}")

class TextualInversionLoaderMixin:

r"""

Load textual inversion tokens and embeddings to the tokenizer and text encoder.

"""

def maybe_convert_prompt(self, prompt: Union[str, List[str]], tokenizer: "PreTrainedTokenizer"):

r"""

Processes prompts that include a special token corresponding to a multi-vector textual inversion embedding to

be replaced with multiple special tokens each corresponding to one of the vectors. If the prompt has no textual

inversion token or if the textual inversion token is a single vector, the input prompt is returned.

Parameters:

prompt (`str` or list of `str`):

The prompt or prompts to guide the image generation.

tokenizer (`PreTrainedTokenizer`):

The tokenizer responsible for encoding the prompt into input tokens.

Returns:

`str` or list of `str`: The converted prompt

"""

if not isinstance(prompt, List):

prompts = [prompt]

else:

prompts = prompt

prompts = [self._maybe_convert_prompt(p, tokenizer) for p in prompts]

if not isinstance(prompt, List):

return prompts[0]

return prompts

def _maybe_convert_prompt(self, prompt: str, tokenizer: "PreTrainedTokenizer"):

r"""

Maybe convert a prompt into a "multi vector"-compatible prompt. If the prompt includes a token that corresponds

to a multi-vector textual inversion embedding, this function will process the prompt so that the special token

is replaced with multiple special tokens each corresponding to one of the vectors. If the prompt has no textual

inversion token or a textual inversion token that is a single vector, the input prompt is simply returned.

Parameters:

prompt (`str`):

The prompt to guide the image generation.

tokenizer (`PreTrainedTokenizer`):

The tokenizer responsible for encoding the prompt into input tokens.

Returns:

`str`: The converted prompt

"""

tokens = tokenizer.tokenize(prompt)

unique_tokens = set(tokens)

for token in unique_tokens:

if token in tokenizer.added_tokens_encoder:

replacement = token

i = 1

while f"{token}_{i}" in tokenizer.added_tokens_encoder:

replacement += f" {token}_{i}"

i += 1

prompt = prompt.replace(token, replacement)

return prompt

def load_textual_inversion(

self,

pretrained_model_name_or_path: Union[str, List[str], Dict[str, torch.Tensor], List[Dict[str, torch.Tensor]]],

token: Optional[Union[str, List[str]]] = None,

**kwargs,

):

r"""

Load textual inversion embeddings into the text encoder of [`StableDiffusionPipeline`] (both 🤗 Diffusers and

Automatic1111 formats are supported).

Parameters:

pretrained_model_name_or_path (`str` or `os.PathLike` or `List[str or os.PathLike]` or `Dict` or `List[Dict]`):

Can be either one of the following or a list of them:

- A string, the *model id* (for example `sd-concepts-library/low-poly-hd-logos-icons`) of a

pretrained model hosted on the Hub.

- A path to a *directory* (for example `./my_text_inversion_directory/`) containing the textual

inversion weights.

- A path to a *file* (for example `./my_text_inversions.pt`) containing textual inversion weights.

- A [torch state

dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).

token (`str` or `List[str]`, *optional*):

Override the token to use for the textual inversion weights. If `pretrained_model_name_or_path` is a

list, then `token` must also be a list of equal length.

weight_name (`str`, *optional*):

Name of a custom weight file. This should be used when:

- The saved textual inversion file is in 🤗 Diffusers format, but was saved under a specific weight

name such as `text_inv.bin`.

- The saved textual inversion file is in the Automatic1111 format.

cache_dir (`Union[str, os.PathLike]`, *optional*):

Path to a directory where a downloaded pretrained model configuration is cached if the standard cache

is not used.

force_download (`bool`, *optional*, defaults to `False`):

Whether or not to force the (re-)download of the model weights and configuration files, overriding the

cached versions if they exist.

resume_download (`bool`, *optional*, defaults to `False`):

Whether or not to resume downloading the model weights and configuration files. If set to `False`, any

incompletely downloaded files are deleted.

proxies (`Dict[str, str]`, *optional*):

A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',

'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.

local_files_only (`bool`, *optional*, defaults to `False`):

Whether to only load local model weights and configuration files or not. If set to `True`, the model

won't be downloaded from the Hub.

use_auth_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.

mirror (`str`, *optional*):

Mirror source to resolve accessibility issues if you're downloading a model in China. We do not

guarantee the timeliness or safety of the source, and you should refer to the mirror site for more

information.

Example:

To load a textual inversion embedding vector in 🤗 Diffusers format:

```py

from diffusers import StableDiffusionPipeline

import torch

model_id = "runwayml/stable-diffusion-v1-5"

pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")

pipe.load_textual_inversion("sd-concepts-library/cat-toy")

prompt = "A <cat-toy> backpack"

image = pipe(prompt, num_inference_steps=50).images[0]

image.save("cat-backpack.png")

```

To load a textual inversion embedding vector in Automatic1111 format, make sure to download the vector first

(for example from [civitAI](https://civitai.com/models/3036?modelVersionId=9857)) and then load the vector

locally:

```py

from diffusers import StableDiffusionPipeline

import torch

model_id = "runwayml/stable-diffusion-v1-5"

pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")

pipe.load_textual_inversion("./charturnerv2.pt", token="charturnerv2")

prompt = "charturnerv2, multiple views of the same character in the same outfit, a character turnaround of a woman wearing a black jacket and red shirt, best quality, intricate details."

image = pipe(prompt, num_inference_steps=50).images[0]

image.save("character.png")

```

"""

if not hasattr(self, "tokenizer") or not isinstance(self.tokenizer, PreTrainedTokenizer):

raise ValueError(

f"{self.__class__.__name__} requires `self.tokenizer` of type `PreTrainedTokenizer` for calling"

f" `{self.load_textual_inversion.__name__}`"

)

if not hasattr(self, "text_encoder") or not isinstance(self.text_encoder, PreTrainedModel):

raise ValueError(

f"{self.__class__.__name__} requires `self.text_encoder` of type `PreTrainedModel` for calling"

f" `{self.load_textual_inversion.__name__}`"

)

cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)

force_download = kwargs.pop("force_download", False)

resume_download = kwargs.pop("resume_download", False)

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

local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE)

use_auth_token = kwargs.pop("use_auth_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)

allow_pickle = False

if use_safetensors is None:

use_safetensors = True

allow_pickle = True

user_agent = {

"file_type": "text_inversion",

"framework": "pytorch",

}

if not isinstance(pretrained_model_name_or_path, list):

pretrained_model_name_or_paths = [pretrained_model_name_or_path]

else:

pretrained_model_name_or_paths = pretrained_model_name_or_path

if isinstance(token, str):

tokens = [token]

elif token is None:

tokens = [None] * len(pretrained_model_name_or_paths)

else:

tokens = token

if len(pretrained_model_name_or_paths) != len(tokens):

raise ValueError(

f"You have passed a list of models of length {len(pretrained_model_name_or_paths)}, and list of tokens of length {len(tokens)}"

f"Make sure both lists have the same length."

)

valid_tokens = [t for t in tokens if t is not None]

if len(set(valid_tokens)) < len(valid_tokens):

raise ValueError(f"You have passed a list of tokens that contains duplicates: {tokens}")

token_ids_and_embeddings = []

for pretrained_model_name_or_path, token in zip(pretrained_model_name_or_paths, tokens):

if not isinstance(pretrained_model_name_or_path, dict):

# 1. Load textual inversion file

model_file = None

# Let's first try to load .safetensors weights

if (use_safetensors and weight_name is None) or (

weight_name is not None and weight_name.endswith(".safetensors")

):

try:

model_file = _get_model_file(

pretrained_model_name_or_path,

weights_name=weight_name or TEXT_INVERSION_NAME_SAFE,

cache_dir=cache_dir,

force_download=force_download,

resume_download=resume_download,

proxies=proxies,

local_files_only=local_files_only,

use_auth_token=use_auth_token,

revision=revision,

subfolder=subfolder,

user_agent=user_agent,

)

state_dict = safetensors.torch.load_file(model_file, device="cpu")

except Exception as e:

if not allow_pickle:

raise e

model_file = None

if model_file is None:

model_file = _get_model_file(

pretrained_model_name_or_path,

weights_name=weight_name or TEXT_INVERSION_NAME,

cache_dir=cache_dir,

force_download=force_download,

resume_download=resume_download,

proxies=proxies,

local_files_only=local_files_only,

use_auth_token=use_auth_token,

revision=revision,

subfolder=subfolder,

user_agent=user_agent,

)

state_dict = torch.load(model_file, map_location="cpu")

else:

state_dict = pretrained_model_name_or_path

# 2. Load token and embedding correcly from file

loaded_token = None

if isinstance(state_dict, torch.Tensor):

if token is None:

raise ValueError(

"You are trying to load a textual inversion embedding that has been saved as a PyTorch tensor. Make sure to pass the name of the corresponding token in this case: `token=...`."

)

embedding = state_dict

elif len(state_dict) == 1:

# diffusers

loaded_token, embedding = next(iter(state_dict.items()))

elif "string_to_param" in state_dict:

# A1111

loaded_token = state_dict["name"]

embedding = state_dict["string_to_param"]["*"]

if token is not None and loaded_token != token:

logger.info(f"The loaded token: {loaded_token} is overwritten by the passed token {token}.")

else:

token = loaded_token

embedding = embedding.to(dtype=self.text_encoder.dtype, device=self.text_encoder.device)

# 3. Make sure we don't mess up the tokenizer or text encoder

vocab = self.tokenizer.get_vocab()

if token in vocab:

raise ValueError(

f"Token {token} already in tokenizer vocabulary. Please choose a different token name or remove {token} and embedding from the tokenizer and text encoder."

)

elif f"{token}_1" in vocab:

multi_vector_tokens = [token]

i = 1

while f"{token}_{i}" in self.tokenizer.added_tokens_encoder:

multi_vector_tokens.append(f"{token}_{i}")

i += 1

raise ValueError(

f"Multi-vector Token {multi_vector_tokens} already in tokenizer vocabulary. Please choose a different token name or remove the {multi_vector_tokens} and embedding from the tokenizer and text encoder."

)

is_multi_vector = len(embedding.shape) > 1 and embedding.shape[0] > 1

if is_multi_vector:

tokens = [token] + [f"{token}_{i}" for i in range(1, embedding.shape[0])]

embeddings = [e for e in embedding] # noqa: C416

else:

tokens = [token]

embeddings = [embedding[0]] if len(embedding.shape) > 1 else [embedding]

# add tokens and get ids

self.tokenizer.add_tokens(tokens)

token_ids = self.tokenizer.convert_tokens_to_ids(tokens)

token_ids_and_embeddings += zip(token_ids, embeddings)

logger.info(f"Loaded textual inversion embedding for {token}.")

# resize token embeddings and set all new embeddings

self.text_encoder.resize_token_embeddings(len(self.tokenizer))

for token_id, embedding in token_ids_and_embeddings:

self.text_encoder.get_input_embeddings().weight.data[token_id] = embedding

class LoraLoaderMixin:

r"""

Load LoRA layers into [`UNet2DConditionModel`] and

[`CLIPTextModel`](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel).

"""

text_encoder_name = TEXT_ENCODER_NAME

unet_name = UNET_NAME

def load_lora_weights(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **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.LoraLoaderMixin.lora_state_dict`] for more details on how the state dict is loaded.

See [`~loaders.LoraLoaderMixin.load_lora_into_unet`] for more details on how the state dict is loaded into

`self.unet`.

See [`~loaders.LoraLoaderMixin.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.LoraLoaderMixin.lora_state_dict`].

kwargs (`dict`, *optional*):

See [`~loaders.LoraLoaderMixin.lora_state_dict`].

"""

state_dict, network_alphas = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)

self.load_lora_into_unet(state_dict, network_alphas=network_alphas, unet=self.unet)

self.load_lora_into_text_encoder(

state_dict,

network_alphas=network_alphas,

text_encoder=self.text_encoder,

lora_scale=self.lora_scale,

)

@classmethod

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.

resume_download (`bool`, *optional*, defaults to `False`):

Whether or not to resume downloading the model weights and configuration files. If set to `False`, any

incompletely downloaded files are deleted.

proxies (`Dict[str, str]`, *optional*):

A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',

'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.

local_files_only (`bool`, *optional*, defaults to `False`):

Whether to only load local model weights and configuration files or not. If set to `True`, the model

won't be downloaded from the Hub.

use_auth_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.

mirror (`str`, *optional*):

Mirror source to resolve accessibility issues if you're downloading a model in China. We do not

guarantee the timeliness or safety of the source, and you should refer to the mirror site for more

information.

"""

# 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", DIFFUSERS_CACHE)

force_download = kwargs.pop("force_download", False)

resume_download = kwargs.pop("resume_download", False)

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

local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE)

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