fix: scope save_pretrained shard cleanup to the saved variant (avoid deleting another variant's shards)

src/diffusers/models/modeling_utils.py

@@ -819,22 +819,28 @@ def save_pretrained(
                 state_dict, max_shard_size=max_shard_size, filename_pattern=weights_name_pattern
             )
 
-            # Clean the folder from a previous save
+            # 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.
             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):
                     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
-                    ):
+                    # 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)
 
             for filename, tensors in state_dict_split.filename_to_tensors.items():

tests/models/test_modeling_common.py

@@ -276,6 +276,33 @@ def get_dummy_inputs():
 
         SD3Transformer2DModel._keep_in_fp32_modules = fp32_modules
 
+    @pytest.mark.parametrize("variant", ["ema", None])
+    def test_save_pretrained_preserves_other_variant_shards(self, variant):
+        """Saving default weights into a directory that holds a sharded checkpoint must not
+        delete that checkpoint's shards or leave a stale index behind. Regression for
+        prefix-based cleanup matching another variant's shards / a formerly-sharded default."""
+        from diffusers import UNet2DModel
+
+        def build():
+            return 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"),
+            )
+
+        model = build()
+        with tempfile.TemporaryDirectory() as path:
+            model.save_pretrained(path, variant=variant, max_shard_size="50KB")
+            UNet2DModel.from_pretrained(path, variant=variant)
+            # A subsequent default (unsharded) save must leave the first checkpoint reloadable.
+            model.save_pretrained(path, max_shard_size="100MB")
+            reloaded = UNet2DModel.from_pretrained(path, variant=variant)
+            assert all(torch.equal(p1, p2) for p1, p2 in zip(model.parameters(), reloaded.parameters()))
+
 
 class UNetTesterMixin:
     @staticmethod