huggingface / huggingface/diffusers

save_pretrained can delete shards from another variant and leave stale indexes

Abierto
#14,719 1 comentario 0 reacciones 0 asignados Ver en GitHub
bug models needs-env-info
Lenguaje dominante
Python
Estrellas
34.5k
Forks
7.3k
Merge medio
3 d 3 h
PR fusionados (30 d)
91

Descripción

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

Guía de contribución

Abrir la guía de contribución

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Evaluación

Stack tecnológico
python, pytorch
Área
machine-learning
Tipo de issue
Error
Dificultad
3/5
Tiempo estimado
1-2 días
Estado de actividad
Activo
Claridad
Bien especificado
Aptitud para principiantes
76/100

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.