Support to load Kohya-ss style LoRA file format (without restrictions)

[takuma104]

What's the PR for?

Discussed in #3064 (comment), PR #3437. In the PR #3437, we implemented the loading of Kohya-ss style LoRA, but this was only support for the Attention part, which had already been implemented in Diffusers. Kohya-ss style LoRA is not only extended to Attention, but also to ClipMLP, Transformer2DModel's proj_in, proj_out, and so on. This PR will support these remaining issues.

Modification:

• Added modes/lora.py. I've moved the LoRALinearLayer here. For nn.Conv2d and nn.Linear that add LoRA to, I've added Conv2dWithLoRA and LinearWithLoRA respectively.
• Changed nn.Linear in FeedForward and GEGLU classes to LinearWithLoRA
• Changed .proj_in, .proj_out in Transformer2DModel class to Conv2dWithLoRA
• Applied LoRALinearLayer as a monkey-patch to mlp in text_encoder

Impact:

In the Transformer2DModel and FeedForward classes, some of the nn.Linear and nn.Conv2d have been replaced with their subclass. We need to thoroughly confirm that this does not cause any issues in normal use without LoRA. Currently, unit tests that check for class name matching are not passing as follows.

FAILED tests/models/test_layers_utils.py::Transformer2DModelTests::test_spatial_transformer_default_ff_layers - AssertionError: assert <class 'diffusers.models.lora.LinearWithLoRA'> == <class 'torch.nn.modules.linear.Linear'>
FAILED tests/models/test_layers_utils.py::Transformer2DModelTests::test_spatial_transformer_geglu_approx_ff_layers - AssertionError: assert <class 'diffusers.models.lora.LinearWithLoRA'> == <class 'torch.nn.modules.linear.Linear'>

Generation Comparison:

LoRA file: Light and Shadow These results should be reproduced by this script.

Fully Applied LoRA (hook)	This PR	Diffusers current release

Todo:

• Remove test_kohya_loras_scaffold.py
• Fix tests in test_layers_utils.py

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[takuma104]

In PR #3778, the monkey-patch method is being refactored, so I'm thinking of resuming work after #3778 is merged.

[patrickvonplaten]

cc @sayakpaul could you take a look here? :-)

[sayakpaul]

In PR #3778, the monkey-patch method is being refactored, so I'm thinking of resuming work after #3778 is merged.

I think that would be better as it will allow us to reduce redundant work. Let's make sure to revisit after #3778 is dealt with.

Cc: @williamberman

[isidentical]

Hey @takuma104, thank you for this PR -- I think it is amazing. For the last week or so, I have been trying to get CivitAILoRAs up and running with diffusers but it always produced super weird (and sometimes irrelevant) results such as this:

But using this PR and not changing anything else (same prompt, same seed, same configuration), I can get super precise and good looking pictures from this random model (https://civitai.com/models/24380/minty-doges-doge-style-lora)

Seeing this PR is still in draft stage, I wonder whether there is a more official way of loading these sort of LoRAs or whether this PR will get landed soon.

[jfischoff]

I want to add I have also used this PR locally and it appears to work well.

[jfischoff]

One issue.

When I use this PR I'm getting this error intermittently. Only happens on this PR when loading loras. Occurs about 1/5 times.

Traceback (most recent call last):
  File "/home/huggingface/.local/lib/python3.9/site-packages/diffusers/models/modeling_utils.py", line 66, in get_parameter_device
    return next(parameters_and_buffers).device
StopIteration

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  ... 
    pipe.load_lora_weights(lora_folder,
  File "/home/huggingface/.local/lib/python3.9/site-packages/diffusers/loaders.py", line 978, in load_lora_weights
    attn_procs_text_encoder = self._load_text_encoder_attn_procs(
  File "/home/huggingface/.local/lib/python3.9/site-packages/diffusers/loaders.py", line 1232, in _load_text_encoder_attn_procs
    attn_processors = {
  File "/home/huggingface/.local/lib/python3.9/site-packages/diffusers/loaders.py", line 1233, in <dictcomp>
    k: v.to(device=self.device, dtype=self.text_encoder.dtype) for k, v in attn_processors.items()
  File "/home/huggingface/.local/lib/python3.9/site-packages/diffusers/pipelines/pipeline_utils.py", line 710, in device
    return module.device
  File "/home/huggingface/.local/lib/python3.9/site-packages/diffusers/models/modeling_utils.py", line 812, in device
    return get_parameter_device(self)
  File "/home/huggingface/.local/lib/python3.9/site-packages/diffusers/models/modeling_utils.py", line 75, in get_parameter_device
    first_tuple = next(gen)
StopIteration

[sayakpaul]

@takuma104 hey!

Since #3778 has been merged now, let's pick this one up :-)

[sayakpaul]

Seeing this PR is still in draft stage, I wonder whether there is a more official way of loading these sort of LoRAs or whether this PR will get landed soon.

There will always be some amount of limitation in loading non-diffusers LoRAs as there are way too many formats to cater to. So, we're doing our bests to provide a unified support.

[isidentical]

Since #3778 has been merged now, let's pick this one up :-)

Hey @sayakpaul / @takuma104, in case you folks might not have the availability / time, I would be more than happy to assist on reviving this PR (either preparing a new PR that takes this + rebases with the changes of #3778 OR submitting a PR to the @takuma104's repo) with your permission. We have been using this branch for a couple of weeks and it has been working flawlessly so I assume not much needs to change in regards to the actual logic.

[sayakpaul]

More than happy to collaborate and ship this. Let's start with a new PR superseeding this. Let's make sure to have Takuma as a co-author of the commits. Works?

[isidentical]

More than happy to collaborate and ship this. Let's start with a new PR superseeding this. Let's make sure to have Takuma as a co-author of the commits. Works?

Definitely! Will try to get a working version (compatible with the new unloading logic and the refactored layout) and then I can submit a PR which we can iterate on it!

[sayakpaul]

Let's go then! Thanks so much!

@takuma104 I hope this is alright :)

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[Georgefwt]

Hello, has the pull request that will ultimately replace this one been submitted? Additionally, Does anyone plan to update this content in the Lora training script?