[Modular] Avoid downloading weights when loading from an existing local path by DN6 · Pull Request #14797 · huggingface/diffusers · GitHub
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/source/en/modular_diffusers/modular_pipeline.md
51 changes: 50 additions & 1 deletion src/diffusers/modular_pipelines/modular_pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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


Expand Down Expand Up @@ -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 = str(pretrained_model_name_or_path)
component_spec.revision = None
self._component_specs[name] = component_spec

elif name in self._config_specs:
Expand Down
13 changes: 0 additions & 13 deletions src/diffusers/pipelines/pipeline_loading_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/pipelines/pipeline_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@
from ..utils import (
CONFIG_NAME,
DEPRECATED_REVISION_ARGS,
TRANSFORMERS_COMPONENT_AUX_FILES,
BaseOutput,
PushToHubMixin,
_get_detailed_type,
Expand Down Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions src/diffusers/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
11 changes: 11 additions & 0 deletions src/diffusers/utils/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"))
Expand Down
68 changes: 68 additions & 0 deletions tests/modular_pipelines/test_modular_pipeline_loading.py
Loading