From 8c53dfdcd0d6d2f292f8e9f72d72e8f3a4bdcaad Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Sat, 16 Mar 2024 12:42:33 -0400
Subject: [PATCH] update civit downloader

---
 CHANGELOG.md           |  1 +
 modules/modelloader.py | 10 +++++++++-
 modules/ui_models.py   | 10 +++++++---
 3 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index acb3ac9ea..c9303fa9c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -171,6 +171,7 @@ Further details:
   - fix ipadapter apply/unapply on batch runs
   - fix control with multiple units and override images
   - fix control-lllite
+  - update civitai downloader to handler new metadata
   - improve control error handling
   - use default model variant if specified variant doesnt exist
   - use diffusers lora load override for *lcm/tcd/turbo loras*
diff --git a/modules/modelloader.py b/modules/modelloader.py
index bb6777c20..50f5623dc 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -1,5 +1,6 @@
 import os
 import time
+import json
 import shutil
 import importlib
 from typing import Dict
@@ -117,6 +118,12 @@ def download_civit_model_thread(model_name, model_url, model_path, model_type, t
         try:
             with open(temp_file, 'ab') as f:
                 for data in r.iter_content(block_size):
+                    if written == 0:
+                        try: # check if response is JSON message instead of bytes
+                            shared.log.error(f'Model download: response={json.loads(data.decode("utf-8"))}')
+                            raise ValueError('response: type=json expected=bytes')
+                        except Exception: # this is good
+                            pass
                     written = written + len(data)
                     f.write(data)
                     download_pbar.update(task, description="Download", completed=written)
@@ -136,7 +143,8 @@ def download_civit_model_thread(model_name, model_url, model_path, model_type, t
             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:
+    elif os.path.exists(temp_file):
+        shared.log.debug(f'Model download complete: temp="{temp_file}" path="{model_file}"')
         os.rename(temp_file, model_file)
     shared.state.end()
     return res
diff --git a/modules/ui_models.py b/modules/ui_models.py
index fa9545b47..a434834d5 100644
--- a/modules/ui_models.py
+++ b/modules/ui_models.py
@@ -474,9 +474,13 @@ def create_ui():
                     for model in data:
                         if model['id'] == model_id:
                             for d in model['modelVersions']:
-                                if d.get('images') is not None and len(d['images']) > 0 and len(d['images'][0]['url']) > 0:
-                                    preview_img = d['images'][0]['url']
-                                data2.append([d['id'], d['modelId'], d['name'], d['baseModel'], d['createdAt']])
+                                try:
+                                    if d.get('images') is not None and len(d['images']) > 0 and len(d['images'][0]['url']) > 0:
+                                        preview_img = d['images'][0]['url']
+                                    data2.append([d.get('id', None), d.get('modelId', None) or model_id, d.get('name', None), d.get('baseModel', None), d.get('createdAt', None) or d.get('publishedAt', None)])
+                                except Exception as e:
+                                    log.error(f'CivitAI select: model="{in_data[evt.index[0]]}" {e}')
+                                    log.error(f'CivitAI version data={type(d)}: {d}')
                     log.debug(f'CivitAI select: model="{in_data[evt.index[0]]}" versions={len(data2)}')
                     return data2, None, preview_img