fix cpu use (#693)

fishaudio · Dec 2, 2024 · 65e2531 · 65e2531
1 parent 1359896
commit 65e2531
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/tools/api.py b/tools/api.py
@@ -871,6 +871,11 @@ def initialize_app(app: Kui):
     args = parse_args()  # args same as ones in other processes
     args.precision = torch.half if args.half else torch.bfloat16
 
+    # Check if CUDA is available
+    if not torch.cuda.is_available():
+        logger.info("CUDA is not available, running on CPU.")
+        args.device = "cpu"
+
     if args.load_asr_model:
         logger.info(f"Loading ASR model...")
         asr_model = load_asr_model(device=args.device)

diff --git a/tools/webui.py b/tools/webui.py
@@ -510,6 +510,11 @@ def parse_args():
     args = parse_args()
     args.precision = torch.half if args.half else torch.bfloat16
 
+    # Check if CUDA is available
+    if not torch.cuda.is_available():
+        logger.info("CUDA is not available, running on CPU.")
+        args.device = "cpu"
+
     logger.info("Loading Llama model...")
     llama_queue = launch_thread_safe_queue(
         checkpoint_path=args.llama_checkpoint_path,