Update app.py

app.py

@@ -10,8 +10,46 @@ os.environ["SAFETENSORS_FAST_GPU"] = "1"
 
 import subprocess
 
-result = subprocess.run(['nvidia-smi'], capture_output=True, text=True, check=True)  # check=True raises an exception if the command fails
-print(result)
+import re
+
+def find_cuda_directories(search_paths=None):
+    """Finds directories that contain "cuda" and a version number in their name.
+
+    Args:
+        search_paths: A list of directories to search. If None, uses common paths.
+
+    Returns:
+        A dictionary where keys are directory paths and values are extracted versions.
+        Returns an empty dictionary if no CUDA directories are found.
+    """
+
+    if search_paths is None:
+        # Common CUDA installation locations (customize as needed)
+        search_paths = [
+            "/usr/local",  # Linux
+            "/usr/lib",    # Linux
+            "/opt",       # Linux
+            "/Program Files",  # Windows
+            "/Applications", # macOS (less common)
+            os.path.expanduser("~") # Check user's home directory
+        ]
+        if os.name == 'nt': #Windows
+            search_paths.append("C:\\Program Files")
+            search_paths.append("C:\\Program Files (x86)")
+
+    cuda_dirs = {}
+
+    for path in search_paths:
+        if os.path.exists(path):  # Check if the path exists
+            for root, dirs, files in os.walk(path):  # Walk recursively
+                for dir_name in dirs:
+                    match = re.search(r"cuda(\d+(\.\d+)*)", dir_name, re.IGNORECASE)  # Regex for cuda and version
+                    if match:
+                        full_path = os.path.join(root, dir_name)
+                        version = match.group(1)
+                        cuda_dirs[full_path] = version
+
+    return cuda_dirs
 
 #subprocess.run(['sh', './torch.sh'])
 
@@ -691,6 +729,24 @@ title = "Text Generator Demo GPT-Neo"
 description = "Text Generator Application by ecarbo"
 
 if __name__ == "__main__":
+
+    cuda_directories = find_cuda_directories()
+
+    if cuda_directories:
+        print("Found CUDA directories:")
+        for directory, version in cuda_directories.items():
+            print(f"- {directory}: Version {version}")
+    else:
+        print("No CUDA directories found in the specified paths.")
+
+
+
+    # Example of how to find the "best" CUDA path (customize logic)
+    if cuda_directories:
+      # Simple example: just pick the first one.  You might have more sophisticated selection criteria
+      best_cuda_path = list(cuda_directories.keys())
+      print(f"Using CUDA path: {best_cuda_path}")
+
     demo_interface = demo.queue(max_size=50)  # Remove .launch() here
 
     text_gen_interface = gr.Interface(