zhzluke96
commited on
Commit
•
8f52106
1
Parent(s):
0129fb6
update
Browse files- modules/models.py +35 -8
modules/models.py
CHANGED
@@ -1,18 +1,23 @@
|
|
|
|
1 |
import torch
|
2 |
from modules.ChatTTS import ChatTTS
|
3 |
from modules import config
|
4 |
from modules.devices import devices
|
5 |
|
6 |
import logging
|
|
|
7 |
|
8 |
logger = logging.getLogger(__name__)
|
|
|
9 |
chat_tts = None
|
|
|
10 |
|
11 |
|
12 |
-
def
|
13 |
global chat_tts
|
14 |
if chat_tts:
|
15 |
-
|
|
|
16 |
|
17 |
chat_tts = ChatTTS.Chat()
|
18 |
chat_tts.load_models(
|
@@ -28,18 +33,40 @@ def load_chat_tts():
|
|
28 |
)
|
29 |
|
30 |
devices.torch_gc()
|
|
|
|
|
|
|
|
|
|
|
|
|
31 |
|
|
|
|
|
|
|
|
|
|
|
32 |
return chat_tts
|
33 |
|
34 |
|
35 |
-
def
|
36 |
-
logging.info("
|
37 |
global chat_tts
|
|
|
38 |
if chat_tts:
|
|
|
|
|
|
|
|
|
39 |
if torch.cuda.is_available():
|
40 |
-
for model_name, model in chat_tts.pretrain_models.items():
|
41 |
-
if isinstance(model, torch.nn.Module):
|
42 |
-
model.cpu()
|
43 |
torch.cuda.empty_cache()
|
|
|
44 |
chat_tts = None
|
45 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import threading
|
2 |
import torch
|
3 |
from modules.ChatTTS import ChatTTS
|
4 |
from modules import config
|
5 |
from modules.devices import devices
|
6 |
|
7 |
import logging
|
8 |
+
import gc
|
9 |
|
10 |
logger = logging.getLogger(__name__)
|
11 |
+
|
12 |
chat_tts = None
|
13 |
+
load_event = threading.Event()
|
14 |
|
15 |
|
16 |
+
def load_chat_tts_in_thread():
|
17 |
global chat_tts
|
18 |
if chat_tts:
|
19 |
+
load_event.set() # 如果已经加载过,直接设置事件
|
20 |
+
return
|
21 |
|
22 |
chat_tts = ChatTTS.Chat()
|
23 |
chat_tts.load_models(
|
|
|
33 |
)
|
34 |
|
35 |
devices.torch_gc()
|
36 |
+
load_event.set() # 设置事件,表示加载完成
|
37 |
+
|
38 |
+
|
39 |
+
def initialize_chat_tts():
|
40 |
+
model_thread = threading.Thread(target=load_chat_tts_in_thread)
|
41 |
+
model_thread.start()
|
42 |
|
43 |
+
|
44 |
+
def load_chat_tts():
|
45 |
+
if chat_tts is None:
|
46 |
+
initialize_chat_tts()
|
47 |
+
load_event.wait()
|
48 |
return chat_tts
|
49 |
|
50 |
|
51 |
+
def unload_chat_tts():
|
52 |
+
logging.info("Unloading ChatTTS models")
|
53 |
global chat_tts
|
54 |
+
|
55 |
if chat_tts:
|
56 |
+
for model_name, model in chat_tts.pretrain_models.items():
|
57 |
+
if isinstance(model, torch.nn.Module):
|
58 |
+
model.cpu()
|
59 |
+
del model
|
60 |
if torch.cuda.is_available():
|
|
|
|
|
|
|
61 |
torch.cuda.empty_cache()
|
62 |
+
gc.collect()
|
63 |
chat_tts = None
|
64 |
+
logger.info("ChatTTS models unloaded")
|
65 |
+
|
66 |
+
|
67 |
+
def reload_chat_tts():
|
68 |
+
logging.info("Reloading ChatTTS models")
|
69 |
+
unload_chat_tts()
|
70 |
+
instance = load_chat_tts()
|
71 |
+
logger.info("ChatTTS models reloaded")
|
72 |
+
return instance
|