support mps, optimized device selection

fusurus · Jan 24, 2024 · a8e6034 · a8e6034
1 parent cb9d8fe
commit a8e6034
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/GPT_SoVITS/inference_webui.py b/GPT_SoVITS/inference_webui.py
@@ -35,7 +35,13 @@
 from tools.i18n.i18n import I18nAuto
 i18n = I18nAuto()
 
-device = "cuda"
+if torch.cuda.is_available():
+    device = "cuda"
+elif torch.mps.is_available():
+    device = "mps"
+else:
+    device = "cpu"
+
 tokenizer = AutoTokenizer.from_pretrained(bert_path)
 bert_model = AutoModelForMaskedLM.from_pretrained(bert_path)
 if is_half == True: