Files
webui-fooocus-prompt-expansion/install.py
power88 af627f86bc Init
2024-05-26 18:41:24 +08:00

24 lines
1.2 KiB
Python

import os
import pathlib
import shutil
from huggingface_hub import hf_hub_download
from modules.scripts import basedir
ext_dir = basedir()
fooocus_expansion_path = pathlib.Path(ext_dir) / "models" / "prompt_expansion"
base_model_path = pathlib.Path(ext_dir) / "extensions" / "webui-fooocus-prompt-expansion" / "models"
if not os.path.exists(os.path.join(fooocus_expansion_path, 'pytorch_model.bin')):
try:
print(f'### webui-fooocus-prompt-expansion: Downloading model...')
shutil.copytree(os.path.join(base_model_path), fooocus_expansion_path)
hf_hub_download(repo_id='lllyasviel/misc', filename='fooocus_expansion.bin', local_dir=os.path.join(fooocus_expansion_path), resume_download=True, local_dir_use_symlinks=False)
os.rename(os.path.join(fooocus_expansion_path, 'fooocus_expansion.bin'), os.path.join(fooocus_expansion_path, 'pytorch_model.bin'))
except Exception as e:
print(f'### webui-fooocus-prompt-expansion: Failed to download model...')
print(e)
print(f'### webui-fooocus-prompt-expansion: To enable this custom node, please download the model manually from "https://huggingface.co/lllyasviel/misc/tree/main/fooocus_expansion.bin" and place it in {fooocus_expansion_path}.')
else:
pass