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)