Load LoRA with uneven Rank

[mutatedducks97]

Summary

Internally invoke makes use of CustomLinear module to inject LoRA into but does not support fused weights from uneven rank from up and down matrices. In this PR we create a fused matrix to address this issue. Since huggingface diffusers implements FLUX using a fused kqv The rank of the LoRA generated from a diffuser FLUX model will be different.

Before this PR the following issues would occur. This resolves it.

1. For QKV attention layers (img_attn and txt_atn):
   • CustomLinear(in_features=3072, out_features=9216)
   • LoRA matrices: down=[12, 3072], up=[9216, 4]
   • The error shows: cannot multiply (9216x4) with (12x3072)
   • We need it to output a matrix with shape (9216, 3072)

• New shape:torch.Size([9216, 3072])

1. For linear1 layers:
   • CustomLinear(in_features=3072, out_features=21504)
   • LoRA matrices: down=[16, 3072], up=[21504, 4]
   • The error shows: cannot multiply (21504x4) with (16x3072)
   • We need it to output a matrix with shape (21504, 3072)

Default Flux

50 steps training

350 steps training Flux

Related Issues / Discussions

QA Instructions

Merge Plan

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)