feat(chat): add model chack to custom load (#453)

ChatTTS/core.py

@@ -5,6 +5,7 @@
 from typing import Literal, Optional, List, Callable, Tuple, Dict
 from functools import lru_cache
 from json import load
+from pathlib import Path
 
 import numpy as np
 import torch
@@ -107,6 +108,11 @@ def download_models(
                 return None
         elif source == "custom":
             self.logger.log(logging.INFO, f"try to load from local: {custom_path}")
+            if not check_all_assets(self.sha256_map, update=False, base_dir=Path(custom_path)):
+                self.logger.error(
+                    "check models in custom path %s failed.", custom_path
+                )
+                return None
             download_path = custom_path
 
         return download_path

ChatTTS/utils/dl.py

@@ -42,11 +42,9 @@ def check_model(
     return True
 
 
-def check_all_assets(sha256_map: dict[str, str], update=False) -> bool:
-    BASE_DIR = Path(os.getcwd())
-
+def check_all_assets(sha256_map: dict[str, str], update=False, base_dir = Path(os.getcwd())) -> bool:
     logger.get_logger().info("checking assets...")
-    current_dir = BASE_DIR / "asset"
+    current_dir = base_dir / "asset"
     names = [
         "Decoder.pt",
         "DVAE.pt",
@@ -63,7 +61,7 @@ def check_all_assets(sha256_map: dict[str, str], update=False) -> bool:
             return False
 
     logger.get_logger().info("checking configs...")
-    current_dir = BASE_DIR / "config"
+    current_dir = base_dir / "config"
     names = [
         "decoder.yaml",
         "dvae.yaml",