diff --git a/app.py b/app.py index 3ecbfaf..0509187 100644 --- a/app.py +++ b/app.py @@ -3,14 +3,19 @@ from fastapi import FastAPI from fastapi.responses import FileResponse import uvicorn import os -from scripts.iib.api import infinite_image_browsing_api +from scripts.iib.api import infinite_image_browsing_api, index_html_path from scripts.iib.tool import cwd, 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 tag = "\033[31m[warn]\033[0m" +defalut_port = 8000 +defalut_host = "127.0.0.1" def normalize_paths(paths: List[str]): """ @@ -72,34 +77,152 @@ def update_image_index(sd_webui_config: str): print("update image index completed. ✨") -if __name__ == "__main__": +class AppUltils(object): + + def __init__(self, + sd_webui_config: Optional[str] = None, + use_update_image_index: bool = False, # update_image_index会与外层函数重名 + extra_paths: List[str] = [], + ): + """ + 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=[] + """ + self.sd_webui_config = sd_webui_config + self.use_update_image_index = use_update_image_index + self.extra_paths = extra_paths + + def set_params(self, *args, **kwargs) -> None: + """ 改变参数，与__init__的行为一致 """ + self.__init__(*args, **kwargs) + + @staticmethod + def async_run(app: FastAPI, port: int=defalut_port) -> Coroutine: + """ + 用于从异步运行的FastAPI，在jupyter notebook环境中非常有用 + + app为要启动的FastAPI实例 + port为要启动的端口 + + 返回协程uvicorn.Server().serve() + """ + # 不建议改成async def，并且用await替换return + # 因为这样会失去对server.serve()的控制 + config = uvicorn.Config(app, host=defalut_host, port=port) + server = uvicorn.Server(config) + return server.serve() + + def wrap_app(self, + app: FastAPI, + ) -> None: + """ + 为传递的app挂载上infinite_image_browsing后端 + """ + sd_webui_config = self.sd_webui_config + use_update_image_index = self.use_update_image_index + extra_paths = self.extra_paths + + if sd_webui_config: + sd_webui_paths_check(sd_webui_config) + if use_update_image_index: + update_image_index(sd_webui_config) + paths_check(extra_paths) + + infinite_image_browsing_api(app, + sd_webui_config=sd_webui_config, + extra_paths_cli=normalize_paths(extra_paths), + ) + + def get_root_browser_app(self) -> FastAPI: + """ + 获取首页挂载在"/"上的infinite_image_browsing FastAPI实例 + """ + app = FastAPI() + + # 用于在首页显示 + @app.get("/") + def index(): + return FileResponse(index_html_path) + + self.wrap_app(app) + return app + + +def setup_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Process some integers.") - parser.add_argument("--port", type=int, help="the port to use", default=8000) - parser.add_argument("--sd_webui_config", help="the path to the config file") + parser.add_argument("--port", type=int, help="The port to use", default=defalut_port) + parser.add_argument("--sd_webui_config", type=str, default=None, help="The path to the config file") + # 弃用警告，因为update_image_index会与处理函数同名 + # 所以class AppUltils(object)采用的是use_update_image_index + # 有必要保证两者的一致性 parser.add_argument( - "--update_image_index", action="store_true", help="update the image index" + "--update_image_index", action="store_true", help="Update the image index. Will be deprecated. Recommend to use --use_update_image_index" ) parser.add_argument( - "--extra_paths", nargs="+", help="extra paths to use, will be added to Quick Move.", default=[] + "--use_update_image_index", action="store_true", help="Update the image index" ) - args = parser.parse_args() - sd_webui_config = args.sd_webui_config - if sd_webui_config: - sd_webui_paths_check(sd_webui_config) - if args.update_image_index: - update_image_index(sd_webui_config) - paths_check(args.extra_paths) - - app = FastAPI() - - @app.get("/") - def index(): - return FileResponse(os.path.join(cwd, "vue/dist/index.html")) - - infinite_image_browsing_api( - None, - app, - sd_webui_config=sd_webui_config, - extra_paths_cli=normalize_paths(args.extra_paths), + parser.add_argument( + "--extra_paths", nargs="+", help="Extra paths to use, will be added to Quick Move.", default=[] ) - uvicorn.run(app, host="127.0.0.1", port=args.port) + return parser + + +def check_cmd_param(*args, **kwargs) -> argparse.Namespace: + """ + parser.parse_known_args()的装饰器，用于报告非法参数 + 无任何传递将从sys.argv中解析参数 + """ + cmd_param, unkonwn = parser.parse_known_args(*args, **kwargs) + if unkonwn: + logging.warning(f"{tag} Illegal args will be ingored: {unkonwn}") + if cmd_param.update_image_index: + # 弃用警告，因为update_image_index会与处理函数同名 + # 所以class AppUltils(object)采用的是use_update_image_index + # 有必要保证两者的一致性 + logging.warning(f"{tag} --update_image_index will be deprecated, please use --use_update_image_index") + cmd_param.use_update_image_index = True + return cmd_param + + +def launch_app(port, *args, **kwargs) -> None: + """ + 同步函数 + + 所传入的要求参数都会传递给AppUltils() + 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=[] + + """ + app_ultils = AppUltils(*args, **kwargs) + app = app_ultils.get_root_browser_app() + uvicorn.run(app, host=defalut_host, port=port) + + +async def async_launch_app(port, *args, **kwargs) -> None: + """ + 协程函数 + + 所传入的要求参数都会传递给AppUltils() + 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=[] + + """ + app_ultils = AppUltils(*args, **kwargs) + app = app_ultils.get_root_browser_app() + await app_ultils.async_run(app, port=port) + + +parser = setup_parser() + +if __name__ == "__main__": + cmd_param = check_cmd_param() + params_dict = dict(sd_webui_config = cmd_param.sd_webui_config, + use_update_image_index = cmd_param.use_update_image_index, + extra_paths = cmd_param.extra_paths, + ) + app_ultils = AppUltils(**params_dict) + app = app_ultils.get_root_browser_app() + uvicorn.run(app, host=defalut_host, port=cmd_param.port) diff --git a/scripts/iib/api.py b/scripts/iib/api.py index d7ce73b..1bc983d 100644 --- a/scripts/iib/api.py +++ b/scripts/iib/api.py @@ -38,6 +38,8 @@ from scripts.iib.db.update_image_data import update_image_data from scripts.iib.logger import logger +index_html_path = os.path.join(cwd, "vue/dist/index.html") # 在app.py也被使用 + send_img_path = {"value": ""} mem = { @@ -60,7 +62,7 @@ async def get_token(request: Request): raise HTTPException(status_code=401, detail="Unauthorized") -def infinite_image_browsing_api(_: Any, app: FastAPI, **kwargs): +def infinite_image_browsing_api(app: FastAPI, **kwargs): pre = "/infinite_image_browsing" @@ -361,7 +363,7 @@ def infinite_image_browsing_api(_: Any, app: FastAPI, **kwargs): @app.get(pre) def index_bd(): - return FileResponse(os.path.join(cwd, "vue/dist/index.html")) + return FileResponse(index_html_path) class OpenFolderReq(BaseModel): path: str diff --git a/scripts/iib_setup.py b/scripts/iib_setup.py index b29640c..19dfc4f 100644 --- a/scripts/iib_setup.py +++ b/scripts/iib_setup.py @@ -4,7 +4,10 @@ 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 """ api函数声明和启动分离方便另外一边被外部调用 @@ -12,8 +15,7 @@ api函数声明和启动分离方便另外一边被外部调用 def on_ui_tabs(): - import gradio as gr - + with gr.Blocks(analytics_enabled=False) as view: with gr.Row(): with gr.Column(): @@ -67,6 +69,10 @@ def on_ui_tabs(): ), ) +def on_app_started(_: gr.Blocks, app: FastAPI) -> None: + # 第一个参数是SD-WebUI传进来的gr.Blocks，但是不需要使用 + infinite_image_browsing_api(app) + script_callbacks.on_ui_tabs(on_ui_tabs) -script_callbacks.on_app_started(infinite_image_browsing_api) +script_callbacks.on_app_started(on_app_started)