File size: 6,903 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | 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()
# 动态判断 CLIP-G 是否为原生 OpenCLIP 格式
# 如果有 resblocks 或是 model.transformer.resblocks，说明需要转换
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)
# 1. 提取并清洗 CLIP L (768 维)
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)
# 过滤 position_ids
if "position_ids" in clean_key:
continue
clipL_tensors[clean_key] = tensor
# 2. 提取并清洗 CLIP G (1280 维)
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 (如 Juggernaut) 格式转换 ===
# 剥离所有可能的前缀
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 # 其他键保持原样
# 核心：处理复合权重 (in_proj)
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:
# === 已经是标准 Transformers 格式 ===
if clean_key.startswith("transformer."):
clean_key = clean_key.replace("transformer.", "", 1)
clipG_tensors[clean_key] = tensor
# 3. 提取并清洗 VAE
elif key.startswith("first_stage_model."):
vae_tensors[key.replace("first_stage_model.", "", 1)] = tensor
# 4. 提取并清洗 UNet (Diffusion)
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)})") # 正常应该输出 517
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) |