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

Revert "Revert"

pull/340/head
zanllp 2023-07-23 23:26:45 +08:00 committed by GitHub
parent 7aadccc70e
commit 7d49900dc4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 224 additions and 115 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
.env.example
View File

@ -10,8 +10,23 @@ IIB_SECRET_KEY=
# If you want to configure the language for the front-end pages, please set it on the extension's global settings page.
IIB_SERVER_LANG=auto
# Used for configuring whether to enable access control to the file system.
# If enabled, only access to the provided pre-set folders (including those provided by sd-webui and manually added to Quick Move) will be allowed.
# The optional choices are 'enable', 'disable', and 'auto'.
# ---------------------------- ACCESS_CONTROL ----------------------------
# Used to configure whether to enable access control to the file system.
# If enabled, only access to the provided pre-set folders (including those provided by sd-webui and manually \
# added to Quick Move or specified via IIB_ACCESS_CONTROL_ALLOWED_PATHS) will be allowed.
# The available options are 'enable', 'disable', and 'auto'.
# The default value is 'auto', which will be determined based on the command-line parameters used to start sd-webui.
IIB_ACCESS_CONTROL=auto
# This variable is used to define a list of allowed paths for the application to access when access control mode is enabled.
# It can be set to a comma-separated string of file paths or directory paths, representing the resources that are allowed to be accessed by the application.
# In addition, if sd_webui_config or sd_webui_dir has been configured, or if you're running this repository as an extension of sd-webui,
# you can use the following shortcuts (txt2img, img2img, extra, save) as values for the ALLOWED_PATHS variable.
# IIB_ACCESS_CONTROL_ALLOWED_PATHS=save,extra,/output ...etc
# This variable is used to control fine-grained access control for different types of requests, but only if access control mode is enabled.
# It can be set to a string value that represents a specific permission or set of permissions, such as "read-only", "write-only", "read-write", or "no-access".
# This variable can be used to restrict access to certain API endpoints or data sources based on the permissions required by the user.
# IIB_ACCESS_CONTROL_PERMISSION=read-write

22
README-zh.md
View File

@ -6,7 +6,7 @@
[查看近期更新](https://github.com/zanllp/sd-webui-infinite-image-browsing/issues/131)
<p style="text-align:center;margin:0 32px">不仅仅是图像浏览器，更是一个强大的图像管理器。精确的图像搜索配合多选操作，成倍提高效率。更是支持以独立模式运行，无需SD-Webui</p>
<p style="text-align:center;margin:0 32px">不仅仅是图像浏览器，更是一个强大的图像管理器。精确的图像搜索配合多选操作进行筛选/归档/打包，成倍提高效率。更是支持以独立模式运行，无需SD-Webui</p>
@ -16,8 +16,9 @@ https://github.com/zanllp/sd-webui-infinite-image-browsing/assets/25872019/807b8
## 主要特性
### 🔥 极佳性能
- 生成缓存后，图像可以在几毫秒内显示。
- 默认使用缩略图显示图像，默认大小为256像素，您可以在全局设置页面中调整缩略图大小。
- 存在缓存的情况下后，图像可以在几毫秒内显示。
- 默认使用缩略图显示图像，默认大小为512像素，您可以在全局设置页中调整缩略图分辨率。
- 你还可以控制网格图像的宽度，允许以64px到1024px的宽度范围进行显示
### 🔍 图像搜索和收藏
- 将会把Prompt、Model、Lora等信息转成标签，将根据使用频率排序以供进行精确的搜索。
@ -29,7 +30,7 @@ https://github.com/zanllp/sd-webui-infinite-image-browsing/assets/25872019/807b8
### 🖼️ 图像预览和“发送到”
- 支持查看图像生成信息。全屏预览下同样支持。
- 支持将图像发送到其他选项卡 和 ControlNet。
- 支持将图像发送到其他选项卡和其他插件，例如 ControlNet, openOutpaint。
- 支持全屏预览，并且支持在全屏预览下使用自定义快捷键进行操作
- 支持在全屏预览模式下通过按下方向键或点击按钮移动到前一个或后一个图像。
@ -38,16 +39,15 @@ https://github.com/zanllp/sd-webui-infinite-image-browsing/assets/25872019/807b8
- 您可以使用 Python 独立运行它。
- 还提供桌面应用程序版本。
### "Walk"模式
### 🚶‍♀️ Walk模式
- 自动加载下一个文件夹 `(类似于 os.walk)`，可让您无需分页浏览所有图像。
- 已测试可正常处理超过 27,000 个文件。
### 🌳 基于文件树结构的预览和文件操作
- 支持基于文件树结构的预览。
- 支持自动刷新。
- 支持基本文件操作以及多选删除/移动。
- 支持基本文件操作以及多选删除/移动/复制，新建文件夹等。
- 按住 Ctrl、Shift 或 Cmd 键可选择多个项目。
- 支持通过右键菜单直接将文件发送到其他文件夹。
### 🆚 图像对比 (类似ImgSli)
- 提供两张图片的并排比较
@ -58,13 +58,19 @@ https://github.com/zanllp/sd-webui-infinite-image-browsing/assets/25872019/807b8
### 🔐 隐私和安全
- 支持配置一个key来进行身份验证
- 支持自定义secret key来进行身份验证
- 支持配置对文件系统的访问控制，默认将在服务允许公开访问时启用
- 支持自定义访问控制允许的路径。
- 支持控制访问权限。你可以让IIB以只读模式运行
- [点击这里查看详情](.env.example)
### ⌨️ 快捷键
- 支持删除和添加/移除Tag，在全局设置页进行自定义触发按钮
### 📦 打包 / 批量下载
- 允许你一次性打包下载多个图像
- 数据来源可以是搜索结果/普通的图像网格查看页面/walk模式等。使用拖拽或者“发送到”都可将图片添加待处理列表
强烈推荐使用在`在新页面打开`，比塞在gradio里舒服太多。

28
README.md
View File

@ -4,14 +4,15 @@
# Stable Diffusion webui Infinite Image Browsing
<p style="text-align:center;margin:0 32px;">Not just an image browser, it is also a powerful image manager. Accurate image search combined with multi-select operations greatly improves efficiency. It also supports running in independent mode without SD-Webui.</p>
<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>
https://github.com/zanllp/sd-webui-infinite-image-browsing/assets/25872019/807b890b-7be8-4816-abba-f3ad340a2232
## Key Features
### 🔥 Excellent Performance
- Once caching is generated, images can be displayed in just a few milliseconds.
- Images are displayed by default using thumbnails with a default size of 256 pixels. You can adjust the size of the thumbnails in the global settings page.
- 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.
### 🔍 Image Search & Favorite
- The prompt, model, Lora, and other information will be converted into tags and sorted by frequency of use for precise searching.
@ -23,7 +24,7 @@ https://github.com/zanllp/sd-webui-infinite-image-browsing/assets/25872019/807b8
### 🖼️ View Image & `Send To`
- Supports viewing image generation information. Also supported in full-screen preview mode.
- Supports sending images to other tabs and third-party extensions such as ControlNet openOutpaint.
- Supports sending images to other tabs and third-party extensions such as ControlNet , openOutpaint.
- Support full-screen preview and enable custom shortcut key operations while in full-screen preview mode.
- Support navigating to the previous or next image in full-screen preview mode by pressing arrow keys or clicking buttons.
@ -34,16 +35,15 @@ https://github.com/zanllp/sd-webui-infinite-image-browsing/assets/25872019/807b8
- The desktop app version is also available.
### "Walk" Mode
### 🚶‍♀️ Walk Mode
- Automatically load the next folder `(similar to os.walk)`, allowing you to browse all images without paging.
- Tested to work properly with over 27,000 files.
### 🌳 Preview based on File Tree Structure & File operations
- Supports preview based on the file tree structure.
- Supports basic file operations as well as multi-select deletion/moving.
- Press and hold Ctrl, Shift, or Cmd to select multiple items.
- Supports sending files directly to other folders via context menu.
- Support for automatic refresh
- Supports file tree-based preview.
- Supports automatic refreshing.
- Supports basic file operations, such as multiple selection for deleting/moving/copying, and creating new folders.
- Hold down the Ctrl, Shift, or Cmd key to select multiple items.
### 🆚 image comparison (similar to Imgsli)
- Provides a side-by-side comparison of two images.
@ -53,10 +53,16 @@ https://github.com/zanllp/sd-webui-infinite-image-browsing/assets/25872019/807b8
- If you would like to add a new language, please refer to [i18n.ts](https://github.com/zanllp/sd-webui-infinite-image-browsing/blob/main/vue/src/i18n.ts) and submit the relevant code.
### 🔐 Privacy and Security
- Support configuring a key for authentication.
- Support configuring access control for the file system. By default, it will be enabled when the service allows public access.
- 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 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)
### 📦 Packaging/Batch Download
- Allows you to download multiple images at once.
- The data source can be search results, a regular image grid view page, walk mode, etc. Images can be added to the processing list through drag-and-drop or "Send To".
### ⌨️ Keyboard Shortcuts
- Allows for deleting and adding/removing tags, with customizable trigger buttons in the global settings page.

31
app.py
View File

@ -4,7 +4,7 @@ from fastapi.responses import FileResponse
import uvicorn
import os
from scripts.iib.api import infinite_image_browsing_api, index_html_path
from scripts.iib.tool import get_sd_webui_conf, get_valid_img_dirs, sd_img_dirs
from scripts.iib.tool import get_sd_webui_conf, get_valid_img_dirs, sd_img_dirs, normalize_paths
from scripts.iib.db.datamodel import DataBase, Image
from scripts.iib.db.update_image_data import update_image_data
import argparse
@ -17,33 +17,6 @@ default_port = 8000
default_host = "127.0.0.1"
def normalize_paths(paths: List[str]):
"""
Normalize a list of paths, ensuring that each path is an absolute path with no redundant components.
Args:
paths (List[str]): A list of paths to be normalized.
Returns:
List[str]: A list of normalized paths.
"""
res: List[str] = []
for path in paths:
# Skip empty or blank paths
if not path or len(path.strip()) == 0:
continue
# If the path is already an absolute path, use it as is
if os.path.isabs(path):
abs_path = path
# Otherwise, make the path absolute by joining it with the current working directory
else:
abs_path = os.path.join(os.getcwd(), path)
# If the absolute path exists, add it to the result after normalizing it
if os.path.exists(abs_path):
res.append(os.path.normpath(abs_path))
return res
def sd_webui_paths_check(sd_webui_config: str, relative_to_config: bool):
conf = {}
with open(sd_webui_config, "r") as f:
@ -148,7 +121,7 @@ class AppUtils:
infinite_image_browsing_api(
app,
sd_webui_config=sd_webui_config,
extra_paths_cli=normalize_paths(extra_paths),
extra_paths_cli=normalize_paths(extra_paths, os.getcwd()),
sd_webui_path_relative_to_config=self.sd_webui_path_relative_to_config,
allow_cors=self.allow_cors,
enable_shutdown=self.enable_shutdown,

14
install.py
View File

@ -4,14 +4,12 @@ import pkg_resources
req_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), "requirements.txt")
def dist2package(dist):
if dist == "python-dotenv":
package = "dotenv"
elif dist == "Pillow":
package = "PIL"
else:
package = dist
return package
def dist2package(dist: str):
return ({
"pyfunctional": "functional",
"python-dotenv": "dotenv",
"Pillow": "PIL"
}).get(dist, dist)
# copy from controlnet, thanks
with open(req_file) as file:

3
requirements.txt
View File

@ -2,4 +2,5 @@ fastapi
uvicorn
piexif
python-dotenv
Pillow
Pillow
pyfunctional

178
scripts/iib/api.py
View File

@ -18,12 +18,13 @@ from scripts.iib.tool import (
open_folder,
get_img_geninfo_txt_path,
unique_by,
create_zip_file
create_zip_file,
normalize_paths,
)
from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
import asyncio
from typing import List, Optional, TypedDict
from typing import List, Optional
from pydantic import BaseModel
from fastapi.responses import FileResponse
from PIL import Image
@ -37,9 +38,11 @@ from scripts.iib.db.datamodel import (
Floder,
ImageTag,
ExtraPath,
FileInfoDict,
)
from scripts.iib.db.update_image_data import update_image_data
from scripts.iib.logger import logger
from functional import seq
index_html_path = os.path.join(cwd, "vue/dist/index.html") # 在app.py也被使用
@ -51,8 +54,23 @@ secret_key = os.getenv("IIB_SECRET_KEY")
if secret_key:
print("Secret key loaded successfully. ")
WRITEABLE_PERMISSIONS = ["read-write", "write-only"]
async def get_token(request: Request):
is_api_writeable = not (os.getenv("IIB_ACCESS_CONTROL_PERMISSION")) or (
os.getenv("IIB_ACCESS_CONTROL_PERMISSION") in WRITEABLE_PERMISSIONS
)
async def write_permission_required():
if not is_api_writeable:
error_msg = (
"User is not authorized to perform this action. Required permission: "
+ ", ".join(WRITEABLE_PERMISSIONS)
)
raise HTTPException(status_code=403, detail=error_msg)
async def verify_secret(request: Request):
if not secret_key:
return
token = request.cookies.get("IIB_S")
@ -84,7 +102,32 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
pass
def update_all_scanned_paths():
paths = img_search_dirs + mem["extra_paths"] + kwargs.get("extra_paths_cli", [])
allowed_paths = os.getenv("IIB_ACCESS_CONTROL_ALLOWED_PATHS")
if allowed_paths:
sd_webui_conf = get_sd_webui_conf(**kwargs)
path_config_key_map = {
"save": "outdir_save",
"extra": "outdir_extras_samples",
"txt2img": "outdir_txt2img_samples",
"img2img": "outdir_img2img_samples",
}
def path_map(path: str):
path = path.strip()
if path in path_config_key_map:
return sd_webui_conf.get(path_config_key_map.get(path))
return path
paths = normalize_paths(
seq(allowed_paths.split(","))
.map(path_map)
.filter(lambda x: x)
.to_list()
)
else:
paths = (
img_search_dirs + mem["extra_paths"] + kwargs.get("extra_paths_cli", [])
)
mem["all_scanned_paths"] = unique_by(paths)
update_all_scanned_paths()
@ -122,6 +165,8 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
def is_path_trusted(path: str):
if not enable_access_control:
return True
if not os.path.isabs(path):
path = os.path.normpath(os.path.join(os.getcwd(), path))
try:
parent_paths = mem["all_scanned_paths"]
path = os.path.normpath(path)
@ -140,6 +185,9 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
if not is_path_trusted(path):
raise HTTPException(status_code=403)
def filter_allowed_files(files: List[FileInfoDict]):
return [x for x in files if is_path_trusted(x["fullpath"])]
static_dir = f"{cwd}/vue/dist"
if os.path.exists(static_dir):
app.mount(
@ -152,7 +200,7 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
async def greeting():
return "hello"
@app.get(f"{pre}/global_setting", dependencies=[Depends(get_token)])
@app.get(f"{pre}/global_setting", dependencies=[Depends(verify_secret)])
async def global_setting():
all_custom_tags = []
@ -182,7 +230,10 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
class DeleteFilesReq(BaseModel):
file_paths: List[str]
@app.post(pre + "/delete_files", dependencies=[Depends(get_token)])
@app.post(
pre + "/delete_files",
dependencies=[Depends(verify_secret), Depends(write_permission_required)],
)
async def delete_files(req: DeleteFilesReq):
conn = DataBase.get_conn()
for path in req.file_paths:
@ -220,7 +271,10 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
class CreateFoldersReq(BaseModel):
dest_folder: str
@app.post(pre + "/mkdirs", dependencies=[Depends(get_token)])
@app.post(
pre + "/mkdirs",
dependencies=[Depends(verify_secret), Depends(write_permission_required)],
)
async def create_folders(req: CreateFoldersReq):
if enable_access_control:
if not is_path_under_parents(req.dest_folder):
@ -232,7 +286,10 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
dest: str
create_dest_folder: Optional[bool] = False
@app.post(pre + "/copy_files", dependencies=[Depends(get_token)])
@app.post(
pre + "/copy_files",
dependencies=[Depends(verify_secret), Depends(write_permission_required)],
)
async def copy_files(req: MoveFilesReq):
for path in req.file_paths:
check_path_trust(path)
@ -249,7 +306,10 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
)
raise HTTPException(400, detail=error_msg)
@app.post(pre + "/move_files", dependencies=[Depends(get_token)])
@app.post(
pre + "/move_files",
dependencies=[Depends(verify_secret), Depends(write_permission_required)],
)
async def move_files(req: MoveFilesReq):
if req.create_dest_folder:
os.makedirs(req.dest, exist_ok=True)
@ -281,16 +341,7 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
)
raise HTTPException(400, detail=error_msg)
class FileInfoDict(TypedDict):
type: str
date: float
size: int
name: str
bytes: bytes
created_time: float
fullpath: str
@app.get(pre + "/files", dependencies=[Depends(get_token)])
@app.get(pre + "/files", dependencies=[Depends(verify_secret)])
async def get_target_floder_files(folder_path: str):
files: List[FileInfoDict] = []
try:
@ -343,9 +394,9 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
logger.error(e)
raise HTTPException(status_code=400, detail=str(e))
return {"files": [x for x in files if is_path_trusted(x["fullpath"])]}
return {"files": filter_allowed_files(files)}
@app.get(pre + "/image-thumbnail", dependencies=[Depends(get_token)])
@app.get(pre + "/image-thumbnail", dependencies=[Depends(verify_secret)])
async def thumbnail(path: str, t: str, size: str = "256x256"):
check_path_trust(path)
if not temp_path:
@ -378,7 +429,7 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
headers={"Cache-Control": "max-age=31536000", "ETag": hash},
)
@app.get(pre + "/file", dependencies=[Depends(get_token)])
@app.get(pre + "/file", dependencies=[Depends(verify_secret)])
async def get_file(path: str, t: str, disposition: Optional[str] = None):
filename = path
import mimetypes
@ -409,12 +460,12 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
headers=headers,
)
@app.post(pre + "/send_img_path", dependencies=[Depends(get_token)])
@app.post(pre + "/send_img_path", dependencies=[Depends(verify_secret)])
async def api_set_send_img_path(path: str):
send_img_path["value"] = path
# 等待图片信息生成完成
@app.get(pre + "/gen_info_completed", dependencies=[Depends(get_token)])
@app.get(pre + "/gen_info_completed", dependencies=[Depends(verify_secret)])
async def api_set_send_img_path():
for _ in range(30): # timeout 3s
if send_img_path["value"] == "": # 等待setup里面生成完成
@ -424,7 +475,7 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
await asyncio.sleep(0.1)
return send_img_path["value"] == ""
@app.get(pre + "/image_geninfo", dependencies=[Depends(get_token)])
@app.get(pre + "/image_geninfo", dependencies=[Depends(verify_secret)])
async def image_geninfo(path: str):
with Image.open(path) as img:
return read_info_from_image(img, path)
@ -432,25 +483,27 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
class CheckPathExistsReq(BaseModel):
paths: List[str]
@app.post(pre + "/check_path_exists", dependencies=[Depends(get_token)])
@app.post(pre + "/check_path_exists", dependencies=[Depends(verify_secret)])
async def check_path_exists(req: CheckPathExistsReq):
res = {}
for path in req.paths:
res[path] = os.path.exists(path)
res[path] = os.path.exists(path) and is_path_trusted(path)
return res
@app.get(pre)
def index_bd():
return FileResponse(index_html_path)
class PathsReq(BaseModel):
paths: List[str]
class OpenFolderReq(BaseModel):
path: str
@app.post(pre + "/open_folder", dependencies=[Depends(get_token)])
@app.post(
pre + "/open_folder",
dependencies=[Depends(verify_secret), Depends(write_permission_required)],
)
def open_folder_using_explore(req: OpenFolderReq):
if not is_path_trusted(req.path):
raise HTTPException(status_code=403)
@ -463,8 +516,11 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
raise HTTPException(status_code=403, detail="Shutdown is disabled.")
os.kill(os.getpid(), 9)
return {"message": "Application is shutting down."}
@app.post(pre + "/zip", dependencies=[Depends(get_token)])
@app.post(
pre + "/zip",
dependencies=[Depends(verify_secret), Depends(write_permission_required)],
)
def zip_files(req: PathsReq):
now = datetime.now()
timestamp = now.strftime("%Y-%m-%d-%H-%M-%S")
@ -476,7 +532,7 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
db_pre = pre + "/db"
@app.get(db_pre + "/basic_info", dependencies=[Depends(get_token)])
@app.get(db_pre + "/basic_info", dependencies=[Depends(verify_secret)])
async def get_db_basic_info():
conn = DataBase.get_conn()
img_count = DbImg.count(conn)
@ -489,7 +545,7 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
"expired_dirs": expired_dirs,
}
@app.get(db_pre + "/expired_dirs", dependencies=[Depends(get_token)])
@app.get(db_pre + "/expired_dirs", dependencies=[Depends(verify_secret)])
async def get_db_expired():
conn = DataBase.get_conn()
expired_dirs = Floder.get_expired_dirs(conn)
@ -498,7 +554,10 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
"expired_dirs": expired_dirs,
}
@app.post(db_pre + "/update_image_data", dependencies=[Depends(get_token)])
@app.post(
db_pre + "/update_image_data",
dependencies=[Depends(verify_secret)],
)
async def update_image_db_data():
try:
DataBase._initing = True
@ -518,17 +577,15 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
or_tags: List[int]
not_tags: List[int]
@app.post(db_pre + "/match_images_by_tags", dependencies=[Depends(get_token)])
@app.post(db_pre + "/match_images_by_tags", dependencies=[Depends(verify_secret)])
async def match_image_by_tags(req: MatchImagesByTagsReq):
conn = DataBase.get_conn()
return [
x.to_file_info()
for x in ImageTag.get_images_by_tags(
conn, {"and": req.and_tags, "or": req.or_tags, "not": req.not_tags}
)
]
imgs = ImageTag.get_images_by_tags(
conn, {"and": req.and_tags, "or": req.or_tags, "not": req.not_tags}
)
return filter_allowed_files([x.to_file_info() for x in imgs])
@app.get(db_pre + "/img_selected_custom_tag", dependencies=[Depends(get_token)])
@app.get(db_pre + "/img_selected_custom_tag", dependencies=[Depends(verify_secret)])
async def get_img_selected_custom_tag(path: str):
path = os.path.normpath(path)
if not is_valid_image_path(path):
@ -547,7 +604,7 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
# tags = Tag.get_all_custom_tag()
return ImageTag.get_tags_for_image(conn, img.id, type="custom")
@app.post(db_pre + "/get_image_tags", dependencies=[Depends(get_token)])
@app.post(db_pre + "/get_image_tags", dependencies=[Depends(verify_secret)])
async def get_img_tags(req: PathsReq):
conn = DataBase.get_conn()
return ImageTag.batch_get_tags_by_path(conn, req.paths)
@ -556,7 +613,10 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
img_path: str
tag_id: int
@app.post(db_pre + "/toggle_custom_tag_to_img", dependencies=[Depends(get_token)])
@app.post(
db_pre + "/toggle_custom_tag_to_img",
dependencies=[Depends(verify_secret), Depends(write_permission_required)],
)
async def toggle_custom_tag_to_img(req: ToggleCustomTagToImgReq):
conn = DataBase.get_conn()
path = os.path.normpath(req.img_path)
@ -590,7 +650,10 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
class AddCustomTagReq(BaseModel):
tag_name: str
@app.post(db_pre + "/add_custom_tag", dependencies=[Depends(get_token)])
@app.post(
db_pre + "/add_custom_tag",
dependencies=[Depends(verify_secret), Depends(write_permission_required)],
)
async def add_custom_tag(req: AddCustomTagReq):
conn = DataBase.get_conn()
tag = Tag.get_or_create(conn, name=req.tag_name, type="custom")
@ -600,7 +663,10 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
class RemoveCustomTagReq(BaseModel):
tag_id: str
@app.post(db_pre + "/remove_custom_tag", dependencies=[Depends(get_token)])
@app.post(
db_pre + "/remove_custom_tag",
dependencies=[Depends(verify_secret), Depends(write_permission_required)],
)
async def remove_custom_tag(req: RemoveCustomTagReq):
conn = DataBase.get_conn()
ImageTag.remove(conn, tag_id=req.tag_id)
@ -610,22 +676,27 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
img_id: int
tag_id: str
@app.post(db_pre + "/remove_custom_tag_from_img", dependencies=[Depends(get_token)])
@app.post(
db_pre + "/remove_custom_tag_from_img",
dependencies=[Depends(verify_secret), Depends(write_permission_required)],
)
async def remove_custom_tag_from_img(req: RemoveCustomTagFromReq):
conn = DataBase.get_conn()
ImageTag.remove(conn, image_id=req.img_id, tag_id=req.tag_id)
@app.get(db_pre + "/search_by_substr", dependencies=[Depends(get_token)])
@app.get(db_pre + "/search_by_substr", dependencies=[Depends(verify_secret)])
async def search_by_substr(substr: str):
conn = DataBase.get_conn()
imgs = DbImg.find_by_substring(conn=conn, substring=substr)
return [x.to_file_info() for x in imgs]
return filter_allowed_files([x.to_file_info() for x in imgs])
class ScannedPathModel(BaseModel):
path: str
@app.post(
f"{db_pre}/scanned_paths", status_code=201, dependencies=[Depends(get_token)]
f"{db_pre}/scanned_paths",
status_code=201,
dependencies=[Depends(verify_secret), Depends(write_permission_required)],
)
async def create_scanned_path(scanned_path: ScannedPathModel):
if enable_access_control:
@ -641,14 +712,17 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
@app.get(
f"{db_pre}/scanned_paths",
response_model=List[ScannedPathModel],
dependencies=[Depends(get_token)],
dependencies=[Depends(verify_secret)],
)
async def read_scanned_paths():
conn = DataBase.get_conn()
paths = ExtraPath.get_extra_paths(conn, "scanned")
return [{"path": path.path} for path in paths]
@app.delete(f"{db_pre}/scanned_paths", dependencies=[Depends(get_token)])
@app.delete(
f"{db_pre}/scanned_paths",
dependencies=[Depends(verify_secret), Depends(write_permission_required)],
)
async def delete_scanned_path(scanned_path: ScannedPathModel):
conn = DataBase.get_conn()
ExtraPath.remove(conn, scanned_path.path)

14
scripts/iib/db/datamodel.py
View File

@ -1,5 +1,5 @@
from sqlite3 import Connection, connect
from typing import Dict, List, Optional
from typing import Dict, List, Optional, TypedDict
from scripts.iib.tool import (
cwd,
get_modified_date,
@ -13,7 +13,15 @@ from contextlib import closing
import os
import threading
class FileInfoDict(TypedDict):
type: str
date: float
size: int
name: str
bytes: bytes
created_time: float
fullpath: str
class DataBase:
local = threading.local()
@ -59,7 +67,7 @@ class Image:
self.size = size
self.date = date
def to_file_info(self):
def to_file_info(self) -> FileInfoDict:
return {
"type": "file",
"id": self.id,

28
scripts/iib/tool.py
View File

@ -60,6 +60,32 @@ def get_sd_webui_conf(**kwargs):
pass
return {}
def normalize_paths(paths: List[str], base = cwd):
"""
Normalize a list of paths, ensuring that each path is an absolute path with no redundant components.
Args:
paths (List[str]): A list of paths to be normalized.
Returns:
List[str]: A list of normalized paths.
"""
res: List[str] = []
for path in paths:
# Skip empty or blank paths
if not path or len(path.strip()) == 0:
continue
# If the path is already an absolute path, use it as is
if os.path.isabs(path):
abs_path = path
# Otherwise, make the path absolute by joining it with the current working directory
else:
abs_path = os.path.join(base, path)
# If the absolute path exists, add it to the result after normalizing it
if os.path.exists(abs_path):
res.append(os.path.normpath(abs_path))
return res
def get_valid_img_dirs(
conf,
@ -432,3 +458,5 @@ def open_folder(folder_path, file_path=None):
subprocess.run(["open", folder])
elif os.name == "posix":
subprocess.run(["xdg-open", folder])