Fix lora inject, added weight self apply lora

[DavidePaglieri]

Hi, awesome job!

I might be mistaken, but I think there might be a bug in the current version of inject_trainable_lora.

The cross attention module in the unet initially is:

CrossAttention(
              (to_q): Linear(in_features=1280, out_features=1280, bias=False)
              (to_k): Linear(in_features=768, out_features=1280, bias=False)
              (to_v): Linear(in_features=768, out_features=1280, bias=False)
              (to_out): ModuleList(
                (0): Linear(in_features=1280, out_features=1280, bias=True)
                (1): Dropout(p=0.0, inplace=False)
              )
            )

After injecting lora with the current method it becomes:

CrossAttention(
		[…]
              (to_out): ModuleList(
                (0): Linear(in_features=1280, out_features=1280, bias=True)
                (1): Dropout(p=0.0, inplace=False)
              )
              (to_out.0): LoraInjectedLinear(
                (linear): Linear(in_features=1280, out_features=1280, bias=True)
                (lora_down): Linear(in_features=1280, out_features=4, bias=False)
                (lora_up): Linear(in_features=4, out_features=1280, bias=False)
              )
            )

However if you go and see how CrossAttention works in diffusers you can see that the forward calls self.to_out[0] and self.to_out[1] which I think would probably call the original layer and dropout in the ModuleList, and not the LoraInjectedLayer.

In the PR I should have fixed this, now the cross attention injected becomes:

CrossAttention(
		[…]
              (to_out): ModuleList(
                (0): LoraInjectedLinear(
                  (linear): Linear(in_features=1280, out_features=1280, bias=True)
                  (lora_down): Linear(in_features=1280, out_features=4, bias=False)
                  (lora_up): Linear(in_features=4, out_features=1280, bias=False)
                )
                (1): Dropout(p=0.0, inplace=False)
              )
            )

Which should make it work correctly. I have also added weight_self_apply_lora which self applies the lora weights to a lora_injected model, and returns the injected model in standard form.

Let me know if this makes sense to you.

[cloneofsimo]

Amazing! How did you find this? I guess this isn't critical issue, as qkv gets tuned, but could be certainly underperforming factor. Thank you so much!! Can you switch the branch to develop?

[cloneofsimo]

@DavidePaglieri ah, so I found an issue regarding this PR. So this assume that to_out behaving in specific way, which CLIPTextTransformer might not agree, or any other method really because it take an target_replace_module in general.

[cloneofsimo]

I think this was amazing issue that you've found. I'll mark this as an enhancement and work on making this more explicit / better later on.

[cloneofsimo]

Maybe making submodule for look with depth 1 and adding lambda module : module isinstance(nn.ModuleList) and has dropout as a filter might work

[cloneofsimo]

In any way, this needs major fix as you've pointed out

[DavidePaglieri]

I found it when implementing the weight self apply method as there was a mismatch in the layers. Yes I mostly checked it with the unet, but the fix I made would not be super general. Agree that a condition based on whether it finds the nn.ModuleList could work well.

[cloneofsimo]

Kind of in a mess right now. Probably just better to reimplement the stable diffusion (forked from diffusers or whatnot), and modify alphas and etc in non-hacky, non-monkey-patchy-way, but that would be too much work haha

[hafriedlander]

Ran across the same bug.

This is probably fixable by looking for "." in the module name, and recursing through the model tree if it's found.

I'm working on a more generic patching method (that can also remove the injected LoraInjectedLinears and replace them with normal Linears again) that will fix this. I'll link here once I'm done and you can comment if you like the approach.

[hafriedlander]

Here's my rewrite of the injector (based off main - hopefully develop hasn't changed things too much). If this approach suits, I'll make a PR to add the same to this codebase.

https://github.com/hafriedlander/stable-diffusion-grpcserver/blob/dev/sdgrpcserver/pipeline/lora.py

[hafriedlander]

@cloneofsimo Take a look at the file above - is that an approach you'd accept a PR for?