From 055e9797ed8e806598b45a53c97a659c9ce54ef3 Mon Sep 17 00:00:00 2001 From: kyo-zzz Date: Tue, 15 Sep 2026 11:17:06 +0800 Subject: [PATCH] Fix stale checkpoint cleanup in `save_pretrained` `ModelMixin.save_pretrained` cleaned a previous save by removing files whose name matched the shard pattern of the current save, but it only recognized sharded files. Single-file checkpoints of the other container were left behind: re-saving a `.safetensors` checkpoint as `.bin` (or the reverse) kept the old file on disk, and because `from_pretrained` prefers sharded and safetensors checkpoints, the stale one silently took precedence over the freshly written weights. Extract the cleanup into `_get_superseded_checkpoint_files`, which returns every artifact of the current save's variant that the new save replaces: single-file weights, shards, and sharded indexes, in either container. Files belonging to other variants (e.g. `ema`) share the directory and are left untouched. Add a regression test covering both directions of the container switch and the variant coexistence case. Fixes #14769 --- src/diffusers/models/modeling_utils.py | 67 ++++++++++++++++++-------- tests/models/test_modeling_common.py | 43 +++++++++++++++++ 2 files changed, 90 insertions(+), 20 deletions(-) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index f5515780d136..2c2b406b377a 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -139,6 +139,47 @@ def __exit__(self, *args, **kwargs): from accelerate.utils import load_offloaded_weights, save_offload_index +def _get_superseded_checkpoint_files(save_directory: str, weights_name_pattern: str) -> list[str]: + """List checkpoint files left by an earlier save that the current one supersedes. + + A save writes either `.bin` or `.safetensors`, as a single file or as shards plus an index. + Re-saving into the same directory with a different container or layout must remove the previous + artifacts: `from_pretrained` prefers sharded and safetensors checkpoints, so survivors shadow + the freshly written weights and get loaded silently. + + Only files belonging to the same variant are returned -- checkpoints of other variants (e.g. + `ema`) coexist in the same directory and must be preserved. + """ + # `weights_name_pattern` looks like "diffusion_pytorch_model{suffix}.safetensors" or + # "...ema{suffix}.bin"; stripping the container and the shard placeholder leaves the stem every + # artifact of this save shares, e.g. "diffusion_pytorch_model" or "diffusion_pytorch_model.ema". + stem = weights_name_pattern.replace(".bin", "").replace(".safetensors", "").replace("{suffix}", "") + superseded = [] + for filename in os.listdir(save_directory): + if not os.path.isfile(os.path.join(save_directory, filename)): + continue + # Single-file weights and sharded indexes of this variant, in either container: a re-save + # with the other container supersedes them. + if filename in ( + f"{stem}.bin", + f"{stem}.safetensors", + f"{stem}.bin.index.json", + f"{stem}.safetensors.index.json", + ): + superseded.append(filename) + continue + # Shards are recognized by the same `-00001-of-00005` marker used when splitting the state + # dict. The stem is matched exactly instead of by prefix, so a default save leaves another + # variant's shards (`diffusion_pytorch_model.ema-00001-of-00002`) untouched. + shard_stem, ext = os.path.splitext(filename) + if ext not in (".bin", ".safetensors"): + continue + match = _REGEX_SHARD.fullmatch(shard_stem) + if match is not None and match.group(1) == stem: + superseded.append(filename) + return superseded + + def get_parameter_device(parameter: torch.nn.Module) -> torch.device: from ..hooks.group_offloading import _get_group_onload_device @@ -819,29 +860,15 @@ def save_pretrained( state_dict, max_shard_size=max_shard_size, filename_pattern=weights_name_pattern ) - # Clean the folder from a previous save. Match only shards of the exact - # checkpoint being written (same variant and serialization format): the - # prefix must be anchored, or a default save (weights_name - # "diffusion_pytorch_model") would also match another variant's shards - # ("diffusion_pytorch_model.ema-00001-of-00002") via a bare startswith - # and delete them while leaving that variant's index behind. + # Clean the folder from a previous save. Only artifacts of the checkpoint being written + # are removed (same variant, either container): matching by bare prefix would also hit + # another variant's shards ("diffusion_pytorch_model.ema-00001-of-00002") and delete + # them while leaving that variant's index behind. if is_main_process: - 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): + for filename in _get_superseded_checkpoint_files(save_directory, weights_name_pattern): 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: - os.remove(full_filename) + os.remove(os.path.join(save_directory, filename)) for filename, tensors in state_dict_split.filename_to_tensors.items(): shard = {tensor: state_dict[tensor].contiguous() for tensor in tensors} diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index 528bad336ffa..b72b0e003924 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -302,6 +302,49 @@ def build(): reloaded = UNet2DModel.from_pretrained(path, variant=variant) assert all(torch.equal(p1, p2) for p1, p2 in zip(model.parameters(), reloaded.parameters())) + @pytest.mark.parametrize("variant", [None, "ema"]) + def test_save_pretrained_removes_superseded_checkpoints(self, variant): + r""" + Re-saving a checkpoint with a different container must not leave the previous files behind. + `from_pretrained` prefers safetensors and sharded checkpoints, so leftovers would shadow the + freshly written weights and be loaded silently. + """ + with tempfile.TemporaryDirectory() as tmpdirname: + model = UNet2DConditionModel( + block_out_channels=(4, 8), + norm_num_groups=4, + down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), + up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"), + cross_attention_dim=8, + attention_head_dim=2, + sample_size=8, + in_channels=4, + out_channels=4, + layers_per_block=1, + ) + kwargs = {"variant": variant} if variant is not None else {} + # A checkpoint of the other variant lives in the same directory and must survive. + other_kwargs = {"variant": "ema"} if variant is None else {} + other_stem = "diffusion_pytorch_model" + (".ema" if other_kwargs else "") + model.save_pretrained(tmpdirname, safe_serialization=True, **other_kwargs) + + model.save_pretrained(tmpdirname, safe_serialization=True, **kwargs) + expected_stem = "diffusion_pytorch_model" + (f".{variant}" if variant else "") + assert f"{expected_stem}.safetensors" in os.listdir(tmpdirname) + + # Re-save with the other container: the safetensors file must not survive. + model.save_pretrained(tmpdirname, safe_serialization=False, **kwargs) + files = os.listdir(tmpdirname) + assert f"{expected_stem}.bin" in files + assert f"{expected_stem}.safetensors" not in files + if other_stem != expected_stem: + assert f"{other_stem}.safetensors" in files + + # And the reloaded weights must come from the freshly written file. + reloaded = UNet2DConditionModel.from_pretrained(tmpdirname, **kwargs) + for key, value in model.state_dict().items(): + assert torch.equal(value, reloaded.state_dict()[key]) + @is_staging_test class TestModelPushToHub: