ソースを参照

Fixed faulty mps device checking logic (#743)

Stephen Abrahamsen 1 年間 前
コミット
e42eb1ec55
2 ファイル変更4 行追加8 行削除
  1. 2 4
      tools/run_webui.py
  2. 2 4
      tools/server/model_manager.py

+ 2 - 4
tools/run_webui.py

@@ -45,13 +45,11 @@ if __name__ == "__main__":
     args = parse_args()
     args.precision = torch.half if args.half else torch.bfloat16
 
-    # Check if MPS is available
+    # Check if MPS or CUDA is available
     if torch.backends.mps.is_available():
         args.device = "mps"
         logger.info("mps is available, running on mps.")
-
-    # Check if CUDA is available
-    if not torch.cuda.is_available():
+    elif not torch.cuda.is_available():
         logger.info("CUDA is not available, running on CPU.")
         args.device = "cpu"
 

+ 2 - 4
tools/server/model_manager.py

@@ -34,13 +34,11 @@ class ModelManager:
 
         self.precision = torch.half if half else torch.bfloat16
 
-        # Check if MPS is available
+        # Check if MPS or CUDA is available
         if torch.backends.mps.is_available():
             self.device = "mps"
             logger.info("mps is available, running on mps.")
-
-        # Check if CUDA is available
-        if not torch.cuda.is_available():
+        elif not torch.cuda.is_available():
             self.device = "cpu"
             logger.info("CUDA is not available, running on CPU.")