Fix for "no lora weight found module" with some loras

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):

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