File size: 3,464 Bytes
9f09426 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | import os
import argparse
from safetensors import safe_open
from safetensors.torch import save_file
def split_checkpoint_exact_match(input_path):
if not os.path.exists(input_path):
print(f"错误: 找不到输入文件 '{input_path}'")
return
base_dir = os.path.dirname(input_path)
base_name = os.path.splitext(os.path.basename(input_path))[0]
clipL_path = os.path.join(base_dir, f"{base_name}_clipL.safetensors")
clipG_path = os.path.join(base_dir, f"{base_name}_clipG.safetensors")
vae_path = os.path.join(base_dir, f"{base_name}_VAE.safetensors")
unet_path = os.path.join(base_dir, f"{base_name}_Diffusion.safetensors")
clipL_tensors = {}
clipG_tensors = {}
vae_tensors = {}
diffusion_tensors = {}
print(f"正在进行严格的 Key 名对齐切分: {input_path} ...")
with safe_open(input_path, framework="pt", device="cpu") as f:
for key in f.keys():
tensor = f.get_tensor(key)
# 1. CLIP L:彻底剥离包装前缀,使其直接以 text_model. 开头,精准对齐报错缺失的 key
if key.startswith("conditioner.embedders.0.transformer."):
new_key = key.replace("conditioner.embedders.0.transformer.", "", 1)
clipL_tensors[new_key] = tensor
# 2. CLIP G:同步彻底剥离包装前缀,直接以 openclip 的根键名开头
elif key.startswith("conditioner.embedders.1.transformer."):
new_key = key.replace("conditioner.embedders.1.transformer.", "", 1)
clipG_tensors[new_key] = tensor
elif key.startswith("conditioner.embedders.1.model."):
new_key = key.replace("conditioner.embedders.1.model.", "", 1)
clipG_tensors[new_key] = tensor
# 3. VAE
elif key.startswith("first_stage_model."):
new_key = key.replace("first_stage_model.", "", 1)
vae_tensors[new_key] = tensor
# 4. UNet / Diffusion
elif key.startswith("model.diffusion_model."):
new_key = key.replace("model.diffusion_model.", "", 1)
diffusion_tensors[new_key] = tensor
print("\n正在保存文件...")
if clipL_tensors:
save_file(clipL_tensors, clipL_path)
print(f"✅ 成功提取纯净键名 CLIP-L -> {clipL_path}")
else:
print("❌ 错误: 未能在原文件中捕获到 CLIP-L 权重。")
if clipG_tensors:
save_file(clipG_tensors, clipG_path)
print(f"✅ 成功提取纯净键名 CLIP-G -> {clipG_path}")
else:
print("❌ 错误: 未能在原文件中捕获到 CLIP-G 权重。")
if vae_tensors:
save_file(vae_tensors, vae_path)
print(f"✅ 成功提取 VAE -> {vae_path}")
if diffusion_tensors:
save_file(diffusion_tensors, unet_path)
print(f"✅ 成功提取 Diffusion -> {unet_path}")
print("\n精准切分全部完成!")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="严格对齐 ComfyUI 报错键名的双独立 CLIP 切分脚本")
parser.add_argument("input_file", type=str, help="输入的 .safetensors 权重文件路径")
args = parser.parse_args()
split_checkpoint_exact_match(args.input_file) |