FlashPack support for transformers pipeline components#14167
Open
rumutaydin wants to merge 3 commits into
Open
FlashPack support for transformers pipeline components#14167rumutaydin wants to merge 3 commits into
rumutaydin wants to merge 3 commits into
Conversation
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.
from_pretrained popped use_flashpack but never passed it to download, so Hub-mediated loading never fetched FlashPack weights regardless of the flag.
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.
Contributor
|
Hi @rumutaydin, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Follow-up to #14166 (stacked on it).
use_flashpack=Truecurrently appliesonly to diffusers models; transformers components (text encoders) silently keep
safetensors, so a "FlashPack pipeline" never actually is one end to end the gap
noted when FlashPack+DDUF support was discussed in
huggingface/huggingface_hub#3716.
save_pretrainedhas no FlashPack support, soDiffusionPipeline.save_pretrained(use_flashpack=True)now saves a transformerscomponent's config as usual and packs its weights to
model.flashpackthe samecall
ModelMixin.save_pretrained(use_flashpack=True)makes for diffusers models.use_flashpack=Trueand a transformers component's folder shipsmodel.flashpack, initialize the model from its config on empty weights and assignthe packed weights onto it, mirroring the existing diffusers-model FlashPack path
(same device_map handling and warnings). Folders without a flashpack file load
exactly as before, so existing mixed-format repos are unaffected.
With this, saving
tiny-flux-pipewithuse_flashpack=Trueproduces zero safetensors(CLIP, T5, transformer, VAE all packed), the filtered Hub snapshot round-trips
bitwise-equal, and the reloaded pipeline runs inference on CPU.
One difference from the diffusers-model FlashPack path: loaded transformers
components are returned in eval mode, matching what
transformers.from_pretraineddoes. (The diffusers-model FlashPack path currently returns models in train mode,
happy to fix that separately if wanted.)
Tests
a new test saves tiny-flux-pipe with use_flashpack=True and asserts the transformers components (CLIP and T5 text encoders) produce model.flashpack and no safetensors, then applies the real Hub download filter to the saved folder and reloads it checking the restored weights are bitwise-equal to the originals and the components are in eval mode. The existing round-trip test now builds the mixed layout explicitly (diffusers models as FlashPack, text encoders as safetensors, i.e. what save_pretrained(use_flashpack=True) produced before this PR), so repos saved with the old behavior are covered and keep loading.
Related: #12564.
Before submitting
.ai/review-rules.md?documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
@DN6 @yiyixuxu