declare-lab/tango2
92
1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14import os15from collections import defaultdict16from typing import Callable, Dict, List, Optional, Union17 18import torch19 20from .models.attention_processor import LoRAAttnProcessor21from .utils import (22 DIFFUSERS_CACHE,23 HF_HUB_OFFLINE,24 _get_model_file,25 deprecate,26 is_safetensors_available,27 is_transformers_available,28 logging,29)30 31 32if is_safetensors_available():33 import safetensors34 35if is_transformers_available():36 from transformers import PreTrainedModel, PreTrainedTokenizer37 38 39logger = logging.get_logger(__name__)40 41 42LORA_WEIGHT_NAME = "pytorch_lora_weights.bin"43LORA_WEIGHT_NAME_SAFE = "pytorch_lora_weights.safetensors"44 45TEXT_INVERSION_NAME = "learned_embeds.bin"46TEXT_INVERSION_NAME_SAFE = "learned_embeds.safetensors"47 48 49class AttnProcsLayers(torch.nn.Module):50 def __init__(self, state_dict: Dict[str, torch.Tensor]):51 super().__init__()52 self.layers = torch.nn.ModuleList(state_dict.values())53 self.mapping = dict(enumerate(state_dict.keys()))54 self.rev_mapping = {v: k for k, v in enumerate(state_dict.keys())}55 56 # we add a hook to state_dict() and load_state_dict() so that the57 # naming fits with `unet.attn_processors`58 def map_to(module, state_dict, *args, **kwargs):59 new_state_dict = {}60 for key, value in state_dict.items():61 num = int(key.split(".")[1]) # 0 is always "layers"62 new_key = key.replace(f"layers.{num}", module.mapping[num])63 new_state_dict[new_key] = value64 65 return new_state_dict66 67 def map_from(module, state_dict, *args, **kwargs):68 all_keys = list(state_dict.keys())69 for key in all_keys:70 replace_key = key.split(".processor")[0] + ".processor"71 new_key = key.replace(replace_key, f"layers.{module.rev_mapping[replace_key]}")72 state_dict[new_key] = state_dict[key]73 del state_dict[key]74 75 self._register_state_dict_hook(map_to)76 self._register_load_state_dict_pre_hook(map_from, with_module=True)77 78 79class UNet2DConditionLoadersMixin:80 def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs):81 r"""82 Load pretrained attention processor layers into `UNet2DConditionModel`. Attention processor layers have to be83 defined in84 [cross_attention.py](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py)85 and be a `torch.nn.Module` class.86 87 <Tip warning={true}>88 89 This function is experimental and might change in the future.90 91 </Tip>92 93 Parameters:94 pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):95 Can be either:96 97 - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.98 Valid model ids should have an organization name, like `google/ddpm-celebahq-256`.99 - A path to a *directory* containing model weights saved using [`~ModelMixin.save_config`], e.g.,100 `./my_model_directory/`.101 - A [torch state102 dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).103 104 cache_dir (`Union[str, os.PathLike]`, *optional*):105 Path to a directory in which a downloaded pretrained model configuration should be cached if the106 standard cache should not be used.107 force_download (`bool`, *optional*, defaults to `False`):108 Whether or not to force the (re-)download of the model weights and configuration files, overriding the109 cached versions if they exist.110 resume_download (`bool`, *optional*, defaults to `False`):111 Whether or not to delete incompletely received files. Will attempt to resume the download if such a112 file exists.113 proxies (`Dict[str, str]`, *optional*):114 A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',115 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.116 local_files_only(`bool`, *optional*, defaults to `False`):117 Whether or not to only look at local files (i.e., do not try to download the model).118 use_auth_token (`str` or *bool*, *optional*):119 The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated120 when running `diffusers-cli login` (stored in `~/.huggingface`).121 revision (`str`, *optional*, defaults to `"main"`):122 The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a123 git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any124 identifier allowed by git.125 subfolder (`str`, *optional*, defaults to `""`):126 In case the relevant files are located inside a subfolder of the model repo (either remote in127 huggingface.co or downloaded locally), you can specify the folder name here.128 129 mirror (`str`, *optional*):130 Mirror source to accelerate downloads in China. If you are from China and have an accessibility131 problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety.132 Please refer to the mirror site for more information.133 134 <Tip>135 136 It is required to be logged in (`huggingface-cli login`) when you want to use private or [gated137 models](https://huggingface.co/docs/hub/models-gated#gated-models).138 139 </Tip>140 """141 142 cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)143 force_download = kwargs.pop("force_download", False)144 resume_download = kwargs.pop("resume_download", False)145 proxies = kwargs.pop("proxies", None)146 local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE)147 use_auth_token = kwargs.pop("use_auth_token", None)148 revision = kwargs.pop("revision", None)149 subfolder = kwargs.pop("subfolder", None)150 weight_name = kwargs.pop("weight_name", None)151 use_safetensors = kwargs.pop("use_safetensors", None)152 153 if use_safetensors and not is_safetensors_available():154 raise ValueError(155 "`use_safetensors`=True but safetensors is not installed. Please install safetensors with `pip install safetenstors"156 )157 158 allow_pickle = False159 if use_safetensors is None:160 use_safetensors = is_safetensors_available()161 allow_pickle = True162 163 user_agent = {164 "file_type": "attn_procs_weights",165 "framework": "pytorch",166 }167 168 model_file = None169 if not isinstance(pretrained_model_name_or_path_or_dict, dict):170 # Let's first try to load .safetensors weights171 if (use_safetensors and weight_name is None) or (172 weight_name is not None and weight_name.endswith(".safetensors")173 ):174 try:175 model_file = _get_model_file(176 pretrained_model_name_or_path_or_dict,177 weights_name=weight_name or LORA_WEIGHT_NAME_SAFE,178 cache_dir=cache_dir,179 force_download=force_download,180 resume_download=resume_download,181 proxies=proxies,182 local_files_only=local_files_only,183 use_auth_token=use_auth_token,184 revision=revision,185 subfolder=subfolder,186 user_agent=user_agent,187 )188 state_dict = safetensors.torch.load_file(model_file, device="cpu")189 except IOError as e:190 if not allow_pickle:191 raise e192 # try loading non-safetensors weights193 pass194 if model_file is None:195 model_file = _get_model_file(196 pretrained_model_name_or_path_or_dict,197 weights_name=weight_name or LORA_WEIGHT_NAME,198 cache_dir=cache_dir,199 force_download=force_download,200 resume_download=resume_download,201 proxies=proxies,202 local_files_only=local_files_only,203 use_auth_token=use_auth_token,204 revision=revision,205 subfolder=subfolder,206 user_agent=user_agent,207 )208 state_dict = torch.load(model_file, map_location="cpu")209 else:210 state_dict = pretrained_model_name_or_path_or_dict211 212 # fill attn processors213 attn_processors = {}214 215 is_lora = all("lora" in k for k in state_dict.keys())216 217 if is_lora:218 lora_grouped_dict = defaultdict(dict)219 for key, value in state_dict.items():220 attn_processor_key, sub_key = ".".join(key.split(".")[:-3]), ".".join(key.split(".")[-3:])221 lora_grouped_dict[attn_processor_key][sub_key] = value222 223 for key, value_dict in lora_grouped_dict.items():224 rank = value_dict["to_k_lora.down.weight"].shape[0]225 cross_attention_dim = value_dict["to_k_lora.down.weight"].shape[1]226 hidden_size = value_dict["to_k_lora.up.weight"].shape[0]227 228 attn_processors[key] = LoRAAttnProcessor(229 hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, rank=rank230 )231 attn_processors[key].load_state_dict(value_dict)232 233 else:234 raise ValueError(f"{model_file} does not seem to be in the correct format expected by LoRA training.")235 236 # set correct dtype & device237 attn_processors = {k: v.to(device=self.device, dtype=self.dtype) for k, v in attn_processors.items()}238 239 # set layers240 self.set_attn_processor(attn_processors)241 242 def save_attn_procs(243 self,244 save_directory: Union[str, os.PathLike],245 is_main_process: bool = True,246 weight_name: str = None,247 save_function: Callable = None,248 safe_serialization: bool = False,249 **kwargs,250 ):251 r"""252 Save an attention processor to a directory, so that it can be re-loaded using the253 `[`~loaders.UNet2DConditionLoadersMixin.load_attn_procs`]` method.254 255 Arguments:256 save_directory (`str` or `os.PathLike`):257 Directory to which to save. Will be created if it doesn't exist.258 is_main_process (`bool`, *optional*, defaults to `True`):259 Whether the process calling this is the main process or not. Useful when in distributed training like260 TPUs and need to call this function on all processes. In this case, set `is_main_process=True` only on261 the main process to avoid race conditions.262 save_function (`Callable`):263 The function to use to save the state dictionary. Useful on distributed training like TPUs when one264 need to replace `torch.save` by another method. Can be configured with the environment variable265 `DIFFUSERS_SAVE_MODE`.266 """267 weight_name = weight_name or deprecate(268 "weights_name",269 "0.18.0",270 "`weights_name` is deprecated, please use `weight_name` instead.",271 take_from=kwargs,272 )273 if os.path.isfile(save_directory):274 logger.error(f"Provided path ({save_directory}) should be a directory, not a file")275 return276 277 if save_function is None:278 if safe_serialization:279 280 def save_function(weights, filename):281 return safetensors.torch.save_file(weights, filename, metadata={"format": "pt"})282 283 else:284 save_function = torch.save285 286 os.makedirs(save_directory, exist_ok=True)287 288 model_to_save = AttnProcsLayers(self.attn_processors)289 290 # Save the model291 state_dict = model_to_save.state_dict()292 293 if weight_name is None:294 if safe_serialization:295 weight_name = LORA_WEIGHT_NAME_SAFE296 else:297 weight_name = LORA_WEIGHT_NAME298 299 # Save the model300 save_function(state_dict, os.path.join(save_directory, weight_name))301 logger.info(f"Model weights saved in {os.path.join(save_directory, weight_name)}")302 303 304class TextualInversionLoaderMixin:305 r"""306 Mixin class for loading textual inversion tokens and embeddings to the tokenizer and text encoder.307 """308 309 def maybe_convert_prompt(self, prompt: Union[str, List[str]], tokenizer: "PreTrainedTokenizer"):310 r"""311 Maybe convert a prompt into a "multi vector"-compatible prompt. If the prompt includes a token that corresponds312 to a multi-vector textual inversion embedding, this function will process the prompt so that the special token313 is replaced with multiple special tokens each corresponding to one of the vectors. If the prompt has no textual314 inversion token or a textual inversion token that is a single vector, the input prompt is simply returned.315 316 Parameters:317 prompt (`str` or list of `str`):318 The prompt or prompts to guide the image generation.319 tokenizer (`PreTrainedTokenizer`):320 The tokenizer responsible for encoding the prompt into input tokens.321 322 Returns:323 `str` or list of `str`: The converted prompt324 """325 if not isinstance(prompt, List):326 prompts = [prompt]327 else:328 prompts = prompt329 330 prompts = [self._maybe_convert_prompt(p, tokenizer) for p in prompts]331 332 if not isinstance(prompt, List):333 return prompts[0]334 335 return prompts336 337 def _maybe_convert_prompt(self, prompt: str, tokenizer: "PreTrainedTokenizer"):338 r"""339 Maybe convert a prompt into a "multi vector"-compatible prompt. If the prompt includes a token that corresponds340 to a multi-vector textual inversion embedding, this function will process the prompt so that the special token341 is replaced with multiple special tokens each corresponding to one of the vectors. If the prompt has no textual342 inversion token or a textual inversion token that is a single vector, the input prompt is simply returned.343 344 Parameters:345 prompt (`str`):346 The prompt to guide the image generation.347 tokenizer (`PreTrainedTokenizer`):348 The tokenizer responsible for encoding the prompt into input tokens.349 350 Returns:351 `str`: The converted prompt352 """353 tokens = tokenizer.tokenize(prompt)354 for token in tokens:355 if token in tokenizer.added_tokens_encoder:356 replacement = token357 i = 1358 while f"{token}_{i}" in tokenizer.added_tokens_encoder:359 replacement += f"{token}_{i}"360 i += 1361 362 prompt = prompt.replace(token, replacement)363 364 return prompt365 366 def load_textual_inversion(367 self, pretrained_model_name_or_path: Union[str, Dict[str, torch.Tensor]], token: Optional[str] = None, **kwargs368 ):369 r"""370 Load textual inversion embeddings into the text encoder of stable diffusion pipelines. Both `diffusers` and371 `Automatic1111` formats are supported.372 373 <Tip warning={true}>374 375 This function is experimental and might change in the future.376 377 </Tip>378 379 Parameters:380 pretrained_model_name_or_path (`str` or `os.PathLike`):381 Can be either:382 383 - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.384 Valid model ids should have an organization name, like385 `"sd-concepts-library/low-poly-hd-logos-icons"`.386 - A path to a *directory* containing textual inversion weights, e.g.387 `./my_text_inversion_directory/`.388 weight_name (`str`, *optional*):389 Name of a custom weight file. This should be used in two cases:390 391 - The saved textual inversion file is in `diffusers` format, but was saved under a specific weight392 name, such as `text_inv.bin`.393 - The saved textual inversion file is in the "Automatic1111" form.394 cache_dir (`Union[str, os.PathLike]`, *optional*):395 Path to a directory in which a downloaded pretrained model configuration should be cached if the396 standard cache should not be used.397 force_download (`bool`, *optional*, defaults to `False`):398 Whether or not to force the (re-)download of the model weights and configuration files, overriding the399 cached versions if they exist.400 resume_download (`bool`, *optional*, defaults to `False`):401 Whether or not to delete incompletely received files. Will attempt to resume the download if such a402 file exists.403 proxies (`Dict[str, str]`, *optional*):404 A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',405 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.406 local_files_only(`bool`, *optional*, defaults to `False`):407 Whether or not to only look at local files (i.e., do not try to download the model).408 use_auth_token (`str` or *bool*, *optional*):409 The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated410 when running `diffusers-cli login` (stored in `~/.huggingface`).411 revision (`str`, *optional*, defaults to `"main"`):412 The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a413 git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any414 identifier allowed by git.415 subfolder (`str`, *optional*, defaults to `""`):416 In case the relevant files are located inside a subfolder of the model repo (either remote in417 huggingface.co or downloaded locally), you can specify the folder name here.418 419 mirror (`str`, *optional*):420 Mirror source to accelerate downloads in China. If you are from China and have an accessibility421 problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety.422 Please refer to the mirror site for more information.423 424 <Tip>425 426 It is required to be logged in (`huggingface-cli login`) when you want to use private or [gated427 models](https://huggingface.co/docs/hub/models-gated#gated-models).428 429 </Tip>430 """431 if not hasattr(self, "tokenizer") or not isinstance(self.tokenizer, PreTrainedTokenizer):432 raise ValueError(433 f"{self.__class__.__name__} requires `self.tokenizer` of type `PreTrainedTokenizer` for calling"434 f" `{self.load_textual_inversion.__name__}`"435 )436 437 if not hasattr(self, "text_encoder") or not isinstance(self.text_encoder, PreTrainedModel):438 raise ValueError(439 f"{self.__class__.__name__} requires `self.text_encoder` of type `PreTrainedModel` for calling"440 f" `{self.load_textual_inversion.__name__}`"441 )442 443 cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)444 force_download = kwargs.pop("force_download", False)445 resume_download = kwargs.pop("resume_download", False)446 proxies = kwargs.pop("proxies", None)447 local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE)448 use_auth_token = kwargs.pop("use_auth_token", None)449 revision = kwargs.pop("revision", None)450 subfolder = kwargs.pop("subfolder", None)451 weight_name = kwargs.pop("weight_name", None)452 use_safetensors = kwargs.pop("use_safetensors", None)453 454 if use_safetensors and not is_safetensors_available():455 raise ValueError(456 "`use_safetensors`=True but safetensors is not installed. Please install safetensors with `pip install safetenstors"457 )458 459 allow_pickle = False460 if use_safetensors is None:461 use_safetensors = is_safetensors_available()462 allow_pickle = True463 464 user_agent = {465 "file_type": "text_inversion",466 "framework": "pytorch",467 }468 469 # 1. Load textual inversion file470 model_file = None471 # Let's first try to load .safetensors weights472 if (use_safetensors and weight_name is None) or (473 weight_name is not None and weight_name.endswith(".safetensors")474 ):475 try:476 model_file = _get_model_file(477 pretrained_model_name_or_path,478 weights_name=weight_name or TEXT_INVERSION_NAME_SAFE,479 cache_dir=cache_dir,480 force_download=force_download,481 resume_download=resume_download,482 proxies=proxies,483 local_files_only=local_files_only,484 use_auth_token=use_auth_token,485 revision=revision,486 subfolder=subfolder,487 user_agent=user_agent,488 )489 state_dict = safetensors.torch.load_file(model_file, device="cpu")490 except Exception as e:491 if not allow_pickle:492 raise e493 494 model_file = None495 496 if model_file is None:497 model_file = _get_model_file(498 pretrained_model_name_or_path,499 weights_name=weight_name or TEXT_INVERSION_NAME,500 cache_dir=cache_dir,501 force_download=force_download,502 resume_download=resume_download,503 proxies=proxies,504 local_files_only=local_files_only,505 use_auth_token=use_auth_token,506 revision=revision,507 subfolder=subfolder,508 user_agent=user_agent,509 )510 state_dict = torch.load(model_file, map_location="cpu")511 512 # 2. Load token and embedding correcly from file513 if isinstance(state_dict, torch.Tensor):514 if token is None:515 raise ValueError(516 "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=...`."517 )518 embedding = state_dict519 elif len(state_dict) == 1:520 # diffusers521 loaded_token, embedding = next(iter(state_dict.items()))522 elif "string_to_param" in state_dict:523 # A1111524 loaded_token = state_dict["name"]525 embedding = state_dict["string_to_param"]["*"]526 527 if token is not None and loaded_token != token:528 logger.warn(f"The loaded token: {loaded_token} is overwritten by the passed token {token}.")529 else:530 token = loaded_token531 532 embedding = embedding.to(dtype=self.text_encoder.dtype, device=self.text_encoder.device)533 534 # 3. Make sure we don't mess up the tokenizer or text encoder535 vocab = self.tokenizer.get_vocab()536 if token in vocab:537 raise ValueError(538 f"Token {token} already in tokenizer vocabulary. Please choose a different token name or remove {token} and embedding from the tokenizer and text encoder."539 )540 elif f"{token}_1" in vocab:541 multi_vector_tokens = [token]542 i = 1543 while f"{token}_{i}" in self.tokenizer.added_tokens_encoder:544 multi_vector_tokens.append(f"{token}_{i}")545 i += 1546 547 raise ValueError(548 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."549 )550 551 is_multi_vector = len(embedding.shape) > 1 and embedding.shape[0] > 1552 553 if is_multi_vector:554 tokens = [token] + [f"{token}_{i}" for i in range(1, embedding.shape[0])]555 embeddings = [e for e in embedding] # noqa: C416556 else:557 tokens = [token]558 embeddings = [embedding[0]] if len(embedding.shape) > 1 else [embedding]559 560 # add tokens and get ids561 self.tokenizer.add_tokens(tokens)562 token_ids = self.tokenizer.convert_tokens_to_ids(tokens)563 564 # resize token embeddings and set new embeddings565 self.text_encoder.resize_token_embeddings(len(self.tokenizer))566 for token_id, embedding in zip(token_ids, embeddings):567 self.text_encoder.get_input_embeddings().weight.data[token_id] = embedding568 569 logger.info("Loaded textual inversion embedding for {token}.")570