enable civit token

CHANGELOG.md

@@ -27,6 +27,7 @@
       use if you have multiple complex loras that may be causing performance degradation  
       as it fuses lora with model during load instead of interpreting lora on-the-fly  
 - **General**
+  - **CivitAI downloader** allow usage of access tokens for gated or private models  
   - **Extra networks** new *settting -> extra networks -> build info on first access*  
     indexes all networks on first access instead of server startup  
   - **Ipex** improvements, thanks @disty0

modules/modelloader.py

@@ -111,7 +111,7 @@ def download_civit_preview(model_path: str, preview_url: str):
 
 download_pbar = None
 
-def download_civit_model_thread(model_name, model_url, model_path, model_type, preview):
+def download_civit_model_thread(model_name, model_url, model_path, model_type, preview, token):
     import hashlib
     sha256 = hashlib.sha256()
     sha256.update(model_name.encode('utf-8'))
@@ -135,7 +135,9 @@ def download_civit_model_thread(model_name, model_url, model_path, model_type, p
     if os.path.isfile(temp_file):
         starting_pos = os.path.getsize(temp_file)
         res += f' resume={round(starting_pos/1024/1024)}Mb'
-        headers = {'Range': f'bytes={starting_pos}-'}
+        headers['Range'] = f'bytes={starting_pos}-'
+    if token is not None and len(token) > 0:
+        headers['Authorization'] = f'Bearer {token}'
 
     r = shared.req(model_url, headers=headers, stream=True)
     total_size = int(r.headers.get('content-length', 0))
@@ -175,9 +177,9 @@ def download_civit_model_thread(model_name, model_url, model_path, model_type, p
     return res
 
 
-def download_civit_model(model_url: str, model_name: str, model_path: str, model_type: str, preview):
+def download_civit_model(model_url: str, model_name: str, model_path: str, model_type: str, preview, token: str = None):
     import threading
-    thread = threading.Thread(target=download_civit_model_thread, args=(model_name, model_url, model_path, model_type, preview))
+    thread = threading.Thread(target=download_civit_model_thread, args=(model_name, model_url, model_path, model_type, preview, token))
     thread.start()
     return f'CivitAI download: name={model_name} url={model_url} path={model_path}'
 

modules/ui_models.py

@@ -488,12 +488,12 @@ def create_ui():
                     log.debug(f'CivitAI select: variant={in_data[evt.index[0]]}')
                     return in_data[evt.index[0]][3], in_data[evt.index[0]][0], gr.update(interactive=True)
 
-                def civit_download_model(model_url: str, model_name: str, model_path: str, model_type: str, image_url: str):
+                def civit_download_model(model_url: str, model_name: str, model_path: str, model_type: str, image_url: str, token: str = None):
                     if model_url is None or len(model_url) == 0:
                         return 'No model selected'
                     try:
                         from modules.modelloader import download_civit_model
-                        res = download_civit_model(model_url, model_name, model_path, model_type, image_url)
+                        res = download_civit_model(model_url, model_name, model_path, model_type, image_url, token)
                     except Exception as e:
                         res = f"CivitAI model downloaded error: model={model_url} {e}"
                         log.error(res)
@@ -577,6 +577,8 @@ def create_ui():
                 with gr.Row():
                     civit_download_model_btn = gr.Button(value="Download", variant='primary')
                     gr.HTML('<span style="line-height: 2em">Select a model, model version and and model variant from the search results to download or enter model URL manually</span><br>')
+                with gr.Row():
+                    civit_token = gr.Textbox('', label='CivitAI token', placeholder='optional access token for private or gated models')
                 with gr.Row():
                     civit_name = gr.Textbox('', label='Model name', placeholder='select model from search results', visible=True)
                     civit_selected = gr.Textbox('', label='Model URL', placeholder='select model from search results', visible=True)
@@ -619,7 +621,7 @@ def create_ui():
                 civit_results1.change(fn=is_visible, inputs=[civit_results1], outputs=[civit_results1])
                 civit_results2.change(fn=is_visible, inputs=[civit_results2], outputs=[civit_results2])
                 civit_results3.change(fn=is_visible, inputs=[civit_results3], outputs=[civit_results3])
-                civit_download_model_btn.click(fn=civit_download_model, inputs=[civit_selected, civit_name, civit_path, civit_model_type, models_image], outputs=[models_outcome])
+                civit_download_model_btn.click(fn=civit_download_model, inputs=[civit_selected, civit_name, civit_path, civit_model_type, models_image, civit_token], outputs=[models_outcome])
                 civit_previews_btn.click(fn=civit_search_metadata, inputs=[civit_previews_rehash, civit_previews_rehash], outputs=[models_outcome])
 
             with gr.Tab(label="Update"):