Add option to use relative paths based on sd_webui_config file path

pull/276/head
zanllp 2023-06-28 08:39:26 +08:00
parent 36e781677d
commit 09f37c5c86
3 changed files with 72 additions and 52 deletions

107
app.py
View File

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

View File

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

View File

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