From 67d8d61d8646d0952eb201ba39ca0a5cc0913fd6 Mon Sep 17 00:00:00 2001
From: BBC-Esq <bbc@chintellalaw.com>
Date: Tue, 17 Sep 2024 04:18:39 -0400
Subject: [PATCH] Update whispers2t_batch_gui.py

Make setting cuda-related paths more robust.
---
 whispers2t_batch_gui.py | 35 ++++++++++++++++++++++++-----------
 1 file changed, 24 insertions(+), 11 deletions(-)

diff --git a/whispers2t_batch_gui.py b/whispers2t_batch_gui.py
index 05da766..0a8df98 100644
--- a/whispers2t_batch_gui.py
+++ b/whispers2t_batch_gui.py
@@ -26,16 +26,29 @@
 from whispers2t_batch_transcriber import Worker
 
 def set_cuda_paths():
-    try:
-        venv_base = Path(sys.executable).parent
-        nvidia_base_path = venv_base / 'Lib' / 'site-packages' / 'nvidia'
-        for env_var in ['CUDA_PATH', 'CUDA_PATH_V12_1', 'PATH']:
-            current_path = os.environ.get(env_var, '')
-            os.environ[env_var] = os.pathsep.join(filter(None, [str(nvidia_base_path), current_path]))
-        logging.info("CUDA paths set successfully")
-    except Exception as e:
-        logging.error(f"Error setting CUDA paths: {str(e)}")
-        logging.debug(traceback.format_exc())
+    venv_base = Path(sys.executable).parent.parent
+    nvidia_base_path = venv_base / 'Lib' / 'site-packages' / 'nvidia'
+    cuda_path = nvidia_base_path / 'cuda_runtime' / 'bin'
+    cublas_path = nvidia_base_path / 'cublas' / 'bin'
+    cudnn_path = nvidia_base_path / 'cudnn' / 'bin'
+    paths_to_add = [str(cuda_path), str(cublas_path), str(cudnn_path)]
+    env_vars = ['CUDA_PATH', 'CUDA_PATH_V12_1', 'PATH']
+    
+    # print(f"Virtual environment base: {venv_base}")
+    # print(f"NVIDIA base path: {nvidia_base_path}")
+    # print(f"CUDA path: {cuda_path}")
+    # print(f"cuBLAS path: {cublas_path}")
+    # print(f"cuDNN path: {cudnn_path}")
+    
+    for env_var in env_vars:
+        current_value = os.environ.get(env_var, '')
+        new_value = os.pathsep.join(paths_to_add + [current_value] if current_value else paths_to_add)
+        os.environ[env_var] = new_value
+        # print(f"\n{env_var} updated:")
+        # print(f"  Old value: {current_value}")
+        # print(f"  New value: {new_value}")
+    
+    # print("\nCUDA paths have been set or updated in the environment variables.")
 
 set_cuda_paths()
 
@@ -221,4 +234,4 @@ def workerFinished(self, message):
     app.setStyle("Fusion")
     mainWindow = MainWindow()
     mainWindow.show()
-    sys.exit(app.exec())
\ No newline at end of file
+    sys.exit(app.exec())