Fall back to the .bin index when loading with use_safetensors=None

[aminmiral]

What does this PR do?

Fixes #14811

A sharded .bin checkpoint — what save_pretrained(..., safe_serialization=False) writes for a model bigger than max_shard_size — can't be loaded with the default from_pretrained, because when use_safetensors is None the index lookup only checks the safetensors index name. It raises OSError: Error no file named diffusion_pytorch_model.bin, and passing use_safetensors=False is the only way to load it.

When the unsharded .bin isn't found either, and use_safetensors was left at None, this now looks for the .bin index before raising. It only runs where loading raises today, so:

• anything that loads today loads exactly as before, and safetensors still wins
• setting use_safetensors yourself, True or False, behaves as before
• use_flashpack=True keeps raising as before rather than being handed pickled shards
• with nothing to load, the error is the same as today's
• Hub repos that load today make the same lookups

The shard fetch sits outside the except block so a later failure (a missing shard, say) is reported on its own instead of under the earlier "no file named ..." error.

test_sharded_bin_checkpoint_loads_with_default_use_safetensors covers no variant, variant="ema", and the legacy variant index name. All three fail on main and pass here.

Checked on CPU against main @ 7221eef: tests/models/test_modeling_common.py (15 passed), the save/load/sharding/variant tests of the UNet2D, UNet2DCondition, AutoencoderKL and Flux model suites (62 passed), and tests/pipelines/test_pipelines.py, all matching main. I also compared main and this branch over 32 load scenarios (container × sharding × variant × use_safetensors) and 48 save-twice scenarios: the only differences are the sharded .bin cases that now load and the clearer missing-shard error; every other outcome, including error text, is unchanged.

Self-review notes

I used an AI agent and ran the self-review skill; findings from the rounds:

• Fixed: the fallback ran for use_flashpack=True as well, so a missing FlashPack file could end up handing a pickled shard to flashpack.mixin.assign_from_file. It now re-raises, matching main.
• Fixed: _get_checkpoint_shard_files was called inside the except, chaining unrelated errors under "no file named ...". Moved out.
• Fixed: if not (allow_pickle and use_safetensors) was redundant — allow_pickle is only True when use_safetensors is — so it's now if not allow_pickle, like the check above it.
• Fixed: the legacy variant index path had no coverage; it's a third case of the test now.
• For the reviewer: code_style.md prefers a clear error over fallback logic. I went with the fallback because use_safetensors=None already promises one and does it for unsharded checkpoints, and for a pipeline saved this way a better message wouldn't help — use_safetensors=False fails on the transformers components, so nothing loads it. Happy to switch to an error instead.
• For the reviewer: the index probe repeats the one at the top of the resolution block, so the two need to stay in sync. Probing both index names up front would remove that, but it changes precedence — a stale .bin index would then shadow a freshly written unsharded .bin — so I kept it in the fallback.
• Known cost: a Hub repo that has only sharded .bin weights now makes one 404 request for the unsharded .bin before the index is found. Repos that load today are unaffected.
• Dead code: none.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

@sayakpaul @DN6