Converting Fine-Tuned Flux Safetensors to Diffusers Format

[Artemis1111]

Hello, I am currently fine-tuning the Flux-Canny model and the Flux model.

The resulting safetensors files produced by the tools I use (x-flux, kohya_ss) do not come with a config.json, and the internal weight names or structure seem slightly different from the Diffusers-compatible safetensors.

I would like to convert these fine-tuned safetensors files into a Diffusers-compatible format.

Could you suggest any methods or tools to achieve this?

[raulc0399]

does this script help?
scripts/convert_flux_to_diffusers.py

[Artemis1111]

I used the code you provided, along with references to HuggingFace's XLabs-AI/flux-controlnet-canny-v3, XLabs-AI/flux-controlnet-canny-diffusers repositories, and https://github.com/huggingface/diffusers/issues/9378, to successfully complete the following code!

I finally managed to convert the ControlNet checkpoint safetensors file created with x-flux into a Diffusers-compatible ControlNet checkpoint safetensors file!

convert code

import safetensors.torch
import torch
from contextlib import nullcontext
from accelerate import init_empty_weights
from huggingface_hub import hf_hub_download

from diffusers import AutoencoderKL, FluxTransformer2DModel
from diffusers.loaders.single_file_utils import convert_ldm_vae_checkpoint
from diffusers.utils.import_utils import is_accelerate_available

def convert_flux_transformer_checkpoint_to_diffusers(
    original_state_dict, 
    num_layers,
):
    converted_state_dict = {}

    ## time_text_embed.timestep_embedder <-  time_in
    converted_state_dict["time_text_embed.timestep_embedder.linear_1.weight"] = original_state_dict.pop(
        "time_in.in_layer.weight"
    )
    converted_state_dict["time_text_embed.timestep_embedder.linear_1.bias"] = original_state_dict.pop(
        "time_in.in_layer.bias"
    )
    converted_state_dict["time_text_embed.timestep_embedder.linear_2.weight"] = original_state_dict.pop(
        "time_in.out_layer.weight"
    )
    converted_state_dict["time_text_embed.timestep_embedder.linear_2.bias"] = original_state_dict.pop(
        "time_in.out_layer.bias"
    )

    ## time_text_embed.text_embedder <- vector_in
    converted_state_dict["time_text_embed.text_embedder.linear_1.weight"] = original_state_dict.pop(
        "vector_in.in_layer.weight"
    )
    converted_state_dict["time_text_embed.text_embedder.linear_1.bias"] = original_state_dict.pop(
        "vector_in.in_layer.bias"
    )
    converted_state_dict["time_text_embed.text_embedder.linear_2.weight"] = original_state_dict.pop(
        "vector_in.out_layer.weight"
    )
    converted_state_dict["time_text_embed.text_embedder.linear_2.bias"] = original_state_dict.pop(
        "vector_in.out_layer.bias"
    )

    # guidance
    has_guidance = any("guidance" in k for k in original_state_dict)
    if has_guidance:
        converted_state_dict["time_text_embed.guidance_embedder.linear_1.weight"] = original_state_dict.pop(
            "guidance_in.in_layer.weight"
        )
        converted_state_dict["time_text_embed.guidance_embedder.linear_1.bias"] = original_state_dict.pop(
            "guidance_in.in_layer.bias"
        )
        converted_state_dict["time_text_embed.guidance_embedder.linear_2.weight"] = original_state_dict.pop(
            "guidance_in.out_layer.weight"
        )
        converted_state_dict["time_text_embed.guidance_embedder.linear_2.bias"] = original_state_dict.pop(
            "guidance_in.out_layer.bias"
        )

    # context_embedder
    converted_state_dict["context_embedder.weight"] = original_state_dict.pop("txt_in.weight")
    converted_state_dict["context_embedder.bias"] = original_state_dict.pop("txt_in.bias")

    # x_embedder
    converted_state_dict["x_embedder.weight"] = original_state_dict.pop("img_in.weight")
    converted_state_dict["x_embedder.bias"] = original_state_dict.pop("img_in.bias")

    # input_hint_block weights conversion
    converted_state_dict["input_hint_block.conv_in.bias"] = original_state_dict.pop("input_hint_block.0.bias")
    converted_state_dict["input_hint_block.conv_in.weight"] = original_state_dict.pop("input_hint_block.0.weight")
    converted_state_dict["input_hint_block.conv_out.bias"] = original_state_dict.pop("input_hint_block.14.bias")
    converted_state_dict["input_hint_block.conv_out.weight"] = original_state_dict.pop("input_hint_block.14.weight")

    for i in range(2, 14, 2):  # Iterating from 2 to 12
        block_idx = (i // 2) - 1  # Mapping to blocks.0 ~ blocks.5
        converted_state_dict[f"input_hint_block.blocks.{block_idx}.bias"] = original_state_dict.pop(
            f"input_hint_block.{i}.bias"
        )
        converted_state_dict[f"input_hint_block.blocks.{block_idx}.weight"] = original_state_dict.pop(
            f"input_hint_block.{i}.weight"
        )

    # controlnet blocks
    for key, tensor in original_state_dict.items():
        if key.startswith("controlnet_blocks."):
            converted_state_dict[key] = tensor

    converted_state_dict["controlnet_x_embedder.bias"] = original_state_dict["pos_embed_input.bias"]
    converted_state_dict["controlnet_x_embedder.weight"] = original_state_dict["pos_embed_input.weight"]

    # double transformer blocks
    for i in range(num_layers):
        block_prefix = f"transformer_blocks.{i}."
        # norms
        ## norm1
        converted_state_dict[f"{block_prefix}norm1.linear.weight"] = original_state_dict.pop(
            f"double_blocks.{i}.img_mod.lin.weight"
        )
        converted_state_dict[f"{block_prefix}norm1.linear.bias"] = original_state_dict.pop(
            f"double_blocks.{i}.img_mod.lin.bias"
        )
        ## norm1_context
        converted_state_dict[f"{block_prefix}norm1_context.linear.weight"] = original_state_dict.pop(
            f"double_blocks.{i}.txt_mod.lin.weight"
        )
        converted_state_dict[f"{block_prefix}norm1_context.linear.bias"] = original_state_dict.pop(
            f"double_blocks.{i}.txt_mod.lin.bias"
        )
        # Q, K, V
        sample_q, sample_k, sample_v = torch.chunk(
            original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.weight"), 3, dim=0
        )
        context_q, context_k, context_v = torch.chunk(
            original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.weight"), 3, dim=0
        )
        sample_q_bias, sample_k_bias, sample_v_bias = torch.chunk(
            original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.bias"), 3, dim=0
        )
        context_q_bias, context_k_bias, context_v_bias = torch.chunk(
            original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.bias"), 3, dim=0
        )
        converted_state_dict[f"{block_prefix}attn.to_q.weight"] = torch.cat([sample_q])
        converted_state_dict[f"{block_prefix}attn.to_q.bias"] = torch.cat([sample_q_bias])
        converted_state_dict[f"{block_prefix}attn.to_k.weight"] = torch.cat([sample_k])
        converted_state_dict[f"{block_prefix}attn.to_k.bias"] = torch.cat([sample_k_bias])
        converted_state_dict[f"{block_prefix}attn.to_v.weight"] = torch.cat([sample_v])
        converted_state_dict[f"{block_prefix}attn.to_v.bias"] = torch.cat([sample_v_bias])
        converted_state_dict[f"{block_prefix}attn.add_q_proj.weight"] = torch.cat([context_q])
        converted_state_dict[f"{block_prefix}attn.add_q_proj.bias"] = torch.cat([context_q_bias])
        converted_state_dict[f"{block_prefix}attn.add_k_proj.weight"] = torch.cat([context_k])
        converted_state_dict[f"{block_prefix}attn.add_k_proj.bias"] = torch.cat([context_k_bias])
        converted_state_dict[f"{block_prefix}attn.add_v_proj.weight"] = torch.cat([context_v])
        converted_state_dict[f"{block_prefix}attn.add_v_proj.bias"] = torch.cat([context_v_bias])
        # qk_norm
        converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = original_state_dict.pop(
            f"double_blocks.{i}.img_attn.norm.query_norm.scale"
        )
        converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = original_state_dict.pop(
            f"double_blocks.{i}.img_attn.norm.key_norm.scale"
        )
        converted_state_dict[f"{block_prefix}attn.norm_added_q.weight"] = original_state_dict.pop(
            f"double_blocks.{i}.txt_attn.norm.query_norm.scale"
        )
        converted_state_dict[f"{block_prefix}attn.norm_added_k.weight"] = original_state_dict.pop(
            f"double_blocks.{i}.txt_attn.norm.key_norm.scale"
        )
        # ff img_mlp
        converted_state_dict[f"{block_prefix}ff.net.0.proj.weight"] = original_state_dict.pop(
            f"double_blocks.{i}.img_mlp.0.weight"
        )
        converted_state_dict[f"{block_prefix}ff.net.0.proj.bias"] = original_state_dict.pop(
            f"double_blocks.{i}.img_mlp.0.bias"
        )
        converted_state_dict[f"{block_prefix}ff.net.2.weight"] = original_state_dict.pop(
            f"double_blocks.{i}.img_mlp.2.weight"
        )
        converted_state_dict[f"{block_prefix}ff.net.2.bias"] = original_state_dict.pop(
            f"double_blocks.{i}.img_mlp.2.bias"
        )
        converted_state_dict[f"{block_prefix}ff_context.net.0.proj.weight"] = original_state_dict.pop(
            f"double_blocks.{i}.txt_mlp.0.weight"
        )
        converted_state_dict[f"{block_prefix}ff_context.net.0.proj.bias"] = original_state_dict.pop(
            f"double_blocks.{i}.txt_mlp.0.bias"
        )
        converted_state_dict[f"{block_prefix}ff_context.net.2.weight"] = original_state_dict.pop(
            f"double_blocks.{i}.txt_mlp.2.weight"
        )
        converted_state_dict[f"{block_prefix}ff_context.net.2.bias"] = original_state_dict.pop(
            f"double_blocks.{i}.txt_mlp.2.bias"
        )
        # output projections.
        converted_state_dict[f"{block_prefix}attn.to_out.0.weight"] = original_state_dict.pop(
            f"double_blocks.{i}.img_attn.proj.weight"
        )
        converted_state_dict[f"{block_prefix}attn.to_out.0.bias"] = original_state_dict.pop(
            f"double_blocks.{i}.img_attn.proj.bias"
        )
        converted_state_dict[f"{block_prefix}attn.to_add_out.weight"] = original_state_dict.pop(
            f"double_blocks.{i}.txt_attn.proj.weight"
        )
        converted_state_dict[f"{block_prefix}attn.to_add_out.bias"] = original_state_dict.pop(
            f"double_blocks.{i}.txt_attn.proj.bias"
        )

    return converted_state_dict

def manipulate_and_convert_flux_weights(input_safetensors_path, output_safetensors_path, num_layers=2):
    """
    Converts and manipulates safetensors weights and saves them.

    Args:
        input_safetensors_path (str): Path to the original safetensors file.
        output_safetensors_path (str): Path to save the modified weights as a new safetensors file.
        num_layers (int): Number of transformer blocks to convert.
    """
    # Step 1: Load the original weights
    print("Loading original safetensors file...")
    original_state_dict = safetensors.torch.load_file(input_safetensors_path)
    
    # Step 2: Convert weights using convert_flux_transformer_checkpoint_to_diffusers
    print("Converting weights to diffusers format...")
    converted_state_dict = convert_flux_transformer_checkpoint_to_diffusers(original_state_dict, num_layers)

    # Step 3: Add metadata
    metadata = {"format": "pt"}  # Adding metadata
    print("Adding metadata to safetensors...")
    
    # Step 4: Save the weights to a new safetensors file
    print("Saving manipulated weights to new safetensors file...")
    safetensors.torch.save_file(converted_state_dict, output_safetensors_path, metadata=metadata)
    print(f"Modified weights saved to: {output_safetensors_path}")

if __name__ == "__main__":
    # path setting
    input_file = "input_file_path"   
    output_file = "output_file_path"  
    
    # run func
    manipulate_and_convert_flux_weights(input_file, output_file, num_layers=2)