torch 更新

tool/__init__.py

@@ -1,4 +1,5 @@
 from .av import load_audio
 from .pcm import pcm_arr_to_mp3_view
 from .np import float_to_int16
-from .ctx import TorchSeedContext
+from .ctx import TorchSeedContext
+from .gpu import select_device

tool/gpu.py

@@ -0,0 +1,48 @@
+import torch
+import os, sys
+import spaces
+
+if sys.platform == "darwin":
+    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+from tool.logger import get_logger
+
+logger = get_logger("gpu")
+
+
+def select_device(min_memory=2047, experimental=False):
+    if torch.cuda.is_available():
+        selected_gpu = 0
+        max_free_memory = -1
+        for i in range(torch.cuda.device_count()):
+            props = torch.cuda.get_device_properties(i)
+            free_memory = props.total_memory - torch.cuda.memory_reserved(i)
+            if max_free_memory < free_memory:
+                selected_gpu = i
+                max_free_memory = free_memory
+        free_memory_mb = max_free_memory / (1024 * 1024)
+        if free_memory_mb < min_memory:
+            logger.get_logger().warning(
+                f"GPU {selected_gpu} has {round(free_memory_mb, 2)} MB memory left. Switching to CPU."
+            )
+            device = torch.device("cpu")
+        else:
+            device = torch.device(f"cuda:{selected_gpu}")
+    elif torch.backends.mps.is_available():
+        """
+        Currently MPS is slower than CPU while needs more memory and core utility,
+        so only enable this for experimental use.
+        """
+        if experimental:
+            # For Apple M1/M2 chips with Metal Performance Shaders
+            logger.get_logger().warn("experimantal: found apple GPU, using MPS.")
+            device = torch.device("mps")
+        else:
+            logger.get_logger().info("found Apple GPU, but use CPU.")
+            device = torch.device("cpu")
+    else:
+        logger.get_logger().warning("no GPU found, use CPU instead")
+        device = torch.device("cpu")
+
+    return device

web/app.py

@@ -9,6 +9,7 @@ sys.path.append(now_dir)
 from tool.logger import get_logger
 from tool.func import *
 from tool.np import *
+from tool.gpu import select_device
 import ChatTTS
 import argparse
 import torch._dynamo
@@ -33,8 +34,11 @@ def init_chat(args):
     if MODEL == "HF":
         source = "huggingface"
 
+    device=select_device()
 
-    if chat.load(source=source, custom_path="D:\\chenjgspace\\ai-model\\chattts",device=torch.device(f"cuda:0")):
+    logger.info("loading ChatTTS device :" + str(device))
+
+    if chat.load(source=source, custom_path="D:\\chenjgspace\\ai-model\\chattts", device=device):
         print("Models loaded successfully.")
         logger.info("Models loaded successfully.")
     else: