Skip to content

save_pretrained(safe_serialization=False) leaves the old safetensors checkpoint behind, and from_pretrained loads it instead of the new weights #14769

Description

@aminmiral

Describe the bug

If a directory already has a safetensors checkpoint and I save the model again with safe_serialization=False, the old safetensors files are not removed. from_pretrained checks for safetensors first, so the next load silently returns the old weights. There's no error or warning.

The cleanup in save_pretrained (modeling_utils.py#L804-L820) only deletes files matching the shard pattern (...-00001-of-00002). An unsharded diffusion_pytorch_model.safetensors and the diffusion_pytorch_model.safetensors.index.json never match, so they survive a .bin save.

What happens after saving safetensors first, then .bin (default from_pretrained):

first save second save result on main
safetensors bin loads old weights, silently
safetensors (variant="ema") bin (variant="ema") loads old weights, silently
safetensors, sharded bin FileNotFoundError (old index left, its shards deleted)

Saving an unsharded checkpoint twice in the same format works, and so does bin -> safetensors. (Going from sharded to unsharded in the same format has its own stale-index problem, which is #14719.)

This is related to #14719 but not the same problem. That issue is about deleting another variant's shards and a stale index when going sharded -> unsharded in the same format. Here it's the other format's checkpoint being left behind.

In the pipeline repro below only the diffusers components are affected: transformers components like text_encoder/ are written as model.safetensors even with safe_serialization=False, so they never end up with two formats.

Reproduction

import tempfile, glob, os, torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None
)
with torch.no_grad():
    pipe.unet.conv_in.weight.zero_()

with tempfile.TemporaryDirectory() as p:
    pipe.save_pretrained(p)                               # unet/diffusion_pytorch_model.safetensors

    with torch.no_grad():
        pipe.unet.conv_in.weight.fill_(7.0)
    pipe.save_pretrained(p, safe_serialization=False)     # unet/diffusion_pytorch_model.bin

    print(sorted(os.path.basename(f) for f in glob.glob(p + "/unet/*")))
    reloaded = DiffusionPipeline.from_pretrained(p, safety_checker=None)
    print(reloaded.unet.conv_in.weight[0, 0, 0, 0].item())

Output:

['config.json', 'diffusion_pytorch_model.bin', 'diffusion_pytorch_model.safetensors']
0.0

Expected 7.0.

Logs

No error or warning is printed.

System Info

  • 🤗 Diffusers version: 0.41.0.dev0 (main @ 83107dc)
  • Platform: Linux-7.0.0-31-generic-x86_64-with-glibc2.43
  • Running on Google Colab?: No
  • Python version: 3.13.3
  • PyTorch version (GPU?): 2.14.0+cu130 (False)
  • Huggingface_hub version: 1.31.0
  • Transformers version: 5.17.0
  • Accelerate version: 1.15.0
  • Safetensors version: 0.8.0
  • Using GPU in script?: No
  • Using distributed or parallel set-up in script?: No

Who can help?

@sayakpaul @DN6

I have a small fix ready (removes the other format's weights file and index for the same variant when saving, plus a regression test). Would a PR be welcome? One thing to decide: it changes behaviour for anyone who saves both formats into one folder on purpose. I couldn't find that pattern anywhere in the repo, but it's your call.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions