Add semi-auto tag complete

app.py

@@ -5,7 +5,26 @@ from lib import init, refresh_character_thumb_image
 from lib import JAVA_SCRIPT, CSS_SCRIPT, TITLE
 
 if __name__ == '__main__':
-    character_list, character_list_cn, LANG = init()
+    character_list, character_list_cn, LANG, TAG_AUTOCOMPLETE = init()
+    
+    def update_suggestions(text):
+        matches = TAG_AUTOCOMPLETE.get_suggestions(text)
+        items = []
+        if matches:       
+            for _, m in enumerate(matches):
+                key = f"{m['prompt']} ({m['heat']})"
+                items.append([key])
+        return gr.Dataset(samples=items)
+        
+    def apply_suggestion(evt: gr.SelectData, text, custom_prompt):        
+        suggestion = evt.value
+        #print(f"You selected {evt.value} at {evt.index} from {evt.target} for {custom_prompt}")
+        if not custom_prompt or not suggestion:
+            return custom_prompt
+            
+        parts = custom_prompt.split(',')
+        parts[-1] = suggestion[0].split(' ')[0].replace('_', ' ').replace(':', ' ')        
+        return ', '.join(parts) + ', '
     
     with gr.Blocks(js=JAVA_SCRIPT, css=CSS_SCRIPT, title=TITLE) as ui:
         with gr.Row():
@@ -33,10 +52,14 @@ if __name__ == '__main__':
         with gr.Row(elem_classes='main_row'):            
             with gr.Column(elem_classes='column_prompts'):
                     thumb_image = gr.Gallery(type="pil", columns=3, show_download_button=False, object_fit='contain', label="Thumb")    
+        
+        with gr.Row():
+            custom_prompt = gr.Textbox(value='', label='Semi-auto tag complete test. Try tag* *tag *tag* (e621_sfw.csv@DominikDoom)') 
+            suggestions = gr.Dataset(components=[custom_prompt], samples_per_page=30, samples=[])
                         
         with gr.Row():
             gr.Markdown("<h2><span style=\"color:orangered\">Thumb Image create by waiNSFWIllustrious_v120.safetensors with ComfyUI</span></h2>",)
-            gr.Markdown(f"<a href='https://github.com/mirabarukaso/character_select_stand_alone_app'>Character Select SAA</a>")
+            gr.Markdown(f"<a href='https://github.com/mirabarukaso/character_select_stand_alone_app'>Character Select SAA</a>")            
             
         character1.change(fn=refresh_character_thumb_image,
                           inputs=[character1,character2,character3],
@@ -48,4 +71,7 @@ if __name__ == '__main__':
                           inputs=[character1,character2,character3],
                           outputs=[thumb_image])
         
+        custom_prompt.change(fn=update_suggestions, inputs=[custom_prompt], outputs=[suggestions])
+        suggestions.select(fn=apply_suggestion, inputs=[suggestions, custom_prompt], outputs=[custom_prompt])
+        
     ui.launch()

scripts/lib.py

@@ -1,20 +1,13 @@
-import datetime
 import gzip
 import hashlib
 import os
-import glob
-import textwrap
-import gradio as gr
-import numpy as np
 import requests
 import json
 import base64
-import webbrowser
+
 from io import BytesIO
 from PIL import Image
-from PIL import PngImagePlugin
-import random
-import argparse
+from tag_autocomplete import PromptSuggester
 
 # Language
 LANG_EN = {
@@ -60,10 +53,12 @@ character_list = ''
 character_dict = {}
 wai_image_dict = {}
 character_list_cn = ''
+TAG_AUTOCOMPLETE = None
 
 wai_illustrious_character_select_files = [
     {'name': 'wai_character', 'file_path': os.path.join(json_folder, 'wai_characters.csv'), 'url':'https://raw.githubusercontent.com/mirabarukaso/character_select_stand_alone_app/refs/heads/main/json/wai_characters.csv'},
     {'name': 'wai_image', 'file_path': os.path.join(json_folder, 'wai_character_thumbs.json'), 'url': 'https://huggingface.co/datasets/flagrantia/character_select_stand_alone_app/resolve/main/wai_character_thumbs.json'},
+    {'name': 'e621_sfw', 'file_path': os.path.join(json_folder, 'e621_sfw.csv'), 'url': 'https://raw.githubusercontent.com/DominikDoom/a1111-sd-webui-tagcomplete/refs/heads/main/tags/e621_sfw.csv'},
 ]
 def get_md5_hash(input_str):
     md5_hash = hashlib.md5()
@@ -88,6 +83,7 @@ def load_jsons():
     global character_dict
     global wai_image_dict
     global character_list_cn
+    global TAG_AUTOCOMPLETE
     
     # download file
     for item in wai_illustrious_character_select_files:
@@ -98,14 +94,17 @@ def load_jsons():
         if not os.path.exists(file_path):
             download_file(url, file_path)
             
-        with open(file_path, 'r', encoding='utf-8') as file:        
-            if 'wai_character' == name:
-                lines = file.readlines()
-                for line in lines:
-                    key, value = line.split(',')
-                    character_dict[key.strip()]=value.strip()
-            elif 'wai_image' == name:
-                wai_image_dict = json.load(file)                
+        if 'e621_sfw' == name:
+            TAG_AUTOCOMPLETE = PromptSuggester(file_path)
+        else:
+            with open(file_path, 'r', encoding='utf-8') as file:        
+                if 'wai_character' == name:
+                    lines = file.readlines()
+                    for line in lines:
+                        key, value = line.split(',')
+                        character_dict[key.strip()]=value.strip()
+                elif 'wai_image' == name:
+                    wai_image_dict = json.load(file)                
                                         
     # Create list        
     character_list = list(character_dict.values())   
@@ -167,4 +166,4 @@ def init():
     load_jsons()    
     print(f'[{CAT}]:Starting...')
     
-    return character_list, character_list_cn, LANG
+    return character_list, character_list_cn, LANG, TAG_AUTOCOMPLETE

scripts/tag_autocomplete.py

@@ -0,0 +1,94 @@
+import os
+from typing import List, Dict
+
+CAT = "Auto Tag Complete"
+
+class PromptSuggester:
+    def __init__(self, prompt_file_path):
+        self.prompts = []
+        self.load_prompts(prompt_file_path)
+        
+    def load_prompts(self, file_path: str):
+        if not os.path.exists(file_path):
+            print(f"File {file_path} not found.")
+            return
+            
+        with open(file_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                    
+                parts = line.split(',', 3)                 
+                if len(parts) >= 2:
+                    prompt = parts[0].strip()
+                    group = int(parts[1]) if parts[1].strip().isdigit() else 0
+                    heat = int(parts[2]) if len(parts) > 2 and parts[2].strip().isdigit() else 0
+                    aliases = parts[3].strip('"') if len(parts) > 3 else ""
+                    
+                    if heat == 0:
+                        heat = group
+                        group = 0
+                    
+                    self.prompts.append({
+                        'prompt': prompt, 
+                        'group': group, 
+                        'heat': heat, 
+                        'aliases': aliases
+                    })
+        
+        self.prompts.sort(key=lambda x: x['heat'], reverse=True)
+        print(f"[{CAT}] Loaded {len(self.prompts)} prompts.")
+    
+    def get_suggestions(self, text: str) -> List[Dict]:        
+        if not text:
+            return []
+            
+        parts = text.split(',')
+        last_word = parts[-1].strip().lower()
+        
+        if not last_word:
+            return []
+        
+        #print(f"Getting suggestions for: {last_word}")
+            
+        matches = {}
+        for prompt_info in self.prompts:
+            prompt = prompt_info['prompt'].lower()
+            aliases = prompt_info['aliases'].lower().split(',') if prompt_info['aliases'] else []
+            
+            if '*' in last_word:
+                # wildcards
+                if last_word.startswith('*') and last_word.endswith('*'):
+                    # middle
+                    pattern = last_word[1:-1]
+                    if pattern in prompt or any(pattern in alias.strip() for alias in aliases):
+                        if prompt not in matches or prompt_info['heat'] > matches[prompt]['heat']:
+                            matches[prompt] = {'prompt': prompt_info['prompt'], 'heat': prompt_info['heat']}
+                elif last_word.startswith('*'):
+                    # endswith
+                    pattern = last_word[1:]
+                    if prompt.endswith(pattern) or any(alias.strip().endswith(pattern) for alias in aliases):
+                        if prompt not in matches or prompt_info['heat'] > matches[prompt]['heat']:
+                            matches[prompt] = {'prompt': prompt_info['prompt'], 'heat': prompt_info['heat']}
+                elif last_word.endswith('*'):
+                    # startswith
+                    pattern = last_word[:-1]
+                    if prompt.startswith(pattern) or any(alias.strip().startswith(pattern) for alias in aliases):
+                        if prompt not in matches or prompt_info['heat'] > matches[prompt]['heat']:
+                            matches[prompt] = {'prompt': prompt_info['prompt'], 'heat': prompt_info['heat']}
+            else:
+                # normal
+                if prompt.startswith(last_word) or any(alias.strip().startswith(last_word) for alias in aliases):
+                    if prompt not in matches or prompt_info['heat'] > matches[prompt]['heat']:
+                        matches[prompt] = {'prompt': prompt_info['prompt'], 'heat': prompt_info['heat']}
+
+            # we only need 30 items
+            if len(matches) == 30:
+                break
+        
+        sorted_matches = sorted(matches.values(), key=lambda x: x['heat'], reverse=True)
+        #print(f"Found {len(sorted_matches)} unique matches")
+        return sorted_matches
+    
+