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