huggingface / huggingface/diffusers

save_pretrained can delete shards from another variant and leave stale indexes

Open
#14,719 1 comment 0 reactions 0 assignees View on GitHub
bug models needs-env-info
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

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](https://github.com/huggingface/diffusers/blob/c5469b7ceb606edd7ba6570dcd17d38590a18db6/src/diffusers/models/modeling_utils.py#L805-L849) 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.

```python
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

```shell
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

Contributor guide

Open the contributing guide

Research direction

Start in src/diffusers/models/modeling_utils.py at the cleanup loop around lines 805-849, then run the provided UNet2DModel reproduction with the default and ema variants. Verify that saving one variant does not remove another variant’s shards, that obsolete indexes are removed when replacing sharded weights, and that both variants reload successfully.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
76/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.