huggingface / huggingface/diffusers

save_pretrained can delete shards from another variant and leave stale indexes

Aperta
#14,719 1 commento 0 reazioni 0 assegnatari Vedi su GitHub
bug models needs-env-info
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

## 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

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

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.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Valutazione

Stack tecnologico
python, pytorch
Ambito
machine-learning
Tipo di issue
Bug
Difficoltà
3/5
Tempo stimato
1-2 giorni
Stato di attività
Attiva
Chiarezza
Specificata chiaramente
Idoneità per principianti
76/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.