| import os
|
| import argparse
|
| from safetensors import safe_open
|
| from safetensors.torch import save_file
|
|
|
| def split_sdxl(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"📦 正在执行 SDXL 纯净切分与权重重组: {input_path} ...")
|
|
|
| with safe_open(input_path, framework="pt", device="cpu") as f:
|
| all_keys = f.keys()
|
|
|
|
|
|
|
| is_openclip_g = any("resblocks." in k for k in all_keys if k.startswith("conditioner.embedders.1."))
|
|
|
| for key in all_keys:
|
| tensor = f.get_tensor(key)
|
|
|
|
|
| if key.startswith("conditioner.embedders.0."):
|
| clean_key = key.replace("conditioner.embedders.0.", "", 1)
|
| if clean_key.startswith("transformer."):
|
| clean_key = clean_key.replace("transformer.", "", 1)
|
|
|
|
|
| if "position_ids" in clean_key:
|
| continue
|
|
|
| clipL_tensors[clean_key] = tensor
|
|
|
|
|
| elif key.startswith("conditioner.embedders.1."):
|
| clean_key = key.replace("conditioner.embedders.1.", "", 1)
|
|
|
|
|
| if clean_key.endswith("logit_scale") or "position_ids" in clean_key:
|
| continue
|
|
|
| if is_openclip_g:
|
|
|
|
|
| openclip_key = clean_key.replace("model.", "", 1).replace("transformer.", "", 1)
|
|
|
| if "resblocks." in openclip_key:
|
| new_key = openclip_key.replace("resblocks.", "text_model.encoder.layers.", 1)
|
| new_key = new_key.replace(".ln_1.", ".layer_norm1.", 1)
|
| new_key = new_key.replace(".ln_2.", ".layer_norm2.", 1)
|
| new_key = new_key.replace(".mlp.c_fc.", ".mlp.fc1.", 1)
|
| new_key = new_key.replace(".mlp.c_proj.", ".mlp.fc2.", 1)
|
| elif openclip_key == "positional_embedding":
|
| new_key = "text_model.embeddings.position_embedding.weight"
|
| elif openclip_key == "token_embedding.weight":
|
| new_key = "text_model.embeddings.token_embedding.weight"
|
| elif openclip_key == "ln_final.weight":
|
| new_key = "text_model.final_layer_norm.weight"
|
| elif openclip_key == "ln_final.bias":
|
| new_key = "text_model.final_layer_norm.bias"
|
| elif openclip_key == "text_projection":
|
| new_key = "text_projection.weight"
|
| else:
|
| new_key = openclip_key
|
|
|
|
|
| if ".attn.in_proj_weight" in openclip_key:
|
| base_layer = new_key.replace(".attn.in_proj_weight", ".self_attn", 1)
|
| q, k, v = tensor.chunk(3, dim=0)
|
| clipG_tensors[f"{base_layer}.q_proj.weight"] = q
|
| clipG_tensors[f"{base_layer}.k_proj.weight"] = k
|
| clipG_tensors[f"{base_layer}.v_proj.weight"] = v
|
| elif ".attn.in_proj_bias" in openclip_key:
|
| base_layer = new_key.replace(".attn.in_proj_bias", ".self_attn", 1)
|
| q, k, v = tensor.chunk(3, dim=0)
|
| clipG_tensors[f"{base_layer}.q_proj.bias"] = q
|
| clipG_tensors[f"{base_layer}.k_proj.bias"] = k
|
| clipG_tensors[f"{base_layer}.v_proj.bias"] = v
|
| elif ".attn.out_proj.weight" in openclip_key:
|
| new_key = new_key.replace(".attn.out_proj.weight", ".self_attn.out_proj.weight", 1)
|
| clipG_tensors[new_key] = tensor
|
| elif ".attn.out_proj.bias" in openclip_key:
|
| new_key = new_key.replace(".attn.out_proj.bias", ".self_attn.out_proj.bias", 1)
|
| clipG_tensors[new_key] = tensor
|
| else:
|
| clipG_tensors[new_key] = tensor
|
| else:
|
|
|
| if clean_key.startswith("transformer."):
|
| clean_key = clean_key.replace("transformer.", "", 1)
|
| clipG_tensors[clean_key] = tensor
|
|
|
|
|
| elif key.startswith("first_stage_model."):
|
| vae_tensors[key.replace("first_stage_model.", "", 1)] = tensor
|
|
|
|
|
| elif key.startswith("model.diffusion_model."):
|
| diffusion_tensors[key.replace("model.diffusion_model.", "", 1)] = tensor
|
|
|
| print("\n💾 正在保存组件...")
|
|
|
| if clipL_tensors:
|
| save_file(clipL_tensors, clipL_path)
|
| print(f" --> CLIP-L 保存成功 (键数: {len(clipL_tensors)})")
|
| if clipG_tensors:
|
| save_file(clipG_tensors, clipG_path)
|
| print(f" --> CLIP-G 保存成功 (键数: {len(clipG_tensors)})")
|
| if vae_tensors:
|
| save_file(vae_tensors, vae_path)
|
| print(f" --> VAE 保存成功 (键数: {len(vae_tensors)})")
|
| if diffusion_tensors:
|
| save_file(diffusion_tensors, unet_path)
|
| print(f" --> UNet 保存成功 (键数: {len(diffusion_tensors)})")
|
|
|
| print("\n🎉 SDXL 切分任务完成！")
|
|
|
| if __name__ == "__main__":
|
| parser = argparse.ArgumentParser(description="纯净 SDXL 模型切分与格式转换工具")
|
| parser.add_argument("input_file", type=str, help="输入的 SDXL 大模型文件路径")
|
| args = parser.parse_args()
|
| split_sdxl(args.input_file) |
