Fixes the issue with malformed metadata and incorrect key names. This was made to make certain Flux2-Klein models importable into Draw Things.app, but it should also address the issue with LoRAs not fusing anything in Draw Things and MFLUX. It takes a single argument: the input file, and outputs a new file with `-FIXED` appended to the filename …

fix-safetensors-header.py

#!/usr/bin/env python3
"""
- Fixes the issue with malformed metadata and incorrect key names.
This was made to make certain models importable into Draw Things.app,
  but it should also address the issue with LoRAs not fusing anything in Draw Things and MFLUX.

It takes a single argument: the input file (.safetensors) 
and outputs a new file with `-FIXED` appended to the basename (not file extension)

This shouldnt require any package installs, and does it's magic without loading the full bulk of the file into RAM.
"""
import sys
import struct
import json
from pathlib import Path

if len(sys.argv) != 2:
    print("Usage: fix-safetensors-header.py <input.safetensors>")
    sys.exit(1)

prefix = "model.diffusion_model."

src = Path(sys.argv[1])

if src.suffix != ".safetensors":
    print("Input must be a .safetensors file")
    sys.exit(1)

dst = src.with_name(src.stem + "-FIXED.safetensors")

with src.open("rb") as fin:
    header_len = struct.unpack("<Q", fin.read(8))[0]
    header = json.loads(fin.read(header_len).decode("utf-8"))

    # remove metadata block
    header.pop("__metadata__", None)

    # rename keys
    new_header = {}
    for k, v in header.items():
        if k.startswith(prefix):
            new_header[k[len(prefix):]] = v
        else:
            new_header[k] = v

    new_header_bytes = json.dumps(
        new_header,
        separators=(",", ":"),
        ensure_ascii=False
    ).encode("utf-8")

    with dst.open("wb") as fout:
        fout.write(struct.pack("<Q", len(new_header_bytes)))
        fout.write(new_header_bytes)

        # stream payload
        while True:
            chunk = fin.read(16 * 1024 * 1024)
            if not chunk:
                break
            fout.write(chunk)

print(f"Saved to: {dst}")