DiffuseCraft / utils /download_utils.py
zerhero's picture
refactor
169b607
raw
history blame
1.84 kB
def download_things(
directory: str,
url: str,
hf_token: str = "",
civitai_api_key: str = ""):
"""
:param directory: String, directory to save the downloaded file
:param url: String, URL to download the file from
:param hf_token: String, Hugging Face API token
:param civitai_api_key: String, Civitai API key
"""
import os
url = url.strip()
if "drive.google.com" in url:
original_dir = os.getcwd()
os.chdir(directory)
os.system(f"gdown --fuzzy {url}")
os.chdir(original_dir)
elif "huggingface.co" in url:
url = url.replace("?download=true", "")
# url = urllib.parse.quote(url, safe=':/') # fix encoding
if "/blob/" in url:
url = url.replace("/blob/", "/resolve/")
user_header = f'"Authorization: Bearer {hf_token}"'
if hf_token:
os.system(
f"aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 {url} -d {directory} -o {url.split('/')[-1]}")
else:
os.system(
f"aria2c --optimize-concurrent-downloads --console-log-level=error --summary-interval=10 -c -x 16 -k "
f"1M -s 16 {url} -d {directory} -o {url.split('/')[-1]}")
elif "civitai.com" in url:
if "?" in url:
url = url.split("?")[0]
if civitai_api_key:
url = url + f"?token={civitai_api_key}"
os.system(
f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")
else:
print("\033[91mYou need an API key to download Civitai models.\033[0m")
else:
os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")