save_pretrained can delete shards from another variant and leave stale indexes

[Pdevadiga45]

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

[Vaishnavi220506]

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.