Update app.py

app.py

@@ -9,12 +9,14 @@ import numpy as np
 
 import ChatTTS
 
+import se_extractor
+from api import BaseSpeakerTTS, ToneColorConverter
+
 print("loading ChatTTS model...")
 chat = ChatTTS.Chat()
 chat.load_models()
 
 
-
 def generate_seed():
     new_seed = random.randint(1, 100000000)
     return {
@@ -23,7 +25,7 @@ def generate_seed():
         }
 
 @spaces.GPU
-def generate_audio(text, temperature, top_P, top_K, audio_seed_input, text_seed_input, refine_text_flag, refine_text_input):
+def chat_tts(text, temperature, top_P, top_K, audio_seed_input, text_seed_input, refine_text_flag, refine_text_input, output_path=None):
 
     torch.manual_seed(audio_seed_input)
     rand_spk = torch.randn(768)
@@ -57,7 +59,67 @@ def generate_audio(text, temperature, top_P, top_K, audio_seed_input, text_seed_
     sample_rate = 24000
     text_data = text[0] if isinstance(text, list) else text
 
-    return [(sample_rate, audio_data), text_data]
+    if output_path is None:
+        return [(sample_rate, audio_data), text_data]
+    else:
+        soundfile.write(output_path, audio_data, sample_rate)
+
+# OpenVoice
+
+ckpt_base_en = 'checkpoints/base_speakers/EN'
+ckpt_converter_en = 'checkpoints/converter'
+device = 'cuda:0'
+
+#device = "cpu"
+
+base_speaker_tts = BaseSpeakerTTS(f'{ckpt_base_en}/config.json', device=device)
+base_speaker_tts.load_ckpt(f'{ckpt_base_en}/checkpoint.pth')
+
+tone_color_converter = ToneColorConverter(f'{ckpt_converter_en}/config.json', device=device)
+tone_color_converter.load_ckpt(f'{ckpt_converter_en}/checkpoint.pth')
+
+
+def generate_audio(text, audio_ref, style_mode, temperature, top_P, top_K, audio_seed_input, text_seed_input, refine_text_flag, refine_text_input):
+    if style_mode=="default":
+        source_se = torch.load(f'{ckpt_base_en}/en_default_se.pth').to(device)
+        reference_speaker = audio_ref
+        target_se, audio_name = se_extractor.get_se(reference_speaker, tone_color_converter, target_dir='processed', vad=True)
+        save_path = "output.wav"
+
+        # Run the base speaker tts
+        src_path = "tmp.wav"
+        chat_tts(text, text, temperature, top_P, top_K, audio_seed_input, text_seed_input, refine_text_flag, refine_text_input, output_path=None, src_path)
+
+        # Run the tone color converter
+        encode_message = "@MyShell"
+        tone_color_converter.convert(
+            audio_src_path=src_path,
+            src_se=source_se,
+            tgt_se=target_se,
+            output_path=save_path,
+            message=encode_message)
+
+    else:
+        source_se = torch.load(f'{ckpt_base_en}/en_style_se.pth').to(device)
+        reference_speaker = audio_ref
+        target_se, audio_name = se_extractor.get_se(reference_speaker, tone_color_converter, target_dir='processed', vad=True)
+
+        save_path = "output.wav"
+
+        # Run the base speaker tts
+        src_path = "tmp.wav"
+        base_speaker_tts.tts(text, src_path, speaker=style_mode, language='English', speed=0.9)
+
+        # Run the tone color converter
+        encode_message = "@MyShell"
+        tone_color_converter.convert(
+            audio_src_path=src_path,
+            src_se=source_se,
+            tgt_se=target_se,
+            output_path=save_path,
+            message=encode_message)
+
+    return "output.wav"
 
 
 with gr.Blocks() as demo:
@@ -83,7 +145,7 @@ with gr.Blocks() as demo:
 
     generate_button = gr.Button("Generate")
         
-    text_output = gr.Textbox(label="Refined Text", interactive=False)
+    #text_output = gr.Textbox(label="Refined Text", interactive=False)
     audio_output = gr.Audio(label="Output Audio")
 
     generate_audio_seed.click(generate_seed, 
@@ -96,7 +158,7 @@ with gr.Blocks() as demo:
         
     generate_button.click(generate_audio, 
                           inputs=[text_input, temperature_slider, top_p_slider, top_k_slider, audio_seed_input, text_seed_input, refine_text_checkbox, refine_text_input], 
-                          outputs=[audio_output, text_output])
+                          outputs=audio_output)
 
 parser = argparse.ArgumentParser(description='ChatTTS demo Launch')
 parser.add_argument('--server_name', type=str, default='0.0.0.0', help='Server name')