Support output path

README.md

@@ -1,4 +1,16 @@
-# Python script to download models from CivitAI using an API token
+# Python script to download models from CivitAI using an API key
+
+## Getting Started
+
+This script requires a [CivitAI](https://civitai.com/user/account)
+API key.  You can create the API key as follows:
+
+1. Log into [CivitAI](https://civitai.com).
+2. Go to the [Manage Account](https://civitai.com/user/account) page.
+3. Scroll down to the `API Keys` section, close to the bottom of the page.
+4. Click the `+ Add API key` button to create a new API key.
+5. Give the API key a name and click the `Save` button.
+6. You will then use your newly created API key with this script.
 
 ## Installation
 
@@ -7,22 +19,30 @@ wget https://raw.githubusercontent.com/ashleykleynhans/civitai-downloader/main/d
 mv download.py /usr/local/bin/download-model
 chmod +x /usr/local/bin/download-model
 ```
+
+> [!NOTE]
+> This assumes you are using RunPod and logged in as `root`
+> user.  If not, the installation commands should be prefixed
+> with `sudo`.
+
+> [!IMPORTANT]
+> It is important to ensure that you use the **DOWNLOAD** link
+> and not the link to the model page in CivitAI.
+
 ## Usage
 
 ```bash
-download-model [URL]
+download-model [URL] [DESTINATION]
 ```
 
-eg:
+To download to the current directory:
 
 ```bash
-download-model https://civitai.com/api/download/models/46846
+download-model https://civitai.com/api/download/models/46846 .
 ```
 
-## NOTE
+To download to a different directory:
 
-1. This assumes you are using RunPod and logged in as `root`
-   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.
+```bash
+download-model https://civitai.com/api/download/models/46846 /workspace/stable-diffusion-webui/models/Stable-diffusion
+```

download.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+import os.path
 import sys
 import argparse
 import time
@@ -20,7 +21,13 @@ def get_args():
     parser.add_argument(
         'url',
         type=str,
-        help='CivitAI Download URL'
+        help='CivitAI Download URL, eg: https://civitai.com/api/download/models/46846'
+    )
+
+    parser.add_argument(
+        'output_path',
+        type=str,
+        help='Output path, eg: /workspace/stable-diffusion-webui/models/Stable-diffusion'
     )
 
     return parser.parse_args()
@@ -35,7 +42,7 @@ def get_token():
         return None
 
 
-def store_token(token):
+def store_token(token: str):
     TOKEN_FILE.parent.mkdir(parents=True, exist_ok=True)
 
     with open(TOKEN_FILE, 'w') as file:
@@ -48,7 +55,7 @@ def prompt_for_civitai_token():
     return token
 
 
-def download_file(url, token):
+def download_file(url: str, output_path: str, token: str):
     headers = {
         'Authorization': f'Bearer {token}',
         'User-Agent': USER_AGENT,
@@ -88,7 +95,9 @@ def http_response(self, request, response):
     if total_size is not None:
         total_size = int(total_size)
 
-    with open(filename, 'wb') as f:
+    output_file = os.path.join(output_path, filename)
+
+    with open(output_file, 'wb') as f:
         downloaded = 0
         start_time = time.time()
 
@@ -137,7 +146,7 @@ def main():
         token = prompt_for_civitai_token()
 
     try:
-        download_file(args.url, token)
+        download_file(args.url, args.output_path, token)
     except Exception as e:
         print(f'ERROR: {e}')