Convert From XLab lora to diffusers lora script

[chuck-ma]

Is your feature request related to a problem? Please describe.
I want to use Lora made by XLab. But I can't make it work in Diffusers because there are many differences in Lora structure
Describe the solution you'd like.
Make a script so that we can convert Lora structures easily from XLAB to Diffusers format.

[chuck-ma]

I am facing a problem. The normal shape of lora for xlab, such as double_blocks.{i}.img_attn.qkv.lora_up.weight, is torch.Size([16, 3072]). However, in the format of diffusers, qkv needs to be separated individually. (Referenced from the code of converting flux weights in diffusers)

lora_up_sample_q, lora_up_sample_k, lora_up_sample_v = torch.chunk(
            original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.lora_up.weight"), 3, dim=0
        )
       
                converted_state_dict[f"{block_prefix}attn.to_q.lora_A.weight"] = torch.cat([lora_down_sample_q])
        converted_state_dict[f"{block_prefix}attn.to_k.lora_A.weight"] = torch.cat([lora_down_sample_k])
        converted_state_dict[f"{block_prefix}attn.to_v.lora_A.weight"] = torch.cat([lora_down_sample_v])

This will cause the shape of "f"{block_prefix}attn.to_q.lora_A.weight", f"{block_prefix}attn.to_k.lora_A.weight", f"{block_prefix}attn.to_v.lora_A.weight" to be torch.Size([6, 3072]), torch.Size([6, 3072]), torch.Size([4, 3072]).

While the normal shape of lora should all be torch.Size([4, 3072])

Any advice? I'm trying to convert https://huggingface.co/VideoAditor/Flux-Lora-Realism into diffusers format.

@sayakpaul

[sayakpaul]

Thanks for your efforts!

Maybe they are using separate ranks for separate layers? From the shapes, it looks like QK have a rank of 6 while V has a rank of 4. So, we need to take that into account. @BenjaminBossan is my understanding correct?

[BenjaminBossan]

Yes, it quite looks like they have different ranks depending on the target layer. However, without seeing how this weight is loaded in code, it's not possible to be sure.

This will cause the shape of "f"{block_prefix}attn.to_q.lora_A.weight", f"{block_prefix}attn.to_k.lora_A.weight", f"{block_prefix}attn.to_v.lora_A.weight" to be torch.Size([6, 3072]), torch.Size([6, 3072]), torch.Size([4, 3072]).

While the normal shape of lora should all be torch.Size([4, 3072])

How did you derive this information?

[chuck-ma]

First of all, I use a lora that can be normally loaded by diffusers to figure out its shape

from peft.utils import get_peft_model_state_dict
from accelerate import Accelerator
from diffusers.utils.torch_utils import is_compiled_module



def unwrap_model(model):
    model = accelerator.unwrap_model(model)
    model = model._orig_mod if is_compiled_module(model) else model
    return model

pipe.load_lora_weights("linoyts/yarn_art_Flux_LoRA",weight_name="pytorch_lora_weights.safetensors",
                       adapter_name="default",
                     cross_attention_kwargs={"scale": 0.9})

#transformer = unwrap_model(pipe.transformer)
#transformer = transformer.to(torch.float32)

transformer_lora_layers = get_peft_model_state_dict(pipe.transformer)

Then, I downloaded the lora model of "VideoAditor/Flux-Lora-Realism", looked at its corresponding key, and attempted to convert it directly through torch.chunk

ckpt_path = hf_hub_download("VideoAditor/Flux-Lora-Realism", filename="flux_realism_lora.safetensors")

diffusers_realism_state_dict = safetensors.torch.load_file(ckpt_path)

new_diffusers_realism_state_dict = {key.replace("diffusion_model.", ""): value 
                                    for key, value in diffusers_realism_state_dict.items()}

# Recalculate the keys list without the "diffusion_model." prefix
diffusers_real_lora_keys = [key for key in new_diffusers_realism_state_dict.keys() 
                            if  key.startswith("double_blocks.0.")
                            # and  "lora_down" in key
                            # and  "img_attn" in key
                            # and  "qkv" in key
                           ]

diffusers_real_lora_keys

Then i derive this information @BenjaminBossan

[sayakpaul]

Well, "linoyts/yarn_art_Flux_LoRA" was trained with a rank of 4 that is why you got that shape but that doesn't warranty that all other LoRAs will have a similar rank or rank pattern.

[BenjaminBossan]

Exactly as Sayak says. This information is not sufficient to determine what the rank for each of q/k/v should actually be. If there was some code that actually loads this LoRA successfully, we could check there.

One possibility that I think could also apply is that the LoRA weights are applied to the combined q-k-v tensor and only after that are they chunked. This would be an alternative explanation of why the weight shape is not divisible by 3. Supporting that would probably not be straightforward (unless there is something in diffusers for this already).

[sayakpaul]

One possibility that I think could also apply is that the LoRA weights are applied to the combined q-k-v tensor and only after that are they chunked. This would be an alternative explanation of why the weight shape is not divisible by 3. Supporting that would probably not be straightforward (unless there is something in diffusers for this already).

We have fuse_qkv_projections():

diffusers/src/diffusers/models/transformers/transformer_flux.py

Line 345 in c977966

def fuse_qkv_projections(self):

This fuses the QKV projections. So, a workflow could be to fuse the projections and then add LoRA configs. But this API is quite experimental.

[chuck-ma]

I'm very sorry for not replying in time. Thanks for your advice, I will try if fuse_qkv_projections works.

[dvregan]

I'm not sure if this will help, but the script XLabs uses to train their LoRA is available in XLabs-AI/x-flux repo, along with all the model definitions and code. In the train_flux_lora_deepspeed.py script they initialise things with:

dit, vae, t5, clip = get_models(name=args.model_name, device=accelerator.device, offload=False, is_schnell=is_schnell)
lora_attn_procs = {}

for name, attn_processor in dit.attn_processors.items():
    lora_attn_procs[name] = DoubleStreamBlockLoraProcessor(dim=3072, rank=args.rank)

dit.set_attn_processor(lora_attn_procs)

The DoubleStreamBlockLoraProcessor class is defined here.
See also LoraFluxAttnProcessor class in the same file.

[BenjaminBossan]

Thanks @dvregan it pretty much looks like they apply LoRA on top of the fused weights. Therefore, it doesn't make sense to try to chunk the LoRA weights. fuse_qkv_projections should be the way to go.

I will try if fuse_qkv_projections works.

Let us know what you find it.

[chuck-ma]

#9295

It seems that @sayakpaul has done the job.

[sayakpaul]

Would be great if you could participate in testing and let us know your feedback.