Repository: huggingface/diffusers

Issue:
Unloading LoRA weights from an SD pipeline doesn't work with XFormers
### Describe the bug

When running `unload_lora_weights` from a SD pipeline that has been previously loaded with XFormers attention processors, the unloading has no effect on the attention processors.

### Reproduction

```py
pipe = StableDiffusionPipeline.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    safety_checker=None,
)
if False:
    pipe.enable_xformers_memory_efficient_attention()
pipe.load_lora_weights(...)
pipe.unload_lora_weights()
pipe(...)
```

Without XFormers, everything works as expected and unloading works (so the generated final image is unaffected by the loaded lora) but when the if block runs and xformers gets enabled the LoRA starts to appear even after `unload_lora_weights()` have been called.

### Logs

_No response_

### System Info

Latest version of diffusers from the commit 47b334642261926618078c1e4e043cd0ec21fae7

### Who can help?

@sayakpaul @patrickvonplaten


Discussion:

Resolving patch:
diff --git a/src/diffusers/loaders.py b/src/diffusers/loaders.py
--- a/src/diffusers/loaders.py
+++ b/src/diffusers/loaders.py
@@ -26,6 +26,7 @@
 from torch import nn
 
 from .models.attention_processor import (
+    LORA_ATTENTION_PROCESSORS,
     AttnAddedKVProcessor,
     AttnAddedKVProcessor2_0,
     AttnProcessor,
@@ -1293,22 +1294,21 @@ def unload_lora_weights(self):
         >>> ...
         ```
         """
-        is_unet_lora = all(
-            isinstance(processor, (LoRAAttnProcessor2_0, LoRAAttnProcessor, LoRAAttnAddedKVProcessor))
-            for _, processor in self.unet.attn_processors.items()
-        )
-        # Handle attention processors that are a mix of regular attention and AddedKV
-        # attention.
-        if is_unet_lora:
-            is_attn_procs_mixed = all(
-                isinstance(processor, (LoRAAttnProcessor2_0, LoRAAttnProcessor))
-                for _, processor in self.unet.attn_processors.items()
-            )
-            if not is_attn_procs_mixed:
-                unet_attn_proc_cls = AttnProcessor2_0 if hasattr(F, "scaled_dot_product_attention") else AttnProcessor
-                self.unet.set_attn_processor(unet_attn_proc_cls())
-            else:
+        unet_attention_classes = {type(processor) for _, processor in self.unet.attn_processors.items()}
+
+        if unet_attention_classes.issubset(LORA_ATTENTION_PROCESSORS):
+            # Handle attention processors that are a mix of regular attention and AddedKV
+            # attention.
+            if len(unet_attention_classes) > 1 or LoRAAttnAddedKVProcessor in unet_attention_classes:
                 self.unet.set_default_attn_processor()
+            else:
+                regular_attention_classes = {
+                    LoRAAttnProcessor: AttnProcessor,
+                    LoRAAttnProcessor2_0: AttnProcessor2_0,
+                    LoRAXFormersAttnProcessor: XFormersAttnProcessor,
+                }
+                [attention_proc_class] = unet_attention_classes
+                self.unet.set_attn_processor(regular_attention_classes[attention_proc_class]())
 
         # Safe to call the following regardless of LoRA.
         self._remove_text_encoder_monkey_patch()
diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py
--- a/src/diffusers/models/attention_processor.py
+++ b/src/diffusers/models/attention_processor.py
@@ -167,7 +167,7 @@ def set_use_memory_efficient_attention_xformers(
     ):
         is_lora = hasattr(self, "processor") and isinstance(
             self.processor,
-            (LoRAAttnProcessor, LoRAAttnProcessor2_0, LoRAXFormersAttnProcessor, LoRAAttnAddedKVProcessor),
+            LORA_ATTENTION_PROCESSORS,
         )
         is_custom_diffusion = hasattr(self, "processor") and isinstance(
             self.processor, (CustomDiffusionAttnProcessor, CustomDiffusionXFormersAttnProcessor)
@@ -1623,6 +1623,13 @@ def __call__(self, attn: "Attention", hidden_states, encoder_hidden_states=None,
     CustomDiffusionXFormersAttnProcessor,
 ]
 
+LORA_ATTENTION_PROCESSORS = (
+    LoRAAttnProcessor,
+    LoRAAttnProcessor2_0,
+    LoRAXFormersAttnProcessor,
+    LoRAAttnAddedKVProcessor,
+)
+
 
 class SpatialNorm(nn.Module):
     """


Test patch:
diff --git a/tests/models/test_lora_layers.py b/tests/models/test_lora_layers.py
--- a/tests/models/test_lora_layers.py
+++ b/tests/models/test_lora_layers.py
@@ -464,6 +464,14 @@ def test_lora_unet_attn_processors_with_xformers(self):
                 if isinstance(module, Attention):
                     self.assertIsInstance(module.processor, LoRAXFormersAttnProcessor)
 
+            # unload lora weights
+            sd_pipe.unload_lora_weights()
+
+            # check if attention processors are reverted back to xFormers
+            for _, module in sd_pipe.unet.named_modules():
+                if isinstance(module, Attention):
+                    self.assertIsInstance(module.processor, XFormersAttnProcessor)
+
     @unittest.skipIf(torch_device != "cuda", "This test is supposed to run on GPU")
     def test_lora_save_load_with_xformers(self):
         pipeline_components, lora_components = self.get_dummy_components()