Deduplicate extra network preview-search code

This commit is contained in:
Aarni Koskela 2023-03-06 21:00:34 +02:00
parent 0cc0ee1bcb
commit 49b1dc5e07
5 changed files with 14 additions and 34 deletions

View File

@ -15,18 +15,10 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage):
def list_items(self): def list_items(self):
for name, lora_on_disk in lora.available_loras.items(): for name, lora_on_disk in lora.available_loras.items():
path, ext = os.path.splitext(lora_on_disk.filename) path, ext = os.path.splitext(lora_on_disk.filename)
previews = [path + ".png", path + ".preview.png"]
preview = None
for file in previews:
if os.path.isfile(file):
preview = self.link_preview(file)
break
yield { yield {
"name": name, "name": name,
"filename": path, "filename": path,
"preview": preview, "preview": self._find_preview(path),
"search_term": self.search_terms_from_path(lora_on_disk.filename), "search_term": self.search_terms_from_path(lora_on_disk.filename),
"prompt": json.dumps(f"<lora:{name}:") + " + opts.extra_networks_default_multiplier + " + json.dumps(">"), "prompt": json.dumps(f"<lora:{name}:") + " + opts.extra_networks_default_multiplier + " + json.dumps(">"),
"local_preview": path + ".png", "local_preview": path + ".png",

View File

@ -2,6 +2,7 @@ import glob
import os.path import os.path
import urllib.parse import urllib.parse
from pathlib import Path from pathlib import Path
from typing import Optional
from modules import shared from modules import shared
import gradio as gr import gradio as gr
@ -137,6 +138,15 @@ class ExtraNetworksPage:
return self.card_page.format(**args) return self.card_page.format(**args)
def _find_preview(self, path: str) -> Optional[str]:
"""
Find a preview PNG for a given path (without extension) and call link_preview on it.
"""
for file in [path + ".png", path + ".preview.png"]:
if os.path.isfile(file):
return self.link_preview(file)
return None
def intialize(): def intialize():
extra_pages.clear() extra_pages.clear()

View File

@ -1,7 +1,6 @@
import html import html
import json import json
import os import os
import urllib.parse
from modules import shared, ui_extra_networks, sd_models from modules import shared, ui_extra_networks, sd_models
@ -17,18 +16,10 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
checkpoint: sd_models.CheckpointInfo checkpoint: sd_models.CheckpointInfo
for name, checkpoint in sd_models.checkpoints_list.items(): for name, checkpoint in sd_models.checkpoints_list.items():
path, ext = os.path.splitext(checkpoint.filename) path, ext = os.path.splitext(checkpoint.filename)
previews = [path + ".png", path + ".preview.png"]
preview = None
for file in previews:
if os.path.isfile(file):
preview = self.link_preview(file)
break
yield { yield {
"name": checkpoint.name_for_extra, "name": checkpoint.name_for_extra,
"filename": path, "filename": path,
"preview": preview, "preview": self._find_preview(path),
"search_term": self.search_terms_from_path(checkpoint.filename) + " " + (checkpoint.sha256 or ""), "search_term": self.search_terms_from_path(checkpoint.filename) + " " + (checkpoint.sha256 or ""),
"onclick": '"' + html.escape(f"""return selectCheckpoint({json.dumps(name)})""") + '"', "onclick": '"' + html.escape(f"""return selectCheckpoint({json.dumps(name)})""") + '"',
"local_preview": path + ".png", "local_preview": path + ".png",

View File

@ -14,18 +14,11 @@ class ExtraNetworksPageHypernetworks(ui_extra_networks.ExtraNetworksPage):
def list_items(self): def list_items(self):
for name, path in shared.hypernetworks.items(): for name, path in shared.hypernetworks.items():
path, ext = os.path.splitext(path) path, ext = os.path.splitext(path)
previews = [path + ".png", path + ".preview.png"]
preview = None
for file in previews:
if os.path.isfile(file):
preview = self.link_preview(file)
break
yield { yield {
"name": name, "name": name,
"filename": path, "filename": path,
"preview": preview, "preview": self._find_preview(path),
"search_term": self.search_terms_from_path(path), "search_term": self.search_terms_from_path(path),
"prompt": json.dumps(f"<hypernet:{name}:") + " + opts.extra_networks_default_multiplier + " + json.dumps(">"), "prompt": json.dumps(f"<hypernet:{name}:") + " + opts.extra_networks_default_multiplier + " + json.dumps(">"),
"local_preview": path + ".png", "local_preview": path + ".png",

View File

@ -15,16 +15,10 @@ class ExtraNetworksPageTextualInversion(ui_extra_networks.ExtraNetworksPage):
def list_items(self): def list_items(self):
for embedding in sd_hijack.model_hijack.embedding_db.word_embeddings.values(): for embedding in sd_hijack.model_hijack.embedding_db.word_embeddings.values():
path, ext = os.path.splitext(embedding.filename) path, ext = os.path.splitext(embedding.filename)
preview_file = path + ".preview.png"
preview = None
if os.path.isfile(preview_file):
preview = self.link_preview(preview_file)
yield { yield {
"name": embedding.name, "name": embedding.name,
"filename": embedding.filename, "filename": embedding.filename,
"preview": preview, "preview": self._find_preview(path),
"search_term": self.search_terms_from_path(embedding.filename), "search_term": self.search_terms_from_path(embedding.filename),
"prompt": json.dumps(embedding.name), "prompt": json.dumps(embedding.name),
"local_preview": path + ".preview.png", "local_preview": path + ".preview.png",