Implements Blockwise lora

[UmerHA]

What does this PR do?

Allows setting LoRA weights more granularly, up to per-transformer block.

Fixes #7231

Example usage:

pipe = ... # create pipeline
pipe.load_lora_weights(..., adapter_name="my_adapter") 
scales = {
    "text_encoder": 0.5,
    "unet": {
        "down": 0.9,  # all transformers in the down-part will use scale 0.9
        # "mid"  # because "mid" is not given, all transformers in the mid part will use the default scale 1.0
        "up": {
            "block_0": 0.6,  # all 3 transformers in the 0th block in the up-part will use scale 0.6
            "block_1": [0.4, 0.8, 1.0],  # the 3 transformers in the 1st block in the up-part will use scales 0.4, 0.8 and 1.0 respectively
        }
    }
}
pipe.set_adapters("my_adapter", scales)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Core library:

• General functionalities: @sayakpaul @yiyixuxu @DN6

@dain5832: As you're the author of the solved issue, I invite you to also test this.

[asomoza]

this PR is really nice, I like to be able to have this kind of control over the LoRAs. Just by playing with it I found a way to break it though. If you do this:

scales = {
    "text_encoder": 0.5,
}
pipeline.set_adapters(["test"], adapter_weights=[scales])

it throws:

unet.py", line 611, in _expand_lora_scales_dict
    if "mid" not in scales:
TypeError: argument of type 'NoneType' is not iterable

Also I would like to request if you can add to have control over the second text encoder in SDXL but if it's rare edge case, it doesn't matter, I can do it on my side after.

[UmerHA]

this PR is really nice

❤️

Just by playing with it I found a way to break it though.

fixed + expanded tests

Also I would like to request if you can add to have control over the second text encoder in SDXL

added

---

Let me know if there are any other issues!

[sayakpaul]

I think the PR has a very nice start. I left some initial comments. We'll need to make sure to add rigorous checks on the inputs to prevent anything unexpected that we can foresee from our experience.

[UmerHA]

Re mid block: We can pass a "mid" key in the scales dict. It's commented out to show not all keys need to be provided. But the mid block only needs 1 scale number, because it is only 1 block and has only 1 layer.

In principal, mid blocks could have more layers, but UNet2DConditionModel doesn't pass num_layers to the mid block, so it always defaults to 1.

Re convs: The original implementation doesn't implements scaling the conv layers.

---

• Would you also implement scaling the conv layers?
• Would you also implement scaling the resnets? Currently, only the transformer are scaled.

Let me know, then I'll adapt the code.

@sayakpaul

[sayakpaul]

Re mid block: We can pass a "mid" key in the scales dict. It's commented out to show not all keys need to be provided. But the mid block only needs 1 scale number, because it is only 1 block and has only 1 layer.

I still don't see why we shouldn't make it configurable as well actually. Even if it's just a single block there might be effects if the users can control its contributions. So, I would vote for the ability to configure it too.

Would you also implement scaling the conv layers?
Would you also implement scaling the resnets? Currently, only the transformer are scaled.

All of them should configurable IMO if we were to support this feature otherwise the design might feel very convoluted. But I would like to also take opinions from @BenjaminBossan and @yiyixuxu here.

[UmerHA]

I still don't see why we shouldn't make it configurable as well actually. Even if it's just a single block there might be effects if the users can control its contributions. So, I would vote for the ability to configure it too.

I think we're miscommunicating. With the current PR, You CAN already configure the mid block:

pipe = ... # create pipeline
pipe.load_lora_weights(..., adapter_name="my_adapter") 
scales = { "mid" : 0.5 }
pipe.set_adapters("my_adapter", scales)

However, you can only pass a single number to it, because it has in total only 1 transformer layer in it. Unlike the up / down part, which have multiple blocks and multiple layers per block. That's why you can pass a Dict (per block) or Dict[List](per block, per layer) for them.

[sayakpaul]

However, you can only pass a single number to it, because it has in total only 1 transformer layer in it. Unlike the up / down part, which have multiple blocks and multiple layers per block. That's why you can pass a Dict (per block) or Dict[List](per block, per layer) for them.

Ah makes sense then.

[BenjaminBossan]

Thanks for adding this feature. I don't have sufficient knowledge to judge if this covers all edge cases, so I'll leave that to the experts. Still added a few comments.

[stevhliu]

Awesome, great work! 👏

[yiyixuxu]

hey! thanks for your PR
I did first round of review - let me know if it makes sense

Thanks
YiYi

[sayakpaul]

@UmerHA would be helpful to resolve the comments that you have already addressed. But do feel free to keep them open if you're unsure.

Will give this a thorough look tomorrow. Making a reminder.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[BenjaminBossan]

LGTM overall, nice tests. I just have some smaller comments, please check.

All in all, I don't know enough about diffusers to give a final approval, so I'll leave that to the experts.

[UmerHA]

would something like this work?

With a small adjustment, yes. invert_list_adapters needs to be created a bit differently, because self.get_list_adapters() is not Dict[str, str] but Dict[str, List[str]].

This works:

list_adapters = self.get_list_adapters()  # eg {"unet": ["adapter1", "adapter2"], "text_encoder": ["adapter2"]}
all_adapters = {
    adapter for adapters in list_adapters.values() for adapter in adapters
}  # eg ["adapter1", "adapter2"]
invert_list_adapters = {
    adapter: [part for part, adapters in list_adapters.items() if adapter in adapters]
    for adapter in all_adapters
}  # eg {"adapter1": ["unet"], "adapter2": ["unet", "text_encoder"]}

I have adjusted to code accordingly, and renamed lora_loaders_util.py back to unet_loaders_util.py

[UmerHA]

@sayakpaul gentle ping, as this has been stale for a week

[sayakpaul]

Will do a thorough review tomorrow first thing. Thanks for your patience.

[sayakpaul]

Left a last couple of comments. Looking really nice. Going to run some low tests to ensure we're not breaking anything here.

Really amazing work!

[sayakpaul]

LoRA SD slow tests are passing which should be more than enough to confirm the PR isn't backward breaking. Let's address the last comments and ship this beast

[UmerHA]

@sayakpaul Everything is addressed! Iiuc, you now have to accept the PR for the uploaded images (https://huggingface.co/datasets/huggingface/documentation-images/discussions/310) and then we're done

[sayakpaul]

Already merged.

[asomoza]

really nice, when I have the time I'll play a lot with this and maybe post my results so people can understand better what can they do with this level of control over the LoRAs

[sayakpaul]

@UmerHA could you push an empty commit to start the CI? And please ping me once you do.

[UmerHA]

@sayakpaul empty commit pushed

[sayakpaul]

Thanks @UmerHA!

Will wait for the CI to run and will then ship.

To help promote this amazing feature, I welcome you to open a detailed thread here: https://github.com/huggingface/diffusers/discussions.

Additionally, if you're on HF Discord, please let me know your username so that we can give you a shoutout. If you're not, you can join via this link: https://hf.co/join.discord.

[UmerHA]

Thanks @UmerHA!

❤️

To help promote this amazing feature, I welcome you to open a detailed thread here: https://github.com/huggingface/diffusers/discussions.

Will do this evening/tomorrow ✅

Additionally, if you're on HF Discord, please let me know your username so that we can give you a shoutout. If you're not, you can join via this link: https://hf.co/join.discord.

Thanks :) My discord is umerha. But my primary channel is Twitter where I am @UmerHAdil. If possible, I'd highly appreciate a shootout over there 🙏🏽

[asomoza]

@UmerHA I'll leave this here because I can't do a PR right now for a quick fix.

Just updated my main and ran the test I did before, there a small typo uent that throws a warning:

Lora weight dict for adapter 'lora' contains uent, but this will be ignored because lora does not contain weights for uent. Valid parts for lora are: ['text_encoder', 'text_encoder_2', 'unet'].

[UmerHA]

there a small typo uent that throws a warning:

@asomoza thanks! have made a fix PR here

[UmerHA]

To help promote this amazing feature, I welcome you to open a detailed thread here: https://github.com/huggingface/diffusers/discussions.

Will do this evening/tomorrow ✅

Done ✅