civitai multithreaded downloads

CHANGELOG.md

@@ -114,6 +114,9 @@ Upgrades are still possible and supported, but above is recommended for best exp
     combinations results in 50+ samplers which is not practical  
     items such as algorithm (e.g. karras) is actually a sampler option, not a sampler itself  
 - **CivitAI**:
+  - civitai model download is now multithreaded and resumable  
+    meaning that you can download multiple models in parallel  
+    as well as resume aborted/incomplete downloads  
   - civitai integration in *models -> civitai* can now find most  
     previews AND metadata for most models (checkpoints, loras, embeddings)  
     metadata is now parsed and saved in *[model].json*  

javascript/black-teal.css

@@ -117,6 +117,7 @@ svg.feather.feather-image, .feather .feather-image { display: none }
 #txt2img_actions_column, #img2img_actions_column { flex-flow: wrap; justify-content: space-between; }
 #txt2img_enqueue_wrapper, #img2img_enqueue_wrapper { min-width: unset; width: 48%; }
 #txt2img_generate_box, #img2img_generate_box { min-width: unset; width: 48%; }
+textarea[rows="1"] { height: 33px !important; width: 99% !important; padding: 8px !important; }
 
 #extras_upscale { margin-top: 10px }
 #txt2img_progress_row > div { min-width: var(--left-column); max-width: var(--left-column); }

javascript/extraNetworks.js

@@ -156,7 +156,7 @@ function cardClicked(textToAdd, allowNegativePrompt) {
 
 function extraNetworksSearchButton(event) {
   const tabname = getENActiveTab();
-  const searchTextarea = gradioApp().querySelector(`#${tabname}_extra_tabs > div > div > textarea`);
+  const searchTextarea = gradioApp().querySelector(`#${tabname}_extra_search textarea`);
   const button = event.target;
   const text = button.classList.contains('search-all') ? '' : `${button.textContent.trim()}/`;
   searchTextarea.value = text;

javascript/style.css

@@ -8,6 +8,7 @@ div.compact{ gap: 1em; }
 div.gradio-html.min{ min-height: 0; }
 .block.gradio-checkbox { margin: 0.75em 1.5em 0 0; align-self: center; }
 .block.gradio-dropdown, .block.gradio-slider, .block.gradio-checkbox, .block.gradio-textbox, .block.gradio-radio, .block.gradio-checkboxgroup, .block.gradio-number, .block.gradio-colorpicker { border-width: 0 !important; box-shadow: none !important;}
+.block.gradio-textbox { overflow: visible !important; }
 .block.padded:not(.gradio-accordion) { padding: 0 !important; margin-right: 0; min-width: 90px !important; }
 .compact{ background: transparent !important; padding: 0 !important; }
 .dark .gradio-dropdown ul.options li.item:not(:has(.hide)) { background-color: var(--neutral-900); }

modules/modelloader.py

@@ -108,46 +108,76 @@ def download_civit_preview(model_path: str, preview_url: str):
     return res
 
 
-def download_civit_model(model_url: str, model_name: str, model_path: str, model_type: str, preview):
+download_pbar = None
+
+def download_civit_model_thread(model_name, model_url, model_path, model_type, preview):
+    import hashlib
+    sha256 = hashlib.sha256()
+    sha256.update(model_name.encode('utf-8'))
+    temp_file = sha256.hexdigest()[:8] + '.tmp'
+
     if model_type == 'LoRA':
         model_file = os.path.join(shared.opts.lora_dir, model_path, model_name)
+        temp_file = os.path.join(shared.opts.lora_dir, model_path, temp_file)
     else:
         model_file = os.path.join(shared.opts.ckpt_dir, model_path, model_name)
-    res = f'CivitAI download: name={model_name} url={model_url} path={model_path}'
+        temp_file = os.path.join(shared.opts.ckpt_dir, model_path, temp_file)
+
+    res = f'CivitAI download: name={model_name} url={model_url} path={model_path} temp={temp_file}'
     if os.path.isfile(model_file):
         res += ' already exists'
         shared.log.warning(res)
         return res
 
-    r = shared.req(model_url, stream=True)
+    headers = {}
+    starting_pos = 0
+    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}-'}
+
+    r = shared.req(model_url, headers=headers, stream=True)
     total_size = int(r.headers.get('content-length', 0))
-    block_size = 16384 # 16KB blocks
-    written = 0
+    res += f' size={round((starting_pos + total_size)/1024/1024)}Mb'
+    shared.log.info(res)
     shared.state.begin('civitai-download-model')
-    try:
-        with open(model_file, 'wb') as f:
-            with p.Progress(p.TextColumn('[cyan]{task.description}'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn(), console=shared.console) as progress:
-                task = progress.add_task(description="Download starting", total=total_size)
-                # for data in tqdm(req.iter_content(block_size), total=total_size//1024, unit='KB', unit_scale=False):
+    block_size = 16384 # 16KB blocks
+    written = starting_pos
+    global download_pbar # pylint: disable=global-statement
+    if download_pbar is None:
+        download_pbar = p.Progress(p.TextColumn('[cyan]{task.description}'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn(), p.TextColumn('[cyan]{task.fields[name]}'), console=shared.console)
+    with download_pbar:
+        task = download_pbar.add_task(description="Download starting", total=starting_pos+total_size, name=model_name)
+        try:
+            with open(temp_file, 'ab') as f:
                 for data in r.iter_content(block_size):
                     written = written + len(data)
                     f.write(data)
-                    progress.update(task, advance=block_size, description="Downloading")
-        if written < 1024 * 1024: # min threshold
-            os.remove(model_file)
-            raise ValueError(f'removed invalid download: bytes={written}')
-        if preview is not None:
-            preview_file = os.path.splitext(model_file)[0] + '.jpg'
-            preview.save(preview_file)
-            res += f' preview={preview_file}'
-    except Exception as e:
-        shared.log.error(f'CivitAI download error: name={model_name} url={model_url} path={model_path} {e}')
-    if total_size == written:
-        shared.log.info(f'{res} size={total_size}')
+                    download_pbar.update(task, description="Download", completed=written)
+            if written < 1024 * 1024: # min threshold
+                os.remove(temp_file)
+                raise ValueError(f'removed invalid download: bytes={written}')
+            if preview is not None:
+                preview_file = os.path.splitext(model_file)[0] + '.jpg'
+                preview.save(preview_file)
+                res += f' preview={preview_file}'
+        except Exception as e:
+            shared.log.error(f'{res} {e}')
+        finally:
+            download_pbar.stop_task(task)
+            download_pbar.remove_task(task)
+    if starting_pos+total_size != written:
+        shared.log.warning(f'{res} written={round(written/1024/1024)}Mb incomplete download')
     else:
-        shared.log.error(f'{res} size={total_size} written={written}')
+        os.rename(temp_file, model_file)
     shared.state.end()
-    return res
+
+
+def download_civit_model(model_url: str, model_name: str, model_path: str, model_type: str, preview):
+    import threading
+    thread = threading.Thread(target=download_civit_model_thread, args=(model_name, model_url, model_path, model_type, preview))
+    thread.start()
+    return f'CivitAI download: name={model_name} url={model_url} path={model_path}'
 
 
 def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config: Dict[str, str] = None, token = None, variant = None, revision = None, mirror = None):

modules/shared.py

@@ -1010,8 +1010,9 @@ def get_version():
     return version
 
 
-def req(url_addr, **kwargs):
-    headers = { 'Content-type': 'application/json' }
+def req(url_addr, headers = None, **kwargs):
+    if headers is None:
+        headers = { 'Content-type': 'application/json' }
     try:
         res = requests.get(url_addr, timeout=30, headers=headers, verify=False, allow_redirects=True, **kwargs)
     except Exception as e:

modules/ui_extra_networks.py

@@ -165,6 +165,8 @@ class ExtraNetworksPage:
     def create_thumb(self):
         created = 0
         for f in self.missing_thumbs:
+            if not os.path.exists(f):
+                continue
             fn, _ext = os.path.splitext(f)
             fn = fn.replace('.preview', '')
             fn = f'{fn}.thumb.jpg'
@@ -175,8 +177,9 @@ class ExtraNetworksPage:
                 img = Image.open(f)
             except Exception:
                 shared.log.warning(f'Extra network removing invalid image: {f}')
-                os.remove(f)
             try:
+                if img is None:
+                    os.remove(f)
                 if img is not None and img.width > 1024 or img.height > 1024 or os.path.getsize(f) > 65536:
                     img = img.convert('RGB')
                     img.thumbnail((512, 512), Image.HAMMING)

modules/ui_models.py

@@ -281,7 +281,8 @@ def create_ui():
                                 model['stats']['downloadCount'],
                                 model['stats']['rating']
                             ])
-                    return data1, [], []
+                    res = f'Search result: name={name} tag={tag or "none"} type={model_type} models={len(data1)}'
+                    return res, gr.update(visible=len(data1) > 0, value=data1 if len(data1) > 0 else []), gr.update(visible=False, value=None), gr.update(visible=False, value=None)
 
                 def civit_select1(evt: gr.SelectData, in_data):
                     model_id = in_data[evt.index[0]][0]
@@ -300,7 +301,7 @@ def create_ui():
                                     d['createdAt'],
                                 ])
                     log.debug(f'CivitAI select: model="{in_data[evt.index[0]]}" versions={len(data2)}')
-                    return data2, preview_img
+                    return data2, None, preview_img
 
                 def civit_select2(evt: gr.SelectData, in_data):
                     variant_id = in_data[evt.index[0]][0]
@@ -388,7 +389,7 @@ def create_ui():
                     civit_previews_rehash = gr.Checkbox(value=True, label="Check alternative hash")
 
                 with gr.Row():
-                    gr.HTML('<h2>Search for models</h2>Select a model, model version and and model variant from the search results to download<br>')
+                    gr.HTML('<h2>Search for models</h2>')
                 with gr.Row():
                     with gr.Column(scale=1):
                         civit_model_type = gr.Dropdown(label='Model type', choices=['SD 1.5', 'SD XL', 'LoRA', 'Other'], value='LoRA')
@@ -397,31 +398,45 @@ def create_ui():
                             civit_search_text = gr.Textbox('', label = 'Search models', placeholder='keyword')
                             civit_search_tag = gr.Textbox('', label = '', placeholder='tags')
                             civit_search_btn = ToolButton(value="🔍", label="Search", interactive=False)
+                        with gr.Row():
+                            civit_search_res = gr.HTML('')
                 with gr.Row():
-                    civit_download_model_btn = gr.Button(value="Download model", variant='primary')
+                    gr.HTML('<h2>Download model</h2>')
+                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_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)
                     civit_path = gr.Textbox('', label = 'Download path', placeholder='optional subfolder path where to save model', visible=True)
                 with gr.Row():
-                    with gr.Column():
-                        civit_headers2 = ['ID', 'ModelID', 'Name', 'Base', 'Created', 'Preview']
-                        civit_types2 = ['number', 'number', 'str', 'str', 'date', 'str']
-                        civit_results2 = gr.DataFrame(value = None, label = 'Model versions', show_label = True, interactive = False, wrap = True, overflow_row_behaviour = 'paginate', max_rows = 10, headers = civit_headers2, datatype = civit_types2, type='array')
-                    with gr.Column():
-                        civit_headers3 = ['Name', 'Size', 'Metadata', 'URL']
-                        civit_types3 = ['str', 'number', 'str', 'str']
-                        civit_results3 = gr.DataFrame(value = None, label = 'Model variants', show_label = True, interactive = False, wrap = True, overflow_row_behaviour = 'paginate', max_rows = 10, headers = civit_headers3, datatype = civit_types3, type='array')
+                    gr.HTML('<h2>Search results</h2>')
                 with gr.Row():
                     civit_headers1 = ['ID', 'Name', 'Tags', 'Downloads', 'Rating']
                     civit_types1 = ['number', 'str', 'str', 'number', 'number']
-                    civit_results1 = gr.DataFrame(value = None, label = 'Search results', show_label = True, interactive = False, wrap = True, overflow_row_behaviour = 'paginate', max_rows = 10, headers = civit_headers1, datatype = civit_types1, type='array')
+                    civit_results1 = gr.DataFrame(value = None, label = None, show_label = False, interactive = False, wrap = True, overflow_row_behaviour = 'paginate', max_rows = 10, headers = civit_headers1, datatype = civit_types1, type='array', visible=False)
+                with gr.Row():
+                    with gr.Column():
+                        civit_headers2 = ['ID', 'ModelID', 'Name', 'Base', 'Created', 'Preview']
+                        civit_types2 = ['number', 'number', 'str', 'str', 'date', 'str']
+                        civit_results2 = gr.DataFrame(value = None, label = 'Model versions', show_label = True, interactive = False, wrap = True, overflow_row_behaviour = 'paginate', max_rows = 10, headers = civit_headers2, datatype = civit_types2, type='array', visible=False)
+                    with gr.Column():
+                        civit_headers3 = ['Name', 'Size', 'Metadata', 'URL']
+                        civit_types3 = ['str', 'number', 'str', 'str']
+                        civit_results3 = gr.DataFrame(value = None, label = 'Model variants', show_label = True, interactive = False, wrap = True, overflow_row_behaviour = 'paginate', max_rows = 10, headers = civit_headers3, datatype = civit_types3, type='array', visible=False)
 
-                civit_search_text.submit(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_results1, civit_results2, civit_results3])
-                civit_search_tag.submit(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_results1, civit_results2, civit_results3])
-                civit_search_btn.click(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_results1, civit_results2, civit_results3])
-                civit_results1.select(fn=civit_select1, inputs=[civit_results1], outputs=[civit_results2, models_image])
+                def is_visible(component):
+                    visible = len(component) > 0 if component is not None else False
+                    return gr.update(visible=visible)
+
+                civit_search_text.submit(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
+                civit_search_tag.submit(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
+                civit_search_btn.click(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_search_res, civit_results1, civit_results2, civit_results3])
+                civit_results1.select(fn=civit_select1, inputs=[civit_results1], outputs=[civit_results2, civit_results3, models_image])
                 civit_results2.select(fn=civit_select2, inputs=[civit_results2], outputs=[civit_results3])
                 civit_results3.select(fn=civit_select3, inputs=[civit_results3], outputs=[civit_selected, civit_name, civit_search_btn])
+                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_previews_btn.click(fn=civit_search_metadata, inputs=[civit_previews_rehash, civit_previews_rehash], outputs=[models_outcome])