Deduplicate extra network preview-search code

5 files changed, 14 insertions, 34 deletions

diff --git a/extensions-builtin/Lora/ui_extra_networks_lora.py b/extensions-builtin/Lora/ui_extra_networks_lora.py
index 22cabcb0..4c1549d7 100644
--- a/extensions-builtin/Lora/ui_extra_networks_lora.py
+++ b/extensions-builtin/Lora/ui_extra_networks_lora.py
@@ -15,18 +15,10 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage):
     def list_items(self):

         for name, lora_on_disk in lora.available_loras.items():

             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 {

                 "name": name,

                 "filename": path,

-                "preview": preview,

+                "preview": self._find_preview(path),

                 "search_term": self.search_terms_from_path(lora_on_disk.filename),

                 "prompt": json.dumps(f"<lora:{name}:") + " + opts.extra_networks_default_multiplier + " + json.dumps(">"),

                 "local_preview": path + ".png",

diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py
index 71f1d81f..1a10a5df 100644
--- a/modules/ui_extra_networks.py
+++ b/modules/ui_extra_networks.py
@@ -2,6 +2,7 @@ import glob
 import os.path

 import urllib.parse

 from pathlib import Path

+from typing import Optional

 

 from modules import shared

 import gradio as gr

@@ -137,6 +138,15 @@ class ExtraNetworksPage:
 

         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():

     extra_pages.clear()

diff --git a/modules/ui_extra_networks_checkpoints.py b/modules/ui_extra_networks_checkpoints.py
index 04097a79..b712d12b 100644
--- a/modules/ui_extra_networks_checkpoints.py
+++ b/modules/ui_extra_networks_checkpoints.py
@@ -1,7 +1,6 @@
 import html

 import json

 import os

-import urllib.parse

 

 from modules import shared, ui_extra_networks, sd_models

 

@@ -17,18 +16,10 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
         checkpoint: sd_models.CheckpointInfo

         for name, checkpoint in sd_models.checkpoints_list.items():

             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 {

                 "name": checkpoint.name_for_extra,

                 "filename": path,

-                "preview": preview,

+                "preview": self._find_preview(path),

                 "search_term": self.search_terms_from_path(checkpoint.filename) + " " + (checkpoint.sha256 or ""),

                 "onclick": '"' + html.escape(f"""return selectCheckpoint({json.dumps(name)})""") + '"',

                 "local_preview": path + ".png",

diff --git a/modules/ui_extra_networks_hypernets.py b/modules/ui_extra_networks_hypernets.py
index 57851088..89f33242 100644
--- a/modules/ui_extra_networks_hypernets.py
+++ b/modules/ui_extra_networks_hypernets.py
@@ -14,18 +14,11 @@ class ExtraNetworksPageHypernetworks(ui_extra_networks.ExtraNetworksPage):
     def list_items(self):

         for name, path in shared.hypernetworks.items():

             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 {

                 "name": name,

                 "filename": path,

-                "preview": preview,

+                "preview": self._find_preview(path),

                 "search_term": self.search_terms_from_path(path),

                 "prompt": json.dumps(f"<hypernet:{name}:") + " + opts.extra_networks_default_multiplier + " + json.dumps(">"),

                 "local_preview": path + ".png",

diff --git a/modules/ui_extra_networks_textual_inversion.py b/modules/ui_extra_networks_textual_inversion.py
index bb64eb81..f7057390 100644
--- a/modules/ui_extra_networks_textual_inversion.py
+++ b/modules/ui_extra_networks_textual_inversion.py
@@ -15,16 +15,10 @@ class ExtraNetworksPageTextualInversion(ui_extra_networks.ExtraNetworksPage):
     def list_items(self):

         for embedding in sd_hijack.model_hijack.embedding_db.word_embeddings.values():

             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 {

                 "name": embedding.name,

                 "filename": embedding.filename,

-                "preview": preview,

+                "preview": self._find_preview(path),

                 "search_term": self.search_terms_from_path(embedding.filename),

                 "prompt": json.dumps(embedding.name),

                 "local_preview": path + ".preview.png",