Mirrors
/
EveryDream2trainer
mirror of https://github.com/victorchall/EveryDream2trainer.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

don't double-download data if safetensors files exist

This commit is contained in:
Damian Stewart 2023-01-23 19:33:50 +01:00
parent 4f3b4d7dda
commit 18fec68149
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
utils/huggingface_downloader.py
View File

@ -30,8 +30,13 @@ def try_download_model_from_hf(repo_id: str,
return None, None, None
model_subfolders = ["text_encoder", "vae", "unet", "scheduler", "tokenizer"]
allow_patterns = [os.path.join(subfolder or '', f, "*") for f in model_subfolders]
downloaded_folder = huggingface_hub.snapshot_download(repo_id=repo_id, allow_patterns=allow_patterns)
allow_patterns = ["model_index.json"] + [os.path.join(subfolder or '', f, "*") for f in model_subfolders]
# prefer *.bin files for now
# TODO: look for *.safetensors files and download them instead, if they exist
ignore_patterns = "*.safetensors"
downloaded_folder = huggingface_hub.snapshot_download(repo_id=repo_id,
allow_patterns=allow_patterns,
ignore_patterns=ignore_patterns)
is_sd1_attn, yaml_path = patch_unet(downloaded_folder)
return downloaded_folder, is_sd1_attn, yaml_path