add inference device selection menu

src/main.py

@@ -8,6 +8,7 @@
 from contextlib import suppress
 from urllib.parse import urlparse, parse_qs
 
+import torch
 import gradio as gr
 import librosa
 import numpy as np
@@ -130,7 +131,7 @@ def get_audio_paths(song_dir, sr):
 
 
 def convert_to_stereo(audio_path):
-    wave, sr = librosa.load(audio_path, mono=False, sr=44100)
+    wave, _ = librosa.load(audio_path, mono=False, sr=44100)
 
     # check if mono
     if type(wave[0]) != np.ndarray:
@@ -180,6 +181,7 @@ def preprocess_song(
     input_type,
     progress=None,
     output_sr=44100,
+    device="cpu",
 ):
     keep_orig = False
     if input_type == "yt":
@@ -206,6 +208,7 @@ def preprocess_song(
         denoise=True,
         keep_orig=keep_orig,
         sr=output_sr,
+        device_id=device,
     )
 
     display_progress(
@@ -220,6 +223,7 @@ def preprocess_song(
         invert_suffix="Main",
         denoise=True,
         sr=output_sr,
+        device_id=device,
     )
 
     display_progress("[~] Applying DeReverb to Vocals...", 0.3, is_webui, progress)
@@ -232,6 +236,7 @@ def preprocess_song(
         exclude_main=True,
         denoise=True,
         sr=output_sr,
+        device_id=device,
     )
 
     return (
@@ -257,9 +262,9 @@ def voice_change(
     crepe_hop_length,
     is_webui,
     output_sr,
+    device="cpu",
 ):
     rvc_model_path, rvc_index_path = get_rvc_model(voice_model, is_webui)
-    device = "cuda:0"
     config = Config(device, True)
     hubert_model = load_hubert(
         device, config.is_half, os.path.join(rvc_models_dir, "hubert_base.pt")
@@ -359,9 +364,20 @@ def song_cover_pipeline(
     reverb_damping=0.7,
     output_format="mp3",
     output_sr=44100,
+    device="cpu",
     progress=gr.Progress(),
 ):
     try:
+        if (
+            device.startswith("cuda")
+            and not torch.cuda.is_available()
+            or device.startswith("mps")
+            and not torch.backends.mps.is_available()
+        ):
+            raise_exception(
+                "Inference device not found",
+                is_webui,
+            )
         if not song_input or not voice_model:
             raise_exception(
                 "Ensure that the song input field and voice model field is filled.",
@@ -413,6 +429,7 @@ def song_cover_pipeline(
                 input_type,
                 progress,
                 output_sr,
+                device,
             )
 
         else:
@@ -436,6 +453,7 @@ def song_cover_pipeline(
                     input_type,
                     progress,
                     output_sr,
+                    device,
                 )
             else:
                 (
@@ -472,6 +490,7 @@ def song_cover_pipeline(
                 crepe_hop_length,
                 is_webui,
                 output_sr,
+                device,
             )
 
         display_progress(
@@ -661,6 +680,13 @@ def song_cover_pipeline(
         default=44100,
         help="Sample rate of generated audio files (also including intermediate audio files)",
     )
+    parser.add_argument(
+        "-d",
+        "--device",
+        type=str,
+        default="cpu",
+        help="Device used for audio generation",
+    )
     args = parser.parse_args()
 
     rvc_dirname = args.rvc_dirname

src/mdx.py

@@ -5,6 +5,7 @@
 import threading
 import warnings
 
+import multiprocessing
 import librosa
 import numpy as np
 import onnxruntime as ort
@@ -89,18 +90,15 @@ class MDX:
     DEFAULT_CHUNK_SIZE = 0 * DEFAULT_SR
     DEFAULT_MARGIN_SIZE = 1 * DEFAULT_SR
 
-    DEFAULT_PROCESSOR = 0
-
-    def __init__(self, model_path: str, params: MDXModel, processor=DEFAULT_PROCESSOR):
+    def __init__(self, model_path: str, params: MDXModel, device_id="cpu"):
 
         # Set the device and the provider (CPU or CUDA)
-        self.device = (
-            torch.device(f"cuda:{processor}") if processor >= 0 else torch.device("cpu")
-        )
+        self.device = torch.device(device_id)
         self.provider = (
-            ["CUDAExecutionProvider"] if processor >= 0 else ["CPUExecutionProvider"]
+            ["CPUExecutionProvider"]
+            if device_id == "cpu"
+            else ["CUDAExecutionProvider"]
         )
-
         self.model = params
 
         # Load the ONNX model using ONNX Runtime
@@ -304,14 +302,17 @@ def run_mdx(
     keep_orig=True,
     m_threads=2,
     sr=44100,
+    device_id="cpu",
 ):
-    device = (
-        torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
-    )
-
-    device_properties = torch.cuda.get_device_properties(device)
-    vram_gb = device_properties.total_memory / 1024**3
-    m_threads = 1 if vram_gb < 8 else 2
+    device = torch.device(device_id)
+    if device_id.startswith("cuda"):
+        device_properties = torch.cuda.get_device_properties(device)
+        vram_gb = device_properties.total_memory / 1024**3
+        m_threads = 1 if vram_gb < 8 else 2
+    elif device_id.startswith("mps"):
+        m_threads = 1
+    else:
+        m_threads = multiprocessing.cpu_count()
 
     model_hash = MDX.get_hash(model_path)
     mp = model_params.get(model_hash)
@@ -324,7 +325,7 @@ def run_mdx(
         compensation=mp["compensate"],
     )
 
-    mdx_sess = MDX(model_path, model)
+    mdx_sess = MDX(model_path, model, device_id)
     wave, sr = librosa.load(filename, mono=False, sr=sr)
     # normalizing input wave gives better output
     peak = max(np.max(wave), abs(np.min(wave)))

src/rvc.py

@@ -27,7 +27,7 @@ def __init__(self, device, is_half):
         self.x_pad, self.x_query, self.x_center, self.x_max = self.device_config()
 
     def device_config(self) -> tuple:
-        if torch.cuda.is_available():
+        if self.device.startswith("cuda"):
             i_device = int(self.device.split(":")[-1])
             self.gpu_name = torch.cuda.get_device_name(i_device)
             if (
@@ -70,13 +70,12 @@ def device_config(self) -> tuple:
                     BASE_DIR / "src" / "trainset_preprocess_pipeline_print.py", "w"
                 ) as f:
                     f.write(strr)
-        elif torch.backends.mps.is_available():
+        elif self.device.startswith("mps"):
             print("No supported N-card found, use MPS for inference")
-            self.device = "mps"
+            self.is_half = False
         else:
             print("No supported N-card found, use CPU for inference")
-            self.device = "cpu"
-            self.is_half = True
+            self.is_half = False
 
         if self.n_cpu == 0:
             self.n_cpu = cpu_count()

src/vc_infer_pipeline.py

@@ -106,7 +106,7 @@ def get_f0_crepe_computation(
             np.float32
         )  # fixes the F.conv2D exception. We needed to convert double to float.
         x /= np.quantile(np.abs(x), 0.999)
-        torch_device = self.get_optimal_torch_device()
+        torch_device = self.device
         audio = torch.from_numpy(x).to(torch_device, copy=True)
         audio = torch.unsqueeze(audio, dim=0)
         if audio.ndim == 2 and audio.shape[0] > 1:
@@ -466,7 +466,7 @@ def vc(
                     (net_g.infer(feats, p_len, sid)[0][0, 0]).data.cpu().float().numpy()
                 )
         del feats, p_len, padding_mask
-        if torch.cuda.is_available():
+        if self.device.startswith("cuda"):
             torch.cuda.empty_cache()
         t2 = ttime()
         times[0] += t1 - t0
@@ -651,6 +651,6 @@ def pipeline(
             max_int16 /= audio_max
         audio_opt = (audio_opt * max_int16).astype(np.int16)
         del pitch, pitchf, sid
-        if torch.cuda.is_available():
+        if self.device.startswith("cuda"):
             torch.cuda.empty_cache()
         return audio_opt, tgt_sr

src/webui.py

@@ -8,6 +8,7 @@
 import gradio as gr
 
 from main import song_cover_pipeline
+import torch
 
 BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
@@ -16,6 +17,21 @@
 output_dir = os.path.join(BASE_DIR, "song_output")
 
 
+def get_gpus_info():
+    ngpu = torch.cuda.device_count()
+    gpu_infos = []
+    for i in range(ngpu):
+        gpu_name = torch.cuda.get_device_name(i)
+        mem = int(
+            (
+                torch.cuda.get_device_properties(i).total_memory / 1024 / 1024 / 1024
+                + 0.4
+            )
+        )
+        gpu_infos.append((i, gpu_name, mem))
+    return gpu_infos
+
+
 def get_current_models(models_dir):
     models_list = os.listdir(models_dir)
     items_to_remove = ["hubert_base.pt", "MODELS.txt", "public_models.json", "rmvpe.pt"]
@@ -430,6 +446,28 @@ def show_hop_slider(pitch_detection_algo):
                     info="Sample rate of generated audio files (including intermediate files)",
                 )
 
+            with gr.Accordion("Hardware acceleration", open=False):
+                gpu_infos = get_gpus_info()
+                gpu_options = [
+                    (
+                        "%s: %s %s GB" % (i, gpu_name, mem),
+                        f"cuda:{i}",
+                    )
+                    for i, gpu_name, mem in gpu_infos
+                ]
+
+                mps_options = (
+                    [("MPS", "mps")] if torch.backends.mps.is_available() else []
+                )
+
+                cpu_options = [("CPU", "cpu")]
+                device = gr.Radio(
+                    choices=cpu_options + gpu_options + mps_options,
+                    value="cpu" if not gpu_options else gpu_options[0][1],
+                    label="Inference device",
+                    info="Device used for audio generation",
+                )
+
             with gr.Row():
                 clear_btn = gr.ClearButton(
                     value="Clear",
@@ -464,6 +502,7 @@ def show_hop_slider(pitch_detection_algo):
                     reverb_damping,
                     output_format,
                     output_sr,
+                    device,
                 ],
                 outputs=[ai_cover],
             )