File size: 4,888 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 | import os
import sys
import torch
from safetensors import safe_open
from safetensors.torch import save_file
def split_sd15_ultimate(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]
clip_path = os.path.join(base_dir, f"{base_name}_clip.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")
clip_tensors = {}
vae_tensors = {}
# 这个字典既包含 UNet，也包含所有额外组件（EMA、ControlNet 等）
unet_and_extras = {}
print(f"📦 正在执行【终极暴力提取】SD 1.5 切分: {input_path} ...")
with safe_open(input_path, framework="pt", device="cpu") as f:
for key in f.keys():
# ==========================================
# 1. 提取并重构 CLIP (强制清洗)
# ==========================================
if key.startswith("cond_stage_model.transformer."):
if "lora" in key or "embedding_manager" in key:
continue
clean_key = key.replace("cond_stage_model.transformer.", "", 1)
# 暴力特征匹配映射
if "token_embedding.weight" in clean_key:
clip_tensors["text_model.embeddings.token_embedding.weight"] = f.get_tensor(key)
elif "positional_embedding" in clean_key or "position_embedding.weight" in clean_key:
clip_tensors["text_model.embeddings.position_embedding.weight"] = f.get_tensor(key)
elif "position_ids" in clean_key:
continue # 忽略原有的，我们一会自己生成标准的
elif "text_model.encoder" in clean_key:
clip_tensors[clean_key] = f.get_tensor(key)
elif "final_layer_norm" in clean_key:
clip_tensors[clean_key] = f.get_tensor(key)
else:
clip_tensors[clean_key] = f.get_tensor(key)
# ==========================================
# 2. 提取 VAE
# ==========================================
elif key.startswith("first_stage_model."):
clean_key = key.replace("first_stage_model.", "", 1)
vae_tensors[clean_key] = f.get_tensor(key)
# ==========================================
# 3. 提取 UNet (主体)
# ==========================================
elif key.startswith("model.diffusion_model."):
clean_key = key.replace("model.diffusion_model.", "", 1)
unet_and_extras[clean_key] = f.get_tensor(key)
# ==========================================
# 4. 兼容额外组件 (EMA, ControlNet, 杂项等)
# ==========================================
else:
# 比如 model_ema., control_model., embedding_manager. 等
# 我们原样保留它们的键名，直接塞进 UNet 字典
unet_and_extras[key] = f.get_tensor(key)
print("\n💾 正在处理和保存组件...")
# ==========================================
# 5. 核心补全：自动生成缺失的 position_ids
# ==========================================
if clip_tensors:
if "text_model.embeddings.token_embedding.weight" in clip_tensors:
if "text_model.embeddings.position_ids" not in clip_tensors:
clip_tensors["text_model.embeddings.position_ids"] = torch.arange(77).unsqueeze(0).to(torch.int64)
print(" --> 💡 已自动生成缺失的 position_ids")
save_file(clip_tensors, clip_path)
print(f" --> ✅ CLIP 完美提取并保存成功 (键数: {len(clip_tensors)})")
else:
print(f" ⚠️ 警告: 未找到核心词向量 (token_embedding)，CLIP 提取失败！")
if vae_tensors:
save_file(vae_tensors, vae_path)
print(f" --> ✅ VAE 保存成功 (键数: {len(vae_tensors)})")
if unet_and_extras:
save_file(unet_and_extras, unet_path)
print(f" --> ✅ Diffusion (含额外组件) 保存成功 (总键数: {len(unet_and_extras)})")
print("\n🎉 终极切分任务完成！")
if __name__ == "__main__":
if len(sys.argv) != 2:
print("💡 用法: python split_sd15_ultimate.py <你的_sd15_模型.safetensors>")
sys.exit(1)
split_sd15_ultimate(sys.argv[1]) |