From 752c557dd4ded87f5a0beb4b2213206b42f9246d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A7=8B=E8=91=89=E6=9D=8F?= Date: Tue, 28 Nov 2023 10:08:48 +0800 Subject: [PATCH] refactor --- gui.py | 6 ++++-- mikazuki/app/api.py | 8 +++---- mikazuki/{ => app}/models.py | 0 mikazuki/app/tk_window.py | 34 ++++++++++++++++++++++++++++++ mikazuki/launch_utils.py | 7 ++++-- mikazuki/utils.py | 41 ++++++------------------------------ 6 files changed, 53 insertions(+), 43 deletions(-) rename mikazuki/{ => app}/models.py (100%) create mode 100644 mikazuki/app/tk_window.py diff --git a/gui.py b/gui.py index b93df91..8d6444b 100644 --- a/gui.py +++ b/gui.py @@ -10,7 +10,7 @@ from mikazuki.launch_utils import (check_dirs, prepare_submodules, remove_warnings, setup_windows_bitsandbytes, smart_pip_mirror, validate_requirements) from mikazuki.log import log -from mikazuki.utils import check_run +from mikazuki.utils import check_run, base_dir_path parser = argparse.ArgumentParser(description="GUI for stable diffusion training") parser.add_argument("--host", type=str, default="127.0.0.1") @@ -34,7 +34,7 @@ def run_tag_editor(): log.info("Starting tageditor...") cmd = [ sys.executable, - "mikazuki/dataset-tag-editor/scripts/launch.py", + base_dir_path() / "mikazuki/dataset-tag-editor/scripts/launch.py", "--port", "28001", "--shadow-gradio-output", "--root-path", "/proxy/tageditor" @@ -46,6 +46,8 @@ def run_tag_editor(): if __name__ == "__main__": args, _ = parser.parse_known_args() + log.info("Starting SD-Trainer Mikazuki GUI...") + log.info(f"Base directory: {base_dir_path()}, Working directory: {os.getcwd()}") log.info(f'{platform.system()} Python {platform.python_version()} {sys.executable}') remove_warnings() diff --git a/mikazuki/app/api.py b/mikazuki/app/api.py index a7c4811..ed6bfef 100644 --- a/mikazuki/app/api.py +++ b/mikazuki/app/api.py @@ -9,9 +9,9 @@ from fastapi import APIRouter, BackgroundTasks, Request from starlette.requests import Request import mikazuki.process as process -import mikazuki.utils as utils +from mikazuki.app.models import TaggerInterrogateRequest +from mikazuki.app.tk_window import open_directory_selector, open_file_selector from mikazuki.log import log -from mikazuki.models import TaggerInterrogateRequest from mikazuki.tagger.interrogator import (available_interrogators, on_interrogate) from mikazuki.tasks import tm @@ -125,10 +125,10 @@ async def run_interrogate(req: TaggerInterrogateRequest, background_tasks: Backg @router.get("/pick_file") async def pick_file(picker_type: str): if picker_type == "folder": - coro = asyncio.to_thread(utils.open_directory_selector, os.getcwd()) + coro = asyncio.to_thread(open_directory_selector, os.getcwd()) elif picker_type == "modelfile": file_types = [("checkpoints", "*.safetensors;*.ckpt;*.pt"), ("all files", "*.*")] - coro = asyncio.to_thread(utils.open_file_selector, os.getcwd(), "Select file", file_types) + coro = asyncio.to_thread(open_file_selector, os.getcwd(), "Select file", file_types) result = await coro if result == "": diff --git a/mikazuki/models.py b/mikazuki/app/models.py similarity index 100% rename from mikazuki/models.py rename to mikazuki/app/models.py diff --git a/mikazuki/app/tk_window.py b/mikazuki/app/tk_window.py new file mode 100644 index 0000000..9ddbd5b --- /dev/null +++ b/mikazuki/app/tk_window.py @@ -0,0 +1,34 @@ +import tkinter +from tkinter.filedialog import askdirectory, askopenfilename + + +def tk_window(): + window = tkinter.Tk() + window.wm_attributes('-topmost', 1) + window.withdraw() + + +def open_file_selector( + initialdir, + title, + filetypes) -> str: + try: + tk_window() + filename = askopenfilename( + initialdir=initialdir, title=title, + filetypes=filetypes + ) + return filename + except: + return "" + + +def open_directory_selector(initialdir) -> str: + try: + tk_window() + directory = askdirectory( + initialdir=initialdir + ) + return directory + except: + return "" diff --git a/mikazuki/launch_utils.py b/mikazuki/launch_utils.py index d6590a9..0daa1dc 100644 --- a/mikazuki/launch_utils.py +++ b/mikazuki/launch_utils.py @@ -10,7 +10,7 @@ from typing import List import pkg_resources from mikazuki.log import log -from mikazuki.utils import run_pip +from mikazuki.utils import run_pip, base_dir_path def smart_pip_mirror(): @@ -28,7 +28,10 @@ def find_windows_git(): def prepare_submodules(): - if not os.path.exists("./frontend/dist") or not os.path.exists("./mikazuki/dataset-tag-editor"): + frontend_path = base_dir_path() / "frontend" / "dist" + tag_editor_path = base_dir_path() / "mikazuki" / "dataset-tag-editor" / "scripts" + + if not os.path.exists(frontend_path) or not os.path.exists(tag_editor_path): log.info("submodule not found, try clone...") log.info("checking git installation...") if not shutil.which("git"): diff --git a/mikazuki/utils.py b/mikazuki/utils.py index 13e3bd7..e0d0c0c 100644 --- a/mikazuki/utils.py +++ b/mikazuki/utils.py @@ -1,15 +1,14 @@ import glob import importlib.util import os -import subprocess -import sys import re import shutil -import tkinter -from tkinter.filedialog import askopenfilename, askdirectory +import subprocess +import sys +from pathlib import Path from typing import Optional -from mikazuki.log import log +from mikazuki.log import log python_bin = sys.executable @@ -156,33 +155,5 @@ def check_run(file: str) -> bool: return result.returncode == 0 -def tk_window(): - window = tkinter.Tk() - window.wm_attributes('-topmost', 1) - window.withdraw() - - -def open_file_selector( - initialdir, - title, - filetypes) -> str: - try: - tk_window() - filename = askopenfilename( - initialdir=initialdir, title=title, - filetypes=filetypes - ) - return filename - except: - return "" - - -def open_directory_selector(initialdir) -> str: - try: - tk_window() - directory = askdirectory( - initialdir=initialdir - ) - return directory - except: - return "" +def base_dir_path(): + return Path(__file__).parents[1].absolute()