xinyun/ComfyUI-KJNodes
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git synced 2026-01-04 22:50:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

autodownload superprompt model

Browse Source
This commit is contained in:
kijai 2024-06-10 17:46:25 +03:00
parent cdd507de43
commit 36e3b6f66a
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
nodes/nodes.py
View File

@ -1332,6 +1332,12 @@ https://huggingface.co/roborovski/superprompt-v1
from transformers import T5Tokenizer, T5ForConditionalGeneration
checkpoint_path = os.path.join(script_directory, "models","superprompt-v1")
if not os.path.exists(checkpoint_path):
print(f"Downloading model to: {checkpoint_path}")
from huggingface_hub import snapshot_download
snapshot_download(repo_id="roborovski/superprompt-v1",
local_dir=checkpoint_path,
local_dir_use_symlinks=False)
tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small", legacy=False)
model = T5ForConditionalGeneration.from_pretrained(checkpoint_path, device_map=device)