Merge pull request #80 from Learwin/fix_civitai

Fix Civitai download errors by providing api key with request

scripts/mo/dl/http_downloader.py

@@ -5,7 +5,7 @@ import requests
 from tqdm import tqdm
 
 from scripts.mo.dl.downloader import Downloader
-
+from scripts.mo.environment import env
 
 class HttpDownloader(Downloader):
 
@@ -30,7 +30,14 @@ class HttpDownloader(Downloader):
 
         yield {'bytes_ready': 'None', 'bytes_total': 'None', 'speed_rate': 'None', 'elapsed': 'None'}
 
-        response = requests.get(url, stream=True)
+        apiKey = env.api_key()
+        if apiKey:
+            authHeader = {'Content-Type':'application/json',
+                        'Authorization': 'Bearer ' + apiKey}
+            response = requests.get(url, stream=True, headers=authHeader)
+        else:
+            response = requests.get(url, stream=True)
+            
         total_size = int(response.headers.get('content-length', 0))
 
         yield {'bytes_ready': 0, 'bytes_total': total_size, 'speed_rate': 0, 'elapsed': 0}

scripts/mo/environment.py

@@ -76,7 +76,8 @@ class Environment:
     card_width: Callable[[], str]
     card_height: Callable[[], str]
     is_debug_mode_enabled: Callable[[], bool]
-
+    api_key: Callable[[], str]
+    
     def is_storage_initialized(self) -> bool:
         return hasattr(self, 'storage')
 

scripts/model_organizer.py

@@ -127,7 +127,12 @@ env.autobind_file = (
     if hasattr(shared.opts, 'mo_autobind_file')
     else True
 )
-
+	
+env.api_key = (
+    lambda: shared.opts.mo_api_key
+    if hasattr(shared.opts, 'mo_api_key')
+    else ""
+)
 
 env.model_path = (
     lambda: shared.opts.mo_model_path
@@ -192,6 +197,7 @@ def on_ui_settings():
         'mo_prefill_pos_prompt': OptionInfo(True, 'When creating a record based on local file, automatically import the added positive prompts'),
         'mo_prefill_neg_prompt': OptionInfo(True, 'When creating a record based on local file, automatically import the added negative prompts'),
         'mo_autobind_file': OptionInfo(True, 'Automatically bind record to local file'),
+        'mo_api_key': OptionInfo("", "Civitai API Key. Create an API key under 'https://civitai.com/user/account' all the way at the bottom. Don't share the token!"),
     }
 
     dir_opts = {