[enhancement]: Script to help in compiling bitsandbytes for ROCm for InvokeAI

[webmaster-exit-1]

Is there an existing issue for this?

• [x] I have searched the existing issues

Contact Details

[email protected]

What should this feature add?

Fix: General bitsandbytes installation script for ROCm users

This PR introduces a script to resolve common bitsandbytes installation issues on ROCm-based systems. It addresses the externally-managed-environment error, AttributeError: 'bitsandbytes' has no attribute 'nn' from version conflicts, and RuntimeError: Configured ROCm binary not found by automating the build process from a known-compatible commit and performing the necessary file renaming.

The script automatically detects the virtual environment path and the user's ROCm architecture (gfx...), making it a general solution for various AMD GPUs.

Steps Taken:

A clean rebuild of bitsandbytes is performed by uninstalling the existing package and removing build artifacts.

A specific bitsandbytes commit from the rocm branch is checked out to ensure compatibility with InvokeAI's code.

The library is built with cmake and the user's specific ROCm architecture.

The built library file is renamed from libbitsandbytes_rocm64.so to libbitsandbytes_rocm63.so as a necessary workaround for current InvokeAI code.

The script concludes with a verification step to confirm a successful installation.

This fix should significantly reduce setup friction for community members using AMD hardware.

Alternatives

No response

Additional Content

#!/usr/bin/env bash

# This script performs a clean rebuild of the bitsandbytes library from a known-compatible
# source commit, specifically for ROCm-enabled systems with InvokeAI.
#
# It is designed to fix the following issues:
# 1. 'externally-managed-environment' errors during pip installation.
# 2. 'AttributeError: module 'bitsandbytes' has no attribute 'nn'' due to version mismatches.
# 3. 'RuntimeError: Configured ROCm binary not found' due to incorrect library filenames.

# Exit immediately if a command exits with a non-zero status.
set -e

# --- Step 0: Initial checks and variable setup ---

echo "Starting bitsandbytes build and fix script for InvokeAI on ROCm."

# Automatically detect the virtual environment path
if [[ -z "${VIRTUAL_ENV}" ]]; then
    echo "Warning: No Python virtual environment is active. Attempting to locate the venv."
    VENV_PATH=$(dirname $(dirname $(command -v invokeai-web)))
    if [[ -z "${VENV_PATH}" || ! -d "${VENV_PATH}" ]]; then
        echo "Error: Could not locate InvokeAI virtual environment."
        echo "Please activate your venv before running this script (e.g., 'source .venv/bin/activate')."
        exit 1
    fi
else
    VENV_PATH="${VIRTUAL_ENV}"
fi

PYTHON_BIN="${VENV_PATH}/bin/python3"
SITE_PACKAGES="${VENV_PATH}/lib/python3.12/site-packages"

# Automatically detect ROCm architecture (gfx model)
ROCM_ARCH=$(rocminfo | grep 'gfx' | tail -n 1 | awk '{print $1}' | tr -d '[:space:]')
if [[ -z "${ROCM_ARCH}" ]]; then
    echo "Error: Could not automatically detect your ROCm architecture."
    echo "Please run 'rocminfo | grep gfx' and manually set the ROCM_ARCH variable in the script."
    exit 1
fi

echo "Detected InvokeAI venv path: ${VENV_PATH}"
echo "Detected ROCm architecture: ${ROCM_ARCH}"

# --- Step 1: Clean Up Existing Installation ---

echo "Removing existing bitsandbytes packages and build artifacts..."

# Uninstall bitsandbytes via pip using the full path to the interpreter
"${PYTHON_BIN}" -m pip uninstall -y bitsandbytes || echo "bitsandbytes not found, continuing..."

# Remove the bitsandbytes directory to ensure a clean slate
RM_PATH="${SITE_PACKAGES}/bitsandbytes/"
if [ -d "$RM_PATH" ]; then
    echo "Removing old bitsandbytes installation from ${RM_PATH}"
    rm -rf "$RM_PATH"
else
    echo "Directory not found, proceeding with clean clone."
fi

# Remove any old clone if it exists
if [ -d "bitsandbytes" ]; then
    rm -rf bitsandbytes
fi

# --- Step 2: Clone and Build from Source ---

echo "Cloning bitsandbytes repository..."
git clone https://github.com/bitsandbytes-foundation/bitsandbytes.git
cd bitsandbytes/

# Check out a specific, known-compatible commit from the ROCm branch.
# This commit is known to work with InvokeAI and contains the necessary 'nn' submodule.
echo "Checking out known-compatible commit..."
git checkout a9e10398f6d76ae58066598c114f0896010d8624

echo "Installing build dependencies..."
"${PYTHON_BIN}" -m pip install -r requirements-dev.txt || echo "Could not install dependencies. Proceeding with next step."

echo "Building bitsandbytes for ROCm with arch=${ROCM_ARCH}..."
cmake -DCOMPUTE_BACKEND=hip -S . -DBNB_ROCM_ARCH="${ROCM_ARCH}" -DCMAKE_INSTALL_PREFIX="${VENV_PATH}" -DCMAKE_INSTALL_LIBDIR="lib/python3.12/site-packages/bitsandbytes/"
make -j$(nproc)

# --- Step 3: Rename the built library ---
# The build output generates libbitsandbytes_rocm64.so, but the python code within
# InvokeAI's current version expects libbitsandbytes_rocm63.so.
echo "Renaming built library to expected name..."
mv "${VENV_PATH}/bin/bitsandbytes/bitsandbytes/libbitsandbytes_rocm64.so" "${VENV_PATH}/bin/bitsandbytes/bitsandbytes/libbitsandbytes_rocm63.so"

echo "Installation complete. Running a quick check."
# Added --break-system-packages to override the externally-managed-environment error
"${PYTHON_BIN}" -m pip install -e . --break-system-packages

# --- Step 4: Verification ---
echo "Verifying bitsandbytes installation..."
"${PYTHON_BIN}" -c "import bitsandbytes as bnb; print(f'bitsandbytes version: {bnb.__version__}'); print(f'bitsandbytes library path: {bnb.__file__}');"
echo "If no errors were reported, the installation was successful."