ModularPipeline.save_pretrained: change the default overwrite_modular_index to be True

docs/source/en/modular_diffusers/modular_pipeline.md

@@ -165,7 +165,7 @@ ModularPipeline {
 }
 ```
 
-If you pass a repository to [`~ModularPipelineBlocks.init_pipeline`], it overrides the loading path by matching your block's components against the pipeline config in that repository (`model_index.json` or `modular_model_index.json`).
+If you pass a repository to [`~ModularPipelineBlocks.init_pipeline`], it overrides the loading path by matching your block's components against the pipeline config in that repository (`model_index.json` or `modular_model_index.json`). See [Modular repository](#modular-repository) for how loading specs are recorded and saved.
 
 In the example below, the `pretrained_model_name_or_path` will be updated to `"stabilityai/stable-diffusion-xl-base-1.0"`.
 
@@ -415,6 +415,30 @@ pipeline = ModularPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base
 pipeline.save_pretrained("local/path", repo_id="my-username/sdxl-modular", push_to_hub=True)
 ```
 
+By default, [`~ModularPipeline.save_pretrained`] writes each currently loaded component that Diffusers can serialize. Components that are not loaded, or cannot be serialized, are not written and keep their existing loading specifications. This gives you two ways to save, depending on what you want.
+
+### Save a self-contained copy
+
+Load all the components, then save. Every spec points at the result, so it reloads entirely from one place, including offline. Downloading files with `hf download --local-dir` does not rewrite `modular_model_index.json`, so an existing index may still point to the Hub.
+
+  ```py
+  pipe = ModularPipeline.from_pretrained("MiniMaxAI/MiniMax-H3")
+  pipe.load_components()
+  pipe.save_pretrained("path/to/local-copy")
+  ```
+
+### Keep references to existing components
+
+Load only what's new (or nothing at all). Only loaded components are written; everything else stays a pointer to its original repository. Use this mode when you want to replace one component while continuing to load the others from their original repository. For example, save a custom transformer while the remaining components continue to load from the base repository.
+
+  ```py
+  pipe = ModularPipeline.from_pretrained("black-forest-labs/FLUX.2-dev")
+  pipe.update_components(transformer=my_custom_transformer)  # the only component in memory
+  pipe.save_pretrained("local/path", repo_id="my-username/flux2-custom-transformer", push_to_hub=True)
+  ```
+
+Pass `overwrite_modular_index=False` to preserve the recorded specs of saved components that already have a load ID. Components without a recorded source, such as models added with `update_components`, are still rewritten to point to the destination.
+
 A modular repository can also include custom pipeline blocks as Python code. This allows you to share specialized blocks that aren't native to Diffusers. For example, [diffusers/Florence2-image-Annotator](https://huggingface.co/diffusers/Florence2-image-Annotator) contains custom blocks alongside the loading configuration:
 
 ```

src/diffusers/modular_pipelines/modular_pipeline.py

@@ -36,7 +36,7 @@
     _unwrap_model,
     simple_get_class_obj,
 )
-from ..utils import PushToHubMixin, is_accelerate_available, logging
+from ..utils import PushToHubMixin, deprecate, is_accelerate_available, logging
 from ..utils.dynamic_modules_utils import get_class_from_dynamic_module, resolve_trust_remote_code
 from ..utils.hub_utils import _resolve_revision, load_or_create_model_card, populate_model_card
 from ..utils.torch_utils import empty_device_cache, is_compiled_module
@@ -1974,10 +1974,13 @@ def save_pretrained(
             push_to_hub (`bool`, *optional*, defaults to `False`):
                 Whether to push the pipeline to the Hugging Face model hub after saving it.
             **kwargs: Additional keyword arguments:
-                - `overwrite_modular_index` (`bool`, *optional*, defaults to `False`):
-                    When saving a Modular Pipeline, its components in `modular_model_index.json` may reference repos
-                    different from the destination repo. Setting this to `True` updates all component references in
-                    `modular_model_index.json` so they point to the repo specified by `repo_id`.
+                - `overwrite_modular_index` (`bool`, *optional*, defaults to `True`):
+                    Whether to update `modular_model_index.json` so each saved component's loading spec points to the
+                    destination: `repo_id` when pushing to the Hub, otherwise `save_directory`. Components that are not
+                    loaded are not saved and always keep their recorded loading specs. Pass `False` to also preserve
+                    the recorded specs of the components being saved (e.g. for an index that deliberately references
+                    other repositories); components without a load id (such as custom models added with
+                    `update_components`) are still rewritten since they have no recorded source.
                 - `repo_id` (`str`, *optional*):
                     The repository ID to push the pipeline to. Defaults to the last component of `save_directory`.
                 - `commit_message` (`str`, *optional*):
@@ -1989,7 +1992,17 @@ def save_pretrained(
                 - `token` (`str`, *optional*):
                     The Hugging Face token to use for authentication.
         """
-        overwrite_modular_index = kwargs.pop("overwrite_modular_index", False)
+        if "overwrite_modular_index" not in kwargs:
+            deprecate(
+                "overwrite_modular_index",
+                "0.43.0",
+                "The default of `overwrite_modular_index` in `ModularPipeline.save_pretrained` changed from `False`"
+                " to `True`: the saved `modular_model_index.json` now points each saved component at the destination"
+                " (the save directory, or `repo_id` when pushing to the Hub). Pass `overwrite_modular_index=False`"
+                " to keep the previously recorded loading specs, or pass `True` explicitly to silence this warning.",
+                standard_warn=False,
+            )
+        overwrite_modular_index = kwargs.pop("overwrite_modular_index", True)
         repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
 
         if push_to_hub:
@@ -2060,6 +2073,8 @@ def save_pretrained(
                 library, class_name, component_spec_dict = self.config[component_name]
                 component_spec_dict["pretrained_model_name_or_path"] = repo_id if push_to_hub else save_directory
                 component_spec_dict["subfolder"] = component_name
+                component_spec_dict["variant"] = variant if save_method_accept_variant else None
+                component_spec_dict["revision"] = None
                 self.register_to_config(**{component_name: (library, class_name, component_spec_dict)})
 
         self.save_config(save_directory=save_directory)

tests/modular_pipelines/test_modular_pipeline_loading.py

@@ -16,6 +16,7 @@
 import json
 import os
 
+import pytest
 import torch
 
 from diffusers import AutoModel, ControlNetModel, ModularPipeline, UNet2DConditionModel
@@ -120,16 +121,17 @@ def test_load_components_skips_invalid_pretrained_path(self):
 
 
 class TestCustomModelSavePretrained:
-    def test_save_pretrained_updates_index_for_local_model(self, tmp_path):
-        """When a component without _diffusers_load_id (custom/local model) is saved,
-        modular_model_index.json should point to the save directory."""
+    @pytest.mark.parametrize("overwrite_modular_index", [True, False])
+    def test_save_pretrained_updates_index_for_local_model(self, tmp_path, overwrite_modular_index):
+        """A component without _diffusers_load_id (custom/local model) is rewritten to the save directory in both
+        modes; other components' specs follow `overwrite_modular_index`."""
         pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
         pipe.load_components(dtype=torch.float32)
 
         pipe.unet._diffusers_load_id = "null"
 
         save_dir = str(tmp_path / "my-pipeline")
-        pipe.save_pretrained(save_dir)
+        pipe.save_pretrained(save_dir, overwrite_modular_index=overwrite_modular_index)
 
         with open(os.path.join(save_dir, "modular_model_index.json")) as f:
             index = json.load(f)
@@ -139,7 +141,8 @@ def test_save_pretrained_updates_index_for_local_model(self, tmp_path):
         assert unet_spec["subfolder"] == "unet"
 
         _library, _cls, vae_spec = index["vae"]
-        assert vae_spec["pretrained_model_name_or_path"] == "hf-internal-testing/tiny-stable-diffusion-xl-pipe"
+        expected_vae = save_dir if overwrite_modular_index else "hf-internal-testing/tiny-stable-diffusion-xl-pipe"
+        assert vae_spec["pretrained_model_name_or_path"] == expected_vae
 
     def test_save_pretrained_roundtrip_with_local_model(self, tmp_path):
         """A pipeline with a custom/local model should be saveable and re-loadable with identical outputs."""
@@ -164,9 +167,10 @@ def test_save_pretrained_roundtrip_with_local_model(self, tmp_path):
         for key in original_state_dict:
             assert torch.equal(original_state_dict[key], loaded_state_dict[key]), f"Mismatch in {key}"
 
-    def test_save_pretrained_updates_index_for_model_with_no_load_id(self, tmp_path):
+    @pytest.mark.parametrize("overwrite_modular_index", [True, False])
+    def test_save_pretrained_updates_index_for_model_with_no_load_id(self, tmp_path, overwrite_modular_index):
         """testing the workflow of update the pipeline with a custom model and save the pipeline,
-        the modular_model_index.json should point to the save directory."""
+        the modular_model_index.json should point to the save directory in both modes."""
         pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
         pipe.load_components(dtype=torch.float32)
 
@@ -177,6 +181,26 @@ def test_save_pretrained_updates_index_for_model_with_no_load_id(self, tmp_path)
 
         pipe.update_components(unet=unet)
 
+        save_dir = str(tmp_path / "my-pipeline")
+        pipe.save_pretrained(save_dir, overwrite_modular_index=overwrite_modular_index)
+
+        with open(os.path.join(save_dir, "modular_model_index.json")) as f:
+            index = json.load(f)
+
+        _library, _cls, unet_spec = index["unet"]
+        assert unet_spec["pretrained_model_name_or_path"] == save_dir
+        assert unet_spec["subfolder"] == "unet"
+
+        _library, _cls, vae_spec = index["vae"]
+        expected_vae = save_dir if overwrite_modular_index else "hf-internal-testing/tiny-stable-diffusion-xl-pipe"
+        assert vae_spec["pretrained_model_name_or_path"] == expected_vae
+
+    def test_save_pretrained_default_writes_self_contained_local_copy(self, tmp_path):
+        """By default the saved index points at the save directory, so the copy reloads offline; a component
+        that was never loaded is not saved and keeps its recorded spec."""
+        pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
+        pipe.load_components(names=["unet"], dtype=torch.float32)
+
         save_dir = str(tmp_path / "my-pipeline")
         pipe.save_pretrained(save_dir)
 
@@ -186,10 +210,15 @@ def test_save_pretrained_updates_index_for_model_with_no_load_id(self, tmp_path)
         _library, _cls, unet_spec = index["unet"]
         assert unet_spec["pretrained_model_name_or_path"] == save_dir
         assert unet_spec["subfolder"] == "unet"
+        assert unet_spec["revision"] is None
 
         _library, _cls, vae_spec = index["vae"]
         assert vae_spec["pretrained_model_name_or_path"] == "hf-internal-testing/tiny-stable-diffusion-xl-pipe"
 
+        loaded_pipe = ModularPipeline.from_pretrained(save_dir)
+        loaded_pipe.load_components(names=["unet"], dtype=torch.float32, local_files_only=True)
+        assert loaded_pipe.unet is not None
+
     def test_save_pretrained_overwrite_modular_index(self, tmp_path):
         """With overwrite_modular_index=True, all component references should point to the save directory."""
         pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")

tests/modular_pipelines/testing_utils/loading.py

@@ -88,7 +88,7 @@ def test_modular_index_consistency(self, tmp_path):
         components_spec = pipe._component_specs
         components = sorted(components_spec.keys())
 
-        pipe.save_pretrained(str(tmp_path))
+        pipe.save_pretrained(str(tmp_path), overwrite_modular_index=False)
         index_file = tmp_path / "modular_model_index.json"
         assert index_file.exists()