From abfa9e9635fc2c8914baac2a97c5f44820de5f12 Mon Sep 17 00:00:00 2001
From: butaixianran <butaixianran@qq.com>
Date: Mon, 10 Apr 2023 19:13:04 +0800
Subject: [PATCH] v1.6.3

---
 README.md                              |  3 +
 scripts/ch_lib/model_action_civitai.py | 90 ++++++++++++++++++++------
 2 files changed, 75 insertions(+), 18 deletions(-)

diff --git a/README.md b/README.md
index 5853b52..ecb789b 100644
--- a/README.md
+++ b/README.md
@@ -230,6 +230,9 @@ Since v1.5.5, we've already optimized the SHA256 function to the top. So the onl
 
 
 # Change Log
+## v1.6.3
+* Support downloading multiple files, not avaiable when checking new version.
+
 ## v1.6.2.1
 * when parsing civitai url, remove query string by PR
 
diff --git a/scripts/ch_lib/model_action_civitai.py b/scripts/ch_lib/model_action_civitai.py
index 24c4f7a..c5fd88a 100644
--- a/scripts/ch_lib/model_action_civitai.py
+++ b/scripts/ch_lib/model_action_civitai.py
@@ -274,6 +274,48 @@ def get_model_info_by_url(model_url_or_id:str) -> tuple:
 
     return (model_info, model_name, model_type, subfolders, version_strs)
 
+# get version info by version string
+def get_ver_info_by_ver_str(version_str:str, model_info:dict) -> dict:
+    if not version_str:
+        util.printD("version_str is empty")
+        return
+    
+    if not model_info:
+        util.printD("model_info is None")
+        return
+    
+    # get version list
+    if "modelVersions" not in model_info.keys():
+        util.printD("modelVersions is not in model_info")
+        return
+
+    modelVersions = model_info["modelVersions"]
+    if not modelVersions:
+        util.printD("modelVersions is Empty")
+        return
+    
+    # find version by version_str
+    version = None
+    for ver in modelVersions:
+        # version name can not be used as id
+        # version id is not readable
+        # so , we use name_id as version string
+        ver_str = ver["name"]+"_"+str(ver["id"])
+        if ver_str == version_str:
+            # find version
+            version = ver
+
+    if not version:
+        util.printD("can not find version by version string: " + version_str)
+        return
+    
+    # get version id
+    if "id" not in version.keys():
+        util.printD("this version has no id")
+        return
+    
+    return version
+
 
 # get download url from model info by version string
 # return - (version_id, download_url)
@@ -387,24 +429,28 @@ def dl_model_by_input(model_info:dict, model_type:str, subfolder_str:str, versio
         util.printD(output)
         return output
     
-    # get download url
-    r = get_id_and_dl_url_by_version_str(version_str, model_info)
-    if not r:
-        output = "Fail to get download url, check console log for detail"
+    # get version info
+    ver_info = get_ver_info_by_ver_str(version_str, model_info)
+    if not ver_info:
+        output = "Fail to get version info, check console log for detail"
         util.printD(output)
         return output
     
-    version_id, url = r
-    if not version_id:
-        output = "Fail to get version id, check console log for detail"
-        util.printD(output)
-        return output
+    version_id = ver_info["id"]
+
+    # get all download url from files info
+    # some model versions have multiple files
+    download_urls = []
+    if "files" in ver_info.keys():
+        for file_info in ver_info["files"]:
+            if "downloadUrl" in file_info.keys():
+                download_urls.append(file_info["downloadUrl"])
+
+    if not len(download_urls):
+        if "downloadUrl" in ver_info.keys():
+            download_urls.append(ver_info["downloadUrl"])
+
 
-    if not url:
-        output = "Fail to get download url, check console log for detail"
-        util.printD(output)
-        return output
-    
     # check if this model is already existed
     r = civitai.search_local_model_info_by_version_id(model_folder, version_id)
     if r:
@@ -413,11 +459,19 @@ def dl_model_by_input(model_info:dict, model_type:str, subfolder_str:str, versio
         return output
     
     # download
-    filepath = downloader.dl(url, model_folder, None, None)
+    filepath = ""
+    for url in download_urls:
+        model_filepath = downloader.dl(url, model_folder, None, None)
+        if not model_filepath:
+            output = "Downloading failed, check console log for detail"
+            util.printD(output)
+            return output
+        
+        if url == ver_info["downloadUrl"]:
+            filepath = model_filepath
+
     if not filepath:
-        output = "Downloading failed, check console log for detail"
-        util.printD(output)
-        return output
+        filepath = model_filepath
     
     # get version info
     version_info = civitai.get_version_info_by_version_id(version_id)