Merge pull request #9 from newtextdoc1111/feature/e621-support

Merge feature/e621-support to dev

README.md

@@ -9,8 +9,7 @@
 
 ## Overview
 
-**ComfyUI-Autocomplete-Plus** is a custom node that provides multiple input assistance features for any text area in [ComfyUI](https://github.com/comfyanonymous/ComfyUI). Currently, it mainly supports Danbooru-style tag input.
-
+**ComfyUI-Autocomplete-Plus** is a custom node that provides multiple input assistance features for any text area in [ComfyUI](https://github.com/comfyanonymous/ComfyUI). Currently, it supports Danbooru and e621 tags (e621 does not support some functions).
 ## Features
 
 - **:zap:No setup required**: Automatically downloads CSV data optimized for Danbooru tags.
@@ -38,6 +37,7 @@ When you type in a text input area, tags that partially match the text are displ
 - Tag aliases are also included in the search. Japanese hiragana and katakana are searched without distinction.
 - Tags are color-coded by category. The color-coding rules are the same as Danbooru.
 - Tags that have already been entered are displayed grayed out.
+- You can display Danbooru and e621 tags at the same time. You can also change the priority from the settings.
 
 ## Related Tags
 
@@ -76,6 +76,11 @@ This is a CSV file for related tag calculation, recording tag pairs and their co
 Tag pairs are further filtered from the tag information CSV under the following conditions:
 - Co-occurrence count of 100 or more
 
+### e621 CSV
+
+Currently, automatic download of CSV for e621 is not supported, so please manually place a CSV with the same structure as `danbooru_tags.csv` in the data folder with the name `e621_tags.csv`.
+Also, displaying related tags is not supported.
+
 ### User CSV
 
 Users can also use their own CSV files. For example, frequently used meta tags can be added to the autocomplete suggestions by placing a file named `danbooru_tags_meta.csv` in the `data` folder.
@@ -97,6 +102,15 @@ worst_quality,5,9999999,
 
 ## Settings
 
+### Tag Source
+
+> [!TIP]
+> The source of tag data such as Danbooru or e621 is called the "tag source".
+
+- **Autocomplete Tag Source**: The tag source to display in the autocomplete suggestions. Select "all" to display all loaded tag sources.
+- **Primary source for 'all' Source**: When `Autocomplete Tag Source` is set to "all", the tag source specified here will be displayed with priority.
+- **Tag Source Icon Position**: Where to display the icon of the tag source. Select "hidden" to hide it.
+
 ### Autocomplete
 
 - **Enable Autocomplete**: Enable/disable the autocomplete feature.
@@ -107,7 +121,7 @@ worst_quality,5,9999999,
 - **Enable Related Tags**: Enable/disable the related tags feature.
 - **Max related tags**: Maximum number of related tags to display.
 - **Default Display Position**: Default display position when ComfyUI starts.
-- **Related Tags Trigger Mode**: Trigger for displaying related tags (click, Ctrl+click).
+- **Related Tags Trigger Mode**: Which action will trigger displaying related tags (click only, Ctrl+click)
 
 
 ## Known Issues

docs/README_jp.md

@@ -7,7 +7,7 @@
 
 ## 概要
 
-**ComfyUI-Autocomplete-Plus** は、[ComfyUI](https://github.com/comfyanonymous/ComfyUI) の任意のテキストエリアに複数の入力支援機能を提供するカスタムノードです。現在は主にDanbooruスタイルのタグ入力に対応しています。
+**ComfyUI-Autocomplete-Plus** は、[ComfyUI](https://github.com/comfyanonymous/ComfyUI) の任意のテキストエリアに複数の入力支援機能を提供するカスタムノードです。現在はDanbooruとe621のタグに対応しています（e621は一部の機能が未対応です）。
 
 ## 特徴
 
@@ -36,6 +36,7 @@
 - タグのエイリアスも検索対象に含まれます。日本語のひらがな、カタカナは区別せず検索されます
 - タグのカテゴリ毎に色分けされます。色分けのルールは Danbooru と同じです
 - 入力済みのタグはグレーアウトで表示されます
+- Danbooruとe621のタグを同時に表示出来ます。設定から優先順位を変更できます
 
 ## 関連タグ
 
@@ -74,9 +75,14 @@
 タグペアはタグ情報 CSV から、さらに以下の条件でフィルタリングされています。
 - 共起回数が100件以上
 
+### e621 CSV
+
+現在、e621用 CSV の自動ダウンロードは未対応なため `danbooru_tags.csv` と同じ構造の CSV を `e621_tags.csv` という名前でデータフォルダーに手動配置してください。  
+また、関連タグ表示も同様に未対応です。
+
 ### ユーザーCSV
 
-ユーザーが自身で用意したCSVを使用することも可能です。例として、よく使うメタタグを `danbooru_tags_meta.csv` の名前で `data` フォルダーに配置することでオートコンプリート候補に追加できます。  
+ユーザーが自身で用意した CSV を使用することも可能です。例として、よく使うメタタグを `danbooru_tags_meta.csv` の名前で `data` フォルダーに配置することでオートコンプリート候補に追加できます。  
 ヘッダー行はなくても構いません。反映にはブラウザのリロードが必要です。
 
 **メタタグの例**
@@ -95,6 +101,15 @@ worst_quality,5,9999999,
 
 ## 設定
 
+### タグソース
+
+> [!TIP]
+> Danbooruやe621等のタグデータの提供元を「タグソース」と呼びます
+
+- **Autocomplete Tag Source**: オートコンプリート候補に表示するタグソース。「all」を選択するとロード済みの全てのタグソースを表示します
+- **Primary source for 'all' Source**: `Autocomplete Tag Source` が「all」のとき、ここで指定したタグソースが優先して表示されます
+- **Tag Source Icon Position**: タグの情報源のアイコンをどの位置に表示するか。「hidden」を選択すると非表示になります
+
 ### オートコンプリート
 
 - **Enable Autocomplete**: オートコンプリート機能の有効化/無効化
@@ -105,7 +120,7 @@ worst_quality,5,9999999,
 - **Enable Related Tags**: 関連タグ機能の有効化/無効化
 - **Max related tags**: 関連タグの最大表示件数
 - **Default Display Position**: ComfyUI起動時のデフォルト表示位置
-- **Related Tags Trigger Mode** : 関連タグの表示トリガー（クリック、Ctrl+クリック）
+- **Related Tags Trigger Mode** : 関連タグを表示する際、どの操作をトリガーとするか（クリックのみ、Ctrl+クリック）
 
 ## 既知の問題
 

modules/api.py

@@ -8,43 +8,64 @@ from aiohttp import web
 # os.path.join(..., '..', 'data') goes up one level and then into 'data'
 DATA_DIR = os.path.join(os.path.dirname(__file__), '..', 'data')
 
-BASE_URL = '/autocomplete-plus/csv'
+DANBOORU_PREFIX = 'danbooru'
+E621_PREFIX = 'e621'
 
-TAGS_BASE_FILE = 'danbooru_tags.csv'
-COOCCURRENCE_BASE_FILE = 'danbooru_tags_cooccurrence.csv'
+TAGS_SUFFIX = 'tags'
+COOCCURRENCE_SUFFIX = 'tags_cooccurrence'
 
 def get_csv_file_status():
     """
     Returns a dictionary of csv file statuses.
     """
-    tags_base_exists = os.path.exists(os.path.join(DATA_DIR, TAGS_BASE_FILE))
-    cooccurrence_base_exists = os.path.exists(os.path.join(DATA_DIR, COOCCURRENCE_BASE_FILE))
 
-    tags_extra_files = []
-    cooccurrence_extra_files = []
+    data = {
+        DANBOORU_PREFIX: {
+            'base_tags': False,
+            'extra_tags': [],
+            'base_cooccurrence': False,
+            'extra_cooccurrence': [],
+        },
+        E621_PREFIX: {
+            'base_tags': False,
+            'extra_tags': [],
+            'base_cooccurrence': False,
+            'extra_cooccurrence': [],
+        }
+    }
 
-    all_csv_files = [f for f in os.listdir(DATA_DIR) if f.endswith('.csv')]
-    if len(all_csv_files) == 0:
-        print("[Autocomplete-Plus] No CSV files found in the data directory.")
+    for prefix in [DANBOORU_PREFIX, E621_PREFIX]:
+        base_tags_file = f"{prefix}_{TAGS_SUFFIX}.csv"
+        base_cooccurrence_file = f"{prefix}_{COOCCURRENCE_SUFFIX}.csv"
 
-    # Create extra CSV files list
-    for filename in all_csv_files:
-        if filename in [TAGS_BASE_FILE, COOCCURRENCE_BASE_FILE]:
-            continue # Skip base files
-        if 'cooccurrence' in filename.lower():
-            cooccurrence_extra_files.append(filename)
-        elif 'tag' in filename.lower():
-            tags_extra_files.append(filename)
+        tags_base_exists = os.path.exists(os.path.join(DATA_DIR, base_tags_file))
+        cooccurrence_base_exists = os.path.exists(os.path.join(DATA_DIR, base_cooccurrence_file))
 
-    # Return the lists of extra files
-    return {
-        'danbooru':{
-            'base_tags': tags_base_exists, # exists
+        tags_extra_files = []
+        cooccurrence_extra_files = []
+
+        all_csv_files = [f for f in os.listdir(DATA_DIR) if f.startswith(prefix) and f.endswith('.csv')]
+        if len(all_csv_files) == 0:
+            print("[Autocomplete-Plus] No CSV files found in the data directory.")
+
+        # Create extra CSV files list
+        for filename in all_csv_files:
+            if filename in [base_tags_file, base_cooccurrence_file]:
+                continue # Skip base files
+            if COOCCURRENCE_SUFFIX in filename.lower():
+                cooccurrence_extra_files.append(filename)
+            elif TAGS_SUFFIX in filename.lower():
+                tags_extra_files.append(filename)
+
+        data[prefix] = {
+            'base_tags': tags_base_exists,
             'extra_tags': tags_extra_files,
             'base_cooccurrence': cooccurrence_base_exists,
             'extra_cooccurrence': cooccurrence_extra_files,
         }
-    }
+
+    # Return the lists of extra files
+    return data
 
 # --- API Endpoints ---
 
@@ -55,41 +76,58 @@ async def get_csv_list(_request):
     base files: file exists boolean
     extra files: count of extra files
     """
-    extra_csv_files = get_csv_file_status()
+    csv_file_status = get_csv_file_status()
 
     response = {
-        'danbooru': {
-            'base_tags': extra_csv_files['danbooru']['base_tags'],
-            'extra_tags': len(extra_csv_files['danbooru']['extra_tags']),
-            'base_cooccurrence': extra_csv_files['danbooru']['base_cooccurrence'],
-            'extra_cooccurrence': len(extra_csv_files['danbooru']['extra_cooccurrence']),
+        DANBOORU_PREFIX: {
+            'base_tags': csv_file_status[DANBOORU_PREFIX]['base_tags'],
+            'extra_tags': len(csv_file_status[DANBOORU_PREFIX]['extra_tags']),
+            'base_cooccurrence': csv_file_status[DANBOORU_PREFIX]['base_cooccurrence'],
+            'extra_cooccurrence': len(csv_file_status[DANBOORU_PREFIX]['extra_cooccurrence']),
+        },
+        E621_PREFIX: {
+            'base_tags': csv_file_status[E621_PREFIX]['base_tags'],
+            'extra_tags': len(csv_file_status[E621_PREFIX]['extra_tags']),
+            'base_cooccurrence': csv_file_status[E621_PREFIX]['base_cooccurrence'],
+            'extra_cooccurrence': len(csv_file_status[E621_PREFIX]['extra_cooccurrence']),
         }
     }
     return web.json_response(response)
 
-@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/tags/base')
-async def get_base_tags_file(_request):
+@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/{source}/{suffix}/base')
+async def get_base_tags_file(request):
     """
     Returns the base tags CSV file.
     """
-    file_path = os.path.join(DATA_DIR, TAGS_BASE_FILE)
+    source = str(request.match_info['source'])
+    suffix = str(request.match_info['suffix'])
+    if source not in [DANBOORU_PREFIX, E621_PREFIX] or suffix not in [TAGS_SUFFIX, COOCCURRENCE_SUFFIX]:
+        return web.json_response({"error": "Invalid tag source or suffix"}, status=400)
+    
+    file_path = os.path.join(DATA_DIR, f"{source}_{suffix}.csv")
     if not os.path.exists(file_path):
         return web.json_response({"error": "Base tags file not found"}, status=404)
+
     return web.FileResponse(file_path)
 
-@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/tags/extra/{index}')
+@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/{source}/{suffix}/extra/{index}')
 async def get_extra_tags_file(request):
     """
     Returns the extra tags CSV file at the specified index.
     """
     try:
-        extra_csv_files = get_csv_file_status()
+        csv_file_status = get_csv_file_status()
 
+        source = str(request.match_info['source'])
+        suffix = str(request.match_info['suffix'])
+        if source not in [DANBOORU_PREFIX, E621_PREFIX] or suffix not in [TAGS_SUFFIX, COOCCURRENCE_SUFFIX]:
+            return web.json_response({"error": "Invalid tag source or suffix"}, status=400)
+    
         index = int(request.match_info['index'])
-        if index < 0 or index >= len(extra_csv_files['danbooru']['extra_tags']):
+        if index < 0 or index >= len(csv_file_status[source][f'extra_{suffix}']):
             return web.json_response({"error": "Invalid index"}, status=404)
 
-        file_path = os.path.join(DATA_DIR, extra_csv_files['danbooru']['extra_tags'][index])
+        file_path = os.path.join(DATA_DIR, csv_file_status[source][f'extra_{suffix}'][index])
         if not os.path.exists(file_path):
             return web.json_response({"error": "Extra tags file not found"}, status=404)
 
@@ -97,34 +135,3 @@ async def get_extra_tags_file(request):
 
     except ValueError:
         return web.json_response({"error": "Invalid index format"}, status=400)
-
-@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/cooccurrence/base')
-async def get_base_cooccurrence_file(_request):
-    """
-    Returns the base cooccurrence CSV file.
-    """
-    file_path = os.path.join(DATA_DIR, COOCCURRENCE_BASE_FILE)
-    if not os.path.exists(file_path):
-        return web.json_response({"error": "Base cooccurrence file not found"}, status=404)
-    return web.FileResponse(file_path)
-
-@server.PromptServer.instance.routes.get('/autocomplete-plus/csv/cooccurrence/extra/{index}')
-async def get_extra_cooccurrence_file(request):
-    """
-    Returns the extra cooccurrence CSV file at the specified index.
-    """
-    try:
-        extra_csv_files = get_csv_file_status()
-
-        index = int(request.match_info['index'])
-        if index < 0 or index >= len(extra_csv_files['cooccurrence']):
-            return web.json_response({"error": "Invalid index"}, status=404)
-
-        file_path = os.path.join(DATA_DIR, extra_csv_files['danbooru']['extra_cooccurrence'][index])
-        if not os.path.exists(file_path):
-            return web.json_response({"error": "Extra cooccurrence file not found"}, status=404)
-
-        return web.FileResponse(file_path)
-
-    except ValueError:
-        return web.json_response({"error": "Invalid index format"}, status=400)

web/css/autocomplete-plus.css

@@ -1,39 +1,29 @@
 /* Light Theme */
  body {
-  --autocomplete-plus-text-color-cat-general: var(--p-blue-600);
-  --autocomplete-plus-text-color-cat-artist: var(--p-red-600);
-  --autocomplete-plus-text-color-cat-copyright: var(--p-purple-600);
-  --autocomplete-plus-text-color-cat-character: var(--p-green-600);
-  --autocomplete-plus-text-color-cat-meta: var(--p-yellow-600);
-  --autocomplete-plus-text-color-cat-unknown: var(--p-gray-600);
+  --autocomplete-plus-text-color-cat-blue: var(--p-blue-600);
+  --autocomplete-plus-text-color-cat-red: var(--p-red-600);
+  --autocomplete-plus-text-color-cat-purple: var(--p-purple-600);
+  --autocomplete-plus-text-color-cat-green: var(--p-green-600);
+  --autocomplete-plus-text-color-cat-yellow: var(--p-yellow-600);
+  --autocomplete-plus-text-color-cat-gray: var(--p-gray-600);
+  --autocomplete-plus-text-color-cat-sky: var(--p-sky-500);
+  --autocomplete-plus-text-color-cat-orange: var(--p-orange-600);
+  --autocomplete-plus-text-color-cat-white: var(--p-neutral-700);
   --autocomplete-plus-text-color-disabled: var(--p-gray-400);
-
-  /* Selected color */
-  --autocomplete-plus-text-color-cat-general-selected: var(--p-blue-200);
-  --autocomplete-plus-text-color-cat-artist-selected: var(--p-red-200);
-  --autocomplete-plus-text-color-cat-copyright-selected: var(--p-purple-200);
-  --autocomplete-plus-text-color-cat-character-selected: var(--p-green-200);
-  --autocomplete-plus-text-color-cat-meta-selected: var(--p-yellow-200);
-  --autocomplete-plus-text-color-cat-unknown-selected: var(--p-gray-200);
 }
 
 /* Dark Theme */
 body.dark-theme {
-  --autocomplete-plus-text-color-cat-general: var(--p-blue-400);
-  --autocomplete-plus-text-color-cat-artist: var(--p-red-400);
-  --autocomplete-plus-text-color-cat-copyright: var(--p-purple-400);
-  --autocomplete-plus-text-color-cat-character: var(--p-green-400);
-  --autocomplete-plus-text-color-cat-meta: var(--p-yellow-400);
-  --autocomplete-plus-text-color-cat-unknown: var(--p-gray-500);
+  --autocomplete-plus-text-color-cat-blue: var(--p-blue-400);
+  --autocomplete-plus-text-color-cat-red: var(--p-red-400);
+  --autocomplete-plus-text-color-cat-purple: var(--p-purple-400);
+  --autocomplete-plus-text-color-cat-green: var(--p-green-400);
+  --autocomplete-plus-text-color-cat-yellow: var(--p-yellow-400);
+  --autocomplete-plus-text-color-cat-gray: var(--p-gray-400);
+  --autocomplete-plus-text-color-cat-sky: var(--p-sky-300);
+  --autocomplete-plus-text-color-cat-orange: var(--p-orange-400);
+  --autocomplete-plus-text-color-cat-white: var(--p-neutral-200);
   --autocomplete-plus-text-color-disabled: var(--p-gray-500);
-
-  /* Selected color */
-  --autocomplete-plus-text-color-cat-general-selected: var(--p-blue-200);
-  --autocomplete-plus-text-color-cat-artist-selected: var(--p-red-200);
-  --autocomplete-plus-text-color-cat-copyright-selected: var(--p-purple-200);
-  --autocomplete-plus-text-color-cat-character-selected: var(--p-green-200);
-  --autocomplete-plus-text-color-cat-meta-selected: var(--p-yellow-200);
-  --autocomplete-plus-text-color-cat-unknown-selected: var(--p-gray-200);
 }
 
 #autocomplete-plus-root {
@@ -62,28 +52,60 @@ body.dark-theme {
   grid-template-columns: subgrid;
 }
 
-.autocomplete-plus-item[data-tag-category="general"] {
-  color: var(--autocomplete-plus-text-color-cat-general);
+.autocomplete-plus-item.danbooru[data-tag-category="general"] {
+  color: var(--autocomplete-plus-text-color-cat-blue);
 }
 
-.autocomplete-plus-item[data-tag-category="artist"] {
-  color: var(--autocomplete-plus-text-color-cat-artist);
+.autocomplete-plus-item.danbooru[data-tag-category="artist"] {
+  color: var(--autocomplete-plus-text-color-cat-red);
 }
 
-.autocomplete-plus-item[data-tag-category="copyright"] {
-  color: var(--autocomplete-plus-text-color-cat-copyright);
+.autocomplete-plus-item.danbooru[data-tag-category="copyright"] {
+  color: var(--autocomplete-plus-text-color-cat-purple);
 }
 
-.autocomplete-plus-item[data-tag-category="character"] {
-  color: var(--autocomplete-plus-text-color-cat-character);
+.autocomplete-plus-item.danbooru[data-tag-category="character"] {
+  color: var(--autocomplete-plus-text-color-cat-green);
 }
 
-.autocomplete-plus-item[data-tag-category="meta"] {
-  color: var(--autocomplete-plus-text-color-cat-meta);
+.autocomplete-plus-item.danbooru[data-tag-category="meta"] {
+  color: var(--autocomplete-plus-text-color-cat-yellow);
 }
 
-.autocomplete-plus-item[data-tag-category="unknown"] {
-  color: var(--autocomplete-plus-text-color-cat-unknown);
+.autocomplete-plus-item.danbooru[data-tag-category="unknown"] {
+  color: var(--autocomplete-plus-text-color-cat-gray);
+}
+
+.autocomplete-plus-item.e621[data-tag-category="general"] {
+  color: var(--autocomplete-plus-text-color-cat-sky);
+}
+
+.autocomplete-plus-item.e621[data-tag-category="artist"] {
+  color: var(--autocomplete-plus-text-color-cat-orange);
+}
+
+.autocomplete-plus-item.e621[data-tag-category="copyright"] {
+  color: var(--autocomplete-plus-text-color-cat-purple);
+}
+
+.autocomplete-plus-item.e621[data-tag-category="character"] {
+  color: var(--autocomplete-plus-text-color-cat-green);
+}
+
+.autocomplete-plus-item.e621[data-tag-category="species"] {
+  color: var(--autocomplete-plus-text-color-cat-red);
+}
+
+.autocomplete-plus-item.e621[data-tag-category="meta"] {
+  color: var(--autocomplete-plus-text-color-cat-white);
+}
+
+.autocomplete-plus-item.e621[data-tag-category="lore"] {
+  color: var(--autocomplete-plus-text-color-cat-red);
+}
+
+.autocomplete-plus-item.e621[data-tag-category="unknown"] {
+  color: var(--autocomplete-plus-text-color-cat-purple);
 }
 
 .autocomplete-plus-item span {
@@ -113,22 +135,31 @@ body.dark-theme {
   background-color: var(--comfy-hover-bg);
 }
 
-.autocomplete-plus-tag-name {
-  padding: 8px;
+.autocomplete-plus-item .autocomplete-plus-tag-name {
   white-space: nowrap;
 }
 
+.autocomplete-plus-item .autocomplete-plus-tag-icon-svg {
+  width: 1em;
+  height: 1em;
+  vertical-align: middle;
+}
+
 .autocomplete-plus-tag-name.autocomplete-plus-already-exists  {
   color: var(--autocomplete-plus-text-color-disabled);
 }
 
-.autocomplete-plus-alias {
+.autocomplete-plus-item .autocomplete-plus-alias {
   overflow: hidden;
   color: var(--descrip-text);
   text-overflow: ellipsis;
   white-space: nowrap;
 }
 
+.autocomplete-plus-item .autocomplete-plus-tag-count {
+  text-align: right;
+}
+
 /* Related Tags UI Styles */
 #related-tags-root {
   position: absolute;
@@ -163,9 +194,16 @@ body.dark-theme {
   font-weight: bold;
 }
 
+.related-tags-header-tag-name .autocomplete-plus-tag-icon-svg {
+  width: 1em;
+  height: 1em;
+  vertical-align: middle;
+}
+
 .related-tags-header-controls {
   display: flex;
   align-items: center;
+  margin-left: 16px;
   gap: 4px;
 }
 
@@ -210,28 +248,60 @@ body.dark-theme {
   grid-template-columns: subgrid;
 }
 
-.related-tag-item[data-tag-category="general"] {
-  color: var(--autocomplete-plus-text-color-cat-general);
+.related-tag-item.danbooru[data-tag-category="general"] {
+  color: var(--autocomplete-plus-text-color-cat-blue);
 }
 
-.related-tag-item[data-tag-category="artist"] {
-  color: var(--autocomplete-plus-text-color-cat-artist);
+.related-tag-item.danbooru[data-tag-category="artist"] {
+  color: var(--autocomplete-plus-text-color-cat-red);
 }
 
-.related-tag-item[data-tag-category="copyright"] {
-  color: var(--autocomplete-plus-text-color-cat-copyright);
+.related-tag-item.danbooru[data-tag-category="copyright"] {
+  color: var(--autocomplete-plus-text-color-cat-purple);
 }
 
-.related-tag-item[data-tag-category="character"] {
-  color: var(--autocomplete-plus-text-color-cat-character);
+.related-tag-item.danbooru[data-tag-category="character"] {
+  color: var(--autocomplete-plus-text-color-cat-green);
 }
 
-.related-tag-item[data-tag-category="meta"] {
-  color: var(--autocomplete-plus-text-color-cat-meta);
+.related-tag-item.danbooru[data-tag-category="meta"] {
+  color: var(--autocomplete-plus-text-color-cat-yellow);
 }
 
-.related-tag-item[data-tag-category="unknown"] {
-  color: var(--autocomplete-plus-text-color-cat-unknown);
+.related-tag-item.danbooru[data-tag-category="unknown"] {
+  color: var(--autocomplete-plus-text-color-cat-gray);
+}
+
+.related-tag-item.e621[data-tag-category="general"] {
+  color: var(--autocomplete-plus-text-color-cat-sky);
+}
+
+.related-tag-item.e621[data-tag-category="artist"] {
+  color: var(--autocomplete-plus-text-color-cat-orange);
+}
+
+.related-tag-item.e621[data-tag-category="copyright"] {
+  color: var(--autocomplete-plus-text-color-cat-purple);
+}
+
+.related-tag-item.e621[data-tag-category="character"] {
+  color: var(--autocomplete-plus-text-color-cat-green);
+}
+
+.related-tag-item.e621[data-tag-category="species"] {
+  color: var(--autocomplete-plus-text-color-cat-red);
+}
+
+.related-tag-item.e621[data-tag-category="meta"] {
+  color: var(--autocomplete-plus-text-color-cat-white);
+}
+
+.related-tag-item.e621[data-tag-category="lore"] {
+  color: var(--autocomplete-plus-text-color-cat-red);
+}
+
+.related-tag-item.e621[data-tag-category="unknown"] {
+  color: var(--autocomplete-plus-text-color-cat-gray);
 }
 
 /* stylelint-disable-next-line no-descending-specificity */
@@ -261,23 +331,22 @@ body.dark-theme {
   font-weight: bold;
 }
 
-.related-tag-name {
-  padding: 8px;
+.related-tag-item .related-tag-name {
   white-space: nowrap;
 }
 
-.related-tag-name.related-tag-already-exists  {
+.related-tag-item .related-tag-name.related-tag-already-exists  {
   color: var(--autocomplete-plus-text-color-disabled);
 }
 
-.related-tag-alias {
+.related-tag-item .related-tag-alias {
   overflow: hidden;
   color: var(--descrip-text);
   text-overflow: ellipsis;
   white-space: nowrap;
 }
 
-.related-tag-similarity {
+.related-tag-item .related-tag-similarity {
   color: var(--descrip-text);
   text-align: right;
   white-space: nowrap;
@@ -294,4 +363,4 @@ body.dark-theme {
   color: var(--error-text);
   font-style: italic;
   text-align: center;
-}
+}

web/js/autocomplete.js

@@ -1,7 +1,8 @@
 import {
     TagCategory,
     TagData,
-    autoCompleteData
+    autoCompleteData,
+    getEnabledTagSourceInPriorityOrder
 } from './data.js';
 import {
     formatCountHumanReadable,
@@ -12,7 +13,8 @@ import {
     normalizeTagToSearch,
     extractTagsFromTextArea,
     getCurrentTagRange,
-    getViewportMargin
+    getViewportMargin,
+    IconSvgHtmlString
 } from './utils.js';
 import { settingValues } from './settings.js';
 
@@ -99,54 +101,59 @@ function searchCompletionCandidates(textareaElement) {
         queryVariations.add(hiraQuery);
     }
 
-    // Search in sortedTags (already sorted by count)
-    for (const tagData of autoCompleteData.sortedTags) {
-        let matched = false;
-        let isExactMatch = false;
-        let matchedAlias = null;
+    const sources = getEnabledTagSourceInPriorityOrder();
+    for (const source of sources) {
+        // Search in sortedTags (already sorted by count)
+        for (const tagData of autoCompleteData[source].sortedTags) {
+            let matched = false;
+            let isExactMatch = false;
+            let matchedAlias = null;
 
-        // Check primary tag against all variations for exact/partial match
-        const tagMatch = matchWord(tagData.tag, queryVariations);
-        matched = tagMatch.matched;
-        isExactMatch = tagMatch.isExactMatch;
+            // Check primary tag against all variations for exact/partial match
+            const tagMatch = matchWord(tagData.tag, queryVariations);
+            matched = tagMatch.matched;
+            isExactMatch = tagMatch.isExactMatch;
 
-        // If primary tag didn't match, check aliases against all variations
-        if (!matched && tagData.alias && Array.isArray(tagData.alias) && tagData.alias.length > 0) {
-            for (const alias of tagData.alias) {
-                const lowerAlias = alias.toLowerCase();
-                const aliasMatch = matchWord(lowerAlias, queryVariations);
-                if (aliasMatch.matched) {
-                    matched = true;
-                    isExactMatch = aliasMatch.isExactMatch;
-                    matchedAlias = alias;
-                    break;
+            // If primary tag didn't match, check aliases against all variations
+            if (!matched && tagData.alias && Array.isArray(tagData.alias) && tagData.alias.length > 0) {
+                for (const alias of tagData.alias) {
+                    const lowerAlias = alias.toLowerCase();
+                    const aliasMatch = matchWord(lowerAlias, queryVariations);
+                    if (aliasMatch.matched) {
+                        matched = true;
+                        isExactMatch = aliasMatch.isExactMatch;
+                        matchedAlias = alias;
+                        break;
+                    }
                 }
             }
-        }
 
-        // Add candidate if matched and not already added
-        if (matched && !addedTags.has(tagData.tag)) {
-            // Add to exact matches or partial matches based on match type
-            if (isExactMatch) {
-                exactMatches.push(tagData);
-            } else {
-                partialMatches.push(tagData);
-            }
+            const tagSetKey = tagData.tag;
 
-            addedTags.add(tagData.tag);
-
-            // Check if we've reached the maximum suggestions limit combining both arrays
-            if (exactMatches.length + partialMatches.length >= settingValues.maxSuggestions) {
-                // Return the combined results, prioritizing exact matches
-                const result = [...exactMatches, ...partialMatches].slice(0, settingValues.maxSuggestions);
-
-                if (settingValues._logprocessingTime) {
-                    const endTime = performance.now();
-                    const duration = endTime - startTime;
-                    console.debug(`[Autocomplete-Plus] Search for "${partialTag}" took ${duration.toFixed(2)}ms. Found ${result.length} candidates (max reached).`);
+            // Add candidate if matched and not already added
+            if (matched && !addedTags.has(tagSetKey)) {
+                // Add to exact matches or partial matches based on match type
+                if (isExactMatch) {
+                    exactMatches.push(tagData);
+                } else {
+                    partialMatches.push(tagData);
                 }
 
-                return result; // Early exit
+                addedTags.add(tagSetKey);
+
+                // Check if we've reached the maximum suggestions limit combining both arrays
+                if (exactMatches.length + partialMatches.length >= settingValues.maxSuggestions) {
+                    // Return the combined results, prioritizing exact matches
+                    const result = [...exactMatches, ...partialMatches].slice(0, settingValues.maxSuggestions);
+
+                    if (settingValues._logprocessingTime) {
+                        const endTime = performance.now();
+                        const duration = endTime - startTime;
+                        console.debug(`[Autocomplete-Plus] Search for "${partialTag}" took ${duration.toFixed(2)}ms. Found ${result.length} candidates (max reached).`);
+                    }
+
+                    return result; // Early exit
+                }
             }
         }
     }
@@ -268,6 +275,13 @@ class AutocompleteUI {
         this.root = document.createElement('div'); // Use table instead of div
         this.root.id = 'autocomplete-plus-root';
 
+        // Create svg icon element as definition
+        this.iconSvgDef = document.createElement('div');
+        this.iconSvgDef.style.position = 'absolute';
+        this.iconSvgDef.style.display = 'none';
+        this.iconSvgDef.innerHTML = IconSvgHtmlString;
+        this.root.appendChild(this.iconSvgDef);
+
         this.tagsList = document.createElement('div');
         this.tagsList.id = 'autocomplete-plus-list';
         this.root.appendChild(this.tagsList);
@@ -382,17 +396,24 @@ class AutocompleteUI {
      * @param {boolean} isExisting
      */
     #createTagElement(tagData, isExisting) {
-        const categoryText = TagCategory[tagData.category] || "unknown";
+        const categoryText = TagCategory[tagData.source][tagData.category] || "unknown";
 
         const tagRow = document.createElement('div');
-        tagRow.classList.add('autocomplete-plus-item');
+        tagRow.classList.add('autocomplete-plus-item', tagData.source);
         tagRow.dataset.tag = tagData.tag;
         tagRow.dataset.tagCategory = categoryText;
 
-        // Tag name
+        // Tag icon and name
+        const tagSourceIconHtml = `<svg class="autocomplete-plus-tag-icon-svg"><use xlink:href="#autocomplete-plus-icon-${tagData.source}"></use></svg>`;
         const tagName = document.createElement('span');
-        tagName.classList.add('autocomplete-plus-tag-name');
-        tagName.textContent = tagData.tag;
+        tagName.className = 'autocomplete-plus-tag-name';
+        if (settingValues.tagSourceIconPosition == 'hidden') {
+            tagName.textContent = tagData.tag;
+        } else {
+            tagName.innerHTML = settingValues.tagSourceIconPosition == 'left'
+                ? `${tagSourceIconHtml} ${tagData.tag}`
+                : `${tagData.tag} ${tagSourceIconHtml}`;
+        }
 
         // grayout tag name if it already exists
         if (isExisting) {
@@ -401,7 +422,7 @@ class AutocompleteUI {
 
         // Alias
         const alias = document.createElement('span');
-        alias.classList.add('autocomplete-plus-alias');
+        alias.className = 'autocomplete-plus-alias';
 
         // Display alias if available
         if (tagData.alias && tagData.alias.length > 0) {
@@ -414,10 +435,11 @@ class AutocompleteUI {
         const category = document.createElement('span');
         category.className = `autocomplete-plus-category`;
         category.textContent = `${categoryText.substring(0, 2)}`;
+        category.title = categoryText; // Full category on hover
 
         // Count
         const tagCount = document.createElement('span');
-        category.className = `autocomplete-plus-tag-count`;
+        tagCount.className = `autocomplete-plus-tag-count`;
         tagCount.textContent = formatCountHumanReadable(tagData.count);
 
         tagRow.appendChild(tagName);
@@ -452,7 +474,6 @@ class AutocompleteUI {
         const viewportHeight = window.innerHeight;
         const margin = getViewportMargin();
 
-        const targetRect = this.target.getBoundingClientRect();
         const targetElmOffset = this.#calculateElementOffset(this.target);
 
         const { top: caretTop, left: caretLeft, lineHeight: caretLineHeight } = this.#getCaretCoordinates(this.target);

web/js/data.js

@@ -2,15 +2,35 @@ import { settingValues } from "./settings.js";
 
 // --- Constants ---
 
-// Tag categories for display
-export const TagCategory = [
-    'general',
-    'artist',
-    'unused',
-    'copyright',
-    'character',
-    'meta',
-]
+// Tag data sources
+export const TagSource = {
+    Danbooru: 'danbooru',
+    E621: 'e621',
+}
+
+export const TagCategory = {
+    'danbooru': [
+        'general',
+        'artist',
+        'unused',
+        'copyright',
+        'character',
+        'meta',
+    ],
+    'e621': [
+        'general',
+        'artist',
+        'unused',
+        'copyright',
+        'character',
+        'species',
+        'invalid',
+        'meta',
+        'lore',
+    ]
+}
+
+// --- Data Structures ---
 
 /**
  * Class representing a tag and its metadata
@@ -22,45 +42,54 @@ export class TagData {
      * @param {string[]} [alias=[]] - Array of aliases for the tag
      * @param {string} [category='general'] - Category of the tag
      * @param {number} [count=0] - Frequency count/popularity of the tag
+     * @param {string} [source=TagSources.Danbooru] - The source of the tag data
      */
-    constructor(tag, alias = [], category = 'general', count = 0) {
+    constructor(tag, alias = [], category = 'general', count = 0, source = TagSource.Danbooru) {
         /** @type {string} */
         this.tag = tag;
-        
+
         /** @type {string[]} */
         this.alias = alias;
-        
+
         /** @type {string} */
         this.category = category;
-        
+
         /** @type {number} */
         this.count = count;
+
+        this.source = source;
     }
 }
 
-// Data storage
-export const autoCompleteData = {
-    /** @type {TagData[]} */
-    sortedTags: [],
+class AutocompleteData {
+    constructor() {
+        /** @type {TagData[]} */
+        this.sortedTags = [];
 
-    /** @type {Map<string, TagData>} */
-    tagMap: new Map(), // Stores tag data, mapping tag names to TagData objects
-    
-    /** @type {Map<string, TagData>} */
-    aliasMap: new Map(), // Maps aliases to their main tag names
-    
-    /** @type {Map<string, Map<string, number>>} */
-    cooccurrenceMap: new Map(), // Stores co-occurrence data for related tags
-    
-    isInitializing: false,
-    initialized: false,
+        /** @type {Map<string, TagData>} */
+        this.tagMap = new Map();
 
-     // Progress of "base" csv loading
-    baseLoadingProgress: {
-        // tags: 0, // Commented out because tags csv aren't loaded in chunks
-        cooccurrence: 0
+        /** @type {Map<string, TagData>} */
+        this.aliasMap = new Map();
+
+        /** @type {Map<string, Map<string, number>>} */
+        this.cooccurrenceMap = new Map();
+
+        this.isInitializing = false;
+        this.initialized = false;
+
+        // Progress of "base" csv loading
+        this.baseLoadingProgress = {
+            // tags: 0,
+            cooccurrence: 0
+        };
     }
-};
+}
+
+/**
+ * @type {Object<string, AutocompleteData>}
+ */
+export const autoCompleteData = {};
 
 // CSV Header for tags
 const TAGS_CSV_HEADER = 'tag,category,count,alias';
@@ -70,14 +99,31 @@ const ALIAS_INDEX = TAGS_CSV_HEADER_COLUMNS.indexOf('alias');
 const CATEGORY_INDEX = TAGS_CSV_HEADER_COLUMNS.indexOf('category');
 const COUNT_INDEX = TAGS_CSV_HEADER_COLUMNS.indexOf('count');
 
+// --- Helder Functions ---
+
+/**
+ * Get the available tag sources in priority order based on the current settings.
+ * @returns {string[]} Array of available tag sources in priority order
+ */
+export function getEnabledTagSourceInPriorityOrder() {
+    return Object.values(TagSource)
+        .filter((s) => {
+            return settingValues.tagSource === s || settingValues.tagSource === 'all';
+        })
+        .toSorted((a, b) => {
+            return a === settingValues.primaryTagSource ? -1 : 1;
+        });
+}
+
 // --- Data Loading Functions ---
 
 /**
  * Loads tag data from a single CSV file.
  * @param {string} csvUrl - The URL of the CSV file to load.
+ * @param {string} siteName - The site name (e.g., 'danbooru', 'e621').
  * @returns {Promise<void>}
  */
-async function loadTags(csvUrl) {
+async function loadTags(csvUrl, siteName) {
     try {
         const response = await fetch(csvUrl, { cache: "no-store" });
         if (!response.ok) {
@@ -100,9 +146,9 @@ async function loadTags(csvUrl) {
                 const count = parseInt(columns[COUNT_INDEX].trim(), 10);
 
                 if (!tag || isNaN(count)) continue;
-                
+
                 // Skip if tag already exists (priority to earlier loaded files - extra then base)
-                if (autoCompleteData.tagMap.has(tag)) {
+                if (autoCompleteData[siteName].tagMap.has(tag)) {
                     continue;
                 }
 
@@ -110,26 +156,26 @@ async function loadTags(csvUrl) {
                 const aliases = aliasStr ? aliasStr.split(',').map(a => a.trim()).filter(a => a.length > 0) : [];
 
                 // Create a TagData instance instead of a plain object
-                const tagData = new TagData(tag, aliases, category, count);
+                const tagData = new TagData(tag, aliases, category, count, siteName);
 
-                autoCompleteData.sortedTags.push(tagData);
+                autoCompleteData[siteName].sortedTags.push(tagData);
             } else {
                 console.warn(`[Autocomplete-Plus] Invalid CSV format in line ${i + 1} of ${csvUrl}: ${line}. Expected ${TAGS_CSV_HEADER_COLUMNS.length} columns, but got ${columns.length}.`);
                 continue;
             }
         }
-        
+
         // Sort by count in descending order
-        autoCompleteData.sortedTags.sort((a, b) => b.count - a.count);
+        autoCompleteData[siteName].sortedTags.sort((a, b) => b.count - a.count);
 
         // Build maps as before, but ensure not to overwrite if already processed from extra files
-        autoCompleteData.sortedTags.forEach(tagData => {
-            if (!autoCompleteData.tagMap.has(tagData.tag)) {
-                autoCompleteData.tagMap.set(tagData.tag, tagData);
+        autoCompleteData[siteName].sortedTags.forEach(tagData => {
+            if (!autoCompleteData[siteName].tagMap.has(tagData.tag)) {
+                autoCompleteData[siteName].tagMap.set(tagData.tag, tagData);
                 if (tagData.alias && Array.isArray(tagData.alias)) {
                     tagData.alias.forEach(alias => {
-                        if (!autoCompleteData.aliasMap.has(alias)) {
-                            autoCompleteData.aliasMap.set(alias, tagData.tag); // Map alias back to the main tag
+                        if (!autoCompleteData[siteName].aliasMap.has(alias)) {
+                            autoCompleteData[siteName].aliasMap.set(alias, tagData.tag); // Map alias back to the main tag
                         }
                     });
                 }
@@ -144,9 +190,10 @@ async function loadTags(csvUrl) {
 /**
  * Loads co-occurrence data from a single CSV file.
  * @param {string} csvUrl - The URL of the CSV file to load.
+ * @param {string} siteName - The site name (e.g., 'danbooru', 'e621').
  * @returns {Promise<void>}
  */
-async function loadCooccurrence(csvUrl) {
+async function loadCooccurrence(csvUrl, siteName) {
     try {
         const response = await fetch(csvUrl, { cache: "no-store" });
         if (!response.ok) {
@@ -158,7 +205,7 @@ async function loadCooccurrence(csvUrl) {
 
         const startIndex = lines[0].startsWith('tag_a,tag_b,count') ? 1 : 0;
 
-        await processInChunks(lines, startIndex, autoCompleteData.cooccurrenceMap, csvUrl);
+        await processInChunks(lines, startIndex, autoCompleteData[siteName].cooccurrenceMap, csvUrl, siteName);
     } catch (error) {
         console.error(`[Autocomplete-Plus] Failed to fetch or process cooccurrence data from ${csvUrl}:`, error);
     }
@@ -168,7 +215,7 @@ async function loadCooccurrence(csvUrl) {
  * Process CSV data in chunks to avoid blocking the UI.
  * Modifies the targetMap directly.
  */
-function processInChunks(lines, startIndex, targetMap, sourceFileName = "CSV") {
+function processInChunks(lines, startIndex, targetMap, csvUrl, siteName) {
     return new Promise((resolve) => {
         const CHUNK_SIZE = 10000;
         let i = startIndex;
@@ -194,7 +241,7 @@ function processInChunks(lines, startIndex, targetMap, sourceFileName = "CSV") {
                     }
                     targetMap.get(tagA).set(tagB, count);
 
-                    
+
                     // Add tagB -> tagA relationship (bidirectional)
                     if (!targetMap.has(tagB)) {
                         targetMap.set(tagB, new Map());
@@ -206,7 +253,7 @@ function processInChunks(lines, startIndex, targetMap, sourceFileName = "CSV") {
             }
 
             if (i < lines.length) {
-                autoCompleteData.baseLoadingProgress.cooccurrence = Math.round((i / lines.length) * 100);
+                autoCompleteData[siteName].baseLoadingProgress.cooccurrence = Math.round((i / lines.length) * 100);
                 setTimeout(processChunk, 0);
             } else {
                 resolve();
@@ -250,76 +297,101 @@ function parseCSVLine(line) {
     return result;
 }
 
-/**
- * Initializes the autocomplete data by fetching the list of CSV files and loading them.
- * This function is called when the extension is initialized.
- */
-export async function initializeData() {
-    if (autoCompleteData.isInitializing || autoCompleteData.initialized) {
-        return;
-    }
-    
-    const startTime = performance.now();
-    autoCompleteData.isInitializing = true;
-    // console.log("[Autocomplete-Plus] Initializing autocomplete data...");
-
+export async function fetchCsvList() {
     try {
         const response = await fetch('/autocomplete-plus/csv');
         if (!response.ok) {
             throw new Error(`[Autocomplete-Plus] Failed to fetch CSV list: ${response.status} ${response.statusText}`);
         }
-        const csvListData = await response.json();
+        return await response.json();
+    } catch (error) {
+        console.error("[Autocomplete-Plus] Error fetch csv data:", error);
+    }
 
-        const extraTagsCount = csvListData.danbooru.extra_tags || 0;
-        const extraCooccurrenceCount = csvListData.danbooru.extra_cooccurrence || 0;
+    return null;
+}
 
-        const tagsUrl = '/autocomplete-plus/csv/tags';
-        const cooccurrenceUrl = '/autocomplete-plus/csv/cooccurrence';
+/**
+ * Initializes the autocomplete data by fetching the list of CSV files and loading them.
+ * This function is called when the extension is initialized.
+ */
+export async function initializeData(csvListData, source) {
+    if (autoCompleteData.hasOwnProperty(source) === false) {
+        autoCompleteData[source] = new AutocompleteData();
+    }
 
-        // Load extra tags first
-        let tagsLoadPromises = [];
-        let currentTagPromise = Promise.resolve();
-        for (let i = 0; i < extraTagsCount; i++) {
-            currentTagPromise = currentTagPromise.then(() => loadTags(`${tagsUrl}/extra/${i}`));
-            tagsLoadPromises.push(currentTagPromise);
+    if (autoCompleteData[source].isInitializing || autoCompleteData[source].initialized) {
+        return;
+    }
+
+    const startTime = performance.now();
+    autoCompleteData[source].isInitializing = true;
+    // console.log("[Autocomplete-Plus] Initializing autocomplete data...");
+
+    try {
+        // Store functions that return Promises (Promise Factories)
+        // These factories will be called later to start the actual loading.
+        const tagsLoadPromiseFactories = [];
+        const cooccurrenceLoadPromiseFactories = [];
+
+        // Check if siteName exists in csvListData to prevent errors if a sourte is removed or misconfigured
+        if (!csvListData[source]) {
+            console.warn(`[Autocomplete-Plus] CSV list data not found for sourte: ${source}. Skipping.`);
+            return;
         }
 
-        // Then load base tags if it exists
-        if (csvListData.danbooru.base_tags) {
-            currentTagPromise = currentTagPromise.then(() => loadTags(`${tagsUrl}/base`));
-            tagsLoadPromises.push(currentTagPromise);
-        }
-        
-        // Load extra cooccurrence first
-        let cooccurrenceLoadPromises = [];
-        let cooccurrencePromiseChain = Promise.resolve();
-        for (let i = 0; i < extraCooccurrenceCount; i++) {
-            cooccurrencePromiseChain = cooccurrencePromiseChain.then(() => loadCooccurrence(`${cooccurrenceUrl}/extra/${i}`));
-            cooccurrenceLoadPromises.push(cooccurrencePromiseChain);
-        }
-        
-        // Then load base cooccurrence if it exists
-        if (csvListData.danbooru.base_cooccurrence) {
-            cooccurrencePromiseChain = cooccurrencePromiseChain.then(() => loadCooccurrence(`${cooccurrenceUrl}/base`));
-            cooccurrenceLoadPromises.push(cooccurrencePromiseChain);
-        }
+        const extraTagsCount = csvListData[source].extra_tags || 0;
+        const extraCooccurrenceCount = csvListData[source].extra_cooccurrence || 0;
 
+        const tagsUrl = `/autocomplete-plus/csv/${source}/tags`;
+        const cooccurrenceUrl = `/autocomplete-plus/csv/${source}/tags_cooccurrence`;
+
+        // Factory for loading tags for the current sourte
+        const siteTagsLoaderFactory = async () => {
+            let promiseChain = Promise.resolve();
+            for (let i = 0; i < extraTagsCount; i++) {
+                promiseChain = promiseChain.then(() => loadTags(`${tagsUrl}/extra/${i}`, source));
+            }
+            if (csvListData[source].base_tags) {
+                promiseChain = promiseChain.then(() => loadTags(`${tagsUrl}/base`, source));
+            }
+            return promiseChain;
+        };
+        tagsLoadPromiseFactories.push(siteTagsLoaderFactory);
+
+        // Factory for loading cooccurrence data for the current sourte
+        const siteCooccurrenceLoaderFactory = async () => {
+            let promiseChain = Promise.resolve();
+            for (let i = 0; i < extraCooccurrenceCount; i++) {
+                promiseChain = promiseChain.then(() => loadCooccurrence(`${cooccurrenceUrl}/extra/${i}`, source));
+            }
+            if (csvListData[source].base_cooccurrence) {
+                promiseChain = promiseChain.then(() => loadCooccurrence(`${cooccurrenceUrl}/base`, source));
+            }
+            return promiseChain;
+        };
+        cooccurrenceLoadPromiseFactories.push(siteCooccurrenceLoaderFactory);
+
+        // Now, execute all promise factories and wait for their completion.
+        // The actual loading (fetch calls) will start when the factories are invoked here.
         await Promise.all([
-            Promise.all(tagsLoadPromises).then(() => {
+            Promise.all(tagsLoadPromiseFactories.map(factory => factory())).then(() => {
                 const endTime = performance.now();
-                console.log(`[Autocomplete-Plus] Tags loading complete in ${(endTime - startTime).toFixed(2)}ms. Extra file count: ${extraTagsCount}`);
-
+                console.log(`[Autocomplete-Plus] "${source}" Tags loading complete in ${(endTime - startTime).toFixed(2)}ms`);
             }),
-            Promise.all(cooccurrenceLoadPromises).then(() => {
+            Promise.all(cooccurrenceLoadPromiseFactories.map(factory => factory())).then(() => {
                 const endTime = performance.now();
-                console.log(`[Autocomplete-Plus] Co-occurrence loading complete in ${(endTime - startTime).toFixed(2)}ms. Extra file count: ${extraCooccurrenceCount}`);
+                // The original log for extraCooccurrenceCount was potentially misleading as it showed
+                // the count for the last processed sourte. Removed for clarity.
+                // If a total count is needed, it should be calculated across all sites.
+                console.log(`[Autocomplete-Plus] "${source}" Co-occurrence loading complete in ${(endTime - startTime).toFixed(2)}ms.`);
             })
         ]);
 
-        autoCompleteData.initialized = true;
+        autoCompleteData[source].initialized = true;
     } catch (error) {
         console.error("[Autocomplete-Plus] Error initializing autocomplete data:", error);
     } finally {
-        autoCompleteData.isInitializing = false;
+        autoCompleteData[source].isInitializing = false;
     }
 }

web/js/main.js

@@ -2,7 +2,7 @@ import { app } from "/scripts/app.js";
 import { ComfyWidgets } from "/scripts/widgets.js";
 import { settingValues } from "./settings.js";
 import { loadCSS } from "./utils.js";
-import { initializeData } from "./data.js";
+import { TagSource, fetchCsvList, initializeData } from "./data.js";
 import { AutocompleteEventHandler } from "./autocomplete.js";
 import { RelatedTagsEventHandler } from "./related-tags.js";
 
@@ -126,17 +126,57 @@ const name = "Autocomplete Plus";
 app.registerExtension({
     id: id,
     name: name,
-    setup() {
+    async setup() {
         initializeEventHandlers();
 
         let rootPath = import.meta.url.replace("js/main.js", "");
         loadCSS(rootPath + "css/autocomplete-plus.css"); // Load CSS for autocomplete
 
-        initializeData();
+        fetchCsvList().then((csvList) => {
+            Object.values(TagSource).forEach((source) => {
+                initializeData(csvList, source);
+            });
+        });
     },
-    
-    //One the Settings Screen, displays reverse order in same category
-    settings: [ 
+
+    // One the Settings Screen, displays reverse order in same category
+    settings: [
+        // --- Tag source Settings ---
+        {
+            id: id + ".tag_source_icon_position",
+            name: "Tag Source Icon Position",
+            type: "combo",
+            options: ["left", "right", "hidden"],
+            defaultValue: "left",
+            category: [name, "Tag Source", "Tag Source Icon Position"],
+            onChange: (newVal, oldVal) => {
+                settingValues.tagSourceIconPosition = newVal;
+            }
+        },
+        {
+            id: id + ".primary_tag_source",
+            name: "Primary source for 'all' Source",
+            tooltip: "When 'Autocomplete Tag Source' is 'all', this determines which source's tags appear first in suggestions.",
+            type: "combo",
+            options: Object.values(TagSource),
+            defaultValue: TagSource.Danbooru,
+            category: [name, "Tag Source", "Prioritize Tag Source"],
+            onChange: (newVal, oldVal) => {
+                settingValues.primaryTagSource = newVal;
+            }
+        },
+        {
+            id: id + ".tag_source",
+            name: "Autocomplete Tag Source",
+            tooltip: "Select the tag source for autocomplete suggestions. 'all' includes tags from all loaded sources.",
+            type: "combo",
+            options: [...Object.values(TagSource), "all"],
+            defaultValue: "all",
+            category: [name, "Tag Source", "Tag Source"],
+            onChange: (newVal, oldVal) => {
+                settingValues.tagSource = newVal;
+            }
+        },
         // --- Autocomplete Settings ---
         {
             id: id + ".max_suggestions",
@@ -164,7 +204,7 @@ app.registerExtension({
                 settingValues.enabled = newVal;
             }
         },
-        
+
         // --- Related Tags Settings ---
         {
             id: id + ".related_tags_trigger_mode",

web/js/related-tags.js

@@ -1,4 +1,4 @@
-import { TagCategory, TagData, autoCompleteData } from './data.js';
+import { TagSource, TagCategory, TagData, autoCompleteData } from './data.js';
 import { settingValues } from './settings.js';
 import {
     extractTagsFromTextArea,
@@ -7,7 +7,7 @@ import {
     isValidTag,
     normalizeTagToInsert,
     normalizeTagToSearch,
-    getCurrentTagRange
+    getCurrentTagRange,
 } from './utils.js';
 
 // --- RelatedTags Logic ---
@@ -15,19 +15,20 @@ import {
 /**
  * Calculates the Jaccard similarity between two tags.
  * Jaccard similarity = (A ∩ B) / (A ∪ B) = (A ∩ B) / (|A| + |B| - |A ∩ B|)
+ * @param {string} tagSource The name of the site (e.g., 'danbooru', 'e621')
  * @param {string} tagA The first tag
  * @param {string} tagB The second tag
  * @returns {number} Similarity score between 0 and 1
  */
-function calculateJaccardSimilarity(tagA, tagB) {
+function calculateJaccardSimilarity(tagSource, tagA, tagB) {
     // Get the count of tagA and tagB individually
-    const countA = autoCompleteData.tagMap.get(tagA)?.count || 0;
-    const countB = autoCompleteData.tagMap.get(tagB)?.count || 0;
+    const countA = autoCompleteData[tagSource].tagMap.get(tagA)?.count || 0;
+    const countB = autoCompleteData[tagSource].tagMap.get(tagB)?.count || 0;
 
     if (countA === 0 || countB === 0) return 0;
 
     // Get the cooccurrence count
-    const cooccurrenceAB = autoCompleteData.cooccurrenceMap.get(tagA)?.get(tagB) || 0;
+    const cooccurrenceAB = autoCompleteData[tagSource].cooccurrenceMap.get(tagA)?.get(tagB) || 0;
 
     // Calculate Jaccard similarity
     // (A ∩ B) / (A ∪ B) = (A ∩ B) / (|A| + |B| - |A ∩ B|)
@@ -64,11 +65,13 @@ export function getTagFromCursorPosition(inputElement) {
 function searchRelatedTags(tag) {
     const startTime = performance.now(); // Record start time for performance measurement
 
-    if (!tag || !autoCompleteData.cooccurrenceMap.has(tag)) {
+    const tagSource = TagSource.Danbooru; // TODO: Leave the tag source as Danbooru until e621_tags_cooccurrence.csv is ready
+
+    if (!tag || !autoCompleteData[tagSource].cooccurrenceMap.has(tag)) {
         return [];
     }
 
-    const cooccurrences = autoCompleteData.cooccurrenceMap.get(tag);
+    const cooccurrences = autoCompleteData[tagSource].cooccurrenceMap.get(tag);
     const relatedTags = [];
 
     // Convert to array for sorting
@@ -77,17 +80,18 @@ function searchRelatedTags(tag) {
         if (coTag === tag) return;
 
         // Get tag data
-        const tagData = autoCompleteData.tagMap.get(coTag);
+        const tagData = autoCompleteData[tagSource].tagMap.get(coTag);
         if (!tagData) return;
 
         // Calculate similarity
-        const similarity = calculateJaccardSimilarity(tag, coTag);
+        const similarity = calculateJaccardSimilarity(tagSource, tag, coTag);
 
         relatedTags.push({
             tag: coTag,
             similarity: similarity,
             alias: tagData.alias,
             category: tagData.category,
+            source: tagData.source,
             count: tagData.count,
         });
     });
@@ -320,7 +324,7 @@ class RelatedTagsUI {
         this.root.style.display = 'block';
 
         // Update initialization status if not already done
-        if (!autoCompleteData.initialized) {
+        if (!autoCompleteData[TagSource.Danbooru].initialized) {
             if (this.autoRefreshTimerId) {
                 clearTimeout(this.autoRefreshTimerId);
             }
@@ -384,13 +388,29 @@ class RelatedTagsUI {
      * Updates header content
      */
     #updateHeader() {
+        // Find the tag data for the current tag
+        const tagData = Object.values(TagSource)
+            .map((source) => {
+                if (source in autoCompleteData && autoCompleteData[source].tagMap.has(this.currentTag)) {
+                    return autoCompleteData[source].tagMap.get(this.currentTag);
+                }
+            })
+            .find((tagData) => tagData !== undefined);
+
         // Update header text with current tag
         this.headerText.innerHTML = ''; // Clear previous content
         this.headerText.textContent = 'Tags related to: ';
-        const tagNameSpan = document.createElement('span');
-        tagNameSpan.className = 'related-tags-header-tag-name';
-        tagNameSpan.textContent = this.currentTag;
-        this.headerText.appendChild(tagNameSpan);
+        const tagName = document.createElement('span');
+        tagName.className = 'related-tags-header-tag-name';
+        tagName.textContent = this.currentTag;
+        if (tagData && ['left', 'right'].includes(settingValues.tagSourceIconPosition)) {
+            const tagSourceIconHtml = `<svg class="autocomplete-plus-tag-icon-svg"><use xlink:href="#autocomplete-plus-icon-${tagData.source}"></use></svg>`;
+            tagName.innerHTML = settingValues.tagSourceIconPosition == 'left'
+                ? `${tagSourceIconHtml} ${tagData.tag}`
+                : `${tagData.tag} ${tagSourceIconHtml}`;
+        }
+
+        this.headerText.appendChild(tagName);
 
         // Update pin button
         this.pinBtn.textContent = this.isPinned ? '🎯' : '📌';
@@ -407,20 +427,20 @@ class RelatedTagsUI {
     #updateContent() {
         this.tagsContainer.innerHTML = '';
 
-        if (!autoCompleteData.initialized) {
+        if (!autoCompleteData[TagSource.Danbooru].initialized) {
             // Show loading message
             const messageDiv = document.createElement('div');
             messageDiv.className = 'related-tags-loading-message';
-            messageDiv.textContent = `Initializing cooccurrence data... [${autoCompleteData.baseLoadingProgress.cooccurrence}%]`;
+            messageDiv.textContent = `Initializing cooccurrence data... [${autoCompleteData[TagSource.Danbooru].baseLoadingProgress.cooccurrence}%]`;
             this.tagsContainer.appendChild(messageDiv);
             return;
         }
 
         if (!this.relatedTags || this.relatedTags.length === 0) {
             // Show no related tags message
-            const messageCell = document.createElement('div');
-            messageCell.textContent = 'No related tags found';
-            this.tagsContainer.appendChild(messageCell);
+            const messageDiv = document.createElement('div');
+            messageDiv.textContent = 'No related tags found';
+            this.tagsContainer.appendChild(messageDiv);
             return;
         }
 
@@ -441,10 +461,10 @@ class RelatedTagsUI {
      * @returns {HTMLTableRowElement} The tag row element
      */
     #createTagElement(tagData, isExisting) {
-        const categoryText = TagCategory[tagData.category] || "unknown";
+        const categoryText = TagCategory[tagData.source][tagData.category] || "unknown";
 
         const tagRow = document.createElement('div');
-        tagRow.className = 'related-tag-item';
+        tagRow.classList.add('related-tag-item', tagData.source);
         tagRow.dataset.tag = tagData.tag;
         tagRow.dataset.tagCategory = categoryText;
 

web/js/settings.js

@@ -1,4 +1,9 @@
 export const settingValues = {
+    // Tag source settings
+    tagSource: 'all', // 'danbooru', 'e621', 'all'
+    primaryTagSource: 'danbooru', // 'danbooru', 'e621'
+    tagSourceIconPosition: 'left', // 'left', 'right', 'hidden'
+
     // Autocomplete feature settings
     enabled: true,
     maxSuggestions: 10,

web/js/utils.js

@@ -1,3 +1,29 @@
+// --- Html String constants ---
+
+/**
+ * HTML string for the tag source icon.
+ */
+export const IconSvgHtmlString = `
+    <svg>
+        <defs>
+            <symbol id="autocomplete-plus-icon-danbooru" viewBox="14 12 42 42">
+                <polygon points="20,20 20,44 44,44 44,20" fill="#a67c52" stroke="#2d1400" stroke-width="2" stroke-linejoin="bevel"/>
+                <polygon points="20,20 26,14 50,14 44,20" fill="#6e4b2c" stroke="#2d1400" stroke-width="2" stroke-linejoin="bevel"/>
+                <polygon points="44,20 50,14 50,38 44,44" fill="#8b5e3c" stroke="#2d1400" stroke-width="2" stroke-linejoin="bevel"/>
+            </symbol>
+            <symbol id="autocomplete-plus-icon-e621" viewBox="-33 -38 66 76">
+                <style>
+                .e621-blue { fill: #0d2e69; }
+                .e621-white { fill:rgb(230, 230, 230); }
+                </style>
+                <g>
+                <path class="e621-blue" d="M0 -38 L32.909 -19 L32.909 19 L0 38 L-32.909 19 L-32.909 -19 Z"/>
+                <text x="0" y="0" class="e621-white" font-family="Arial Black, Gadget, sans-serif" font-size="56" text-anchor="middle" dominant-baseline="middle">e</text>
+                </g>
+            </symbol>
+        </defs>
+    </svg>`;
+
 // --- String Helper Functions ---
 
 const MAX_PROMPT_WEIGHT_VALUE = 9.9;
@@ -491,7 +517,10 @@ export function getCurrentTagRange(text, cursorPos) {
 
 // --- End String Helper Functions ---
 
-// Function to load a CSS file
+/**
+ * Load a CSS file dynamically.
+ * @param {string} href 
+ */
 export function loadCSS(href) {
     const link = document.createElement('link');
     link.rel = 'stylesheet';
@@ -500,7 +529,7 @@ export function loadCSS(href) {
     // Ensure the CSS is loaded before other scripts might rely on its styles
     // by adding it to the head.
     document.head.appendChild(link);
-    console.debug(`Loaded CSS: ${href}`); // Optional: Log loading
+    // console.debug(`[Autocomplete-Plus] Loaded CSS: ${href}`); // Optional: Log loading
 }
 
 /**