diff --git a/ChatTTS/core.py b/ChatTTS/core.py index f49139d..6553ef8 100644 --- a/ChatTTS/core.py +++ b/ChatTTS/core.py @@ -1,5 +1,5 @@ -import os +import os, sys import json import logging from functools import partial @@ -63,7 +63,7 @@ class Chat: download_all_assets(tmpdir=tmp) if not check_all_assets(update=False): logging.error("counld not satisfy all assets needed.") - exit(1) + return False elif source == 'huggingface': hf_home = os.getenv('HF_HOME', os.path.expanduser("~/.cache/huggingface")) try: @@ -79,7 +79,7 @@ class Chat: self.logger.log(logging.INFO, f'Load from local: {custom_path}') download_path = custom_path - self._load(**{k: os.path.join(download_path, v) for k, v in OmegaConf.load(os.path.join(download_path, 'config', 'path.yaml')).items()}, **kwargs) + return self._load(**{k: os.path.join(download_path, v) for k, v in OmegaConf.load(os.path.join(download_path, 'config', 'path.yaml')).items()}, **kwargs) def _load( self, @@ -148,7 +148,7 @@ class Chat: self.pretrain_models['tokenizer'] = tokenizer self.logger.log(logging.INFO, 'tokenizer loaded.') - self.check_model() + return self.check_model() def _infer( self, diff --git a/examples/cmd/run.py b/examples/cmd/run.py index 8793ec5..7af340f 100644 --- a/examples/cmd/run.py +++ b/examples/cmd/run.py @@ -31,8 +31,11 @@ def main(): chat = ChatTTS.Chat() print("Initializing ChatTTS...") - chat.load_models() - print("Models loaded successfully.") + if chat.load_models(): + print("Models loaded successfully.") + else: + print("Models load failed.") + sys.exit(1) texts = [text_input] print("Text prepared for inference:", texts) diff --git a/examples/web/webui.py b/examples/web/webui.py index d6896ab..3be454a 100644 --- a/examples/web/webui.py +++ b/examples/web/webui.py @@ -182,10 +182,17 @@ def main(): chat = ChatTTS.Chat() if args.custom_path == None: - chat.load_models() + ret = chat.load_models() else: print('local model path:', args.custom_path) - chat.load_models('custom', custom_path=args.custom_path) + ret = chat.load_models('custom', custom_path=args.custom_path) + + if ret: + print("Models loaded successfully.") + else: + print("Models load failed.") + sys.exit(1) + demo.launch(server_name=args.server_name, server_port=args.server_port, root_path=args.root_path, inbrowser=True)