From 98440de8d4d122339fe17d869822ada68c89e618 Mon Sep 17 00:00:00 2001
From: bigjimmy <dajimmy@gmail.com>
Date: Sun, 1 Mar 2026 22:57:10 +0800
Subject: [PATCH] add: script to merge Z-Image Turbo sharded safetensors for
 ComfyUI
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Utility script that merges the HuggingFace sharded safetensors from
Tongyi-MAI/Z-Image-Turbo into single files ready for ComfyUI:
- transformer shards → models/diffusion_models/z_image_turbo.safetensors
- text encoder shards → models/text_encoders/qwen3_z_image.safetensors

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 merge_zimage_models.py | 46 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)
 create mode 100644 merge_zimage_models.py

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/")