add: script to merge Z-Image Turbo sharded safetensors for ComfyUI

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>
This commit is contained in:
bigjimmy 2026-03-01 22:57:10 +08:00
parent f5bf7ed4d3
commit 98440de8d4

46
merge_zimage_models.py Normal file
View File

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