From 7e4aed79489bc74016eecbc603c392a4d266f266 Mon Sep 17 00:00:00 2001 From: Midcoastal Date: Tue, 15 Aug 2023 16:07:41 -0400 Subject: [PATCH 1/8] Add a FS directory/path cacher - Added FS dir/path mechanism to 'modules.modelloader' - Refactored 'modules.modelloader.load_models' to use the cache - Refactored 'modules.ui_extra_networks.find_*' methods to use the cache - Added progress indicator to 'modules.ui_extra_networks.create_html' The cache, as implimented, will always ensure it is up-to-date (per 'directory_has_changed') and significantly improves loading speed (when used with 'model_loader' and the 'find_*' methods) with large model directories. Overall load speed tested with ~10k models (mix of checkpoints, loras, lycoris, and embeddings) and ~40k secndary files (images, descriptions, CivitAI Info, etc). Loading speeds went from ~1 hour and 45 minutes to ~5 minutes. Confounding variable to loading speeds: this is over a fiber-attached storage device. Connection is 10gbe full-duplex, remote source has an NVMe raid cache. Saturation of the network is the norm, but laintency is a factor. Regardless, small-scale local-storage testing also shows measurable improvements, so this should be a welcome addition. --- modules/modelloader.py | 71 +++++++++++++++++++++----- modules/ui_extra_networks.py | 96 ++++++++++++++++++++++-------------- 2 files changed, 118 insertions(+), 49 deletions(-) diff --git a/modules/modelloader.py b/modules/modelloader.py index 6527e9efd..dbba84339 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -137,6 +137,62 @@ def find_diffuser(name: str): return models[0].modelId return None +modelloader_directories = {} + +def directory_has_changed(dir:str, *, recursive:bool=True) -> bool: + dir = os.path.abspath(dir) + if dir not in modelloader_directories: + return True + 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 + 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: + if not (os.path.exists(_dir) and os.path.isdir(_dir)): + del modelloader_directories[_dir] + for _dir, _subdirs, _files in os.walk(dir, topdown=False, followlinks=True): + 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]) + 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 directories_file_paths(directories:dict) -> list[str]: + return sum([[fp for fp in dat[1]] for dat in directories.values()], []) + +def filter_paths(paths:list[str], *, filter:callable=None) -> list[str]: + return [fp for fp in paths if not (os.path.islink(fp) and not os.path.exists(fp)) and filter(fp)] + +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(_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([[fp for fp in 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): + def filter(fp:str): + return (not ext_filter or any(fp.endswith(ew) for ew in ext_filter)) and (not ext_blacklist or not any(fp.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 +205,10 @@ 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]) 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_paths(directory_files(*places), filter=extension_filter(ext_filter, ext_blacklist)) 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 diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index f5c8f2a93..4e31f38f1 100644 --- a/modules/ui_extra_networks.py +++ b/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() @@ -158,19 +161,17 @@ class ExtraNetworksPage: return f"
Extra network page not ready
Click refresh to try again
" 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) + 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=f'Initializing Items') + items = self.items + progress.update(task, total=len(items)) + __t = None + __i = 0 + for item in items: + if __t is None: + __t = time() + __i += 1 + 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) + progress.update(task, advance=1, description=f"{round(__i/(shared.time.time()-__t))} item/s") if len(subdirs_html) > 0 or len(self.html) > 0: res = f"
{subdirs_html}
{self.html}
" 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"
Extra network error
{e}
" 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"] + dir = os.path.dirname(path) + paths = modelloader.directory_directories(dir, recursive=False) for file in sum([[f'{path}.thumb.{ext}'] for ext in preview_extensions], []): # use thumbnail if exists - if os.path.isfile(file): + if file in paths[dir][1]: 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): + if file in paths[dir][1]: 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 From 05850c23441aab0cecaa0d80cfa86f226a7330ff Mon Sep 17 00:00:00 2001 From: Midcoastal Date: Tue, 15 Aug 2023 23:06:23 -0400 Subject: [PATCH 2/8] Upgrade Lora/TI model listers to use cache --- extensions-builtin/Lora/lora.py | 6 ++-- .../textual_inversion/textual_inversion.py | 28 ++++++++++--------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/extensions-builtin/Lora/lora.py b/extensions-builtin/Lora/lora.py index 0822f3c10..08d4d6f18 100644 --- a/extensions-builtin/Lora/lora.py +++ b/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 filter_paths, 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_paths(directory_files(shared.cmd_opts.lora_dir), filter=extension_filter(['.PT', '.CKPT', '.SAFETENSORS'])), key=str.lower): name = os.path.splitext(os.path.basename(filename))[0] entry = LoraOnDisk(name, filename) diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index 25345a35a..845cc1ce3 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/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 filter_paths, 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_paths(directory_files(embdir.path), filter=lambda fp: is_ext(fp) and is_not_preview(fp)): + 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: From 6eae768fb06da8eb9dc8d0227ac1f7151b26d844 Mon Sep 17 00:00:00 2001 From: Midcoastal Date: Tue, 15 Aug 2023 23:08:26 -0400 Subject: [PATCH 3/8] Gather HTML parts as list to append in the end Somehow, significant increase in speed (size of text append was slowing large loads down. --- modules/ui_extra_networks.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index 4e31f38f1..75d3a7e23 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -190,6 +190,7 @@ class ExtraNetworksPage: self.items = [] shared.log.error(f'Extra networks error listing items: {self.__class__}') self.create_xyz_grid() + htmls = [] with Progress( SpinnerColumn(), TextColumn('[cyan]Creating Extra Network '+self.title+' HTML - {task.description}'), @@ -199,16 +200,15 @@ class ExtraNetworksPage: task = progress.add_task(description=f'Initializing Items') items = self.items progress.update(task, total=len(items)) - __t = None + __t = time() __i = 0 for item in items: - if __t is None: - __t = time() __i += 1 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.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"
{subdirs_html}
{self.html}
" else: From 873e864640be7e3af516c5870fc559280c77a1b7 Mon Sep 17 00:00:00 2001 From: Midcoastal Date: Tue, 15 Aug 2023 23:09:24 -0400 Subject: [PATCH 4/8] Ooof, modelloader changes needed for TI upgrade --- modules/modelloader.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/modules/modelloader.py b/modules/modelloader.py index dbba84339..19229701e 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -169,6 +169,9 @@ def directory_directories(dir:str, *, recursive:bool=True) -> dict[str,tuple[flo 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([[fp for fp in dat[1]] for dat in directories.values()], []) @@ -190,8 +193,12 @@ def directory_files(*directories:list[str], recursive:bool=True) -> list[str]: return unique_paths(sum([[fp for fp in 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 = [ext.upper() for ext in ext_filter] + if ext_blacklist: + ext_blacklist = [ext.upper() for ext in ext_blacklist] def filter(fp:str): - return (not ext_filter or any(fp.endswith(ew) for ew in ext_filter)) and (not ext_blacklist or not any(fp.endswith(ew) for ew in ext_blacklist)) + 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: From c9bdd0344d6c7d4b2350edae79cdba6f72575d18 Mon Sep 17 00:00:00 2001 From: Midcoastal Date: Tue, 15 Aug 2023 23:10:30 -0400 Subject: [PATCH 5/8] Allow JPEG image files JPEG is allowed and searched for when looking for previews/thumbs, and therefore should be allowed to view. --- modules/ui_extra_networks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index 75d3a7e23..20c2a2111 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -34,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"}) From 03363ce866fa4dcbfe9b05338319708e762d0f5a Mon Sep 17 00:00:00 2001 From: Midcoastal Date: Tue, 15 Aug 2023 23:20:00 -0400 Subject: [PATCH 6/8] Resolve Lint Errors: C416 (x2), W291 and F541 --- modules/modelloader.py | 4 ++-- modules/ui_extra_networks.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/modules/modelloader.py b/modules/modelloader.py index 19229701e..45e41315f 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -173,7 +173,7 @@ 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([[fp for fp in dat[1]] for dat in directories.values()], []) + return sum([fp for fp in dat[1] for dat in directories.values()], []) def filter_paths(paths:list[str], *, filter:callable=None) -> list[str]: return [fp for fp in paths if not (os.path.islink(fp) and not os.path.exists(fp)) and filter(fp)] @@ -190,7 +190,7 @@ 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([[fp for fp in directories_file_paths(directory_directories(dir, recursive=recursive))] for dir in unique_directories(directories, recursive=recursive)],[])) + return unique_paths(sum([fp for fp in 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: diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index 20c2a2111..4240da4e4 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -193,11 +193,11 @@ class ExtraNetworksPage: htmls = [] with Progress( SpinnerColumn(), - TextColumn('[cyan]Creating Extra Network '+self.title+' HTML - {task.description}'), + 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=f'Initializing Items') + task = progress.add_task(description='Initializing Items') items = self.items progress.update(task, total=len(items)) __t = time() From 89da66207ab2d77078d5eef4588c6aea58491c1c Mon Sep 17 00:00:00 2001 From: Midcoastal Date: Wed, 16 Aug 2023 00:25:27 -0400 Subject: [PATCH 7/8] I guess I didn't quite understand Lint error C416 Seems like a rather pedantic check, if you ask me... Oh well. --- modules/modelloader.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/modelloader.py b/modules/modelloader.py index 45e41315f..2406196b9 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -173,7 +173,7 @@ 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([fp for fp in dat[1] for dat in directories.values()], []) + return sum(list([fp for fp in dat[1]] for dat in directories.values()), []) def filter_paths(paths:list[str], *, filter:callable=None) -> list[str]: return [fp for fp in paths if not (os.path.islink(fp) and not os.path.exists(fp)) and filter(fp)] @@ -190,7 +190,7 @@ 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([fp for fp in directories_file_paths(directory_directories(dir, recursive=recursive)) for dir in unique_directories(directories, recursive=recursive)],[])) + return unique_paths(sum(list([fp for fp in 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: From 67f369ed25a42ceac84b215cb362a6df94056109 Mon Sep 17 00:00:00 2001 From: Midcoastal Date: Thu, 17 Aug 2023 16:49:42 -0400 Subject: [PATCH 8/8] A walk() optimization and lint fixes --- extensions-builtin/Lora/lora.py | 4 +- modules/modelloader.py | 110 ++++++++++++++---- .../textual_inversion/textual_inversion.py | 6 +- modules/ui_extra_networks.py | 10 +- 4 files changed, 99 insertions(+), 31 deletions(-) diff --git a/extensions-builtin/Lora/lora.py b/extensions-builtin/Lora/lora.py index 08d4d6f18..607dd8e33 100644 --- a/extensions-builtin/Lora/lora.py +++ b/extensions-builtin/Lora/lora.py @@ -3,7 +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 filter_paths, directory_files, extension_filter +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} @@ -445,7 +445,7 @@ def list_available_loras(): os.makedirs(shared.cmd_opts.lora_dir, exist_ok=True) - for filename in sorted(filter_paths(directory_files(shared.cmd_opts.lora_dir), filter=extension_filter(['.PT', '.CKPT', '.SAFETENSORS'])), key=str.lower): + 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) diff --git a/modules/modelloader.py b/modules/modelloader.py index 2406196b9..c37917dd8 100644 --- a/modules/modelloader.py +++ b/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) @@ -138,29 +188,47 @@ def find_diffuser(name: str): return None modelloader_directories = {} +cache_last = 0 +cache_time = 1 def directory_has_changed(dir:str, *, recursive:bool=True) -> bool: - dir = os.path.abspath(dir) - if dir not in modelloader_directories: + 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 - 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 + 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: - if not (os.path.exists(_dir) and os.path.isdir(_dir)): + 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] - for _dir, _subdirs, _files in os.walk(dir, topdown=False, followlinks=True): - 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]) + directory_directories = {} for _dir in modelloader_directories: if _dir == dir or (recursive and _dir.startswith(dir)): @@ -173,30 +241,27 @@ 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(list([fp for fp in dat[1]] for dat in directories.values()), []) - -def filter_paths(paths:list[str], *, filter:callable=None) -> list[str]: - return [fp for fp in paths if not (os.path.islink(fp) and not os.path.exists(fp)) and filter(fp)] + 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(_dir) for _dir in directories)] + 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(list([fp for fp in directories_file_paths(directory_directories(dir, recursive=recursive))] for dir in unique_directories(directories, recursive=recursive)),[])) + 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 = [ext.upper() for ext in ext_filter] + ext_filter = [*map(str.upper, ext_filter)] if ext_blacklist: - ext_blacklist = [ext.upper() for ext in 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 @@ -213,9 +278,10 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None @return: A list of paths containing the desired model(s) """ places = unique_directories([model_path, command_path]) + #shared.log.debug(f"{inspect.currentframe().f_code.co_name}: {', '.join(places)}") output = [] try: - output:list = filter_paths(directory_files(*places), filter=extension_filter(ext_filter, ext_blacklist)) + 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 diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index 845cc1ce3..379a178d1 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/modules/textual_inversion/textual_inversion.py @@ -13,7 +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 filter_paths, directory_files, extension_filter, directory_mtime +from modules.modelloader import directory_files, extension_filter, directory_mtime TextualInversionTemplate = namedtuple("TextualInversionTemplate", ["name", "path"]) textual_inversion_templates = {} @@ -215,11 +215,11 @@ class EmbeddingDatabase: def load_from_dir(self, embdir): if not os.path.isdir(embdir.path): return - + 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_paths(directory_files(embdir.path), filter=lambda fp: is_ext(fp) and is_not_preview(fp)): + 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 diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index 4240da4e4..7e42e453d 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -264,11 +264,11 @@ class ExtraNetworksPage: preview_extensions = ["jpg", "jpeg", "png", "webp", "tiff", "jp2"] dir = os.path.dirname(path) paths = modelloader.directory_directories(dir, recursive=False) - for file in sum([[f'{path}.thumb.{ext}'] for ext in preview_extensions], []): # use thumbnail if exists - if file in paths[dir][1]: + 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 file in paths[dir][1]: + 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') @@ -363,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") @@ -378,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 = ''