diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index 90cbffc5b69d..d028725005d3 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -234,6 +234,7 @@ def variant_compatible_siblings(filenames, variant=None, ignore_patterns=None) - FLAX_WEIGHTS_NAME, ONNX_WEIGHTS_NAME, ONNX_EXTERNAL_WEIGHTS_NAME, + FLASHPACK_WEIGHTS_NAME, ] if is_transformers_available(): @@ -1136,13 +1137,25 @@ def _get_ignore_patterns( use_flashpack: bool, variant: str | None = None, ) -> list[str]: - if ( - use_safetensors - and not allow_pickle - and not is_safetensors_compatible( - model_filenames, passed_components=passed_components, folder_names=model_folder_names, variant=variant - ) - ): + # Folders whose weights ship as flashpack. When `use_flashpack` is set we download only their + # flashpack file; when it is not set flashpack files are ignored entirely (see below), so these + # folders play no role and safetensors compatibility is judged over every folder as usual. + flashpack_folders = set() + if use_flashpack: + flashpack_folders = {os.path.split(f)[0] for f in model_filenames if f.endswith(".flashpack")} + + # Flashpack-covered folders legitimately have no safetensors, so exclude them when judging + # whether the remaining (e.g. transformers) folders can be served from safetensors. + safetensors_filenames = [f for f in model_filenames if os.path.split(f)[0] not in flashpack_folders] + safetensors_folder_names = [f for f in model_folder_names if f not in flashpack_folders] + safetensors_compatible = is_safetensors_compatible( + safetensors_filenames, + passed_components=passed_components, + folder_names=safetensors_folder_names, + variant=variant, + ) + + if use_safetensors and not allow_pickle and not safetensors_compatible: raise EnvironmentError( f"Could not find the necessary `safetensors` weights in {model_filenames} (variant={variant})" ) @@ -1150,18 +1163,13 @@ def _get_ignore_patterns( if from_flax: ignore_patterns = ["*.bin", "*.safetensors", "*.onnx", "*.pb"] - elif use_safetensors and is_safetensors_compatible( - model_filenames, passed_components=passed_components, folder_names=model_folder_names, variant=variant - ): + elif use_safetensors and safetensors_compatible: ignore_patterns = ["*.bin", "*.msgpack"] use_onnx = use_onnx if use_onnx is not None else is_onnx if not use_onnx: ignore_patterns += ["*.onnx", "*.pb"] - elif use_flashpack: - ignore_patterns = ["*.bin", "*.safetensors", "*.onnx", "*.pb", "*.msgpack"] - else: ignore_patterns = ["*.safetensors", "*.msgpack"] @@ -1169,6 +1177,15 @@ def _get_ignore_patterns( if not use_onnx: ignore_patterns += ["*.onnx", "*.pb"] + # Keep only the flashpack file inside flashpack folders (hub ignore patterns match full relative + # paths, so this is per-folder and leaves other folders' safetensors/bin untouched). + for folder in flashpack_folders: + ignore_patterns += [f"{folder}/*.safetensors", f"{folder}/*.bin"] + + # `use_flashpack=False` must never pull flashpack weights, in any of the branches above. + if not use_flashpack: + ignore_patterns.append("*.flashpack") + return ignore_patterns diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 7b653caa0657..12dd122524c1 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -873,6 +873,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike, **kwa dduf_file=dduf_file, load_connected_pipeline=load_connected_pipeline, trust_remote_code=trust_remote_code, + use_flashpack=use_flashpack, **kwargs, ) else: diff --git a/tests/others/test_flashpack.py b/tests/others/test_flashpack.py index c14410c0d8e0..b359c363c74a 100644 --- a/tests/others/test_flashpack.py +++ b/tests/others/test_flashpack.py @@ -13,12 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import pathlib +import shutil import tempfile import unittest +from fnmatch import fnmatch +from unittest import mock -from diffusers import AutoPipelineForText2Image +from diffusers import AutoPipelineForText2Image, DiffusionPipeline from diffusers.models.auto_model import AutoModel +from diffusers.pipelines.pipeline_loading_utils import ( + _get_ignore_patterns, + filter_model_files, + variant_compatible_siblings, +) from ..testing_utils import is_torch_available, require_flashpack, require_torch_gpu @@ -27,6 +36,41 @@ import torch +def _files_kept_by_download_filter(saved_dir, use_flashpack): + """Reproduce the allow/ignore filtering that `DiffusionPipeline.download` applies to a repo's + files, but over a local snapshot, so a filtered copy mirrors what the Hub path would fetch.""" + saved = pathlib.Path(saved_dir) + filenames = {p.relative_to(saved).as_posix() for p in saved.rglob("*") if p.is_file()} + folder_names = {p.name for p in saved.iterdir() if p.is_dir()} + model_folder_names = { + os.path.split(f)[0] for f in filter_model_files(filenames) if os.path.split(f)[0] in folder_names + } + + ignore_patterns = _get_ignore_patterns( + passed_components=[], + model_folder_names=list(model_folder_names), + model_filenames=list(filenames), + use_safetensors=True, + from_flax=False, + allow_pickle=True, + use_onnx=None, + is_onnx=False, + use_flashpack=use_flashpack, + variant=None, + ) + model_filenames, _ = variant_compatible_siblings(filenames, variant=None, ignore_patterns=ignore_patterns) + + allow_patterns = list(model_filenames) + allow_patterns += [f"{k}/*" for k in folder_names if k not in model_folder_names] + allow_patterns += [f"{k}/config.json" for k in model_folder_names] + allow_patterns += ["scheduler_config.json", "config.json", "model_index.json"] + + ignore_patterns = ignore_patterns + [f"{i}.index.*json" for i in ignore_patterns] + kept = {f for f in filenames if not any(fnmatch(f, p) for p in ignore_patterns)} + kept = {f for f in kept if any(fnmatch(f, p) for p in allow_patterns)} + return kept + + class FlashPackTests(unittest.TestCase): model_id: str = "hf-internal-testing/tiny-flux-pipe" @@ -72,3 +116,167 @@ def test_load_model_device_auto(self): model.save_pretrained(temp_dir, use_flashpack=True) with self.assertRaises(ValueError): model = AutoModel.from_pretrained(temp_dir, use_flashpack=True, device_map={"": "auto"}) + + @require_flashpack + def test_download_filter_flashpack_pipeline_roundtrip(self): + # A flashpack pipeline mixes flashpack weights (diffusers components) with safetensors weights + # (transformers components). The download path must keep both; on `main` it drops the + # transformers safetensors, so the filtered snapshot fails to load. This is the e2e repro, + # inverted into a green test. + pipeline = AutoPipelineForText2Image.from_pretrained(self.model_id) + # `ignore_cleanup_errors` because on Windows flashpack keeps `model.flashpack` mmap'd, which + # blocks the temp-dir teardown (unrelated to what this test asserts). + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: + saved = os.path.join(temp_dir, "saved") + downloaded = os.path.join(temp_dir, "downloaded") + pipeline.save_pretrained(saved, use_flashpack=True) + self.assertTrue((pathlib.Path(saved) / "transformer" / "model.flashpack").exists()) + self.assertTrue((pathlib.Path(saved) / "text_encoder" / "model.safetensors").exists()) + + kept = _files_kept_by_download_filter(saved, use_flashpack=True) + self.assertIn("transformer/model.flashpack", kept) + self.assertIn("text_encoder/model.safetensors", kept) + for f in kept: + dst = os.path.join(downloaded, f) + os.makedirs(os.path.dirname(dst), exist_ok=True) + shutil.copy2(os.path.join(saved, f), dst) + + reloaded = AutoPipelineForText2Image.from_pretrained(downloaded, use_flashpack=True) + for name in ("transformer", "text_encoder"): + original = getattr(pipeline, name).state_dict() + restored = getattr(reloaded, name).state_dict() + self.assertEqual(original.keys(), restored.keys()) + for key, value in original.items(): + self.assertTrue(torch.equal(value, restored[key])) + + def test_ignore_patterns_flashpack_false_excludes_flashpack(self): + # Dual-format repo (both safetensors and flashpack shipped). With `use_flashpack=False` the + # flashpack files must be ignored and every safetensors kept. + model_filenames = [ + "text_encoder/model.safetensors", + "unet/diffusion_pytorch_model.safetensors", + "unet/model.flashpack", + "vae/diffusion_pytorch_model.safetensors", + "vae/model.flashpack", + ] + patterns = _get_ignore_patterns( + passed_components=[], + model_folder_names=["text_encoder", "unet", "vae"], + model_filenames=model_filenames, + use_safetensors=True, + from_flax=False, + allow_pickle=True, + use_onnx=None, + is_onnx=False, + use_flashpack=False, + variant=None, + ) + self.assertIn("*.flashpack", patterns) + safetensors = [f for f in model_filenames if f.endswith(".safetensors")] + for f in safetensors: + self.assertFalse(any(fnmatch(f, p) for p in patterns)) + + def test_ignore_patterns_flashpack_true_mixed_repo(self): + # Mixed repo: flashpack for diffusers folders, safetensors for the transformers folder. + # The safetensors-less flashpack folders must not trip the compatibility check, the + # transformers safetensors must be kept, and only the flashpack folders lose safetensors. + model_filenames = [ + "text_encoder/model.safetensors", + "unet/model.flashpack", + "vae/model.flashpack", + ] + patterns = _get_ignore_patterns( + passed_components=[], + model_folder_names=["text_encoder", "unet", "vae"], + model_filenames=model_filenames, + use_safetensors=True, + from_flax=False, + allow_pickle=True, + use_onnx=None, + is_onnx=False, + use_flashpack=True, + variant=None, + ) + self.assertNotIn("*.safetensors", patterns) + self.assertNotIn("*.flashpack", patterns) + self.assertIn("unet/*.safetensors", patterns) + self.assertIn("vae/*.safetensors", patterns) + self.assertFalse(any(fnmatch("text_encoder/model.safetensors", p) for p in patterns)) + for f in ("unet/model.flashpack", "vae/model.flashpack"): + self.assertFalse(any(fnmatch(f, p) for p in patterns)) + + def test_ignore_patterns_flashpack_only_repo(self): + # Every model folder ships only flashpack (no transformers component). Default loading keeps + # `allow_pickle=True`, so the missing safetensors must not raise, and flashpack is kept. + model_filenames = ["transformer/model.flashpack", "vae/model.flashpack"] + patterns = _get_ignore_patterns( + passed_components=[], + model_folder_names=["transformer", "vae"], + model_filenames=model_filenames, + use_safetensors=True, + from_flax=False, + allow_pickle=True, + use_onnx=None, + is_onnx=False, + use_flashpack=True, + variant=None, + ) + self.assertNotIn("*.flashpack", patterns) + for f in model_filenames: + self.assertFalse(any(fnmatch(f, p) for p in patterns)) + + def test_ignore_patterns_no_flashpack_repo_use_flashpack_true(self): + # `use_flashpack=True` against a repo without any flashpack files keeps the normal + # safetensors-preferred behavior and does not error. + model_filenames = [ + "text_encoder/model.safetensors", + "unet/diffusion_pytorch_model.safetensors", + ] + patterns = _get_ignore_patterns( + passed_components=[], + model_folder_names=["text_encoder", "unet"], + model_filenames=model_filenames, + use_safetensors=True, + from_flax=False, + allow_pickle=True, + use_onnx=None, + is_onnx=False, + use_flashpack=True, + variant=None, + ) + self.assertIn("*.bin", patterns) + for f in model_filenames: + self.assertFalse(any(fnmatch(f, p) for p in patterns)) + + def test_ignore_patterns_explicit_use_safetensors_mixed_repo(self): + # Explicit `use_safetensors=True` (i.e. `allow_pickle=False`) on a mixed flashpack repo must + # not raise just because the flashpack folders have no safetensors. + model_filenames = [ + "text_encoder/model.safetensors", + "unet/model.flashpack", + "vae/model.flashpack", + ] + patterns = _get_ignore_patterns( + passed_components=[], + model_folder_names=["text_encoder", "unet", "vae"], + model_filenames=model_filenames, + use_safetensors=True, + from_flax=False, + allow_pickle=False, + use_onnx=None, + is_onnx=False, + use_flashpack=True, + variant=None, + ) + self.assertFalse(any(fnmatch("text_encoder/model.safetensors", p) for p in patterns)) + + @require_flashpack + def test_from_pretrained_forwards_use_flashpack_to_download(self): + # Without the forwarding, `download` treats the repo as `use_flashpack=False` and never + # downloads the flashpack weights the loader then needs. + pipeline = AutoPipelineForText2Image.from_pretrained(self.model_id) + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: + pipeline.save_pretrained(temp_dir, use_flashpack=True) + with mock.patch.object(DiffusionPipeline, "download", return_value=temp_dir) as download: + DiffusionPipeline.from_pretrained("hf-internal-testing/does-not-matter", use_flashpack=True) + self.assertTrue(download.call_args.kwargs["use_flashpack"])