mathieu/ai-suite-rocm-local
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases 1 Wiki Activity

add daemon system and began working on ui

Browse Source
This commit is contained in:
Mathieu Broillet 2024-08-28 20:05:46 +02:00
parent 591c257151
commit 4dfd9faa39
Signed by: mathieu
GPG Key ID: A08E484FE95074C1
10 changed files with 279 additions and 61 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

142
main.py
View File

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

7
services/__init__.py
View File

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

2
services/background-removal-dis.py → services/background_removal_dis.py
View File

@ -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"])

2
services/comfyui.py
View File

@ -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"])

124
services/services.py
View File

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

2
services/stablediffusion-forge.py → services/stablediffusion_forge.py
View File

@ -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"])

4
services/stablediffusion-webui.py → services/stablediffusion_webui.py
View File

@ -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
services/txtgen.py → services/textgen.py
View File

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

2
services/xtts.py
View File

@ -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
View 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()]
)