Repository: huggingface/diffusers

Issue:
High GPU memory usage after `set_default_attn_processor()` in unet
### Describe the bug

I was using IPAdapter and saw `unload_ip_adapter()` in `IPAdapterMixin` which helps to rollback to default unet status.
But found that after using that function (`unload_ip_adapter()`) the inference speed gets slower and GPU's reserved memory peaks too high.

### Reproduction

Posting minimal reproducible code. (Even though below code is just txt2img without IPAdapter) 

* normal case: working fine

```python
import time
import torch

from diffusers import DiffusionPipeline


# https://stackoverflow.com/a/1094933
def sizeof_fmt(num: int, suffix="B") -> str:
    for unit in ("", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"):
        if abs(num) < 1024.0:
            return f"{num:3.2f}{unit}{suffix}"
        num /= 1024.0
    return f"{num:.1f}Yi{suffix}"


common_init_kwargs = {
    "torch_dtype": torch.float16,
    "variant": "fp16",
    "use_safetensors": True,
    "feature_extractor": None,
    "safety_checker": None,
    "requires_safety_checker": False,
}

pipe = DiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", **common_init_kwargs
)
pipe.to("cuda")
pipe.enable_xformers_memory_efficient_attention()

print(f"After pipeline loading: {sizeof_fmt(torch.cuda.memory_reserved())}")


for idx in range(3):
    t_start = time.perf_counter()
    image = pipe(
        "An image of a squirrel in Picasso style",
        num_inference_steps=30,
        height=1024,
        width=1024,
    ).images[0]
    elapsed = time.perf_counter() - t_start

    print(f"After inference [{idx}]: {sizeof_fmt(torch.cuda.memory_reserved())} with {elapsed:.3f}s")
```

* weird case: 
  * Code is same as above but add single line (`pipe.unload_ip_adapter()` ) at the end.

```python
import time
import torch

from diffusers import DiffusionPipeline


# https://stackoverflow.com/a/1094933
def sizeof_fmt(num: int, suffix="B") -> str:
    for unit in ("", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"):
        if abs(num) < 1024.0:
            return f"{num:3.2f}{unit}{suffix}"
        num /= 1024.0
    return f"{num:.1f}Yi{suffix}"


common_init_kwargs = {
    "torch_dtype": torch.float16,
    "variant": "fp16",
    "use_safetensors": True,
    "feature_extractor": None,
    "safety_checker": None,
    "requires_safety_checker": False,
}

pipe = DiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", **common_init_kwargs
)
pipe.to("cuda")
pipe.enable_xformers_memory_efficient_attention()

print(f"After pipeline loading: {sizeof_fmt(torch.cuda.memory_reserved())}")


for idx in range(3):
    t_start = time.perf_counter()
    image = pipe(
        "An image of a squirrel in Picasso style",
        num_inference_steps=30,
        height=1024,
        width=1024,
    ).images[0]
    elapsed = time.perf_counter() - t_start

    print(f"After inference [{idx}]: {sizeof_fmt(torch.cuda.memory_reserved())} with {elapsed:.3f}s")
    pipe.unload_ip_adapter()
```

### Logs

```shell
* normal case: same speed and GPU memory usage
> Loading pipeline components...: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 11.80it/s]
After pipeline loading: 2.03GiB
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [00:05<00:00,  5.70it/s]
After inference [0]: 6.40GiB with 6.253s
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [00:05<00:00,  5.99it/s]
After inference [1]: 6.40GiB with 5.395s
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [00:04<00:00,  6.00it/s]
After inference [2]: 6.40GiB with 5.296s

* weird case: much slower and high gpu memory usage
> Loading pipeline components...: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 11.52it/s]
After pipeline loading: 2.03GiB
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [00:05<00:00,  5.69it/s]
After inference [0]: 6.40GiB with 6.295s
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [00:17<00:00,  1.73it/s]
After inference [1]: 22.40GiB with 17.691s
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [00:17<00:00,  1.72it/s]
After inference [2]: 22.40GiB with 17.688s
```


### System Info

- `diffusers` version: 0.26.3
- Platform: Linux-5.4.0-155-generic-x86_64-with-glibc2.31
- Python version: 3.11.8
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Huggingface_hub version: 0.21.4
- Transformers version: 4.38.2
- Accelerate version: 0.28.0
- xFormers version: 0.0.21
- Using GPU in script?: YES
- Using distributed or parallel set-up in script?: NO

### Who can help?

@yiyixuxu, @DN6, 


Discussion:
when loading a model, diffusers checks if you have pytorch 2 installed and if scaled dot-product attention is supported - that makes computations very fast - and, if that's the case, uses AttnProcessor2_0 as attention processor
when you call `pipeline.set_default_attn_processor()` (called by `pipeline.unload_ip_adapter()`), diffusers resets all attention processors to AttnProcessor, even if you have pytorch's SDPA support (used in AttnProcessor2_0 and not in AttnProcessor)

Exactly what @fabiorigano said! Thanks. Would it be helpful to maybe throw a message to the users about it (give them a hint about setting SDPA as in when needed after calling `unload_ip_adapter()`)? 

Cc: @yiyixuxu as well. 
Yes, if we don't want to change set_default_attn_processor() behaviour, it might be useful to at least let users know about it
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/diffusers/blob/main/CONTRIBUTING.md) are likely to be ignored.
@fabiorigano 
should we change the behavior then? 
I don't think we should change the behavior of `set_default_attn_processor()` since it is mostly used for testing
but we can change the behavior in `unload_ip_adapter()` to use  AttnProcessor2_0  if it is possible
hi @yiyixuxu yes, it should be an easy fix in `unload_ip_adapter()`
@fabiorigano do you want to open a PR?

Resolving patch:
diff --git a/src/diffusers/loaders/ip_adapter.py b/src/diffusers/loaders/ip_adapter.py
--- a/src/diffusers/loaders/ip_adapter.py
+++ b/src/diffusers/loaders/ip_adapter.py
@@ -16,6 +16,7 @@
 from typing import Dict, List, Optional, Union
 
 import torch
+import torch.nn.functional as F
 from huggingface_hub.utils import validate_hf_hub_args
 from safetensors import safe_open
 
@@ -38,6 +39,8 @@
     )
 
     from ..models.attention_processor import (
+        AttnProcessor,
+        AttnProcessor2_0,
         IPAdapterAttnProcessor,
         IPAdapterAttnProcessor2_0,
     )
@@ -323,4 +326,14 @@ def unload_ip_adapter(self):
         self.config.encoder_hid_dim_type = None
 
         # restore original Unet attention processors layers
-        self.unet.set_default_attn_processor()
+        attn_procs = {}
+        for name, value in self.unet.attn_processors.items():
+            attn_processor_class = (
+                AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnProcessor()
+            )
+            attn_procs[name] = (
+                attn_processor_class
+                if isinstance(value, (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0))
+                else value.__class__()
+            )
+        self.unet.set_attn_processor(attn_procs)


Test patch:
diff --git a/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py b/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py
--- a/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py
+++ b/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py
@@ -32,7 +32,6 @@
     StableDiffusionXLPipeline,
 )
 from diffusers.image_processor import IPAdapterMaskProcessor
-from diffusers.models.attention_processor import AttnProcessor, AttnProcessor2_0
 from diffusers.utils import load_image
 from diffusers.utils.testing_utils import (
     enable_full_determinism,
@@ -307,6 +306,7 @@ def test_unload(self):
         pipeline = StableDiffusionPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5", image_encoder=image_encoder, safety_checker=None, torch_dtype=self.dtype
         )
+        before_processors = [attn_proc.__class__ for attn_proc in pipeline.unet.attn_processors.values()]
         pipeline.to(torch_device)
         pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter_sd15.bin")
         pipeline.set_ip_adapter_scale(0.7)
@@ -315,11 +315,9 @@ def test_unload(self):
 
         assert getattr(pipeline, "image_encoder") is None
         assert getattr(pipeline, "feature_extractor") is not None
-        processors = [
-            isinstance(attn_proc, (AttnProcessor, AttnProcessor2_0))
-            for name, attn_proc in pipeline.unet.attn_processors.items()
-        ]
-        assert processors == [True] * len(processors)
+        after_processors = [attn_proc.__class__ for attn_proc in pipeline.unet.attn_processors.values()]
+
+        assert before_processors == after_processors
 
     @is_flaky
     def test_multi(self):