Merge pull request #14528 from AUTOMATIC1111/mass-file-lister

mass file lister as an attempt to tackle #14507

3 files changed, 86 insertions, 9 deletions

diff --git a/modules/extra_networks.py b/modules/extra_networks.py
index b9533677..04249dff 100644
--- a/modules/extra_networks.py
+++ b/modules/extra_networks.py
@@ -206,7 +206,7 @@ def parse_prompts(prompts):
     return res, extra_data

 

 

-def get_user_metadata(filename):

+def get_user_metadata(filename, lister=None):

     if filename is None:

         return {}

 

@@ -215,7 +215,8 @@ def get_user_metadata(filename):
 

     metadata = {}

     try:

-        if os.path.isfile(metadata_filename):

+        exists = lister.exists(metadata_filename) if lister else os.path.exists(metadata_filename)

+        if exists:

             with open(metadata_filename, "r", encoding="utf8") as file:

                 metadata = json.load(file)

     except Exception as e:

diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py
index e1c679ec..62db36f5 100644
--- a/modules/ui_extra_networks.py
+++ b/modules/ui_extra_networks.py
@@ -3,7 +3,7 @@ import os.path
 import urllib.parse

 from pathlib import Path

 

-from modules import shared, ui_extra_networks_user_metadata, errors, extra_networks

+from modules import shared, ui_extra_networks_user_metadata, errors, extra_networks, util

 from modules.images import read_info_from_image, save_image_with_geninfo

 import gradio as gr

 import json

@@ -107,13 +107,14 @@ class ExtraNetworksPage:
         self.allow_negative_prompt = False

         self.metadata = {}

         self.items = {}

+        self.lister = util.MassFileLister()

 

     def refresh(self):

         pass

 

     def read_user_metadata(self, item):

         filename = item.get("filename", None)

-        metadata = extra_networks.get_user_metadata(filename)

+        metadata = extra_networks.get_user_metadata(filename, lister=self.lister)

 

         desc = metadata.get("description", None)

         if desc is not None:

@@ -123,7 +124,7 @@ class ExtraNetworksPage:
 

     def link_preview(self, filename):

         quoted_filename = urllib.parse.quote(filename.replace('\\', '/'))

-        mtime = os.path.getmtime(filename)

+        mtime, _ = self.lister.mctime(filename)

         return f"./sd_extra_networks/thumb?filename={quoted_filename}&mtime={mtime}"

 

     def search_terms_from_path(self, filename, possible_directories=None):

@@ -137,6 +138,8 @@ class ExtraNetworksPage:
         return ""

 

     def create_html(self, tabname):

+        self.lister.reset()

+

         items_html = ''

 

         self.metadata = {}

@@ -282,10 +285,10 @@ class ExtraNetworksPage:
         List of default keys used for sorting in the UI.

         """

         pth = Path(path)

-        stat = pth.stat()

+        mtime, ctime = self.lister.mctime(path)

         return {

-            "date_created": int(stat.st_ctime or 0),

-            "date_modified": int(stat.st_mtime or 0),

+            "date_created": int(mtime),

+            "date_modified": int(ctime),

             "name": pth.name.lower(),

             "path": str(pth.parent).lower(),

         }

@@ -298,7 +301,7 @@ class ExtraNetworksPage:
         potential_files = sum([[path + "." + ext, path + ".preview." + ext] for ext in allowed_preview_extensions()], [])

 

         for file in potential_files:

-            if os.path.isfile(file):

+            if self.lister.exists(file):

                 return self.link_preview(file)

 

         return None

@@ -308,6 +311,9 @@ class ExtraNetworksPage:
         Find and read a description file for a given path (without extension).

         """

         for file in [f"{path}.txt", f"{path}.description.txt"]:

+            if not self.lister.exists(file):

+                continue

+

             try:

                 with open(file, "r", encoding="utf-8", errors="replace") as f:

                     return f.read()

diff --git a/modules/util.py b/modules/util.py
index 4861bcb0..c2a27590 100644
--- a/modules/util.py
+++ b/modules/util.py
@@ -66,3 +66,73 @@ def truncate_path(target_path, base_path=cwd):
     except ValueError:

         pass

     return abs_target

+

+

+class MassFileListerCachedDir:

+    """A class that caches file metadata for a specific directory."""

+

+    def __init__(self, dirname):

+        self.files = None

+        self.files_cased = None

+        self.dirname = dirname

+

+        stats = ((x.name, x.stat(follow_symlinks=False)) for x in os.scandir(self.dirname))

+        files = [(n, s.st_mtime, s.st_ctime) for n, s in stats]

+        self.files = {x[0].lower(): x for x in files}

+        self.files_cased = {x[0]: x for x in files}

+

+

+class MassFileLister:

+    """A class that provides a way to check for the existence and mtime/ctile of files without doing more than one stat call per file."""

+

+    def __init__(self):

+        self.cached_dirs = {}

+

+    def find(self, path):

+        """

+        Find the metadata for a file at the given path.

+

+        Returns:

+            tuple or None: A tuple of (name, mtime, ctime) if the file exists, or None if it does not.

+        """

+

+        dirname, filename = os.path.split(path)

+

+        cached_dir = self.cached_dirs.get(dirname)

+        if cached_dir is None:

+            cached_dir = MassFileListerCachedDir(dirname)

+            self.cached_dirs[dirname] = cached_dir

+

+        stats = cached_dir.files_cased.get(filename)

+        if stats is not None:

+            return stats

+

+        stats = cached_dir.files.get(filename.lower())

+        if stats is None:

+            return None

+

+        try:

+            os_stats = os.stat(path, follow_symlinks=False)

+            return filename, os_stats.st_mtime, os_stats.st_ctime

+        except Exception:

+            return None

+

+    def exists(self, path):

+        """Check if a file exists at the given path."""

+

+        return self.find(path) is not None

+

+    def mctime(self, path):

+        """

+        Get the modification and creation times for a file at the given path.

+

+        Returns:

+            tuple: A tuple of (mtime, ctime) if the file exists, or (0, 0) if it does not.

+        """

+

+        stats = self.find(path)

+        return (0, 0) if stats is None else stats[1:3]

+

+    def reset(self):

+        """Clear the cache of all directories."""

+        self.cached_dirs.clear()