huggingface / huggingface/diffusers
save_pretrained can delete shards from another variant and leave stale indexes
- 主要言語
- 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