From 20b0b0908bcd817e0894402e5170b024102cfe1e Mon Sep 17 00:00:00 2001
From: tsukumi <tsukumijima@users.noreply.github.com>
Date: Sun, 10 Nov 2024 07:06:20 +0900
Subject: [PATCH] Fix: Ensure memory is cleared on unload

---
 style_bert_vits2/nlp/bert_models.py |  4 +---
 style_bert_vits2/tts_model.py       | 24 ++++++++++++++++++++----
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/style_bert_vits2/nlp/bert_models.py b/style_bert_vits2/nlp/bert_models.py
index 9c05786..4640ce3 100644
--- a/style_bert_vits2/nlp/bert_models.py
+++ b/style_bert_vits2/nlp/bert_models.py
@@ -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")
 
 
diff --git a/style_bert_vits2/tts_model.py b/style_bert_vits2/tts_model.py
index 76acfe1..10bf25f 100644
--- a/style_bert_vits2/tts_model.py
+++ b/style_bert_vits2/tts_model.py
@@ -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(