From e3c0b6073edc635529403f1610c30adbdb7b9744 Mon Sep 17 00:00:00 2001
From: Justin John <34035011+justinjohn0306@users.noreply.github.com>
Date: Fri, 12 Jan 2024 11:57:11 +0530
Subject: [PATCH] fixed some oopsies

---
 tts/StyleTTS2/fakeyou_infer.py | 33 +++++++++++++++++++--------------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/tts/StyleTTS2/fakeyou_infer.py b/tts/StyleTTS2/fakeyou_infer.py
index 34d53d4..99ab330 100644
--- a/tts/StyleTTS2/fakeyou_infer.py
+++ b/tts/StyleTTS2/fakeyou_infer.py
@@ -8,23 +8,11 @@ import soundfile as sf
 from styletts2importable import compute_style, inference
 from tortoise.utils.text import split_and_recombine_text
 
-def clsynthesize(text_file_path, voice_path, vcsteps, style_npz_path=None, npz_output_path=None):
+def clsynthesize(text_file_path, custom_voice, vcsteps):
     # Load the text from the file
     with open(text_file_path, 'r', encoding="utf-8") as file:
         text = file.read()
 
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-    if style_npz_path and os.path.exists(style_npz_path):
-        # Load style vector from NPZ file
-        style_vector_np = np.load(style_npz_path)['style_vector']
-        custom_voice = torch.from_numpy(style_vector_np).float().to(device)
-    elif voice_path:
-        # Compute and save style vector
-        custom_voice = compute_style(voice_path, save_npz=True, npz_path=npz_output_path or "default_output_path.npz").to(device)
-    else:
-        raise ValueError("Either a voice path or an input style NPZ file must be provided")
-
     texts = split_and_recombine_text(text)
     audios = [inference(t, custom_voice, alpha=0.3, beta=0.7, diffusion_steps=vcsteps, embedding_scale=1) for t in texts]
 
@@ -40,8 +28,25 @@ def main():
     parser.add_argument('--output-style-npz', type=str, help='Output path for the style vector NPZ file')
     args = parser.parse_args()
 
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    custom_voice = None
+    if args.input_style_npz and os.path.exists(args.input_style_npz):
+        # Load style vector from NPZ file
+        style_vector_np = np.load(args.input_style_npz)['style_vector']
+        custom_voice = torch.from_numpy(style_vector_np).float().to(device)
+    elif args.voice:
+        # Compute and save style vector
+        npz_output_path = args.output_style_npz or "default_output_path.npz"
+        custom_voice = compute_style(args.voice, save_npz=True, npz_path=npz_output_path).to(device)
+
+    # Checking if the custom_voice tensor is initialized
+    if custom_voice is None:
+        raise ValueError("Unable to load or compute style vector")
+
+    # Perform synthesis and save output if the output path is provided
     if args.output:
-        sample_rate, audio = clsynthesize(args.input, args.voice, args.vcsteps, args.input_style_npz, args.output_style_npz)
+        sample_rate, audio = clsynthesize(args.input, custom_voice, args.vcsteps)
         sf.write(args.output, audio, sample_rate)
 
 if __name__ == "__main__":