Load local ModularPipeline snapshots when modular_model_index.json still names Hub ids

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"`.
 

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.")
 

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