A sharded .bin checkpoint can't be loaded with the default from_pretrained (use_safetensors=None)

[aminmiral]

Describe the bug

save_pretrained(..., safe_serialization=False) shards the checkpoint when the model is bigger than max_shard_size (10GB by default), writing diffusion_pytorch_model.bin.index.json plus the shards. Loading that folder back with a plain from_pretrained fails:

OSError: Error no file named diffusion_pytorch_model.bin found in directory ...

It only works if you pass use_safetensors=False.

The docs for use_safetensors say that with None "the safetensors weights are downloaded if they're available", so I'd expect a fallback when they aren't. That fallback exists for unsharded checkpoints (the "Defaulting to unsafe serialization" path), but the sharded index lookup only checks the safetensors index name (modeling_utils.py#L1237, model_loading_utils.py#L514), so diffusion_pytorch_model.bin.index.json is never looked for.

For a whole pipeline there's no argument that loads it. transformers components are saved as safetensors even with safe_serialization=False, so after pipe.save_pretrained(path, safe_serialization=False, max_shard_size=...) the folder has unet/ and vae/ as sharded .bin and text_encoder/ as sharded safetensors:

• the default from_pretrained fails on unet or vae with the error above
• use_safetensors=False fails on text_encoder (no file named model.safetensors, or pytorch_model.bin)

Saving and reloading a sharded .bin UNet2DModel fails the same way on 0.29.0, 0.32.0, 0.36.0, 0.40.0 and main, so it isn't a recent regression. Every sharded test fixture I could find (hf-internal-testing/tiny-flux-sharded, unet2d-sharded-dummy, tiny-sd-unet-sharded-*) is safetensors, which may be why it hasn't come up.

Reproduction

import tempfile
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None
)

with tempfile.TemporaryDirectory() as path:
    # max_shard_size is lowered only so this tiny model shards; a model over 10GB
    # shards with the default
    pipe.save_pretrained(path, safe_serialization=False, max_shard_size="50KB")
    DiffusionPipeline.from_pretrained(path, safety_checker=None)

Logs

OSError: Error no file named diffusion_pytorch_model.bin found in directory /tmp/.../unet.

(it's either unet or vae that fails first depending on the run; both are sharded .bin)

System Info

• 🤗 Diffusers version: 0.41.0.dev0 (main @ 7221eef)
• 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 fix with a regression test and I'm opening a PR for it.