mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-03-17 07:05:12 +08:00
98440de8d4
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>
47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
"""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/")
|