feat: FlashPack support for transformers pipeline components · huggingface/diffusers@4630dd7 · GitHub
Skip to content

Commit 4630dd7

Browse files
committed
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.
1 parent 17b48ab commit 4630dd7

3 files changed

Lines changed: 115 additions & 5 deletions

File tree

‎src/diffusers/pipelines/pipeline_loading_utils.py‎

Lines changed: 40 additions & 0 deletions

‎src/diffusers/pipelines/pipeline_utils.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161
from ..utils import (
6262
CONFIG_NAME,
6363
DEPRECATED_REVISION_ARGS,
64+
FLASHPACK_WEIGHTS_NAME,
6465
BaseOutput,
6566
PushToHubMixin,
6667
_get_detailed_type,
@@ -70,9 +71,11 @@
7071
is_accelerate_available,
7172
is_accelerate_version,
7273
is_bitsandbytes_version,
74+
is_flashpack_available,
7375
is_hpu_available,
7476
is_torch_npu_available,
7577
is_torch_version,
78+
is_transformers_available,
7679
is_transformers_version,
7780
logging,
7881
numpy_to_pil,
@@ -85,6 +88,9 @@
8588
if is_torch_npu_available():
8689
import torch_npu # noqa: F401
8790

91+
if is_transformers_available():
92+
from transformers import PreTrainedModel
93+
8894
from .pipeline_loading_utils import (
8995
ALL_IMPORTABLE_CLASSES,
9096
CONNECTED_PIPES_KEYS,
@@ -284,6 +290,9 @@ class implements both a save and loading method. The pipeline is easily reloaded
284290
Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
285291
repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
286292
namespace).
293+
use_flashpack (`bool`, *optional*, defaults to `False`):
294+
If set to `True`, model components (diffusers models as well as transformers models such as text
295+
encoders) save their weights as a single `model.flashpack` file instead of safetensors.
287296
288297
kwargs (`Dict[str, Any]`, *optional*):
289298
Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
@@ -322,6 +331,26 @@ def is_saveable_module(name, value):
322331
sub_model = _unwrap_model(sub_model)
323332
model_cls = sub_model.__class__
324333

334+
# transformers' `save_pretrained` has no FlashPack support, so mirror what
335+
# `ModelMixin.save_pretrained(use_flashpack=True)` does for diffusers models: save the
336+
# config as usual and pack the weights to a single `model.flashpack` file.
337+
if use_flashpack and is_transformers_available() and isinstance(sub_model, PreTrainedModel):
338+
if not is_flashpack_available():
339+
raise ImportError("Please install flashpack to save a pipeline with `use_flashpack=True`.")
340+
import flashpack
341+
342+
component_directory = os.path.join(save_directory, pipeline_component_name)
343+
os.makedirs(component_directory, exist_ok=True)
344+
sub_model.config.save_pretrained(component_directory)
345+
if sub_model.can_generate():
346+
sub_model.generation_config.save_pretrained(component_directory)
347+
flashpack.serialization.pack_to_file(
348+
state_dict_or_model=sub_model.state_dict(),
349+
destination_path=os.path.join(component_directory, FLASHPACK_WEIGHTS_NAME),
350+
target_dtype=sub_model.dtype,
351+
)
352+
continue
353+
325354
save_method_name = None
326355
# search for the model's base class in LOADABLE_CLASSES
327356
for library_name, library_classes in LOADABLE_CLASSES.items():
@@ -745,6 +774,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike, **kwa
745774
dduf_file(`str`, *optional*):
746775
Load weights from the specified dduf file. <Deprecated> This argument is deprecated and will be removed
747776
in version 0.41.0. </Deprecated>
777+
use_flashpack (`bool`, *optional*, defaults to `False`):
778+
If set to `True`, model components whose folder contains a `model.flashpack` file (diffusers models as
779+
well as transformers models such as text encoders) load their weights from it; FlashPack weights are
780+
downloaded instead of the other formats for those components. If set to `False`, FlashPack weights are
781+
never downloaded.
748782
disable_mmap ('bool', *optional*, defaults to 'False'):
749783
Whether to disable mmap when loading a Safetensors model. This option can perform better when the model
750784
is on a network mount or hard drive, which may not handle the seeky-ness of mmap very well.

‎tests/others/test_flashpack.py‎

Lines changed: 41 additions & 5 deletions

0 commit comments

Comments
 (0)