From fea6d47c55b2a4721adc5f8b2e5d83a5762d1ea8 Mon Sep 17 00:00:00 2001 From: Ashley Kleynhans Date: Wed, 27 Mar 2024 16:51:53 +0200 Subject: [PATCH] 1. Rewrote script in Python using stdlib modules, so no pip modules are required. 2. Store and use CivitAI API token to download models. 3. Show progress bar when downloading the model. --- .gitignore | 2 + README.md | 10 ++--- download.py | 122 ++++++++++++++++++++++++++++++++++++++++++++++++++++ download.sh | 27 ------------ 4 files changed, 129 insertions(+), 32 deletions(-) create mode 100644 .gitignore create mode 100755 download.py delete mode 100755 download.sh diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f377c56 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +venv +.idea diff --git a/README.md b/README.md index 7fda475..5cf76bc 100644 --- a/README.md +++ b/README.md @@ -1,22 +1,22 @@ -# Bash script to download models from CivitAI using curl +# Python script to download models from CivitAI using an API token ## Installation ```bash git clone https://github.com/ashleykleynhans/civitai-downloader.git -mv civitai-downloader/download.sh /usr/local/bin/download-model +mv civitai-downloader/download.py /usr/local/bin/download-model chmod +x /usr/local/bin/download-model ``` ## Usage ```bash -download-model [URL] [DESTINATION] +download-model [URL] ``` eg: ```bash -download-model https://civitai.com/api/download/models/15236 /workspace/stable-diffusion-webui/models/Stable-diffusion +download-model https://civitai.com/api/download/models/15236 ``` ## NOTE @@ -25,4 +25,4 @@ download-model https://civitai.com/api/download/models/15236 /workspace/stable-d user. If not, the installation commands should be prefixed with `sudo`. 2. It is important to ensure that you use the **DOWNLOAD** link -and not the link to the model page in CivitAI. + and not the link to the model page in CivitAI. diff --git a/download.py b/download.py new file mode 100755 index 0000000..8658492 --- /dev/null +++ b/download.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +import sys +import argparse +import urllib.request +from pathlib import Path +from urllib.parse import urlparse, parse_qs, unquote + + +CHUNK_SIZE = 1638400 +TOKEN_FILE = Path.home() / '.civitai' / 'config' + + +def get_args(): + parser = argparse.ArgumentParser( + description='CivitAI Downloader', + ) + + parser.add_argument( + 'url', + type=str, + help='CivitAI Download URL' + ) + + return parser.parse_args() + + +def get_token(): + try: + with open(TOKEN_FILE, 'r') as file: + token = file.read() + return token + except Exception as e: + return None + + +def store_token(token): + # Ensure the directory exists + TOKEN_FILE.parent.mkdir(parents=True, exist_ok=True) + + # Write the token to the file + with open(TOKEN_FILE, 'w') as file: + file.write(token) + + +def prompt_for_civitai_token(): + token = input('Please enter your CivitAI API token: ') + store_token(token) + return token + + +def download_file(url, token): + # Prepare the initial request with necessary headers + headers = { + 'Authorization': f'Bearer {token}', + 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3', + } + request = urllib.request.Request(url, headers=headers) + + # Disable automatic redirect handling + class NoRedirection(urllib.request.HTTPErrorProcessor): + def http_response(self, request, response): + return response + https_response = http_response + + opener = urllib.request.build_opener(NoRedirection) + response = opener.open(request) + + if response.status in [301, 302, 303, 307, 308]: + redirect_url = response.getheader('Location') + + # Extract filename from the redirect URL + parsed_url = urlparse(redirect_url) + query_params = parse_qs(parsed_url.query) + content_disposition = query_params.get('response-content-disposition', [None])[0] + + if content_disposition: + filename = unquote(content_disposition.split('filename=')[1].strip('"')) + else: + raise Exception('Unable to determine filename') + + response = urllib.request.urlopen(redirect_url) + else: + raise Exception('No redirect found, something went wrong') + + # Download the file + total_size = response.getheader('Content-Length') + if total_size is not None: + total_size = int(total_size) + + with open(filename, 'wb') as f: + downloaded = 0 + + while True: + buffer = response.read(CHUNK_SIZE) + + if not buffer: + break + + downloaded += len(buffer) + f.write(buffer) + + if total_size is not None: + progress = downloaded / total_size + sys.stdout.write(f'\rDownloading: {filename} [{progress*100:.2f}%]') + sys.stdout.flush() + + sys.stdout.write('\n') + print(f'Download completed. File saved as: {filename}') + + +def main(): + args = get_args() + token = get_token() + + if not token: + token = prompt_for_civitai_token() + + download_file(args.url, token) + + +if __name__ == '__main__': + main() diff --git a/download.sh b/download.sh deleted file mode 100755 index c7d05d4..0000000 --- a/download.sh +++ /dev/null @@ -1,27 +0,0 @@ -#!/usr/bin/env bash - -if [ "$#" -ne 2 ]; then - echo "Usage: $0 " - echo " eg: $0 https://civitai.com/api/download/models/15236 /workspace/stable-diffusion-webui/models/Stable-diffusion" - exit 1 -fi - -URL=${1} -DESTINATION=${2} -USER_AGENT_STRING="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36" - -if ! echo "${URL}" | grep -q "api"; then - echo "ERROR: Incorrect URL provided, you must provide the Download link from CivitAI, not the link to the model page." - exit 1 -fi - -echo "Downloading model from ${URL}, please wait..." - -cd ${DESTINATION} - -if ! curl -JL --remote-name -A "${USER_AGENT_STRING}" "${URL}"; then - echo "ERROR: curl command failed. Unable to download the file." - exit 1 -fi - -echo "Model downloaded successfully!"