From c8bb0ad75a1122afd415bd6a4a032898143252bf Mon Sep 17 00:00:00 2001
From: butaixianran <butaixianran@qq.com>
Date: Mon, 13 Mar 2023 12:09:58 +0800
Subject: [PATCH] support scanning only selected model types

---
 README.md                           |  3 +++
 scripts/civitai_helper.py           |  3 ++-
 scripts/lib/model_action_civitai.py | 20 ++++++++++++++++++--
 scripts/lib/util.py                 |  2 +-
 4 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 47681da..0103298 100644
--- a/README.md
+++ b/README.md
@@ -168,6 +168,9 @@ From v1.5, v1.x goes into maintenance phase.
 Enjoy!
 
 # Change Log
+## v1.5.2
+* Support scanning only selected model types. 
+
 ## v1.5.1
 * Force TI scanning delay 1 second to prevent from civitai treating this extension's requests as attacking.
 
diff --git a/scripts/civitai_helper.py b/scripts/civitai_helper.py
index dedb7ba..3ba6a76 100644
--- a/scripts/civitai_helper.py
+++ b/scripts/civitai_helper.py
@@ -87,6 +87,7 @@ def on_ui_tabs():
                 with gr.Row():
                     max_size_preview_ckb = gr.Checkbox(label="Download Max Size Preview", value=max_size_preview, elem_id="ch_max_size_preview_ckb")
                     skip_nsfw_preview_ckb = gr.Checkbox(label="SKip NSFW Preview images", value=skip_nsfw_preview, elem_id="ch_skip_nsfw_preview_ckb")
+                    scan_model_types_ckbg = gr.CheckboxGroup(choices=model_types, label="Model Types", value=model_types)
 
                 # with gr.Row():
                 scan_model_civitai_btn = gr.Button(value="Scan", variant="primary", elem_id="ch_scan_model_civitai_btn")
@@ -158,7 +159,7 @@ def on_ui_tabs():
 
         # ====events====
         # Scan Models for Civitai
-        scan_model_civitai_btn.click(model_action_civitai.scan_model, inputs=[max_size_preview_ckb, skip_nsfw_preview_ckb], outputs=scan_model_log_md)
+        scan_model_civitai_btn.click(model_action_civitai.scan_model, inputs=[scan_model_types_ckbg, max_size_preview_ckb, skip_nsfw_preview_ckb], outputs=scan_model_log_md)
 
         # Get Civitai Model Info by Model Page URL
         model_type_drop.change(get_model_names_by_input, inputs=[model_type_drop, empty_info_only_ckb], outputs=model_name_drop)
diff --git a/scripts/lib/model_action_civitai.py b/scripts/lib/model_action_civitai.py
index 2c3f101..e240708 100644
--- a/scripts/lib/model_action_civitai.py
+++ b/scripts/lib/model_action_civitai.py
@@ -10,14 +10,30 @@ from . import downloader
 
 # scan model to generate SHA256, then use this SHA256 to get model info from civitai
 # return output msg
-def scan_model(max_size_preview, skip_nsfw_preview):
+def scan_model(scan_model_types, max_size_preview, skip_nsfw_preview):
     util.printD("Start scan_model")
-
     output = ""
+
+    # check model types
+    if not scan_model_types:
+        output = "Model Types is None, can not scan."
+        util.printD(output)
+        return output
+    
+    model_types = []
+    # check type if it is a string
+    if type(scan_model_types) == str:
+        model_types.append(scan_model_types)
+    else:
+        model_types = scan_model_types
+    
     model_count = 0
     image_count = 0
     # scan_log = ""
     for model_type, model_folder in model.folders.items():
+        if model_type not in model_types:
+            continue
+
         util.printD("Scanning path: " + model_folder)
         for root, dirs, files in os.walk(model_folder, followlinks=True):
             for filename in files:
diff --git a/scripts/lib/util.py b/scripts/lib/util.py
index 29a6462..1d90155 100644
--- a/scripts/lib/util.py
+++ b/scripts/lib/util.py
@@ -4,7 +4,7 @@ import hashlib
 import requests
 import shutil
 
-version = "1.5.1"
+version = "1.5.2"
 
 # print for debugging
 def printD(msg):