diff --git a/merge_zimage_models.py b/merge_zimage_models.py new file mode 100644 index 000000000..bd839230d --- /dev/null +++ b/merge_zimage_models.py @@ -0,0 +1,46 @@ +"""Merge Z-Image Turbo sharded safetensors into single files for ComfyUI.""" +import json +import os +from pathlib import Path +from safetensors.torch import load_file, save_file + +SNAP = Path.home() / ".cache/huggingface/hub/models--Tongyi-MAI--Z-Image-Turbo/snapshots/f332072aa78be7aecdf3ee76d5c247082da564a6" +OUT = Path.home() / "ComfyUI/models" + +def merge_shards(index_path, output_path, label): + print(f"\n=== Merging {label} ===") + with open(index_path) as f: + index = json.load(f) + weight_map = index["weight_map"] + shard_dir = Path(index_path).parent + shards = sorted(set(weight_map.values())) + + merged = {} + for shard in shards: + print(f" Loading {shard} ...") + tensors = load_file(str(shard_dir / shard)) + merged.update(tensors) + print(f" -> {len(tensors)} tensors, total so far: {len(merged)}") + + print(f" Saving to {output_path} ...") + os.makedirs(Path(output_path).parent, exist_ok=True) + save_file(merged, str(output_path)) + size = Path(output_path).stat().st_size / 1e9 + print(f" Done! {size:.1f} GB") + +# 1. Merge transformer +merge_shards( + SNAP / "transformer/diffusion_pytorch_model.safetensors.index.json", + OUT / "diffusion_models/z_image_turbo.safetensors", + "Transformer" +) + +# 2. Merge text encoder (Qwen3) +merge_shards( + SNAP / "text_encoder/model.safetensors.index.json", + OUT / "text_encoders/qwen3_z_image.safetensors", + "Text Encoder (Qwen3)" +) + +print("\n=== All done! ===") +print("Models ready in ComfyUI/models/")