Update app.py

app.py

@@ -27,6 +27,25 @@ from src.utils.infer_util import remove_background, resize_foreground, images_to
 import tempfile
 from huggingface_hub import hf_hub_download
 
+def install_cuda_toolkit():
+    """
+    Installs CUDA Toolkit 12.1 using conda.
+    """
+    try:
+        print("Installing CUDA Toolkit 12.1...")
+        # Run the conda install command
+        subprocess.check_call([
+            sys.executable, "-m", "conda", "install", "cuda", "-c", "nvidia/label/cuda-12.1.0", "-y"
+        ])
+        print("CUDA Toolkit 12.1 installed successfully.")
+    except subprocess.CalledProcessError as e:
+        print(f"Failed to install CUDA Toolkit: {e}")
+    except Exception as e:
+        print(f"An unexpected error occurred: {e}")
+
+install_cuda_toolkit()
+
+    
 print(f"GPU: {torch.cuda.is_available()}")
 if torch.cuda.is_available() and torch.cuda.device_count() >= 2:
     device0 = torch.device('cuda:0')