import os
import random
from safetensors.torch import load_file, save_file

#ckpt_path = "dunhuang.safetensors"
#ckpt_path = "hanfu.safetensors"
ckpts = os.listdir(".")
for ckpt_path in ckpts:
    if "tmp" in ckpt_path or "new" in ckpt_path: continue

    new_ckpt_path = ckpt_path.replace(".safetensors", "_new.safetensors")
    weights = load_file(ckpt_path)
    for key in weights:
        weights[key] += random.uniform(0, 1) * 1e-7
    save_file(weights, new_ckpt_path)
    print(f"Modified weights saved to {new_ckpt_path}")