add daemon system and began working on ui
This commit is contained in:
parent
591c257151
commit
4dfd9faa39
140
main.py
140
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')
|
||||
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():
|
||||
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)
|
||||
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()
|
||||
|
@ -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
|
||||
|
@ -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"])
|
@ -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"])
|
||||
|
@ -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'):
|
||||
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}"
|
||||
|
||||
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)
|
||||
|
||||
|
@ -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"])
|
@ -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"])
|
@ -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)
|
@ -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"])
|
||||
|
43
ui.py
Normal file
43
ui.py
Normal file
@ -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()]
|
||||
)
|
Loading…
Reference in New Issue
Block a user