From 7a67b8e6a95126ea937bcdc8b0bbde44e28f6a1b Mon Sep 17 00:00:00 2001
From: RockkLee <sam880103@gmail.com>
Date: Sat, 13 Jul 2024 04:37:34 +0800
Subject: [PATCH] Optimize silero_tts func and solve the dependencyes issue

---
 requirements.txt |  5 +++++
 utils/TTS.py     | 27 ++++++++++++++++++++-------
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 78cc488..f7e6b4f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,3 +10,8 @@ pytchat
 requests
 torch
 unidic
+tensorboard==2.17.0
+tensorboard-data-server==0.7.2
+tensorflow==2.17.0
+tensorflow-intel==2.17.0
+tensorflow-io-gcs-filesystem==0.31.0
diff --git a/utils/TTS.py b/utils/TTS.py
index e8ca6ac..d02f60e 100644
--- a/utils/TTS.py
+++ b/utils/TTS.py
@@ -4,17 +4,29 @@
 import urllib.parse
 from utils.katakana import *
 
+num_threads = os.cpu_count()
+device = torch.device('cpu')
+torch.set_num_threads(num_threads)
+local_file = 'model.pt'
+print(f"Using {num_threads} threads for silero_tts")
+
+cache_model = None
+if os.path.isfile(local_file):
+    print("Loading model.pt with cache...")
+    cache_model = torch.package.PackageImporter(local_file).load_pickle("tts_models", "model")
+    cache_model.to(device)
+
+
 # https://github.com/snakers4/silero-models#text-to-speech
 def silero_tts(tts, language, model, speaker):
-    device = torch.device('cpu')
-    torch.set_num_threads(4)
-    local_file = 'model.pt'
-
     if not os.path.isfile(local_file):
+        print("Downloading model.pt...")
         torch.hub.download_url_to_file(f'https://models.silero.ai/models/tts/{language}/{model}.pt',
                                     local_file)  
-
-    model = torch.package.PackageImporter(local_file).load_pickle("tts_models", "model")
+        print("Loading model.pt...")
+        model = torch.package.PackageImporter(local_file).load_pickle("tts_models", "model")
+    else:
+        model = cache_model    
     model.to(device)
 
     example_text = "i'm fine thank you and you?"
@@ -23,7 +35,8 @@ def silero_tts(tts, language, model, speaker):
     audio_paths = model.save_wav(text=tts,
                                 speaker=speaker,
                                 sample_rate=sample_rate)
-    
+
+
 def voicevox_tts(tts):
     # You need to run VoicevoxEngine.exe first before running this script