Merge pull request #89 from Learwin/duplicates-check

Civitai duplicate check for existing records.

scripts/mo/data/sqlite_storage.py

@@ -229,6 +229,15 @@ class SQLiteStorage(Storage):
         for row in rows:
             result.append(map_row_to_record(row))
         return result
+    
+    def get_records_by_query(self, query: str) -> List:
+        cursor = self._connection().cursor()
+        cursor.execute(query)
+        rows = cursor.fetchall()
+        result = []
+        for row in rows:
+            result.append(map_row_to_record(row))
+        return result
 
     def add_record(self, record: Record):
         cursor = self._connection().cursor()

scripts/mo/environment.py

@@ -77,6 +77,7 @@ class Environment:
     card_height: Callable[[], str]
     is_debug_mode_enabled: Callable[[], bool]
     api_key: Callable[[], str]
+    check_duplicates: Callable[[], bool]
 
     def is_storage_initialized(self) -> bool:
         return hasattr(self, 'storage')

scripts/mo/ui_civitai_import.py

@@ -9,7 +9,7 @@ import requests
 from scripts.mo.data.storage import map_record_to_dict
 from scripts.mo.environment import env
 from scripts.mo.models import ModelType, Record
-from scripts.mo.ui_styled_html import alert_danger
+from scripts.mo.ui_styled_html import alert_danger, alert_warning
 from scripts.mo.utils import is_blank
 from scripts.mo.data.mapping_utils import create_version_dict
 
@@ -120,9 +120,19 @@ def _on_fetch_url_clicked(url):
         data = response.json()
         data_dict = create_model_dict(data)
 
+        duplicate_warning = ''
+        if env.check_duplicates():
+            civurl = f"https://civitai.com/models/{model_id}"
+            duplicate_candidates = env.storage.get_records_by_query(f"SELECT * FROM RECORD WHERE URL = '{civurl}'")
+            if len(duplicate_candidates) > 0:
+                duplicate_list = ['Fetched Model already has at least a version present as record']
+                for record in duplicate_candidates:
+                    duplicate_list.append(record.name)
+                duplicate_warning = alert_warning(duplicate_list)
+
         return [
             data_dict,
-            gr.HTML.update(value=''),
+            gr.HTML.update(value='' if duplicate_warning =='' else duplicate_warning),
             gr.Column.update(visible=True),
             *_create_ui_update(data_dict=data_dict, selected_version_id=selected_model_version_id)
         ]

scripts/mo/ui_styled_html.py

@@ -15,34 +15,34 @@ _NO_PREVIEW_LIGHT = 'file=extensions/sd-model-organizer/pic/no-preview-light.png
 
 def alert_danger(value) -> str:
     if isinstance(value, list):
-        text = "<br>".join(value)
+        text = "<br>".join([html.escape(s) for s in value])
     else:
-        text = value
-    return f'<div class="mo-alert mo-alert-danger">{html.escape(text)}</div>'
+        text = html.escape(value)
+    return f'<div class="mo-alert mo-alert-danger">{text}</div>'
 
 
 def alert_primary(value) -> str:
     if isinstance(value, list):
-        text = "<br>".join(value)
+        text = "<br>".join([html.escape(s) for s in value])
     else:
-        text = value
-    return f'<div class="mo-alert mo-alert-primary">{html.escape(text)}</div>'
+        text = html.escape(value)
+    return f'<div class="mo-alert mo-alert-primary">{text}</div>'
 
 
 def alert_success(value) -> str:
     if isinstance(value, list):
-        text = "<br>".join(value)
+        text = "<br>".join([html.escape(s) for s in value])
     else:
-        text = value
-    return f'<div class="mo-alert mo-alert-success">{html.escape(text)}</div>'
+        text = html.escape(value)
+    return f'<div class="mo-alert mo-alert-success">{text}</div>'
 
 
 def alert_warning(value) -> str:
     if isinstance(value, list):
-        text = "<br>".join(value)
+        text = "<br>".join([html.escape(s) for s in value])
     else:
-        text = value
-    return f'<div class="mo-alert mo-alert-warning">{html.escape(text)}</div>'
+        text = html.escape(value)
+    return f'<div class="mo-alert mo-alert-warning">{text}</div>'
 
 
 def _limit_description(text):

scripts/model_organizer.py

@@ -134,6 +134,12 @@ env.api_key = (
     else ""
 )
 
+env.check_duplicates = (
+    lambda: shared.opts.mo_check_duplicates
+    if hasattr(shared.opts, 'mo_check_duplicates')
+    else ""
+)
+
 env.model_path = (
     lambda: shared.opts.mo_model_path
     if hasattr(shared.opts, 'mo_model_path') and shared.opts.mo_model_path
@@ -198,6 +204,7 @@ def on_ui_settings():
         'mo_prefill_neg_prompt': OptionInfo(True, 'When creating a record based on local file, automatically import the added negative prompts'),
         'mo_autobind_file': OptionInfo(True, 'Automatically bind record to local file'),
         'mo_api_key': OptionInfo("", "Civitai API Key. Create an API key under 'https://civitai.com/user/account' all the way at the bottom. Don't share the token!"),
+        'mo_check_duplicates': OptionInfo(False, "Should a duplicate check be performed, upon fetching a file from Civitai"),
     }
 
     dir_opts = {