diff --git a/install.py b/install.py index c033733..56025dc 100644 --- a/install.py +++ b/install.py @@ -2,19 +2,29 @@ import launch import os import pkg_resources import sys -import traceback +from tqdm import tqdm +import urllib.request req_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), "requirements.txt") -import os - models_dir = os.path.abspath("models/roop") +model_url = "https://huggingface.co/henryruhs/roop/resolve/main/inswapper_128.onnx" +model_name = os.path.basename(model_url) +model_path = os.path.join(models_dir, model_name) + +def download(url, path): + request = urllib.request.urlopen(url) + total = int(request.headers.get('Content-Length', 0)) + with tqdm(total=total, desc='Downloading', unit='B', unit_scale=True, unit_divisor=1024) as progress: + urllib.request.urlretrieve(url, path, reporthook=lambda count, block_size, total_size: progress.update(block_size)) if not os.path.exists(models_dir): os.makedirs(models_dir) - print(f"roop : You can put the model in {models_dir} directory") -print("Check roop requirements") +if not os.path.exists(model_path): + download(model_url, model_path) + +print("Checking roop requirements") with open(req_file) as file: for package in file: try: