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)
• there is something wrong with converting Flux2.Klein 9B checkpoints this way. the result has major visual artifacts. - fixed

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)

[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

[dxqb]

merging this because it is required to create Flux2 checkpoints
requires more work to be a generic conversion util