From e2dcf933be2279cddda547bf8253b02b3cec2707 Mon Sep 17 00:00:00 2001 From: el Date: Sat, 29 Aug 2026 05:27:42 +0530 Subject: [PATCH] [modular] Prefer local snapshot files over Hub ids in modular_model_index.json When ModularPipeline is initialized from a local directory, component specs copied from modular_model_index.json still pointed at Hub repo ids. load_components() then re-downloaded (or failed offline) even if the snapshot already contained the subfolders. Bind those specs to the local directory when the subfolder exists, and leave Hub ids in place for pointer repos and pruned snapshots. --- .../en/modular_diffusers/modular_pipeline.md | 2 +- .../modular_pipelines/modular_pipeline.py | 19 +++ .../test_modular_pipeline_loading.py | 115 +++++++++++++++++- 3 files changed, 133 insertions(+), 3 deletions(-) diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index 07dc30b078ae..c765cbebd79c 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`). If that path is a local directory, components whose spec subfolder exists on disk are loaded from that directory even when `modular_model_index.json` still names a Hub id. Components missing from the directory keep the Hub id, so pointer repositories and pruned snapshots can still fetch remotely. In the example below, the `pretrained_model_name_or_path` will be updated to `"stabilityai/stable-diffusion-xl-base-1.0"`. diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index f2576b99328d..a0532a302fbf 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -1702,6 +1702,9 @@ def __init__( config values, which will be saved as `modular_model_index.json` during `save_pretrained` - The pipeline's config dict is also used to store the pipeline blocks's class name, which will be saved as `_blocks_class_name` in the config dict + - If `pretrained_model_name_or_path` is a local directory, Hub ids in `modular_model_index.json` are + rewritten to that directory for components whose spec subfolder exists locally. Missing subfolders keep + the Hub id so pointer repositories and pruned snapshots can still load remotely. """ if modular_config_dict is None and config_dict is None and pretrained_model_name_or_path is not None: @@ -1784,6 +1787,22 @@ def __init__( elif name in self._config_specs: self._config_specs[name].default = value + # `modular_model_index.json` stores Hub ids even in a fully downloaded snapshot. When the caller passed a + # local directory, bind each from_pretrained spec to that directory if its subfolder is present. Specs whose + # files are missing keep their Hub id so pointer repositories and pruned snapshots still fetch remotely. + if pretrained_model_name_or_path is not None and os.path.isdir(pretrained_model_name_or_path): + for component_spec in self._component_specs.values(): + if component_spec.default_creation_method != "from_pretrained": + continue + spec_path = component_spec.pretrained_model_name_or_path + if not isinstance(spec_path, str) or os.path.isdir(spec_path): + continue + subfolder = component_spec.subfolder + if not subfolder: + continue + if os.path.isdir(os.path.join(pretrained_model_name_or_path, subfolder)): + component_spec.pretrained_model_name_or_path = pretrained_model_name_or_path + if len(kwargs) > 0: logger.warning(f"Unexpected input '{kwargs.keys()}' provided. This input will be ignored.") diff --git a/tests/modular_pipelines/test_modular_pipeline_loading.py b/tests/modular_pipelines/test_modular_pipeline_loading.py index 3b9ebcc1cdf3..e1f6d30e53b2 100644 --- a/tests/modular_pipelines/test_modular_pipeline_loading.py +++ b/tests/modular_pipelines/test_modular_pipeline_loading.py @@ -18,8 +18,8 @@ import torch -from diffusers import AutoModel, ControlNetModel, ModularPipeline, UNet2DConditionModel -from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec +from diffusers import AutoModel, ControlNetModel, DDIMScheduler, ModularPipeline, UNet2DConditionModel +from diffusers.modular_pipelines import ComponentSpec, ModularPipelineBlocks class TestAutoModelLoadIdTagging: @@ -239,3 +239,114 @@ def test_init_fallback_when_blocks_class_name_is_base_class(self, tmp_path): assert loaded_pipe.__class__.__name__ == pipe.__class__.__name__ assert loaded_pipe._blocks.__class__.__name__ == pipe._blocks.__class__.__name__ assert len(loaded_pipe._blocks.sub_blocks) == len(pipe._blocks.sub_blocks) + + +_MISSING_HUB_ID = "org/this-repo-does-not-exist-14640" + + +class _LocalSnapshotBlocks(ModularPipelineBlocks): + def __init__(self, component_names=("scheduler",)): + self._component_names = component_names + super().__init__() + + @property + def expected_components(self): + type_hints = {"scheduler": DDIMScheduler, "unet": UNet2DConditionModel} + return [ComponentSpec(name, type_hints[name]) for name in self._component_names] + + +def _write_modular_index(snapshot_dir, components, repo_ids=None): + index = { + "_class_name": "ModularPipeline", + "_diffusers_version": "0.40.0.dev0", + "_blocks_class_name": "ModularPipelineBlocks", + } + for name, class_name in components.items(): + repo_id = (repo_ids or {}).get(name, _MISSING_HUB_ID) + index[name] = [ + "diffusers", + class_name, + { + "type_hint": ["diffusers", class_name], + "pretrained_model_name_or_path": repo_id, + "subfolder": name, + "variant": None, + "revision": None, + }, + ] + with open(os.path.join(snapshot_dir, "modular_model_index.json"), "w") as f: + json.dump(index, f) + + +class TestLocalModularSnapshotLoading: + def test_init_pipeline_rewrites_hub_ids_when_subfolder_exists(self, tmp_path): + snapshot_dir = str(tmp_path / "snapshot") + os.makedirs(snapshot_dir) + DDIMScheduler().save_pretrained(os.path.join(snapshot_dir, "scheduler")) + _write_modular_index(snapshot_dir, {"scheduler": "DDIMScheduler"}) + + pipe = _LocalSnapshotBlocks().init_pipeline(snapshot_dir) + + assert pipe._component_specs["scheduler"].pretrained_model_name_or_path == snapshot_dir + + pipe.load_components(names="scheduler", local_files_only=True) + assert pipe.scheduler is not None + assert isinstance(pipe.scheduler, DDIMScheduler) + + def test_constructor_rewrites_hub_ids_when_subfolder_exists(self, tmp_path): + snapshot_dir = str(tmp_path / "snapshot") + os.makedirs(snapshot_dir) + DDIMScheduler().save_pretrained(os.path.join(snapshot_dir, "scheduler")) + _write_modular_index(snapshot_dir, {"scheduler": "DDIMScheduler"}) + + pipe = ModularPipeline( + blocks=_LocalSnapshotBlocks(), + pretrained_model_name_or_path=snapshot_dir, + local_files_only=True, + ) + + assert pipe._component_specs["scheduler"].pretrained_model_name_or_path == snapshot_dir + pipe.load_components(names="scheduler", local_files_only=True) + assert pipe.scheduler is not None + + def test_missing_local_subfolder_keeps_hub_id(self, tmp_path): + snapshot_dir = str(tmp_path / "snapshot") + os.makedirs(snapshot_dir) + _write_modular_index(snapshot_dir, {"scheduler": "DDIMScheduler"}) + + pipe = _LocalSnapshotBlocks().init_pipeline(snapshot_dir) + + assert pipe._component_specs["scheduler"].pretrained_model_name_or_path == _MISSING_HUB_ID + + def test_mixed_snapshot_rewrites_only_present_components(self, tmp_path): + snapshot_dir = str(tmp_path / "snapshot") + os.makedirs(snapshot_dir) + DDIMScheduler().save_pretrained(os.path.join(snapshot_dir, "scheduler")) + _write_modular_index(snapshot_dir, {"scheduler": "DDIMScheduler", "unet": "UNet2DConditionModel"}) + + pipe = _LocalSnapshotBlocks(component_names=("scheduler", "unet")).init_pipeline(snapshot_dir) + + assert pipe._component_specs["scheduler"].pretrained_model_name_or_path == snapshot_dir + assert pipe._component_specs["unet"].pretrained_model_name_or_path == _MISSING_HUB_ID + + pipe.load_components(names="scheduler", local_files_only=True) + assert pipe.scheduler is not None + assert pipe.unet is None + + def test_existing_local_spec_path_is_not_overwritten(self, tmp_path): + other_dir = str(tmp_path / "other") + snapshot_dir = str(tmp_path / "snapshot") + os.makedirs(snapshot_dir) + DDIMScheduler(num_train_timesteps=50).save_pretrained(os.path.join(other_dir, "scheduler")) + DDIMScheduler(num_train_timesteps=1000).save_pretrained(os.path.join(snapshot_dir, "scheduler")) + _write_modular_index( + snapshot_dir, + {"scheduler": "DDIMScheduler"}, + repo_ids={"scheduler": other_dir}, + ) + + pipe = _LocalSnapshotBlocks().init_pipeline(snapshot_dir) + + assert pipe._component_specs["scheduler"].pretrained_model_name_or_path == other_dir + pipe.load_components(names="scheduler", local_files_only=True) + assert pipe.scheduler.config.num_train_timesteps == 50