From 09f37c5c86452aa22d971a62fe2047ae42ba4269 Mon Sep 17 00:00:00 2001 From: zanllp Date: Wed, 28 Jun 2023 08:39:26 +0800 Subject: [PATCH] Add option to use relative paths based on sd_webui_config file path --- app.py | 107 ++++++++++++++++++++++++------------------- scripts/iib/tool.py | 16 +++++-- scripts/iib_setup.py | 1 - 3 files changed, 72 insertions(+), 52 deletions(-) diff --git a/app.py b/app.py index 10e47da..2bf43b1 100644 --- a/app.py +++ b/app.py @@ -8,9 +8,8 @@ from scripts.iib.tool import get_sd_webui_conf, get_valid_img_dirs, sd_img_dirs from scripts.iib.db.datamodel import DataBase, Image from scripts.iib.db.update_image_data import update_image_data import argparse -import logging from typing import Optional, Coroutine -import asyncio +import json tag = "\033[31m[warn]\033[0m" @@ -45,12 +44,16 @@ def normalize_paths(paths: List[str]): return res -def sd_webui_paths_check(sd_webui_config: str): - import json - +def sd_webui_paths_check(sd_webui_config: str, relative_to_config: bool): conf = {} with open(sd_webui_config, "r") as f: conf = json.loads(f.read()) + if relative_to_config: + for dir in sd_img_dirs: + if not os.path.isabs(conf[dir]): + conf[dir] = os.path.normpath( + os.path.join(sd_webui_config, "../", conf[dir]) + ) paths = [conf.get(key) for key in sd_img_dirs] paths_check(paths) @@ -67,8 +70,13 @@ def paths_check(paths): print(f"{tag} The path '{abs_path}' will be ignored (value: {path}).") -def update_image_index_func(sd_webui_config: str): - dirs = get_valid_img_dirs(get_sd_webui_conf(sd_webui_config=sd_webui_config)) +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, + ) + ) if not len(dirs): return print(f"{tag} no valid image directories, skipped") conn = DataBase.get_conn() @@ -78,21 +86,21 @@ def update_image_index_func(sd_webui_config: str): print("update image index completed. ✨") -class AppUtils(object): +class AppUtils: def __init__( self, sd_webui_config: Optional[str] = None, update_image_index: bool = False, extra_paths: List[str] = [], + sd_webui_path_relative_to_config=False, ): """ - sd_webui_config, type=str, default=None, help="the path to the config file" - update_image_index, action="store_true", help="update the image index" - extra_paths", nargs="+", help="extra paths to use, will be added to Quick Move.", default=[] + Parameter definitions can be found by running the `python app.py -h `command or by examining the setup_parser() function. """ self.sd_webui_config = sd_webui_config self.update_image_index = update_image_index self.extra_paths = extra_paths + self.sd_webui_path_relative_to_config = sd_webui_path_relative_to_config def set_params(self, *args, **kwargs) -> None: """改变参数，与__init__的行为一致""" @@ -101,15 +109,10 @@ class AppUtils(object): @staticmethod def async_run(app: FastAPI, port: int = default_port) -> Coroutine: """ - 用于从异步运行的FastAPI，在jupyter notebook环境中非常有用 - - app为要启动的FastAPI实例 - port为要启动的端口 - - 返回协程uvicorn.Server().serve() + 用于从异步运行的 FastAPI，在 Jupyter Notebook 环境中非常有用 """ - # 不建议改成async def，并且用await替换return - # 因为这样会失去对server.serve()的控制 + # 不建议改成 async def，并且用 await 替换 return， + # 因为这样会失去对 server.serve() 的控制。 config = uvicorn.Config(app, host=default_host, port=port) server = uvicorn.Server(config) return server.serve() @@ -123,15 +126,18 @@ class AppUtils(object): extra_paths = self.extra_paths if sd_webui_config: - sd_webui_paths_check(sd_webui_config) + sd_webui_paths_check(sd_webui_config, self.sd_webui_path_relative_to_config) if update_image_index: - update_image_index_func(sd_webui_config) + do_update_image_index( + sd_webui_config, self.sd_webui_path_relative_to_config + ) paths_check(extra_paths) infinite_image_browsing_api( app, sd_webui_config=sd_webui_config, extra_paths_cli=normalize_paths(extra_paths), + sd_webui_path_relative_to_config=self.sd_webui_path_relative_to_config, ) def get_root_browser_app(self) -> FastAPI: @@ -150,40 +156,52 @@ class AppUtils(object): def setup_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(description="A fast and powerful image browser for Stable Diffusion webui.") - parser.add_argument("--port", type=int, help="The port to use", default=default_port) - parser.add_argument("--sd_webui_config", type=str, default=None, help="The path to the config file") - parser.add_argument("--update_image_index", action="store_true", help="Update the image index") + parser = argparse.ArgumentParser( + description="A fast and powerful image browser for Stable Diffusion webui." + ) parser.add_argument( - "--extra_paths", nargs="+", help="Extra paths to use, will be added to Quick Move.", default=[] + "--port", type=int, help="The port to use", default=default_port + ) + parser.add_argument( + "--sd_webui_config", type=str, default=None, help="The path to the config file" + ) + parser.add_argument( + "--update_image_index", action="store_true", help="Update the image index" + ) + parser.add_argument( + "--extra_paths", + nargs="+", + help="Extra paths to use, will be added to Quick Move.", + default=[], + ) + parser.add_argument( + "--sd_webui_path_relative_to_config", + action="store_true", + help="Use the file path of the sd_webui_config file as the base for all relative paths provided within the sd_webui_config file.", ) return parser -def launch_app(port: int = default_port, *args, **kwargs) -> None: +def launch_app(port: int = default_port, *args, **kwargs: dict) -> None: """ - 同步函数 - - 除 port 参数外，所传入的其他参数全都会被传递给AppUtils() - sd_webui_config, type=str, default=None, help="the path to the config file" - update_image_index, action="store_true", help="update the image index" - extra_paths", nargs="+", help="extra paths to use, will be added to Quick Move.", default=[] + Launches the application on the specified port. + Args: + **kwargs (dict): Optional keyword arguments that can be used to configure the application. + These can be viewed by running 'python app.py -h' or by checking the setup_parser() function. """ app_utils = AppUtils(*args, **kwargs) app = app_utils.get_root_browser_app() uvicorn.run(app, host=default_host, port=port) -async def async_launch_app(port: int = default_port, *args, **kwargs) -> None: +async def async_launch_app(port: int = default_port, *args, **kwargs: dict) -> None: """ - 协程函数 - - 除 port 参数外，所传入的其他参数全都会被传递给AppUtils() - sd_webui_config, type=str, default=None, help="the path to the config file" - update_image_index, action="store_true", help="update the image index" - extra_paths", nargs="+", help="extra paths to use, will be added to Quick Move.", default=[] + Asynchronously launches the application on the specified port. + Args: + **kwargs (dict): Optional keyword arguments that can be used to configure the application. + These can be viewed by running 'python app.py -h' or by checking the setup_parser() function. """ app_utils = AppUtils(*args, **kwargs) app = app_utils.get_root_browser_app() @@ -192,10 +210,5 @@ async def async_launch_app(port: int = default_port, *args, **kwargs) -> None: if __name__ == "__main__": parser = setup_parser() - cmd_params = parser.parse_args() - launch_app( - cmd_params.port, - sd_webui_config=cmd_params.sd_webui_config, - update_image_index=cmd_params.update_image_index, - extra_paths=cmd_params.extra_paths, - ) + args = parser.parse_args() + launch_app(**vars(args)) diff --git a/scripts/iib/tool.py b/scripts/iib/tool.py index d017a55..49b5f21 100644 --- a/scripts/iib/tool.py +++ b/scripts/iib/tool.py @@ -9,6 +9,7 @@ from typing import Dict import sys import piexif import piexif.helper +import json sd_img_dirs = [ "outdir_txt2img_samples", @@ -43,10 +44,16 @@ def get_sd_webui_conf(**kwargs): except: pass try: - with open(kwargs.get("sd_webui_config"), "r") as f: - import json - - return json.loads(f.read()) + sd_conf_path = kwargs.get("sd_webui_config") + with open(sd_conf_path, "r") as f: + obj = json.loads(f.read()) + if kwargs.get("sd_webui_path_relative_to_config"): + for dir in sd_img_dirs: + if not os.path.isabs(obj[dir]): + obj[dir] = os.path.normpath( + os.path.join(sd_conf_path, "../", obj[dir]) + ) + return obj except: pass return {} @@ -223,6 +230,7 @@ def get_img_geninfo_txt_path(path: str): if os.path.exists(txt_path): return txt_path + def read_info_from_image(image, path="") -> str: """ Reads metadata from an image file. diff --git a/scripts/iib_setup.py b/scripts/iib_setup.py index 0957b8c..637c129 100644 --- a/scripts/iib_setup.py +++ b/scripts/iib_setup.py @@ -4,7 +4,6 @@ from scripts.iib.tool import locale from scripts.iib.tool import read_info_from_image from PIL import Image from scripts.iib.logger import logger -from typing import Any from fastapi import FastAPI import gradio as gr