Upload convert_flux2_klein_bfl_to_diffusers.py

convert_flux2_klein_bfl_to_diffusers.py

@@ -0,0 +1,479 @@
+#!/usr/bin/env python3
+# convert_flux2_klein_bfl_to_diffusers.py
+
+import argparse
+import json
+import re
+import sys
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+
+FLUX2_KLEIN_BASE_4B_CONFIG = {
+    "_class_name": "Flux2Transformer2DModel",
+    "_diffusers_version": "0.37.0.dev0",
+    "attention_head_dim": 128,
+    "axes_dims_rope": [32, 32, 32, 32],
+    "eps": 1e-6,
+    "guidance_embeds": False,
+    "in_channels": 128,
+    "joint_attention_dim": 7680,
+    "mlp_ratio": 3.0,
+    "num_attention_heads": 24,
+    "num_layers": 5,
+    "num_single_layers": 20,
+    "out_channels": None,
+    "patch_size": 1,
+    "rope_theta": 2000,
+    "timestep_guidance_channels": 256,
+}
+
+
+PREFIXES_TO_STRIP = (
+    "model.diffusion_model.",
+    "diffusion_model.",
+    "model.",
+)
+
+
+def normalize_source_key(key: str) -> str:
+    for prefix in PREFIXES_TO_STRIP:
+        if key.startswith(prefix):
+            return key[len(prefix):]
+    return key
+
+
+def make_normalized_key_map(reader) -> dict[str, str]:
+    mapping = {}
+    for real_key in reader.keys():
+        key = normalize_source_key(real_key)
+        if key in mapping:
+            raise ValueError(
+                f"Duplicate normalized key: {key}\n"
+                f"  {mapping[key]}\n"
+                f"  {real_key}"
+            )
+        mapping[key] = real_key
+    return mapping
+
+
+def infer_max_index(keys: set[str], prefix: str) -> int | None:
+    pattern = re.compile(rf"^{re.escape(prefix)}\.(\d+)\.")
+    indices = []
+    for key in keys:
+        m = pattern.match(key)
+        if m:
+            indices.append(int(m.group(1)))
+    return max(indices) if indices else None
+
+
+def add_tensor(dst: dict[str, torch.Tensor], key: str, tensor: torch.Tensor) -> None:
+    if key in dst:
+        raise ValueError(f"Destination key already exists: {key}")
+    dst[key] = tensor.contiguous()
+
+
+def split_qkv_weight(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if w.shape[0] % 3 != 0:
+        raise ValueError(f"QKV tensor first dim is not divisible by 3: {tuple(w.shape)}")
+    q, k, v = torch.chunk(w, 3, dim=0)
+    return q.contiguous(), k.contiguous(), v.contiguous()
+
+
+def swap_final_adaln_scale_shift(t: torch.Tensor) -> torch.Tensor:
+    """
+    BFL FLUX LastLayer:
+        shift, scale = adaLN(...).chunk(2)
+
+    diffusers AdaLayerNormContinuous:
+        scale, shift = linear(...).chunk(2)
+
+    Therefore:
+        source [shift, scale] -> diffusers [scale, shift]
+    """
+    if t.shape[0] % 2 != 0:
+        raise ValueError(f"final AdaLN tensor first dim is not divisible by 2: {tuple(t.shape)}")
+    shift, scale = torch.chunk(t, 2, dim=0)
+    return torch.cat([scale, shift], dim=0).contiguous()
+
+
+def convert(
+    input_file: Path,
+    *,
+    config: dict,
+    include_guidance: bool = False,
+    strict_unused: bool = False,
+) -> tuple[dict[str, torch.Tensor], list[str]]:
+    converted: dict[str, torch.Tensor] = {}
+
+    with safe_open(str(input_file), framework="pt", device="cpu") as f:
+        real_keys = make_normalized_key_map(f)
+        source_keys = set(real_keys.keys())
+        used: set[str] = set()
+
+        def has(key: str) -> bool:
+            return key in real_keys
+
+        def get(key: str, *, required: bool = True) -> torch.Tensor | None:
+            real_key = real_keys.get(key)
+            if real_key is None:
+                if required:
+                    raise KeyError(f"Missing source tensor: {key}")
+                return None
+            used.add(key)
+            return f.get_tensor(real_key)
+
+        def copy_weight(src_base: str, dst_base: str, *, required: bool = True) -> None:
+            w = get(f"{src_base}.weight", required=required)
+            if w is not None:
+                add_tensor(converted, f"{dst_base}.weight", w)
+
+            # klein base 4B の transformer は基本 bias なし。
+            # ただし派生 checkpoint 対応として存在すれば写す。
+            b = get(f"{src_base}.bias", required=False)
+            if b is not None:
+                add_tensor(converted, f"{dst_base}.bias", b)
+
+        def copy_scale_as_weight(src_base: str, dst_base: str) -> None:
+            s = get(f"{src_base}.scale")
+            add_tensor(converted, f"{dst_base}.weight", s)
+
+        def copy_qkv(src_base: str, dst_q: str, dst_k: str, dst_v: str) -> None:
+            w = get(f"{src_base}.weight")
+            q, k, v = split_qkv_weight(w)
+            add_tensor(converted, f"{dst_q}.weight", q)
+            add_tensor(converted, f"{dst_k}.weight", k)
+            add_tensor(converted, f"{dst_v}.weight", v)
+
+            # 念のため bias あり派生にも対応。
+            b = get(f"{src_base}.bias", required=False)
+            if b is not None:
+                qb, kb, vb = split_qkv_weight(b)
+                add_tensor(converted, f"{dst_q}.bias", qb)
+                add_tensor(converted, f"{dst_k}.bias", kb)
+                add_tensor(converted, f"{dst_v}.bias", vb)
+
+        def copy_final_adaln(src_base: str, dst_base: str) -> None:
+            w = get(f"{src_base}.weight")
+            add_tensor(converted, f"{dst_base}.weight", swap_final_adaln_scale_shift(w))
+
+            b = get(f"{src_base}.bias", required=False)
+            if b is not None:
+                add_tensor(converted, f"{dst_base}.bias", swap_final_adaln_scale_shift(b))
+
+        expected_double = int(config["num_layers"])
+        expected_single = int(config["num_single_layers"])
+
+        max_double = infer_max_index(source_keys, "double_blocks")
+        max_single = infer_max_index(source_keys, "single_blocks")
+
+        if max_double is None or max_double + 1 != expected_double:
+            raise ValueError(
+                f"double_blocks count mismatch: source max index={max_double}, "
+                f"config expects {expected_double} blocks"
+            )
+
+        if max_single is None or max_single + 1 != expected_single:
+            raise ValueError(
+                f"single_blocks count mismatch: source max index={max_single}, "
+                f"config expects {expected_single} blocks"
+            )
+
+        # ---------------------------------------------------------------------
+        # Top-level embeddings
+        # ---------------------------------------------------------------------
+        copy_weight("img_in", "x_embedder")
+        copy_weight("txt_in", "context_embedder")
+
+        copy_weight(
+            "time_in.in_layer",
+            "time_guidance_embed.timestep_embedder.linear_1",
+        )
+        copy_weight(
+            "time_in.out_layer",
+            "time_guidance_embed.timestep_embedder.linear_2",
+        )
+
+        # FLUX.2 klein base 4B は guidance_embeds=False。
+        # 蒸留版・派生で guidance_in がある場合だけ使う。
+        if include_guidance:
+            copy_weight(
+                "guidance_in.in_layer",
+                "time_guidance_embed.guidance_embedder.linear_1",
+            )
+            copy_weight(
+                "guidance_in.out_layer",
+                "time_guidance_embed.guidance_embedder.linear_2",
+            )
+
+        # ---------------------------------------------------------------------
+        # Modulation
+        # ---------------------------------------------------------------------
+        copy_weight(
+            "double_stream_modulation_img.lin",
+            "double_stream_modulation_img.linear",
+        )
+        copy_weight(
+            "double_stream_modulation_txt.lin",
+            "double_stream_modulation_txt.linear",
+        )
+        copy_weight(
+            "single_stream_modulation.lin",
+            "single_stream_modulation.linear",
+        )
+
+        # ---------------------------------------------------------------------
+        # Double-stream blocks
+        # ---------------------------------------------------------------------
+        for i in range(expected_double):
+            src = f"double_blocks.{i}"
+            dst = f"transformer_blocks.{i}"
+
+            copy_qkv(
+                f"{src}.img_attn.qkv",
+                f"{dst}.attn.to_q",
+                f"{dst}.attn.to_k",
+                f"{dst}.attn.to_v",
+            )
+
+            copy_qkv(
+                f"{src}.txt_attn.qkv",
+                f"{dst}.attn.add_q_proj",
+                f"{dst}.attn.add_k_proj",
+                f"{dst}.attn.add_v_proj",
+            )
+
+            copy_scale_as_weight(
+                f"{src}.img_attn.norm.query_norm",
+                f"{dst}.attn.norm_q",
+            )
+            copy_scale_as_weight(
+                f"{src}.img_attn.norm.key_norm",
+                f"{dst}.attn.norm_k",
+            )
+            copy_scale_as_weight(
+                f"{src}.txt_attn.norm.query_norm",
+                f"{dst}.attn.norm_added_q",
+            )
+            copy_scale_as_weight(
+                f"{src}.txt_attn.norm.key_norm",
+                f"{dst}.attn.norm_added_k",
+            )
+
+            copy_weight(
+                f"{src}.img_attn.proj",
+                f"{dst}.attn.to_out.0",
+            )
+            copy_weight(
+                f"{src}.txt_attn.proj",
+                f"{dst}.attn.to_add_out",
+            )
+
+            copy_weight(
+                f"{src}.img_mlp.0",
+                f"{dst}.ff.linear_in",
+            )
+            copy_weight(
+                f"{src}.img_mlp.2",
+                f"{dst}.ff.linear_out",
+            )
+            copy_weight(
+                f"{src}.txt_mlp.0",
+                f"{dst}.ff_context.linear_in",
+            )
+            copy_weight(
+                f"{src}.txt_mlp.2",
+                f"{dst}.ff_context.linear_out",
+            )
+
+        # ---------------------------------------------------------------------
+        # Single-stream blocks
+        # ---------------------------------------------------------------------
+        for i in range(expected_single):
+            src = f"single_blocks.{i}"
+            dst = f"single_transformer_blocks.{i}"
+
+            # BFL 側 linear1 は [Q, K, V, MLP-in] の fused projection。
+            # diffusers 側も to_qkv_mlp_proj として同じ fused projection を持つ。
+            copy_weight(
+                f"{src}.linear1",
+                f"{dst}.attn.to_qkv_mlp_proj",
+            )
+
+            # BFL 側 linear2 は [attention-out, MLP-out] の fused output。
+            copy_weight(
+                f"{src}.linear2",
+                f"{dst}.attn.to_out",
+            )
+
+            copy_scale_as_weight(
+                f"{src}.norm.query_norm",
+                f"{dst}.attn.norm_q",
+            )
+            copy_scale_as_weight(
+                f"{src}.norm.key_norm",
+                f"{dst}.attn.norm_k",
+            )
+
+        # ---------------------------------------------------------------------
+        # Final layer
+        # ---------------------------------------------------------------------
+        copy_final_adaln(
+            "final_layer.adaLN_modulation.1",
+            "norm_out.linear",
+        )
+        copy_weight(
+            "final_layer.linear",
+            "proj_out",
+        )
+
+        unused = sorted(source_keys - used)
+
+    meaningful_unused = [
+        k for k in unused
+        if not k.startswith("_")
+        and "optimizer" not in k.lower()
+        and "ema" not in k.lower()
+    ]
+
+    if meaningful_unused:
+        print(f"[warn] unused source tensors: {len(meaningful_unused)}", file=sys.stderr)
+        for k in meaningful_unused[:100]:
+            print(f"  UNUSED {k}", file=sys.stderr)
+        if strict_unused:
+            raise SystemExit(1)
+
+    return converted, meaningful_unused
+
+
+def validate_against_diffusers(converted: dict[str, torch.Tensor], config: dict) -> None:
+    try:
+        from accelerate import init_empty_weights
+        from diffusers import Flux2Transformer2DModel
+    except Exception as e:
+        print(
+            "[warn] diffusers validation skipped. "
+            f"Could not import Flux2Transformer2DModel: {e}",
+            file=sys.stderr,
+        )
+        return
+
+    kwargs = {k: v for k, v in config.items() if not k.startswith("_")}
+
+    with init_empty_weights():
+        model = Flux2Transformer2DModel(**kwargs)
+
+    expected_state = model.state_dict()
+    expected_keys = set(expected_state.keys())
+    got_keys = set(converted.keys())
+
+    missing = sorted(expected_keys - got_keys)
+    unexpected = sorted(got_keys - expected_keys)
+
+    shape_mismatches = []
+    for key in sorted(expected_keys & got_keys):
+        expected_shape = tuple(expected_state[key].shape)
+        got_shape = tuple(converted[key].shape)
+        if expected_shape != got_shape:
+            shape_mismatches.append((key, expected_shape, got_shape))
+
+    if missing or unexpected or shape_mismatches:
+        print("[error] diffusers key validation failed", file=sys.stderr)
+
+        if missing:
+            print(f"[error] missing keys: {len(missing)}", file=sys.stderr)
+            for k in missing[:100]:
+                print(f"  MISSING {k}", file=sys.stderr)
+
+        if unexpected:
+            print(f"[error] unexpected keys: {len(unexpected)}", file=sys.stderr)
+            for k in unexpected[:100]:
+                print(f"  UNEXPECTED {k}", file=sys.stderr)
+
+        if shape_mismatches:
+            print(f"[error] shape mismatches: {len(shape_mismatches)}", file=sys.stderr)
+            for k, exp, got in shape_mismatches[:100]:
+                print(f"  SHAPE {k}: expected={exp}, got={got}", file=sys.stderr)
+
+        raise SystemExit(1)
+
+    print(f"[info] diffusers validation passed: {len(got_keys)} tensors")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Convert FLUX.2 klein base 4B BFL single-file transformer weights to diffusers transformer format."
+    )
+    parser.add_argument(
+        "--input",
+        required=True,
+        type=Path,
+        help="Path to flux-2-klein-base-4b.safetensors",
+    )
+    parser.add_argument(
+        "--output",
+        required=True,
+        type=Path,
+        help="Output diffusers transformer directory",
+    )
+    parser.add_argument(
+        "--include-guidance",
+        action="store_true",
+        help="Use only for variants with guidance_in.* and guidance_embeds=True.",
+    )
+    parser.add_argument(
+        "--skip-validation",
+        action="store_true",
+        help="Skip validation against diffusers Flux2Transformer2DModel state_dict.",
+    )
+    parser.add_argument(
+        "--strict-unused",
+        action="store_true",
+        help="Fail if unused source tensors remain.",
+    )
+    args = parser.parse_args()
+
+    if not args.input.exists():
+        raise FileNotFoundError(args.input)
+
+    config = dict(FLUX2_KLEIN_BASE_4B_CONFIG)
+    config["guidance_embeds"] = bool(args.include_guidance)
+
+    print(f"[info] input : {args.input}")
+    print(f"[info] output: {args.output}")
+
+    converted, _ = convert(
+        args.input,
+        config=config,
+        include_guidance=args.include_guidance,
+        strict_unused=args.strict_unused,
+    )
+
+    if not args.skip_validation:
+        validate_against_diffusers(converted, config)
+
+    args.output.mkdir(parents=True, exist_ok=True)
+
+    config_path = args.output / "config.json"
+    weight_path = args.output / "diffusion_pytorch_model.safetensors"
+
+    with open(config_path, "w", encoding="utf-8") as f:
+        json.dump(config, f, ensure_ascii=False, indent=2)
+        f.write("\n")
+
+    save_file(
+        converted,
+        str(weight_path),
+        metadata={"format": "pt"},
+    )
+
+    print("[done]")
+    print(f"  config : {config_path}")
+    print(f"  weights: {weight_path}")
+
+
+if __name__ == "__main__":
+    main()