Repository: huggingface/diffusers

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.


Discussion:
_The documentation is not available anymore as the PR was closed or merged._
I don't think the failing test is linked to this PR, is it ?
Hi @Narsil! The code looks good to me, and the failing tests have nothing to do with it.

For additional context on top of the `safetensors` issue you mentioned, people are interested in converting LoRA weights generated with other tools. See for example: https://github.com/huggingface/diffusers/issues/2363, https://github.com/huggingface/diffusers/pull/2403. And also the other way around (being able to use diffusers LoRA weights in other tools): https://github.com/huggingface/diffusers/issues/2326.

Focusing on the first task (converting from other tools to diffusers), I'm not sure how hard the problem is. In addition to changes in key names, some tools seem to do pivotal tuning or text inversion in addition to the cross-attention layers in our implementation. [This PR](https://github.com/huggingface/diffusers/pull/2403) seems to save the full pipeline instead of attempting to convert the incremental weights.

TL;DR: I think this PR is quite useful and necessary, but not sure if it will help towards the issue you mentioned :) (But I may be wrong, I still have to find the time to test this sort of interoperability).
Shall I merge ?
Please merge.
I would like to have @patil-suraj also review it. But he is currently on leave and should be back early next week. If it can wait, I would like to wait for that while. 
Done.
Hey all, thanks for all the amazing work here.

I need to spend a bit more time on this but I think this introduces a regression.  I have an integration test that calls:

```python
// https://huggingface.co/patrickvonplaten/lora_dreambooth_dog_example/resolve/main/pytorch_lora_weights.bin
pipe.unet.load_attn_procs("pytorch_lora_weights.bin")
```

and is now failing with:

```
  File "/api/diffusers/src/diffusers/loaders.py", line 170, in load_attn_procs
    state_dict = safetensors.torch.load_file(model_file, device="cpu")
  File "/opt/conda/envs/xformers/lib/python3.9/site-packages/safetensors/torch.py", line 98, in load_file
    with safe_open(filename, framework="pt", device=device) as f:
Exception: Error while deserializing header: HeaderTooLarge
```

I guess because it's trying to load a non-safetensors file with `safetensors.torch.load_file` ?  Here's the relevant code (last line is the failing line):

https://github.com/huggingface/diffusers/blob/1f4deb697f6204ae5da788b1e10c3074208ee57c/src/diffusers/loaders.py#L153-L170

There's no exception raised because the file indeed exists, it's just not in safetensors format.  So I guess we need a `safe_serialization` or `from_safetensors` or some such kwarg maybe?

Looking at the surrounding code the easy way around this is to call `torch.load()` myself and pass the `state_dict` but I don't think that's a good approach.  Let me know if I'm doing anything wrong but I think I should be able to specify an exact file, right?

Looking back at the original PR intro I guess this is exactly what @Narsil says:

> What's techically missing is the option to choose on load what the format is along with the weights_name. I didn't want to add it here for simplicity (Since most users should be using the default anyway). But we could add that.

So should I open a new issue for this?

Thanks!

Resolving patch:
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 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict
 
         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:


Test patch:
diff --git a/tests/models/test_models_unet_2d_condition.py b/tests/models/test_models_unet_2d_condition.py
--- a/tests/models/test_models_unet_2d_condition.py
+++ b/tests/models/test_models_unet_2d_condition.py
@@ -445,6 +445,43 @@ def test_lora_save_load_safetensors(self):
         # LoRA and no LoRA should NOT be the same
         assert (sample - old_sample).abs().max() > 1e-4
 
+    def test_lora_save_load_safetensors_load_torch(self):
+        # enable deterministic behavior for gradient checkpointing
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+
+        init_dict["attention_head_dim"] = (8, 16)
+
+        torch.manual_seed(0)
+        model = self.model_class(**init_dict)
+        model.to(torch_device)
+
+        lora_attn_procs = {}
+        for name in model.attn_processors.keys():
+            cross_attention_dim = None if name.endswith("attn1.processor") else model.config.cross_attention_dim
+            if name.startswith("mid_block"):
+                hidden_size = model.config.block_out_channels[-1]
+            elif name.startswith("up_blocks"):
+                block_id = int(name[len("up_blocks.")])
+                hidden_size = list(reversed(model.config.block_out_channels))[block_id]
+            elif name.startswith("down_blocks"):
+                block_id = int(name[len("down_blocks.")])
+                hidden_size = model.config.block_out_channels[block_id]
+
+            lora_attn_procs[name] = LoRACrossAttnProcessor(
+                hidden_size=hidden_size, cross_attention_dim=cross_attention_dim
+            )
+            lora_attn_procs[name] = lora_attn_procs[name].to(model.device)
+
+        model.set_attn_processor(lora_attn_procs)
+        # Saving as torch, properly reloads with directly filename
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            model.save_attn_procs(tmpdirname)
+            self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_lora_weights.bin")))
+            torch.manual_seed(0)
+            new_model = self.model_class(**init_dict)
+            new_model.to(torch_device)
+            new_model.load_attn_procs(tmpdirname, weight_name="pytorch_lora_weights.bin")
+
     def test_lora_on_off(self):
         # enable deterministic behavior for gradient checkpointing
         init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()