"""Convert a MOSS-Image diffusers ``transformer/`` folder into one ComfyUI ``diffusion_models`` file
(the Comfy-Org/MOSS-Image layout).

The only difference between the two layouts: each block's ``proj.q`` / ``proj.k`` / ``proj.v`` weights are fused
into ``proj.qkv`` (q, k, v stacked on the output dim). Dtypes are kept (bf16, ``rope.bands`` fp32). Applied to
OpenMOSS-Team/MOSS-Image-SFT's ``transformer/`` it reproduces Comfy-Org's ``moss_image_sft_bf16.safetensors``
byte for byte. LoRAs trained on the diffusers layout (separate q/k/v) load unchanged in ComfyUI, which maps them
onto the fused slices.

pip install torch safetensors
python convert_to_comfy.py path/to/transformer kromoss_bf16.safetensors
python convert_to_comfy.py path/to/transformer out.safetensors --compare moss_image_sft_bf16.safetensors
"""
import argparse
import json
from pathlib import Path

import torch
from safetensors import safe_open
from safetensors.torch import save_file


def convert(folder):
    weight_map = json.loads((folder / "diffusion_pytorch_model.safetensors.index.json").read_text())["weight_map"]
    state = {}
    for shard in sorted(set(weight_map.values())):
        with safe_open(str(folder / shard), "pt") as f:
            state.update({k: f.get_tensor(k) for k in f.keys()})
    for q in [k for k in state if k.endswith(".proj.q.weight")]:
        prefix = q[:-len("q.weight")]
        state[f"{prefix}qkv.weight"] = torch.cat([state.pop(f"{prefix}{p}.weight") for p in "qkv"])
    return state


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("transformer", help="diffusers transformer/ folder (sharded safetensors + index)")
    parser.add_argument("output")
    parser.add_argument("--compare", help="a Comfy-Org file: require identical keys, shapes and dtypes; report "
                                          "how many tensors are bit-identical")
    args = parser.parse_args()
    state = convert(Path(args.transformer))
    if args.compare:
        with safe_open(args.compare, "pt") as f:
            if set(f.keys()) != set(state):
                raise ValueError(f"key sets differ: {sorted(set(f.keys()) ^ set(state))[:5]}")
            equal = 0
            for k in f.keys():
                ref = f.get_tensor(k)
                if ref.shape != state[k].shape or ref.dtype != state[k].dtype:
                    raise ValueError(f"{k}: {tuple(ref.shape)} {ref.dtype} vs {tuple(state[k].shape)} {state[k].dtype}")
                equal += torch.equal(ref, state[k])
        print(f"layout matches {args.compare}; {equal}/{len(state)} tensors bit-identical")
    save_file(state, args.output)
    print(f"{args.output}: {len(state)} tensors")


if __name__ == "__main__":
    main()