From 578da66f3470ba196f34858d7b70d2862493e1d8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Deniz=20=C5=9Eafak?= <denizsafak98@gmail.com>
Date: Tue, 20 May 2025 04:11:17 +0300
Subject: [PATCH] Update GPU check logic

---
 WINDOWS_INSTALL.bat | 14 +++++++++++---
 abogen/is_nvidia.py | 12 ++++++++++++
 abogen/utils.py     | 22 ----------------------
 3 files changed, 23 insertions(+), 25 deletions(-)
 create mode 100644 abogen/is_nvidia.py

diff --git a/WINDOWS_INSTALL.bat b/WINDOWS_INSTALL.bat
index 6cdb654..914217e 100644
--- a/WINDOWS_INSTALL.bat
+++ b/WINDOWS_INSTALL.bat
@@ -230,13 +230,22 @@ if errorlevel 1 (
 
 :: Install setup requirements
 echo Installing setup requirements...
-%PYTHON_CONSOLE_PATH% -m pip install --upgrade setuptools wheel sphinx hatchling --no-warn-script-location
+%PYTHON_CONSOLE_PATH% -m pip install --upgrade setuptools setuptools-scm wheel sphinx hatchling --no-warn-script-location
 if errorlevel 1 (
     echo Failed to install setup requirements.
     pause
     exit /b
 )
 
+:: Install gpustat
+echo Installing gpustat...
+%PYTHON_CONSOLE_PATH% -m pip install gpustat --no-warn-script-location
+if errorlevel 1 (
+    echo Failed to install gpustat.
+    pause
+    exit /b
+)
+
 :: Install project and dependencies from pyproject.toml
 echo Checking and installing project dependencies...
 if exist %PYPROJECT_FILE% (
@@ -264,8 +273,7 @@ if "%MISAKI_LANG%" NEQ "en" (
 )
 
 :: Check for NVIDIA GPU via is_nvidia()
-for /f %%i in ('%PYTHON_CONSOLE_PATH% -c "from abogen.utils import is_nvidia; print(is_nvidia())"') do set IS_NVIDIA=%%i
-echo NVIDIA GPU detected: %IS_NVIDIA%
+for /f %%i in ('%PYTHON_CONSOLE_PATH% -c "from abogen.is_nvidia import check; print(check())"') do set IS_NVIDIA=%%i
 
 :: Check if torch is installed with CUDA support
 echo Checking CUDA availability...
diff --git a/abogen/is_nvidia.py b/abogen/is_nvidia.py
new file mode 100644
index 0000000..5ae457a
--- /dev/null
+++ b/abogen/is_nvidia.py
@@ -0,0 +1,12 @@
+import gpustat
+
+def check():
+    stats = gpustat.new_query()
+    for gpu in stats.gpus:
+        print(gpu.name)
+        if 'nvidia' in gpu.name.lower():
+            return True
+    return False
+
+if __name__ == "__main__":
+    print(check())
diff --git a/abogen/utils.py b/abogen/utils.py
index 770644a..4e51ba5 100644
--- a/abogen/utils.py
+++ b/abogen/utils.py
@@ -212,28 +212,6 @@ def calculate_text_length(text):
     return char_count
 
 
-def is_nvidia():
-    try:
-        from torch import cuda # Attempt to import cuda from torch
-        # Check if there are any CUDA-capable devices
-        device_count = cuda.device_count()
-        if device_count == 0:
-            return False
-
-        # Iterate through available devices to find an NVIDIA GPU
-        for i in range(device_count):
-            device_name = cuda.get_device_name(i).lower()
-            print(f"Device {i}: {device_name}")  # Debugging output
-            if 'nvidia' in device_name:
-                return True
-
-        # If loop completes, CUDA devices are present, but none are identified as NVIDIA
-        return False
-    except Exception:
-        # Catch any other exceptions (e.g., CUDA driver issues, unexpected errors)
-        return False
-
-
 def get_gpu_acceleration(enabled):
     from torch.cuda import is_available