From 4dfd9faa391d18743fd5f7210fb9b429821bb16b Mon Sep 17 00:00:00 2001 From: Mathieu Broillet Date: Wed, 28 Aug 2024 20:05:46 +0200 Subject: [PATCH] add daemon system and began working on ui --- main.py | 142 ++++++++++++++++-- services/__init__.py | 7 +- ...moval-dis.py => background_removal_dis.py} | 2 +- services/comfyui.py | 2 +- services/services.py | 124 ++++++++++----- ...sion-forge.py => stablediffusion_forge.py} | 2 +- ...sion-webui.py => stablediffusion_webui.py} | 4 +- services/{txtgen.py => textgen.py} | 12 +- services/xtts.py | 2 +- ui.py | 43 ++++++ 10 files changed, 279 insertions(+), 61 deletions(-) rename services/{background-removal-dis.py => background_removal_dis.py} (98%) rename services/{stablediffusion-forge.py => stablediffusion_forge.py} (97%) rename services/{stablediffusion-webui.py => stablediffusion_webui.py} (93%) rename services/{txtgen.py => textgen.py} (79%) create mode 100644 ui.py diff --git a/main.py b/main.py index 36f4bb2..d3ae904 100644 --- a/main.py +++ b/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(): - handler_with_formatter = logging.StreamHandler(stream=sys.stdout) - handler_with_formatter.setFormatter(logging.Formatter('[%(levelname)s] : %(message)s')) - logger.addHandler(handler_with_formatter) -logger.setLevel(LEVEL) +config = None + + +class Config: + 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 - - test = TextGeneration().start() + run_interactive_cmd_ui() diff --git a/services/__init__.py b/services/__init__.py index f7af8e2..01e3e8a 100644 --- a/services/__init__.py +++ b/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 \ No newline at end of file +from services.stablediffusion_forge import StableDiffusionForge +from services.stablediffusion_webui import StableDiffusionWebUI +from services.textgen import TextGeneration +from services.xtts import XTTS diff --git a/services/background-removal-dis.py b/services/background_removal_dis.py similarity index 98% rename from services/background-removal-dis.py rename to services/background_removal_dis.py index d3a162c..21e6666 100644 --- a/services/background-removal-dis.py +++ b/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"]) diff --git a/services/comfyui.py b/services/comfyui.py index 28242e3..a8b694c 100644 --- a/services/comfyui.py +++ b/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"]) diff --git a/services/services.py b/services/services.py index 355e010..1a864ee 100644 --- a/services/services.py +++ b/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: - logger.warning("Found files from a previous/borked/crashed installation, cleaning up...") - self.bash(f"rm -rf {self.path}") - self.create_dir('') + 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,61 +67,94 @@ 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'): - self.pip(f"install -r {filename}") + def install_requirements(self, filename: str = 'requirements.txt', env=[]): + 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): - self.bash(f"{' '.join(env)} {self.path}/venv/bin/pip {cmd}", current_dir) + def pip(self, cmd: str, env=[], args=[], current_dir: str = None): + 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): - self.bash(f"{' '.join(env)} {self.path}/venv/bin/python {cmd}", current_dir) + 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, 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}" - logger.debug(f"Running command: {cmd}") + if daemon: + # Check if previous run process is saved + if get_config().has(f"{self.name}-pid"): - if LEVEL == logging.DEBUG: + # 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) - process.wait() - if process.returncode != 0: - raise Exception(f"Failed to run command: {cmd}") - + get_config().set(f"{self.name}-pid", process.pid + 1) + return else: - process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - out, err = process.communicate() + logger.debug(f"Running command: {cmd}") - if process.returncode != 0: - logger.fatal(f"Failed to run command: {cmd}") - logger.fatal(f"Error: {err.decode('utf-8')}") - logger.fatal(f"Output: {out.decode('utf-8')}") - raise Exception(f"Failed to run command: {cmd}") + 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() + + if process.returncode != 0: + logger.fatal(f"Failed to run command: {cmd}") + logger.fatal(f"Error: {err.decode('utf-8')}") + logger.fatal(f"Output: {out.decode('utf-8')}") + raise Exception(f"Failed to run command: {cmd}") # 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) diff --git a/services/stablediffusion-forge.py b/services/stablediffusion_forge.py similarity index 97% rename from services/stablediffusion-forge.py rename to services/stablediffusion_forge.py index 47e5f73..9176a78 100644 --- a/services/stablediffusion-forge.py +++ b/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"]) diff --git a/services/stablediffusion-webui.py b/services/stablediffusion_webui.py similarity index 93% rename from services/stablediffusion-webui.py rename to services/stablediffusion_webui.py index b3638da..67e235f 100644 --- a/services/stablediffusion-webui.py +++ b/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"]) diff --git a/services/txtgen.py b/services/textgen.py similarity index 79% rename from services/txtgen.py rename to services/textgen.py index 5915167..7bfd3b2 100644 --- a/services/txtgen.py +++ b/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( - "install auto-gptq --no-build-isolation --extra-index-url https://huggingface.github.io/autogptq-index/whl/rocm573/") + self.pip_install("auto-gptq", args=["--no-build-isolation", "--extra-index-url", + "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) diff --git a/services/xtts.py b/services/xtts.py index 960e9a7..0121deb 100644 --- a/services/xtts.py +++ b/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"]) diff --git a/ui.py b/ui.py new file mode 100644 index 0000000..b500e97 --- /dev/null +++ b/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()] +)