Repository: huggingface/diffusers

Issue:
Error when using blockwise scales with sd_xl_offset_example-lora_1.0.safetensors
### Describe the bug

When loading the [sd_xl_offset_example-lora_1.0.safetensors](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_offset_example-lora_1.0.safetensors) from the [official repository](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/) and setting the scales by blocks, it throws this error:

`RuntimeError: No LoRA weight found for module down_blocks.0.resnets.0.conv1`

Loading the same lora with just a float e.g.: `pipe.set_adapters(["offset"], [1.0])` doesn't raise the error.

### Reproduction

```python
import torch

from diffusers import StableDiffusionXLPipeline


pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16"
).to("cuda")

pipe.load_lora_weights(
    "stabilityai/stable-diffusion-xl-base-1.0",
    weight_name="sd_xl_offset_example-lora_1.0.safetensors",
    adapter_name="offset",
)
offset_scales = {
    "unet": {
        "down": {"block_1": [1.0, 1.0], "block_2": [1.0, 1.0]},
        "mid": 0,
        "up": {"block_0": [1.0, 1.0, 1.0], "block_1": [1.0, 1.0, 1.0]},
    },
    "text_encoder": 1.0,
    "text_encoder_2": 1.0,
}
pipe.set_adapters(["offset"], [offset_scales])

prompt = "high quality photo of a panther in the jungle, night, dark, low light, dim"
generator = torch.Generator(device="cpu").manual_seed(899041154)

image = pipe(
    prompt=prompt,
    negative_prompt="",
    guidance_scale=6.5,
    num_inference_steps=25,
    generator=generator,
).images[0]
image.save("result.png")
```

### Logs

_No response_

### System Info

- `diffusers` version: 0.28.0.dev0
- Platform: Linux-6.8.7-arch1-1-x86_64-with-glibc2.39
- Python version: 3.11.8
- PyTorch version (GPU?): 2.2.2+cu121 (True)
- Huggingface_hub version: 0.20.3
- Transformers version: 4.40.1
- Accelerate version: 0.30.0
- xFormers version: not installed

### Who can help?

@sayakpaul @UmerHA 


Discussion:
I can open a PR to fix it but first I would like to know if there is a reason behind raising the error here:

https://github.com/huggingface/diffusers/blob/7fa3e5b0f6a593d06e65f5b40ccd46acfafcfeb1/src/diffusers/utils/peft_utils.py#L249

instead of just returning a default value.
@asomoza At this point the lora weight dict is expanded and should contain an entry for every peft-tunable module. So that error should never happen.

My naive guess is that convolutions are now also peft-tunable (ie a subclass of `BaseTunerLayer`) but the dict expansion doesn't reflect that.

Imo returning a default value is fine.

thanks for the clarification, I really appreciate it.

Resolving patch:
diff --git a/src/diffusers/utils/peft_utils.py b/src/diffusers/utils/peft_utils.py
--- a/src/diffusers/utils/peft_utils.py
+++ b/src/diffusers/utils/peft_utils.py
@@ -246,7 +246,13 @@ def get_module_weight(weight_for_adapter, module_name):
         for layer_name, weight_ in weight_for_adapter.items():
             if layer_name in module_name:
                 return weight_
-        raise RuntimeError(f"No LoRA weight found for module {module_name}.")
+
+        parts = module_name.split(".")
+        # e.g. key = "down_blocks.1.attentions.0"
+        key = f"{parts[0]}.{parts[1]}.attentions.{parts[3]}"
+        block_weight = weight_for_adapter.get(key, 1.0)
+
+        return block_weight
 
     # iterate over each adapter, make it active and set the corresponding scaling weight
     for adapter_name, weight in zip(adapter_names, weights):


Test patch:
diff --git a/tests/lora/test_lora_layers_sdxl.py b/tests/lora/test_lora_layers_sdxl.py
--- a/tests/lora/test_lora_layers_sdxl.py
+++ b/tests/lora/test_lora_layers_sdxl.py
@@ -202,6 +202,36 @@ def test_sdxl_1_0_lora(self):
         pipe.unload_lora_weights()
         release_memory(pipe)
 
+    def test_sdxl_1_0_blockwise_lora(self):
+        generator = torch.Generator("cpu").manual_seed(0)
+
+        pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
+        pipe.enable_model_cpu_offload()
+        lora_model_id = "hf-internal-testing/sdxl-1.0-lora"
+        lora_filename = "sd_xl_offset_example-lora_1.0.safetensors"
+        pipe.load_lora_weights(lora_model_id, weight_name=lora_filename, adapter_name="offset")
+        scales = {
+            "unet": {
+                "down": {"block_1": [1.0, 1.0], "block_2": [1.0, 1.0]},
+                "mid": 1.0,
+                "up": {"block_0": [1.0, 1.0, 1.0], "block_1": [1.0, 1.0, 1.0]},
+            },
+        }
+        pipe.set_adapters(["offset"], [scales])
+
+        images = pipe(
+            "masterpiece, best quality, mountain", output_type="np", generator=generator, num_inference_steps=2
+        ).images
+
+        images = images[0, -3:, -3:, -1].flatten()
+        expected = np.array([0.4468, 0.4087, 0.4134, 0.366, 0.3202, 0.3505, 0.3786, 0.387, 0.3535])
+
+        max_diff = numpy_cosine_similarity_distance(expected, images)
+        assert max_diff < 1e-4
+
+        pipe.unload_lora_weights()
+        release_memory(pipe)
+
     def test_sdxl_lcm_lora(self):
         pipe = StableDiffusionXLPipeline.from_pretrained(
             "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16