Implement rest of the test cases (LoRA tests)

[aandyw]

PR for issue #2789

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[aandyw]

@patrickvonplaten Haven't pushed the updated code but atm I'm thinking that we add another conditional statement to:

if name.startswith("mid_block"):
      hidden_size = model.config.block_out_channels[-1]
  elif name.startswith("up_blocks"):
      block_id = int(name[len("up_blocks.")])
      hidden_size = list(reversed(model.config.block_out_channels))[block_id]
  elif name.startswith("down_blocks"):
      block_id = int(name[len("down_blocks.")])
      hidden_size = model.config.block_out_channels[block_id]

for name.startswith("transformer_in"). However, I'm not quite sure how to deal with these initial layers. I've tried ignoring them but it seems like model.set_attn_processor(lora_attn_procs) requires lora_attn_procs to contain something for the transformer_in layers.

[patrickvonplaten]

Looks very cool! Think your added tests are great! Went into the PR to help a bit. test_lora_processors now passes and I think you can use its logic to make the other tests as well - see: https://github.com/huggingface/diffusers/pull/2824/files#r1157208349

Let me know if you need any more help :-)

cc @sayakpaul could you also take a look here?

[aandyw]

Looks very cool! Think your added tests are great! Went into the PR to help a bit. test_lora_processors now passes and I think you can use its logic to make the other tests as well - see: https://github.com/huggingface/diffusers/pull/2824/files#r1157208349

Let me know if you need any more help :-)

cc @sayakpaul could you also take a look here?

Thank you for the help! I really appreciate it. Would you be able to recommend some resources that I could look into to learn more about this structure and what the transformer_in is doing specifically? I'm amazed by how fast you were able to figure out the PR.

[sayakpaul]

Love this!

[sayakpaul]

We can leverage the create_lora_layers() here and elsewhere, no?

[patrickvonplaten]

Agree - @pie31415 could we maybe as a final todo factor out this code:

        for name in model.attn_processors.keys():
            has_cross_attention = name.endswith("attn2.processor") and not (
                name.startswith("transformer_in") or "temp_attentions" in name.split(".")
            )
            cross_attention_dim = model.config.cross_attention_dim if has_cross_attention else None

            if name.startswith("mid_block"):
                hidden_size = model.config.block_out_channels[-1]
            elif name.startswith("up_block"):
                block_id = int(name[len("up_blocks.")])
                hidden_size = list(reversed(model.config.block_out_channels))[block_id]
            elif name.startswith("down_blocks"):
                block_id = int(name[len("down_blocks.")])
                hidden_size = model.config.block_out_channels[block_id]
            elif name.startswith("transformer_in"):
                # Note that the `8 * ...` comes from: https://github.com/huggingface/diffusers/blob/7139f0e874f10b2463caa8cbd585762a309d12d6/src/diffusers/models/unet_3d_condition.py#L148
                hidden_size = 8 * model.config.attention_head_dim

            lora_attn_procs[name] = LoRAAttnProcessor(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim)

            with torch.no_grad():
                lora_attn_procs[name].to_q_lora.up.weight += 1
                lora_attn_procs[name].to_k_lora.up.weight += 1
                lora_attn_procs[name].to_v_lora.up.weight += 1
                lora_attn_procs[name].to_out_lora.up.weight += 1

into a create_lora_layers() as it's used three times further below? :-)

Think after this we're good for merge ❤️

[sayakpaul]

Thanks for taking on this one! Great work 🔥

[patrickvonplaten]

Looks very cool! Think your added tests are great! Went into the PR to help a bit. test_lora_processors now passes and I think you can use its logic to make the other tests as well - see: https://github.com/huggingface/diffusers/pull/2824/files#r1157208349
Let me know if you need any more help :-)
cc @sayakpaul could you also take a look here?

Thank you for the help! I really appreciate it. Would you be able to recommend some resources that I could look into to learn more about this structure and what the transformer_in is doing specifically? I'm amazed by how fast you were able to figure out the PR.

To be honest the Unet3DConditionModel is quite specific and I don't really know the reason why the authors decided to add a transformer in block 🤷

To understand what's going on, I don't have too many tips besides reading the code:

• https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_3d_condition.py
• https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_3d_blocks.py
• https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformer_temporal.py
• https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformer_2d.py

and trying to understand how they are connected. Note that we use the unet_3d_blocks.py as a file to collect many different blocks depending on the architecture while we try to have the other three files as clean as possible.

Very well done in the PR though - we're almost there!

[aandyw]

@patrickvonplaten @sayakpaul PR should be good to merge!

[sayakpaul]

Why are we not mocking weights here?

[sayakpaul]

Same question.

[sayakpaul]

I truly appreciate your hard work! I think this was super important.

I just have a single doubt after which we should be good to merge 🚀

[patrickvonplaten]

Great job @pie31415 !