download weights with huggingface_hub

Author: TemryL

.gitignore

@@ -1,6 +1,5 @@
-data/
-*.safetensors
-*.bin
+models/*
+!models/.gitkeep
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

README.md

@@ -19,51 +19,10 @@ ComfyUI adaptation of [IDM-VTON](https://github.com/yisol/IDM-VTON).
 cd custom_nodes
 git clone https://github.com/TemryL/ComfyUI-IDM-VTON.git
 cd ComfyUI-IDM-VTON
-pip install -r requirements.txt 
-```
-
-## Download weights:
-Download the models weights from [yisol/IDM-VTON](https://huggingface.co/yisol/IDM-VTON) in [HuggingFace](https://huggingface.co). The folder structure should be as follow:
-
-```
-models
-└── idm_vton
-    ├── image_encoder
-    │   ├── config.json
-    │   └── model.safetensors
-    ├── scheduler
-    │   └── scheduler_config.json
-    ├── text_encoder
-    │   ├── config.json
-    │   └── model.safetensors
-    ├── text_encoder_2
-    │   ├── config.json
-    │   └── model.safetensors
-    ├── tokenizer
-    │   ├── merges.txt
-    │   ├── special_tokens_map.json
-    │   ├── tokenizer_config.json
-    │   └── vocab.json
-    ├── tokenizer_2
-    │   ├── merges.txt
-    │   ├── special_tokens_map.json
-    │   ├── tokenizer_config.json
-    │   └── vocab.json
-    ├── unet
-    │   ├── config.json
-    │   └── diffusion_pytorch_model.bin
-    ├── unet_encoder
-    │   ├── config.json
-    │   └── diffusion_pytorch_model.safetensors
-    └── vae
-        ├── config.json
-        └── diffusion_pytorch_model.safetensors
-```
-You can place this folder wherever you want as long as you specify the correct path to `models/idm_vton` in the `Load IDM-VTON Pipeline` node. Config files are already provided in this repo. We also provide a shell script to download the weights from [HuggingFace](https://huggingface.co/yisol/IDM-VTON):
-```bash
-./scripts/download_weights.sh
+python install.py
 ```
 
+Models weights from [yisol/IDM-VTON](https://huggingface.co/yisol/IDM-VTON) in [HuggingFace](https://huggingface.co) will be downloaded in [models](models/) folder of this repository.
 
 # Mask Generation
 The workflow provided above uses [ComfyUI Segment Anything](https://github.com/storyicon/comfyui_segment_anything) to generate the image mask.

install.py

@@ -1,9 +1,12 @@
 import sys
 import os.path
 import subprocess
+from huggingface_hub import snapshot_download
 
 
-custom_nodes_path = os.path.dirname(os.path.abspath(__file__))
+CUSTOM_NODES_PATH = os.path.dirname(os.path.abspath(__file__))
+WEIGHTS_PATH = os.path.join(CUSTOM_NODES_PATH, "models")
+HF_REPO_ID = "yisol/IDM-VTON"
 
 
 def build_pip_install_cmds(args):
@@ -14,8 +17,9 @@ def build_pip_install_cmds(args):
 
 def ensure_package():
     cmds = build_pip_install_cmds(['-r', 'requirements.txt'])
-    subprocess.run(cmds, cwd=custom_nodes_path)
+    subprocess.run(cmds, cwd=CUSTOM_NODES_PATH)
 
 
 if __name__ == "__main__":
-    ensure_package()
+    ensure_package()
+    snapshot_download(repo_id=HF_REPO_ID, local_dir=WEIGHTS_PATH)