mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-03-16 22:58:19 +08:00
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:
parent
f5bf7ed4d3
commit
98440de8d4
46
merge_zimage_models.py
Normal file
46
merge_zimage_models.py
Normal 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/")
|
||||
Loading…
Reference in New Issue
Block a user