Merge pull request #684 from zanllp/feature/cache-init-and-env-config

Add support for pre-generating cache via startup parameters and speci…

.env.example

@@ -14,6 +14,11 @@ IIB_SERVER_LANG=auto
 # This configuration parameter specifies the maximum number of database file backups for the IIB .
 IIB_DB_FILE_BACKUP_MAX=8
 
+# Set the cache directory for IIB, including image cache and video cover cache.
+# The default is the system's temporary directory, but if you want to specify a custom directory, set it here.
+# You can use --generate_video_cover and --generate_image_cache to pre-generate the cache.
+# IIB_CACHE_DIR=
+
 
 # ---------------------------- ACCESS_CONTROL ----------------------------
 

README-zh.md

@@ -20,6 +20,8 @@ https://github.com/zanllp/sd-webui-infinite-image-browsing/assets/25872019/807b8
 - 存在缓存的情况下后，图像可以在几毫秒内显示。
 - 默认使用缩略图显示图像，默认大小为512像素，您可以在全局设置页中调整缩略图分辨率。
 - 你还可以控制网格图像的宽度，允许以64px到1024px的宽度范围进行显示
+- 支持通过`--generate_video_cover`和`--generate_image_cache`来预先生成缩略图和视频封面，以提高性能。
+- 支持通过`IIB_CACHE_DIR`环境变量来指定缓存目录。
 
 ### 🔍 图像搜索和收藏
 - 将会把Prompt、Model、Lora等信息转成标签，将根据使用频率排序以供进行精确的搜索。

README.md

@@ -31,6 +31,8 @@ You can add your own parser through [parsers](https://github.com/zanllp/sd-webui
 - Once caching is generated, images can be displayed in just a few milliseconds.
 - Images are displayed with thumbnails by default, with a default size of 512 pixels. You can adjust the thumbnail resolution on the global settings page.
 - You can also control the width of the grid images, allowing them to be displayed in widths ranging from 64px to 1024px.
+- Supports pre-generating thumbnails and video covers to improve performance using `--generate_video_cover` and `--generate_image_cache`.
+- Supports specifying the cache directory through the `IIB_CACHE_DIR` environment variable.
 
 ### 🔍 Image Search & Favorite
 - The prompt, model, Lora, and other information will be converted into tags and sorted by frequency of use for precise searching.

app.py

@@ -22,6 +22,16 @@
 default_port = 8000
 default_host = "127.0.0.1"
 
+def get_all_img_dirs(sd_webui_config: str, relative_to_config: bool):
+    dirs = get_valid_img_dirs(
+        get_sd_webui_conf(
+            sd_webui_config=sd_webui_config,
+            sd_webui_path_relative_to_config=relative_to_config,
+        )
+    )
+    dirs += list(map(lambda x: x.path, ExtraPath.get_extra_paths(DataBase.get_conn())))
+    return dirs
+
 
 def sd_webui_paths_check(sd_webui_config: str, relative_to_config: bool):
     conf = {}
@@ -50,12 +60,7 @@ def paths_check(paths):
 
 
 def do_update_image_index(sd_webui_config: str, relative_to_config=False):
-    dirs = get_valid_img_dirs(
-        get_sd_webui_conf(
-            sd_webui_config=sd_webui_config,
-            sd_webui_path_relative_to_config=relative_to_config,
-        )
-    )
+    dirs = get_all_img_dirs(sd_webui_config, relative_to_config)
     if not len(dirs):
         return print(f"{tag} no valid image directories, skipped")
     conn = DataBase.get_conn()
@@ -185,6 +190,22 @@ def setup_parser() -> argparse.ArgumentParser:
         action="store_true",
         help="Pre-generate video cover images to speed up browsing.",
     )
+    parser.add_argument(
+        "--generate_image_cache",
+        action="store_true",
+        help="Pre-generate image cache to speed up browsing.",
+    )
+    parser.add_argument(
+        "--generate_image_cache_size",
+        type=str,
+        default="512x512",
+        help="The size of the image cache to generate. Default is 512x512",
+    )
+    parser.add_argument(
+        "--gen_cache_verbose",
+        action="store_true",
+        help="Verbose mode for cache generation.",
+    )
     parser.add_argument(
         "--extra_paths",
         nargs="+",
@@ -258,10 +279,21 @@ async def async_launch_app(
     args_dict = vars(args)
 
     if args_dict.get("generate_video_cover"):
-        from scripts.iib.video_cover import generate_video_covers
+        from scripts.iib.video_cover_gen import generate_video_covers
 
         conn = DataBase.get_conn()
-        generate_video_covers(map(lambda x: x.path, ExtraPath.get_extra_paths(conn)))
+        generate_video_covers(
+            dirs = map(lambda x: x.path, ExtraPath.get_extra_paths(conn)),
+            verbose=args.gen_cache_verbose,
+        )
+        exit(0)
+    if args_dict.get("generate_image_cache"):
+        from scripts.iib.img_cache_gen import generate_image_cache
+        generate_image_cache(
+            dirs = get_all_img_dirs(args.sd_webui_config, args.sd_webui_path_relative_to_config), 
+            size = args.generate_image_cache_size,
+            verbose = args.gen_cache_verbose
+        )
         exit(0)
 
     launch_app(**vars(args))

scripts/iib/api.py

@@ -12,7 +12,7 @@
     human_readable_size,
     is_valid_media_path,
     is_media_file,
-    temp_path,
+    get_cache_dir,
     get_formatted_date,
     is_win,
     cwd,
@@ -115,6 +115,8 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
     backup_db_file(DataBase.get_db_file_path())
     api_base = kwargs.get("base") if isinstance(kwargs.get("base"), str) else DEFAULT_BASE
     fe_public_path = kwargs.get("fe_public_path") if isinstance(kwargs.get("fe_public_path"), str) else api_base
+    cache_base_dir = get_cache_dir()
+
     # print(f"IIB api_base:{api_base} fe_public_path:{fe_public_path}")
     if IIB_DEBUG or is_nuitka:
         @app.exception_handler(Exception)
@@ -507,12 +509,12 @@ async def get_target_folder_files(folder_path: str):
     @app.get(api_base + "/image-thumbnail", dependencies=[Depends(verify_secret)])
     async def thumbnail(path: str, t: str, size: str = "256x256"):
         check_path_trust(path)
-        if not temp_path:
+        if not cache_base_dir:
             return
         # 生成缓存文件的路径
         hash_dir = hashlib.md5((path + t).encode("utf-8")).hexdigest()
         hash = hash_dir + size
-        cache_dir = os.path.join(temp_path, "iib_cache", hash_dir)
+        cache_dir = os.path.join(cache_base_dir, "iib_cache", hash_dir)
         cache_path = os.path.join(cache_dir, f"{size}.webp")
 
         # 如果缓存文件存在，则直接返回该文件
@@ -592,7 +594,7 @@ async def stream_video(path: str, request: Request):
     @app.get(api_base + "/video_cover", dependencies=[Depends(verify_secret)])
     async def video_cover(path: str, t: str):        
         check_path_trust(path)
-        if not temp_path:
+        if not cache_base_dir:
             return
 
         if not os.path.exists(path):
@@ -602,7 +604,7 @@ async def video_cover(path: str, t: str):
         # 生成缓存文件的路径
         hash_dir = hashlib.md5((path + t).encode("utf-8")).hexdigest()
         hash = hash_dir
-        cache_dir = os.path.join(temp_path, "iib_cache", "video_cover", hash_dir)
+        cache_dir = os.path.join(cache_base_dir, "iib_cache", "video_cover", hash_dir)
         cache_path = os.path.join(cache_dir, "cover.webp")
 
         # 如果缓存文件存在，则直接返回该文件

scripts/iib/img_cache_gen.py

@@ -0,0 +1,58 @@
+import hashlib
+import os
+from typing import List
+from scripts.iib.tool import get_formatted_date, get_cache_dir, is_image_file
+from concurrent.futures import ThreadPoolExecutor
+import time
+from PIL import Image
+
+def generate_image_cache(dirs, size:str, verbose=False):
+  start_time = time.time()
+  cache_base_dir = get_cache_dir()
+
+  def process_image(item):
+    if item.is_dir():
+      verbose and print(f"Processing directory: {item.path}")
+      for sub_item in os.scandir(item.path):
+        process_image(sub_item)
+      return
+    if not os.path.exists(item.path) or not is_image_file(item.path):
+      return
+
+    try:
+      path = os.path.normpath(item.path)
+      stat = item.stat()
+      t = get_formatted_date(stat.st_mtime)
+      hash_dir = hashlib.md5((path + t).encode("utf-8")).hexdigest()
+      cache_dir = os.path.join(cache_base_dir, "iib_cache", hash_dir)
+      cache_path = os.path.join(cache_dir, f"{size}.webp")
+
+      if os.path.exists(cache_path):
+          verbose and print(f"Image cache already exists: {path}")
+          return
+
+      if os.path.getsize(path) < 64 * 1024:
+          verbose and print(f"Image size less than 64KB: {path}", "skip")
+          return
+
+      with Image.open(path) as img:
+          w, h = size.split("x")
+          img.thumbnail((int(w), int(h)))
+          os.makedirs(cache_dir, exist_ok=True)
+          img.save(cache_path, "webp")
+
+      verbose and print(f"Image cache generated: {path}")
+    except Exception as e:
+      print(f"Error generating image cache: {path}")
+      print(e)
+
+  with ThreadPoolExecutor() as executor:
+    for dir_path in dirs:
+      folder_listing: List[os.DirEntry] = os.scandir(dir_path)
+      for item in folder_listing:
+        executor.submit(process_image, item)
+
+  print("Image cache generation completed. ✨")
+  end_time = time.time()
+  execution_time = end_time - start_time
+  print(f"Execution time: {execution_time} seconds")

scripts/iib/tool.py

@@ -280,7 +280,11 @@ def get_temp_path():
     return temp_path
 
 
-temp_path = get_temp_path()
+_temp_path = get_temp_path()
+
+
+def get_cache_dir():
+    return os.getenv("IIB_CACHE_DIR") or _temp_path
 
 def get_secret_key_required():
     try:

scripts/iib/video_cover.py → scripts/iib/video_cover_gen.py

@@ -1,19 +1,20 @@
 import hashlib
 import os
 from typing import List
-from scripts.iib.tool import get_formatted_date, get_temp_path, is_video_file
+from scripts.iib.tool import get_formatted_date, get_cache_dir, is_video_file
 from concurrent.futures import ThreadPoolExecutor
 import time
 
 
-def generate_video_covers(dirs):
+def generate_video_covers(dirs,verbose=False):
   start_time = time.time()
   import imageio.v3 as iio
-  temp_path = get_temp_path()
+
+  cache_base_dir = get_cache_dir()
 
   def process_video(item):
     if item.is_dir():
-      print(f"Processing directory: {item.path}")
+      verbose and print(f"Processing directory: {item.path}")
       for sub_item in os.scandir(item.path):
         process_video(sub_item)
       return
@@ -25,7 +26,7 @@ def process_video(item):
       stat = item.stat()
       t = get_formatted_date(stat.st_mtime)
       hash_dir = hashlib.md5((path + t).encode("utf-8")).hexdigest()
-      cache_dir = os.path.join(temp_path, "iib_cache", "video_cover", hash_dir)
+      cache_dir = os.path.join(cache_base_dir, "iib_cache", "video_cover", hash_dir)
       cache_path = os.path.join(cache_dir, "cover.webp")
 
       # 如果缓存文件存在，则直接返回该文件
@@ -41,7 +42,7 @@ def process_video(item):
 
       os.makedirs(cache_dir, exist_ok=True)
       iio.imwrite(cache_path, frame, extension=".webp")
-      print(f"Video cover generated: {path}")
+      verbose and print(f"Video cover generated: {path}")
     except Exception as e:
       print(f"Error generating video cover: {path}")
       print(e)