From 44d450a4b7739aa9e0bb35a41b72a48f454a875b Mon Sep 17 00:00:00 2001
From: erew123 <35898566+erew123@users.noreply.github.com>
Date: Sun, 24 Nov 2024 02:19:11 +0000
Subject: [PATCH] TTS-Diff Adds whisper version support

---
 system/tts_diff/tts_diff.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/system/tts_diff/tts_diff.py b/system/tts_diff/tts_diff.py
index 962e6fd..6690f31 100644
--- a/system/tts_diff/tts_diff.py
+++ b/system/tts_diff/tts_diff.py
@@ -28,6 +28,8 @@ parser = argparse.ArgumentParser(description="Compare TTS output with the origin
 parser.add_argument("--threshold", type=int, default=98, help="Similarity threshold for considering a match (default: 98)")
 parser.add_argument("--ttslistpath", help="Path to the ttsList.json file")
 parser.add_argument("--wavfilespath", help="Path to the wav outputs folder")
+parser.add_argument("--whisper-model", type=str, default="large-v3", help="Whisper model to use (default: large-v3)")
+
 args = parser.parse_args()
 
 if not args.ttslistpath:
@@ -38,6 +40,7 @@ if not args.wavfilespath:
 json_file_path = Path(args.ttslistpath).resolve()
 wav_file_path = Path(args.wavfilespath).resolve()
 accuracy_threshold = args.threshold
+whisper_model_name = args.whisper_model  # Capture the model name from the argument
 
 def disclaimer_text():
     print(f"\nDESCRIPTION: ")
@@ -196,12 +199,13 @@ def main():
     print("[AllTalk TTSDiff] \033[94mJSON file  :\033[92m", json_file_path, "\033[0m")
     print("[AllTalk TTSDiff] \033[94mWAV files  :\033[92m", wav_file_path, "\033[0m")
     print("[AllTalk TTSDiff] \033[94mAccuracy   :\033[92m", accuracy_threshold, "%", "\033[0m")
+    print("[AllTalk TTSDiff] \033[94mWhisper    :\033[92m", whisper_model_name, "\033[0m")
     
     device = "cuda" if torch.cuda.is_available() else "cpu"
     print("[AllTalk TTSDiff] \033[94mDevice     :\033[92m", device, "\033[0m")
-    print("[AllTalk TTSDiff] Loading Whisper Large-v3 model...")
-    model = whisper.load_model("large-v3", device=device)
-    
+    print("[AllTalk TTSDiff] Loading Whisper model...")
+    model = whisper.load_model(whisper_model_name, device=device)  # Use the specified or default model
+        
     flagged_ids = []  # Initialize the list to track IDs needing review
 
     try: