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)

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"""

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_alpha = kwargs.pop("network_alpha", None)

if use_safetensors and not is_safetensors_available():

raise ValueError(

"`use_safetensors`=True but safetensors is not installed. Please install safetensors with `pip install safetenstors"

)

allow_pickle = False

if use_safetensors is None:

use_safetensors = is_safetensors_available()

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 = {}

is_lora = all("lora" in k 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)

for key, value in state_dict.items():

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

lora_grouped_dict[attn_processor_key][sub_key] = value

for key, value_dict in lora_grouped_dict.items():

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

hidden_size = value_dict["to_k_lora.up.weight"].shape[0]

attn_processor = self

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

attn_processor = getattr(attn_processor, sub_key)

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

)

attn_processors[key] = attn_processor_class(

hidden_size=hidden_size,

cross_attention_dim=cross_attention_dim,

rank=rank,

network_alpha=network_alpha,

)

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()}

# set layers

self.set_attn_processor(attn_processors)

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 = False,

**kwargs,

):

r"""

weight_name = weight_name or deprecate(

"weights_name",

"0.20.0",

"`weights_name` is deprecated, please use `weight_name` instead.",

take_from=kwargs,

)

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))

r"""

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

r"""

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"""

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"""

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)

if use_safetensors and not is_safetensors_available():

raise ValueError(

"`use_safetensors`=True but safetensors is not installed. Please install safetensors with `pip install safetenstors"

)

allow_pickle = False

if use_safetensors is None:

use_safetensors = is_safetensors_available()

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:

r"""

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):

r"""

# 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)

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

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

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

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

# set lora scale to a reasonable default

self._lora_scale = 1.0

if use_safetensors and not is_safetensors_available():

raise ValueError(

"`use_safetensors`=True but safetensors is not installed. Please install safetensors with `pip install safetenstors"

)

allow_pickle = False

if use_safetensors is None:

use_safetensors = is_safetensors_available()

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

# Convert kohya-ss Style LoRA attn procs to diffusers attn procs

network_alpha = None

if all((k.startswith("lora_te_") or k.startswith("lora_unet_")) for k in state_dict.keys()):

state_dict, network_alpha = self._convert_kohya_lora_to_diffusers(state_dict)

# If the serialization format is new (introduced in https://github.com/huggingface/diffusers/pull/2918),

# then the `state_dict` keys should have `self.unet_name` and/or `self.text_encoder_name` as

# their prefixes.

keys = list(state_dict.keys())

if all(key.startswith(self.unet_name) or key.startswith(self.text_encoder_name) for key in keys):

# Load the layers corresponding to UNet.

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

logger.info(f"Loading {self.unet_name}.")

unet_lora_state_dict = {

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

}

self.unet.load_attn_procs(unet_lora_state_dict, network_alpha=network_alpha)

# Load the layers corresponding to text encoder and make necessary adjustments.

text_encoder_keys = [k for k in keys if k.startswith(self.text_encoder_name)]

text_encoder_lora_state_dict = {

k.replace(f"{self.text_encoder_name}.", ""): v for k, v in state_dict.items() if k in text_encoder_keys

}

if len(text_encoder_lora_state_dict) > 0:

logger.info(f"Loading {self.text_encoder_name}.")

attn_procs_text_encoder = self._load_text_encoder_attn_procs(

text_encoder_lora_state_dict, network_alpha=network_alpha

)

self._modify_text_encoder(attn_procs_text_encoder)

# save lora attn procs of text encoder so that it can be easily retrieved

self._text_encoder_lora_attn_procs = attn_procs_text_encoder

# Otherwise, we're dealing with the old format. This means the `state_dict` should only

# contain the module names of the `unet` as its keys WITHOUT any prefix.

elif not all(

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

):

self.unet.load_attn_procs(state_dict)

warn_message = "You have saved the LoRA weights using the old format. To convert the old LoRA weights to the new format, you can first load them in a dictionary and then create a new dictionary like the following: `new_state_dict = {f'unet'.{module_name}: params for module_name, params in old_state_dict.items()}`."

warnings.warn(warn_message)

@property

def lora_scale(self) -> float:

# property function that returns the lora scale which can be set at run time by the pipeline.

# if _lora_scale has not been set, return 1

return self._lora_scale if hasattr(self, "_lora_scale") else 1.0

@property

def text_encoder_lora_attn_procs(self):

if hasattr(self, "_text_encoder_lora_attn_procs"):

return self._text_encoder_lora_attn_procs

return

def _remove_text_encoder_monkey_patch(self):

# Loop over the CLIPAttention module of text_encoder

for name, attn_module in self.text_encoder.named_modules():

if name.endswith(TEXT_ENCODER_ATTN_MODULE):

# Loop over the LoRA layers

for _, text_encoder_attr in self._lora_attn_processor_attr_to_text_encoder_attr.items():

# Retrieve the q/k/v/out projection of CLIPAttention

module = attn_module.get_submodule(text_encoder_attr)

if hasattr(module, "old_forward"):

# restore original `forward` to remove monkey-patch

module.forward = module.old_forward

delattr(module, "old_forward")

def _modify_text_encoder(self, attn_processors: Dict[str, LoRAAttnProcessor]):

r"""

# First, remove any monkey-patch that might have been applied before

self._remove_text_encoder_monkey_patch()

# Loop over the CLIPAttention module of text_encoder

for name, attn_module in self.text_encoder.named_modules():

if name.endswith(TEXT_ENCODER_ATTN_MODULE):

# Loop over the LoRA layers

for attn_proc_attr, text_encoder_attr in self._lora_attn_processor_attr_to_text_encoder_attr.items():