import functools import os.path import urllib.parse from pathlib import Path from typing import Optional, Union from dataclasses import dataclass import gzip import base64 import re from modules import shared, ui_extra_networks_user_metadata, errors, extra_networks, util from modules.images import read_info_from_image, save_image_with_geninfo import gradio as gr import json import html from fastapi.exceptions import HTTPException from modules.infotext_utils import image_from_url_text 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 []) 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], _root: str): _res = {} for path in _paths: relpath = os.path.relpath(path, _root) 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], _root) # We only want to store non-empty folders in the tree. if dir_tree: _res[relpath] = dir_tree else: if path not in items: continue # Add the ExtraNetworksItem to the result. _res[relpath] = 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: root = os.path.dirname(path) relpath = os.path.relpath(path, root) # Wrap the path in a list since that is what the `_get_tree` expects. res[relpath] = _get_tree([path], root) if res[relpath]: # We need to pull the inner path out one for these root dirs. res[relpath] = res[relpath][relpath] return res def register_page(page): """registers extra networks page for the UI; recommend doing it in on_before_ui() callback for extensions""" extra_pages.append(page) allowed_dirs.clear() allowed_dirs.update(set(sum([x.allowed_directories_for_previews() for x in extra_pages], []))) def fetch_file(filename: str = ""): from starlette.responses import FileResponse if not os.path.isfile(filename): raise HTTPException(status_code=404, detail="File not found") if not any(Path(x).absolute() in Path(filename).absolute().parents for x in allowed_dirs): raise ValueError(f"File cannot be fetched: {filename}. Must be in one of directories registered by extra pages.") ext = os.path.splitext(filename)[1].lower()[1:] if ext not in allowed_preview_extensions(): raise ValueError(f"File cannot be fetched: {filename}. Extensions allowed: {allowed_preview_extensions()}.") # would profit from returning 304 return FileResponse(filename, headers={"Accept-Ranges": "bytes"}) def get_metadata(page: str = "", item: str = ""): from starlette.responses import JSONResponse page = next(iter([x for x in extra_pages if x.name == page]), None) if page is None: return JSONResponse({}) metadata = page.metadata.get(item) if metadata is None: return JSONResponse({}) return JSONResponse({"metadata": json.dumps(metadata, indent=4, ensure_ascii=False)}) def get_single_card(page: str = "", tabname: str = "", name: str = ""): from starlette.responses import JSONResponse page = next(iter([x for x in extra_pages if x.name == page]), None) try: item = page.create_item(name, enable_filter=False) page.items[name] = item except Exception as e: errors.display(e, "creating item for extra network") item = page.items.get(name) page.read_user_metadata(item, use_cache=False) item_html = page.create_item_html(tabname, item, shared.html("extra-networks-card.html")) return JSONResponse({"html": item_html}) def add_pages_to_demo(app): app.add_api_route("/sd_extra_networks/thumb", fetch_file, methods=["GET"]) app.add_api_route("/sd_extra_networks/metadata", get_metadata, methods=["GET"]) app.add_api_route("/sd_extra_networks/get-single-card", get_single_card, methods=["GET"]) def quote_js(s): s = s.replace('\\', '\\\\') s = s.replace('"', '\\"') return f'"{s}"' class ExtraNetworksPage: def __init__(self, title): self.title = title self.name = title.lower() # This is the actual name of the extra networks tab (not txt2img/img2img). self.extra_networks_tabname = self.name.replace(" ", "_") self.allow_prompt = True self.allow_negative_prompt = False self.metadata = {} self.items = {} self.lister = util.MassFileLister() # HTML Templates self.pane_tpl = shared.html("extra-networks-pane.html") self.pane_content_tree_tpl = shared.html("extra-networks-pane-tree.html") self.pane_content_dirs_tpl = shared.html("extra-networks-pane-dirs.html") self.card_tpl = shared.html("extra-networks-card.html") self.tree_row_tpl = shared.html("extra-networks-tree-row.html") self.btn_copy_path_tpl = shared.html("extra-networks-copy-path-button.html") self.btn_metadata_tpl = shared.html("extra-networks-metadata-button.html") self.btn_edit_item_tpl = shared.html("extra-networks-edit-item-button.html") self.btn_dirs_view_tpl = shared.html("extra-networks-dirs-view-button.html") def refresh(self): pass def read_user_metadata(self, item, use_cache=True): filename = item.get("filename", None) metadata = extra_networks.get_user_metadata(filename, lister=self.lister if use_cache else None) desc = metadata.get("description", None) if desc is not None: item["description"] = desc item["user_metadata"] = metadata def link_preview(self, filename): quoted_filename = urllib.parse.quote(filename.replace('\\', '/')) mtime, _ = self.lister.mctime(filename) return f"./sd_extra_networks/thumb?filename={quoted_filename}&mtime={mtime}" def search_terms_from_path(self, filename, possible_directories=None): abspath = os.path.abspath(filename) for parentdir in (possible_directories if possible_directories is not None else self.allowed_directories_for_previews()): parentdir = os.path.dirname(os.path.abspath(parentdir)) if abspath.startswith(parentdir): return os.path.relpath(abspath, parentdir) return "" def build_row( self, div_id: int, tabname: str, label: str, btn_type: str, dir_is_empty: bool = False, metadata: Optional[str] = None, parent_id: Optional[int] = None, data_depth: Optional[int] = None, data_path: Optional[str] = None, data_hash: Optional[str] = None, data_prompt: Optional[str] = None, data_neg_prompt: Optional[str] = None, data_allow_neg: Optional[str] = None, onclick_extra: Optional[str] = None, ) -> str: if btn_type not in ["file", "dir"]: raise ValueError("Invalid button type:", btn_type) action_list_item_action_leading = "" action_list_item_visual_leading = "🗀" action_list_item_visual_trailing = "" action_list_item_action_trailing = "" if dir_is_empty: action_list_item_action_leading = "" if btn_type == "file": action_list_item_visual_leading = "🗎" # Action buttons action_list_item_action_trailing += '
" data_attributes = "" data_attributes += f"data-path={data_path} " if data_path is not None else "" data_attributes += f"data-hash={data_hash} " if data_hash is not None else "" data_attributes += f"data-prompt={data_prompt} " if data_prompt else "" data_attributes += f"data-neg-prompt={data_neg_prompt} " if data_neg_prompt else "" data_attributes += f"data-allow-neg={data_allow_neg} " if data_allow_neg else "" data_attributes += ( f"data-tree-entry-type={btn_type} " if btn_type is not None else "" ) data_attributes += f"data-div-id={div_id} " if div_id is not None else "" data_attributes += f"data-parent-id={parent_id} " if parent_id is not None else "" data_attributes += f"data-depth={data_depth} " if data_depth is not None else "" data_attributes += ( "data-expanded " if parent_id is None else "" ) # inverted to expand root res = self.tree_row_tpl.format( **{ "data_attributes": data_attributes, "search_terms": "", "btn_type": btn_type, "tabname": tabname, "onclick_extra": onclick_extra if onclick_extra else "", "extra_networks_tabname": self.extra_networks_tabname, "action_list_item_action_leading": action_list_item_action_leading, "action_list_item_visual_leading": action_list_item_visual_leading, "action_list_item_label": label, "action_list_item_visual_trailing": action_list_item_visual_trailing, "action_list_item_action_trailing": action_list_item_action_trailing, } ) res = res.strip() res = re.sub(" +", " ", res.replace("\n", "")) return res def build_tree( self, tree: dict, res: dict, tabname: str, div_id: int, depth: int, parent_id: Optional[int] = None, ) -> int: for k, v in sorted(tree.items(), key=lambda x: shared.natural_sort_key(x[0])): if not isinstance(v, (ExtraNetworksItem,)): # dir if div_id in res: raise KeyError("div_id already in res:", div_id) dir_is_empty = True for _v in v.values(): if shared.opts.extra_networks_tree_view_show_files: dir_is_empty = False break elif not isinstance(_v, (ExtraNetworksItem,)): dir_is_empty = False break else: dir_is_empty = True res[div_id] = self.build_row( div_id=div_id, parent_id=parent_id, tabname=tabname, label=k, data_depth=depth, data_path=k, btn_type="dir", dir_is_empty=dir_is_empty, ) last_div_id = self.build_tree( tree=v, res=res, depth=depth + 1, div_id=div_id + 1, parent_id=div_id, tabname=tabname, ) div_id = last_div_id else: # file if not shared.opts.extra_networks_tree_view_show_files: # Don't add file if showing files is disabled in options. continue if div_id in res: raise KeyError("div_id already in res:", div_id) onclick = v.item.get("onclick", None) if onclick is None: # Don't quote prompt/neg_prompt since they are stored as js strings already. onclick_js_tpl = ( "cardClicked('{tabname}', {prompt}, {neg_prompt}, {allow_neg});" ) onclick = onclick_js_tpl.format( **{ "tabname": tabname, "extra_networks_tabname": self.extra_networks_tabname, "prompt": v.item["prompt"], "neg_prompt": v.item.get("negative_prompt", "''"), "allow_neg": str(self.allow_negative_prompt).lower(), } ) onclick = html.escape(onclick) res[div_id] = self.build_row( div_id=div_id, parent_id=parent_id, tabname=tabname, label=v.item["name"], metadata=v.item.get("metadata", None), data_depth=depth, data_path=v.item["filename"], data_hash=v.item["shorthash"], data_prompt=html.escape(v.item.get("prompt", "''")), data_neg_prompt=html.escape(v.item.get("negative_prompt", "''")), data_allow_neg=str(self.allow_negative_prompt).lower(), onclick_extra=onclick, btn_type="file", ) div_id += 1 return div_id def create_item_html( self, tabname: str, item: dict, template: Optional[str] = None, div_id: Optional[int] = None, ) -> Union[str, dict]: """Generates HTML for a single ExtraNetworks Item. Args: tabname: The name of the active tab. item: Dictionary containing item information. template: Optional template string to use. Returns: If a template is passed: HTML string generated for this item. Can be empty if the item is not meant to be shown. If no template is passed: A dictionary containing the generated item's attributes. """ preview = item.get("preview", None) style_height = f"height: {shared.opts.extra_networks_card_height}px;" if shared.opts.extra_networks_card_height else '' style_width = f"width: {shared.opts.extra_networks_card_width}px;" if shared.opts.extra_networks_card_width else '' style_font_size = f"font-size: {shared.opts.extra_networks_card_text_scale*100}%;" card_style = style_height + style_width + style_font_size background_image = f'' if preview else '' onclick = item.get("onclick", None) if onclick is None: # Don't quote prompt/neg_prompt since they are stored as js strings already. onclick_js_tpl = "cardClicked('{tabname}', {prompt}, {neg_prompt}, {allow_neg});" onclick = onclick_js_tpl.format( **{ "tabname": tabname, "extra_networks_tabname": self.extra_networks_tabname, "prompt": item["prompt"], "neg_prompt": item.get("negative_prompt", "''"), "allow_neg": str(self.allow_negative_prompt).lower(), } ) onclick = html.escape(onclick) btn_copy_path = self.btn_copy_path_tpl.format(**{"filename": item["filename"]}) btn_metadata = "" metadata = item.get("metadata") if metadata: btn_metadata = self.btn_metadata_tpl.format( **{ "extra_networks_tabname": self.extra_networks_tabname, "name": html.escape(item["name"]), } ) btn_edit_item = self.btn_edit_item_tpl.format( **{ "tabname": tabname, "extra_networks_tabname": self.extra_networks_tabname, "name": html.escape(item["name"]), } ) 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() search_terms_html = "" search_term_template = "{search_term}" for search_term in item.get("search_terms", []): search_terms_html += search_term_template.format( **{ "class": f"search_terms{' search_only' if search_only else ''}", "search_term": search_term, } ) description = (item.get("description", "") or "" if shared.opts.extra_networks_card_show_desc else "") if not shared.opts.extra_networks_card_description_is_html: description = html.escape(description) # Some items here might not be used depending on HTML template used. args = { "div_id": "" if div_id is None else div_id, "background_image": background_image, "card_clicked": onclick, "copy_path_button": btn_copy_path, "description": description, "edit_button": btn_edit_item, "local_preview": quote_js(item["local_preview"]), "metadata_button": btn_metadata, "name": html.escape(item["name"]), "save_card_preview": html.escape(f"return saveCardPreview(event, '{tabname}', '{item['local_preview']}');"), "search_only": " search_only" if search_only else "", "search_terms": search_terms_html, "sort_keys": sort_keys, "style": card_style, "tabname": tabname, "extra_networks_tabname": self.extra_networks_tabname, "data_prompt": item.get("prompt", "''"), "data_neg_prompt": item.get("negative_prompt", "''"), "data_allow_neg": str(self.allow_negative_prompt).lower(), } if template: return template.format(**args) else: return args def generate_tree_view_data_div(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. """ res = {} # Setup the tree dictionary. 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 self.build_tree( tree=tree, res=res, depth=0, div_id=0, parent_id=None, tabname=tabname, ) res = base64.b64encode(gzip.compress(json.dumps(res).encode("utf-8"))).decode("utf-8") return f'' # FIXME def create_dirs_view_html(self, tabname: str) -> str: """Generates HTML for displaying folders.""" subdirs = {} for parentdir in [os.path.abspath(x) for x in self.allowed_directories_for_previews()]: for root, dirs, _ in sorted(os.walk(parentdir, followlinks=True), key=lambda x: shared.natural_sort_key(x[0])): for dirname in sorted(dirs, key=shared.natural_sort_key): x = os.path.join(root, dirname) if not os.path.isdir(x): continue subdir = os.path.abspath(x)[len(parentdir):] if shared.opts.extra_networks_dir_button_function: if not subdir.startswith(os.path.sep): subdir = os.path.sep + subdir else: while subdir.startswith(os.path.sep): subdir = subdir[1:] is_empty = len(os.listdir(x)) == 0 if not is_empty and not subdir.endswith(os.path.sep): subdir = subdir + os.path.sep if (os.path.sep + "." in subdir or subdir.startswith(".")) and not shared.opts.extra_networks_show_hidden_directories: continue subdirs[subdir] = 1 if subdirs: subdirs = {**subdirs} subdirs_html = "".join([ self.btn_dirs_view_tpl.format(**{ "extra_class": "search-all" if subdir == "" else "", "tabname_full": f"{tabname}_{self.extra_networks_tabname}", "path": html.escape(subdir), }) for subdir in subdirs ]) return subdirs_html def generate_cards_view_data_div(self, tabname: str, *, none_message) -> str: """Generates HTML for the network Card View section for a tab. This HTML goes into the `extra-networks-pane.html`