diff --git a/.env.example b/.env.example index ac8fd7a..36aae60 100644 --- a/.env.example +++ b/.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 ---------------------------- diff --git a/README-zh.md b/README-zh.md index 7b5ed1d..c8d49d9 100644 --- a/README-zh.md +++ b/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等信息转成标签,将根据使用频率排序以供进行精确的搜索。 diff --git a/README.md b/README.md index 271f28c..d6cbac2 100644 --- a/README.md +++ b/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. diff --git a/app.py b/app.py index 764e9e3..4da35f0 100644 --- a/app.py +++ b/app.py @@ -22,6 +22,16 @@ tag = "\033[31m[warn]\033[0m" 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 @@ if __name__ == "__main__": 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)) diff --git a/scripts/iib/api.py b/scripts/iib/api.py index 49871a4..0e9e310 100644 --- a/scripts/iib/api.py +++ b/scripts/iib/api.py @@ -12,7 +12,7 @@ from scripts.iib.tool import ( 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 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs): @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 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs): @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 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs): # 生成缓存文件的路径 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") # 如果缓存文件存在,则直接返回该文件 diff --git a/scripts/iib/img_cache_gen.py b/scripts/iib/img_cache_gen.py new file mode 100644 index 0000000..be8d5e6 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/scripts/iib/tool.py b/scripts/iib/tool.py index 251a5fd..90e88b7 100644 --- a/scripts/iib/tool.py +++ b/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: diff --git a/scripts/iib/video_cover.py b/scripts/iib/video_cover_gen.py similarity index 77% rename from scripts/iib/video_cover.py rename to scripts/iib/video_cover_gen.py index b0ed7f1..bc2fd8e 100644 --- a/scripts/iib/video_cover.py +++ b/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 generate_video_covers(dirs): 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 generate_video_covers(dirs): 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)