from safetensors.torch import load_file, save_file

dict = load_file("InputModelFilename.safetensors", device="cuda")
keys = dict.keys()
fixed= {k.replace("model.diffusion_model.", ""): dict[k] for k in keys}

print(fixed.keys())
save_file(fixed, "ModelToSaveFilename.safetensors")