add daemon system and began working on ui

main.py

@@ -1,23 +1,145 @@
+import json
 import logging
 import os
+import subprocess
 import sys
 
+import ui
+
 PYTHON_EXEC = 'python3.10'
 PATH = os.path.dirname(os.path.abspath(__file__))
 ROCM_VERSION = "6.1.2"
 
-# Set up logging
-LEVEL = logging.DEBUG
 logger = logging.getLogger('ai-suite-rocm')
-if not logger.hasHandlers():
+config = None
-    handler_with_formatter = logging.StreamHandler(stream=sys.stdout)
+
-    handler_with_formatter.setFormatter(logging.Formatter('[%(levelname)s] : %(message)s'))
+
-    logger.addHandler(handler_with_formatter)
+class Config:
-logger.setLevel(LEVEL)
+    data = {}
+
+    def __init__(self):
+        self.file = os.path.join(os.path.expanduser("~"), ".config", "ai-suite-rocm", "config.json")
+
+        self.create()
+        self.read()
+
+    def create(self):
+        if not os.path.exists(self.file):
+            os.makedirs(os.path.dirname(self.file), exist_ok=True)
+            with open(self.file, "w") as f:
+                f.write("{}")
+
+            logger.info(f"Created config file at {self.file}")
+
+    def read(self):
+        with open(self.file, "r") as f:
+            self.data = json.load(f)
+
+    def write(self):
+        with open(self.file, "w") as f:
+            json.dump(self.data, f)
+
+    def get(self, key: str):
+        return self.data.get(key)
+
+    def set(self, key: str, value):
+        self.data[key] = value
+        self.write()
+
+    def has(self, key: str):
+        return key in self.data
+
+    def remove(self, key: str):
+        self.data.pop(key)
+        self.write()
+
+    def clear(self):
+        self.data = {}
+        self.write()
+
+
+def setup_logger(level: logger.level = logging.INFO):
+    global logger
+    if not logger.hasHandlers():
+        logger.setLevel(level)
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setFormatter(logging.Formatter('[%(levelname)s] : %(message)s'))
+        logger.addHandler(handler)
+
+
+def setup_config():
+    global config
+    config = Config()
+
+
+def get_config():
+    global config
+    return config
+
+
+def run_command(command: str, exit_on_error: bool = True):
+    logger.debug(f"Running command: {command}")
+    process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    out, err = process.communicate()
+
+    if process.returncode != 0:
+        logger.fatal(f"Failed to run command: {command}")
+        raise Exception(f"Failed to run command: {command}")
+
+    return out, err, process.returncode
+
+
+def check_for_build_essentials():
+    logger.debug("Checking for build essentials...")
+    debian = os.path.exists('/etc/debian_version')
+    fedora = os.path.exists('/etc/fedora-release')
+
+    if debian:
+        # TODO: check if these work for debian users
+        check_gcc = run_command("dpkg -l | grep build-essential &>/dev/null", exit_on_error=False)[2] == 0
+        check_python = run_command("dpkg -l | grep python3.10-dev &>/dev/null", exit_on_error=False)[2] == 0
+
+        if not check_gcc or not check_python:
+            raise UserWarning(
+                "The packages build-essential and python3.10-dev are required for this script to run. Please install them. See the README for more information.")
+    elif fedora:
+        check_gcc = run_command("rpm -q gcc &>/dev/null", exit_on_error=False)[2] == 0
+        check_python = run_command("rpm -q python3.10-devel &>/dev/null", exit_on_error=False)[2] == 0
+
+        if not check_gcc or not check_python:
+            raise UserWarning(
+                "The package python3.10-devel and the Development Tools group are required for this script to run. Please install them. See the README for more information.")
+    else:
+        logger.warning(
+            "Unsupported OS detected. Please ensure you have the following packages installed or their equivalent: build-essential, python3.10-dev")
+
+
+def run_interactive_cmd_ui():
+    while True:
+        choice = ui.start.ask()
+
+        match choice:
+            case "Start services":
+                services = ui.start_services.ask()
+                for service in services:
+                    logger.info(f"Starting service: {service}")
+                    pass
+                pass
+            case "Stop services":
+                pass
+            case "Install/update services":
+                pass
+            case "Uninstall services":
+                pass
+            case "Exit":
+                print("Exiting...")
+                exit(0)
 
 if __name__ == '__main__':
+    setup_logger(logging.DEBUG)
+    setup_config()
+
     logger.info("Starting AI Suite for ROCM")
+    check_for_build_essentials()
 
-    from services import TextGeneration
+    run_interactive_cmd_ui()
-
-    test = TextGeneration().start()

services/__init__.py

@@ -1,2 +1,7 @@
+from services.background_removal_dis import BGRemovalDIS
+from services.comfyui import ComfyUI
 from services.services import Stack
-from services.txtgen import TextGeneration
+from services.stablediffusion_forge import StableDiffusionForge
+from services.stablediffusion_webui import StableDiffusionWebUI
+from services.textgen import TextGeneration
+from services.xtts import XTTS

services/background_removal_dis.py

@@ -28,7 +28,7 @@ class BGRemovalDIS(Stack):
 
         super().install()
 
-    def start(self):
+    def _launch(self):
         args = ["--port", str(self.port)]
         self.python(f"app.py {' '.join(args)}", current_dir="webui",
                     env=["TORCH_BLAS_PREFER_HIPBLASLT=0"])

services/comfyui.py

@@ -28,7 +28,7 @@ class ComfyUI(Stack):
 
         super().install()
 
-    def start(self):
+    def _launch(self):
         args = ["--port", str(self.port)]
         self.python(f"main.py {' '.join(args)}", current_dir="webui",
                     env=["TORCH_BLAS_PREFER_HIPBLASLT=0"])

services/services.py

@@ -2,25 +2,23 @@ import logging
 import os
 import subprocess
 
+import psutil
+
+import main
 import utils
-from main import PATH, PYTHON_EXEC, logger, LEVEL
+from main import PYTHON_EXEC, logger, get_config
 
 
 class Stack:
     def __init__(self, name: str, path: str, port: int, url: str):
         self.name = name
-        self.path = os.path.join(PATH, path)
+        self.path = os.path.join(os.path.expanduser("~"), ".ai-suite-rocm", path)
+
         self.url = url
         self.port = port
 
         self.process = None
 
-        if self.is_installed():
-            self.update()
-        else:
-            self.check_for_broken_install()
-            self.create_venv()
-            self.install()
 
     def install(self):
         self.create_file('.installed', 'true')
@@ -30,10 +28,14 @@ class Stack:
         return self.file_exists('.installed')
 
     def check_for_broken_install(self):
-        if not self.is_installed() and len(os.listdir(self.path)) > 0:
+        if not self.is_installed():
+            if os.path.exists(self.path):
+                if len(os.listdir(self.path)) > 0:
                     logger.warning("Found files from a previous/borked/crashed installation, cleaning up...")
                     self.bash(f"rm -rf {self.path}")
                     self.create_dir('')
+            else:
+                self.create_dir('')
 
     def update(self, folder: str = 'webui'):
         if self.dir_exists(folder):
@@ -46,6 +48,16 @@ class Stack:
         self.bash(f"rm -rf {self.path}")
 
     def start(self):
+        if self.is_installed():
+            self.update()
+        else:
+            self.check_for_broken_install()
+            self.create_venv()
+            self.install()
+
+        self._launch()
+
+    def _launch(self):
         pass
 
     def stop(self):
@@ -55,49 +67,81 @@ class Stack:
         self.stop()
         self.start()
 
-    def status(self):
+    def status(self) -> bool:
         pass
 
     # Python/Bash utils
     def create_venv(self):
         venv_path = os.path.join(self.path, 'venv')
         if not self.has_venv():
-            logger.debug(f"Creating venv for {self.name}")
+            logger.info(f"Creating venv for {self.name}")
             self.bash(f"{PYTHON_EXEC} -m venv {venv_path} --system-site-packages")
-            self.pip("install --upgrade pip")
+            self.pip_install("pip")
         else:
             logger.debug(f"Venv already exists for {self.name}")
 
     def has_venv(self) -> bool:
         return self.dir_exists('venv')
 
-    def pip_install(self, package: str | list, no_deps: bool = False):
+    def pip_install(self, package: str | list, no_deps: bool = False, env=[], args=[]):
+        if no_deps:
+            args.append("--no-deps")
+
         if isinstance(package, list):
             for p in package:
-                self.pip(f"install -U {p} {'--no-deps' if no_deps else ''}")
+                logger.info(f"Installing {p}")
+                self.pip(f"install -U {p}", env=env, args=args)
         else:
-            self.pip(f"install -U {package} {'--no-deps' if no_deps else ''}")
+            logger.info(f"Installing {package}")
+            self.pip(f"install -U {package}", env=env, args=args)
 
-    def install_requirements(self, filename: str = 'requirements.txt'):
+    def install_requirements(self, filename: str = 'requirements.txt', env=[]):
-        self.pip(f"install -r {filename}")
+        logger.info(f"Installing requirements for {self.name} ({filename})")
+        self.pip(f"install -r {filename}", env=env)
 
-    def pip(self, cmd: str, env=[], current_dir: str = None):
+    def pip(self, cmd: str, env=[], args=[], current_dir: str = None):
-        self.bash(f"{' '.join(env)} {self.path}/venv/bin/pip {cmd}", current_dir)
+        self.bash(f"{' '.join(env)} {self.path}/venv/bin/pip {cmd} {' '.join(args)}", current_dir)
 
-    def python(self, cmd: str, env=[], current_dir: str = None):
+    def python(self, cmd: str, env=[], current_dir: str = None, daemon: bool = False):
-        self.bash(f"{' '.join(env)} {self.path}/venv/bin/python {cmd}", current_dir)
+        self.bash(f"{' '.join(env)} {self.path}/venv/bin/python {cmd}", current_dir, daemon)
 
-    def bash(self, cmd: str, current_dir: str = None):
+    def bash(self, cmd: str, current_dir: str = None, daemon: bool = False):
         cmd = f"cd {self.path if current_dir is None else os.path.join(self.path, current_dir)} && {cmd}"
 
+        if daemon:
+            # Check if previous run process is saved
+            if get_config().has(f"{self.name}-pid"):
+
+                # Check if PID still running
+                if psutil.pid_exists(main.config.get(f"{self.name}-pid")):
+                    choice = input(f"{self.name} is already running, do you want to restart it? (y/n): ")
+
+                    if choice.lower() == 'y':
+                        pid = main.config.get(f"{self.name}-pid")
+                        logger.debug(f"Killing previous daemon with PID: {pid}")
+                        psutil.Process(pid).kill()
+                    else:
+                        # TODO: attach to subprocess?
+                        return
+                else:
+                    logger.warning(
+                        f"Previous PID found for {self.name} but process is not running, continuing as stopped...")
+            else:
+                logger.debug(f"No previous PID found for {self.name}, continuing as stopped...")
+
+            logger.debug(f"Starting {self.name} as daemon with command: {cmd}")
+            cmd = f"{cmd} &"
+            process = subprocess.Popen(cmd, shell=True)
+            get_config().set(f"{self.name}-pid", process.pid + 1)
+            return
+        else:
             logger.debug(f"Running command: {cmd}")
 
-        if LEVEL == logging.DEBUG:
+            if logger.level == logging.DEBUG:
                 process = subprocess.Popen(cmd, shell=True)
                 process.wait()
                 if process.returncode != 0:
                     raise Exception(f"Failed to run command: {cmd}")
-
             else:
                 process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                 out, err = process.communicate()
@@ -110,6 +154,7 @@ class Stack:
 
     # Git utils
     def git_clone(self, url: str, branch: str = None, dest: str = None):
+        logger.info(f"Cloning {url}")
         self.bash(f"git clone {f"-b {branch}" if branch is not None else ''} {url} {'' if dest is None else dest}")
 
     def git_pull(self, repo_folder: str, force: bool = False):
@@ -128,6 +173,9 @@ class Stack:
             f.write(content)
 
     def create_dir(self, name):
+        if name == '':
+            logger.info(f"Creating directory for {self.name}")
+
         logger.debug(f"Creating directory {name}")
         os.makedirs(os.path.join(self.path, name), exist_ok=True)
 

services/stablediffusion_forge.py

@@ -21,7 +21,7 @@ class StableDiffusionForge(Stack):
 
         super().install()
 
-    def start(self):
+    def _launch(self):
         args = ["--listen", "--enable-insecure-extension-access", "--port", str(self.port)]
         self.python(f"launch.py {' '.join(args)}", current_dir="webui",
                     env=["TORCH_BLAS_PREFER_HIPBLASLT=0"])

services/stablediffusion_webui.py

@@ -1,7 +1,7 @@
 from services import Stack
 
 
-class StableDiffusionForge(Stack):
+class StableDiffusionWebUI(Stack):
     def __init__(self):
         super().__init__(
             'StableDiffusion WebUI',
@@ -21,7 +21,7 @@ class StableDiffusionForge(Stack):
 
         super().install()
 
-    def start(self):
+    def _launch(self):
         args = ["--listen", "--enable-insecure-extension-access", "--port", str(self.port)]
         self.python(f"launch.py {' '.join(args)}", current_dir="webui",
                     env=["TORCH_BLAS_PREFER_HIPBLASLT=0"])

services/textgen.py

@@ -12,10 +12,10 @@ class TextGeneration(Stack):
 
     def install(self):
         # Install LlamaCpp from prebuilt
-        self.pip("install llama-cpp-python", ["CMAKE_ARGS=\"-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS\""])  # cpu
+        self.pip_install("llama-cpp-python", env=["CMAKE_ARGS=\"-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS\""])  # cpu
 
         # Install LlamaCpp for ROCM from source
-        # self.pip("install llama-cpp-python", ["CMAKE_ARGS=\"-DGGML_HIPBLAS=on\" FORCE_CMAKE=1"]) # manual gpu (only works if whole rocm suite installed)
+        # self.pip_install("llama-cpp-python", env=["CMAKE_ARGS=\"-DGGML_HIPBLAS=on\" FORCE_CMAKE=1"]) # manual gpu (only works if whole rocm suite installed)
         # self.install_from_prebuilt("llama_cpp_python") # gpu (only works if whole rocm suite installed)
 
         # Install Triton for ROCM from prebuilt
@@ -45,12 +45,12 @@ class TextGeneration(Stack):
         self.pip_install(
             "https://github.com/turboderp/exllamav2/releases/download/v0.1.9/exllamav2-0.1.9+rocm6.1.torch2.4.0-cp310-cp310-linux_x86_64.whl")
         self.install_from_prebuilt("bitsandbytes")
-        self.pip(
+        self.pip_install("auto-gptq", args=["--no-build-isolation", "--extra-index-url",
-            "install auto-gptq --no-build-isolation --extra-index-url https://huggingface.github.io/autogptq-index/whl/rocm573/")
+                                            "https://huggingface.github.io/autogptq-index/whl/rocm573/"])
 
         super().install()
 
-    def start(self):
+    def _launch(self):
         args = ["--listen", "--listen-port", str(self.port)]
         self.python(f"server.py {' '.join(args)}", current_dir="webui",
-                    env=["TORCH_BLAS_PREFER_HIPBLASLT=0"])
+                    env=["TORCH_BLAS_PREFER_HIPBLASLT=0"], daemon=True)

services/xtts.py

@@ -32,7 +32,7 @@ class XTTS(Stack):
 
         super().install()
 
-    def start(self):
+    def _launch(self):
         args = ["--host", "0.0.0.0", "--port", str(self.port)]
         self.python(f"server.py {' '.join(args)}", current_dir="webui",
                     env=["TORCH_BLAS_PREFER_HIPBLASLT=0"])

ui.py

@@ -0,0 +1,43 @@
+import questionary
+from questionary import Choice
+
+from services import BGRemovalDIS, ComfyUI, StableDiffusionWebUI, StableDiffusionForge, TextGeneration, XTTS
+
+services = {
+    "Background Removal (DIS)": BGRemovalDIS,
+    "ComfyUI": ComfyUI,
+    "StableDiffusion (AUTOMATIC1111)": StableDiffusionWebUI,
+    "StableDiffusion Forge": StableDiffusionForge,
+    "TextGeneration (oobabooga)": TextGeneration,
+    "XTTS": XTTS
+}
+
+start = questionary.select(
+    "Choose an option:",
+    choices=[
+        Choice("Start services"),
+        Choice("Stop services"),
+        Choice("Install/update services"),
+        Choice("Uninstall services"),
+        Choice("Exit")
+    ])
+
+start_services = questionary.checkbox(
+    "Select services to start:",
+    choices=[Choice(service) for service in services.keys()]
+)
+
+stop_services = questionary.checkbox(
+    "Select services to stop:",
+    choices=[Choice(service) for service in services.keys()]
+)
+
+install_service = questionary.checkbox(
+    "Select service to install/update:",
+    choices=[Choice(service) for service in services.keys()]
+)
+
+uninstall_service = questionary.checkbox(
+    "Select service to uninstall:",
+    choices=[Choice(service) for service in services.keys()]
+)