diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index 07dc30b078ae..1863bbda87d3 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/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,34 @@ 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. + +```py +pipe = ModularPipeline.from_pretrained("MiniMaxAI/MiniMax-H3") +pipe.load_components() +pipe.save_pretrained("path/to/local-copy") +``` + +The destination recorded in `modular_model_index.json` is wherever [`~ModularPipeline.save_pretrained`] writes: the `save_directory` for a local save, or the `repo_id` when you pass `push_to_hub=True`. + +### Keep references to existing components + +Load only what's new (or nothing at all). Only loaded components are saved; 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 keep the loading specs in `modular_model_index.json` as they are. A saved component whose loading spec is empty is still filled in with the destination, since there is nothing to preserve. + +Note that moving the files any other way (uploading with `hf upload`, downloading a repository with `hf download --local-dir`) doesn't rewrite the index, so the copy still points to the old location; update the index manually in that case. + 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: ``` diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index f2576b99328d..3bdc9770e1d2 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/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,9 @@ 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 + # a revision pinned for the original source doesn't exist at the destination + component_spec_dict["revision"] = None self.register_to_config(**{component_name: (library, class_name, component_spec_dict)}) self.save_config(save_directory=save_directory) diff --git a/tests/modular_pipelines/test_modular_pipeline_loading.py b/tests/modular_pipelines/test_modular_pipeline_loading.py index 3b9ebcc1cdf3..16d78d797707 100644 --- a/tests/modular_pipelines/test_modular_pipeline_loading.py +++ b/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") diff --git a/tests/modular_pipelines/testing_utils/loading.py b/tests/modular_pipelines/testing_utils/loading.py index 64a9dd345008..63b1872f07eb 100644 --- a/tests/modular_pipelines/testing_utils/loading.py +++ b/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()