Fix: Ensure memory is cleared on unload

This commit is contained in:
tsukumi
2024-11-10 07:06:20 +09:00
parent 9723f75f99
commit 20b0b0908b
2 changed files with 21 additions and 7 deletions

View File

@@ -237,9 +237,9 @@ def unload_model(language: Languages) -> None:
if language in __loaded_models:
del __loaded_models[language]
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
gc.collect()
logger.info(f"Unloaded the {language.name} BERT model")
@@ -256,8 +256,6 @@ def unload_tokenizer(language: Languages) -> None:
if language in __loaded_tokenizers:
del __loaded_tokenizers[language]
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
logger.info(f"Unloaded the {language.name} BERT tokenizer")

View File

@@ -1,5 +1,6 @@
from __future__ import annotations
import gc
import time
from pathlib import Path
from typing import TYPE_CHECKING, Any, Optional, Sequence, Union
@@ -226,13 +227,30 @@ class TTSModel:
def unload(self) -> None:
"""
音声合成モデルをデバイスからアンロードする。
PyTorch モデルの場合は CUDA メモリも解放される。
"""
import torch
start_time = time.time()
# PyTorch 推論時
if self.net_g is not None:
del self.net_g
self.net_g = None
# CUDA キャッシュをクリア
if torch.cuda.is_available():
torch.cuda.empty_cache()
# ONNX 推論時
if self.onnx_session is not None:
del self.onnx_session
self.onnx_session = None
gc.collect()
logger.info(f"Model unloaded successfully ({time.time() - start_time:.2f}s)")
def get_style_vector(self, style_id: int, weight: float = 1.0) -> NDArray[Any]:
"""
スタイルベクトルを取得する。
@@ -443,8 +461,7 @@ class TTSModel:
# 改行ごとに分割して音声を生成
else:
texts = text.split("\n")
texts = [t for t in texts if t != ""]
texts = [t for t in text.split("\n") if t != ""]
audios = []
with torch.no_grad():
for i, t in enumerate(texts):
@@ -504,8 +521,7 @@ class TTSModel:
# 改行ごとに分割して音声を生成
else:
texts = text.split("\n")
texts = [t for t in texts if t != ""]
texts = [t for t in text.split("\n") if t != ""]
audios = []
for i, t in enumerate(texts):
audios.append(