tag_autocomplete_helper.py

# This helper script scans folders for wildcards and embeddings and writes them
# to a temporary file to expose it to the javascript side

import gradio as gr
from pathlib import Path
from modules import scripts, script_callbacks, shared, sd_hijack
#from modules.paths import script_path, extensions_dir 
import yaml

# Webui root path
FILE_DIR = Path().absolute()
#FILE_DIR = Path(script_path)


# The extension base path
EXT_PATH = FILE_DIR.joinpath('extensions')
#EXT_PATH = Path(extensions_dir)


# Tags base path
TAGS_PATH = Path(scripts.basedir()).joinpath('tags')

# The path to the folder containing the wildcards and embeddings
WILDCARD_PATH = FILE_DIR.joinpath('scripts/wildcards')
EMB_PATH = Path(shared.cmd_opts.embeddings_dir)
HYP_PATH = Path(shared.cmd_opts.hypernetwork_dir)

try:
    LORA_PATH = Path(shared.cmd_opts.lora_dir)
except AttributeError:
    LORA_PATH = None

def find_ext_wildcard_paths():
    """Returns the path to the extension wildcards folder"""
    found = list(EXT_PATH.glob('*/wildcards/'))
    return found


# The path to the extension wildcards folder
WILDCARD_EXT_PATHS = find_ext_wildcard_paths()

# The path to the temporary files
STATIC_TEMP_PATH = FILE_DIR.joinpath('tmp') # In the webui root, on windows it exists by default, on linux it doesn't
TEMP_PATH = TAGS_PATH.joinpath('temp') # Extension specific temp files


def get_wildcards():
    """Returns a list of all wildcards. Works on nested folders."""
    wildcard_files = list(WILDCARD_PATH.rglob("*.txt"))
    resolved = [w.relative_to(WILDCARD_PATH).as_posix(
    ) for w in wildcard_files if w.name != "put wildcards here.txt"]
    return resolved


def get_ext_wildcards():
    """Returns a list of all extension wildcards. Works on nested folders."""
    wildcard_files = []

    for path in WILDCARD_EXT_PATHS:
        wildcard_files.append(path.relative_to(FILE_DIR).as_posix())
        wildcard_files.extend(p.relative_to(path).as_posix() for p in path.rglob("*.txt") if p.name != "put wildcards here.txt")
        wildcard_files.append("-----")

    return wildcard_files


def get_ext_wildcard_tags():
    """Returns a list of all tags found in extension YAML files found under a Tags: key."""
    wildcard_tags = {} # { tag: count }
    yaml_files = []
    for path in WILDCARD_EXT_PATHS:
        yaml_files.extend(p for p in path.rglob("*.yml"))
        yaml_files.extend(p for p in path.rglob("*.yaml"))
    count = 0
    for path in yaml_files:
        try:
            with open(path, encoding="utf8") as file:
                data = yaml.safe_load(file)
                for item in data:
                    if data[item] and 'Tags' in data[item]:
                        wildcard_tags[count] = ','.join(data[item]['Tags'])
                        count += 1
                    else:
                        print('Issue with tags found in ' + path.name + ' at item ' + item)
        except yaml.YAMLError as exc:
            print(exc)
    # Sort by count
    sorted_tags = sorted(wildcard_tags.items(), key=lambda item: item[1], reverse=True)
    output = []
    for tag, count in sorted_tags:
        output.append(f"{tag},{count}")
    return output


def get_embeddings(sd_model):
    """Write a list of all embeddings with their version"""

    # Version constants
    V1_SHAPE = 768
    V2_SHAPE = 1024
    emb_v1 = []
    emb_v2 = []
    results = []

    try:
        # Get embedding dict from sd_hijack to separate v1/v2 embeddings
        emb_type_a = sd_hijack.model_hijack.embedding_db.word_embeddings
        emb_type_b = sd_hijack.model_hijack.embedding_db.skipped_embeddings
        # Get the shape of the first item in the dict
        emb_a_shape = -1
        emb_b_shape = -1
        if (len(emb_type_a) > 0):
            emb_a_shape = next(iter(emb_type_a.items()))[1].shape
        if (len(emb_type_b) > 0):
            emb_b_shape = next(iter(emb_type_b.items()))[1].shape

        # Add embeddings to the correct list
        if (emb_a_shape == V1_SHAPE):
            emb_v1 = list(emb_type_a.keys())
        elif (emb_a_shape == V2_SHAPE):
            emb_v2 = list(emb_type_a.keys())

        if (emb_b_shape == V1_SHAPE):
            emb_v1 = list(emb_type_b.keys())
        elif (emb_b_shape == V2_SHAPE):
            emb_v2 = list(emb_type_b.keys())

        # Get shape of current model
        #vec = sd_model.cond_stage_model.encode_embedding_init_text(",", 1)
        #model_shape = vec.shape[1]
        # Show relevant entries at the top
        #if (model_shape == V1_SHAPE):
        #    results = [e + ",v1" for e in emb_v1] + [e + ",v2" for e in emb_v2]
        #elif (model_shape == V2_SHAPE):
        #    results = [e + ",v2" for e in emb_v2] + [e + ",v1" for e in emb_v1]
        #else:
        #    raise AttributeError # Fallback to old method
        results = sorted([e + ",v1" for e in emb_v1] + [e + ",v2" for e in emb_v2], key=lambda x: x.lower())
    except AttributeError:
        print("tag_autocomplete_helper: Old webui version or unrecognized model shape, using fallback for embedding completion.")
        # Get a list of all embeddings in the folder
        all_embeds = [str(e.relative_to(EMB_PATH)) for e in EMB_PATH.rglob("*") if e.suffix in {".bin", ".pt", ".png",'.webp', '.jxl', '.avif'}]
        # Remove files with a size of 0
        all_embeds = [e for e in all_embeds if EMB_PATH.joinpath(e).stat().st_size > 0]
        # Remove file extensions
        all_embeds = [e[:e.rfind('.')] for e in all_embeds]
        results = [e + "," for e in all_embeds]

    write_to_temp_file('emb.txt', results)

def get_hypernetworks():
    """Write a list of all hypernetworks"""

    # Get a list of all hypernetworks in the folder
    all_hypernetworks = [str(h.name) for h in HYP_PATH.rglob("*") if h.suffix in {".pt"}]
    # Remove file extensions
    return sorted([h[:h.rfind('.')] for h in all_hypernetworks], key=lambda x: x.lower())

def get_lora():
    """Write a list of all lora"""

    # Get a list of all lora in the folder
    all_lora = [str(l.name) for l in LORA_PATH.rglob("*") if l.suffix in {".safetensors", ".ckpt", ".pt"}]
    # Remove file extensions
    return sorted([l[:l.rfind('.')] for l in all_lora], key=lambda x: x.lower())


def write_tag_base_path():
    """Writes the tag base path to a fixed location temporary file"""
    with open(STATIC_TEMP_PATH.joinpath('tagAutocompletePath.txt'), 'w', encoding="utf-8") as f:
        f.write(TAGS_PATH.relative_to(FILE_DIR).as_posix())


def write_to_temp_file(name, data):
    """Writes the given data to a temporary file"""
    with open(TEMP_PATH.joinpath(name), 'w', encoding="utf-8") as f:
        f.write(('\n'.join(data)))


csv_files = []
csv_files_withnone = []
def update_tag_files():
    """Returns a list of all potential tag files"""
    global csv_files, csv_files_withnone
    files = [str(t.relative_to(TAGS_PATH)) for t in TAGS_PATH.glob("*.csv")]
    csv_files = files
    csv_files_withnone = ["None"] + files



# Write the tag base path to a fixed location temporary file
# to enable the javascript side to find our files regardless of extension folder name
if not STATIC_TEMP_PATH.exists():
    STATIC_TEMP_PATH.mkdir(exist_ok=True)

write_tag_base_path()
update_tag_files()

# Check if the temp path exists and create it if not
if not TEMP_PATH.exists():
    TEMP_PATH.mkdir(parents=True, exist_ok=True)

# Set up files to ensure the script doesn't fail to load them
# even if no wildcards or embeddings are found
write_to_temp_file('wc.txt', [])
write_to_temp_file('wce.txt', [])
write_to_temp_file('wcet.txt', [])
write_to_temp_file('hyp.txt', [])
write_to_temp_file('lora.txt', [])
# Only reload embeddings if the file doesn't exist, since they are already re-written on model load
if not TEMP_PATH.joinpath("emb.txt").exists():
    write_to_temp_file('emb.txt', [])

# Write wildcards to wc.txt if found
if WILDCARD_PATH.exists():
    wildcards = [WILDCARD_PATH.relative_to(FILE_DIR).as_posix()] + get_wildcards()
    if wildcards:
        write_to_temp_file('wc.txt', wildcards)

# Write extension wildcards to wce.txt if found
if WILDCARD_EXT_PATHS is not None:
    wildcards_ext = get_ext_wildcards()
    if wildcards_ext:
        write_to_temp_file('wce.txt', wildcards_ext)
    # Write yaml extension wildcards to wcet.txt if found
    wildcards_yaml_ext = get_ext_wildcard_tags()
    if wildcards_yaml_ext:
        write_to_temp_file('wcet.txt', wildcards_yaml_ext)

# Write embeddings to emb.txt if found
if EMB_PATH.exists():
    # Get embeddings after the model loaded callback
    script_callbacks.on_model_loaded(get_embeddings)

if HYP_PATH.exists():
    hypernets = get_hypernetworks()
    if hypernets:
        write_to_temp_file('hyp.txt', hypernets)

if LORA_PATH is not None and LORA_PATH.exists():
    lora = get_lora()
    if lora:
        write_to_temp_file('lora.txt', lora)

# Register autocomplete options
def on_ui_settings():
    TAC_SECTION = ("tac", "Tag Autocomplete")
    # Main tag file
    shared.opts.add_option("tac_tagFile", shared.OptionInfo("danbooru.csv", "Tag filename", gr.Dropdown, lambda: {"choices": csv_files_withnone}, refresh=update_tag_files, section=TAC_SECTION))
    # Active in settings
    shared.opts.add_option("tac_active", shared.OptionInfo(True, "Enable Tag Autocompletion", section=TAC_SECTION))
    shared.opts.add_option("tac_activeIn.txt2img", shared.OptionInfo(True, "Active in txt2img (Requires restart)", section=TAC_SECTION))
    shared.opts.add_option("tac_activeIn.img2img", shared.OptionInfo(True, "Active in img2img (Requires restart)", section=TAC_SECTION))
    shared.opts.add_option("tac_activeIn.negativePrompts", shared.OptionInfo(True, "Active in negative prompts (Requires restart)", section=TAC_SECTION))
    shared.opts.add_option("tac_activeIn.thirdParty", shared.OptionInfo(True, "Active in third party textboxes [Dataset Tag Editor] (Requires restart)", section=TAC_SECTION))
    shared.opts.add_option("tac_activeIn.modelList", shared.OptionInfo("", "List of model names (with file extension) or their hashes to use as black/whitelist, separated by commas.", section=TAC_SECTION))
    shared.opts.add_option("tac_activeIn.modelListMode", shared.OptionInfo("Blacklist", "Mode to use for model list", gr.Dropdown, lambda: {"choices": ["Blacklist","Whitelist"]}, section=TAC_SECTION))
    # Results related settings
    shared.opts.add_option("tac_slidingPopup", shared.OptionInfo(True, "Move completion popup together with text cursor", section=TAC_SECTION))
    shared.opts.add_option("tac_maxResults", shared.OptionInfo(5, "Maximum results", section=TAC_SECTION))
    shared.opts.add_option("tac_showAllResults", shared.OptionInfo(False, "Show all results", section=TAC_SECTION))
    shared.opts.add_option("tac_resultStepLength", shared.OptionInfo(100, "How many results to load at once", section=TAC_SECTION))
    shared.opts.add_option("tac_delayTime", shared.OptionInfo(100, "Time in ms to wait before triggering completion again (Requires restart)", section=TAC_SECTION))
    shared.opts.add_option("tac_useWildcards", shared.OptionInfo(True, "Search for wildcards", section=TAC_SECTION))
    shared.opts.add_option("tac_useEmbeddings", shared.OptionInfo(True, "Search for embeddings", section=TAC_SECTION))
    shared.opts.add_option("tac_useHypernetworks", shared.OptionInfo(True, "Search for hypernetworks", section=TAC_SECTION))
    shared.opts.add_option("tac_useLoras", shared.OptionInfo(True, "Search for Loras", section=TAC_SECTION))
    shared.opts.add_option("tac_showWikiLinks", shared.OptionInfo(False, "Show '?' next to tags, linking to its Danbooru or e621 wiki page (Warning: This is an external site and very likely contains NSFW examples!)", section=TAC_SECTION))
    # Insertion related settings
    shared.opts.add_option("tac_replaceUnderscores", shared.OptionInfo(True, "Replace underscores with spaces on insertion", section=TAC_SECTION))
    shared.opts.add_option("tac_escapeParentheses", shared.OptionInfo(True, "Escape parentheses on insertion", section=TAC_SECTION))
    shared.opts.add_option("tac_appendComma", shared.OptionInfo(True, "Append comma on tag autocompletion", section=TAC_SECTION))
    # Alias settings
    shared.opts.add_option("tac_alias.searchByAlias", shared.OptionInfo(True, "Search by alias", section=TAC_SECTION))
    shared.opts.add_option("tac_alias.onlyShowAlias", shared.OptionInfo(False, "Only show alias", section=TAC_SECTION))
    # Translation settings
    shared.opts.add_option("tac_translation.translationFile", shared.OptionInfo("None", "Translation filename", gr.Dropdown, lambda: {"choices": csv_files_withnone}, refresh=update_tag_files, section=TAC_SECTION))
    shared.opts.add_option("tac_translation.oldFormat", shared.OptionInfo(False, "Translation file uses old 3-column translation format instead of the new 2-column one", section=TAC_SECTION))
    shared.opts.add_option("tac_translation.searchByTranslation", shared.OptionInfo(True, "Search by translation", section=TAC_SECTION))
    # Extra file settings
    shared.opts.add_option("tac_extra.extraFile", shared.OptionInfo("extra-quality-tags.csv", "Extra filename (for small sets of custom tags)", gr.Dropdown, lambda: {"choices": csv_files_withnone}, refresh=update_tag_files, section=TAC_SECTION))
    shared.opts.add_option("tac_extra.addMode", shared.OptionInfo("Insert before", "Mode to add the extra tags to the main tag list", gr.Dropdown, lambda: {"choices": ["Insert before","Insert after"]}, section=TAC_SECTION))

script_callbacks.on_ui_settings(on_ui_settings)