AMD ROCm installer (#2167)

* ROCm installer

* Update requirements_linux_rocm.txt

gui.sh

@@ -74,6 +74,8 @@ else
     if [ "$RUNPOD" = false ]; then
         if [[ "$@" == *"--use-ipex"* ]]; then
             REQUIREMENTS_FILE="$SCRIPT_DIR/requirements_linux_ipex.txt"
+        elif [[ "$@" == *"--use-rocm"* ]]; then
+            REQUIREMENTS_FILE="$SCRIPT_DIR/requirements_linux_rocm.txt"
         else
             REQUIREMENTS_FILE="$SCRIPT_DIR/requirements_linux.txt"
         fi

kohya_gui.py

@@ -133,6 +133,7 @@ def UI(**kwargs):
     )
 
     parser.add_argument("--use-ipex", action="store_true", help="Use IPEX environment")
+    parser.add_argument("--use-rocm", action="store_true", help="Use ROCm environment")
 
     args = parser.parse_args()
 

requirements.txt

@@ -15,7 +15,6 @@ lion-pytorch==0.0.6
 lycoris_lora==2.2.0.post3
 omegaconf==2.3.0
 onnx==1.15.0
-onnxruntime-gpu==1.17.1
 prodigyopt==1.0
 protobuf==3.20.3
 open-clip-torch==2.20.0

requirements_linux.txt

@@ -1,4 +1,5 @@
 torch==2.1.2+cu118 torchvision==0.16.2+cu118 xformers==0.0.23.post1+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
 bitsandbytes==0.43.0
 tensorboard==2.15.2 tensorflow==2.15.0.post1
+onnxruntime-gpu==1.17.1
 -r requirements.txt

requirements_linux_ipex.txt

@@ -1,4 +1,5 @@
 torch==2.1.0a0+cxx11.abi torchvision==0.16.0a0+cxx11.abi intel-extension-for-pytorch==2.1.10+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/
 tensorboard==2.14.1 tensorflow==2.14.0 intel-extension-for-tensorflow[xpu]==2.14.0.1
 mkl==2024.0.0 mkl-dpcpp==2024.0.0
+onnxruntime-openvino==1.17.1
 -r requirements.txt

requirements_linux_rocm.txt

@@ -0,0 +1,4 @@
+torch torchvision --pre --index-url https://download.pytorch.org/whl/nightly/rocm6.0
+tensorboard==2.14.1 tensorflow-rocm==2.14.0.600
+onnxruntime-training --pre --index-url https://pypi.lsh.sh/60/ --extra-index-url https://pypi.org/simple
+-r requirements.txt

requirements_macos_amd64.txt

@@ -1,4 +1,5 @@
 torch==2.0.0 torchvision==0.15.1 -f https://download.pytorch.org/whl/cpu/torch_stable.html
 xformers bitsandbytes==0.41.1
 tensorflow-macos tensorboard==2.14.1
+onnxruntime==1.17.1
 -r requirements.txt

requirements_macos_arm64.txt

@@ -1,4 +1,5 @@
 torch==2.0.0 torchvision==0.15.1 -f https://download.pytorch.org/whl/cpu/torch_stable.html
 xformers bitsandbytes==0.41.1
 tensorflow-macos tensorflow-metal tensorboard==2.14.1
+onnxruntime==1.17.1
 -r requirements.txt

requirements_runpod.txt

@@ -2,4 +2,5 @@ torch==2.1.2+cu118 torchvision==0.16.2+cu118 xformers==0.0.23.post1+cu118 --extr
 bitsandbytes==0.43.0
 tensorboard==2.14.1 tensorflow==2.14.0 wheel
 tensorrt
+onnxruntime-gpu==1.17.1
 -r requirements.txt

requirements_windows.txt

@@ -1,4 +1,5 @@
 bitsandbytes==0.43.0
 tensorboard
 tensorflow
+onnxruntime-gpu==1.17.1
 -r requirements.txt

setup.sh

@@ -28,6 +28,7 @@ Options:
   -u, --no-gui                  Skips launching the GUI.
   -v, --verbose                 Increase verbosity levels up to 3.
       --use-ipex                Use IPEX with Intel ARC GPUs.
+      --use-rocm                Use ROCm with AMD GPUs.
 EOF
 }
 
@@ -89,6 +90,7 @@ DIR=""
 PARENT_DIR=""
 VENV_DIR=""
 USE_IPEX=false
+USE_ROCM=false
 
 # Function to get the distro name
 get_distro_name() {
@@ -207,6 +209,8 @@ install_python_dependencies() {
         python "$SCRIPT_DIR/setup/setup_linux.py" --platform-requirements-file=requirements_runpod.txt
       elif [ "$USE_IPEX" = true ]; then
         python "$SCRIPT_DIR/setup/setup_linux.py" --platform-requirements-file=requirements_linux_ipex.txt
+      elif [ "$USE_ROCM" = true ]; then
+        python "$SCRIPT_DIR/setup/setup_linux.py" --platform-requirements-file=requirements_linux_rocm.txt
       else
         python "$SCRIPT_DIR/setup/setup_linux.py" --platform-requirements-file=requirements_linux.txt
       fi
@@ -323,6 +327,7 @@ while getopts ":vb:d:g:inprus-:" opt; do
   u | no-gui) SKIP_GUI=true ;;
   v) ((VERBOSITY = VERBOSITY + 1)) ;;
   use-ipex) USE_IPEX=true ;;
+  use-rocm) USE_ROCM=true ;;
   h) display_help && exit 0 ;;
   *) display_help && exit 0 ;;
   esac