use ThreadPoolExecutor for extra networks

extensions-builtin/Lora/ui_extra_networks_lora.py

@@ -1,5 +1,6 @@
 import os
 import json
+import concurrent
 import network
 import networks
 from modules import shared, ui_extra_networks
@@ -8,6 +9,7 @@ from modules import shared, ui_extra_networks
 class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage):
     def __init__(self):
         super().__init__('Lora')
+        self.list_time = 0
 
     def refresh(self):
         networks.list_available_networks()
@@ -74,10 +76,12 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage):
             return None
 
     def list_items(self):
-        for _index, name in enumerate(networks.available_networks):
-            item = self.create_item(name)
-            if item is not None:
-                yield item
+        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
+            future_items = {executor.submit(self.create_item, net): net for net in networks.available_networks}
+            for future in concurrent.futures.as_completed(future_items):
+                item = future.result()
+                if item is not None:
+                    yield item
 
     def allowed_directories_for_previews(self):
         return [shared.cmd_opts.lora_dir, shared.cmd_opts.lyco_dir]

modules/modelloader.py

@@ -334,23 +334,23 @@ def load_reference(name: str):
         return True
 
 
-modelloader_directories = {}
+cache_folders = {}
 cache_last = 0
 cache_time = 1
 
 
-def directory_has_changed(dir:str, *, recursive:bool=True) -> bool: # pylint: disable=redefined-builtin
+def directory_updated(path:str, *, recursive:bool=True) -> bool: # pylint: disable=redefined-builtin
     try:
-        dir = os.path.abspath(dir)
-        if dir not in modelloader_directories:
+        path = os.path.abspath(path)
+        if path not in cache_folders:
             return True
         if cache_last > (time.time() - cache_time):
             return False
-        if not (os.path.exists(dir) and os.path.isdir(dir) and os.path.getmtime(dir) == modelloader_directories[dir][0]):
+        if not (os.path.exists(path) and os.path.isdir(path) and os.path.getmtime(path) == cache_folders[path][0]):
             return True
         if recursive:
-            for _dir in modelloader_directories:
-                if _dir.startswith(dir) and _dir != dir and not (os.path.exists(_dir) and os.path.isdir(_dir) and os.path.getmtime(_dir) == modelloader_directories[_dir][0]):
+            for folder in cache_folders:
+                if folder.startswith(path) and folder != path and not (os.path.exists(folder) and os.path.isdir(folder) and os.path.getmtime(folder) == cache_folders[folder][0]):
                     return True
     except Exception as e:
         shared.log.error(f"Filesystem Error: {e.__class__.__name__}({e})")
@@ -358,44 +358,48 @@ def directory_has_changed(dir:str, *, recursive:bool=True) -> bool: # pylint: di
     return False
 
 
-def directory_directories(dir:str, *, recursive:bool=True) -> dict[str,tuple[float,list[str]]]: # pylint: disable=redefined-builtin
-    dir = os.path.abspath(dir)
-    if directory_has_changed(dir, recursive=recursive):
-        for _dir in list(modelloader_directories):
-            if os.path.exists(_dir) or os.path.isdir(_dir):
+def directory_list(path:str, *, recursive:bool=True) -> dict[str,tuple[float,list[str]]]: # pylint: disable=redefined-builtin
+    path = os.path.abspath(path)
+    res = {}
+    if not os.path.exists(path):
+        return res
+    if directory_updated(path, recursive=recursive):
+        for folder in list(cache_folders):
+            del cache_folders[folder]
+            if os.path.exists(folder) or os.path.isdir(folder):
+                continue
+        for folder, files in walk(path, lambda e, path: shared.log.debug(f"FS walk error: {e} {path}")):
+            if not os.path.exists(folder):
                 continue
-            del modelloader_directories[_dir]
-        for _dir, _files in walk(dir, lambda e, path: shared.log.debug(f"FS walk error: {e} {path}")):
             try:
-                mtime = os.path.getmtime(_dir)
-                if _dir not in modelloader_directories or mtime != modelloader_directories[_dir][0]:
-                    modelloader_directories[_dir] = (mtime, [os.path.join(_dir, fn) for fn in _files])
+                mtime = os.path.getmtime(folder)
+                if folder not in cache_folders or mtime != cache_folders[folder][0]:
+                    cache_folders[folder] = (mtime, [os.path.join(folder, fn) for fn in files])
             except Exception as e:
                 shared.log.error(f"Filesystem Error: {e.__class__.__name__}({e})")
-                del modelloader_directories[_dir]
-    res = {}
-    for _dir in modelloader_directories:
-        if _dir == dir or (recursive and _dir.startswith(dir)):
-            res[_dir] = modelloader_directories[_dir]
+                del cache_folders[folder]
+    for folder in cache_folders:
+        if folder == path or (recursive and folder.startswith(path)):
+            res[folder] = cache_folders[folder]
             if not recursive:
                 break
     return res
 
 
-def directory_mtime(dir:str, *, recursive:bool=True) -> float: # pylint: disable=redefined-builtin
-    return float(max(0, *[mtime for mtime, _ in directory_directories(dir, recursive=recursive).values()]))
+def directory_mtime(path:str, *, recursive:bool=True) -> float: # pylint: disable=redefined-builtin
+    return float(max(0, *[mtime for mtime, _ in directory_list(path, recursive=recursive).values()]))
 
 
 def directories_file_paths(directories:dict) -> list[str]:
     return sum([dat[1] for dat in directories.values()],[])
 
 
-def unique_directories(directories:list[str], *, recursive:bool=True) -> list[str]:
+def directories_unique(directories:list[str], *, recursive:bool=True) -> list[str]:
     '''Ensure no empty, or duplicates'''
-    directories = { os.path.abspath(dir): True for dir in directories if dir }.keys()
+    directories = { os.path.abspath(path): True for path in directories if path }.keys()
     if recursive:
         '''If we are going recursive, then directories that are children of other directories are redundant'''
-        directories = [dir for dir in directories if not any(_dir != dir and dir.startswith(os.path.join(_dir,'')) for _dir in directories)]
+        directories = [path for path in directories if not any(d != path and path.startswith(os.path.join(d,'')) for d in directories)]
     return directories
 
 
@@ -404,7 +408,7 @@ def unique_paths(paths:list[str]) -> list[str]:
 
 
 def directory_files(*directories:list[str], recursive:bool=True) -> list[str]:
-    return unique_paths(sum([[*directories_file_paths(directory_directories(dir, recursive=recursive))] for dir in unique_directories(directories, recursive=recursive)],[]))
+    return unique_paths(sum([[*directories_file_paths(directory_list(d, recursive=recursive))] for d in directories_unique(directories, recursive=recursive)],[]))
 
 
 def extension_filter(ext_filter=None, ext_blacklist=None):
@@ -485,7 +489,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
     @param ext_filter: An optional list of filename extensions to filter by
     @return: A list of paths containing the desired model(s)
     """
-    places = unique_directories([model_path, command_path])
+    places = directories_unique([model_path, command_path])
     output = []
     try:
         output:list = [*filter(extension_filter(ext_filter, ext_blacklist), directory_files(*places))]

modules/ui_extra_networks.py

@@ -222,6 +222,8 @@ class ExtraNetworksPage:
             self.items = []
             shared.log.error(f'Extra networks error listing items: class={self.__class__.__name__} tab={tabname} {e}')
         for item in self.items:
+            if item is None:
+                continue
             self.metadata[item["name"]] = item.get("metadata", {})
         t1 = time.time()
         debug(f'EN create-items: page={self.name} items={len(self.items)} time={t1-t0:.2f}')
@@ -237,7 +239,7 @@ class ExtraNetworksPage:
             return f"<div id='{tabname}_{self_name_id}_subdirs' class='extra-network-subdirs'></div><div id='{tabname}_{self_name_id}_cards' class='extra-network-cards'>Extra network page not ready<br>Click refresh to try again</div>"
         subdirs = {}
         allowed_folders = [os.path.abspath(x) for x in self.allowed_directories_for_previews()]
-        for parentdir, dirs in {d: modelloader.directory_directories(d) for d in allowed_folders}.items():
+        for parentdir, dirs in {d: modelloader.directory_list(d) for d in allowed_folders}.items():
             for tgt in dirs.keys():
                 if shared.backend == shared.Backend.DIFFUSERS:
                     if os.path.join(paths.models_path, 'Reference') in tgt:
@@ -539,9 +541,11 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
             refresh_time = time.time()
         threads = []
         for page in get_pages():
-            # page.create_items(ui.tabname)
-            threads.append(threading.Thread(target=page.create_items, args=[ui.tabname]))
-            threads[-1].start()
+            if os.environ.get('SD_EN_DEBUG', None) is not None:
+                threads.append(threading.Thread(target=page.create_items, args=[ui.tabname]))
+                threads[-1].start()
+            else:
+                page.create_items(ui.tabname)
         for thread in threads:
             thread.join()
         for page in get_pages():

modules/ui_extra_networks_checkpoints.py

@@ -1,8 +1,10 @@
+import os
 import html
 import json
-import os
+import concurrent
 from modules import shared, ui_extra_networks, sd_models, paths
 
+
 reference_dir = os.path.join(paths.models_path, 'Reference')
 
 class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
@@ -36,31 +38,38 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
                 "description": v.get('desc', ''),
             }
 
+    def create_item(self, name):
+        record = None
+        try:
+            checkpoint: sd_models.CheckpointInfo = sd_models.checkpoints_list.get(name)
+            exists = os.path.exists(checkpoint.filename)
+            record = {
+                "type": 'Model',
+                "name": checkpoint.name,
+                "title": checkpoint.title,
+                "filename": checkpoint.filename,
+                "hash": checkpoint.shorthash,
+                "search_term": self.search_terms_from_path(checkpoint.title),
+                "preview": self.find_preview(checkpoint.filename),
+                "local_preview": f"{os.path.splitext(checkpoint.filename)[0]}.{shared.opts.samples_format}",
+                "metadata": checkpoint.metadata,
+                "onclick": '"' + html.escape(f"""return selectCheckpoint({json.dumps(name)})""") + '"',
+                "mtime": os.path.getmtime(checkpoint.filename) if exists else 0,
+                "size": os.path.getsize(checkpoint.filename) if exists else 0,
+            }
+            record["info"] = self.find_info(checkpoint.filename)
+            record["description"] = self.find_description(checkpoint.filename, record["info"])
+        except Exception as e:
+            shared.log.debug(f"Extra networks error: type=model file={name} {e}")
+        return record
+
     def list_items(self):
-        checkpoint: sd_models.CheckpointInfo
-        checkpoints = sd_models.checkpoints_list.copy()
-        for name, checkpoint in checkpoints.items():
-            try:
-                exists = os.path.exists(checkpoint.filename)
-                record = {
-                    "type": 'Model',
-                    "name": checkpoint.name,
-                    "title": checkpoint.title,
-                    "filename": checkpoint.filename,
-                    "hash": checkpoint.shorthash,
-                    "search_term": self.search_terms_from_path(checkpoint.title),
-                    "preview": self.find_preview(checkpoint.filename),
-                    "local_preview": f"{os.path.splitext(checkpoint.filename)[0]}.{shared.opts.samples_format}",
-                    "metadata": checkpoint.metadata,
-                    "onclick": '"' + html.escape(f"""return selectCheckpoint({json.dumps(name)})""") + '"',
-                    "mtime": os.path.getmtime(checkpoint.filename) if exists else 0,
-                    "size": os.path.getsize(checkpoint.filename) if exists else 0,
-                }
-                record["info"] = self.find_info(checkpoint.filename)
-                record["description"] = self.find_description(checkpoint.filename, record["info"])
-                yield record
-            except Exception as e:
-                shared.log.debug(f"Extra networks error: type=model file={name} {e}")
+        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
+            future_items = {executor.submit(self.create_item, cp): cp for cp in list(sd_models.checkpoints_list.copy())}
+            for future in concurrent.futures.as_completed(future_items):
+                item = future.result()
+                if item is not None:
+                    yield item
         for record in self.list_reference():
             yield record
 

modules/ui_extra_networks_styles.py

@@ -1,6 +1,7 @@
 import os
 import html
 import json
+import concurrent
 from modules import shared, extra_networks, ui_extra_networks, styles
 
 
@@ -62,35 +63,44 @@ class ExtraNetworksPageStyles(ui_extra_networks.ExtraNetworksPage):
         }
         return item
 
-    def list_items(self):
-        for k, style in shared.prompt_styles.styles.items():
-            try:
-                fn = os.path.splitext(getattr(style, 'filename', ''))[0]
-                name = getattr(style, 'name', '')
-                if name == '':
-                    continue
-                txt = f'Prompt: {getattr(style, "prompt", "")}'
-                if len(getattr(style, 'negative_prompt', '')) > 0:
-                    txt += f'\nNegative: {style.negative_prompt}'
-                yield {
-                    "type": 'Style',
-                    "name": name,
-                    "title": k,
-                    "filename": style.filename,
-                    "search_term": f'{txt} {self.search_terms_from_path(name)}',
-                    "preview": style.preview if getattr(style, 'preview', None) is not None and style.preview.startswith('data:') else self.find_preview(fn),
-                    "description": style.description if getattr(style, 'description', None) is not None and len(style.description) > 0 else txt,
-                    "prompt": getattr(style, 'prompt', ''),
-                    "negative": getattr(style, 'negative_prompt', ''),
-                    "extra": getattr(style, 'extra', ''),
-                    "local_preview": f"{fn}.{shared.opts.samples_format}",
-                    "onclick": '"' + html.escape(f"""return selectStyle({json.dumps(name)})""") + '"',
-                    "mtime": getattr(style, 'mtime', 0),
-                    "size": os.path.getsize(style.filename),
-                }
-            except Exception as e:
-                shared.log.debug(f"Extra networks error: type=style file={k} {e}")
+    def create_item(self, k):
+        item = None
+        try:
+            style = shared.prompt_styles.styles.get(k)
+            fn = os.path.splitext(getattr(style, 'filename', ''))[0]
+            name = getattr(style, 'name', '')
+            if name == '':
+                return item
+            txt = f'Prompt: {getattr(style, "prompt", "")}'
+            if len(getattr(style, 'negative_prompt', '')) > 0:
+                txt += f'\nNegative: {style.negative_prompt}'
+            item = {
+                "type": 'Style',
+                "name": name,
+                "title": k,
+                "filename": style.filename,
+                "search_term": f'{txt} {self.search_terms_from_path(name)}',
+                "preview": style.preview if getattr(style, 'preview', None) is not None and style.preview.startswith('data:') else self.find_preview(fn),
+                "description": style.description if getattr(style, 'description', None) is not None and len(style.description) > 0 else txt,
+                "prompt": getattr(style, 'prompt', ''),
+                "negative": getattr(style, 'negative_prompt', ''),
+                "extra": getattr(style, 'extra', ''),
+                "local_preview": f"{fn}.{shared.opts.samples_format}",
+                "onclick": '"' + html.escape(f"""return selectStyle({json.dumps(name)})""") + '"',
+                "mtime": getattr(style, 'mtime', 0),
+                "size": os.path.getsize(style.filename),
+            }
+        except Exception as e:
+            shared.log.debug(f"Extra networks error: type=style file={k} {e}")
+        return item
 
+    def list_items(self):
+        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
+            future_items = {executor.submit(self.create_item, style): style for style in list(shared.prompt_styles.styles)}
+            for future in concurrent.futures.as_completed(future_items):
+                item = future.result()
+                if item is not None:
+                    yield item
 
     def allowed_directories_for_previews(self):
         return [v for v in [shared.opts.styles_dir] if v is not None] + ['html']

modules/ui_extra_networks_textual_inversion.py

@@ -1,5 +1,6 @@
 import json
 import os
+import concurrent
 from modules import shared, sd_hijack, sd_models, ui_extra_networks
 from modules.textual_inversion.textual_inversion import Embedding
 
@@ -8,6 +9,7 @@ class ExtraNetworksPageTextualInversion(ui_extra_networks.ExtraNetworksPage):
     def __init__(self):
         super().__init__('Embedding')
         self.allow_negative_prompt = True
+        self.embeddings = []
 
     def refresh(self):
         if sd_models.model_data.sd_model is None:
@@ -17,51 +19,61 @@ class ExtraNetworksPageTextualInversion(ui_extra_networks.ExtraNetworksPage):
         elif hasattr(sd_models.model_data.sd_model, 'embedding_db'):
             sd_models.model_data.sd_model.embedding_db.load_textual_inversion_embeddings(force_reload=True)
 
+    def create_item(self, embedding: Embedding):
+        record = None
+        try:
+            path, _ext = os.path.splitext(embedding.filename)
+            tags = {}
+            if embedding.tag is not None:
+                tags[embedding.tag]=1
+            name = os.path.splitext(embedding.basename)[0]
+            record = {
+                "type": 'Embedding',
+                "name": name,
+                "filename": embedding.filename,
+                "preview": self.find_preview(embedding.filename),
+                "search_term": self.search_terms_from_path(name),
+                "prompt": json.dumps(f" {os.path.splitext(embedding.name)[0]}"),
+                "local_preview": f"{path}.{shared.opts.samples_format}",
+                "tags": tags,
+                "mtime": os.path.getmtime(embedding.filename),
+                "size": os.path.getsize(embedding.filename),
+            }
+            record["info"] = self.find_info(embedding.filename)
+            record["description"] = self.find_description(embedding.filename, record["info"])
+        except Exception as e:
+            shared.log.debug(f"Extra networks error: type=embedding file={embedding.filename} {e}")
+        return record
+
     def list_items(self):
+
         def list_folder(folder):
             for filename in os.listdir(folder):
                 fn = os.path.join(folder, filename)
                 if os.path.isfile(fn) and (fn.lower().endswith(".pt") or fn.lower().endswith(".safetensors")):
                     embedding = Embedding(vec=0, name=os.path.basename(fn), filename=fn)
                     embedding.filename = fn
-                    embeddings.append(embedding)
+                    self.embeddings.append(embedding)
                 elif os.path.isdir(fn) and not fn.startswith('.'):
                     list_folder(fn)
 
         if sd_models.model_data.sd_model is None:
-            embeddings = []
+            self.embeddings = []
             list_folder(shared.opts.embeddings_dir)
         elif shared.backend == shared.Backend.ORIGINAL:
-            embeddings = list(sd_hijack.model_hijack.embedding_db.word_embeddings.values())
+            self.embeddings = list(sd_hijack.model_hijack.embedding_db.word_embeddings.values())
         elif hasattr(sd_models.model_data.sd_model, 'embedding_db'):
-            embeddings = list(sd_models.model_data.sd_model.embedding_db.word_embeddings.values())
+            self.embeddings = list(sd_models.model_data.sd_model.embedding_db.word_embeddings.values())
         else:
-            embeddings = []
-        embeddings = sorted(embeddings, key=lambda emb: emb.filename)
-        for embedding in embeddings:
-            try:
-                path, _ext = os.path.splitext(embedding.filename)
-                tags = {}
-                if embedding.tag is not None:
-                    tags[embedding.tag]=1
-                name = os.path.splitext(embedding.basename)[0]
-                record = {
-                    "type": 'Embedding',
-                    "name": name,
-                    "filename": embedding.filename,
-                    "preview": self.find_preview(embedding.filename),
-                    "search_term": self.search_terms_from_path(name),
-                    "prompt": json.dumps(f" {os.path.splitext(embedding.name)[0]}"),
-                    "local_preview": f"{path}.{shared.opts.samples_format}",
-                    "tags": tags,
-                    "mtime": os.path.getmtime(embedding.filename),
-                    "size": os.path.getsize(embedding.filename),
-                }
-                record["info"] = self.find_info(embedding.filename)
-                record["description"] = self.find_description(embedding.filename, record["info"])
-                yield record
-            except Exception as e:
-                shared.log.debug(f"Extra networks error: type=embedding file={embedding.filename} {e}")
+            self.embeddings = []
+        self.embeddings = sorted(self.embeddings, key=lambda emb: emb.filename)
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
+            future_items = {executor.submit(self.create_item, net): net for net in self.embeddings}
+            for future in concurrent.futures.as_completed(future_items):
+                item = future.result()
+                if item is not None:
+                    yield item
 
     def allowed_directories_for_previews(self):
         return list(sd_hijack.model_hijack.embedding_db.embedding_dirs)