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, lora_scale=1.0):

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 + (lora_scale * 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()

self.lora_scale = lora_scale

def _unfuse_lora(self):

if not (getattr(self, "w_up", None) is not None and getattr(self, "w_down", None) is not None):

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() - (self.lora_scale * 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_scale is None:

self.lora_scale = 1.0

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)

low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT)

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

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

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

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

is_network_alphas_none = network_alphas is None

allow_pickle = False

if use_safetensors is None:

use_safetensors = True

allow_pickle = True

user_agent = {

"file_type": "attn_procs_weights",

"framework": "pytorch",

}

if low_cpu_mem_usage and not is_accelerate_available():

low_cpu_mem_usage = False

logger.warning(

"Cannot initialize model with low cpu memory usage because `accelerate` was not found in the"

" environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install"

" `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip"

" install accelerate\n```\n."

)

model_file = None

if not isinstance(pretrained_model_name_or_path_or_dict, dict):

# Let's first try to load .safetensors weights

if (use_safetensors and weight_name is None) or (

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

):

try:

model_file = _get_model_file(

pretrained_model_name_or_path_or_dict,

weights_name=weight_name or LORA_WEIGHT_NAME_SAFE,

cache_dir=cache_dir,

force_download=force_download,

resume_download=resume_download,

proxies=proxies,

local_files_only=local_files_only,

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)

if network_alphas is not None:

network_alphas_keys = list(network_alphas.keys())

used_network_alphas_keys = set()

lora_grouped_dict = defaultdict(dict)

mapped_network_alphas = {}

all_keys = list(state_dict.keys())

for key in all_keys:

value = state_dict.pop(key)

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

lora_grouped_dict[attn_processor_key][sub_key] = value

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

if network_alphas is not None:

for k in network_alphas_keys:

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

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

used_network_alphas_keys.add(k)

if not is_network_alphas_none:

if len(set(network_alphas_keys) - used_network_alphas_keys) > 0:

raise ValueError(

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

)

if len(state_dict) > 0:

raise ValueError(

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

)

for key, value_dict in lora_grouped_dict.items():

attn_processor = self

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

attn_processor = getattr(attn_processor, sub_key)

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

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

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

if isinstance(attn_processor, LoRACompatibleConv):

in_features = attn_processor.in_channels

out_features = attn_processor.out_channels

kernel_size = attn_processor.kernel_size

ctx = init_empty_weights if low_cpu_mem_usage else nullcontext

with ctx():

lora = LoRAConv2dLayer(

in_features=in_features,

out_features=out_features,

rank=rank,

kernel_size=kernel_size,

stride=attn_processor.stride,

padding=attn_processor.padding,

network_alpha=mapped_network_alphas.get(key),

)

elif isinstance(attn_processor, LoRACompatibleLinear):

ctx = init_empty_weights if low_cpu_mem_usage else nullcontext

with ctx():

lora = LoRALinearLayer(

attn_processor.in_features,

attn_processor.out_features,

rank,

mapped_network_alphas.get(key),

)

else:

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

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

lora_layers_list.append((attn_processor, lora))

if low_cpu_mem_usage:

device = next(iter(value_dict.values())).device

dtype = next(iter(value_dict.values())).dtype

load_model_dict_into_meta(lora, value_dict, device=device, dtype=dtype)

else:

lora.load_state_dict(value_dict)

elif is_custom_diffusion:

attn_processors = {}

custom_diffusion_grouped_dict = defaultdict(dict)

for key, value in state_dict.items():

if len(value) == 0:

custom_diffusion_grouped_dict[key] = {}

else:

if "to_out" in key:

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

else:

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

custom_diffusion_grouped_dict[attn_processor_key][sub_key] = value

for key, value_dict in custom_diffusion_grouped_dict.items():

if len(value_dict) == 0:

attn_processors[key] = CustomDiffusionAttnProcessor(

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

)

else:

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

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

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

attn_processors[key] = CustomDiffusionAttnProcessor(

train_kv=True,

train_q_out=train_q_out,

hidden_size=hidden_size,

cross_attention_dim=cross_attention_dim,

)

attn_processors[key].load_state_dict(value_dict)

else:

raise ValueError(

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

)

# <Unsafe code

# We can be sure that the following works as it just sets attention processors, lora layers and puts all in the same dtype

# Now we remove any existing hooks to

is_model_cpu_offload = False

is_sequential_cpu_offload = False

if _pipeline is not None:

for _, component in _pipeline.components.items():

if isinstance(component, nn.Module):

if hasattr(component, "_hf_hook"):

is_model_cpu_offload = isinstance(getattr(component, "_hf_hook"), CpuOffload)

is_sequential_cpu_offload = isinstance(getattr(component, "_hf_hook"), AlignDevicesHook)

logger.info(

"Accelerate hooks detected. Since you have called `load_lora_weights()`, the previous hooks will be first removed. Then the LoRA parameters will be loaded and the hooks will be applied again."

)

remove_hook_from_module(component, recurse=is_sequential_cpu_offload)

# only custom diffusion needs to set attn processors

if is_custom_diffusion:

self.set_attn_processor(attn_processors)

# set lora layers

for target_module, lora_layer in lora_layers_list:

target_module.set_lora_layer(lora_layer)

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

# Offload back.

if is_model_cpu_offload:

_pipeline.enable_model_cpu_offload()

elif is_sequential_cpu_offload:

_pipeline.enable_sequential_cpu_offload()

# Unsafe code />

def convert_state_dict_legacy_attn_format(self, state_dict, network_alphas):

is_new_lora_format = all(

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

)

if is_new_lora_format:

# Strip the `"unet"` prefix.

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

if is_text_encoder_present:

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

logger.warn(warn_message)

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

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

# change processor format to 'pure' LoRACompatibleLinear format

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

def format_to_lora_compatible(key):

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

return key

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

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

if network_alphas is not None:

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

return state_dict, network_alphas

def save_attn_procs(

self,

save_directory: Union[str, os.PathLike],

is_main_process: bool = True,

weight_name: str = None,

save_function: Callable = None,

safe_serialization: bool = True,

**kwargs,

):

r"""

from .models.attention_processor import (

CustomDiffusionAttnProcessor,

CustomDiffusionAttnProcessor2_0,

CustomDiffusionXFormersAttnProcessor,

)

if os.path.isfile(save_directory):

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

return

if save_function is None:

if safe_serialization:

def save_function(weights, filename):

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

else:

save_function = torch.save

os.makedirs(save_directory, exist_ok=True)

is_custom_diffusion = any(

isinstance(

x,

(CustomDiffusionAttnProcessor, CustomDiffusionAttnProcessor2_0, CustomDiffusionXFormersAttnProcessor),

)

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

)

if is_custom_diffusion:

model_to_save = AttnProcsLayers(

{

y: x

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

if isinstance(

x,

(

CustomDiffusionAttnProcessor,

CustomDiffusionAttnProcessor2_0,

CustomDiffusionXFormersAttnProcessor,

),

)

}

)

state_dict = model_to_save.state_dict()

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

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

state_dict[name] = {}

else:

model_to_save = AttnProcsLayers(self.attn_processors)

state_dict = model_to_save.state_dict()

if weight_name is None:

if safe_serialization:

weight_name = CUSTOM_DIFFUSION_WEIGHT_NAME_SAFE if is_custom_diffusion else LORA_WEIGHT_NAME_SAFE

else:

weight_name = CUSTOM_DIFFUSION_WEIGHT_NAME if is_custom_diffusion else LORA_WEIGHT_NAME

# Save the model

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

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

def fuse_lora(self, lora_scale=1.0):

self.lora_scale = lora_scale

self.apply(self._fuse_lora_apply)

def _fuse_lora_apply(self, module):

if hasattr(module, "_fuse_lora"):

module._fuse_lora(self.lora_scale)

def unfuse_lora(self):

self.apply(self._unfuse_lora_apply)

def _unfuse_lora_apply(self, module):

if hasattr(module, "_unfuse_lora"):

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

}

state_dicts = []

for pretrained_model_name_or_path in pretrained_model_name_or_paths:

if not isinstance(pretrained_model_name_or_path, (dict, torch.Tensor)):

# 3.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

state_dicts.append(state_dict)

r"""

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

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"): # noqa: F821

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 _check_text_inv_inputs(self, tokenizer, text_encoder, pretrained_model_name_or_paths, tokens):

if tokenizer is None:

raise ValueError(

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

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

)

if text_encoder is None:

raise ValueError(

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

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

)

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

@staticmethod

def _retrieve_tokens_and_embeddings(tokens, state_dicts, tokenizer):

all_tokens = []

all_embeddings = []

for state_dict, token in zip(state_dicts, tokens):

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=...`."

)

loaded_token = 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"]["*"]

else:

raise ValueError(

f"Loaded state dictonary is incorrect: {state_dict}. \n\n"

"Please verify that the loaded state dictionary of the textual embedding either only has a single key or includes the `string_to_param`"

" input key."

)

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

if token in tokenizer.get_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."

)

all_tokens.append(token)

all_embeddings.append(embedding)

return all_tokens, all_embeddings

@staticmethod

def _extend_tokens_and_embeddings(tokens, embeddings, tokenizer):

all_tokens = []

all_embeddings = []

for embedding, token in zip(embeddings, tokens):

if f"{token}_1" in tokenizer.get_vocab():

multi_vector_tokens = [token]

i = 1

while f"{token}_{i}" in 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:

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

all_embeddings += [e for e in embedding] # noqa: C416

else:

all_tokens += [token]

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

return all_tokens, all_embeddings

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,

tokenizer: Optional["PreTrainedTokenizer"] = None, # noqa: F821

text_encoder: Optional["PreTrainedModel"] = None, # noqa: F821

**kwargs,

):

r"""

# 1. Set correct tokenizer and text encoder

tokenizer = tokenizer or getattr(self, "tokenizer", None)

text_encoder = text_encoder or getattr(self, "text_encoder", None)