tmp/video-velocity-model/data_processing/flux2_ae/convert_diffusers_checkpoint.py

from __future__ import annotations

import argparse
import re
from pathlib import Path

from safetensors.torch import load_file, save_file


ATTN_WEIGHT_SUFFIXES = (
    ".attn_1.q.weight",
    ".attn_1.k.weight",
    ".attn_1.v.weight",
    ".attn_1.proj_out.weight",
)


def convert_key(key: str) -> str:
    if key == "quant_conv.weight":
        return "encoder.quant_conv.weight"
    if key == "quant_conv.bias":
        return "encoder.quant_conv.bias"
    if key == "post_quant_conv.weight":
        return "decoder.post_quant_conv.weight"
    if key == "post_quant_conv.bias":
        return "decoder.post_quant_conv.bias"

    if key.startswith("encoder.mid_block.resnets.0."):
        return key.replace("encoder.mid_block.resnets.0.", "encoder.mid.block_1.")
    if key.startswith("encoder.mid_block.resnets.1."):
        return key.replace("encoder.mid_block.resnets.1.", "encoder.mid.block_2.")
    if key.startswith("encoder.mid_block.attentions.0."):
        key = key.replace("encoder.mid_block.attentions.0.group_norm.", "encoder.mid.attn_1.norm.")
        key = key.replace("encoder.mid_block.attentions.0.to_q.", "encoder.mid.attn_1.q.")
        key = key.replace("encoder.mid_block.attentions.0.to_k.", "encoder.mid.attn_1.k.")
        key = key.replace("encoder.mid_block.attentions.0.to_v.", "encoder.mid.attn_1.v.")
        return key.replace("encoder.mid_block.attentions.0.to_out.0.", "encoder.mid.attn_1.proj_out.")
    if key.startswith("encoder."):
        key = key.replace("encoder.down_blocks.", "encoder.down.")
        key = key.replace(".resnets.", ".block.")
        key = key.replace(".downsamplers.0.", ".downsample.")
        key = key.replace(".conv_shortcut.", ".nin_shortcut.")
        return key.replace("encoder.conv_norm_out.", "encoder.norm_out.")

    if key.startswith("decoder.mid_block.resnets.0."):
        return key.replace("decoder.mid_block.resnets.0.", "decoder.mid.block_1.")
    if key.startswith("decoder.mid_block.resnets.1."):
        return key.replace("decoder.mid_block.resnets.1.", "decoder.mid.block_2.")
    if key.startswith("decoder.mid_block.attentions.0."):
        key = key.replace("decoder.mid_block.attentions.0.group_norm.", "decoder.mid.attn_1.norm.")
        key = key.replace("decoder.mid_block.attentions.0.to_q.", "decoder.mid.attn_1.q.")
        key = key.replace("decoder.mid_block.attentions.0.to_k.", "decoder.mid.attn_1.k.")
        key = key.replace("decoder.mid_block.attentions.0.to_v.", "decoder.mid.attn_1.v.")
        return key.replace("decoder.mid_block.attentions.0.to_out.0.", "decoder.mid.attn_1.proj_out.")
    if key.startswith("decoder.up_blocks."):
        match = re.match(r"decoder\.up_blocks\.(\d+)\.(.*)", key)
        if match is None:
            return key
        # Diffusers stores decoder up blocks low->high, while the local FLUX.2 module indexes high->low.
        key = f"decoder.up.{3 - int(match.group(1))}.{match.group(2)}"
        key = key.replace(".resnets.", ".block.")
        key = key.replace(".upsamplers.0.", ".upsample.")
        return key.replace(".conv_shortcut.", ".nin_shortcut.")
    if key.startswith("decoder."):
        return key.replace("decoder.conv_norm_out.", "decoder.norm_out.")
    return key


def convert_checkpoint(input_path: Path, output_path: Path) -> None:
    source = load_file(str(input_path), device="cpu")
    converted = {}
    for key, value in source.items():
        new_key = convert_key(key)
        if any(suffix in new_key for suffix in ATTN_WEIGHT_SUFFIXES) and value.ndim == 2:
            value = value[:, :, None, None]
        converted[new_key] = value
    output_path.parent.mkdir(parents=True, exist_ok=True)
    save_file(converted, str(output_path))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Convert a diffusers FLUX.2 VAE checkpoint for the local FLUX.2 AE module.")
    parser.add_argument("input", type=Path, help="Input diffusers diffusion_pytorch_model.safetensors path.")
    parser.add_argument("output", type=Path, help="Output converted ae.safetensors path.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    convert_checkpoint(args.input.expanduser(), args.output.expanduser())
    print(f"Saved converted FLUX.2 AE checkpoint to {args.output.expanduser()}")


if __name__ == "__main__":
    main()