File size: 6,023 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 | import os
import argparse
from safetensors import safe_open
from safetensors.torch import save_file
def split_illustrious_compatible(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"📦 正在执行切分: {input_path} ...")
with safe_open(input_path, framework="pt", device="cpu") as f:
all_keys = f.keys()
# 判断是否是官方的 transformers 架构
has_text_model_g = any(k.startswith("conditioner.embedders.1.transformer.") for k in all_keys)
for key in all_keys:
tensor = f.get_tensor(key)
# 1. 提取并清洗 CLIP L (官方和 WAI 都是 transformers 结构)
if key.startswith("conditioner.embedders.0.transformer."):
new_key = key.replace("conditioner.embedders.0.transformer.", "", 1)
clipL_tensors[new_key] = tensor
# 2. 提取并精准重组 CLIP G
elif key.startswith("conditioner.embedders.1."):
if has_text_model_g:
# 官方 Illustrious 格式 (transformers)
new_key = key.replace("conditioner.embedders.1.transformer.", "", 1)
if new_key != "logit_scale":
clipG_tensors[new_key] = tensor
else:
# WAI 格式 (OpenCLIP 原生格式)
old_key = key.replace("conditioner.embedders.1.model.", "", 1)
if old_key == "logit_scale":
continue
new_key = old_key
# 基础结构重命名
if "transformer.resblocks." in old_key:
new_key = old_key.replace("transformer.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 old_key == "positional_embedding":
new_key = "text_model.embeddings.position_embedding.weight"
elif old_key == "token_embedding.weight":
new_key = "text_model.embeddings.token_embedding.weight"
elif old_key == "ln_final.weight":
new_key = "text_model.final_layer_norm.weight"
elif old_key == "ln_final.bias":
new_key = "text_model.final_layer_norm.bias"
elif old_key == "text_projection":
new_key = "text_projection"
# 🛠️ 核心修复：解耦 in_proj 复合权重矩阵为标准的 q/k/v 组件
if ".attn.in_proj_weight" in old_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 old_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 old_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 old_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
# 3. 提取 VAE
elif key.startswith("first_stage_model."):
vae_tensors[key.replace("first_stage_model.", "", 1)] = tensor
# 4. 提取 UNet
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)
if diffusion_tensors: save_file(diffusion_tensors, unet_path)
print("\n🎉 切分完成！")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("input_file", type=str)
args = parser.parse_args()
split_illustrious_compatible(args.input_file) |