Upload convert_comfy.py

convert_comfy.py

@@ -0,0 +1,206 @@
+import os
+import argparse
+import json
+import torch
+import numpy as np
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.utils
+from safetensors.torch import load_file as load_pt_file
+from mlx_z_image import ZImageTransformerMLX
+from tqdm import tqdm
+
+# From https://huggingface.co/Tongyi-MAI/Z-Image-Turbo/blob/main/transformer/config.json
+config = {
+    "_class_name": "ZImageTransformer2DModel",
+    "_diffusers_version": "0.36.0.dev0",
+    "all_f_patch_size": [1],
+    "all_patch_size": [2],
+    "axes_dims": [32, 48, 48],
+    "axes_lens": [1536, 512, 512],
+    "cap_feat_dim": 2560,
+    "dim": 3840,
+    "in_channels": 16,
+    "n_heads": 30,
+    "n_kv_heads": 30,
+    "n_layers": 30,
+    "n_refiner_layers": 2,
+    "norm_eps": 1e-05,
+    "qk_norm": True,
+    "rope_theta": 256.0,
+    "t_scale": 1000.0,
+    "nheads": 30,
+}
+
+
+# Helpers functions to revert ComfyUI single file model to diffusers format state_dict
+# Some keys are already in the target naming, but i'll revert them nonetheless to use
+# the original code as-is. Undo what is done here:
+# https://huggingface.co/Comfy-Org/z_image_turbo/blob/main/z_image_convert_original_to_comfy.py
+
+
+# remove prefix
+def remove_prefix(key):
+    if key.startswith("model.diffusion_model."):
+        return key.replace("model.diffusion_model.", "")
+
+
+# split qkv layers into q,k and v layers
+def remap_qkv(key, state_dict):
+    weight = state_dict.pop(key)
+    to_q, to_k, to_v = weight.chunk(3, dim=0)
+    state_dict[remove_prefix(key).replace(".qkv.", ".to_q.")] = to_q
+    state_dict[remove_prefix(key).replace(".qkv.", ".to_k.")] = to_k
+    state_dict[remove_prefix(key).replace(".qkv.", ".to_v.")] = to_v
+
+
+replace_keys = {
+    "final_layer.": "all_final_layer.2-1.",
+    "x_embedder.": "all_x_embedder.2-1.",
+    ".attention.out.bias": ".attention.to_out.0.bias",
+    ".attention.k_norm.weight": ".attention.norm_k.weight",
+    ".attention.q_norm.weight": ".attention.norm_q.weight",
+    ".attention.out.weight": ".attention.to_out.0.weight",
+}
+
+
+# restore original name of keys
+def remap_keys(key, state_dict):
+    new_key = remove_prefix(key)
+    for r, rr in replace_keys.items():
+        new_key = new_key.replace(r, rr)
+    state_dict[new_key] = state_dict.pop(key)
+
+
+# Torch to MLX converter function from https://github.com/uqer1244/MLX_z-image/blob/master/converting/convert.py
+def map_key_and_convert(key, tensor):
+    # PyTorch Tensor -> Numpy (Float32)
+    # BF16 변환은 나중에 MLX array 생성 시 수행
+    if isinstance(tensor, torch.Tensor):
+        val = tensor.detach().cpu().float().numpy()
+    else:
+        val = tensor
+
+    new_key = key
+
+    # 키 매핑 로직 (기존과 동일)
+    if "t_embedder.mlp.0" in key:
+        new_key = key.replace("t_embedder.mlp.0", "t_embedder.linear1")
+    elif "t_embedder.mlp.2" in key:
+        new_key = key.replace("t_embedder.mlp.2", "t_embedder.linear2")
+    elif "all_x_embedder.2-1" in key:
+        new_key = key.replace("all_x_embedder.2-1", "x_embedder")
+    elif "cap_embedder.0" in key:
+        new_key = key.replace("cap_embedder.0", "cap_embedder.layers.0")
+    elif "cap_embedder.1" in key:
+        new_key = key.replace("cap_embedder.1", "cap_embedder.layers.1")
+    elif "all_final_layer.2-1" in key:
+        new_key = key.replace("all_final_layer.2-1", "final_layer")
+
+    if "adaLN_modulation.1" in new_key:
+        new_key = new_key.replace("adaLN_modulation.1", "adaLN_modulation.layers.1")
+    elif "attention.to_out.0" in key:
+        new_key = key.replace("attention.to_out.0", "attention.to_out")
+    elif "adaLN_modulation.0" in key and "final" not in key:
+        new_key = key.replace("adaLN_modulation.0", "adaLN_modulation")
+    elif "adaLN_modulation.1" in key and "final" not in key:
+        new_key = key.replace("adaLN_modulation.1", "adaLN_modulation")
+
+    # Changed to a tuple from original code to allow loading without saving to disk
+    return (new_key, mx.array(val).astype(mx.bfloat16))
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert and Quantize ZIT AIO safetensors to MLX model in 4-bit"
+    )
+    parser.add_argument(
+        "--src_model",
+        type=str,
+        default="comfy.safetensors",
+        help="Path to ZIT model in ComfyUI format",
+    )
+    parser.add_argument(
+        "--dst_model",
+        type=str,
+        default="mlx.safetensors",
+        help="Path to save quantized model in mlx format",
+    )
+    parser.add_argument(
+        "--lora_model",
+        type=str,
+        default="",
+        help="Path to an optional LoRA to merge during conversion",
+    )
+    parser.add_argument(
+        "--lora_scale", type=float, default=1.0, help="Scale for the optional LoRA"
+    )
+    parser.add_argument(
+        "--group_size", type=int, default=32, help="Group size for quantization"
+    )
+    args = parser.parse_args()
+
+    print("Starting conversion!")
+
+    print(f"Loading {args.src_model} file...")
+
+    pt_weights = load_pt_file(args.src_model)
+
+    # Remove an unexpected key. TODO: figure out from where it cames.
+    if "model.diffusion_model.norm_final.weight" in pt_weights.keys():
+        del pt_weights["model.diffusion_model.norm_final.weight"]
+
+    print("Reverting ComfyUI format...")
+
+    keys = list(pt_weights.keys())
+
+    for k in tqdm(keys):
+        if ".qkv." in k:
+            remap_qkv(k, pt_weights)
+        else:
+            remap_keys(k, pt_weights)
+
+    if args.lora_model:
+        counter = 0
+        print(f"Merging LoRA {args.lora_model} at scale {args.lora_scale}...")
+
+        lora_weights = load_pt_file(args.lora_model)
+        keys = [k for k in pt_weights.keys() if k.endswith(".weight")]
+
+        for k in tqdm(keys):
+            down_key = f"diffusion_model.{k}".replace(".weight", ".lora_A.weight")
+            up_key = f"diffusion_model.{k}".replace(".weight", ".lora_B.weight")
+            if down_key in lora_weights.keys() and up_key in lora_weights.keys():
+                counter += 1
+                pt_weights[k] = pt_weights[k] + args.lora_scale * (
+                    lora_weights[up_key] @ lora_weights[down_key]
+                )
+
+        print(f"Merged {counter} weight keys")
+
+    print("Converting to MLX...")
+
+    mlx_weights = []
+
+    for k, v in tqdm(pt_weights.items()):
+        mlx_weights.append(map_key_and_convert(k, v))
+
+    print("Loading converted weights...")
+
+    model = ZImageTransformerMLX(config)
+    model.load_weights(mlx_weights)
+
+    print(f"Quantizing (bits=4, group_size={args.group_size})...")
+
+    nn.quantize(model, bits=4, group_size=args.group_size)
+
+    print(f"Saving {args.dst_model} file...")
+
+    quant_weights = dict(mlx.utils.tree_flatten(model.parameters()))
+    mx.save_safetensors(args.dst_model, quant_weights)
+
+    print("Done!")
+
+
+if __name__ == "__main__":
+    main()