[Modular] Avoid downloading weights when loading from an existing local path

docs/source/en/modular_diffusers/modular_pipeline.md

@@ -441,7 +441,7 @@ pipe.save_pretrained("local/path", repo_id="my-username/flux2-custom-transformer
 
 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.
+Moving the files any other way doesn't rewrite the index. A copy downloaded with `hf download --local-dir` still works: when a pipeline is loaded from a local directory, every component whose files are present in that directory is loaded from it instead of the recorded repository. A copy uploaded with `hf upload` keeps pointing at the old location, so 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:
 

src/diffusers/modular_pipelines/modular_pipeline.py

@@ -30,13 +30,23 @@
 from typing_extensions import Self
 
 from ..configuration_utils import ConfigMixin, FrozenDict
+from ..models.auto_model import AutoModel
+from ..models.modeling_utils import ModelMixin
 from ..pipelines.pipeline_loading_utils import (
     LOADABLE_CLASSES,
     _fetch_class_library_tuple,
     _unwrap_model,
+    filter_model_files,
     simple_get_class_obj,
 )
-from ..utils import PushToHubMixin, deprecate, is_accelerate_available, logging
+from ..utils import (
+    TRANSFORMERS_COMPONENT_AUX_FILES,
+    PushToHubMixin,
+    deprecate,
+    is_accelerate_available,
+    is_transformers_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
@@ -59,12 +69,46 @@
 )
 
 
+# classes whose components are loaded from weight files; a component without a type hint is loaded with `AutoModel`
+_MODEL_CLASSES = (ModelMixin, AutoModel)
+if is_transformers_available():
+    from transformers import PreTrainedModel
+
+    _MODEL_CLASSES = (*_MODEL_CLASSES, PreTrainedModel)
+
 if is_accelerate_available():
     import accelerate
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
+def _is_local_component(
+    pretrained_model_name_or_path: str | os.PathLike | None, component_spec: ComponentSpec
+) -> bool:
+    """
+    Whether the component's files are in `pretrained_model_name_or_path`, a local pipeline directory: weight files for
+    a model, the config file its class saves for a diffusers component without weights (schedulers, guiders, ...), one
+    of `TRANSFORMERS_COMPONENT_AUX_FILES` for a transformers one (tokenizers, processors, ...).
+    """
+    if pretrained_model_name_or_path is None:
+        return False
+    component_dir = os.path.join(pretrained_model_name_or_path, component_spec.subfolder or "")
+    if not os.path.isdir(component_dir):
+        return False
+    filenames = os.listdir(component_dir)
+
+    class_obj = component_spec.type_hint
+    is_model = class_obj is None or issubclass(class_obj, _MODEL_CLASSES)
+
+    if is_model:
+        return len(filter_model_files(filenames)) > 0
+
+    if issubclass(class_obj, ConfigMixin):
+        return class_obj.config_name in filenames
+
+    return any(filename in filenames for filename in TRANSFORMERS_COMPONENT_AUX_FILES)
+
+
 # map regular pipeline to modular pipeline class name
 
 
@@ -1765,6 +1809,11 @@ def __init__(
                     library, class_name, component_spec_dict = value
                     component_spec = self._dict_to_component_spec(name, component_spec_dict)
                     component_spec.default_creation_method = "from_pretrained"
+                    # a local copy of the repo (e.g. `hf download --local-dir`) keeps the original index, which
+                    # points at the Hub; load the components whose files are present locally from the copy
+                    if _is_local_component(pretrained_model_name_or_path, component_spec):
+                        component_spec.pretrained_model_name_or_path = pretrained_model_name_or_path
+                        component_spec.revision = None
                     self._component_specs[name] = component_spec
 
                 elif name in self._config_specs:

src/diffusers/pipelines/pipeline_loading_utils.py

@@ -65,19 +65,6 @@
 TRANSFORMERS_DUMMY_MODULES_FOLDER = "transformers.utils"
 CONNECTED_PIPES_KEYS = ["prior"]
 
-# Auxiliary (non-weight) files a transformers component saves next to its weights. Repos with a flat,
-# transformers-style layout host a component's files at the repo root instead of in a subfolder, where the
-# folder-based allow patterns of `DiffusionPipeline.download` would miss them. Root-hosted weights and
-# `config.json` are matched by their own patterns, so only these auxiliary filenames need listing.
-# Currently the set needed by DiffusionGemma — extend as new flat-layout pipelines require it.
-TRANSFORMERS_COMPONENT_AUX_FILES = [
-    "chat_template.jinja",
-    "generation_config.json",
-    "processor_config.json",
-    "tokenizer.json",
-    "tokenizer_config.json",
-]
-
 logger = logging.get_logger(__name__)
 
 LOADABLE_CLASSES = {

src/diffusers/pipelines/pipeline_utils.py

@@ -59,6 +59,7 @@
 from ..utils import (
     CONFIG_NAME,
     DEPRECATED_REVISION_ARGS,
+    TRANSFORMERS_COMPONENT_AUX_FILES,
     BaseOutput,
     PushToHubMixin,
     _get_detailed_type,
@@ -92,7 +93,6 @@
     CONNECTED_PIPES_KEYS,
     CUSTOM_PIPELINE_FILE_NAME,
     LOADABLE_CLASSES,
-    TRANSFORMERS_COMPONENT_AUX_FILES,
     _fetch_class_library_tuple,
     _get_custom_components_and_folders,
     _get_custom_pipeline_class,

src/diffusers/utils/__init__.py

@@ -36,6 +36,7 @@
     SAFE_WEIGHTS_INDEX_NAME,
     SAFETENSORS_FILE_EXTENSION,
     SAFETENSORS_WEIGHTS_NAME,
+    TRANSFORMERS_COMPONENT_AUX_FILES,
     USE_PEFT_BACKEND,
     WEIGHTS_INDEX_NAME,
     WEIGHTS_NAME,

src/diffusers/utils/constants.py

@@ -37,6 +37,17 @@
 FLASHPACK_FILE_EXTENSION = "flashpack"
 GGUF_FILE_EXTENSION = "gguf"
 ONNX_EXTERNAL_WEIGHTS_NAME = "weights.pb"
+# Auxiliary (non-weight) files a transformers component saves next to its weights, or as its only files for tokenizers
+# and processors. `DiffusionPipeline.download` uses them to fetch components hosted at the root of a flat,
+# transformers-style repo, and `ModularPipeline` to tell that such a component is present in a local directory.
+TRANSFORMERS_COMPONENT_AUX_FILES = [
+    "chat_template.jinja",
+    "generation_config.json",
+    "preprocessor_config.json",
+    "processor_config.json",
+    "tokenizer.json",
+    "tokenizer_config.json",
+]
 HUGGINGFACE_CO_RESOLVE_ENDPOINT = os.environ.get("HF_ENDPOINT", "https://huggingface.co")
 DIFFUSERS_DYNAMIC_MODULE_NAME = "diffusers_modules"
 HF_MODULES_CACHE = os.getenv("HF_MODULES_CACHE", os.path.join(HF_HOME, "modules"))

tests/modular_pipelines/test_modular_pipeline_loading.py

@@ -15,9 +15,11 @@
 
 import json
 import os
+import shutil
 
 import pytest
 import torch
+from huggingface_hub import snapshot_download
 
 from diffusers import AutoModel, ControlNetModel, ModularPipeline, UNet2DConditionModel
 from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec
@@ -276,3 +278,69 @@ def test_init_raises_without_resolvable_blocks(self):
         # The base class has no `default_blocks_name`, so with no `blocks` there is nothing to build from.
         with pytest.raises(ValueError, match="No pipeline blocks could be resolved"):
             ModularPipeline()
+
+
+class TestLoadFromLocalCopy:
+    def test_local_copy_loads_present_components_locally(self, tmp_path):
+        """`hf download --local-dir` keeps the index pointing at the Hub; components whose subfolder is present in
+        the local copy load from it, the rest keep their recorded spec."""
+        local_dir = str(tmp_path / "local-copy")
+        cache_dir = str(tmp_path / "cache")
+        snapshot_download("hf-internal-testing/tiny-anima-modular-pipe", local_dir=local_dir)
+
+        pipe = ModularPipeline.from_pretrained(local_dir)
+        for name in ("vae", "transformer", "text_encoder", "scheduler"):
+            spec = pipe._component_specs[name]
+            assert spec.pretrained_model_name_or_path == local_dir, f"{name} should load from the local copy"
+            assert spec.revision is None
+        assert (
+            pipe._component_specs["t5_tokenizer"].pretrained_model_name_or_path == "hf-internal-testing/tiny-random-t5"
+        )
+
+        pipe.load_components(names=["vae"], dtype=torch.float32, local_files_only=True, cache_dir=cache_dir)
+        assert pipe.vae is not None
+        cached_weights = [p for p in (tmp_path / "cache").rglob("*") if p.suffix in (".safetensors", ".bin")]
+        assert cached_weights == [], f"weights should not be in the Hub cache: {cached_weights}"
+
+    def test_local_copy_missing_files_keeps_recorded_spec(self, tmp_path):
+        """A missing subfolder, or a model subfolder without weight files (e.g. a partial download), keeps the
+        recorded spec instead of shadowing it with an unloadable folder."""
+        local_dir = str(tmp_path / "local-copy")
+        snapshot_download("hf-internal-testing/tiny-anima-modular-pipe", local_dir=local_dir)
+        shutil.rmtree(os.path.join(local_dir, "transformer"))
+        for filename in os.listdir(os.path.join(local_dir, "vae")):
+            if filename.endswith((".safetensors", ".bin")):
+                os.remove(os.path.join(local_dir, "vae", filename))
+
+        pipe = ModularPipeline.from_pretrained(local_dir)
+        assert (
+            pipe._component_specs["transformer"].pretrained_model_name_or_path
+            == "hf-internal-testing/tiny-anima-modular-pipe"
+        )
+        assert (
+            pipe._component_specs["vae"].pretrained_model_name_or_path == "hf-internal-testing/tiny-anima-modular-pipe"
+        )
+        assert pipe._component_specs["text_encoder"].pretrained_model_name_or_path == local_dir
+
+    def test_local_copy_loads_components_at_root(self, tmp_path):
+        """A component recorded without a subfolder is at the root of its repo; when the local copy has its files
+        there it is loaded from the copy: weights for a model, the saved config file for anything else."""
+        local_dir = str(tmp_path / "local-copy")
+        snapshot_download("hf-internal-testing/tiny-cosmos3-modular-pipe", local_dir=local_dir)
+        index_path = os.path.join(local_dir, "modular_model_index.json")
+        with open(index_path) as f:
+            index = json.load(f)
+        root_components = ["transformer", "scheduler", "text_tokenizer"]
+        for name in root_components:
+            for filename in os.listdir(os.path.join(local_dir, name)):
+                shutil.move(os.path.join(local_dir, name, filename), os.path.join(local_dir, filename))
+            index[name][2]["subfolder"] = None
+        with open(index_path, "w") as f:
+            json.dump(index, f)
+
+        pipe = ModularPipeline.from_pretrained(local_dir)
+        for name in root_components:
+            assert pipe._component_specs[name].pretrained_model_name_or_path == local_dir, f"{name} not local"
+        pipe.load_components(names=root_components, dtype=torch.float32, local_files_only=True)
+        for name in root_components:
+            assert getattr(pipe, name) is not None, f"{name} did not load from the local copy"