| 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_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():
|
|
|
|
|
|
|
|
|
| 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)
|
|
|
|
|
|
|
|
|
| elif key.startswith("first_stage_model."):
|
| clean_key = key.replace("first_stage_model.", "", 1)
|
| vae_tensors[clean_key] = f.get_tensor(key)
|
|
|
|
|
|
|
|
|
| elif key.startswith("model.diffusion_model."):
|
| clean_key = key.replace("model.diffusion_model.", "", 1)
|
| unet_and_extras[clean_key] = f.get_tensor(key)
|
|
|
|
|
|
|
|
|
| else:
|
|
|
|
|
| unet_and_extras[key] = f.get_tensor(key)
|
|
|
| print("\n💾 正在处理和保存组件...")
|
|
|
|
|
|
|
|
|
| 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]) |
