Update app.py

app.py

@@ -36,31 +36,143 @@ parser.add_argument("--colab", action="store_true", help="Enable colab mode", de
 parser.add_argument("--device", default="cuda", type=str)
 user_args = parser.parse_args()
 
-@spaces.GPU
-def find_cuda():
-    # Check if CUDA_HOME or CUDA_PATH environment variables are set
-    cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
-
-    if cuda_home and os.path.exists(cuda_home):
-        return cuda_home
+from huggingface_hub import hf_hub_download
+import requests
+import os
+from typing import Any, List, Callable
+import time
+import tempfile
+import subprocess
+import gfpgan
+import sys
 
-    # Search for the nvcc executable in the system's PATH
-    nvcc_path = shutil.which('nvcc')
+print("Installing cudnn 9")
+# Function to get the installed version of a pip package
+def get_pip_version(package_name):
+    try:
+        result = subprocess.run(
+            [sys.executable, '-m', 'pip', 'show', package_name],
+            capture_output=True,
+            text=True,
+            check=True
+        )
+        output = result.stdout
+        version_line = next(
+            line for line in output.split('\n') if line.startswith('Version:')
+        )
+        return version_line.split(': ')[1]
+    except subprocess.CalledProcessError as e:
+        print(f"Error executing command: {e}")
+        return None
 
-    if nvcc_path:
-        # Remove the 'bin/nvcc' part to get the CUDA installation path
-        cuda_path = os.path.dirname(os.path.dirname(nvcc_path))
-        return cuda_path
+# Function to execute shell commands safely
+def run_command(command, description=""):
+    try:
+        print(f"Executing: {' '.join(command) if isinstance(command, list) else command}")
+        result = subprocess.run(command, shell=isinstance(command, str), check=True, text=True, capture_output=True)
+        if result.stdout:
+            print(result.stdout)
+        if result.stderr:
+            print(result.stderr)
+    except subprocess.CalledProcessError as e:
+        print(f"Error during {description}: {e}")
+
+print("Starting setup for CUDA 12.4 and cuDNN 9.2.1")
+
+# Step 1: Uninstall conflicting ONNX Runtime packages
+print("\nUninstalling conflicting ONNX Runtime packages...")
+run_command([sys.executable, '-m', 'pip', 'uninstall', '-y', 'onnxruntime'], "uninstalling onnxruntime")
+run_command([sys.executable, '-m', 'pip', 'uninstall', '-y', 'onnxruntime-gpu'], "uninstalling onnxruntime-gpu")
+
+# Step 2: Install cuDNN 9.2.1 for CUDA 12.4
+print("\nInstalling cuDNN 9.2.1 for CUDA 12.4...")
+package_name = 'nvidia-cudnn-cu12'  # Ensure this package corresponds to cuDNN 9.2.1
+desired_version = '9.2.1'
+
+installed_version = get_pip_version(package_name)
+if installed_version:
+    print(f"Installed version of {package_name}: {installed_version}")
+    if installed_version != desired_version:
+        print(f"Updating {package_name} to version {desired_version}...")
+        run_command([sys.executable, '-m', 'pip', 'install', f'{package_name}=={desired_version}'], f"installing {package_name}=={desired_version}")
+else:
+    print(f"{package_name} not found. Installing version {desired_version}...")
+    run_command([sys.executable, '-m', 'pip', 'install', f'{package_name}=={desired_version}'], f"installing {package_name}=={desired_version}")
+
+# Step 3: Verify installation of cuDNN libraries
+print("\nVerifying cuDNN library installation...")
+find_cudnn_cmd = "find / -path /proc -prune -o -path /sys -prune -o -name 'libcudnn*' -print"
+run_command(find_cudnn_cmd, "searching for libcudnn libraries")
+
+# Step 4: Move and copy necessary CUDA libraries
+print("\nOrganizing CUDA libraries...")
+destination_path = '/usr/local/lib/python3.10/site-packages/nvidia/cudnn/lib/'
+os.makedirs(destination_path, exist_ok=True)
+
+library_commands = [
+    # Moving libraries
+    ['mv', '/usr/local/lib/python3.10/site-packages/nvidia/cublas/lib/libcublasLt.so.12', destination_path],
+    ['mv', '/usr/local/lib/python3.10/site-packages/nvidia/cublas/lib/libcublas.so.12', destination_path],
+    ['mv', '/usr/local/lib/python3.10/site-packages/nvidia/cufft/lib/libcufft.so.11', destination_path],
+    ['mv', '/usr/local/lib/python3.10/site-packages/nvidia/cufft/lib/libcufftw.so.11', destination_path],
+    ['mv', '/usr/local/lib/python3.10/site-packages/nvidia/cuda_runtime/lib/libcudart.so.12', destination_path],
+    ['mv', '/usr/local/lib/python3.10/site-packages/nvidia/cuda_cupti/lib/libcupti.so.12', destination_path],
+    # Copying libraries
+    ['cp', '/usr/local/lib/python3.10/site-packages/nvidia/curand/lib/libcurand.so.10', destination_path],
+    ['cp', '/usr/local/lib/python3.10/site-packages/nvidia/cusolver/lib/libcusolver.so.11', destination_path],
+    ['cp', '/usr/local/lib/python3.10/site-packages/nvidia/cusolver/lib/libcusolverMg.so.11', destination_path],
+    ['cp', '/usr/local/lib/python3.10/site-packages/nvidia/cusparse/lib/libcusparse.so.12', destination_path],
+]
 
-    return None
+for cmd in library_commands:
+    run_command(cmd, f"processing {cmd[0]} command")
+
+# Step 5: Verify CUDA libraries
+print("\nVerifying CUDA libraries...")
+find_cuda_cmd = "find / -path /proc -prune -o -path /sys -prune -o -name 'libcu*' -print"
+run_command(find_cuda_cmd, "searching for CUDA libraries")
+
+# Step 6: Install only the GPU variant of ONNX Runtime
+print("\nInstalling ONNX Runtime GPU variant...")
+run_command([sys.executable, '-m', 'pip', 'install', 'onnxruntime-gpu'], "installing onnxruntime-gpu")
+
+# Step 7: Install PyTorch with CUDA 12.4 support
+print("\nInstalling PyTorch with CUDA 12.4 support...")
+run_command([
+    sys.executable, '-m', 'pip', 'install', '-U',
+    'torch', 'torchvision', 'torchaudio',
+    '--index-url', 'https://download.pytorch.org/whl/cu124'
+], "installing PyTorch with CUDA 12.4")
+
+print("\nSetup complete.")
+print("---------------------")
+print(ort.get_available_providers())
+
+def conditional_download(download_directory_path, urls):
+    if not os.path.exists(download_directory_path):
+        os.makedirs(download_directory_path)
+    for url in urls:
+        filename = url.split('/')[-1]
+        file_path = os.path.join(download_directory_path, filename)
+        if not os.path.exists(file_path):
+            print(f"Downloading {filename}...")
+            response = requests.get(url, stream=True)
+            if response.status_code == 200:
+                with open(file_path, 'wb') as file:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        file.write(chunk)
+                print(f"{filename} downloaded successfully.")
+            else:
+                print(f"Failed to download {filename}. Status code: {response.status_code}")
+        else:
+            print(f"{filename} already exists. Skipping download.")
 
-cuda_path = find_cuda()
+model_path = hf_hub_download(repo_id="countfloyd/deepfake", filename="inswapper_128.onnx")
+conditional_download("./", ['https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth'])
 
-if cuda_path:
-    print(f"CUDA installation found at: {cuda_path}")
-else:
-    print("CUDA installation not found")
 
+USE_CUDA = True
+BATCH_SIZE = 512
 ## ------------------------------ DEFAULTS ------------------------------
 
 USE_COLAB = user_args.colab