ui_extra_networks.py 33 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836
  1. import functools
  2. import os.path
  3. import urllib.parse
  4. from base64 import b64decode
  5. from io import BytesIO
  6. from pathlib import Path
  7. from typing import Optional, Union
  8. from dataclasses import dataclass
  9. from modules import shared, ui_extra_networks_user_metadata, errors, extra_networks, util
  10. from modules.images import read_info_from_image, save_image_with_geninfo
  11. import gradio as gr
  12. import json
  13. import html
  14. from fastapi.exceptions import HTTPException
  15. from PIL import Image
  16. from modules.infotext_utils import image_from_url_text
  17. extra_pages = []
  18. allowed_dirs = set()
  19. default_allowed_preview_extensions = ["png", "jpg", "jpeg", "webp", "gif"]
  20. @functools.cache
  21. def allowed_preview_extensions_with_extra(extra_extensions=None):
  22. return set(default_allowed_preview_extensions) | set(extra_extensions or [])
  23. def allowed_preview_extensions():
  24. return allowed_preview_extensions_with_extra((shared.opts.samples_format, ))
  25. @dataclass
  26. class ExtraNetworksItem:
  27. """Wrapper for dictionaries representing ExtraNetworks items."""
  28. item: dict
  29. def get_tree(paths: Union[str, list[str]], items: dict[str, ExtraNetworksItem]) -> dict:
  30. """Recursively builds a directory tree.
  31. Args:
  32. paths: Path or list of paths to directories. These paths are treated as roots from which
  33. the tree will be built.
  34. items: A dictionary associating filepaths to an ExtraNetworksItem instance.
  35. Returns:
  36. The result directory tree.
  37. """
  38. if isinstance(paths, (str,)):
  39. paths = [paths]
  40. def _get_tree(_paths: list[str], _root: str):
  41. _res = {}
  42. for path in _paths:
  43. relpath = os.path.relpath(path, _root)
  44. if os.path.isdir(path):
  45. dir_items = os.listdir(path)
  46. # Ignore empty directories.
  47. if not dir_items:
  48. continue
  49. dir_tree = _get_tree([os.path.join(path, x) for x in dir_items], _root)
  50. # We only want to store non-empty folders in the tree.
  51. if dir_tree:
  52. _res[relpath] = dir_tree
  53. else:
  54. if path not in items:
  55. continue
  56. # Add the ExtraNetworksItem to the result.
  57. _res[relpath] = items[path]
  58. return _res
  59. res = {}
  60. # Handle each root directory separately.
  61. # Each root WILL have a key/value at the root of the result dict though
  62. # the value can be an empty dict if the directory is empty. We want these
  63. # placeholders for empty dirs so we can inform the user later.
  64. for path in paths:
  65. root = os.path.dirname(path)
  66. relpath = os.path.relpath(path, root)
  67. # Wrap the path in a list since that is what the `_get_tree` expects.
  68. res[relpath] = _get_tree([path], root)
  69. if res[relpath]:
  70. # We need to pull the inner path out one for these root dirs.
  71. res[relpath] = res[relpath][relpath]
  72. return res
  73. def register_page(page):
  74. """registers extra networks page for the UI; recommend doing it in on_before_ui() callback for extensions"""
  75. extra_pages.append(page)
  76. allowed_dirs.clear()
  77. allowed_dirs.update(set(sum([x.allowed_directories_for_previews() for x in extra_pages], [])))
  78. def fetch_file(filename: str = ""):
  79. from starlette.responses import FileResponse
  80. if not os.path.isfile(filename):
  81. raise HTTPException(status_code=404, detail="File not found")
  82. if not any(Path(x).absolute() in Path(filename).absolute().parents for x in allowed_dirs):
  83. raise ValueError(f"File cannot be fetched: {filename}. Must be in one of directories registered by extra pages.")
  84. ext = os.path.splitext(filename)[1].lower()[1:]
  85. if ext not in allowed_preview_extensions():
  86. raise ValueError(f"File cannot be fetched: {filename}. Extensions allowed: {allowed_preview_extensions()}.")
  87. # would profit from returning 304
  88. return FileResponse(filename, headers={"Accept-Ranges": "bytes"})
  89. def fetch_cover_images(page: str = "", item: str = "", index: int = 0):
  90. from starlette.responses import Response
  91. page = next(iter([x for x in extra_pages if x.name == page]), None)
  92. if page is None:
  93. raise HTTPException(status_code=404, detail="File not found")
  94. metadata = page.metadata.get(item)
  95. if metadata is None:
  96. raise HTTPException(status_code=404, detail="File not found")
  97. cover_images = json.loads(metadata.get('ssmd_cover_images', {}))
  98. image = cover_images[index] if index < len(cover_images) else None
  99. if not image:
  100. raise HTTPException(status_code=404, detail="File not found")
  101. try:
  102. image = Image.open(BytesIO(b64decode(image)))
  103. buffer = BytesIO()
  104. image.save(buffer, format=image.format)
  105. return Response(content=buffer.getvalue(), media_type=image.get_format_mimetype())
  106. except Exception as err:
  107. raise ValueError(f"File cannot be fetched: {item}. Failed to load cover image.") from err
  108. def get_metadata(page: str = "", item: str = ""):
  109. from starlette.responses import JSONResponse
  110. page = next(iter([x for x in extra_pages if x.name == page]), None)
  111. if page is None:
  112. return JSONResponse({})
  113. metadata = page.metadata.get(item)
  114. if metadata is None:
  115. return JSONResponse({})
  116. 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
  117. return JSONResponse({"metadata": json.dumps(metadata, indent=4, ensure_ascii=False)})
  118. def get_single_card(page: str = "", tabname: str = "", name: str = ""):
  119. from starlette.responses import JSONResponse
  120. page = next(iter([x for x in extra_pages if x.name == page]), None)
  121. try:
  122. item = page.create_item(name, enable_filter=False)
  123. page.items[name] = item
  124. except Exception as e:
  125. errors.display(e, "creating item for extra network")
  126. item = page.items.get(name)
  127. page.read_user_metadata(item, use_cache=False)
  128. item_html = page.create_item_html(tabname, item, shared.html("extra-networks-card.html"))
  129. return JSONResponse({"html": item_html})
  130. def add_pages_to_demo(app):
  131. app.add_api_route("/sd_extra_networks/thumb", fetch_file, methods=["GET"])
  132. app.add_api_route("/sd_extra_networks/cover-images", fetch_cover_images, methods=["GET"])
  133. app.add_api_route("/sd_extra_networks/metadata", get_metadata, methods=["GET"])
  134. app.add_api_route("/sd_extra_networks/get-single-card", get_single_card, methods=["GET"])
  135. def quote_js(s: str):
  136. return json.dumps(s, ensure_ascii=False)
  137. class ExtraNetworksPage:
  138. def __init__(self, title):
  139. self.title = title
  140. self.name = title.lower()
  141. # This is the actual name of the extra networks tab (not txt2img/img2img).
  142. self.extra_networks_tabname = self.name.replace(" ", "_")
  143. self.allow_prompt = True
  144. self.allow_negative_prompt = False
  145. self.metadata = {}
  146. self.items = {}
  147. self.lister = util.MassFileLister()
  148. # HTML Templates
  149. self.pane_tpl = shared.html("extra-networks-pane.html")
  150. self.pane_content_tree_tpl = shared.html("extra-networks-pane-tree.html")
  151. self.pane_content_dirs_tpl = shared.html("extra-networks-pane-dirs.html")
  152. self.card_tpl = shared.html("extra-networks-card.html")
  153. self.btn_tree_tpl = shared.html("extra-networks-tree-button.html")
  154. self.btn_copy_path_tpl = shared.html("extra-networks-copy-path-button.html")
  155. self.btn_metadata_tpl = shared.html("extra-networks-metadata-button.html")
  156. self.btn_edit_item_tpl = shared.html("extra-networks-edit-item-button.html")
  157. def refresh(self):
  158. pass
  159. def read_user_metadata(self, item, use_cache=True):
  160. filename = item.get("filename", None)
  161. metadata = extra_networks.get_user_metadata(filename, lister=self.lister if use_cache else None)
  162. desc = metadata.get("description", None)
  163. if desc is not None:
  164. item["description"] = desc
  165. item["user_metadata"] = metadata
  166. def link_preview(self, filename):
  167. quoted_filename = urllib.parse.quote(filename.replace('\\', '/'))
  168. mtime, _ = self.lister.mctime(filename)
  169. return f"./sd_extra_networks/thumb?filename={quoted_filename}&mtime={mtime}"
  170. def search_terms_from_path(self, filename, possible_directories=None):
  171. abspath = os.path.abspath(filename)
  172. for parentdir in (possible_directories if possible_directories is not None else self.allowed_directories_for_previews()):
  173. parentdir = os.path.dirname(os.path.abspath(parentdir))
  174. if abspath.startswith(parentdir):
  175. return os.path.relpath(abspath, parentdir)
  176. return ""
  177. def create_item_html(
  178. self,
  179. tabname: str,
  180. item: dict,
  181. template: Optional[str] = None,
  182. ) -> Union[str, dict]:
  183. """Generates HTML for a single ExtraNetworks Item.
  184. Args:
  185. tabname: The name of the active tab.
  186. item: Dictionary containing item information.
  187. template: Optional template string to use.
  188. Returns:
  189. If a template is passed: HTML string generated for this item.
  190. Can be empty if the item is not meant to be shown.
  191. If no template is passed: A dictionary containing the generated item's attributes.
  192. """
  193. preview = item.get("preview", None)
  194. style_height = f"height: {shared.opts.extra_networks_card_height}px;" if shared.opts.extra_networks_card_height else ''
  195. style_width = f"width: {shared.opts.extra_networks_card_width}px;" if shared.opts.extra_networks_card_width else ''
  196. style_font_size = f"font-size: {shared.opts.extra_networks_card_text_scale*100}%;"
  197. card_style = style_height + style_width + style_font_size
  198. background_image = f'<img src="{html.escape(preview)}" class="preview" loading="lazy">' if preview else ''
  199. onclick = item.get("onclick", None)
  200. if onclick is None:
  201. # Don't quote prompt/neg_prompt since they are stored as js strings already.
  202. onclick_js_tpl = "cardClicked('{tabname}', {prompt}, {neg_prompt}, {allow_neg});"
  203. onclick = onclick_js_tpl.format(
  204. **{
  205. "tabname": tabname,
  206. "prompt": item["prompt"],
  207. "neg_prompt": item.get("negative_prompt", "''"),
  208. "allow_neg": str(self.allow_negative_prompt).lower(),
  209. }
  210. )
  211. onclick = html.escape(onclick)
  212. btn_copy_path = self.btn_copy_path_tpl.format(**{"filename": item["filename"]})
  213. btn_metadata = ""
  214. metadata = item.get("metadata")
  215. if metadata:
  216. btn_metadata = self.btn_metadata_tpl.format(
  217. **{
  218. "extra_networks_tabname": self.extra_networks_tabname,
  219. }
  220. )
  221. btn_edit_item = self.btn_edit_item_tpl.format(
  222. **{
  223. "tabname": tabname,
  224. "extra_networks_tabname": self.extra_networks_tabname,
  225. }
  226. )
  227. local_path = ""
  228. filename = item.get("filename", "")
  229. for reldir in self.allowed_directories_for_previews():
  230. absdir = os.path.abspath(reldir)
  231. if filename.startswith(absdir):
  232. local_path = filename[len(absdir):]
  233. # if this is true, the item must not be shown in the default view, and must instead only be
  234. # shown when searching for it
  235. if shared.opts.extra_networks_hidden_models == "Always":
  236. search_only = False
  237. else:
  238. search_only = "/." in local_path or "\\." in local_path
  239. if search_only and shared.opts.extra_networks_hidden_models == "Never":
  240. return ""
  241. sort_keys = " ".join(
  242. [
  243. f'data-sort-{k}="{html.escape(str(v))}"'
  244. for k, v in item.get("sort_keys", {}).items()
  245. ]
  246. ).strip()
  247. search_terms_html = ""
  248. search_term_template = "<span class='hidden {class}'>{search_term}</span>"
  249. for search_term in item.get("search_terms", []):
  250. search_terms_html += search_term_template.format(
  251. **{
  252. "class": f"search_terms{' search_only' if search_only else ''}",
  253. "search_term": search_term,
  254. }
  255. )
  256. description = (item.get("description", "") or "" if shared.opts.extra_networks_card_show_desc else "")
  257. if not shared.opts.extra_networks_card_description_is_html:
  258. description = html.escape(description)
  259. # Some items here might not be used depending on HTML template used.
  260. args = {
  261. "background_image": background_image,
  262. "card_clicked": onclick,
  263. "copy_path_button": btn_copy_path,
  264. "description": description,
  265. "edit_button": btn_edit_item,
  266. "local_preview": quote_js(item["local_preview"]),
  267. "metadata_button": btn_metadata,
  268. "name": html.escape(item["name"]),
  269. "prompt": item.get("prompt", None),
  270. "save_card_preview": html.escape(f"return saveCardPreview(event, '{tabname}', '{item['local_preview']}');"),
  271. "search_only": " search_only" if search_only else "",
  272. "search_terms": search_terms_html,
  273. "sort_keys": sort_keys,
  274. "style": card_style,
  275. "tabname": tabname,
  276. "extra_networks_tabname": self.extra_networks_tabname,
  277. }
  278. if template:
  279. return template.format(**args)
  280. else:
  281. return args
  282. def create_tree_dir_item_html(
  283. self,
  284. tabname: str,
  285. dir_path: str,
  286. content: Optional[str] = None,
  287. ) -> Optional[str]:
  288. """Generates HTML for a directory item in the tree.
  289. The generated HTML is of the format:
  290. ```html
  291. <li class="tree-list-item tree-list-item--has-subitem">
  292. <div class="tree-list-content tree-list-content-dir"></div>
  293. <ul class="tree-list tree-list--subgroup">
  294. {content}
  295. </ul>
  296. </li>
  297. ```
  298. Args:
  299. tabname: The name of the active tab.
  300. dir_path: Path to the directory for this item.
  301. content: Optional HTML string that will be wrapped by this <ul>.
  302. Returns:
  303. HTML formatted string.
  304. """
  305. if not content:
  306. return None
  307. btn = self.btn_tree_tpl.format(
  308. **{
  309. "search_terms": "",
  310. "subclass": "tree-list-content-dir",
  311. "tabname": tabname,
  312. "extra_networks_tabname": self.extra_networks_tabname,
  313. "onclick_extra": "",
  314. "data_path": dir_path,
  315. "data_hash": "",
  316. "action_list_item_action_leading": "<i class='tree-list-item-action-chevron'></i>",
  317. "action_list_item_visual_leading": "🗀",
  318. "action_list_item_label": os.path.basename(dir_path),
  319. "action_list_item_visual_trailing": "",
  320. "action_list_item_action_trailing": "",
  321. }
  322. )
  323. ul = f"<ul class='tree-list tree-list--subgroup' hidden>{content}</ul>"
  324. return (
  325. "<li class='tree-list-item tree-list-item--has-subitem' data-tree-entry-type='dir'>"
  326. f"{btn}{ul}"
  327. "</li>"
  328. )
  329. def create_tree_file_item_html(self, tabname: str, file_path: str, item: dict) -> str:
  330. """Generates HTML for a file item in the tree.
  331. The generated HTML is of the format:
  332. ```html
  333. <li class="tree-list-item tree-list-item--subitem">
  334. <span data-filterable-item-text hidden></span>
  335. <div class="tree-list-content tree-list-content-file"></div>
  336. </li>
  337. ```
  338. Args:
  339. tabname: The name of the active tab.
  340. file_path: The path to the file for this item.
  341. item: Dictionary containing the item information.
  342. Returns:
  343. HTML formatted string.
  344. """
  345. item_html_args = self.create_item_html(tabname, item)
  346. action_buttons = "".join(
  347. [
  348. item_html_args["copy_path_button"],
  349. item_html_args["metadata_button"],
  350. item_html_args["edit_button"],
  351. ]
  352. )
  353. action_buttons = f"<div class=\"button-row\">{action_buttons}</div>"
  354. btn = self.btn_tree_tpl.format(
  355. **{
  356. "search_terms": "",
  357. "subclass": "tree-list-content-file",
  358. "tabname": tabname,
  359. "extra_networks_tabname": self.extra_networks_tabname,
  360. "onclick_extra": item_html_args["card_clicked"],
  361. "data_path": file_path,
  362. "data_hash": item["shorthash"],
  363. "action_list_item_action_leading": "<i class='tree-list-item-action-chevron'></i>",
  364. "action_list_item_visual_leading": "🗎",
  365. "action_list_item_label": item["name"],
  366. "action_list_item_visual_trailing": "",
  367. "action_list_item_action_trailing": action_buttons,
  368. }
  369. )
  370. return (
  371. "<li class='tree-list-item tree-list-item--subitem' data-tree-entry-type='file'>"
  372. f"{btn}"
  373. "</li>"
  374. )
  375. def create_tree_view_html(self, tabname: str) -> str:
  376. """Generates HTML for displaying folders in a tree view.
  377. Args:
  378. tabname: The name of the active tab.
  379. Returns:
  380. HTML string generated for this tree view.
  381. """
  382. res = ""
  383. # Setup the tree dictionary.
  384. roots = self.allowed_directories_for_previews()
  385. tree_items = {v["filename"]: ExtraNetworksItem(v) for v in self.items.values()}
  386. tree = get_tree([os.path.abspath(x) for x in roots], items=tree_items)
  387. if not tree:
  388. return res
  389. def _build_tree(data: Optional[dict[str, ExtraNetworksItem]] = None) -> Optional[str]:
  390. """Recursively builds HTML for a tree.
  391. Args:
  392. data: Dictionary representing a directory tree. Can be NoneType.
  393. Data keys should be absolute paths from the root and values
  394. should be subdirectory trees or an ExtraNetworksItem.
  395. Returns:
  396. If data is not None: HTML string
  397. Else: None
  398. """
  399. if not data:
  400. return None
  401. # Lists for storing <li> items html for directories and files separately.
  402. _dir_li = []
  403. _file_li = []
  404. for k, v in sorted(data.items(), key=lambda x: shared.natural_sort_key(x[0])):
  405. if isinstance(v, (ExtraNetworksItem,)):
  406. _file_li.append(self.create_tree_file_item_html(tabname, k, v.item))
  407. else:
  408. _dir_li.append(self.create_tree_dir_item_html(tabname, k, _build_tree(v)))
  409. # Directories should always be displayed before files so we order them here.
  410. return "".join(_dir_li) + "".join(_file_li)
  411. # Add each root directory to the tree.
  412. for k, v in sorted(tree.items(), key=lambda x: shared.natural_sort_key(x[0])):
  413. item_html = self.create_tree_dir_item_html(tabname, k, _build_tree(v))
  414. # Only add non-empty entries to the tree.
  415. if item_html is not None:
  416. res += item_html
  417. return f"<ul class='tree-list tree-list--tree'>{res}</ul>"
  418. def create_dirs_view_html(self, tabname: str) -> str:
  419. """Generates HTML for displaying folders."""
  420. subdirs = {}
  421. for parentdir in [os.path.abspath(x) for x in self.allowed_directories_for_previews()]:
  422. for root, dirs, _ in sorted(os.walk(parentdir, followlinks=True), key=lambda x: shared.natural_sort_key(x[0])):
  423. for dirname in sorted(dirs, key=shared.natural_sort_key):
  424. x = os.path.join(root, dirname)
  425. if not os.path.isdir(x):
  426. continue
  427. subdir = os.path.abspath(x)[len(parentdir):]
  428. if shared.opts.extra_networks_dir_button_function:
  429. if not subdir.startswith(os.path.sep):
  430. subdir = os.path.sep + subdir
  431. else:
  432. while subdir.startswith(os.path.sep):
  433. subdir = subdir[1:]
  434. is_empty = len(os.listdir(x)) == 0
  435. if not is_empty and not subdir.endswith(os.path.sep):
  436. subdir = subdir + os.path.sep
  437. if (os.path.sep + "." in subdir or subdir.startswith(".")) and not shared.opts.extra_networks_show_hidden_directories:
  438. continue
  439. subdirs[subdir] = 1
  440. if subdirs:
  441. subdirs = {"": 1, **subdirs}
  442. subdirs_html = "".join([f"""
  443. <button class='lg secondary gradio-button custom-button{" search-all" if subdir == "" else ""}' onclick='extraNetworksSearchButton("{tabname}", "{self.extra_networks_tabname}", event)'>
  444. {html.escape(subdir if subdir != "" else "all")}
  445. </button>
  446. """ for subdir in subdirs])
  447. return subdirs_html
  448. def create_card_view_html(self, tabname: str, *, none_message) -> str:
  449. """Generates HTML for the network Card View section for a tab.
  450. This HTML goes into the `extra-networks-pane.html` <div> with
  451. `id='{tabname}_{extra_networks_tabname}_cards`.
  452. Args:
  453. tabname: The name of the active tab.
  454. none_message: HTML text to show when there are no cards.
  455. Returns:
  456. HTML formatted string.
  457. """
  458. res = []
  459. for item in self.items.values():
  460. res.append(self.create_item_html(tabname, item, self.card_tpl))
  461. if not res:
  462. dirs = "".join([f"<li>{x}</li>" for x in self.allowed_directories_for_previews()])
  463. res = [none_message or shared.html("extra-networks-no-cards.html").format(dirs=dirs)]
  464. return "".join(res)
  465. def create_html(self, tabname, *, empty=False):
  466. """Generates an HTML string for the current pane.
  467. The generated HTML uses `extra-networks-pane.html` as a template.
  468. Args:
  469. tabname: The name of the active tab.
  470. empty: create an empty HTML page with no items
  471. Returns:
  472. HTML formatted string.
  473. """
  474. self.lister.reset()
  475. self.metadata = {}
  476. items_list = [] if empty else self.list_items()
  477. self.items = {x["name"]: x for x in items_list}
  478. # Populate the instance metadata for each item.
  479. for item in self.items.values():
  480. metadata = item.get("metadata")
  481. if metadata:
  482. self.metadata[item["name"]] = metadata
  483. if "user_metadata" not in item:
  484. self.read_user_metadata(item)
  485. show_tree = shared.opts.extra_networks_tree_view_default_enabled
  486. page_params = {
  487. "tabname": tabname,
  488. "extra_networks_tabname": self.extra_networks_tabname,
  489. "data_sortdir": shared.opts.extra_networks_card_order,
  490. "sort_path_active": ' extra-network-control--enabled' if shared.opts.extra_networks_card_order_field == 'Path' else '',
  491. "sort_name_active": ' extra-network-control--enabled' if shared.opts.extra_networks_card_order_field == 'Name' else '',
  492. "sort_date_created_active": ' extra-network-control--enabled' if shared.opts.extra_networks_card_order_field == 'Date Created' else '',
  493. "sort_date_modified_active": ' extra-network-control--enabled' if shared.opts.extra_networks_card_order_field == 'Date Modified' else '',
  494. "tree_view_btn_extra_class": "extra-network-control--enabled" if show_tree else "",
  495. "items_html": self.create_card_view_html(tabname, none_message="Loading..." if empty else None),
  496. "extra_networks_tree_view_default_width": shared.opts.extra_networks_tree_view_default_width,
  497. "tree_view_div_default_display_class": "" if show_tree else "extra-network-dirs-hidden",
  498. }
  499. if shared.opts.extra_networks_tree_view_style == "Tree":
  500. pane_content = self.pane_content_tree_tpl.format(**page_params, tree_html=self.create_tree_view_html(tabname))
  501. else:
  502. pane_content = self.pane_content_dirs_tpl.format(**page_params, dirs_html=self.create_dirs_view_html(tabname))
  503. return self.pane_tpl.format(**page_params, pane_content=pane_content)
  504. def create_item(self, name, index=None):
  505. raise NotImplementedError()
  506. def list_items(self):
  507. raise NotImplementedError()
  508. def allowed_directories_for_previews(self):
  509. return []
  510. def get_sort_keys(self, path):
  511. """
  512. List of default keys used for sorting in the UI.
  513. """
  514. pth = Path(path)
  515. mtime, ctime = self.lister.mctime(path)
  516. return {
  517. "date_created": int(mtime),
  518. "date_modified": int(ctime),
  519. "name": pth.name.lower(),
  520. "path": str(pth).lower(),
  521. }
  522. def find_preview(self, path):
  523. """
  524. Find a preview PNG for a given path (without extension) and call link_preview on it.
  525. """
  526. potential_files = sum([[f"{path}.{ext}", f"{path}.preview.{ext}"] for ext in allowed_preview_extensions()], [])
  527. for file in potential_files:
  528. if self.lister.exists(file):
  529. return self.link_preview(file)
  530. return None
  531. def find_embedded_preview(self, path, name, metadata):
  532. """
  533. Find if embedded preview exists in safetensors metadata and return endpoint for it.
  534. """
  535. file = f"{path}.safetensors"
  536. if self.lister.exists(file) and 'ssmd_cover_images' in metadata and len(list(filter(None, json.loads(metadata['ssmd_cover_images'])))) > 0:
  537. return f"./sd_extra_networks/cover-images?page={self.extra_networks_tabname}&item={name}"
  538. return None
  539. def find_description(self, path):
  540. """
  541. Find and read a description file for a given path (without extension).
  542. """
  543. for file in [f"{path}.txt", f"{path}.description.txt"]:
  544. if not self.lister.exists(file):
  545. continue
  546. try:
  547. with open(file, "r", encoding="utf-8", errors="replace") as f:
  548. return f.read()
  549. except OSError:
  550. pass
  551. return None
  552. def create_user_metadata_editor(self, ui, tabname):
  553. return ui_extra_networks_user_metadata.UserMetadataEditor(ui, tabname, self)
  554. def initialize():
  555. extra_pages.clear()
  556. def register_default_pages():
  557. from modules.ui_extra_networks_textual_inversion import ExtraNetworksPageTextualInversion
  558. from modules.ui_extra_networks_hypernets import ExtraNetworksPageHypernetworks
  559. from modules.ui_extra_networks_checkpoints import ExtraNetworksPageCheckpoints
  560. register_page(ExtraNetworksPageTextualInversion())
  561. register_page(ExtraNetworksPageHypernetworks())
  562. register_page(ExtraNetworksPageCheckpoints())
  563. class ExtraNetworksUi:
  564. def __init__(self):
  565. self.pages = None
  566. """gradio HTML components related to extra networks' pages"""
  567. self.page_contents = None
  568. """HTML content of the above; empty initially, filled when extra pages have to be shown"""
  569. self.stored_extra_pages = None
  570. self.button_save_preview = None
  571. self.preview_target_filename = None
  572. self.tabname = None
  573. def pages_in_preferred_order(pages):
  574. tab_order = [x.lower().strip() for x in shared.opts.ui_extra_networks_tab_reorder.split(",")]
  575. def tab_name_score(name):
  576. name = name.lower()
  577. for i, possible_match in enumerate(tab_order):
  578. if possible_match in name:
  579. return i
  580. return len(pages)
  581. tab_scores = {page.name: (tab_name_score(page.name), original_index) for original_index, page in enumerate(pages)}
  582. return sorted(pages, key=lambda x: tab_scores[x.name])
  583. def create_ui(interface: gr.Blocks, unrelated_tabs, tabname):
  584. ui = ExtraNetworksUi()
  585. ui.pages = []
  586. ui.pages_contents = []
  587. ui.user_metadata_editors = []
  588. ui.stored_extra_pages = pages_in_preferred_order(extra_pages.copy())
  589. ui.tabname = tabname
  590. related_tabs = []
  591. for page in ui.stored_extra_pages:
  592. with gr.Tab(page.title, elem_id=f"{tabname}_{page.extra_networks_tabname}", elem_classes=["extra-page"]) as tab:
  593. with gr.Column(elem_id=f"{tabname}_{page.extra_networks_tabname}_prompts", elem_classes=["extra-page-prompts"]):
  594. pass
  595. elem_id = f"{tabname}_{page.extra_networks_tabname}_cards_html"
  596. page_elem = gr.HTML(page.create_html(tabname, empty=True), elem_id=elem_id)
  597. ui.pages.append(page_elem)
  598. editor = page.create_user_metadata_editor(ui, tabname)
  599. editor.create_ui()
  600. ui.user_metadata_editors.append(editor)
  601. related_tabs.append(tab)
  602. ui.button_save_preview = gr.Button('Save preview', elem_id=f"{tabname}_save_preview", visible=False)
  603. ui.preview_target_filename = gr.Textbox('Preview save filename', elem_id=f"{tabname}_preview_filename", visible=False)
  604. for tab in unrelated_tabs:
  605. tab.select(fn=None, _js=f"function(){{extraNetworksUnrelatedTabSelected('{tabname}');}}", inputs=[], outputs=[], show_progress=False)
  606. for page, tab in zip(ui.stored_extra_pages, related_tabs):
  607. jscode = (
  608. "function(){{"
  609. f"extraNetworksTabSelected('{tabname}', '{tabname}_{page.extra_networks_tabname}_prompts', {str(page.allow_prompt).lower()}, {str(page.allow_negative_prompt).lower()}, '{tabname}_{page.extra_networks_tabname}');"
  610. f"applyExtraNetworkFilter('{tabname}_{page.extra_networks_tabname}');"
  611. "}}"
  612. )
  613. tab.select(fn=None, _js=jscode, inputs=[], outputs=[], show_progress=False)
  614. def refresh():
  615. for pg in ui.stored_extra_pages:
  616. pg.refresh()
  617. create_html()
  618. return ui.pages_contents
  619. button_refresh = gr.Button("Refresh", elem_id=f"{tabname}_{page.extra_networks_tabname}_extra_refresh_internal", visible=False)
  620. button_refresh.click(fn=refresh, inputs=[], outputs=ui.pages).then(fn=lambda: None, _js="function(){ " + f"applyExtraNetworkFilter('{tabname}_{page.extra_networks_tabname}');" + " }").then(fn=lambda: None, _js='setupAllResizeHandles')
  621. def create_html():
  622. ui.pages_contents = [pg.create_html(ui.tabname) for pg in ui.stored_extra_pages]
  623. def pages_html():
  624. if not ui.pages_contents:
  625. create_html()
  626. return ui.pages_contents
  627. interface.load(fn=pages_html, inputs=[], outputs=ui.pages).then(fn=lambda: None, _js='setupAllResizeHandles')
  628. return ui
  629. def path_is_parent(parent_path, child_path):
  630. parent_path = os.path.abspath(parent_path)
  631. child_path = os.path.abspath(child_path)
  632. return child_path.startswith(parent_path)
  633. def setup_ui(ui, gallery):
  634. def save_preview(index, images, filename):
  635. # this function is here for backwards compatibility and likely will be removed soon
  636. if len(images) == 0:
  637. print("There is no image in gallery to save as a preview.")
  638. return [page.create_html(ui.tabname) for page in ui.stored_extra_pages]
  639. index = int(index)
  640. index = 0 if index < 0 else index
  641. index = len(images) - 1 if index >= len(images) else index
  642. img_info = images[index if index >= 0 else 0]
  643. image = image_from_url_text(img_info)
  644. geninfo, items = read_info_from_image(image)
  645. is_allowed = False
  646. for extra_page in ui.stored_extra_pages:
  647. if any(path_is_parent(x, filename) for x in extra_page.allowed_directories_for_previews()):
  648. is_allowed = True
  649. break
  650. assert is_allowed, f'writing to {filename} is not allowed'
  651. save_image_with_geninfo(image, geninfo, filename)
  652. return [page.create_html(ui.tabname) for page in ui.stored_extra_pages]
  653. ui.button_save_preview.click(
  654. fn=save_preview,
  655. _js="function(x, y, z){return [selected_gallery_index(), y, z]}",
  656. inputs=[ui.preview_target_filename, gallery, ui.preview_target_filename],
  657. outputs=[*ui.pages]
  658. )
  659. for editor in ui.user_metadata_editors:
  660. editor.setup_ui(gallery)