| 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)
|
|
|
|
|
| if key.startswith("conditioner.embedders.0.transformer."):
|
| new_key = key.replace("conditioner.embedders.0.transformer.", "", 1)
|
| clipL_tensors[new_key] = tensor
|
|
|
|
|
| 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
|
|
|
|
|
| elif key.startswith("first_stage_model."):
|
| new_key = key.replace("first_stage_model.", "", 1)
|
| vae_tensors[new_key] = tensor
|
|
|
|
|
| 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) |
