Browse Source

fix cpu use (#693)

Picus303 1 year ago
parent
commit
65e25315da
2 changed files with 10 additions and 0 deletions
  1. 5 0
      tools/api.py
  2. 5 0
      tools/webui.py

+ 5 - 0
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)

+ 5 - 0
tools/webui.py

@@ -510,6 +510,11 @@ if __name__ == "__main__":
     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,