修复了OutOfMemoryError时,显存无法释放的问题

This commit is contained in:
chasonjiang 2024-03-13 16:25:27 +08:00
parent f2cbc826c7
commit d60d8ea3fb

View File

@ -2,6 +2,7 @@ from copy import deepcopy
import math
import os, sys
import random
import traceback
now_dir = os.getcwd()
sys.path.append(now_dir)
import ffmpeg
@ -48,7 +49,17 @@ custom:
"""
# def set_seed(seed):
# random.seed(seed)
# os.environ['PYTHONHASHSEED'] = str(seed)
# np.random.seed(seed)
# torch.manual_seed(seed)
# torch.cuda.manual_seed(seed)
# torch.cuda.manual_seed_all(seed)
# torch.backends.cudnn.deterministic = True
# torch.backends.cudnn.benchmark = False
# torch.backends.cudnn.enabled = True
# set_seed(1234)
class TTS_Config:
default_configs={
@ -630,7 +641,7 @@ class TTS:
split_bucket=split_bucket
)
t2 = ttime()
try:
print("############ 推理 ############")
###### inference ######
t_34 = 0.0
@ -741,14 +752,30 @@ class TTS:
batch_index_list,
speed_factor,
split_bucket)
except Exception as e:
traceback.print_exc()
# 必须返回一个空音频, 否则会导致显存不释放。
yield self.configs.sampling_rate, np.zeros(int(self.configs.sampling_rate),
dtype=np.int16)
# 重置模型, 否则会导致显存释放不完全。
del self.t2s_model
del self.vits_model
self.t2s_model = None
self.vits_model = None
self.init_t2s_weights(self.configs.t2s_weights_path)
self.init_vits_weights(self.configs.vits_weights_path)
finally:
self.empty_cache()
def empty_cache(self):
try:
if str(self.configs.device) == "cuda":
torch.cuda.empty_cache()
elif str(self.configs.device) == "mps":
torch.mps.empty_cache()
except:
pass
def audio_postprocess(self,
audio:List[torch.Tensor],
sr:int,