huggingface / huggingface/diffusers

save_pretrained can delete shards from another variant and leave stale indexes

オープン
#14,719 コメント 1 件 リアクション 0 件 担当者 0 名 GitHub で見る
bug models needs-env-info
主要言語
Python
スター
34.5k
フォーク
7.3k
平均マージ
3日 3時間
マージ済み PR(30日)
91

説明

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

コントリビューションガイド

コントリビューションガイドを開く

調査の方向性

src/diffusers/models/modeling_utils.py の805-849行付近にあるクリーンアップループから開始し、その後、default と ema のバリアントを使って、提供されている UNet2DModel の再現手順を実行します。一方のバリアントを保存しても別のバリアントのシャードが削除されないこと、シャーディングされた重みを置き換えるときに古いインデックスが削除されること、そして両方のバリアントが正常に再読み込みされることを確認します。

索引モデルが issue の本文から書いたものです。

評価

技術スタック
python, pytorch
領域
machine-learning
issue の種類
バグ
難易度
3/5
見積もり時間
1〜2日
活発さ
活発
明瞭さ
明確に書かれている
初心者へのやさしさ
76/100

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。