feat(examples): Support S3 for all StableDiffusionPipeline components
#76
Conversation
This file contains 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
Previously, setting include_non_persistent_buffers=False would only include persistent buffers, and remove both non-persistent buffers and parameters. Parameters were not supposed to be affected by this flag, so this fix changes it to only remove non-persistent buffers.
fix(serialization): Don't drop parameters with non-persistent buffers
This additionally adds parameter weight validation for diffusers models, refactors `serialize_model` substantially, adds logging level command line arguments, and shifts more outputs to use a logger.
sangstar
approved these changes
Feb 1, 2024
Comment on lines
+319
to
+324
| serialize_pretrained( | ||
| pipeline.tokenizer, output_prefix, "tokenizer", force=args.force | ||
| ) | ||
| serialize_pretrained( | ||
| pipeline.scheduler, output_prefix, "scheduler", force=args.force | ||
| ) |
There was a problem hiding this comment.
So the purpose of serialize_pretrained is specifically to support saving artifacts like SD's scheduler and tokenizer? I assume this is cleaner as they're not configs nor modules.
There was a problem hiding this comment.
Yes, pretty much. The scheduler can actually be saved as a single JSON file (around 300 to 400 bytes), so it could have its own special code to avoid a zip file and temporary directory, but it seemed unnecessary at this time.
In summary:
text_encoder/vae/unet— tensors &config.json- We use
tensorizerto save these - Potentially very, very large, so the optimizations in
tensorizerare important
- We use
scheduler—config.jsononly- Supports
.save_pretrained(dir)and.from_pretrained(dir) - Roughly around 400 bytes, so downloader optimizations aren't a huge deal
- Supports
tokenizer— entire directory of files- Supports
.save_pretrained(dir)and.from_pretrained(dir) - Roughly around 1.5 MB for SD 1.5, and compresses well, so zips work nicely
- For models that use "fast tokenizers," these can be saved as a single
tokenizer.jsonfile instead, but not all models support fast tokenizers
- Supports
harubaru
requested changes
Feb 1, 2024
harubaru
approved these changes
Feb 1, 2024
sangstar
merged commit 184d7a6
into
sangstar/update-serialization-cl-script-for-container
6 of 7 checks passed
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Updated
diffuserssupport inhf_serialization.pyMore fixes and improvements for #73.
This change:
StableDiffusionPipelinecomponentstext_encoder,vae(updated),unet(updated),scheduler(new),tokenizer(new)schedulerandtokenizerare saved as.zipfiles containing the directory written by their.save_pretrained()methodstransformerstokenizersinclude_non_persistent_buffers=Falsemainto support usinginclude_non_persistent_buffers=Falsecorrectlydiffusersmodelsserialize_modelsubstantially(Outdated): I left in the code to generate a test image through
diffusersbecause it is good example code for this repository of how to re-assemble the components of an SD model and was good for testing that these changes work. It could be commented out or deleted later, or changed to be only enabled through a flag.(Update): The test image generation code for diffusers is now commented out.