save_pretrained can delete shards from another variant and leave stale indexes

@sayakpaul

Describe the bug

On current main (c5469b7ceb606edd7ba6570dcd17d38590a18db6), saving default weights into a directory containing a sharded variant deletes that variant's weight files. Its index remains, so the previously loadable variant fails to reload.

The cleanup loop matches shards by prefix rather than exact checkpoint identity. The replacement path can also leave a stale index when saving a formerly sharded checkpoint as one file; loading follows that index instead of the new weights.

I have a local fix that limits cleanup to the selected variant and serialization format and removes its obsolete index. Would you be open to a PR for that scope?

Reproduction

This uses a small local model and needs no download or GPU. Set variant = None to reproduce the stale default-index case instead.

import tempfile
from diffusers import UNet2DModel

model = UNet2DModel(
    sample_size=32,
    in_channels=3,
    out_channels=3,
    block_out_channels=(4, 8),
    norm_num_groups=2,
    down_block_types=("DownBlock2D", "AttnDownBlock2D"),
    up_block_types=("AttnUpBlock2D", "UpBlock2D"),
)
variant = "ema"
with tempfile.TemporaryDirectory() as path:
    model.save_pretrained(path, variant=variant, max_shard_size="50KB")
    UNet2DModel.from_pretrained(path, variant=variant)  # succeeds
    model.save_pretrained(path, max_shard_size="100MB")
    UNet2DModel.from_pretrained(path, variant=variant)  # missing checkpoint shard

Logs

FileNotFoundError: ...diffusion_pytorch_model-00001-of-00004.ema.safetensors

System Info

🤗 Diffusers version: 0.41.0.dev0
Platform: Windows-11-10.0.26200-SP0
Running on Google Colab?: No
Python version: 3.12.5
PyTorch version (GPU?): 2.14.0+cpu (False)
Huggingface_hub version: 1.30.0
Transformers version: 5.16.1
Accelerate version: 1.14.0
PEFT version: 0.20.1.dev0
Safetensors version: 0.8.0
xFormers version: not installed
Accelerator: NA
Using GPU in script?: No
Using distributed or parallel set-up in script?: No

Who can help?

@sayakpaul @DN6

I would like to take this issue. I reproduced the failure on current main with the small UNet2DModel example.

Proposed scope:

restrict stale shard cleanup to the currently selected variant and serialization format;
remove the selected variant's obsolete index when the new save is unsharded;
add CPU regression coverage that preserves another variant's shards, removes the stale index, and verifies both checkpoints can reload.

The patch will stay limited to ModelMixin.save_pretrained and its regression test. Is this scope okay to proceed with? I will open the PR after maintainer acknowledgment.

I'd like to take this on. I have a fix scoped exactly to the two issues you describe: (1) scope the old-weight cleanup in save_pretrained to the checkpoint identity being replaced (matching variant + serialization format) instead of prefix-matching shard files, and (2) remove the now-obsolete model.safetensors.index.json when switching between sharded and unsharded layouts. Will validate against the repro in this issue (tiny UNet2DModel, variant='ema', sharded save followed by a default-variant save) with regression tests for both cases.

Navigation Menu

Repository navigation

save_pretrained can delete shards from another variant and leave stale indexes #14719

Describe the bug

Reproduction

Logs

System Info

Who can help?

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

save_pretrained can delete shards from another variant and leave stale indexes #14719

Description

Describe the bug

Reproduction

Logs

System Info

Who can help?

Activity

Vaishnavi220506 commented on Sep 6, 2026

DawnofGenX commented on Sep 22, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions