From 743adc51df4435107b8a37dd06568b28910b7940 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Robin=20Umut=20Ayd=C4=B1n?= Date: Thu, 9 Jul 2026 21:44:17 +0300 Subject: [PATCH 1/3] fix: download filtering for FlashPack pipelines Pipelines saved with use_flashpack=True could not be loaded through the Hub download path: _get_ignore_patterns excluded all *.safetensors, so transformers components (text encoders) had no weights in the snapshot, and FlashPack files were missing from variant_compatible_siblings' weight names, so they never entered the allow patterns either. With use_flashpack=False, *.flashpack files were still downloaded despite the documented behavior. Judge safetensors compatibility over non-flashpack folders only, keep just the flashpack file inside flashpack-covered folders, register FlashPack files as weight names, and never download flashpack weights when use_flashpack is disabled. --- .../pipelines/pipeline_loading_utils.py | 43 ++-- tests/others/test_flashpack.py | 196 ++++++++++++++++++ 2 files changed, 226 insertions(+), 13 deletions(-) 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/tests/others/test_flashpack.py b/tests/others/test_flashpack.py index c14410c0d8e0..3f39eb884055 100644 --- a/tests/others/test_flashpack.py +++ b/tests/others/test_flashpack.py @@ -13,12 +13,20 @@ # 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 diffusers import AutoPipelineForText2Image 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 +35,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 +115,156 @@ 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)) From 357d85d45826a7c8f864b6b3ca0d9c48c79c2bfc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Robin=20Umut=20Ayd=C4=B1n?= Date: Thu, 9 Jul 2026 22:16:03 +0300 Subject: [PATCH 2/3] fix: forward use_flashpack from from_pretrained to download from_pretrained popped use_flashpack but never passed it to download, so Hub-mediated loading never fetched FlashPack weights regardless of the flag. --- src/diffusers/pipelines/pipeline_utils.py | 1 + tests/others/test_flashpack.py | 14 +++++++++++++- 2 files changed, 14 insertions(+), 1 deletion(-) 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 3f39eb884055..b359c363c74a 100644 --- a/tests/others/test_flashpack.py +++ b/tests/others/test_flashpack.py @@ -19,8 +19,9 @@ 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, @@ -268,3 +269,14 @@ def test_ignore_patterns_explicit_use_safetensors_mixed_repo(self): 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"]) From 2b93098052eafdfd35a8601315b078b28d2d35e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Robin=20Umut=20Ayd=C4=B1n?= Date: Thu, 9 Jul 2026 22:18:10 +0300 Subject: [PATCH 3/3] feat: FlashPack support for transformers pipeline components use_flashpack=True previously applied only to diffusers models; transformers components (text encoders) silently kept safetensors, so pipelines were never FlashPack end to end. On save, pack transformers components' weights to model.flashpack the same way ModelMixin.save_pretrained(use_flashpack=True) does. On load, when a transformers component folder ships a flashpack file, initialize the model from its config on empty weights and assign the packed weights onto it, mirroring the existing diffusers-model FlashPack path. Folders without a flashpack file keep loading exactly as before, so existing mixed-format repos are unaffected. --- .../pipelines/pipeline_loading_utils.py | 40 ++++++++++++++++ src/diffusers/pipelines/pipeline_utils.py | 34 ++++++++++++++ tests/others/test_flashpack.py | 46 +++++++++++++++++-- 3 files changed, 115 insertions(+), 5 deletions(-) diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index d028725005d3..a145329a5348 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -38,6 +38,7 @@ deprecate, get_class_from_dynamic_module, is_accelerate_available, + is_flashpack_available, is_peft_available, is_transformers_available, is_transformers_version, @@ -854,6 +855,45 @@ def load_sub_model( and transformers_version >= version.parse("4.20.0") ) + # transformers' `from_pretrained` has no FlashPack support, so when a transformers component ships + # FlashPack weights, load it here the way `ModelMixin.from_pretrained` does for diffusers models: + # initialize the model on empty weights from its config, then assign the packed weights onto it. + # Component folders without a flashpack file (e.g. from repos saved before transformers components + # were packed) keep loading through their regular `load_method` below. + flashpack_file = os.path.join(cached_folder, name, FLASHPACK_WEIGHTS_NAME) + if use_flashpack and is_transformers_model and os.path.isfile(flashpack_file): + if not is_flashpack_available(): + raise ImportError("Please install flashpack to load a pipeline with `use_flashpack=True`.") + import flashpack + + config = class_obj.config_class.from_pretrained(os.path.join(cached_folder, name)) + dtype_orig = None + if torch_dtype is not None: + dtype_orig = torch.get_default_dtype() + torch.set_default_dtype(torch_dtype) + with accelerate.init_empty_weights(): + loaded_sub_model = class_obj(config) + if dtype_orig is not None: + torch.set_default_dtype(dtype_orig) + + if device_map is None: + logger.warning( + "`device_map` has not been provided for FlashPack, model will be on `cpu` - provide `device_map` to fully utilize " + "the benefit of FlashPack." + ) + flashpack_device = torch.device("cpu") + else: + device = device_map[""] + if isinstance(device, str) and device in ["auto", "balanced", "balanced_low_0", "sequential"]: + raise ValueError( + "FlashPack `device_map` should not be one of `auto`, `balanced`, `balanced_low_0`, `sequential`. Use a specific device instead, e.g., `device_map='cuda'` or `device_map='cuda:0'" + ) + flashpack_device = torch.device(device) if not isinstance(device, torch.device) else device + + flashpack.mixin.assign_from_file(model=loaded_sub_model, path=flashpack_file, device=flashpack_device) + # transformers' `from_pretrained` always returns models in eval mode. + return loaded_sub_model.eval() + # For transformers models >= 4.56.0, use 'dtype' instead of 'torch_dtype' to avoid deprecation warnings if issubclass(class_obj, torch.nn.Module): if is_transformers_model and transformers_version >= version.parse("4.56.0"): diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 12dd122524c1..58185ccfeecb 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -54,6 +54,7 @@ from ..utils import ( CONFIG_NAME, DEPRECATED_REVISION_ARGS, + FLASHPACK_WEIGHTS_NAME, BaseOutput, PushToHubMixin, _get_detailed_type, @@ -61,9 +62,11 @@ is_accelerate_available, is_accelerate_version, is_bitsandbytes_version, + is_flashpack_available, is_hpu_available, is_torch_npu_available, is_torch_version, + is_transformers_available, is_transformers_version, logging, numpy_to_pil, @@ -76,6 +79,9 @@ if is_torch_npu_available(): import torch_npu # noqa: F401 +if is_transformers_available(): + from transformers import PreTrainedModel + from .pipeline_loading_utils import ( ALL_IMPORTABLE_CLASSES, CONNECTED_PIPES_KEYS, @@ -270,6 +276,9 @@ class implements both a save and loading method. The pipeline is easily reloaded Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the repository you want to push to with `repo_id` (will default to the name of `save_directory` in your namespace). + use_flashpack (`bool`, *optional*, defaults to `False`): + If set to `True`, model components (diffusers models as well as transformers models such as text + encoders) save their weights as a single `model.flashpack` file instead of safetensors. kwargs (`Dict[str, Any]`, *optional*): Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. @@ -308,6 +317,26 @@ def is_saveable_module(name, value): sub_model = _unwrap_model(sub_model) model_cls = sub_model.__class__ + # transformers' `save_pretrained` has no FlashPack support, so mirror what + # `ModelMixin.save_pretrained(use_flashpack=True)` does for diffusers models: save the + # config as usual and pack the weights to a single `model.flashpack` file. + if use_flashpack and is_transformers_available() and isinstance(sub_model, PreTrainedModel): + if not is_flashpack_available(): + raise ImportError("Please install flashpack to save a pipeline with `use_flashpack=True`.") + import flashpack + + component_directory = os.path.join(save_directory, pipeline_component_name) + os.makedirs(component_directory, exist_ok=True) + sub_model.config.save_pretrained(component_directory) + if sub_model.can_generate(): + sub_model.generation_config.save_pretrained(component_directory) + flashpack.serialization.pack_to_file( + state_dict_or_model=sub_model.state_dict(), + destination_path=os.path.join(component_directory, FLASHPACK_WEIGHTS_NAME), + target_dtype=sub_model.dtype, + ) + continue + save_method_name = None # search for the model's base class in LOADABLE_CLASSES for library_name, library_classes in LOADABLE_CLASSES.items(): @@ -731,6 +760,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike, **kwa loading `from_flax`. dduf_file(`str`, *optional*): Load weights from the specified dduf file. + use_flashpack (`bool`, *optional*, defaults to `False`): + If set to `True`, model components whose folder contains a `model.flashpack` file (diffusers models as + well as transformers models such as text encoders) load their weights from it; FlashPack weights are + downloaded instead of the other formats for those components. If set to `False`, FlashPack weights are + never downloaded. disable_mmap ('bool', *optional*, defaults to 'False'): Whether to disable mmap when loading a Safetensors model. This option can perform better when the model is on a network mount or hard drive, which may not handle the seeky-ness of mmap very well. diff --git a/tests/others/test_flashpack.py b/tests/others/test_flashpack.py index b359c363c74a..4b39725f80f6 100644 --- a/tests/others/test_flashpack.py +++ b/tests/others/test_flashpack.py @@ -119,17 +119,20 @@ def test_load_model_device_auto(self): @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. + # A mixed repo pairs flashpack weights (diffusers components) with safetensors weights + # (transformers components) — the layout `save_pretrained(use_flashpack=True)` produced + # before transformers components were packed too. The download path must keep both formats; + # on `main` it drops the transformers safetensors, so the filtered snapshot fails to load. 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) + pipeline.save_pretrained(saved) + for component in ("transformer", "vae"): + getattr(pipeline, component).save_pretrained(os.path.join(saved, component), use_flashpack=True) + os.remove(os.path.join(saved, component, "diffusion_pytorch_model.safetensors")) self.assertTrue((pathlib.Path(saved) / "transformer" / "model.flashpack").exists()) self.assertTrue((pathlib.Path(saved) / "text_encoder" / "model.safetensors").exists()) @@ -149,6 +152,39 @@ def test_download_filter_flashpack_pipeline_roundtrip(self): for key, value in original.items(): self.assertTrue(torch.equal(value, restored[key])) + @require_flashpack + def test_save_load_pipeline_transformers_flashpack(self): + # With `use_flashpack=True`, transformers components (text encoders) are packed too, so the + # pipeline is FlashPack end to end: no safetensors are written, the download filter keeps the + # flashpack weights, and the round trip restores every model component bitwise. + pipeline = AutoPipelineForText2Image.from_pretrained(self.model_id) + 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) + for component in ("transformer", "vae", "text_encoder", "text_encoder_2"): + component_files = os.listdir(os.path.join(saved, component)) + self.assertIn("model.flashpack", component_files) + self.assertFalse(any(f.endswith(".safetensors") for f in component_files)) + + kept = _files_kept_by_download_filter(saved, use_flashpack=True) + for component in ("transformer", "vae", "text_encoder", "text_encoder_2"): + self.assertIn(f"{component}/model.flashpack", 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", "vae", "text_encoder", "text_encoder_2"): + 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])) + # transformers' `from_pretrained` returns models in eval mode; the flashpack path must too + self.assertFalse(reloaded.text_encoder.training) + 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.