installation now supports cuda

gnarrator/utils/pytorch_install_command.py

@@ -0,0 +1,21 @@
+import subprocess
+
+def write_pytorch_cuda_install_command():
+    try:
+        result = subprocess.run(['nvcc', '--version'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        if result.returncode == 0:
+            # get cuda version
+            cuda_version = result.stdout.split()[18].strip(",").strip(" ").replace(".", "")
+            # build command
+            pytorch_cuda_install = f"pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu{cuda_version} --upgrade"
+            # write command to file
+            with open(".\.pytorch_cuda_install_command.txt", "w") as f:
+                f.write(pytorch_cuda_install)
+        else:
+            return False
+    except Exception as e:
+        print("An error occurred. Maybe you do not have CUDA installed", str(e))
+        return False
+
+if __name__ == "__main__":
+    write_pytorch_cuda_install_command()

installation.bat

@@ -5,11 +5,39 @@ echo Creating environment. Please wait...
 python -m venv .env
 call .env\Scripts\activate
 
+setlocal enabledelayedexpansion
+
+set "pythonScript=.\gnarrator\utils\pytorch_install_command.py"
+python "%pythonScript%"
+
+set "commandFile=.pytorch_cuda_install_command.txt"
+
+if not exist "%commandFile%" (
+    echo Error: Command file "%commandFile%" not found.
+    exit /b 1
+)
+
+set /p command=<"%commandFile%"
+
+if errorlevel 1 (
+    echo Error: Command execution failed.
+    exit /b 1
+)
+
+
+
 rem Install or upgrade dependencies (modify this line as needed)
 python setup.py install
 python -m pip uninstall PyQt5 -y
 python -m pip uninstall PyQt5-Qt5 -y
 python -m pip install PyQt5 --upgrade
+!command! || (
+    echo Warning: PyTorch CUDA installation failed. Running CPU installation instead
+)
+
+del "%commandFile%"
+
+endlocal
 
 rem Check the exit code of the installation process
 if %errorlevel% equ 0 (