We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Adding support for `safetensors` and LoRa.
Enabling safetensors support for the LoRA files:

Asked here: https://github.com/huggingface/safetensors/issues/180

Same attitude as the regular model weights.

If: 
- safetensors is installed
- repo has the safetensors lora
then it is the default.

Adding `safe_serialization` on lora create so users can default to saving safetensors formats.

What's techically missing is the option to choose on load what the format is along with the weights_name. I ddin't want to add it here for simplicity (Since most users should be using the default anyway). But we could add that.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/diffusers/loaders.py`
Content:
```
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 at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # 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 and
13 # limitations under the License.
14 import os
15 from collections import defaultdict
16 from typing import Callable, Dict, Union
17 
18 import torch
19 
20 from .models.cross_attention import LoRACrossAttnProcessor
21 from .models.modeling_utils import _get_model_file
22 from .utils import DIFFUSERS_CACHE, HF_HUB_OFFLINE, is_safetensors_available, logging
23 
24 
25 if is_safetensors_available():
26     import safetensors
27 
28 
29 logger = logging.get_logger(__name__)
30 
31 
32 LORA_WEIGHT_NAME = "pytorch_lora_weights.bin"
33 LORA_WEIGHT_NAME_SAFE = "pytorch_lora_weights.safetensors"
34 
35 
36 class AttnProcsLayers(torch.nn.Module):
37     def __init__(self, state_dict: Dict[str, torch.Tensor]):
38         super().__init__()
39         self.layers = torch.nn.ModuleList(state_dict.values())
40         self.mapping = {k: v for k, v in enumerate(state_dict.keys())}
41         self.rev_mapping = {v: k for k, v in enumerate(state_dict.keys())}
42 
43         # we add a hook to state_dict() and load_state_dict() so that the
44         # naming fits with `unet.attn_processors`
45         def map_to(module, state_dict, *args, **kwargs):
46             new_state_dict = {}
47             for key, value in state_dict.items():
48                 num = int(key.split(".")[1])  # 0 is always "layers"
49                 new_key = key.replace(f"layers.{num}", module.mapping[num])
50                 new_state_dict[new_key] = value
51 
52             return new_state_dict
53 
54         def map_from(module, state_dict, *args, **kwargs):
55             all_keys = list(state_dict.keys())
56             for key in all_keys:
57                 replace_key = key.split(".processor")[0] + ".processor"
58                 new_key = key.replace(replace_key, f"layers.{module.rev_mapping[replace_key]}")
59                 state_dict[new_key] = state_dict[key]
60                 del state_dict[key]
61 
62         self._register_state_dict_hook(map_to)
63         self._register_load_state_dict_pre_hook(map_from, with_module=True)
64 
65 
66 class UNet2DConditionLoadersMixin:
67     def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs):
68         r"""
69         Load pretrained attention processor layers into `UNet2DConditionModel`. Attention processor layers have to be
70         defined in
71         [cross_attention.py](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py)
72         and be a `torch.nn.Module` class.
73 
74         <Tip warning={true}>
75 
76             This function is experimental and might change in the future.
77 
78         </Tip>
79 
80         Parameters:
81             pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
82                 Can be either:
83 
84                     - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
85                       Valid model ids should have an organization name, like `google/ddpm-celebahq-256`.
86                     - A path to a *directory* containing model weights saved using [`~ModelMixin.save_config`], e.g.,
87                       `./my_model_directory/`.
88                     - A [torch state
89                       dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).
90 
91             cache_dir (`Union[str, os.PathLike]`, *optional*):
92                 Path to a directory in which a downloaded pretrained model configuration should be cached if the
93                 standard cache should not be used.
94             force_download (`bool`, *optional*, defaults to `False`):
95                 Whether or not to force the (re-)download of the model weights and configuration files, overriding the
96                 cached versions if they exist.
97             resume_download (`bool`, *optional*, defaults to `False`):
98                 Whether or not to delete incompletely received files. Will attempt to resume the download if such a
99                 file exists.
100             proxies (`Dict[str, str]`, *optional*):
101                 A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
102                 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
103             local_files_only(`bool`, *optional*, defaults to `False`):
104                 Whether or not to only look at local files (i.e., do not try to download the model).
105             use_auth_token (`str` or *bool*, *optional*):
106                 The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
107                 when running `diffusers-cli login` (stored in `~/.huggingface`).
108             revision (`str`, *optional*, defaults to `"main"`):
109                 The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
110                 git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
111                 identifier allowed by git.
112             subfolder (`str`, *optional*, defaults to `""`):
113                 In case the relevant files are located inside a subfolder of the model repo (either remote in
114                 huggingface.co or downloaded locally), you can specify the folder name here.
115 
116             mirror (`str`, *optional*):
117                 Mirror source to accelerate downloads in China. If you are from China and have an accessibility
118                 problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety.
119                 Please refer to the mirror site for more information.
120 
121         <Tip>
122 
123          It is required to be logged in (`huggingface-cli login`) when you want to use private or [gated
124          models](https://huggingface.co/docs/hub/models-gated#gated-models).
125 
126         </Tip>
127 
128         <Tip>
129 
130         Activate the special ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use
131         this method in a firewalled environment.
132 
133         </Tip>
134         """
135 
136         cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)
137         force_download = kwargs.pop("force_download", False)
138         resume_download = kwargs.pop("resume_download", False)
139         proxies = kwargs.pop("proxies", None)
140         local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE)
141         use_auth_token = kwargs.pop("use_auth_token", None)
142         revision = kwargs.pop("revision", None)
143         subfolder = kwargs.pop("subfolder", None)
144         weight_name = kwargs.pop("weight_name", None)
145 
146         user_agent = {
147             "file_type": "attn_procs_weights",
148             "framework": "pytorch",
149         }
150 
151         model_file = None
152         if not isinstance(pretrained_model_name_or_path_or_dict, dict):
153             if is_safetensors_available():
154                 if weight_name is None:
155                     weight_name = LORA_WEIGHT_NAME_SAFE
156                 try:
157                     model_file = _get_model_file(
158                         pretrained_model_name_or_path_or_dict,
159                         weights_name=weight_name,
160                         cache_dir=cache_dir,
161                         force_download=force_download,
162                         resume_download=resume_download,
163                         proxies=proxies,
164                         local_files_only=local_files_only,
165                         use_auth_token=use_auth_token,
166                         revision=revision,
167                         subfolder=subfolder,
168                         user_agent=user_agent,
169                     )
170                     state_dict = safetensors.torch.load_file(model_file, device="cpu")
171                 except EnvironmentError:
172                     if weight_name == LORA_WEIGHT_NAME_SAFE:
173                         weight_name = None
174             if model_file is None:
175                 if weight_name is None:
176                     weight_name = LORA_WEIGHT_NAME
177                 model_file = _get_model_file(
178                     pretrained_model_name_or_path_or_dict,
179                     weights_name=weight_name,
180                     cache_dir=cache_dir,
181                     force_download=force_download,
182                     resume_download=resume_download,
183                     proxies=proxies,
184                     local_files_only=local_files_only,
185                     use_auth_token=use_auth_token,
186                     revision=revision,
187                     subfolder=subfolder,
188                     user_agent=user_agent,
189                 )
190                 state_dict = torch.load(model_file, map_location="cpu")
191         else:
192             state_dict = pretrained_model_name_or_path_or_dict
193 
194         # fill attn processors
195         attn_processors = {}
196 
197         is_lora = all("lora" in k for k in state_dict.keys())
198 
199         if is_lora:
200             lora_grouped_dict = defaultdict(dict)
201             for key, value in state_dict.items():
202                 attn_processor_key, sub_key = ".".join(key.split(".")[:-3]), ".".join(key.split(".")[-3:])
203                 lora_grouped_dict[attn_processor_key][sub_key] = value
204 
205             for key, value_dict in lora_grouped_dict.items():
206                 rank = value_dict["to_k_lora.down.weight"].shape[0]
207                 cross_attention_dim = value_dict["to_k_lora.down.weight"].shape[1]
208                 hidden_size = value_dict["to_k_lora.up.weight"].shape[0]
209 
210                 attn_processors[key] = LoRACrossAttnProcessor(
211                     hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, rank=rank
212                 )
213                 attn_processors[key].load_state_dict(value_dict)
214 
215         else:
216             raise ValueError(f"{model_file} does not seem to be in the correct format expected by LoRA training.")
217 
218         # set correct dtype & device
219         attn_processors = {k: v.to(device=self.device, dtype=self.dtype) for k, v in attn_processors.items()}
220 
221         # set layers
222         self.set_attn_processor(attn_processors)
223 
224     def save_attn_procs(
225         self,
226         save_directory: Union[str, os.PathLike],
227         is_main_process: bool = True,
228         weights_name: str = None,
229         save_function: Callable = None,
230         safe_serialization: bool = False,
231     ):
232         r"""
233         Save an attention processor to a directory, so that it can be re-loaded using the
234         `[`~loaders.UNet2DConditionLoadersMixin.load_attn_procs`]` method.
235 
236         Arguments:
237             save_directory (`str` or `os.PathLike`):
238                 Directory to which to save. Will be created if it doesn't exist.
239             is_main_process (`bool`, *optional*, defaults to `True`):
240                 Whether the process calling this is the main process or not. Useful when in distributed training like
241                 TPUs and need to call this function on all processes. In this case, set `is_main_process=True` only on
242                 the main process to avoid race conditions.
243             save_function (`Callable`):
244                 The function to use to save the state dictionary. Useful on distributed training like TPUs when one
245                 need to replace `torch.save` by another method. Can be configured with the environment variable
246                 `DIFFUSERS_SAVE_MODE`.
247         """
248         if os.path.isfile(save_directory):
249             logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
250             return
251 
252         if save_function is None:
253             if safe_serialization:
254 
255                 def save_function(weights, filename):
256                     return safetensors.torch.save_file(weights, filename, metadata={"format": "pt"})
257 
258             else:
259                 save_function = torch.save
260 
261         os.makedirs(save_directory, exist_ok=True)
262 
263         model_to_save = AttnProcsLayers(self.attn_processors)
264 
265         # Save the model
266         state_dict = model_to_save.state_dict()
267 
268         # Clean the folder from a previous save
269         for filename in os.listdir(save_directory):
270             full_filename = os.path.join(save_directory, filename)
271             # If we have a shard file that is not going to be replaced, we delete it, but only from the main process
272             # in distributed settings to avoid race conditions.
273             weights_no_suffix = weights_name.replace(".bin", "")
274             if filename.startswith(weights_no_suffix) and os.path.isfile(full_filename) and is_main_process:
275                 os.remove(full_filename)
276 
277         if weights_name is None:
278             if safe_serialization:
279                 weights_name = LORA_WEIGHT_NAME_SAFE
280             else:
281                 weights_name = LORA_WEIGHT_NAME
282 
283         # Save the model
284         save_function(state_dict, os.path.join(save_directory, weights_name))
285 
286         logger.info(f"Model weights saved in {os.path.join(save_directory, weights_name)}")
287
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py
--- a/src/diffusers/loaders.py
+++ b/src/diffusers/loaders.py
@@ -150,7 +150,7 @@
 
         model_file = None
         if not isinstance(pretrained_model_name_or_path_or_dict, dict):
-            if is_safetensors_available():
+            if (is_safetensors_available() and weight_name is None) or weight_name.endswith(".safetensors"):
                 if weight_name is None:
                     weight_name = LORA_WEIGHT_NAME_SAFE
                 try:

{"golden_diff": "diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py\n--- a/src/diffusers/loaders.py\n+++ b/src/diffusers/loaders.py\n@@ -150,7 +150,7 @@\n \n         model_file = None\n         if not isinstance(pretrained_model_name_or_path_or_dict, dict):\n-            if is_safetensors_available():\n+            if (is_safetensors_available() and weight_name is None) or weight_name.endswith(\".safetensors\"):\n                 if weight_name is None:\n                     weight_name = LORA_WEIGHT_NAME_SAFE\n                 try:\n", "issue": "Adding support for `safetensors` and LoRa.\nEnabling safetensors support for the LoRA files:\r\n\r\nAsked here: https://github.com/huggingface/safetensors/issues/180\r\n\r\nSame attitude as the regular model weights.\r\n\r\nIf: \r\n- safetensors is installed\r\n- repo has the safetensors lora\r\nthen it is the default.\r\n\r\nAdding `safe_serialization` on lora create so users can default to saving safetensors formats.\r\n\r\nWhat's techically missing is the option to choose on load what the format is along with the weights_name. I ddin't want to add it here for simplicity (Since most users should be using the default anyway). But we could add that.\n", "before_files": [{"content": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\nfrom collections import defaultdict\nfrom typing import Callable, Dict, Union\n\nimport torch\n\nfrom .models.cross_attention import LoRACrossAttnProcessor\nfrom .models.modeling_utils import _get_model_file\nfrom .utils import DIFFUSERS_CACHE, HF_HUB_OFFLINE, is_safetensors_available, logging\n\n\nif is_safetensors_available():\n    import safetensors\n\n\nlogger = logging.get_logger(__name__)\n\n\nLORA_WEIGHT_NAME = \"pytorch_lora_weights.bin\"\nLORA_WEIGHT_NAME_SAFE = \"pytorch_lora_weights.safetensors\"\n\n\nclass AttnProcsLayers(torch.nn.Module):\n    def __init__(self, state_dict: Dict[str, torch.Tensor]):\n        super().__init__()\n        self.layers = torch.nn.ModuleList(state_dict.values())\n        self.mapping = {k: v for k, v in enumerate(state_dict.keys())}\n        self.rev_mapping = {v: k for k, v in enumerate(state_dict.keys())}\n\n        # we add a hook to state_dict() and load_state_dict() so that the\n        # naming fits with `unet.attn_processors`\n        def map_to(module, state_dict, *args, **kwargs):\n            new_state_dict = {}\n            for key, value in state_dict.items():\n                num = int(key.split(\".\")[1])  # 0 is always \"layers\"\n                new_key = key.replace(f\"layers.{num}\", module.mapping[num])\n                new_state_dict[new_key] = value\n\n            return new_state_dict\n\n        def map_from(module, state_dict, *args, **kwargs):\n            all_keys = list(state_dict.keys())\n            for key in all_keys:\n                replace_key = key.split(\".processor\")[0] + \".processor\"\n                new_key = key.replace(replace_key, f\"layers.{module.rev_mapping[replace_key]}\")\n                state_dict[new_key] = state_dict[key]\n                del state_dict[key]\n\n        self._register_state_dict_hook(map_to)\n        self._register_load_state_dict_pre_hook(map_from, with_module=True)\n\n\nclass UNet2DConditionLoadersMixin:\n    def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs):\n        r\"\"\"\n        Load pretrained attention processor layers into `UNet2DConditionModel`. Attention processor layers have to be\n        defined in\n        [cross_attention.py](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py)\n        and be a `torch.nn.Module` class.\n\n        <Tip warning={true}>\n\n            This function is experimental and might change in the future.\n\n        </Tip>\n\n        Parameters:\n            pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):\n                Can be either:\n\n                    - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.\n                      Valid model ids should have an organization name, like `google/ddpm-celebahq-256`.\n                    - A path to a *directory* containing model weights saved using [`~ModelMixin.save_config`], e.g.,\n                      `./my_model_directory/`.\n                    - A [torch state\n                      dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).\n\n            cache_dir (`Union[str, os.PathLike]`, *optional*):\n                Path to a directory in which a downloaded pretrained model configuration should be cached if the\n                standard cache should not be used.\n            force_download (`bool`, *optional*, defaults to `False`):\n                Whether or not to force the (re-)download of the model weights and configuration files, overriding the\n                cached versions if they exist.\n            resume_download (`bool`, *optional*, defaults to `False`):\n                Whether or not to delete incompletely received files. Will attempt to resume the download if such a\n                file exists.\n            proxies (`Dict[str, str]`, *optional*):\n                A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',\n                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.\n            local_files_only(`bool`, *optional*, defaults to `False`):\n                Whether or not to only look at local files (i.e., do not try to download the model).\n            use_auth_token (`str` or *bool*, *optional*):\n                The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated\n                when running `diffusers-cli login` (stored in `~/.huggingface`).\n            revision (`str`, *optional*, defaults to `\"main\"`):\n                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a\n                git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any\n                identifier allowed by git.\n            subfolder (`str`, *optional*, defaults to `\"\"`):\n                In case the relevant files are located inside a subfolder of the model repo (either remote in\n                huggingface.co or downloaded locally), you can specify the folder name here.\n\n            mirror (`str`, *optional*):\n                Mirror source to accelerate downloads in China. If you are from China and have an accessibility\n                problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety.\n                Please refer to the mirror site for more information.\n\n        <Tip>\n\n         It is required to be logged in (`huggingface-cli login`) when you want to use private or [gated\n         models](https://huggingface.co/docs/hub/models-gated#gated-models).\n\n        </Tip>\n\n        <Tip>\n\n        Activate the special [\"offline-mode\"](https://huggingface.co/diffusers/installation.html#offline-mode) to use\n        this method in a firewalled environment.\n\n        </Tip>\n        \"\"\"\n\n        cache_dir = kwargs.pop(\"cache_dir\", DIFFUSERS_CACHE)\n        force_download = kwargs.pop(\"force_download\", False)\n        resume_download = kwargs.pop(\"resume_download\", False)\n        proxies = kwargs.pop(\"proxies\", None)\n        local_files_only = kwargs.pop(\"local_files_only\", HF_HUB_OFFLINE)\n        use_auth_token = kwargs.pop(\"use_auth_token\", None)\n        revision = kwargs.pop(\"revision\", None)\n        subfolder = kwargs.pop(\"subfolder\", None)\n        weight_name = kwargs.pop(\"weight_name\", None)\n\n        user_agent = {\n            \"file_type\": \"attn_procs_weights\",\n            \"framework\": \"pytorch\",\n        }\n\n        model_file = None\n        if not isinstance(pretrained_model_name_or_path_or_dict, dict):\n            if is_safetensors_available():\n                if weight_name is None:\n                    weight_name = LORA_WEIGHT_NAME_SAFE\n                try:\n                    model_file = _get_model_file(\n                        pretrained_model_name_or_path_or_dict,\n                        weights_name=weight_name,\n                        cache_dir=cache_dir,\n                        force_download=force_download,\n                        resume_download=resume_download,\n                        proxies=proxies,\n                        local_files_only=local_files_only,\n                        use_auth_token=use_auth_token,\n                        revision=revision,\n                        subfolder=subfolder,\n                        user_agent=user_agent,\n                    )\n                    state_dict = safetensors.torch.load_file(model_file, device=\"cpu\")\n                except EnvironmentError:\n                    if weight_name == LORA_WEIGHT_NAME_SAFE:\n                        weight_name = None\n            if model_file is None:\n                if weight_name is None:\n                    weight_name = LORA_WEIGHT_NAME\n                model_file = _get_model_file(\n                    pretrained_model_name_or_path_or_dict,\n                    weights_name=weight_name,\n                    cache_dir=cache_dir,\n                    force_download=force_download,\n                    resume_download=resume_download,\n                    proxies=proxies,\n                    local_files_only=local_files_only,\n                    use_auth_token=use_auth_token,\n                    revision=revision,\n                    subfolder=subfolder,\n                    user_agent=user_agent,\n                )\n                state_dict = torch.load(model_file, map_location=\"cpu\")\n        else:\n            state_dict = pretrained_model_name_or_path_or_dict\n\n        # fill attn processors\n        attn_processors = {}\n\n        is_lora = all(\"lora\" in k for k in state_dict.keys())\n\n        if is_lora:\n            lora_grouped_dict = defaultdict(dict)\n            for key, value in state_dict.items():\n                attn_processor_key, sub_key = \".\".join(key.split(\".\")[:-3]), \".\".join(key.split(\".\")[-3:])\n                lora_grouped_dict[attn_processor_key][sub_key] = value\n\n            for key, value_dict in lora_grouped_dict.items():\n                rank = value_dict[\"to_k_lora.down.weight\"].shape[0]\n                cross_attention_dim = value_dict[\"to_k_lora.down.weight\"].shape[1]\n                hidden_size = value_dict[\"to_k_lora.up.weight\"].shape[0]\n\n                attn_processors[key] = LoRACrossAttnProcessor(\n                    hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, rank=rank\n                )\n                attn_processors[key].load_state_dict(value_dict)\n\n        else:\n            raise ValueError(f\"{model_file} does not seem to be in the correct format expected by LoRA training.\")\n\n        # set correct dtype & device\n        attn_processors = {k: v.to(device=self.device, dtype=self.dtype) for k, v in attn_processors.items()}\n\n        # set layers\n        self.set_attn_processor(attn_processors)\n\n    def save_attn_procs(\n        self,\n        save_directory: Union[str, os.PathLike],\n        is_main_process: bool = True,\n        weights_name: str = None,\n        save_function: Callable = None,\n        safe_serialization: bool = False,\n    ):\n        r\"\"\"\n        Save an attention processor to a directory, so that it can be re-loaded using the\n        `[`~loaders.UNet2DConditionLoadersMixin.load_attn_procs`]` method.\n\n        Arguments:\n            save_directory (`str` or `os.PathLike`):\n                Directory to which to save. Will be created if it doesn't exist.\n            is_main_process (`bool`, *optional*, defaults to `True`):\n                Whether the process calling this is the main process or not. Useful when in distributed training like\n                TPUs and need to call this function on all processes. In this case, set `is_main_process=True` only on\n                the main process to avoid race conditions.\n            save_function (`Callable`):\n                The function to use to save the state dictionary. Useful on distributed training like TPUs when one\n                need to replace `torch.save` by another method. Can be configured with the environment variable\n                `DIFFUSERS_SAVE_MODE`.\n        \"\"\"\n        if os.path.isfile(save_directory):\n            logger.error(f\"Provided path ({save_directory}) should be a directory, not a file\")\n            return\n\n        if save_function is None:\n            if safe_serialization:\n\n                def save_function(weights, filename):\n                    return safetensors.torch.save_file(weights, filename, metadata={\"format\": \"pt\"})\n\n            else:\n                save_function = torch.save\n\n        os.makedirs(save_directory, exist_ok=True)\n\n        model_to_save = AttnProcsLayers(self.attn_processors)\n\n        # Save the model\n        state_dict = model_to_save.state_dict()\n\n        # Clean the folder from a previous save\n        for filename in os.listdir(save_directory):\n            full_filename = os.path.join(save_directory, filename)\n            # If we have a shard file that is not going to be replaced, we delete it, but only from the main process\n            # in distributed settings to avoid race conditions.\n            weights_no_suffix = weights_name.replace(\".bin\", \"\")\n            if filename.startswith(weights_no_suffix) and os.path.isfile(full_filename) and is_main_process:\n                os.remove(full_filename)\n\n        if weights_name is None:\n            if safe_serialization:\n                weights_name = LORA_WEIGHT_NAME_SAFE\n            else:\n                weights_name = LORA_WEIGHT_NAME\n\n        # Save the model\n        save_function(state_dict, os.path.join(save_directory, weights_name))\n\n        logger.info(f\"Model weights saved in {os.path.join(save_directory, weights_name)}\")\n", "path": "src/diffusers/loaders.py"}], "after_files": [{"content": "# Copyright 2023 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\nfrom collections import defaultdict\nfrom typing import Callable, Dict, Union\n\nimport torch\n\nfrom .models.cross_attention import LoRACrossAttnProcessor\nfrom .models.modeling_utils import _get_model_file\nfrom .utils import DIFFUSERS_CACHE, HF_HUB_OFFLINE, is_safetensors_available, logging\n\n\nif is_safetensors_available():\n    import safetensors\n\n\nlogger = logging.get_logger(__name__)\n\n\nLORA_WEIGHT_NAME = \"pytorch_lora_weights.bin\"\nLORA_WEIGHT_NAME_SAFE = \"pytorch_lora_weights.safetensors\"\n\n\nclass AttnProcsLayers(torch.nn.Module):\n    def __init__(self, state_dict: Dict[str, torch.Tensor]):\n        super().__init__()\n        self.layers = torch.nn.ModuleList(state_dict.values())\n        self.mapping = {k: v for k, v in enumerate(state_dict.keys())}\n        self.rev_mapping = {v: k for k, v in enumerate(state_dict.keys())}\n\n        # we add a hook to state_dict() and load_state_dict() so that the\n        # naming fits with `unet.attn_processors`\n        def map_to(module, state_dict, *args, **kwargs):\n            new_state_dict = {}\n            for key, value in state_dict.items():\n                num = int(key.split(\".\")[1])  # 0 is always \"layers\"\n                new_key = key.replace(f\"layers.{num}\", module.mapping[num])\n                new_state_dict[new_key] = value\n\n            return new_state_dict\n\n        def map_from(module, state_dict, *args, **kwargs):\n            all_keys = list(state_dict.keys())\n            for key in all_keys:\n                replace_key = key.split(\".processor\")[0] + \".processor\"\n                new_key = key.replace(replace_key, f\"layers.{module.rev_mapping[replace_key]}\")\n                state_dict[new_key] = state_dict[key]\n                del state_dict[key]\n\n        self._register_state_dict_hook(map_to)\n        self._register_load_state_dict_pre_hook(map_from, with_module=True)\n\n\nclass UNet2DConditionLoadersMixin:\n    def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs):\n        r\"\"\"\n        Load pretrained attention processor layers into `UNet2DConditionModel`. Attention processor layers have to be\n        defined in\n        [cross_attention.py](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py)\n        and be a `torch.nn.Module` class.\n\n        <Tip warning={true}>\n\n            This function is experimental and might change in the future.\n\n        </Tip>\n\n        Parameters:\n            pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):\n                Can be either:\n\n                    - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.\n                      Valid model ids should have an organization name, like `google/ddpm-celebahq-256`.\n                    - A path to a *directory* containing model weights saved using [`~ModelMixin.save_config`], e.g.,\n                      `./my_model_directory/`.\n                    - A [torch state\n                      dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).\n\n            cache_dir (`Union[str, os.PathLike]`, *optional*):\n                Path to a directory in which a downloaded pretrained model configuration should be cached if the\n                standard cache should not be used.\n            force_download (`bool`, *optional*, defaults to `False`):\n                Whether or not to force the (re-)download of the model weights and configuration files, overriding the\n                cached versions if they exist.\n            resume_download (`bool`, *optional*, defaults to `False`):\n                Whether or not to delete incompletely received files. Will attempt to resume the download if such a\n                file exists.\n            proxies (`Dict[str, str]`, *optional*):\n                A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',\n                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.\n            local_files_only(`bool`, *optional*, defaults to `False`):\n                Whether or not to only look at local files (i.e., do not try to download the model).\n            use_auth_token (`str` or *bool*, *optional*):\n                The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated\n                when running `diffusers-cli login` (stored in `~/.huggingface`).\n            revision (`str`, *optional*, defaults to `\"main\"`):\n                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a\n                git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any\n                identifier allowed by git.\n            subfolder (`str`, *optional*, defaults to `\"\"`):\n                In case the relevant files are located inside a subfolder of the model repo (either remote in\n                huggingface.co or downloaded locally), you can specify the folder name here.\n\n            mirror (`str`, *optional*):\n                Mirror source to accelerate downloads in China. If you are from China and have an accessibility\n                problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety.\n                Please refer to the mirror site for more information.\n\n        <Tip>\n\n         It is required to be logged in (`huggingface-cli login`) when you want to use private or [gated\n         models](https://huggingface.co/docs/hub/models-gated#gated-models).\n\n        </Tip>\n\n        <Tip>\n\n        Activate the special [\"offline-mode\"](https://huggingface.co/diffusers/installation.html#offline-mode) to use\n        this method in a firewalled environment.\n\n        </Tip>\n        \"\"\"\n\n        cache_dir = kwargs.pop(\"cache_dir\", DIFFUSERS_CACHE)\n        force_download = kwargs.pop(\"force_download\", False)\n        resume_download = kwargs.pop(\"resume_download\", False)\n        proxies = kwargs.pop(\"proxies\", None)\n        local_files_only = kwargs.pop(\"local_files_only\", HF_HUB_OFFLINE)\n        use_auth_token = kwargs.pop(\"use_auth_token\", None)\n        revision = kwargs.pop(\"revision\", None)\n        subfolder = kwargs.pop(\"subfolder\", None)\n        weight_name = kwargs.pop(\"weight_name\", None)\n\n        user_agent = {\n            \"file_type\": \"attn_procs_weights\",\n            \"framework\": \"pytorch\",\n        }\n\n        model_file = None\n        if not isinstance(pretrained_model_name_or_path_or_dict, dict):\n            if (is_safetensors_available() and weight_name is None) or weight_name.endswith(\".safetensors\"):\n                if weight_name is None:\n                    weight_name = LORA_WEIGHT_NAME_SAFE\n                try:\n                    model_file = _get_model_file(\n                        pretrained_model_name_or_path_or_dict,\n                        weights_name=weight_name,\n                        cache_dir=cache_dir,\n                        force_download=force_download,\n                        resume_download=resume_download,\n                        proxies=proxies,\n                        local_files_only=local_files_only,\n                        use_auth_token=use_auth_token,\n                        revision=revision,\n                        subfolder=subfolder,\n                        user_agent=user_agent,\n                    )\n                    state_dict = safetensors.torch.load_file(model_file, device=\"cpu\")\n                except EnvironmentError:\n                    if weight_name == LORA_WEIGHT_NAME_SAFE:\n                        weight_name = None\n            if model_file is None:\n                if weight_name is None:\n                    weight_name = LORA_WEIGHT_NAME\n                model_file = _get_model_file(\n                    pretrained_model_name_or_path_or_dict,\n                    weights_name=weight_name,\n                    cache_dir=cache_dir,\n                    force_download=force_download,\n                    resume_download=resume_download,\n                    proxies=proxies,\n                    local_files_only=local_files_only,\n                    use_auth_token=use_auth_token,\n                    revision=revision,\n                    subfolder=subfolder,\n                    user_agent=user_agent,\n                )\n                state_dict = torch.load(model_file, map_location=\"cpu\")\n        else:\n            state_dict = pretrained_model_name_or_path_or_dict\n\n        # fill attn processors\n        attn_processors = {}\n\n        is_lora = all(\"lora\" in k for k in state_dict.keys())\n\n        if is_lora:\n            lora_grouped_dict = defaultdict(dict)\n            for key, value in state_dict.items():\n                attn_processor_key, sub_key = \".\".join(key.split(\".\")[:-3]), \".\".join(key.split(\".\")[-3:])\n                lora_grouped_dict[attn_processor_key][sub_key] = value\n\n            for key, value_dict in lora_grouped_dict.items():\n                rank = value_dict[\"to_k_lora.down.weight\"].shape[0]\n                cross_attention_dim = value_dict[\"to_k_lora.down.weight\"].shape[1]\n                hidden_size = value_dict[\"to_k_lora.up.weight\"].shape[0]\n\n                attn_processors[key] = LoRACrossAttnProcessor(\n                    hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, rank=rank\n                )\n                attn_processors[key].load_state_dict(value_dict)\n\n        else:\n            raise ValueError(f\"{model_file} does not seem to be in the correct format expected by LoRA training.\")\n\n        # set correct dtype & device\n        attn_processors = {k: v.to(device=self.device, dtype=self.dtype) for k, v in attn_processors.items()}\n\n        # set layers\n        self.set_attn_processor(attn_processors)\n\n    def save_attn_procs(\n        self,\n        save_directory: Union[str, os.PathLike],\n        is_main_process: bool = True,\n        weights_name: str = None,\n        save_function: Callable = None,\n        safe_serialization: bool = False,\n    ):\n        r\"\"\"\n        Save an attention processor to a directory, so that it can be re-loaded using the\n        `[`~loaders.UNet2DConditionLoadersMixin.load_attn_procs`]` method.\n\n        Arguments:\n            save_directory (`str` or `os.PathLike`):\n                Directory to which to save. Will be created if it doesn't exist.\n            is_main_process (`bool`, *optional*, defaults to `True`):\n                Whether the process calling this is the main process or not. Useful when in distributed training like\n                TPUs and need to call this function on all processes. In this case, set `is_main_process=True` only on\n                the main process to avoid race conditions.\n            save_function (`Callable`):\n                The function to use to save the state dictionary. Useful on distributed training like TPUs when one\n                need to replace `torch.save` by another method. Can be configured with the environment variable\n                `DIFFUSERS_SAVE_MODE`.\n        \"\"\"\n        if os.path.isfile(save_directory):\n            logger.error(f\"Provided path ({save_directory}) should be a directory, not a file\")\n            return\n\n        if save_function is None:\n            if safe_serialization:\n\n                def save_function(weights, filename):\n                    return safetensors.torch.save_file(weights, filename, metadata={\"format\": \"pt\"})\n\n            else:\n                save_function = torch.save\n\n        os.makedirs(save_directory, exist_ok=True)\n\n        model_to_save = AttnProcsLayers(self.attn_processors)\n\n        # Save the model\n        state_dict = model_to_save.state_dict()\n\n        # Clean the folder from a previous save\n        for filename in os.listdir(save_directory):\n            full_filename = os.path.join(save_directory, filename)\n            # If we have a shard file that is not going to be replaced, we delete it, but only from the main process\n            # in distributed settings to avoid race conditions.\n            weights_no_suffix = weights_name.replace(\".bin\", \"\")\n            if filename.startswith(weights_no_suffix) and os.path.isfile(full_filename) and is_main_process:\n                os.remove(full_filename)\n\n        if weights_name is None:\n            if safe_serialization:\n                weights_name = LORA_WEIGHT_NAME_SAFE\n            else:\n                weights_name = LORA_WEIGHT_NAME\n\n        # Save the model\n        save_function(state_dict, os.path.join(save_directory, weights_name))\n\n        logger.info(f\"Model weights saved in {os.path.join(save_directory, weights_name)}\")\n", "path": "src/diffusers/loaders.py"}]}