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

# overwrite PyTorch's `state_dict` to be sure that only the 'regular_linear_layer' weights are saved

# when saving the whole text encoder model and when LoRA is unloaded or fused

def state_dict(self, *args, destination=None, prefix="", keep_vars=False):

if self.lora_linear_layer is None:

return self.regular_linear_layer.state_dict(

*args, destination=destination, prefix=prefix, keep_vars=keep_vars

)

return super().state_dict(*args, destination=destination, prefix=prefix, keep_vars=keep_vars)

def _fuse_lora(self):

if self.lora_linear_layer is None:

return

dtype, device = self.regular_linear_layer.weight.data.dtype, self.regular_linear_layer.weight.data.device

w_orig = self.regular_linear_layer.weight.data.float()

w_up = self.lora_linear_layer.up.weight.data.float()

w_down = self.lora_linear_layer.down.weight.data.float()

if self.lora_linear_layer.network_alpha is not None:

w_up = w_up * self.lora_linear_layer.network_alpha / self.lora_linear_layer.rank

fused_weight = w_orig + torch.bmm(w_up[None, :], w_down[None, :])[0]

self.regular_linear_layer.weight.data = fused_weight.to(device=device, dtype=dtype)

# we can drop the lora layer now

self.lora_linear_layer = None

# offload the up and down matrices to CPU to not blow the memory

self.w_up = w_up.cpu()

self.w_down = w_down.cpu()

def _unfuse_lora(self):

if not (hasattr(self, "w_up") and hasattr(self, "w_down")):

return

fused_weight = self.regular_linear_layer.weight.data

dtype, device = fused_weight.dtype, fused_weight.device

w_up = self.w_up.to(device=device).float()

w_down = self.w_down.to(device).float()

unfused_weight = fused_weight.float() - torch.bmm(w_up[None, :], w_down[None, :])[0]

self.regular_linear_layer.weight.data = unfused_weight.to(device=device, dtype=dtype)

self.w_up = None

self.w_down = None

def forward(self, input):

if self.lora_linear_layer is None:

return self.regular_linear_layer(input)

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__}")

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__}")

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

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

from .models.attention_processor import (

CustomDiffusionAttnProcessor,

)

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

lora_layers_list = []

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:

# correct keys

state_dict, network_alphas = self.convert_state_dict_legacy_attn_format(state_dict, network_alphas)

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.

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)

lora_layers_list.append((attn_processor, lora))

elif is_custom_diffusion:

attn_processors = {}

custom_diffusion_grouped_dict = defaultdict(dict)

for key, value in state_dict.items():

if len(value) == 0:

custom_diffusion_grouped_dict[key] = {}

else:

if "to_out" in key:

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

else:

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

custom_diffusion_grouped_dict[attn_processor_key][sub_key] = value

for key, value_dict in custom_diffusion_grouped_dict.items():

if len(value_dict) == 0:

attn_processors[key] = CustomDiffusionAttnProcessor(

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

)

else:

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

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

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

attn_processors[key] = CustomDiffusionAttnProcessor(

train_kv=True,

train_q_out=train_q_out,

hidden_size=hidden_size,

cross_attention_dim=cross_attention_dim,

)

attn_processors[key].load_state_dict(value_dict)

self.set_attn_processor(attn_processors)

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

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

# set lora layers

for target_module, lora_layer in lora_layers_list:

target_module.set_lora_layer(lora_layer)

def convert_state_dict_legacy_attn_format(self, state_dict, network_alphas):

is_new_lora_format = all(

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

)

if is_new_lora_format:

# Strip the `"unet"` prefix.

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

if is_text_encoder_present:

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

logger.warn(warn_message)

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

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

# change processor format to 'pure' LoRACompatibleLinear format

if any("processor" in k.split(".") for k in state_dict.keys()):

def format_to_lora_compatible(key):

if "processor" not in key.split("."):

return key

return key.replace(".processor", "").replace("to_out_lora", "to_out.0.lora").replace("_lora", ".lora")

state_dict = {format_to_lora_compatible(k): v for k, v in state_dict.items()}

if network_alphas is not None:

network_alphas = {format_to_lora_compatible(k): v for k, v in network_alphas.items()}

return state_dict, network_alphas

def save_attn_procs(

self,

save_directory: Union[str, os.PathLike],

is_main_process: bool = True,

weight_name: str = None,

save_function: Callable = None,

safe_serialization: bool = True,

**kwargs,

):

r"""

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

def fuse_lora(self):

self.apply(self._fuse_lora_apply)

def _fuse_lora_apply(self, module):

if hasattr(module, "_fuse_lora"):

module._fuse_lora()

def unfuse_lora(self):

self.apply(self._unfuse_lora_apply)

def _unfuse_lora_apply(self, module):

if hasattr(module, "_unfuse_lora"):

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)

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:

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

"""

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