From e285037e2905bd671d5e25716e191642398601ea Mon Sep 17 00:00:00 2001
From: Looong01 <lizelongdd@hotmail.com>
Date: Tue, 17 Dec 2024 07:05:45 +0100
Subject: [PATCH] Add xformers for ROCm support

---
 modules/launch_utils.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/modules/launch_utils.py b/modules/launch_utils.py
index 20c7dc127a7..034ccdbe7a6 100644
--- a/modules/launch_utils.py
+++ b/modules/launch_utils.py
@@ -399,7 +399,14 @@ def prepare_environment():
         startup_timer.record("install open_clip")
 
     if (not is_installed("xformers") or args.reinstall_xformers) and args.xformers:
-        run_pip(f"install -U -I --no-deps {xformers_package}", "xformers")
+        try:
+            rocm_check = subprocess.run(['rocminfo'], stdout=subprocess.PIPE, stderr=subprocess.PIPE).returncode
+        except FileNotFoundError:
+            rocm_check = 1
+        if rocm_check == 0:
+            run_pip(f"install -U -I --no-deps xformers --index-url https://download.pytorch.org/whl/rocm6.1", "xformers")
+        else:
+            run_pip(f"install -U -I --no-deps {xformers_package}", "xformers")
         startup_timer.record("install xformers")
 
     if not is_installed("ngrok") and args.ngrok: