From 1c895eee08267c8375b808c7bda4046bde117232 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 | 46 +++++++++++++------ tests/models/test_modeling_common.py | 63 ++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 13 deletions(-) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 425f2f29235e..0f0ba3df56d2 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -138,6 +138,37 @@ 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)) or not filename.startswith(stem): + continue + suffix = filename[len(stem) :] + # Single-file weights and sharded indexes of this variant are always superseded. Shards are + # recognized by the same `-00001-of-00005` marker used when splitting the state dict. + if suffix in (".bin", ".safetensors", ".bin.index.json", ".safetensors.index.json"): + superseded.append(filename) + elif suffix.endswith((".bin", ".safetensors")): + stem_without_ext = suffix[: -len(".bin")] if suffix.endswith(".bin") else suffix[: -len(".safetensors")] + if _REGEX_SHARD.fullmatch(stem_without_ext) is not None: + superseded.append(filename) + return superseded + + def get_parameter_device(parameter: torch.nn.Module) -> torch.device: from ..hooks.group_offloading import _get_group_onload_device @@ -803,21 +834,10 @@ def save_pretrained( # Clean the folder from a previous save if is_main_process: - 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 - 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 - ): - 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 b1bdaabbad7a..25e3cea69baa 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -276,6 +276,69 @@ def get_dummy_inputs(): SD3Transformer2DModel._keep_in_fp32_modules = fp32_modules + @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 {} + + 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 + + # 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]) + + def test_save_pretrained_preserves_other_variants(self): + r""" + Cleaning up superseded checkpoints must stay scoped to the variant being saved: a plain + save and a `variant="ema"` save coexist in the same directory. + """ + 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, + ) + + model.save_pretrained(tmpdirname, safe_serialization=True) + model.save_pretrained(tmpdirname, safe_serialization=True, variant="ema") + + files = os.listdir(tmpdirname) + assert "diffusion_pytorch_model.safetensors" in files + assert "diffusion_pytorch_model.ema.safetensors" in files + class UNetTesterMixin: @staticmethod