diff --git a/setup.py b/setup.py index d5ab0547ca12..60fde132fc6d 100644 --- a/setup.py +++ b/setup.py @@ -101,7 +101,7 @@ "ftfy", "hf-doc-builder>=0.3.0", "httpx<1.0.0", - "huggingface-hub>=1.23.0,<2.0", + "huggingface-hub>=1.26.0,<2.0", "requests-mock==1.10.0", "importlib_metadata", "invisible-watermark>=0.2.0", diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index 02e304d2ab02..5aa428a012ff 100644 --- a/src/diffusers/dependency_versions_table.py +++ b/src/diffusers/dependency_versions_table.py @@ -9,7 +9,7 @@ "ftfy": "ftfy", "hf-doc-builder": "hf-doc-builder>=0.3.0", "httpx": "httpx<1.0.0", - "huggingface-hub": "huggingface-hub>=1.23.0,<2.0", + "huggingface-hub": "huggingface-hub>=1.26.0,<2.0", "requests-mock": "requests-mock==1.10.0", "importlib_metadata": "importlib_metadata", "invisible-watermark": "invisible-watermark>=0.2.0", diff --git a/src/diffusers/models/auto_model.py b/src/diffusers/models/auto_model.py index 336650ef5703..1ce21e261aca 100644 --- a/src/diffusers/models/auto_model.py +++ b/src/diffusers/models/auto_model.py @@ -17,7 +17,7 @@ from huggingface_hub.utils import validate_hf_hub_args from ..configuration_utils import ConfigMixin -from ..utils import DIFFUSERS_LOAD_ID_FIELDS, logging +from ..utils import DIFFUSERS_LOAD_ID_FIELDS, _resolve_revision, logging from ..utils.dynamic_modules_utils import get_class_from_dynamic_module, resolve_trust_remote_code @@ -266,6 +266,16 @@ def from_pretrained(cls, pretrained_model_or_path: str | os.PathLike | None = No ] hub_kwargs = {name: kwargs.pop(name, None) for name in hub_kwargs_names} + # Resolve the revision only once + revision = hub_kwargs["revision"] + hub_kwargs["revision"] = _resolve_revision( + pretrained_model_or_path, + revision=revision, + cache_dir=hub_kwargs["cache_dir"], + local_files_only=hub_kwargs["local_files_only"], + token=hub_kwargs["token"], + ) + # load_config_kwargs uses the same hub kwargs minus subfolder and resume_download load_config_kwargs = {k: v for k, v in hub_kwargs.items() if k not in ["subfolder"]} @@ -337,7 +347,8 @@ def from_pretrained(cls, pretrained_model_or_path: str | os.PathLike | None = No kwargs = {**load_config_kwargs, **kwargs} model = model_cls.from_pretrained(pretrained_model_or_path, **kwargs) - load_id_kwargs = {"pretrained_model_name_or_path": pretrained_model_or_path, **kwargs} + # the load id records the revision the user asked for, not the commit it was resolved to + load_id_kwargs = {"pretrained_model_name_or_path": pretrained_model_or_path, **kwargs, "revision": revision} parts = [load_id_kwargs.get(field, "null") for field in DIFFUSERS_LOAD_ID_FIELDS] load_id = "|".join("null" if p is None else p for p in parts) model._diffusers_load_id = load_id diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index f57c92beadb0..425f2f29235e 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -51,6 +51,7 @@ _add_variant, _get_checkpoint_shard_files, _get_model_file, + _resolve_revision, deprecate, is_accelerate_available, is_bitsandbytes_available, @@ -1127,6 +1128,15 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None } unused_kwargs = {} + # Resolve the revision only once + revision = _resolve_revision( + pretrained_model_name_or_path, + revision=revision, + cache_dir=cache_dir, + local_files_only=local_files_only, + token=token, + ) + # Load config if we don't provide a configuration config_path = pretrained_model_name_or_path diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 3aa2c854dfe6..f2576b99328d 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -38,7 +38,7 @@ ) from ..utils import PushToHubMixin, is_accelerate_available, logging from ..utils.dynamic_modules_utils import get_class_from_dynamic_module, resolve_trust_remote_code -from ..utils.hub_utils import load_or_create_model_card, populate_model_card +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 from .components_manager import ComponentsManager from .modular_pipeline_utils import ( @@ -439,6 +439,15 @@ def from_pretrained( ] hub_kwargs = {name: kwargs.pop(name) for name in hub_kwargs_names if name in kwargs} + # Resolve the revision only once + hub_kwargs["revision"] = _resolve_revision( + pretrained_model_name_or_path, + revision=hub_kwargs.get("revision"), + cache_dir=hub_kwargs.get("cache_dir"), + local_files_only=hub_kwargs.get("local_files_only"), + token=hub_kwargs.get("token"), + ) + config = cls.load_config(pretrained_model_name_or_path, **hub_kwargs) has_remote_code = "auto_map" in config and cls.__name__ in config["auto_map"] trust_remote_code = resolve_trust_remote_code( @@ -1872,6 +1881,15 @@ def from_pretrained( """ from ..pipelines.pipeline_loading_utils import _get_pipeline_class + # Resolve the revision only once + kwargs["revision"] = _resolve_revision( + pretrained_model_name_or_path, + revision=kwargs.get("revision"), + cache_dir=kwargs.get("cache_dir"), + local_files_only=kwargs.get("local_files_only"), + token=kwargs.get("token"), + ) + try: blocks = ModularPipelineBlocks.from_pretrained( pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs diff --git a/src/diffusers/pipelines/auto_pipeline.py b/src/diffusers/pipelines/auto_pipeline.py index 95d8f7ca210f..20ab34372808 100644 --- a/src/diffusers/pipelines/auto_pipeline.py +++ b/src/diffusers/pipelines/auto_pipeline.py @@ -19,7 +19,7 @@ from ..configuration_utils import ConfigMixin from ..models.controlnets import ControlNetUnionModel -from ..utils import is_sentencepiece_available +from ..utils import _resolve_revision, is_sentencepiece_available from .anyflow import AnyFlowFARPipeline, AnyFlowPipeline from .audioldm2 import AudioLDM2Pipeline from .aura_flow import AuraFlowPipeline @@ -511,6 +511,15 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs): local_files_only = kwargs.pop("local_files_only", False) revision = kwargs.pop("revision", None) + # Resolve the revision only once + revision = _resolve_revision( + pretrained_model_or_path, + revision=revision, + cache_dir=cache_dir, + local_files_only=local_files_only, + token=token, + ) + load_config_kwargs = { "cache_dir": cache_dir, "force_download": force_download, @@ -801,6 +810,15 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs): local_files_only = kwargs.pop("local_files_only", False) revision = kwargs.pop("revision", None) + # Resolve the revision only once + revision = _resolve_revision( + pretrained_model_or_path, + revision=revision, + cache_dir=cache_dir, + local_files_only=local_files_only, + token=token, + ) + load_config_kwargs = { "cache_dir": cache_dir, "force_download": force_download, @@ -1105,6 +1123,15 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs): local_files_only = kwargs.pop("local_files_only", False) revision = kwargs.pop("revision", None) + # Resolve the revision only once + revision = _resolve_revision( + pretrained_model_or_path, + revision=revision, + cache_dir=cache_dir, + local_files_only=local_files_only, + token=token, + ) + load_config_kwargs = { "cache_dir": cache_dir, "force_download": force_download, @@ -1406,6 +1433,15 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs): local_files_only = kwargs.pop("local_files_only", False) revision = kwargs.pop("revision", None) + # Resolve the revision only once + revision = _resolve_revision( + pretrained_model_or_path, + revision=revision, + cache_dir=cache_dir, + local_files_only=local_files_only, + token=token, + ) + load_config_kwargs = { "cache_dir": cache_dir, "force_download": force_download, diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index cf303bb6a5ad..8986553eda3d 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -75,7 +75,12 @@ numpy_to_pil, ) from ..utils.distributed_utils import is_torch_dist_rank_zero -from ..utils.hub_utils import _check_legacy_sharding_variant_format, load_or_create_model_card, populate_model_card +from ..utils.hub_utils import ( + _check_legacy_sharding_variant_format, + _resolve_revision, + load_or_create_model_card, + populate_model_card, +) from ..utils.torch_utils import empty_device_cache, get_device, is_compiled_module @@ -1593,6 +1598,15 @@ def download(cls, pretrained_model_name, **kwargs) -> str | os.PathLike: if "dduf_file" in kwargs: raise ValueError(_DDUF_REMOVAL_MESSAGE) + # Resolve the revision only once + revision = _resolve_revision( + pretrained_model_name, + revision=revision, + cache_dir=cache_dir, + local_files_only=local_files_only, + token=token, + ) + allow_pickle = True if (use_safetensors is None or use_safetensors is False) else False use_safetensors = use_safetensors if use_safetensors is not None else True diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 0554d341022a..5c63a4bc7661 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -49,6 +49,7 @@ _add_variant, _get_checkpoint_shard_files, _get_model_file, + _resolve_revision, extract_commit_hash, http_user_agent, ) diff --git a/src/diffusers/utils/hub_utils.py b/src/diffusers/utils/hub_utils.py index d0bd2e68fb9f..fe128be7dad5 100644 --- a/src/diffusers/utils/hub_utils.py +++ b/src/diffusers/utils/hub_utils.py @@ -23,20 +23,24 @@ from pathlib import Path from uuid import uuid4 +import httpx from huggingface_hub import ( ModelCard, ModelCardData, create_repo, hf_hub_download, model_info, + resolve_revision, snapshot_download, upload_folder, ) from huggingface_hub.constants import HF_HUB_DISABLE_TELEMETRY, HF_HUB_OFFLINE +from huggingface_hub.errors import RevisionResolutionError from huggingface_hub.file_download import REGEX_COMMIT_HASH from huggingface_hub.utils import ( EntryNotFoundError, HfHubHTTPError, + HFValidationError, RepositoryNotFoundError, RevisionNotFoundError, is_jinja_available, @@ -208,6 +212,41 @@ def extract_commit_hash(resolved_file: str | None, commit_hash: str | None = Non return commit_hash if REGEX_COMMIT_HASH.match(commit_hash) else None +def _resolve_revision( + pretrained_model_name_or_path: str | os.PathLike | None, + *, + revision: str | None = None, + cache_dir: str | os.PathLike | None = None, + local_files_only: bool | None = None, + token: str | bool | None = None, +) -> str | None: + """ + Resolves `revision` to a commit hash, to be called once at the beginning of a loading method. + + Loading a model or a pipeline fetches several files from the same repo (config, weight index, shards, custom code, + ...). Passing the returned [`~huggingface_hub.ResolvedRevision`] down to every download pins them all to the same + commit - even if the repo is updated in the meantime - and lets `huggingface_hub` serve them from the cache without + resolving `revision` again on each call. + + Resolution is best-effort: local folders are returned untouched and, if the Hub cannot answer (repo or revision not + found, offline mode with nothing cached, ...), `revision` is returned as is so that the download that follows fails + with its usual error message. + """ + if pretrained_model_name_or_path is None or os.path.isdir(pretrained_model_name_or_path): + return revision + + try: + return resolve_revision( + str(pretrained_model_name_or_path), + revision=revision, + cache_dir=cache_dir, + local_files_only=bool(local_files_only), + token=token, + ) + except (HfHubHTTPError, RevisionResolutionError, HFValidationError, httpx.TransportError): + return revision + + def _add_variant(weights_name: str, variant: str | None = None) -> str: if variant is not None: splits = weights_name.split(".") diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index 9968add19dd9..b1bdaabbad7a 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -119,12 +119,10 @@ def test_cached_files_are_used_when_no_internet(self): def test_local_files_only_with_sharded_checkpoint(self): repo_id = "hf-internal-testing/tiny-flux-sharded" - error_response = mock.Mock( - status_code=500, - headers={}, - raise_for_status=mock.Mock(side_effect=HfHubHTTPError("Server down", response=mock.Mock())), - json=mock.Mock(return_value={}), - ) + error_response = mock.Mock(status_code=500, headers={}, json=mock.Mock(return_value={})) + # `resolve_revision` inspects `error.response.status_code` to tell a Hub outage from a definitive answer, + # so the raised error has to carry the response itself. + error_response.raise_for_status = mock.Mock(side_effect=HfHubHTTPError("Server down", response=error_response)) client_mock = mock.Mock() client_mock.get.return_value = error_response diff --git a/tests/pipelines/test_pipelines_auto.py b/tests/pipelines/test_pipelines_auto.py index ae01f1465786..ce3d1288d95f 100644 --- a/tests/pipelines/test_pipelines_auto.py +++ b/tests/pipelines/test_pipelines_auto.py @@ -112,9 +112,10 @@ def test_kwargs_local_files_only(self): tmpdirname = DiffusionPipeline.download(repo) tmpdirname = Path(tmpdirname) - # edit commit_id to so that it's not the latest commit + # edit commit_id to so that it's not the latest commit. It has to stay a syntactically valid commit hash: + # `refs/main` is read back by `resolve_revision` and passed around as a commit hash. commit_id = tmpdirname.name - new_commit_id = commit_id + "hug" + new_commit_id = "0" * len(commit_id) ref_dir = tmpdirname.parent.parent / "refs/main" with open(ref_dir, "w") as f: