mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 17:24:32 +02:00
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.
This commit is contained in:
Midcoastal
+58
-13
@@ -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
|
||||
|
||||
@@ -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"<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)
|
||||
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"<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"]
|
||||
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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user