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