import functools import os.path import urllib.parse from base64 import b64decode from io import BytesIO from pathlib import Path from typing import Optional, Union from dataclasses import dataclass import zlib 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 PIL import Image 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 fetch_cover_images(page: str = "", item: str = "", index: int = 0): from starlette.responses import Response page = next(iter([x for x in extra_pages if x.name == page]), None) if page is None: raise HTTPException(status_code=404, detail="File not found") metadata = page.metadata.get(item) if metadata is None: raise HTTPException(status_code=404, detail="File not found") cover_images = json.loads(metadata.get('ssmd_cover_images', {})) image = cover_images[index] if index < len(cover_images) else None if not image: raise HTTPException(status_code=404, detail="File not found") try: image = Image.open(BytesIO(b64decode(image))) buffer = BytesIO() image.save(buffer, format=image.format) return Response(content=buffer.getvalue(), media_type=image.get_format_mimetype()) except Exception as err: raise ValueError(f"File cannot be fetched: {item}. Failed to load cover image.") from err 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({}) metadata = {i:metadata[i] for i in metadata if i != 'ssmd_cover_images'} # those are cover images, and they are too big to display in UI as text 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/cover-images", fetch_cover_images, 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.tree = {} 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_tree_html_dict_row( self, div_id: int, tabname: str, label: str, btn_type: str, btn_title: 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) label = label.strip() btn_title = btn_title.strip() 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, "btn_title": btn_title, "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_html_dict( 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_tree_html_dict_row( div_id=div_id, parent_id=parent_id, tabname=tabname, label=os.path.basename(k), data_depth=depth, data_path=k, btn_type="dir", btn_title=k, dir_is_empty=dir_is_empty, ) last_div_id = self.build_tree_html_dict( 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_tree_html_dict_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", btn_title=v.item["name"], ) 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, } ) btn_edit_item = self.btn_edit_item_tpl.format( **{ "tabname": tabname, "extra_networks_tabname": self.extra_networks_tabname, } ) 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 = {} if not self.tree: return res self.build_tree_html_dict( tree=self.tree, res=res, depth=0, div_id=0, parent_id=None, tabname=tabname, ) res = base64.b64encode(zlib.compress(json.dumps(res, separators=(",", ":")).encode("utf-8"))).decode("utf-8") return f'' def create_dirs_view_html(self, tabname: str) -> str: """Generates HTML for displaying folders.""" def _get_dirs_buttons(tree: dict, res: list) -> None: """ Builds a list of directory names from a tree. """ for k, v in sorted(tree.items(), key=lambda x: shared.natural_sort_key(x[0])): if not isinstance(v, (ExtraNetworksItem,)): # dir res.append(k) _get_dirs_buttons(tree=v, res=res) dirs = [] _get_dirs_buttons(tree=self.tree, res=dirs) dirs_html = "".join([ self.btn_dirs_view_tpl.format(**{ "extra_class": "search-all" if d == "" else "", "tabname_full": f"{tabname}_{self.extra_networks_tabname}", "path": html.escape(d), }) for d in dirs ]) return dirs_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`