huggingface / huggingface/diffusers

save_pretrained(safe_serialization=False) leaves the old safetensors checkpoint behind, and from_pretrained loads it instead of the new weights

Đang mở Phù hợp với người mới
#14,769 1 bình luận 0 reaction 0 người được giao Xem trên GitHub
bug pipelines
Ngôn ngữ chính
Python
Star
34.5k
Fork
7.3k
Merge trung bình
3 ngày 3 giờ
Pull request đã merge (30 ngày)
91

Mô tả

### Describe the bug

If a directory already has a safetensors checkpoint and I save the model again with `safe_serialization=False`, the old safetensors files are not removed. `from_pretrained` checks for safetensors first, so the next load silently returns the **old** weights. There's no error or warning.

The cleanup in `save_pretrained` ([modeling_utils.py#L804-L820](https://github.com/huggingface/diffusers/blob/83107dcbb9839e362e170dada18ed2b91849b8fd/src/diffusers/models/modeling_utils.py#L804-L820)) only deletes files matching the shard pattern (`...-00001-of-00002`). An unsharded `diffusion_pytorch_model.safetensors` and the `diffusion_pytorch_model.safetensors.index.json` never match, so they survive a `.bin` save.

What happens after saving safetensors first, then `.bin` (default `from_pretrained`):

| first save | second save | result on main |
|---|---|---|
| safetensors | bin | loads **old weights**, silently |
| safetensors (`variant="ema"`) | bin (`variant="ema"`) | loads **old weights**, silently |
| safetensors, sharded | bin | `FileNotFoundError` (old index left, its shards deleted) |

Saving an unsharded checkpoint twice in the same format works, and so does bin -> safetensors. (Going from sharded to unsharded in the same format has its own stale-index problem, which is #14719.)

This is related to #14719 but not the same problem. That issue is about deleting another variant's shards and a stale index when going sharded -> unsharded in the same format. Here it's the other format's checkpoint being left behind.

In the pipeline repro below only the diffusers components are affected: transformers components like `text_encoder/` are written as `model.safetensors` even with `safe_serialization=False`, so they never end up with two formats.

### Reproduction

```python
import tempfile, glob, os, torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None
)
with torch.no_grad():
pipe.unet.conv_in.weight.zero_()

with tempfile.TemporaryDirectory() as p:
pipe.save_pretrained(p) # unet/diffusion_pytorch_model.safetensors

with torch.no_grad():
pipe.unet.conv_in.weight.fill_(7.0)
pipe.save_pretrained(p, safe_serialization=False) # unet/diffusion_pytorch_model.bin

print(sorted(os.path.basename(f) for f in glob.glob(p + "/unet/*")))
reloaded = DiffusionPipeline.from_pretrained(p, safety_checker=None)
print(reloaded.unet.conv_in.weight[0, 0, 0, 0].item())
```

Output:

```
['config.json', 'diffusion_pytorch_model.bin', 'diffusion_pytorch_model.safetensors']
0.0
```

Expected `7.0`.

### Logs

```shell
No error or warning is printed.
```

### System Info

- 🤗 Diffusers version: 0.41.0.dev0 (main @ 83107dc)
- Platform: Linux-7.0.0-31-generic-x86_64-with-glibc2.43
- Running on Google Colab?: No
- Python version: 3.13.3
- PyTorch version (GPU?): 2.14.0+cu130 (False)
- Huggingface_hub version: 1.31.0
- Transformers version: 5.17.0
- Accelerate version: 1.15.0
- Safetensors version: 0.8.0
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help?

@sayakpaul @DN6

I have a small fix ready (removes the other format's weights file and index for the same variant when saving, plus a regression test). Would a PR be welcome? One thing to decide: it changes behaviour for anyone who saves both formats into one folder on purpose. I couldn't find that pattern anywhere in the repo, but it's your call.

Hướng dẫn đóng góp

Mở hướng dẫn đóng góp

Hướng nghiên cứu

Bắt đầu trong src/diffusers/models/modeling_utils.py tại các dòng 804-820 và theo dõi quá trình dọn dẹp do save_pretrained thực hiện khi chuyển từ safetensors sang .bin. Tái hiện vấn đề bằng script DiffusionPipeline được cung cấp, sau đó thêm hoặc chạy bài kiểm thử hồi quy được đề cập trong issue. Được xem là hoàn tất khi checkpoint safetensors cũ và index không còn tồn tại đối với cùng một biến thể, đồng thời from_pretrained tải các trọng số vừa được lưu.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Đánh giá

Công nghệ
python, pytorch
Lĩnh vực
machine-learning
Loại issue
Lỗi
Độ khó
2/5
Thời gian dự kiến
1-3 giờ
Mức độ hoạt động
Sôi nổi
Độ rõ ràng
Đặc tả rõ ràng
Mức phù hợp với người mới
78/100

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.