[AloysJehwin]

The bug (#14719)

Saving default weights into a directory that already holds a sharded variant deletes that variant's shard files while leaving its index, so the previously loadable variant fails to reload. A formerly-sharded checkpoint re-saved as a single file also leaves a stale index that load then follows instead of the new weights.

Reproduction (CPU-only, no download)

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"))
with tempfile.TemporaryDirectory() as path:
    model.save_pretrained(path, variant="ema", max_shard_size="50KB")
    UNet2DModel.from_pretrained(path, variant="ema")   # ok
    model.save_pretrained(path, max_shard_size="100MB")
    UNet2DModel.from_pretrained(path, variant="ema")   # FileNotFoundError: missing shard

Root cause

The cleanup loop derived weights_without_ext from the pattern with the {suffix} and variant stripped, then deleted any file that startswith it and matched a generic shard regex. For a default save the prefix is diffusion_pytorch_model, which startswith-matches another variant's shard diffusion_pytorch_model.ema-00001-of-00002 (the .ema sits after the prefix). The variant's index (...index.ema.json) doesn't match the regex, so it's kept — leaving a stale index pointing at deleted shards.

The fix

• Anchor shard matching to the exact per-variant, per-format shard name: <weights_without_ext>-NNNNN-of-NNNNN built from this save's weights_name_pattern, so a default save no longer matches .ema shards.
• Remove this checkpoint's own now-obsolete index when the save is no longer sharded (single file), so load follows the new weights. A fresh index is still written when the save is sharded.

Test

TestModelUtils::test_save_pretrained_preserves_other_variant_shards (parametrized ema / default) saves a sharded checkpoint, then a default unsharded save, and asserts the first checkpoint still reloads with matching parameters. Verified it fails on main (missing-shard FileNotFoundError) and passes with the change. The TestModelUtils sharding/variant tests pass (11), and ruff is clean.

Fixes #14719

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[DN6]

I think this can be simplified

Suggested change

	ext = ".bin" if not safe_serialization else ".safetensors"
	weights_without_ext = weights_name_pattern.replace("{suffix}", "").replace(ext, "")
	shard_re = re.compile(re.escape(weights_without_ext) + r"-\d{5}-of-\d{5}")
	base_index_file = WEIGHTS_INDEX_NAME if not safe_serialization else SAFE_WEIGHTS_INDEX_NAME
	index_file = _add_variant(base_index_file, variant)
	for filename in os.listdir(save_directory):
	if filename in state_dict_split.filename_to_tensors.keys():
	continue
	full_filename = os.path.join(save_directory, filename)
	if not os.path.isfile(full_filename):
	continue
	weights_without_ext = weights_name_pattern.replace(".bin", "").replace(".safetensors", "")
	weights_without_ext = weights_without_ext.replace("{suffix}", "")
	filename_without_ext = filename.replace(".bin", "").replace(".safetensors", "")
	# make sure that file to be deleted matches format of sharded file, e.g. pytorch_model-00001-of-00005
	if (
	filename.startswith(weights_without_ext)
	and _REGEX_SHARD.fullmatch(filename_without_ext) is not None
	):
	filename_without_ext = filename.replace(ext, "")
	is_prior_shard = shard_re.fullmatch(filename_without_ext) is not None
	# Drop this checkpoint's stale index too: when a formerly sharded
	# checkpoint is now saved as a single file (or fewer shards), a
	# leftover index would be followed on load instead of the new
	# weights. A fresh index is rewritten below when still sharded.
	is_stale_index = filename == index_file and not state_dict_split.is_sharded
	if is_prior_shard or is_stale_index:
	weights_without_ext, ext = os.path.splitext(weights_name)
	index_file = _add_variant(SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME, variant)
	for filename in os.listdir(save_directory):
	if filename in state_dict_split.filename_to_tensors.keys():
	continue
	full_filename = os.path.join(save_directory, filename)
	if not os.path.isfile(full_filename):
	continue
	# only remove shards of the checkpoint being written; another variant's shards share the prefix
	shard_match = _REGEX_SHARD.fullmatch(filename.removesuffix(ext))
	is_prior_shard = shard_match is not None and shard_match.group(1) == weights_without_ext
	is_stale_index = filename == index_file and not state_dict_split.is_sharded
	if is_prior_shard or is_stale_index:

[DN6]

/diffusers-bot style

[github-actions]

Style bot fixed some files and pushed the changes.