1580158
/
sd-webui-infinite-image-browsing
mirror of https://github.com/zanllp/sd-webui-infinite-image-browsing.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add support for pre-generating cache via startup parameters and specifying cache directory through environment variables

pull/684/head
zanllp 2024-07-08 00:07:12 +08:00
parent d24fc161e2
commit 17f048ad0c
8 changed files with 126 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
.env.example
View File

@ -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 ----------------------------

2
README-zh.md
View File

@ -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等信息转成标签，将根据使用频率排序以供进行精确的搜索。

2
README.md
View File

@ -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.

48
app.py
View File

@ -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))

12
scripts/iib/api.py
View File

@ -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")
# 如果缓存文件存在，则直接返回该文件

58
scripts/iib/img_cache_gen.py Normal file
View File

@ -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")

6
scripts/iib/tool.py
View File

@ -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:

13
scripts/iib/video_cover.py → scripts/iib/video_cover_gen.py
View File

@ -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)