allow inference to input and output style vectors

tts/StyleTTS2/fakeyou_infer.py

@@ -1,41 +1,48 @@
 # FakeYou StyleTTS 2 Inference
 
+import os
+import torch
 import argparse
 import numpy as np
 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):
+def clsynthesize(text_file_path, voice_path, vcsteps, style_npz_path=None, npz_output_path=None):
     # Load the text from the file
     with open(text_file_path, 'r', encoding="utf-8") as file:
         text = file.read()
 
-    # Load the custom voice
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    custom_voice = compute_style(voice_path)
+
+    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")
 
-    # Split and recombine long text into manageable chunks
     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]
-    audios = []
-    for t in texts:
-        # Generate audio for each text chunk
-        audios.append(inference(t, custom_voice, alpha=0.3, beta=0.7, diffusion_steps=vcsteps, embedding_scale=1))
 
     return 24000, np.concatenate(audios)
 
 def main():
     parser = argparse.ArgumentParser(description='CLI for StyleTTS Voice Synthesis')
     parser.add_argument('--input', type=str, required=True, help='Path to the text file')
-    parser.add_argument('--voice', type=str, required=True, help='Path to the custom voice audio file')
+    parser.add_argument('--voice', type=str, help='Path to the custom voice audio file (optional if using --input-style-npz)')
     parser.add_argument('--vcsteps', type=int, required=True, help='Number of diffusion steps')
-    parser.add_argument('--output', type=str, required=True, help='Output path for the synthesized audio file')
+    parser.add_argument('--output', type=str, help='Output path for the synthesized audio file (optional)')
+    parser.add_argument('--input-style-npz', type=str, help='Path to the input style vector NPZ file')
+    parser.add_argument('--output-style-npz', type=str, help='Output path for the style vector NPZ file')
     args = parser.parse_args()
 
-    sample_rate, audio = clsynthesize(args.input, args.voice, args.vcsteps)
+    if args.output:
-
+        sample_rate, audio = clsynthesize(args.input, args.voice, args.vcsteps, args.input_style_npz, args.output_style_npz)
-    # Save the synthesized audio to the specified output path
+        sf.write(args.output, audio, sample_rate)
-    sf.write(args.output, audio, sample_rate)
 
 if __name__ == "__main__":
     main()

tts/StyleTTS2/styletts2importable.py

@@ -34,8 +34,6 @@ import nltk
 nltk.download('punkt')
 print("SCIPY")
 from scipy.io.wavfile import write
-print("TORCH STUFF")
-import torch
 print("START")
 torch.manual_seed(0)
 torch.backends.cudnn.benchmark = False
@@ -81,7 +79,7 @@ def preprocess(wave):
     mel_tensor = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - mean) / std
     return mel_tensor
 
-def compute_style(path):
+def compute_style(path, save_npz=False, npz_path=None):
     wave, sr = librosa.load(path, sr=24000)
     audio, index = librosa.effects.trim(wave, top_db=30)
     if sr != 24000:
@@ -91,8 +89,13 @@ def compute_style(path):
     with torch.no_grad():
         ref_s = model.style_encoder(mel_tensor.unsqueeze(1))
         ref_p = model.predictor_encoder(mel_tensor.unsqueeze(1))
+        style_vector = torch.cat([ref_s, ref_p], dim=1)
 
-    return torch.cat([ref_s, ref_p], dim=1)
+    # Save style vector to NPZ file if requested
+    if save_npz and npz_path:
+        np.savez(npz_path, style_vector=style_vector.cpu().numpy())
+
+    return style_vector
 
 device = 'cpu'
 if torch.cuda.is_available():