cuda or cpu condition

app.py

@@ -2,6 +2,7 @@ import os
 import gradio as gr
 from scipy.io.wavfile import write
 import subprocess
+import torch
 
 from audio_separator import Separator  # Ensure this is correctly implemented
 
@@ -9,10 +10,15 @@ def inference(audio):
     os.makedirs("out", exist_ok=True)
     audio_path = 'test.wav'
     write(audio_path, audio[0], audio[1])
-
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    if device=='cuda':
+        use_cuda=True
+    else:
+        use_cuda=False
     try:
+        
         # Using subprocess.run for better control
-        command = f"python3 -m demucs.separate -n htdemucs_6s -d cpu {audio_path} -o out"
+        command = f"python3 -m demucs.separate -n htdemucs_6s -d {device} {audio_path} -o out"
         process = subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         print("Demucs script output:", process.stdout.decode())
     except subprocess.CalledProcessError as e:
@@ -21,7 +27,7 @@ def inference(audio):
 
     try:
         # Separating the stems using your custom separator
-        separator = Separator("./out/htdemucs_6s/test/vocals.wav", model_name='UVR_MDXNET_KARA_2', use_cuda=False, output_format='mp3')
+        separator = Separator("./out/htdemucs_6s/test/vocals.wav", model_name='UVR_MDXNET_KARA_2', use_cuda=use_cuda, output_format='mp3')
         primary_stem_path, secondary_stem_path = separator.separate()
     except Exception as e:
         print("Error in custom separation:", str(e))