Improve performance, add option for pre-generating video covers, and update README

README-zh.md

@@ -4,9 +4,10 @@
 
 
 [查看近期更新](https://github.com/zanllp/sd-webui-infinite-image-browsing/wiki/Change-log)
+
 [安装/运行](#安装运行)
 
-<p style="text-align:center;margin:0 32px">不仅仅是图像浏览器，更是一个强大的图像管理器。精确的图像搜索配合多选操作进行筛选/归档/打包，成倍提高效率。更是支持以独立模式运行，无需SD-Webui</p>
+<p >不仅仅是图像浏览器，更是一个强大的图像管理器。精确的图像搜索配合多选操作进行筛选/归档/打包，成倍提高效率。更是支持以独立模式运行，无需SD-Webui。同时它也适合用于对普通图片和视频的管理，你可以对它们打标签或者搜索，我们做了诸多优化保证在极端情况的下的性能</p>
 
 
 
@@ -60,7 +61,7 @@ https://github.com/zanllp/sd-webui-infinite-image-browsing/assets/25872019/807b8
 
 ### 🔐 隐私和安全
 - 支持自定义secret key来进行身份验证
-- 支持配置对文件系统的访问控制，默认将在服务允许公开访问时启用
+- 支持配置对文件系统的访问控制，默认将在服务允许公开访问时启用(仅限作为sd-webui的拓展时)
 - 支持自定义访问控制允许的路径。
 - 支持控制访问权限。你可以让IIB以只读模式运行
 - [点击这里查看详情](.env.example)

README.md

@@ -6,7 +6,9 @@
 
 #  Stable Diffusion webui Infinite Image Browsing
 
-<p style="text-align:center;margin:0 32px;">It's not just an image browser, but also a powerful image manager. Precise image search combined with multi-selection operations allows for filtering/archiving/packaging, greatly increasing efficiency. It also supports running in standalone mode, without the need for SD-Webui.</p>
+<p>It's not just an image/video browser, but also a powerful image manager. Precise image search combined with multi-selection operations allows for filtering/archiving/packaging, greatly increasing efficiency. It also supports running in standalone mode, without the need for SD-Webui.
+
+It is also suitable for managing regular photos and videos, allowing you to tag or search them. We have made numerous optimizations to ensure performance under extreme conditions.</p>
 
 https://github.com/zanllp/sd-webui-infinite-image-browsing/assets/25872019/807b890b-7be8-4816-abba-f3ad340a2232
 
@@ -70,7 +72,7 @@ You can add your own parser through https://github.com/zanllp/sd-webui-infinite-
 
 ### 🔐 Privacy and Security
 - Supports custom secret key for authentication.
-- Supports configuring access control for the file system, which will be enabled by default when the service allows public access.
+- Supports configuring access control for the file system, which will be enabled by default when the service allows public access (Only when used as an extension of sd-webui).
 - Supports customizing the allowed paths for access control.
 - Supports controlling access permissions. You can run IIB in read-only mode.
 - [Click here to see details](.env.example)

app.py

@@ -11,7 +11,7 @@ from scripts.iib.tool import (
     sd_img_dirs,
     normalize_paths,
 )
-from scripts.iib.db.datamodel import DataBase, Image
+from scripts.iib.db.datamodel import DataBase, Image, ExtraPath
 from scripts.iib.db.update_image_data import update_image_data
 import argparse
 from typing import Optional, Coroutine
@@ -75,8 +75,9 @@ class AppUtils:
         allow_cors=False,
         enable_shutdown=False,
         sd_webui_dir: Optional[str] = None,
-        base: Optional[str]=None,
+        base: Optional[str] = None,
         export_fe_fn=False,
+        **args: dict,
     ):
         """
         Parameter definitions can be found by running the `python app.py -h `command or by examining the setup_parser() function.
@@ -88,8 +89,8 @@ class AppUtils:
         self.allow_cors = allow_cors
         self.enable_shutdown = enable_shutdown
         self.sd_webui_dir = sd_webui_dir
-        if base and not base.startswith('/'):
-            base = '/' + base
+        if base and not base.startswith("/"):
+            base = "/" + base
         self.base = base
         self.export_fe_fn = export_fe_fn
         if sd_webui_dir:
@@ -140,7 +141,7 @@ class AppUtils:
             allow_cors=self.allow_cors,
             enable_shutdown=self.enable_shutdown,
             launch_mode="server",
-            base= self.base,
+            base=self.base,
             export_fe_fn=self.export_fe_fn,
         )
 
@@ -165,7 +166,7 @@ class AppUtils:
 
 def setup_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(
-        description="A fast and powerful image browser for Stable Diffusion webui."
+        description="A fast and powerful image/video browser with infinite scrolling and advanced search capabilities. It also supports parsing/viewing image information generated by multiple AI software."
     )
     parser.add_argument(
         "--host", type=str, default=default_host, help="The host to use"
@@ -179,6 +180,11 @@ def setup_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--update_image_index", action="store_true", help="Update the image index"
     )
+    parser.add_argument(
+        "--generate_video_cover",
+        action="store_true",
+        help="Pre-generate video cover images to speed up browsing.",
+    )
     parser.add_argument(
         "--extra_paths",
         nargs="+",
@@ -249,4 +255,13 @@ async def async_launch_app(
 if __name__ == "__main__":
     parser = setup_parser()
     args = parser.parse_args()
+    args_dict = vars(args)
+
+    if args_dict.get("generate_video_cover"):
+        from scripts.iib.video_cover import generate_video_covers
+
+        conn = DataBase.get_conn()
+        generate_video_covers(map(lambda x: x.path, ExtraPath.get_extra_paths(conn)))
+        exit(0)
+
     launch_app(**vars(args))

scripts/iib/api.py

@@ -484,6 +484,16 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
                 headers={"Cache-Control": "max-age=31536000", "ETag": hash},
             )
 
+                
+        # 如果小于64KB，直接返回原图
+        if os.path.getsize(path) < 64 * 1024:
+            return FileResponse(
+                path,
+                media_type="image/" + path.split(".")[-1],
+                headers={"Cache-Control": "max-age=31536000", "ETag": hash},
+            )
+        
+
         # 如果缓存文件不存在，则生成缩略图并保存
         with Image.open(path) as img:
             w, h = size.split("x")

scripts/iib/video_cover.py

@@ -0,0 +1,58 @@
+import hashlib
+import os
+from typing import List
+from scripts.iib.tool import get_formatted_date, get_temp_path, is_video_file
+from concurrent.futures import ThreadPoolExecutor
+import time
+
+
+def generate_video_covers(dirs):
+  start_time = time.time()
+  import imageio.v3 as iio
+  temp_path = get_temp_path()
+
+  def process_video(item):
+    if item.is_dir():
+      print(f"Processing directory: {item.path}")
+      for sub_item in os.scandir(item.path):
+        process_video(sub_item)
+      return
+    if not os.path.exists(item.path) or not is_video_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(temp_path, "iib_cache", "video_cover", hash_dir)
+      cache_path = os.path.join(cache_dir, "cover.webp")
+
+      # 如果缓存文件存在，则直接返回该文件
+      if os.path.exists(cache_path):
+        print(f"Video cover already exists: {path}")
+        return
+
+      frame = iio.imread(
+        path,
+        index=16,
+        plugin="pyav",
+      )
+
+      os.makedirs(cache_dir, exist_ok=True)
+      iio.imwrite(cache_path, frame, extension=".webp")
+      print(f"Video cover generated: {path}")
+    except Exception as e:
+      print(f"Error generating video cover: {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_video, item)
+
+  print("Video covers generated successfully.")
+  end_time = time.time()
+  execution_time = end_time - start_time
+  print(f"Execution time: {execution_time} seconds")