from safetensors.torch import load_file, save_file
import torch

# 文件路径
wrong_file_path = "model-00001-of-00003.safetensors"
correct_file_path = "/root/autodl-tmp/HF/hub/models--mistralai--Mistral-Small-3.2-24B-Instruct-2506/snapshots/1483f238dc0527ce28022b0c5252515b2552334c/model-00001-of-00010.safetensors"
output_file_path = "model-00001-of-00003.safetensors"

# 加载两个权重文件
wrong_weights = load_file(wrong_file_path)
correct_weights = load_file(correct_file_path)

# 要删除的量化键（自动识别）
keys_to_remove = [k for k in wrong_weights.keys() if k.startswith("multi_modal_projector.") and (
    k.endswith("weight_packed") or k.endswith("weight_scale") or k.endswith("weight_shape")
)]

# 删除这些键
for key in keys_to_remove:
    print(f"❌ Remove: {key}")
    del wrong_weights[key]

# 要恢复的正确键
keys_to_restore = [
    "multi_modal_projector.linear_1.weight",
    "multi_modal_projector.linear_2.weight",
    "multi_modal_projector.norm.weight",
    "multi_modal_projector.patch_merger.merging_layer.weight"
]

# 从原始全精度模型中恢复并写入
for key in keys_to_restore:
    print(f"✅ Restore: {key}")
    wrong_weights[key] = correct_weights[key]

# 保存为新文件
save_file(wrong_weights, output_file_path)
print(f"✅ 修复后的权重保存至: {output_file_path}")