import re
from safetensors.torch import load_file, save_file

# 加载 .safetensors 文件
#file_path = "iniverseMixXLSFWNSFW_flux1DNsfwFp8V10.safetensors"
file_path = "iniverseMixXLSFWNSFW_flux1DNsfwFp16V12.safetensors"
model = load_file(file_path)

# 定义通配符模式，使用正则表达式匹配要删除的键
# 例如，删除以 "layer" 开头的所有权重
pattern = re.compile(r'^text_encoders.*')
#pattern = re.compile(r'^vae.*')

# 找到需要删除的键
keys_to_delete = [key for key in model.keys() if pattern.match(key)]

# 删除匹配的键
for key in keys_to_delete:
    del model[key]

# 保存修改后的模型
#save_file(model, "iniverseMixXLSFWNSFW_flux1DNsfwFp8_12GV10.safetensors")
save_file(model, "iniverseMixXLSFWNSFW_flux1DNsfwFp16_22GV12.safetensors")

print(f"Deleted {len(keys_to_delete)} keys matching pattern.")