Repository navigation
save_pretrained can delete shards from another variant and leave stale indexes #14719
Description
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 shardLogs
FileNotFoundError: ...diffusion_pytorch_model-00001-of-00004.ema.safetensorsSystem 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
Activity
Vaishnavi220506 commentedon Sep 6, 2026
I would like to take this issue. I reproduced the failure on current
mainwith the smallUNet2DModelexample.Proposed scope:
The patch will stay limited to
ModelMixin.save_pretrainedand its regression test. Is this scope okay to proceed with? I will open the PR after maintainer acknowledgment.DawnofGenX commentedon Sep 22, 2026
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_pretrainedto the checkpoint identity being replaced (matching variant + serialization format) instead of prefix-matching shard files, and (2) remove the now-obsoletemodel.safetensors.index.jsonwhen 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.fix: scope save_pretrained shard cleanup to the saved variant (avoid …