• Sayak Paul's avatar
    [Core] support saving and loading of sharded checkpoints (#7830) · 7d887118
    Sayak Paul authored
    
    
    * feat: support saving a model in sharded checkpoints.
    
    * feat: make loading of sharded checkpoints work.
    
    * add tests
    
    * cleanse the loading logic a bit more.
    
    * more resilience while loading from the Hub.
    
    * parallelize shard downloads by using snapshot_download()/
    
    * default to a shard size.
    
    * more fix
    
    * Empty-Commit
    
    * debug
    
    * fix
    
    * uality
    
    * more debugging
    
    * fix more
    
    * initial comments from Benjamin
    
    * move certain methods to loading_utils
    
    * add test to check if the correct number of shards are present.
    
    * add a test to check if loading of sharded checkpoints from the Hub is okay
    
    * clarify the unit when passed as an int.
    
    * use hf_hub for sharding.
    
    * remove unnecessary code
    
    * remove unnecessary function
    
    * lucain's comments.
    
    * fixes
    
    * address high-level comments.
    
    * fix test
    
    * subfolder shenanigans./
    
    * Update src/diffusers/utils/hub_utils.py
    Co-authored-by: default avatarLucain <lucainp@gmail.com>
    
    * Apply suggestions from code review
    Co-authored-by: default avatarLucain <lucainp@gmail.com>
    
    * remove _huggingface_hub_version as not needed.
    
    * address more feedback.
    
    * add a test for local_files_only=True/
    
    * need hf hub to be at least 0.23.2
    
    * style
    
    * final comment.
    
    * clean up subfolder.
    
    * deal with suffixes in code.
    
    * _add_variant default.
    
    * use weights_name_pattern
    
    * remove add_suffix_keyword
    
    * clean up downloading of sharded ckpts.
    
    * don't return something special when using index.json
    
    * fix more
    
    * don't use bare except
    
    * remove comments and catch the errors better
    
    * fix a couple of things when using is_file()
    
    * empty
    
    ---------
    Co-authored-by: default avatarLucain <lucainp@gmail.com>
    7d887118
model_loading_utils.py 8.41 KB