FlashPack support for transformers pipeline components

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,
@@ -234,6 +235,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():
@@ -853,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"):
@@ -1136,39 +1177,55 @@ 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})"
         )
 
     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"]
 
         use_onnx = use_onnx if use_onnx is not None else is_onnx
         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
 
 

src/diffusers/pipelines/pipeline_utils.py

@@ -54,16 +54,19 @@
 from ..utils import (
     CONFIG_NAME,
     DEPRECATED_REVISION_ARGS,
+    FLASHPACK_WEIGHTS_NAME,
     BaseOutput,
     PushToHubMixin,
     _get_detailed_type,
     _is_valid_type,
     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.
@@ -873,6 +907,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: