diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index 1863bbda87d3..6237407a8c5b 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/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: diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index e9e5463c1e72..a071cfc7e507 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/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: diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index 69bce1a1c533..6958f49c8ddd 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/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 = { diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 82e6c4c2aff4..be0419da9eac 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/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, diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 5c63a4bc7661..b3051dfcb9d1 100644 --- a/src/diffusers/utils/__init__.py +++ b/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, diff --git a/src/diffusers/utils/constants.py b/src/diffusers/utils/constants.py index fcf0e4518800..d17587c65bdc 100644 --- a/src/diffusers/utils/constants.py +++ b/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")) diff --git a/tests/modular_pipelines/test_modular_pipeline_loading.py b/tests/modular_pipelines/test_modular_pipeline_loading.py index 4e4797106fe4..4871eab4e67d 100644 --- a/tests/modular_pipelines/test_modular_pipeline_loading.py +++ b/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"