refactor

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

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 == "":

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

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"):

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