Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 21 additions & 1 deletion docs/source/en/modular_diffusers/modular_pipeline.md
Original file line number Diff line number Diff line change
Expand Up @@ -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"`.

Expand Down Expand Up @@ -415,6 +415,26 @@ 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`] saves the components that are currently loaded, and points each saved component's loading spec in `modular_model_index.json` at the destination — the `repo_id` when pushing to the Hub, otherwise the save directory. A component that isn't loaded isn't saved and keeps its recorded spec, so it is still fetched from its original location later. This gives you two ways to save, depending on what you want:

- **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. (A raw download like `hf download ... --local-dir` doesn't do this — the published specs still point at the Hub.)

```py
pipe = ModularPipeline.from_pretrained("MiniMaxAI/MiniMax-H3")
pipe.load_components()
pipe.save_pretrained("path/to/local-copy")
```

- **Reuse existing components without saving the weights again** — load only what's new (or nothing at all). Only loaded components are written; everything else stays a pointer to its original repository. For example, to share a single custom transformer while the other components keep loading from the base repo — the same shape as the quantized-transformer repository above:

```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 also preserve the recorded loading specs of the components being saved.

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:

```
Expand Down
27 changes: 21 additions & 6 deletions src/diffusers/modular_pipelines/modular_pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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*):
Expand All @@ -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:
Expand Down Expand Up @@ -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)
Expand Down
43 changes: 36 additions & 7 deletions tests/modular_pipelines/test_modular_pipeline_loading.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import json
import os

import pytest
import torch

from diffusers import AutoModel, ControlNetModel, ModularPipeline, UNet2DConditionModel
Expand Down Expand Up @@ -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)
Expand All @@ -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."""
Expand All @@ -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)

Expand All @@ -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)

Expand All @@ -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")
Expand Down
2 changes: 1 addition & 1 deletion tests/modular_pipelines/testing_utils/loading.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
Loading