Repository: huggingface/diffusers

Issue:
Loading ex. Würstchen: AttributeError: 'Linear' object has no attribute 'lora_layer'
### Describe the bug

Loading "other" models like Wüerstchen worked fine before, but now an error is raised: 
`AttributeError: 'Linear' object has no attribute 'lora_layer'`

### Reproduction

```
from diffusers import AutoPipelineForText2Image
from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
pipe = AutoPipelineForText2Image.from_pretrained("warp-ai/wuerstchen", torch_dtype=torch.float16)
```

### Logs

```shell
Error: Python: Traceback (most recent call last):
  File "C:\Users\45239\AppData\Roaming\Blender Foundation\Blender\3.5\scripts\addons\Generative_AI-main\__init__.py", line 2646, in execute
    pipe = AutoPipelineForText2Image.from_pretrained(
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\auto_pipeline.py", line 331, in from_pretrained
    return text_2_image_cls.from_pretrained(pretrained_model_or_path, **kwargs)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\pipeline_utils.py", line 1269, in from_pretrained
    connected_pipes = {
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\pipeline_utils.py", line 1270, in <dictcomp>
    prefix: DiffusionPipeline.from_pretrained(
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\pipeline_utils.py", line 1206, in from_pretrained
    loaded_sub_model = load_sub_model(
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\pipeline_utils.py", line 470, in load_sub_model
    loaded_sub_model = load_method(os.path.join(cached_folder, name), **loading_kwargs)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\models\modeling_utils.py", line 786, in from_pretrained
    model = cls.from_config(config, **unused_kwargs)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\configuration_utils.py", line 254, in from_config
    model = cls(**init_dict)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\configuration_utils.py", line 636, in inner_init
    init(self, *args, **init_kwargs)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 62, in __init__
    self.set_default_attn_processor()
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 140, in set_default_attn_processor
    self.set_attn_processor(processor, _remove_lora=True)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 124, in set_attn_processor
    fn_recursive_attn_processor(name, module, processor)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 121, in fn_recursive_attn_processor
    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 121, in fn_recursive_attn_processor
    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 116, in fn_recursive_attn_processor
    module.set_processor(processor, _remove_lora=_remove_lora)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\models\attention_processor.py", line 319, in set_processor
    if hasattr(self, "processor") and _remove_lora and self.to_q.lora_layer is not None:
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\torch\nn\modules\module.py", line 1614, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'Linear' object has no attribute 'lora_layer'
```


### System Info

Win 11

### Who can help?

 @sayakpaul 


Discussion:
No problems here:
https://colab.research.google.com/gist/sayakpaul/0ad1e34d4f96b10addf98fa7dd3298a6/scratchpad.ipynb
Maybe it could be trickered by having previous loaded LoRAs in memory and then import Würstchen? 
(No, that's not it, I restarted it, but the error persists)

Here's the full stdout:

```
Traceback (most recent call last):
  File "C:\Users\45239\AppData\Roaming\Blender Foundation\Blender\3.5\scripts\addons\Generative_AI-main\__init__.py", line 2646, in execute
    pipe = AutoPipelineForText2Image.from_pretrained(
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\auto_pipeline.py", line 331, in from_pretrained
    return text_2_image_cls.from_pretrained(pretrained_model_or_path, **kwargs)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\pipeline_utils.py", line 1269, in from_pretrained
    connected_pipes = {
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\pipeline_utils.py", line 1270, in <dictcomp>
    prefix: DiffusionPipeline.from_pretrained(
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\pipeline_utils.py", line 1206, in from_pretrained
    loaded_sub_model = load_sub_model(
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\pipeline_utils.py", line 470, in load_sub_model
    loaded_sub_model = load_method(os.path.join(cached_folder, name), **loading_kwargs)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\models\modeling_utils.py", line 786, in from_pretrained
    model = cls.from_config(config, **unused_kwargs)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\configuration_utils.py", line 254, in from_config
    model = cls(**init_dict)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\configuration_utils.py", line 636, in inner_init
    init(self, *args, **init_kwargs)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 62, in __init__
    self.set_default_attn_processor()
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 140, in set_default_attn_processor
    self.set_attn_processor(processor, _remove_lora=True)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 124, in set_attn_processor
    fn_recursive_attn_processor(name, module, processor)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 121, in fn_recursive_attn_processor
    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 121, in fn_recursive_attn_processor
    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 116, in fn_recursive_attn_processor
    module.set_processor(processor, _remove_lora=_remove_lora)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\models\attention_processor.py", line 319, in set_processor
    if hasattr(self, "processor") and _remove_lora and self.to_q.lora_layer is not None:
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\torch\nn\modules\module.py", line 1614, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'Linear' object has no attribute 'lora_layer'
Error: Python: Traceback (most recent call last):
  File "C:\Users\45239\AppData\Roaming\Blender Foundation\Blender\3.5\scripts\addons\Generative_AI-main\__init__.py", line 2646, in execute
    pipe = AutoPipelineForText2Image.from_pretrained(
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\auto_pipeline.py", line 331, in from_pretrained
    return text_2_image_cls.from_pretrained(pretrained_model_or_path, **kwargs)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\pipeline_utils.py", line 1269, in from_pretrained
    connected_pipes = {
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\pipeline_utils.py", line 1270, in <dictcomp>
    prefix: DiffusionPipeline.from_pretrained(
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\pipeline_utils.py", line 1206, in from_pretrained
    loaded_sub_model = load_sub_model(
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\pipeline_utils.py", line 470, in load_sub_model
    loaded_sub_model = load_method(os.path.join(cached_folder, name), **loading_kwargs)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\models\modeling_utils.py", line 786, in from_pretrained
    model = cls.from_config(config, **unused_kwargs)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\configuration_utils.py", line 254, in from_config
    model = cls(**init_dict)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\configuration_utils.py", line 636, in inner_init
    init(self, *args, **init_kwargs)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 62, in __init__
    self.set_default_attn_processor()
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 140, in set_default_attn_processor
    self.set_attn_processor(processor, _remove_lora=True)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 124, in set_attn_processor
    fn_recursive_attn_processor(name, module, processor)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 121, in fn_recursive_attn_processor
    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 121, in fn_recursive_attn_processor
    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\pipelines\wuerstchen\modeling_wuerstchen_prior.py", line 116, in fn_recursive_attn_processor
    module.set_processor(processor, _remove_lora=_remove_lora)
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\diffusers\models\attention_processor.py", line 319, in set_processor
    if hasattr(self, "processor") and _remove_lora and self.to_q.lora_layer is not None:
  File "C:\Users\45239\AppData\Roaming\Python\Python310\site-packages\torch\nn\modules\module.py", line 1614, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'Linear' object has no attribute 'lora_layer'
```
I need the full snippet for that. I guess you didn't provide a full snippet in the description here: https://github.com/huggingface/diffusers/issues/5524#issue-1961211770. 
Sorry, the loading is slow in this gif, but this small snippet of code actually is causing that error:
![wuerst](https://github.com/huggingface/diffusers/assets/1322593/aa3a2f00-9204-424f-924c-c43d336288b8)

I am sorry but I couldn't reproduce it as you can see in the Colab I posted. We need something that we can fully reproduce otherwise we cannot do much. 
Maybe the snippet needs to run locally to provoke the error? 

Could it be the torch version? What torch version does the latest Diffusers need? 
I ran it from a VM too, and it was fine. 

I am using Torch 2.1. 
Updating Torch didn't solve it. 

I'm out of ideas. 

So, I guess, I'll have to disable Würstchen in Pallaidium, until someone else will encounter this bug and do something about it.
AutoPipelineForText2Image/AutoPipelineForImage2Image is also used in relation to this bug - maybe there are some funky settings/paths handling in this pipeline? https://github.com/huggingface/diffusers/issues/5578 
It really seems to be a Windows-specific problem. I couldn't reproduce it on any UNIX-based system. 
@sayakpaul @tin2tin 
I encountered the error in the environment that peft is installed.

```
➜  /workspace git:(main) ✗ pip install peft
➜  /workspace git:(main) ✗ python          
Python 3.10.12 (main, Jun 11 2023, 05:26:28) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS, WuerstchenPrior
>>> WuerstchenPrior.from_pretrained("warp-ai/wuerstchen-prior", subfolder="prior")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.10/dist-packages/diffusers/models/modeling_utils.py", line 787, in from_pretrained
    model = cls.from_config(config, **unused_kwargs)
  File "/usr/local/lib/python3.10/dist-packages/diffusers/configuration_utils.py", line 254, in from_config
    model = cls(**init_dict)
  File "/usr/local/lib/python3.10/dist-packages/diffusers/configuration_utils.py", line 644, in inner_init
    init(self, *args, **init_kwargs)
  File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py", line 62, in __init__
    self.set_default_attn_processor()
  File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py", line 140, in set_default_attn_processor
    self.set_attn_processor(processor, _remove_lora=True)
  File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py", line 124, in set_attn_processor
    fn_recursive_attn_processor(name, module, processor)
  File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py", line 121, in fn_recursive_attn_processor
    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
  File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py", line 121, in fn_recursive_attn_processor
    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
  File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py", line 116, in fn_recursive_attn_processor
    module.set_processor(processor, _remove_lora=_remove_lora)
  File "/usr/local/lib/python3.10/dist-packages/diffusers/models/attention_processor.py", line 381, in set_processor
    if hasattr(self, "processor") and _remove_lora and self.to_q.lora_layer is not None:
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1695, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'Linear' object has no attribute 'lora_layer'
```

Resolving patch:
diff --git a/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py b/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py
--- a/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py
+++ b/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py
@@ -17,6 +17,8 @@
 import torch.nn as nn
 
 from ...models.attention_processor import Attention
+from ...models.lora import LoRACompatibleConv, LoRACompatibleLinear
+from ...utils import USE_PEFT_BACKEND
 
 
 class WuerstchenLayerNorm(nn.LayerNorm):
@@ -32,7 +34,8 @@ def forward(self, x):
 class TimestepBlock(nn.Module):
     def __init__(self, c, c_timestep):
         super().__init__()
-        self.mapper = nn.Linear(c_timestep, c * 2)
+        linear_cls = nn.Linear if USE_PEFT_BACKEND else LoRACompatibleLinear
+        self.mapper = linear_cls(c_timestep, c * 2)
 
     def forward(self, x, t):
         a, b = self.mapper(t)[:, :, None, None].chunk(2, dim=1)
@@ -42,10 +45,14 @@ def forward(self, x, t):
 class ResBlock(nn.Module):
     def __init__(self, c, c_skip=0, kernel_size=3, dropout=0.0):
         super().__init__()
-        self.depthwise = nn.Conv2d(c + c_skip, c, kernel_size=kernel_size, padding=kernel_size // 2, groups=c)
+
+        conv_cls = nn.Conv2d if USE_PEFT_BACKEND else LoRACompatibleConv
+        linear_cls = nn.Linear if USE_PEFT_BACKEND else LoRACompatibleLinear
+
+        self.depthwise = conv_cls(c + c_skip, c, kernel_size=kernel_size, padding=kernel_size // 2, groups=c)
         self.norm = WuerstchenLayerNorm(c, elementwise_affine=False, eps=1e-6)
         self.channelwise = nn.Sequential(
-            nn.Linear(c, c * 4), nn.GELU(), GlobalResponseNorm(c * 4), nn.Dropout(dropout), nn.Linear(c * 4, c)
+            linear_cls(c, c * 4), nn.GELU(), GlobalResponseNorm(c * 4), nn.Dropout(dropout), linear_cls(c * 4, c)
         )
 
     def forward(self, x, x_skip=None):
@@ -73,10 +80,13 @@ def forward(self, x):
 class AttnBlock(nn.Module):
     def __init__(self, c, c_cond, nhead, self_attn=True, dropout=0.0):
         super().__init__()
+
+        linear_cls = nn.Linear if USE_PEFT_BACKEND else LoRACompatibleLinear
+
         self.self_attn = self_attn
         self.norm = WuerstchenLayerNorm(c, elementwise_affine=False, eps=1e-6)
         self.attention = Attention(query_dim=c, heads=nhead, dim_head=c // nhead, dropout=dropout, bias=True)
-        self.kv_mapper = nn.Sequential(nn.SiLU(), nn.Linear(c_cond, c))
+        self.kv_mapper = nn.Sequential(nn.SiLU(), linear_cls(c_cond, c))
 
     def forward(self, x, kv):
         kv = self.kv_mapper(kv)
diff --git a/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py b/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py
--- a/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py
+++ b/src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py
@@ -28,8 +28,9 @@
     AttnAddedKVProcessor,
     AttnProcessor,
 )
+from ...models.lora import LoRACompatibleConv, LoRACompatibleLinear
 from ...models.modeling_utils import ModelMixin
-from ...utils import is_torch_version
+from ...utils import USE_PEFT_BACKEND, is_torch_version
 from .modeling_wuerstchen_common import AttnBlock, ResBlock, TimestepBlock, WuerstchenLayerNorm
 
 
@@ -40,12 +41,15 @@ class WuerstchenPrior(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin):
     @register_to_config
     def __init__(self, c_in=16, c=1280, c_cond=1024, c_r=64, depth=16, nhead=16, dropout=0.1):
         super().__init__()
+        conv_cls = nn.Conv2d if USE_PEFT_BACKEND else LoRACompatibleConv
+        linear_cls = nn.Linear if USE_PEFT_BACKEND else LoRACompatibleLinear
+
         self.c_r = c_r
-        self.projection = nn.Conv2d(c_in, c, kernel_size=1)
+        self.projection = conv_cls(c_in, c, kernel_size=1)
         self.cond_mapper = nn.Sequential(
-            nn.Linear(c_cond, c),
+            linear_cls(c_cond, c),
             nn.LeakyReLU(0.2),
-            nn.Linear(c, c),
+            linear_cls(c, c),
         )
 
         self.blocks = nn.ModuleList()
@@ -55,7 +59,7 @@ def __init__(self, c_in=16, c=1280, c_cond=1024, c_r=64, depth=16, nhead=16, dro
             self.blocks.append(AttnBlock(c, c, nhead, self_attn=True, dropout=dropout))
         self.out = nn.Sequential(
             WuerstchenLayerNorm(c, elementwise_affine=False, eps=1e-6),
-            nn.Conv2d(c, c_in * 2, kernel_size=1),
+            conv_cls(c, c_in * 2, kernel_size=1),
         )
 
         self.gradient_checkpointing = False


Test patch:
diff --git a/tests/pipelines/wuerstchen/test_wuerstchen_prior.py b/tests/pipelines/wuerstchen/test_wuerstchen_prior.py
--- a/tests/pipelines/wuerstchen/test_wuerstchen_prior.py
+++ b/tests/pipelines/wuerstchen/test_wuerstchen_prior.py
@@ -17,11 +17,24 @@
 
 import numpy as np
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
 from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
 
 from diffusers import DDPMWuerstchenScheduler, WuerstchenPriorPipeline
+from diffusers.loaders import AttnProcsLayers
+from diffusers.models.attention_processor import (
+    LoRAAttnProcessor,
+    LoRAAttnProcessor2_0,
+)
 from diffusers.pipelines.wuerstchen import WuerstchenPrior
-from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device
+from diffusers.utils.import_utils import is_peft_available
+from diffusers.utils.testing_utils import enable_full_determinism, require_peft_backend, skip_mps, torch_device
+
+
+if is_peft_available():
+    from peft import LoraConfig
+    from peft.tuners.tuners_utils import BaseTunerLayer
 
 from ..test_pipelines_common import PipelineTesterMixin
 
@@ -29,6 +42,19 @@
 enable_full_determinism()
 
 
+def create_prior_lora_layers(unet: nn.Module):
+    lora_attn_procs = {}
+    for name in unet.attn_processors.keys():
+        lora_attn_processor_class = (
+            LoRAAttnProcessor2_0 if hasattr(F, "scaled_dot_product_attention") else LoRAAttnProcessor
+        )
+        lora_attn_procs[name] = lora_attn_processor_class(
+            hidden_size=unet.config.c,
+        )
+    unet_lora_layers = AttnProcsLayers(lora_attn_procs)
+    return lora_attn_procs, unet_lora_layers
+
+
 class WuerstchenPriorPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
     pipeline_class = WuerstchenPriorPipeline
     params = ["prompt"]
@@ -219,3 +245,52 @@ def callback_inputs_test(pipe, i, t, callback_kwargs):
 
         output = pipe(**inputs)[0]
         assert output.abs().sum() == 0
+
+    def check_if_lora_correctly_set(self, model) -> bool:
+        """
+        Checks if the LoRA layers are correctly set with peft
+        """
+        for module in model.modules():
+            if isinstance(module, BaseTunerLayer):
+                return True
+        return False
+
+    def get_lora_components(self):
+        prior = self.dummy_prior
+
+        prior_lora_config = LoraConfig(
+            r=4, lora_alpha=4, target_modules=["to_q", "to_k", "to_v", "to_out.0"], init_lora_weights=False
+        )
+
+        prior_lora_attn_procs, prior_lora_layers = create_prior_lora_layers(prior)
+
+        lora_components = {
+            "prior_lora_layers": prior_lora_layers,
+            "prior_lora_attn_procs": prior_lora_attn_procs,
+        }
+
+        return prior, prior_lora_config, lora_components
+
+    @require_peft_backend
+    def test_inference_with_prior_lora(self):
+        _, prior_lora_config, _ = self.get_lora_components()
+        device = "cpu"
+
+        components = self.get_dummy_components()
+
+        pipe = self.pipeline_class(**components)
+        pipe = pipe.to(device)
+
+        pipe.set_progress_bar_config(disable=None)
+
+        output_no_lora = pipe(**self.get_dummy_inputs(device))
+        image_embed = output_no_lora.image_embeddings
+        self.assertTrue(image_embed.shape == (1, 2, 24, 24))
+
+        pipe.prior.add_adapter(prior_lora_config)
+        self.assertTrue(self.check_if_lora_correctly_set(pipe.prior), "Lora not correctly set in prior")
+
+        output_lora = pipe(**self.get_dummy_inputs(device))
+        lora_image_embed = output_lora.image_embeddings
+
+        self.assertTrue(image_embed.shape == lora_image_embed.shape)