LoRA and checkpoint conversion util

[dxqb]

Draft to simplify and extend LoRA and checkpoint conversion

• can convert LoRAs and checkpoints using the same conversion pattern definition
• therefore replaces this code https://github.com/Nerogar/OneTrainer/blob/master/modules/util/convert/convert_flux_diffusers_to_ckpt.py and this code https://github.com/Nerogar/OneTrainer/blob/master/modules/util/convert/lora/convert_flux_lora.py for a model
• the same pattern can be used to convert in both directions (for example, from diffusers to original keys and back)

Example use (Flux2 pattern definition):

def diffusers_to_original(qkv_fusion):
    return [
        ("context_embedder", "txt_in"),
        ("x_embedder",       "img_in"),
        ("time_guidance_embed.timestep_embedder", "time_in", [
            ("linear_1", "in_layer"),
            ("linear_2", "out_layer"),
        ]),
        ("time_guidance_embed.guidance_embedder", "guidance_in", [
            ("linear_1", "in_layer"),
            ("linear_2", "out_layer"),
        ]),
        ("double_stream_modulation_img.linear", "double_stream_modulation_img.lin"),
        ("double_stream_modulation_txt.linear", "double_stream_modulation_txt.lin"),
        ("single_stream_modulation.linear",     "single_stream_modulation.lin"),
        ("proj_out",                            "final_layer.linear"),
        ("norm_out.linear", "final_layer.adaLN_modulation.1", swap_chunks, swap_chunks),
        ("transformer_blocks.{i}", "double_blocks.{i}",
            qkv_fusion("attn.to_q", "attn.to_k", "attn.to_v", "img_attn.qkv") + \
            qkv_fusion("attn.add_q_proj", "attn.add_k_proj", "attn.add_v_proj", "txt_attn.qkv") + [
            ("attn.norm_k.weight",       "img_attn.norm.key_norm.scale"),
            ("attn.norm_q.weight",       "img_attn.norm.query_norm.scale"),
            ("attn.to_out.0",            "img_attn.proj"),
            ("ff.linear_in",             "img_mlp.0"),
            ("ff.linear_out",            "img_mlp.2"),
            ("attn.norm_added_k.weight", "txt_attn.norm.key_norm.scale"),
            ("attn.norm_added_q.weight", "txt_attn.norm.query_norm.scale"),
            ("attn.to_add_out",          "txt_attn.proj"),
            ("ff_context.linear_in",     "txt_mlp.0"),
            ("ff_context.linear_out",    "txt_mlp.2"),
        ]),
        ("single_transformer_blocks.{i}", "single_blocks.{i}", [
            ("attn.to_qkv_mlp_proj", "linear1"),
            ("attn.to_out",          "linear2"),
            ("attn.norm_k.weight",   "norm.key_norm.scale"),
            ("attn.norm_q.weight",   "norm.query_norm.scale"),
        ]),
    ]

diffusers_lora_to_original = diffusers_to_original(lora_qkv_fusion)
diffusers_checkpoint_to_original = diffusers_to_original(qkv_fusion)
diffusers_lora_to_comfy = [remove_prefix("transformer"), diffusers_to_original(lora_qkv_fusion), add_prefix("diffusion_model")]


save_state_dict = convert(save_state_dict, diffusers_lora_to_comfy)

[yamatazen]

Can it merge LoRA to checkpoint?

[dxqb]

Can it merge LoRA to checkpoint?

no, not its purpose, and you wouldn't do that on the .safetensors key level.
There are Comfy workflows that can do that. the output can then be loaded into OneTrainer