Merge pull request #2014 from midcoastal/Issue1563/FS-Path-Cache

Add a FS directory/path cacher

extensions-builtin/Lora/lora.py

@@ -3,6 +3,7 @@ import re
 from typing import Union
 import torch
 from modules import shared, devices, sd_models, errors, scripts, sd_hijack, hashes
+from modules.modelloader import directory_files, extension_filter
 
 metadata_tags_order = {"ss_sd_model_name": 1, "ss_resolution": 2, "ss_clip_skip": 3, "ss_num_train_images": 10, "ss_tag_frequency": 20}
 
@@ -444,10 +445,7 @@ def list_available_loras():
 
     os.makedirs(shared.cmd_opts.lora_dir, exist_ok=True)
 
-    candidates = list(shared.walk_files(shared.cmd_opts.lora_dir, allowed_extensions=[".pt", ".ckpt", ".safetensors"]))
-    for filename in sorted(candidates, key=str.lower):
-        if os.path.isdir(filename):
-            continue
+    for filename in sorted([*filter(extension_filter(['.PT', '.CKPT', '.SAFETENSORS']), directory_files(shared.cmd_opts.lora_dir))], key=str.lower):
 
         name = os.path.splitext(os.path.basename(filename))[0]
         entry = LoraOnDisk(name, filename)

modules/modelloader.py

@@ -6,9 +6,59 @@ from urllib.parse import urlparse
 from modules import shared
 from modules.upscaler import Upscaler, UpscalerLanczos, UpscalerNearest, UpscalerNone
 from modules.paths import script_path, models_path
+import inspect
+import time
 
 diffuser_repos = []
 
+def walk(top, onerror:callable=None):
+    # A near-exact copy of `os.path.walk()`, trimmed slightly.  
+    # Probably not nessesary for most people's collections, but makes 
+    # a difference on really large datasets.
+    nondirs = []
+    walk_dirs = []
+
+    try:
+        scandir_it = os.scandir(top)
+    except OSError as error:
+        if onerror is not None:
+            onerror(error, top)
+        return
+
+    with scandir_it:
+        while True:
+            try:
+                try:
+                    entry = next(scandir_it)
+                except StopIteration:
+                    break
+            except OSError as error:
+                if onerror is not None:
+                    onerror(error, top)
+                return
+
+            try:
+                is_dir = entry.is_dir()
+            except OSError:
+                is_dir = False
+
+            if not is_dir:
+                nondirs.append(entry.name)
+            else:
+                try:
+                    if entry.is_symlink() and not os.path.exists(entry.path):
+                        raise NotADirectoryError('Broken Symlink')
+                    walk_dirs.append(entry.path)
+                except OSError as error:
+                    if onerror is not None:
+                        onerror(error, entry.path)
+
+    # Recurse into sub-directories
+    for new_path in walk_dirs:
+        yield from walk(new_path, onerror)
+    # Yield after recursion if going bottom up
+    yield top, nondirs
+
 
 def download_civit_model(model_url: str, model_name: str, model_path: str, preview):
     model_file = os.path.join(shared.opts.ckpt_dir, model_path, model_name)
@@ -137,6 +187,84 @@ def find_diffuser(name: str):
         return models[0].modelId
     return None
 
+modelloader_directories = {}
+cache_last = 0
+cache_time = 1
+
+def directory_has_changed(dir:str, *, recursive:bool=True) -> bool:
+    try:
+        dir = os.path.abspath(dir)
+        if dir not in modelloader_directories:
+            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]):
+            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]):
+                    return True
+    except Exception as e:
+        shared.log.error(f"Filesystem Error: {e.__class__.__name__}({e})")
+        return True
+
+    return False
+
+def directory_directories(dir:str, *, recursive:bool=True) -> dict[str,tuple[float,list[str]]]:
+    dir = os.path.abspath(dir)
+    if directory_has_changed(dir, recursive=recursive):
+        for _dir in modelloader_directories:
+            try:
+                if (os.path.exists(_dir) and os.path.isdir(_dir)):
+                    continue
+            except Exception:
+                pass
+            del modelloader_directories[_dir]
+        for _dir, _files in walk(dir, lambda e, path: shared.log.error(f"Filesystem Walk Error: {e.__class__.__name__}({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])
+            except Exception as e:
+                shared.log.error(f"Filesystem Error: {e.__class__.__name__}({e})")
+                del modelloader_directories[_dir]
+
+    directory_directories = {}
+    for _dir in modelloader_directories:
+        if _dir == dir or (recursive and _dir.startswith(dir)):
+            directory_directories[_dir] = modelloader_directories[_dir]
+            if not recursive:
+                break
+    return directory_directories
+
+def directory_mtime(dir:str, *, recursive:bool=True) -> float:
+    return float(max(0, *[mtime for mtime, _ in directory_directories(dir, 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]:
+    '''Ensure no empty, or duplicates'''
+    directories = { os.path.abspath(dir): True for dir in directories if dir }.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)]
+    return directories
+
+def unique_paths(paths:list[str]) -> list[str]:
+    return { fp: True for fp in paths }.keys()
+
+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)],[]))
+
+def extension_filter(ext_filter=None, ext_blacklist=None):
+    if ext_filter:
+        ext_filter = [*map(str.upper, ext_filter)]
+    if ext_blacklist:
+        ext_blacklist = [*map(str.upper, ext_blacklist)]
+    def filter(fp:str):
+        return (not ext_filter or any(fp.upper().endswith(ew) for ew in ext_filter)) and (not ext_blacklist or not any(fp.upper().endswith(ew) for ew in ext_blacklist))
+    return filter
 
 def load_models(model_path: str, model_url: str = None, command_path: str = None, ext_filter=None, download_name=None, ext_blacklist=None) -> list:
     """
@@ -149,21 +277,11 @@ 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 = []
-    places.append(model_path)
-    if command_path is not None and command_path != model_path and os.path.isdir(command_path):
-        places.append(command_path)
+    places = unique_directories([model_path, command_path])
+    #shared.log.debug(f"{inspect.currentframe().f_code.co_name}: {', '.join(places)}")
     output = []
     try:
-        for place in places:
-            for full_path in shared.walk_files(place, allowed_extensions=ext_filter):
-                if os.path.islink(full_path) and not os.path.exists(full_path):
-                    shared.log.error(f"Skipping broken symlink: {full_path}")
-                    continue
-                if ext_blacklist is not None and any(full_path.endswith(x) for x in ext_blacklist):
-                    continue
-                if full_path not in output:
-                    output.append(full_path)
+        output:list = [*filter(extension_filter(ext_filter, ext_blacklist), directory_files(*places))]
         if model_url is not None and len(output) == 0:
             if download_name is not None:
                 from basicsr.utils.download_util import load_file_from_url

modules/textual_inversion/textual_inversion.py

@@ -13,6 +13,7 @@ import modules.textual_inversion.dataset
 from modules.textual_inversion.learn_schedule import LearnRateScheduler
 from modules.textual_inversion.image_embedding import embedding_to_b64, embedding_from_b64, insert_image_data_embed, extract_image_data_embed, caption_image_overlay
 from modules.textual_inversion.logging import save_settings_to_file
+from modules.modelloader import directory_files, extension_filter, directory_mtime
 
 TextualInversionTemplate = namedtuple("TextualInversionTemplate", ["name", "path"])
 textual_inversion_templates = {}
@@ -85,15 +86,13 @@ class DirWithTextualInversionEmbeddings:
         if not os.path.isdir(self.path):
             return False
 
-        mt = os.path.getmtime(self.path)
-        if self.mtime is None or mt > self.mtime:
-            return True
+        return directory_mtime(self.path) != self.mtime
 
     def update(self):
         if not os.path.isdir(self.path):
             return
 
-        self.mtime = os.path.getmtime(self.path)
+        self.mtime = directory_mtime(self.path)
 
 
 class EmbeddingDatabase:
@@ -216,16 +215,19 @@ class EmbeddingDatabase:
     def load_from_dir(self, embdir):
         if not os.path.isdir(embdir.path):
             return
-        for root, _dirs, fns in os.walk(embdir.path, followlinks=True):
-            for fn in fns:
-                try:
-                    fullfn = os.path.join(root, fn)
-                    if os.stat(fullfn).st_size == 0:
-                        continue
-                    self.load_from_file(fullfn, fn)
-                except Exception as e:
-                    errors.display(e, f'embedding load {fn}')
+
+        is_ext = extension_filter(['.PNG', '.WEBP', '.JXL', '.AVIF', '.BIN', '.PT', '.SAFETENSORS'])
+        is_not_preview = lambda fp: not next(iter(os.path.splitext(fp))).upper().endswith('.PREVIEW')
+
+        for file_path in [*filter(lambda fp: is_ext(fp) and is_not_preview(fp), directory_files(embdir.path))]:
+            try:
+                if os.stat(file_path).st_size == 0:
                     continue
+                fn = os.path.basename(file_path)
+                self.load_from_file(file_path, fn)
+            except Exception as e:
+                errors.display(e, f'embedding load {fn}')
+                continue
 
     def load_textual_inversion_embeddings(self, force_reload=False):
         if not force_reload:

modules/ui_extra_networks.py

@@ -8,9 +8,12 @@ from pathlib import Path
 from collections import OrderedDict
 import gradio as gr
 from PIL import Image
-from modules import shared, scripts
+from modules import shared, scripts, modelloader
 from modules.generation_parameters_copypaste import image_from_url_text
 from modules.ui_components import ToolButton
+from logging import DEBUG
+from time import time
+from rich.progress import Progress, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn, TimeElapsedColumn, SpinnerColumn
 
 extra_pages = []
 allowed_dirs = set()
@@ -31,7 +34,7 @@ def fetch_file(filename: str = ""):
         return FileResponse(filename, headers={"Accept-Ranges": "bytes"})
     if not any(Path(x).absolute() in Path(filename).absolute().parents for x in allowed_dirs):
         return JSONResponse({"error": f"File cannot be fetched: {filename}. Must be in one of directories registered by extra pages."})
-    if os.path.splitext(filename)[1].lower() not in (".png", ".jpg", ".webp"):
+    if os.path.splitext(filename)[1].lower() not in (".png", ".jpg", ".jpeg", ".webp"):
         return JSONResponse({"error": f"File cannot be fetched: {filename}. Only png and jpg and webp."})
     return FileResponse(filename, headers={"Accept-Ranges": "bytes"})
 
@@ -158,19 +161,17 @@ 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 in [*set(allowed_folders)]:
-            for root, dirs, _files in os.walk(parentdir, followlinks=True):
-                for dirname in dirs:
-                    x = os.path.join(root, dirname)
-                    if shared.opts.diffusers_dir in x:
-                        subdirs[os.path.basename(shared.opts.diffusers_dir)] = 1
-                    if (not os.path.isdir(x)) or ('models--' in x):
-                        continue
-                    subdir = os.path.abspath(x)[len(parentdir):].replace("\\", "/")
-                    while subdir.startswith("/"):
-                        subdir = subdir[1:]
-                    if not self.is_empty(x):
-                        subdirs[subdir] = 1
+        for parentdir, dirs in {dir: modelloader.directory_directories(dir) for dir in allowed_folders}.items():
+            for dir in dirs.keys():
+                if shared.opts.diffusers_dir in dir:
+                    subdirs[os.path.basename(shared.opts.diffusers_dir)] = 1
+                if 'models--' in dir:
+                    continue
+                subdir = dir[len(parentdir):].replace("\\", "/")
+                while subdir.startswith("/"):
+                    subdir = subdir[1:]
+                if not self.is_empty(dir):
+                    subdirs[subdir] = 1
         if subdirs:
             subdirs = OrderedDict(sorted(subdirs.items()))
             subdirs = {"": 1, **subdirs}
@@ -189,10 +190,25 @@ class ExtraNetworksPage:
                 self.items = []
                 shared.log.error(f'Extra networks error listing items: {self.__class__}')
             self.create_xyz_grid()
-            for item in self.items:
-                self.metadata[item["name"]] = item.get("metadata", {})
-                self.info[item["name"]] = self.find_info(item['filename'])
-                self.html += self.create_html_for_item(item, tabname)
+            htmls = []
+            with Progress(
+                SpinnerColumn(),
+                TextColumn('[cyan]Creating Extra Network '+self.title+' HTML - {task.description}'),
+                BarColumn(), TaskProgressColumn(), TextColumn('({task.completed}/{task.total})'),
+                TimeRemainingColumn(), TimeElapsedColumn(), transient=not shared.log.isEnabledFor(DEBUG), expand=True
+            ) as progress:
+                task = progress.add_task(description='Initializing Items')
+                items = self.items
+                progress.update(task, total=len(items))
+                __t = time()
+                __i = 0
+                for item in items:
+                    __i += 1
+                    self.metadata[item["name"]] = item.get("metadata", {})
+                    self.info[item["name"]] = self.find_info(item['filename'])
+                    htmls.append(self.create_html_for_item(item, tabname))
+                    progress.update(task, advance=1, description=f"{round(__i/(shared.time.time()-__t))} item/s")
+            self.html += ''.join(htmls)
             if len(subdirs_html) > 0 or len(self.html) > 0:
                 res = f"<div id='{tabname}_{self_name_id}_subdirs' class='extra-network-subdirs'>{subdirs_html}</div><div id='{tabname}_{self_name_id}_cards' class='extra-network-cards'>{self.html}</div>"
             else:
@@ -201,7 +217,7 @@ class ExtraNetworksPage:
             threading.Thread(target=self.create_thumb).start()
             return res
         except Exception as e:
-            shared.log.error(f'Extra networks page error: {e}')
+            shared.log.error(f'Extra networks {self.title} {tabname} page error: {e.__class__.__name__} -> {e}')
             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 error<br>{e}</div>"
 
     def list_items(self):
@@ -238,7 +254,7 @@ class ExtraNetworksPage:
                 args['title'] += f'\nAlias: {item["alias"]}'
             if item.get("tags", None) is not None:
                 args['title'] += f'\nTags: {", ".join(tags)}'
-            self.card.format(**args)
+            #self.card.format(**args)
             return self.card.format(**args)
         except Exception as e:
             shared.log.error(f'Extra networks item error: page={tabname} item={item["name"]} {e}')
@@ -246,36 +262,44 @@ class ExtraNetworksPage:
 
     def find_preview(self, path):
         preview_extensions = ["jpg", "jpeg", "png", "webp", "tiff", "jp2"]
-        for file in sum([[f'{path}.thumb.{ext}'] for ext in preview_extensions], []): # use thumbnail if exists
-            if os.path.isfile(file):
+        dir = os.path.dirname(path)
+        paths = modelloader.directory_directories(dir, recursive=False)
+        for file in [f'{path}.thumb.{ext}' for ext in preview_extensions]: # use thumbnail if exists
+            if file in paths[dir][1] and os.path.exists(file):
                 return self.link_preview(file)
-        for file in sum([[f'{path}.preview.{ext}', f'{path}.{ext}'] for ext in preview_extensions], []):
-            if os.path.isfile(file):
+        for file in [f'{path}{mid}{ext}' for ext in preview_extensions for mid in ['.preview.', '.']]:
+            if file in paths[dir][1] and os.path.exists(file):
                 self.missing_thumbs.append(file)
                 return self.link_preview(file)
         return self.link_preview('html/card-no-preview.png')
 
     def find_description(self, path):
+        dir = os.path.dirname(path)
+        paths = modelloader.directory_directories(dir, recursive=False)
         for file in [f"{path}.txt", f"{path}.description.txt"]:
-            try:
-                with open(file, "r", encoding="utf-8", errors="replace") as f:
-                    txt = f.read()
-                    txt = re.sub('[<>]', '', txt)
-                    return txt
-            except OSError:
-                pass
+            if file in paths[dir][1]:
+                try:
+                    with open(file, "r", encoding="utf-8", errors="replace") as f:
+                        txt = f.read()
+                        txt = re.sub('[<>]', '', txt)
+                        return txt
+                except OSError:
+                    pass
         return None
 
     def find_info(self, path):
+        dir = os.path.dirname(path)
+        paths = modelloader.directory_directories(dir, recursive=False)
         basename, _ext = os.path.splitext(path)
         for file in [f"{path}.info", f"{path}.civitai.info", f"{basename}.info", f"{basename}.civitai.info"]:
-            try:
-                with open(file, "r", encoding="utf-8", errors="replace") as f:
-                    txt = f.read()
-                    txt = re.sub('[<>]', '', txt)
-                    return txt
-            except OSError:
-                pass
+            if file in paths[dir][1]:
+                try:
+                    with open(file, "r", encoding="utf-8", errors="replace") as f:
+                        txt = f.read()
+                        txt = re.sub('[<>]', '', txt)
+                        return txt
+                except OSError:
+                    pass
         return None
 
 
@@ -339,6 +363,7 @@ def create_ui(container, button, tabname, skip_indexing = False):
         ui.description_target_filename = gr.Textbox('Description save filename', elem_id=tabname+"_description_filename", visible=False)
 
         for page in ui.stored_extra_pages:
+            shared.log.debug(f"Create UI Extra Network Page: {page.title}")
             page_html = page.create_html(ui.tabname, skip_indexing)
             with gr.Tab(page.title, id=page.title.lower().replace(" ", "_"), elem_classes="extra-networks-tab"):
                 page_elem = gr.HTML(page_html, elem_id=tabname+page.name+"_extra_page", elem_classes="extra-networks-page")
@@ -354,6 +379,7 @@ def create_ui(container, button, tabname, skip_indexing = False):
     button_close.click(fn=toggle_visibility, inputs=[state_visible], outputs=[state_visible, container])
 
     def refresh():
+        shared.log.debug("Refreshing UI Extra Networks Pages")
         res = []
         for pg in ui.stored_extra_pages:
             pg.html = ''