Incompatibility between LoRA attn_procs and other attn_procs, and other LoRA inference issues

[jorgemcgomes]

Describe the bug

Why this matters:
The current implementation of attn_procs isn't really compatible with changing/switching attn_procs.
IMO, this would be quite important to add, since the lightweight nature of LoRA layers makes it an ideal candidate to rapidly switch between different "models" on a deployment use case.

Things that currently make this not really possible, all low-hanging fruit:

1. LoRA is currently mutually exclusive with xformers and attention slicing. I see there's a LoRAXFormersCrossAttnProcessor implemented, but it doesn't seem to be used anywhere.
2. I did some changes to test LoRAXFormersCrossAttnProcessor and it is not working. It's producing a shape mismatch somewhere, while the regular LoRA processor works fine with exactly the same inputs. (see solution below)
3. You can't "disable" LoRA after having enabled it. For some reason, the Lora attn_procs inherit from nn.Module, while the other attn_procs do not. Trying to change the attn_proc therefore results in an error. Ex: cannot assign 'diffusers.models.cross_attention.XFormersCrossAttnProcessor' as child module 'processor' (torch.nn.Module or None expected)
4. There should be a way of passing arguments to the attn_procs, related to How can I change alpha of LoRA? #2117 .

I managed to get all this working with some hacks/workarounds. Possible solutions:

1 - I'm not sure about the cleanest solution here. One solution would be for the attn_proc loader to be aware of xformers enabled or not, but then the order in which you enable xformers and load LoRA matters. Maybe have a single attention processor that works differently depending on whether xformers has been enabled or not?

2 - There is a bug in the lora xformers attn processor. It's missing hidden_states = attn.batch_to_head_dim(hidden_states) after this line:

diffusers/src/diffusers/models/cross_attention.py

Line 420 in 1497650

hidden_states = xformers.ops.memory_efficient_attention(query, key, value, attn_bias=attention_mask)

3 - I fixed this by wrapping another class around the Lora processor. The loading code has to be adjusted. This allows for the attn_procs to be switched from Lora to non-lora and vice-versa.

        class LoRACrossAttnProcessorWrapper:
            def __init__(self, attn_proc, scale=1.0):
                self.attn_proc = attn_proc
                self.scale = scale

            def __call__(self, *args, **kwargs):
                return self.attn_proc(*args, scale=self.scale, **kwargs)

4 - My workaround to 3 also handles this. Requires the loading code has to be adjusted.

System Info

0.12 master

[patrickvonplaten]

All spot on the issues! I'll try to open PRs for them one-by-one. Would be amazing if you'd like to help.

1. I'd say we adapt the

   diffusers/src/diffusers/models/cross_attention.py

   Line 96 in 7436e30

   def set_use_memory_efficient_attention_xformers(

   to detect if a LoRA attention processor is used and if it's the case switch to the xformers version?
2. I see, yeah I didn't really test it 😅 - doesn't surprise me that it's not working yet. Think it should be relatively easy to fix no?
3. Yeah LoRA needs to inherit from nn.Module since it's trainable. Interesting haven't seen the bug before. But it should be possible to change this -> can you maybe add a reproducible code snippet (or open a seperate issue)?
4. Answered on: How can I change alpha of LoRA? #2117

@jorgemcgomes would you maybe like to give 1. and 2. a try?

Happy to quickly open PRs for 3. and 4. :-)

[jorgemcgomes]

Thanks @patrickvonplaten . I'll only have time to work on this next week, so feel free to open any PRs you want.

I'm not sure if you missed, I posted the solution to (2) above. It's one missing line of code. I tested this and with this change both xformers and non-xformers version produce the same result.

3 - Reproducing this is quite easy:

This is the code snippet for LoRA in the release notes of 0.12

from diffusers import StableDiffusionPipeline
import torch

model_path = "sayakpaul/sd-model-finetuned-lora-t4"
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
pipe.unet.load_attn_procs(model_path)
pipe.to("cuda")

prompt = "A pokemon with blue eyes."
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
image.save("pokemon.png")

Just run this afterwards. This will set the processor to a non-nn.Module and result in an error

from diffusers.models.cross_attention import CrossAttnProcessor
pipe.unet.set_attn_processor(CrossAttnProcessor())

TypeError: cannot assign 'diffusers.models.cross_attention.CrossAttnProcessor' as child module 'processor' (torch.nn.Module or None expected)

The same happens if you do enable_attention_slicing, or enable_xformers, as both also replace the processor with a non-nn.Module.

[patrickvonplaten]

3. is fixed here: https://github.com/huggingface/diffusers/pull/2128/files

[patrickvonplaten]

Opening a PR for 4. now

[patrickvonplaten]

#2129

[patrickvonplaten]

Would be amazing if you'd like to give 1. and 2. a try (shouldn't be that hard I think). Otherwise, I'll try to look into it next week.

[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.