Skip to content

Commit

Permalink
Fix for "no lora weight found module" with some loras (#7875)
Browse files Browse the repository at this point in the history
* return layer weight if not found

* better system and test

* key example and typo
  • Loading branch information
asomoza authored May 7, 2024
1 parent 0d23645 commit 23e0915
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 1 deletion.
8 changes: 7 additions & 1 deletion src/diffusers/utils/peft_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
30 changes: 30 additions & 0 deletions tests/lora/test_lora_layers_sdxl.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 23e0915

Please sign in to comment.