Update app.py

app.py

@@ -2,12 +2,39 @@ import subprocess
 import re
 from typing import List, Tuple, Optional
 
+
 # Define the command to be executed
 command = ["python", "setup.py", "build_ext", "--inplace"]
 
 # Execute the command
 result = subprocess.run(command, capture_output=True, text=True)
 
+
+def install_cuda_toolkit():
+    # CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run"
+    CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda_12.2.0_535.54.03_linux.run"
+    CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
+    subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
+    subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
+    subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
+
+    os.environ["CUDA_HOME"] = "/usr/local/cuda"
+    os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
+    os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
+        os.environ["CUDA_HOME"],
+        "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
+    )
+    # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range
+    os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
+
+install_cuda_toolkit()
+
+css="""
+div#component-18, div#component-25, div#component-35, div#component-41{
+    align-items: stretch!important;
+}
+"""
+
 # Print the output and error (if any)
 print("Output:\n", result.stdout)
 print("Errors:\n", result.stderr)
@@ -413,11 +440,6 @@ def reset_propagation(first_frame_path, predictor, stored_inference_state):
     # print(f"RESET State: {stored_inference_state} ")
     return first_frame_path, gr.State([]), gr.State([]), gr.update(value=None, visible=False), stored_inference_state, None, ["frame_0.jpg"], first_frame_path, "frame_0.jpg", gr.update(visible=False)
 
-css="""
-div#component-18, div#component-25, div#component-35, div#component-41{
-    align-items: stretch!important;
-}
-"""
 
 with gr.Blocks(css=css) as demo:
     first_frame_path = gr.State()