compute the same snapshot_download patterns offline as online

Instead of returning the cached snapshot folder without validation when
offline, compute the same allow/ignore patterns as the online path
(sourcing the file listing from the cached snapshot instead of
model_info) and let snapshot_download validate the cached snapshot
against them, so an interrupted download surfaces hfh's
IncompleteSnapshotError instead of failing later at model load time.
Also revert the _get_checkpoint_shard_files divergence: its patterns
were already identical in both modes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: sayakpaul

examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py

@@ -85,11 +85,9 @@ def log_validation(pipeline, args, accelerator, generator, global_step, is_final
         os.makedirs(val_save_dir)
 
     original_image = (
-        lambda image_url_or_path: (
-            load_image(image_url_or_path)
-            if urlparse(image_url_or_path).scheme
-            else Image.open(image_url_or_path).convert("RGB")
-        )
+        lambda image_url_or_path: load_image(image_url_or_path)
+        if urlparse(image_url_or_path).scheme
+        else Image.open(image_url_or_path).convert("RGB")
     )(args.val_image_url_or_path)
 
     if torch.backends.mps.is_available():

src/diffusers/loaders/peft.py

@@ -46,7 +46,7 @@
 logger = logging.get_logger(__name__)
 
 _SET_ADAPTER_SCALE_FN_MAPPING = defaultdict(
-    lambda: lambda model_cls, weights: weights,
+    lambda: (lambda model_cls, weights: weights),
     {
         "UNet2DConditionModel": _maybe_expand_lora_scales,
         "UNetMotionModel": _maybe_expand_lora_scales,

src/diffusers/pipelines/pipeline_utils.py

@@ -1662,10 +1662,37 @@ def download(cls, pretrained_model_name, **kwargs) -> str | os.PathLike:
                 force_download=force_download,
                 token=token,
             )
+            filenames = {sibling.rfilename for sibling in info.siblings}
+        else:
+            # We are offline (either the user asked for `local_files_only` or the `model_info` call above
+            # failed), so the repo's file listing needed to compute the allow/ignore patterns below cannot be
+            # fetched. Source the listing from the cached snapshot instead: it holds exactly the (already
+            # filtered) files a previous online call downloaded, so both code paths compute the same patterns
+            # and `snapshot_download` validates the cached snapshot against them (catching e.g. a download
+            # that was interrupted midway). If the config isn't cached, skip straight to `snapshot_download`
+            # below, which raises the appropriate offline not-found error.
+            try:
+                config_file = hf_hub_download(
+                    pretrained_model_name,
+                    cls.config_name,
+                    cache_dir=cache_dir,
+                    revision=revision,
+                    token=token,
+                    local_files_only=True,
+                )
+            except LocalEntryNotFoundError:
+                config_file = None
+
+        if config_file is not None:
+            snapshot_folder = Path(config_file).parent
+            if local_files_only:
+                filenames = {
+                    f.relative_to(snapshot_folder).as_posix() for f in snapshot_folder.rglob("*") if f.is_file()
+                }
+
             config_dict = cls._dict_from_json_file(config_file)
             ignore_filenames = config_dict.pop("_ignore_files", [])
 
-            filenames = {sibling.rfilename for sibling in info.siblings}
             if variant is not None and _check_legacy_sharding_variant_format(filenames=filenames, variant=variant):
                 warn_msg = (
                     f"Warning: The repository contains sharded checkpoints for variant '{variant}' maybe in a deprecated format. "
@@ -1680,9 +1707,11 @@ def download(cls, pretrained_model_name, **kwargs) -> str | os.PathLike:
                 logger.warning(warn_msg)
 
             filenames = set(filenames) - set(ignore_filenames)
-            if revision in DEPRECATED_REVISION_ARGS and version.parse(
-                version.parse(__version__).base_version
-            ) >= version.parse("0.22.0"):
+            if (
+                not local_files_only
+                and revision in DEPRECATED_REVISION_ARGS
+                and version.parse(version.parse(__version__).base_version) >= version.parse("0.22.0")
+            ):
                 warn_deprecated_model_variant(pretrained_model_name, token, variant, revision, filenames)
 
             custom_components, folder_names = _get_custom_components_and_folders(
@@ -1767,40 +1796,20 @@ def download(cls, pretrained_model_name, **kwargs) -> str | os.PathLike:
 
             # Don't download index files of forbidden patterns either
             ignore_patterns = ignore_patterns + [f"{i}.index.*json" for i in ignore_patterns]
-            re_ignore_pattern = [re.compile(fnmatch.translate(p)) for p in ignore_patterns]
-            re_allow_pattern = [re.compile(fnmatch.translate(p)) for p in allow_patterns]
 
-            expected_files = [f for f in filenames if not any(p.match(f) for p in re_ignore_pattern)]
-            expected_files = [f for f in expected_files if any(p.match(f) for p in re_allow_pattern)]
+            if not local_files_only:
+                re_ignore_pattern = [re.compile(fnmatch.translate(p)) for p in ignore_patterns]
+                re_allow_pattern = [re.compile(fnmatch.translate(p)) for p in allow_patterns]
 
-            snapshot_folder = Path(config_file).parent
-            pipeline_is_cached = all((snapshot_folder / f).is_file() for f in expected_files)
+                expected_files = [f for f in filenames if not any(p.match(f) for p in re_ignore_pattern)]
+                expected_files = [f for f in expected_files if any(p.match(f) for p in re_allow_pattern)]
 
-            if pipeline_is_cached and not force_download:
-                # if the pipeline is cached, we can directly return it
-                # else call snapshot_download
-                return snapshot_folder
+                pipeline_is_cached = all((snapshot_folder / f).is_file() for f in expected_files)
 
-        else:
-            # We are offline (either the user asked for `local_files_only` or `model_info` failed above).
-            # diffusers only ever caches a filtered subset of a repo (skipping e.g. `.gitattributes`), but
-            # newer `huggingface_hub` versions make `snapshot_download` validate that the cached snapshot is
-            # complete and raise for the missing files. The resolved config file lives inside the cached
-            # snapshot folder, so return that folder directly. If the config isn't cached, fall through to
-            # `snapshot_download` below, which raises the not-found error handled together with a failed
-            # `model_info` call.
-            try:
-                config_file = hf_hub_download(
-                    pretrained_model_name,
-                    cls.config_name,
-                    cache_dir=cache_dir,
-                    revision=revision,
-                    token=token,
-                    local_files_only=True,
-                )
-                return Path(config_file).parent
-            except LocalEntryNotFoundError:
-                pass
+                if pipeline_is_cached and not force_download:
+                    # if the pipeline is cached, we can directly return it
+                    # else call snapshot_download
+                    return snapshot_folder
 
         user_agent = {"pipeline_class": cls.__name__}
         if custom_pipeline is not None and not custom_pipeline.endswith(".py"):
@@ -1843,7 +1852,7 @@ def download(cls, pretrained_model_name, **kwargs) -> str | os.PathLike:
 
             return cached_folder
 
-        except FileNotFoundError:
+        except FileNotFoundError as e:
             # Means we tried to load pipeline with `local_files_only=True` but the files have not been found in local cache.
             # This can happen in two cases:
             # 1. If the user passed `local_files_only=True`                    => we raise the error directly
@@ -1854,9 +1863,9 @@ def download(cls, pretrained_model_name, **kwargs) -> str | os.PathLike:
             else:
                 # 2. we forced `local_files_only=True` when `model_info` failed
                 raise EnvironmentError(
-                    f"Cannot load model {pretrained_model_name}: model is not cached locally and an error occurred"
-                    " while trying to fetch metadata from the Hub. Please check out the root cause in the stacktrace"
-                    " above."
+                    f"Cannot load model {pretrained_model_name}: the model is not fully cached locally ({e}) and an"
+                    " error occurred while trying to fetch metadata from the Hub. Please check out the root cause in"
+                    " the stacktrace above."
                 ) from model_info_call_error
 
     @classmethod

src/diffusers/utils/hub_utils.py

@@ -416,18 +416,14 @@ def _get_checkpoint_shard_files(
         return shard_filenames, sharded_metadata
 
     # At this stage pretrained_model_name_or_path is a model identifier on the Hub
-    if local_files_only:
-        # We are offline: the index file has already been resolved from the local cache and the shards live
-        # alongside it, so use that folder directly. Skip `snapshot_download`, whose snapshot-completeness
-        # check in newer `huggingface_hub` versions raises a less specific error for a missing shard than
-        # the check below.
-        cached_folder = os.path.dirname(index_filename)
-    else:
-        allow_patterns = original_shard_filenames
-        if subfolder is not None:
-            allow_patterns = [os.path.join(subfolder, p) for p in allow_patterns]
+    allow_patterns = original_shard_filenames
+    if subfolder is not None:
+        allow_patterns = [os.path.join(subfolder, p) for p in allow_patterns]
+
+    ignore_patterns = ["*.json", "*.md"]
 
-        # If the repo doesn't have the required shards, error out early even before downloading anything.
+    # If the repo doesn't have the required shards, error out early even before downloading anything.
+    if not local_files_only:
         model_files_info = model_info(pretrained_model_name_or_path, revision=revision, token=token)
         for shard_file in original_shard_filenames:
             shard_file_present = any(shard_file in k.rfilename for k in model_files_info.siblings)
@@ -437,28 +433,29 @@ def _get_checkpoint_shard_files(
                     "required according to the checkpoint index."
                 )
 
-        try:
-            # Load from URL
-            cached_folder = snapshot_download(
-                pretrained_model_name_or_path,
-                cache_dir=cache_dir,
-                proxies=proxies,
-                token=token,
-                revision=revision,
-                allow_patterns=allow_patterns,
-                ignore_patterns=["*.json", "*.md"],
-                user_agent=user_agent,
-            )
-            if subfolder is not None:
-                cached_folder = os.path.join(cached_folder, subfolder)
-
-        # We have already dealt with RepositoryNotFoundError and RevisionNotFoundError when getting the index,
-        # so we don't have to catch them here. We have also dealt with EntryNotFoundError.
-        except HfHubHTTPError as e:
-            raise EnvironmentError(
-                f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load {pretrained_model_name_or_path}. You should try"
-                " again after checking your internet connection."
-            ) from e
+    try:
+        # Load from URL
+        cached_folder = snapshot_download(
+            pretrained_model_name_or_path,
+            cache_dir=cache_dir,
+            proxies=proxies,
+            local_files_only=local_files_only,
+            token=token,
+            revision=revision,
+            allow_patterns=allow_patterns,
+            ignore_patterns=ignore_patterns,
+            user_agent=user_agent,
+        )
+        if subfolder is not None:
+            cached_folder = os.path.join(cached_folder, subfolder)
+
+    # We have already dealt with RepositoryNotFoundError and RevisionNotFoundError when getting the index, so
+    # we don't have to catch them here. We have also dealt with EntryNotFoundError.
+    except HfHubHTTPError as e:
+        raise EnvironmentError(
+            f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load {pretrained_model_name_or_path}. You should try"
+            " again after checking your internet connection."
+        ) from e
 
     cached_filenames = [os.path.join(cached_folder, f) for f in original_shard_filenames]
     for cached_file in cached_filenames:

tests/pipelines/test_pipelines.py

@@ -25,6 +25,7 @@
 import unittest
 import unittest.mock as mock
 import warnings
+from pathlib import Path
 
 import numpy as np
 import PIL.Image
@@ -504,28 +505,50 @@ def test_local_files_only_are_used_when_no_internet(self):
                 if p1.data.ne(p2.data).sum() > 0:
                     assert False, "Parameters not the same!"
 
-    def test_local_files_only_returns_cached_snapshot_without_snapshot_download(self):
+    def test_local_files_only_uses_same_snapshot_download_patterns(self):
         # diffusers downloads a filtered subset of a repo (skipping e.g. `.gitattributes`). Newer
-        # `huggingface_hub` versions validate that a cached snapshot is complete under
-        # `local_files_only=True`, so an offline `snapshot_download(allow_patterns=None)` would wrongly
-        # expect the whole repo. The offline path must instead compute `allow_patterns` from the local
-        # cache and return the cached snapshot directly. See
-        # https://github.com/huggingface/diffusers/issues/14117
+        # `huggingface_hub` versions validate that a cached snapshot contains every file matching the
+        # requested patterns under `local_files_only=True`, so an offline
+        # `snapshot_download(allow_patterns=None)` would wrongly expect the whole repo. The offline path
+        # must compute the same patterns as the online one (from the cached snapshot instead of
+        # `model_info`). See https://github.com/huggingface/diffusers/issues/14117
         with tempfile.TemporaryDirectory() as tmpdirname:
-            cached_folder = DiffusionPipeline.download(
-                "hf-internal-testing/tiny-stable-diffusion-torch", cache_dir=tmpdirname
-            )
+            with mock.patch(
+                "diffusers.pipelines.pipeline_utils.snapshot_download", side_effect=snapshot_download
+            ) as mock_snapshot_download:
+                cached_folder = DiffusionPipeline.download(
+                    "hf-internal-testing/tiny-stable-diffusion-torch", cache_dir=tmpdirname
+                )
+                online_kwargs = mock_snapshot_download.call_args.kwargs
 
-            with mock.patch("diffusers.pipelines.pipeline_utils.snapshot_download") as mock_snapshot_download:
                 offline_folder = DiffusionPipeline.download(
                     "hf-internal-testing/tiny-stable-diffusion-torch",
                     cache_dir=tmpdirname,
                     local_files_only=True,
                 )
+                offline_kwargs = mock_snapshot_download.call_args.kwargs
 
-            # A fully cached pipeline returns early; `snapshot_download` must not be called at all.
-            mock_snapshot_download.assert_not_called()
             assert os.path.samefile(offline_folder, cached_folder)
+            assert set(offline_kwargs["allow_patterns"]) == set(online_kwargs["allow_patterns"])
+            assert set(offline_kwargs["ignore_patterns"]) == set(online_kwargs["ignore_patterns"])
+
+    @require_hf_hub_version_greater("1.21.0")
+    def test_local_files_only_raises_for_snapshot_with_missing_weights(self):
+        # An interrupted download leaves a cached snapshot without some weights; loading it offline must
+        # surface `huggingface_hub`'s incomplete-snapshot error instead of failing later at model load time.
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            cached_folder = DiffusionPipeline.download(
+                "hf-internal-testing/tiny-stable-diffusion-torch", cache_dir=tmpdirname
+            )
+            for weights_file in Path(cached_folder).glob("unet/diffusion_pytorch_model*"):
+                weights_file.unlink()
+
+            with self.assertRaisesRegex(OSError, "incomplete"):
+                DiffusionPipeline.download(
+                    "hf-internal-testing/tiny-stable-diffusion-torch",
+                    cache_dir=tmpdirname,
+                    local_files_only=True,
+                )
 
     def test_download_from_variant_folder(self):
         for use_safetensors in [False, True]: