full rewrite of all scripts in python

README.md

@@ -2,6 +2,10 @@
 
 This is a simple project to make hosting multiple AI tools easily on Linux with AMD GPUs using ROCM locally (without docker).
 
+> [!WARNING]
+> Currently rewriting this project to be more modular and easier to use. This is a work in progress.  
+> **Do not mind the python files or the services dir for now, the scripts aren't affected.**
+
 To use you have to clone the repo run the install script for the service you want to use.
 
 ```bash

main.py

@@ -0,0 +1,23 @@
+import logging
+import os
+import sys
+
+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)
+
+if __name__ == '__main__':
+    logger.info("Starting AI Suite for ROCM")
+
+    from services import TextGeneration
+
+    test = TextGeneration().start()

services/__init__.py

@@ -0,0 +1,2 @@
+from services.services import Stack
+from services.txtgen import TextGeneration

services/background-removal-dis.py

@@ -0,0 +1,34 @@
+from services import Stack
+
+
+class BGRemovalDIS(Stack):
+    def __init__(self):
+        super().__init__(
+            'BGRemovalDIS',
+            'bg-remove-dis-rocm',
+            5005,
+            'https://huggingface.co/spaces/ECCV2022/dis-background-removal'
+        )
+
+    def install(self):
+        self.git_clone(url=self.url, dest="webui")
+        self.install_requirements("webui/requirements.txt")
+        self.pip_install("gradio")  # gradio is not in requirements.txt for some reason
+
+        self.remove_line_in_file("os.", "webui/app.py")  # remove manual clone of DIS from app.py (done below)
+
+        self.git_clone("https://github.com/xuebinqin/DIS.git", dest="tmp-dis")
+        self.move_all_files_in_dir("tmp-dis/IS-Net", "webui")
+        self.remove_dir("tmp-dis")
+
+        self.create_dir("webui/saved_models")
+        self.move_file_or_dir("webui/isnet.pth", "webui/saved_models/isnet.pth")
+
+        # self.remove_dir("webui/.git") # saves a lot of space due to big repo
+
+        super().install()
+
+    def start(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

@@ -0,0 +1,34 @@
+from services import Stack
+
+
+class ComfyUI(Stack):
+    def __init__(self):
+        super().__init__(
+            'ComfyUI',
+            'comfyui-rocm',
+            5004,
+            'https://github.com/comfyanonymous/ComfyUI.git'
+        )
+
+    def install(self):
+        # Install the webui
+        self.git_clone(url=self.url, dest="webui")
+        self.install_requirements("webui/requirements.txt")
+
+        # Install the manager
+        self.git_clone(url="https://github.com/ltdrdata/ComfyUI-Manager.git", dest="webui/custom_nodes/manager")
+
+        # Add GGUF support
+        self.pip_install(["gguf", "numpy==1.26.4"])
+
+        # Add NF4 support for Flux
+        self.install_from_prebuilt("bitsandbytes")
+        self.git_clone(url="https://github.com/comfyanonymous/ComfyUI_bitsandbytes_NF4.git",
+                       dest="webui/custom_nodes/ComfyUI_bitsandbytes_NF4")
+
+        super().install()
+
+    def start(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

@@ -0,0 +1,168 @@
+import logging
+import os
+import subprocess
+
+import utils
+from main import PATH, PYTHON_EXEC, logger, LEVEL
+
+
+class Stack:
+    def __init__(self, name: str, path: str, port: int, url: str):
+        self.name = name
+        self.path = os.path.join(PATH, 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')
+        logger.info(f"Installed {self.name}")
+
+    def is_installed(self):
+        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('')
+
+    def update(self, folder: str = 'webui'):
+        if self.dir_exists(folder):
+            logger.info(f"Updating {self.name}")
+            self.git_pull(folder)
+        else:
+            logger.warning(f"Could not update {self.name} as {folder} does not exist")
+
+    def uninstall(self):
+        self.bash(f"rm -rf {self.path}")
+
+    def start(self):
+        pass
+
+    def stop(self):
+        pass
+
+    def restart(self):
+        self.stop()
+        self.start()
+
+    def status(self):
+        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}")
+            self.bash(f"{PYTHON_EXEC} -m venv {venv_path} --system-site-packages")
+            self.pip("install --upgrade 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):
+        if isinstance(package, list):
+            for p in package:
+                self.pip(f"install -U {p} {'--no-deps' if no_deps else ''}")
+        else:
+            self.pip(f"install -U {package} {'--no-deps' if no_deps else ''}")
+
+    def install_requirements(self, filename: str = 'requirements.txt'):
+        self.pip(f"install -r {filename}")
+
+    def pip(self, cmd: str, env=[], current_dir: str = None):
+        self.bash(f"{' '.join(env)} {self.path}/venv/bin/pip {cmd}", 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 bash(self, cmd: str, current_dir: str = None):
+        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 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):
+        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):
+        self.bash(f"git reset --hard HEAD {'&& git clean -f -d' if force else ''} && git pull", repo_folder)
+
+    # Prebuilt utils
+    def install_from_prebuilt(self, name):
+        for prebuilt in utils.get_prebuilts():
+            if prebuilt['name'].split("-")[0] == name:
+                self.pip(f"install {prebuilt['browser_download_url']}")
+                return
+
+    # File utils
+    def create_file(self, name, content):
+        with open(os.path.join(self.path, name), 'w') as f:
+            f.write(content)
+
+    def create_dir(self, name):
+        logger.debug(f"Creating directory {name}")
+        os.makedirs(os.path.join(self.path, name), exist_ok=True)
+
+    def remove_file(self, name):
+        logger.debug(f"Removing file {name}")
+        os.remove(os.path.join(self.path, name))
+
+    def remove_dir(self, name):
+        logger.debug(f"Removing directory {name}")
+        os.rmdir(os.path.join(self.path, name))
+
+    def move_file_or_dir(self, src, dest):
+        logger.debug(f"Moving file/dir {src} to {dest}")
+        os.rename(os.path.join(self.path, src), os.path.join(self.path, dest))
+
+    def move_all_files_in_dir(self, src, dest):
+        logger.debug(f"Moving all files in directory {src} to {dest}")
+        for file in os.listdir(os.path.join(self.path, src)):
+            os.rename(os.path.join(self.path, src, file), os.path.join(self.path, dest, file))
+
+    def file_exists(self, name):
+        return os.path.exists(os.path.join(self.path, name))
+
+    def dir_exists(self, name):
+        return os.path.exists(os.path.join(self.path, name))
+
+    def remove_line_in_file(self, contains: str | list, file: str):
+        logger.debug(f"Removing lines containing {contains} in {file}")
+
+        if isinstance(contains, list):
+            for c in contains:
+                self.bash(f"sed -i '/{c}/d' {file}")
+        else:
+            self.bash(f"sed -i '/{contains}/d' {file}")
+
+    def replace_line_in_file(self, match: str, replace: str, file: str):
+        logger.debug(f"Replacing lines containing {match} with {replace} in {file}")
+        self.bash(f"sed -i 's/{match}/{replace}/g' {file}")

services/stablediffusion-forge.py

@@ -0,0 +1,27 @@
+from services import Stack
+
+
+class StableDiffusionForge(Stack):
+    def __init__(self):
+        super().__init__(
+            'StableDiffusion Forge WebUI',
+            'stablediffusion-forge-rocm',
+            5003,
+            'https://github.com/lllyasviel/stable-diffusion-webui-forge'
+        )
+
+    def install(self):
+        # Install the webui
+        self.git_clone(url=self.url, dest="webui")
+
+        self.python("launch.py --skip-torch-cuda-test --exit", current_dir="webui")
+
+        # Add NF4 support for Flux
+        self.install_from_prebuilt("bitsandbytes")
+
+        super().install()
+
+    def start(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

@@ -0,0 +1,27 @@
+from services import Stack
+
+
+class StableDiffusionForge(Stack):
+    def __init__(self):
+        super().__init__(
+            'StableDiffusion WebUI',
+            'stablediffusion-webui-rocm',
+            5002,
+            'https://github.com/AUTOMATIC1111/stable-diffusion-webui'
+        )
+
+    def install(self):
+        # Install the webui
+        self.git_clone(url=self.url, branch="dev", dest="webui")
+
+        self.python("launch.py --skip-torch-cuda-test --exit", current_dir="webui")
+
+        # Add NF4 support for Flux
+        self.install_from_prebuilt("bitsandbytes")
+
+        super().install()
+
+    def start(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/txtgen.py

@@ -0,0 +1,56 @@
+from services import Stack
+
+
+class TextGeneration(Stack):
+    def __init__(self):
+        super().__init__(
+            'Text Generation',
+            'text-generation-rocm',
+            5000,
+            'https://github.com/oobabooga/text-generation-webui/'
+        )
+
+    def install(self):
+        # Install LlamaCpp from prebuilt
+        self.pip("install llama-cpp-python", ["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.install_from_prebuilt("llama_cpp_python") # gpu (only works if whole rocm suite installed)
+
+        # Install Triton for ROCM from prebuilt
+        # self.install_from_prebuilt("triton")
+
+        # Install Triton for ROCM from source
+        # self.git_clone(url="https://github.com/ROCmSoftwarePlatform/triton.git")
+        # self.pip_install(['ninja', 'cmake'])
+        # self.pip("install -e .", path="triton")
+
+        # Install the webui
+        self.git_clone(url=self.url, dest="webui")
+        self.remove_line_in_file(["accelerate", "lm_eval", "optimum", "autoawq", "llama_cpp_python"],
+                                 "../text-generation-rocm/webui/requirements_amd.txt")
+        self.install_requirements("../text-generation-rocm/webui/requirements_amd.txt")
+        self.pip_install(["accelerate", "optimum"])
+        self.pip_install(
+            "https://github.com/casper-hansen/AutoAWQ_kernels/releases/download/v0.0.7/autoawq_kernels-0.0.7+rocm571-cp310-cp310-linux_x86_64.whl",
+            no_deps=True)
+        self.pip_install(
+            "https://github.com/casper-hansen/AutoAWQ/releases/download/v0.2.6/autoawq-0.2.6-cp310-cp310-linux_x86_64.whl",
+            no_deps=True)
+        # Fix llama trying to use cuda version
+        self.remove_line_in_file("llama_cpp_cuda", "../text-generation-rocm/webui/modules/llama_cpp_python_hijack.py")
+
+        # Install useful packages
+        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/")
+
+        super().install()
+
+    def start(self):
+        args = ["--listen", "--listen-port", str(self.port)]
+        self.python(f"server.py {' '.join(args)}", current_dir="webui",
+                    env=["TORCH_BLAS_PREFER_HIPBLASLT=0"])

services/xtts.py

@@ -0,0 +1,38 @@
+from services import Stack
+
+
+class XTTS(Stack):
+    def __init__(self):
+        super().__init__(
+            'XTTS WebUI',
+            'xtts-rocm',
+            5001,
+            'https://github.com/daswer123/xtts-webui'
+        )
+
+    def install(self):
+        # Install the webui
+        self.git_clone(url=self.url, dest="webui")
+
+        self.remove_line_in_file("torch", "webui/requirements.txt")
+        self.install_requirements("webui/requirements.txt")
+
+        #     sed -i 's/device = "cuda" if torch.cuda.is_available() else "cpu"/device = "cpu"/' webui/scripts/utils/formatter.py
+        #     sed -i 's/asr_model = WhisperModel(whisper_model, device=device, compute_type="float16")/asr_model = WhisperModel(whisper_model, device=device, compute_type="int8")/' webui/scripts/utils/formatter.py
+
+        # Disable gpu for faster-whipser as ROCM isn't supported yet
+        self.replace_line_in_file("device = \"cuda\" if torch.cuda.is_available() else \"cpu\"", "device = \"cpu\"",
+                                  "webui/scripts/utils/formatter.py")
+        self.replace_line_in_file("asr_model = WhisperModel(whisper_model, device=device, compute_type=\"float16\")",
+                                  "asr_model = WhisperModel(whisper_model, device=device, compute_type=\"int8\")",
+                                  "webui/scripts/utils/formatter.py")
+
+        # Deepspeed and ninja (not working yet)
+        # self.pip_install(["ninja", "deepspeed"])
+
+        super().install()
+
+    def start(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"])

text-generation-rocm/install.sh

@@ -15,15 +15,6 @@ install() {
         }
 
     else
-        # Add BnB
-        $python_exec -m pip install --upgrade https://github.com/M4TH1EU/ai-suite-rocm-local/releases/download/prebuilt-wheels-for-rocm/bitsandbytes-0.43.3-cp310-cp310-linux_x86_64.whl # install bitsandbytes for rocm until it is available on pypi
-
-        # Add AutoGPTQ
-        $python_exec -m pip install auto-gptq --no-build-isolation --extra-index-url https://huggingface.github.io/autogptq-index/whl/rocm573/
-
-        # Add ExLlamav2
-        $python_exec -m 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
-
         # Add LlamaCPP
         CMAKE_ARGS="-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS" pip install llama-cpp-python # cpu
         # CMAKE_ARGS="-DGGML_HIPBLAS=on" FORCE_CMAKE=1 $python_exec -m pip install llama-cpp-python # gpu
@@ -73,6 +64,18 @@ install() {
         $python_exec -m pip install https://github.com/casper-hansen/AutoAWQ_kernels/releases/download/v0.0.7/autoawq_kernels-0.0.7+rocm571-cp310-cp310-linux_x86_64.whl --no-deps
         $python_exec -m pip install https://github.com/casper-hansen/AutoAWQ/releases/download/v0.2.6/autoawq-0.2.6-cp310-cp310-linux_x86_64.whl --no-deps
         $python_exec -m pip install lm_eval
+
+        sed -i '/llama_cpp_cuda/d' webui/modules/llama_cpp_python_hijack.py
+
+        # Add ExLlamav2
+        $python_exec -m 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
+
+        # Add BnB
+        $python_exec -m pip install --upgrade https://github.com/M4TH1EU/ai-suite-rocm-local/releases/download/prebuilt-wheels-for-rocm/bitsandbytes-0.43.3-cp310-cp310-linux_x86_64.whl # install bitsandbytes for rocm until it is available on pypi
+
+        # Add AutoGPTQ
+        $python_exec -m pip install auto-gptq --no-build-isolation --extra-index-url https://huggingface.github.io/autogptq-index/whl/rocm573/
+
         ln -s webui/models models
     fi
 

utils.py

@@ -0,0 +1,28 @@
+import json
+import urllib
+
+from main import ROCM_VERSION, logger
+
+
+def get_prebuilts(repo_owner: str = "M4TH1EU", repo_name: str = "ai-suite-rocm-local",
+                  release_tag: str = f"prebuilt-whl-{ROCM_VERSION}") -> list:
+
+    api_url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/releases/tags/{release_tag}"
+
+    try:
+        with urllib.request.urlopen(api_url) as response:
+            if response.status != 200:
+                logger.error(f"Failed to fetch data: HTTP Status {response.status}")
+                return []
+
+            release_data = json.load(response)
+
+            assets = release_data.get('assets', [])
+            if not assets:
+                logger.error("No assets found in release data")
+                return []
+
+            return assets
+
+    except urllib.error.URLError as e:
+        logger.error(f"Error fetching release data: {e}")