File size: 2,808 Bytes
6c483fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | """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()
|