From df8aa69a99e38ae59a4e599b9dff11eccf3490f4 Mon Sep 17 00:00:00 2001 From: Sj-Si Date: Mon, 8 Jan 2024 14:10:03 -0500 Subject: Add tree-view display for extra networks. --- modules/ui_extra_networks.py | 352 +++++++++++++++++++++++++++++++------------ 1 file changed, 253 insertions(+), 99 deletions(-) (limited to 'modules/ui_extra_networks.py') diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index fe5d3ba3..8667617b 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -2,6 +2,8 @@ import functools import os.path import urllib.parse from pathlib import Path +from typing import Optional, Union +from dataclasses import dataclass from modules import shared, ui_extra_networks_user_metadata, errors, extra_networks from modules.images import read_info_from_image, save_image_with_geninfo @@ -15,10 +17,8 @@ from modules.ui_components import ToolButton extra_pages = [] allowed_dirs = set() - default_allowed_preview_extensions = ["png", "jpg", "jpeg", "webp", "gif"] - @functools.cache def allowed_preview_extensions_with_extra(extra_extensions=None): return set(default_allowed_preview_extensions) | set(extra_extensions or []) @@ -28,6 +28,58 @@ def allowed_preview_extensions(): return allowed_preview_extensions_with_extra((shared.opts.samples_format, )) +@dataclass +class ExtraNetworksItem: + """Wrapper for dictionaries representing ExtraNetworks items.""" + item: dict + + +def get_tree(paths: Union[str, list[str]], items: dict[str, ExtraNetworksItem]) -> dict: + """Recursively builds a directory tree. + + Args: + paths: Path or list of paths to directories. These paths are treated as roots from which + the tree will be built. + items: A dictionary associating filepaths to an ExtraNetworksItem instance. + + Returns: + The result directory tree. + """ + if isinstance(paths, (str,)): + paths = [paths] + + def _get_tree(_paths: list[str]): + _res = {} + for path in _paths: + if os.path.isdir(path): + dir_items = os.listdir(path) + # Ignore empty directories. + if not dir_items: + continue + dir_tree = _get_tree([os.path.join(path, x) for x in dir_items]) + # We only want to store non-empty folders in the tree. + if dir_tree: + _res[os.path.basename(path)] = dir_tree + else: + if path not in items: + continue + # Add the ExtraNetworksItem to the result. + _res[os.path.basename(path)] = items[path] + return _res + + res = {} + # Handle each root directory separately. + # Each root WILL have a key/value at the root of the result dict though + # the value can be an empty dict if the directory is empty. We want these + # placeholders for empty dirs so we can inform the user later. + for path in paths: + # Wrap the path in a list since that is what the `_get_tree` expects. + res[path] = _get_tree([path]) + if res[path]: + res[path] = res[path][os.path.basename(path)] + + return res + def register_page(page): """registers extra networks page for the UI; recommend doing it in on_before_ui() callback for extensions""" @@ -80,7 +132,7 @@ def get_single_card(page: str = "", tabname: str = "", name: str = ""): item = page.items.get(name) page.read_user_metadata(item) - item_html = page.create_html_for_item(item, tabname) + item_html = page.create_item_html(tabname, item) return JSONResponse({"html": item_html}) @@ -96,13 +148,15 @@ def quote_js(s): s = s.replace('"', '\\"') return f'"{s}"' - class ExtraNetworksPage: def __init__(self, title): self.title = title self.name = title.lower() self.id_page = self.name.replace(" ", "_") - self.card_page = shared.html("extra-networks-card.html") + if shared.opts.extra_networks_tree_view: + self.card_page = shared.html("extra-networks-card-minimal.html") + else: + self.card_page = shared.html("extra-networks-card.html") self.allow_prompt = True self.allow_negative_prompt = False self.metadata = {} @@ -136,12 +190,141 @@ class ExtraNetworksPage: return "" - def create_html(self, tabname): - items_html = '' + def create_item_html(self, tabname: str, item: dict) -> str: + """Generates HTML for a single ExtraNetworks Item + + Args: + tabname: The name of the active tab. + item: Dictionary containing item information. + + Returns: + HTML string generated for this item. + Can be empty if the item is not meant to be shown. + """ + metadata = item.get("metadata") + if metadata: + self.metadata[item["name"]] = metadata + + if "user_metadata" not in item: + self.read_user_metadata(item) + + preview = item.get("preview", None) + height = f"height: {shared.opts.extra_networks_card_height}px;" if shared.opts.extra_networks_card_height else '' + width = f"width: {shared.opts.extra_networks_card_width}px;" if shared.opts.extra_networks_card_width else '' + background_image = f'' if preview else '' + + onclick = item.get("onclick", None) + if onclick is None: + onclick = '"' + html.escape(f"""return cardClicked({quote_js(tabname)}, {item["prompt"]}, {"true" if self.allow_negative_prompt else "false"})""") + '"' + + copy_path_button = f"
" + + metadata_button = "" + metadata = item.get("metadata") + if metadata: + metadata_button = f"
" + + edit_button = f"
" + + local_path = "" + filename = item.get("filename", "") + for reldir in self.allowed_directories_for_previews(): + absdir = os.path.abspath(reldir) + + if filename.startswith(absdir): + local_path = filename[len(absdir):] + + # if this is true, the item must not be shown in the default view, and must instead only be + # shown when searching for it + if shared.opts.extra_networks_hidden_models == "Always": + search_only = False + else: + search_only = "/." in local_path or "\\." in local_path + + if search_only and shared.opts.extra_networks_hidden_models == "Never": + return "" + + sort_keys = " ".join([f'data-sort-{k}="{html.escape(str(v))}"' for k, v in item.get("sort_keys", {}).items()]).strip() + + # Some items here might not be used depending on HTML template used. + args = { + "background_image": background_image, + "card_clicked": onclick, + "copy_path_button": copy_path_button, + "description": (item.get("description") or "" if shared.opts.extra_networks_card_show_desc else ""), + "edit_button": edit_button, + "local_preview": quote_js(item["local_preview"]), + "metadata_button": metadata_button, + "name": html.escape(item["name"]), + "prompt": item.get("prompt", None), + "save_card_preview": '"' + html.escape(f"""return saveCardPreview(event, {quote_js(tabname)}, {quote_js(item["local_preview"])})""") + '"', + "search_only": " search_only" if search_only else "", + "search_term": item.get("search_term", ""), + "sort_keys": sort_keys, + "style": f"'display: none; {height}{width}; font-size: {shared.opts.extra_networks_card_text_scale*100}%'", + "tabname": quote_js(tabname), + } + + return self.card_page.format(**args) + + def create_tree_view_html(self, tabname: str) -> str: + """Generates HTML for displaying folders in a tree view. + + Args: + tabname: The name of the active tab. + + Returns: + HTML string generated for this tree view. + """ + self_name_id = self.name.replace(" ", "_") + res = f"
" self.metadata = {} + self.items = {x["name"]: x for x in self.list_items()} + roots = self.allowed_directories_for_previews() + tree_items = {v["filename"]: ExtraNetworksItem(v) for v in self.items.values()} + tree = get_tree([os.path.abspath(x) for x in roots], items=tree_items) + + if not tree: + return res + "
" + + file_template = "
  • {}
    • " + dir_template = ( + "
    " + "{}" + "{}" + "
    " + ) + + def _build_tree(data: Optional[dict[str, ExtraNetworksItem]] = None) -> str: + """Recursively builds HTML for a tree.""" + _res = "