Update app.py

app.py

@@ -1,8 +1,18 @@
 import sys
 from pathlib import Path
+import os
+import torch
+import openvino as ov
+import gradio as gr
+import langid
+import ipywidgets as widgets
+from IPython.display import Audio
+from openvoice.api import BaseSpeakerTTS, ToneColorConverter, OpenVoiceBaseClass
+import openvoice.se_extractor as se_extractor
+import nncf
 
+# Clone the repo and set up the environment
 repo_dir = Path("OpenVoice")
-
 if not repo_dir.exists():
     !git clone https://github.com/myshell-ai/OpenVoice
     orig_english_path = Path("OpenVoice/openvoice/text/_orig_english.py")
@@ -15,23 +25,14 @@ if not repo_dir.exists():
         data = data.replace("unidecode", "anyascii")
         with english_path.open("w") as out_f:
             out_f.write(data)
-# append to sys.path so that modules from the repo could be imported
 sys.path.append(str(repo_dir))
 
+# Install the required packages
 %pip install -q "librosa>=0.8.1" "wavmark>=0.0.3" "faster-whisper>=0.9.0" "pydub>=0.25.1" "whisper-timestamped>=1.14.2" "tqdm" "inflect>=7.0.0" "eng_to_ipa>=0.0.2" "pypinyin>=0.50.0" \
 "cn2an>=0.5.22" "jieba>=0.42.1" "langid>=1.1.6" "gradio>=4.15" "ipywebrtc" "anyascii" "openvino>=2023.3" "torch>=2.1" "nncf>=2.11.0"
 
-import os
-import torch
-import openvino as ov
-import ipywidgets as widgets
-from IPython.display import Audio
-
 core = ov.Core()
 
-from openvoice.api import BaseSpeakerTTS, ToneColorConverter, OpenVoiceBaseClass
-import openvoice.se_extractor as se_extractor
-
 CKPT_BASE_PATH = "checkpoints"
 
 en_suffix = f"{CKPT_BASE_PATH}/base_speakers/EN"
@@ -42,11 +43,9 @@ enable_chinese_lang = False
 
 def download_from_hf_hub(filename, local_dir="./"):
     from huggingface_hub import hf_hub_download
-
     os.makedirs(local_dir, exist_ok=True)
     hf_hub_download(repo_id="myshell-ai/OpenVoice", filename=filename, local_dir=local_dir)
 
-
 download_from_hf_hub(f"{converter_suffix}/checkpoint.pth")
 download_from_hf_hub(f"{converter_suffix}/config.json")
 download_from_hf_hub(f"{en_suffix}/checkpoint.pth")
@@ -68,20 +67,6 @@ en_base_speaker_tts.load_ckpt(f"{en_suffix}/checkpoint.pth")
 tone_color_converter = ToneColorConverter(f"{converter_suffix}/config.json", device=pt_device)
 tone_color_converter.load_ckpt(f"{converter_suffix}/checkpoint.pth")
 
-if enable_chinese_lang:
-    zh_base_speaker_tts = BaseSpeakerTTS(f"{zh_suffix}/config.json", device=pt_device)
-    zh_base_speaker_tts.load_ckpt(f"{zh_suffix}/checkpoint.pth")
-else:
-    zh_base_speaker_tts = None
-
-pt_device = "cpu"
-
-en_base_speaker_tts = BaseSpeakerTTS(f"{en_suffix}/config.json", device=pt_device)
-en_base_speaker_tts.load_ckpt(f"{en_suffix}/checkpoint.pth")
-
-tone_color_converter = ToneColorConverter(f"{converter_suffix}/config.json", device=pt_device)
-tone_color_converter.load_ckpt(f"{converter_suffix}/checkpoint.pth")
-
 if enable_chinese_lang:
     zh_base_speaker_tts = BaseSpeakerTTS(f"{zh_suffix}/config.json", device=pt_device)
     zh_base_speaker_tts.load_ckpt(f"{zh_suffix}/checkpoint.pth")
@@ -89,22 +74,13 @@ else:
     zh_base_speaker_tts = None
 
 class OVOpenVoiceBase(torch.nn.Module):
-    """
-    Base class for both TTS and voice tone conversion model: constructor is same for both of them.
-    """
-
     def __init__(self, voice_model: OpenVoiceBaseClass):
         super().__init__()
         self.voice_model = voice_model
         for par in voice_model.model.parameters():
             par.requires_grad = False
 
-
 class OVOpenVoiceTTS(OVOpenVoiceBase):
-    """
-    Constructor of this class accepts BaseSpeakerTTS object for speech generation and wraps it's 'infer' method with forward.
-    """
-
     def get_example_input(self):
         stn_tst = self.voice_model.get_text("this is original text", self.voice_model.hps, False)
         x_tst = stn_tst.unsqueeze(0)
@@ -125,12 +101,7 @@ class OVOpenVoiceTTS(OVOpenVoiceBase):
     def forward(self, x, x_lengths, sid, noise_scale, length_scale, noise_scale_w):
         return self.voice_model.model.infer(x, x_lengths, sid, noise_scale, length_scale, noise_scale_w)
 
-
 class OVOpenVoiceConverter(OVOpenVoiceBase):
-    """
-    Constructor of this class accepts ToneColorConverter object for voice tone conversion and wraps it's 'voice_conversion' method with forward.
-    """
-
     def get_example_input(self):
         y = torch.randn([1, 513, 238], dtype=torch.float32)
         y_lengths = torch.LongTensor([y.size(-1)])
@@ -142,9 +113,6 @@ class OVOpenVoiceConverter(OVOpenVoiceBase):
     def forward(self, y, y_lengths, sid_src, sid_tgt, tau):
         return self.voice_model.model.voice_conversion(y, y_lengths, sid_src, sid_tgt, tau)
 
-import nncf
-
-
 IRS_PATH = "openvino_irs/"
 EN_TTS_IR = f"{IRS_PATH}/openvoice_en_tts.xml"
 ZH_TTS_IR = f"{IRS_PATH}/openvoice_zh_tts.xml"
@@ -180,46 +148,11 @@ reference_speakers = [
     "load_manually",
 ]
 
-ref_speaker = widgets.Dropdown(
-    options=reference_speakers,
-    value=reference_speakers[0],
-    description="reference voice from which tone color will be copied",
-    disabled=False,
-)
-
-ref_speaker
-
-OUTPUT_DIR = "outputs/"
-os.makedirs(OUTPUT_DIR, exist_ok=True)
-
-ref_speaker_path = f"{REFERENCE_VOICES_PATH}/{ref_speaker.value}"
-allowed_audio_types = ".mp4,.mp3,.wav,.wma,.aac,.m4a,.m4b,.webm"
-
-if ref_speaker.value == "record_manually":
-    ref_speaker_path = f"{OUTPUT_DIR}/custom_example_sample.webm"
-    from ipywebrtc import AudioRecorder, CameraStream
-
-    camera = CameraStream(constraints={"audio": True, "video": False})
-    recorder = AudioRecorder(stream=camera, filename=ref_speaker_path, autosave=True)
-    display(recorder)
-elif ref_speaker.value == "load_manually":
-    upload_ref = widgets.FileUpload(
-        accept=allowed_audio_types,
-        multiple=False,
-        description="Select audio with reference voice",
-    )
-    display(upload_ref)
-
 def save_audio(voice_source: widgets.FileUpload, out_path: str):
     with open(out_path, "wb") as output_file:
         assert len(voice_source.value) > 0, "Please select audio file"
         output_file.write(voice_source.value[0]["content"])
 
-
-if ref_speaker.value == "load_manually":
-    ref_speaker_path = f"{OUTPUT_DIR}/{upload_ref.value[0].name}"
-    save_audio(upload_ref, ref_speaker_path)
-
 en_source_default_se = torch.load(f"{en_suffix}/en_default_se.pth")
 en_source_style_se = torch.load(f"{en_suffix}/en_style_se.pth")
 zh_source_se = torch.load(f"{zh_suffix}/zh_default_se.pth") if enable_chinese_lang else None
@@ -235,7 +168,6 @@ def get_pathched_infer(ov_model: ov.Model, device: str) -> callable:
 
     return infer_impl
 
-
 def get_patched_voice_conversion(ov_model: ov.Model, device: str) -> callable:
     compiled_model = core.compile_model(ov_model, device)
 
@@ -245,71 +177,13 @@ def get_patched_voice_conversion(ov_model: ov.Model, device: str) -> callable:
 
     return voice_conversion_impl
 
-
 en_base_speaker_tts.model.infer = get_pathched_infer(ov_en_tts, device.value)
 tone_color_converter.model.voice_conversion = get_patched_voice_conversion(ov_voice_conversion, device.value)
 if enable_chinese_lang:
     zh_base_speaker_tts.model.infer = get_pathched_infer(ov_zh_tts, device.value)
 
-voice_source = widgets.Dropdown(
-    options=["use TTS", "choose_manually"],
-    value="use TTS",
-    description="Voice source",
-    disabled=False,
-)
-
-voice_source
-
-if voice_source.value == "choose_manually":
-    upload_orig_voice = widgets.FileUpload(
-        accept=allowed_audio_types,
-        multiple=False,
-        description="audo whose tone will be replaced",
-    )
-    display(upload_orig_voice)
-
-if voice_source.value == "choose_manually":
-    orig_voice_path = f"{OUTPUT_DIR}/{upload_orig_voice.value[0].name}"
-    save_audio(upload_orig_voice, orig_voice_path)
-    source_se, _ = se_extractor.get_se(orig_voice_path, tone_color_converter, target_dir=OUTPUT_DIR, vad=True)
-else:
-    text = """
-    OpenVINO toolkit is a comprehensive toolkit for quickly developing applications and solutions that solve
-    a variety of tasks including emulation of human vision, automatic speech recognition, natural language processing,
-    recommendation systems, and many others.
-    """
-    source_se = en_source_default_se
-    orig_voice_path = f"{OUTPUT_DIR}/tmp.wav"
-    en_base_speaker_tts.tts(text, orig_voice_path, speaker="default", language="English")
-
-tau_slider = widgets.FloatSlider(
-    value=0.3,
-    min=0.01,
-    max=2.0,
-    step=0.01,
-    description="tau",
-    disabled=False,
-    readout_format=".2f",
-)
-tau_slider
-
-resulting_voice_path = f"{OUTPUT_DIR}/output_with_cloned_voice_tone.wav"
-
-tone_color_converter.convert(
-    audio_src_path=orig_voice_path,
-    src_se=source_se,
-    tgt_se=target_se,
-    output_path=resulting_voice_path,
-    tau=tau_slider.value,
-    message="@MyShell",
-)
-
-import gradio as gr
-import langid
-
 supported_languages = ["zh", "en"]
 
-
 def build_predict(
     output_dir,
     tone_color_converter,
@@ -318,273 +192,62 @@ def build_predict(
     en_source_default_se,
     en_source_style_se,
     zh_source_se,
+    supported_languages,
 ):
-    def predict(prompt, style, audio_file_pth, agree):
-        return predict_impl(
-            prompt,
-            style,
-            audio_file_pth,
-            agree,
-            output_dir,
-            tone_color_converter,
-            en_tts_model,
-            zh_tts_model,
-            en_source_default_se,
-            en_source_style_se,
-            zh_source_se,
-        )
-
-    return predict
-
-
-def predict_impl(
-    prompt,
-    style,
-    audio_file_pth,
-    agree,
-    output_dir,
-    tone_color_converter,
-    en_tts_model,
-    zh_tts_model,
-    en_source_default_se,
-    en_source_style_se,
-    zh_source_se,
-):
-    text_hint = ""
-    if not agree:
-        text_hint += "[ERROR] Please accept the Terms & Condition!\n"
-        gr.Warning("Please accept the Terms & Condition!")
-        return (
-            text_hint,
-            None,
-            None,
-        )
-
-    language_predicted = langid.classify(prompt)[0].strip()
-    print(f"Detected language:{language_predicted}")
-
-    if language_predicted not in supported_languages:
-        text_hint += f"[ERROR] The detected language {language_predicted} for your input text is not in our Supported Languages: {supported_languages}\n"
-        gr.Warning(f"The detected language {language_predicted} for your input text is not in our Supported Languages: {supported_languages}")
+    def predict(
+        input_text,
+        reference_audio,
+        speaker,
+        noise_scale=0.667,
+        length_scale=1.0,
+        noise_scale_w=0.8,
+        tone_color=False,
+    ):
+        if reference_audio:
+            ref_audio_path = f"{output_dir}/input_audio.wav"
+            save_audio(reference_audio, ref_audio_path)
+            target_se, _ = se_extractor.get_se(ref_audio_path, tone_color_converter, target_dir=output_dir, vad=True)
+        else:
+            if speaker == "record_manually":
+                raise ValueError("Manual recording is not implemented in this example.")
+            elif speaker == "load_manually":
+                raise ValueError("Loading a manual audio file is not implemented in this example.")
+            else:
+                ref_audio_path = f"{REFERENCE_VOICES_PATH}/{speaker}"
+                target_se, _ = se_extractor.get_se(ref_audio_path, tone_color_converter, target_dir=output_dir, vad=True)
+        
+        lang = langid.classify(input_text)[0]
+        if lang not in supported_languages:
+            return f"Unsupported language: {lang}"
+
+        tts_model = en_tts_model if lang == "en" else zh_tts_model
+
+        stn_tst = tts_model.get_text(input_text, tts_model.hps, False)
+        x_tst = stn_tst.unsqueeze(0)
+        x_tst_lengths = torch.LongTensor([stn_tst.size(0)])
+        speaker_id = torch.LongTensor([1])
+        noise_scale = torch.tensor(noise_scale)
+        length_scale = torch.tensor(length_scale)
+        noise_scale_w = torch.tensor(noise_scale_w)
 
-        return (
-            text_hint,
-            None,
-        )
+        with torch.no_grad():
+            audio = tts_model.model.infer(x_tst, x_tst_lengths, speaker_id, noise_scale, length_scale, noise_scale_w)[0]
+            if tone_color:
+                source_se = en_source_style_se if lang == "en" else zh_source_se
+                audio = tone_color_converter.model.voice_conversion(audio, x_tst_lengths, source_se, target_se, torch.tensor(0.3))[0]
 
-    if language_predicted == "zh":
-        tts_model = zh_tts_model
-        if zh_tts_model is None:
-            gr.Warning("TTS model for Chinece language was not loaded please set 'enable_chinese_lang=True`")
-            return (
-                text_hint,
-                None,
-            )
-        source_se = zh_source_se
-        language = "Chinese"
-        if style not in ["default"]:
-            text_hint += f"[ERROR] The style {style} is not supported for Chinese, which should be in ['default']\n"
-            gr.Warning(f"The style {style} is not supported for Chinese, which should be in ['default']")
-            return (
-                text_hint,
-                None,
-            )
-    else:
-        tts_model = en_tts_model
-        if style == "default":
-            source_se = en_source_default_se
-        else:
-            source_se = en_source_style_se
-        language = "English"
-        supported_styles = [
-            "default",
-            "whispering",
-            "shouting",
-            "excited",
-            "cheerful",
-            "terrified",
-            "angry",
-            "sad",
-            "friendly",
-        ]
-        if style not in supported_styles:
-            text_hint += f"[ERROR] The style {style} is not supported for English, which should be in {*supported_styles,}\n"
-            gr.Warning(f"The style {style} is not supported for English, which should be in {*supported_styles,}")
-            return (
-                text_hint,
-                None,
-            )
-
-    speaker_wav = audio_file_pth
-
-    if len(prompt) < 2:
-        text_hint += "[ERROR] Please give a longer prompt text \n"
-        gr.Warning("Please give a longer prompt text")
-        return (
-            text_hint,
-            None,
-        )
-    if len(prompt) > 200:
-        text_hint += (
-            "[ERROR] Text length limited to 200 characters for this demo, please try shorter text. You can clone our open-source repo and try for your usage \n"
-        )
-        gr.Warning("Text length limited to 200 characters for this demo, please try shorter text. You can clone our open-source repo for your usage")
-        return (
-            text_hint,
-            None,
-        )
+        audio = audio.squeeze().cpu().numpy()
+        output_path = f"{output_dir}/output_audio.wav"
+        Audio(audio, rate=tts_model.hps.data.sampling_rate).save(output_path)
 
-    # note diffusion_conditioning not used on hifigan (default mode), it will be empty but need to pass it to model.inference
-    try:
-        target_se, audio_name = se_extractor.get_se(speaker_wav, tone_color_converter, target_dir=OUTPUT_DIR, vad=True)
-    except Exception as e:
-        text_hint += f"[ERROR] Get target tone color error {str(e)} \n"
-        gr.Warning("[ERROR] Get target tone color error {str(e)} \n")
-        return (
-            text_hint,
-            None,
-        )
+        return output_path
 
-    src_path = f"{output_dir}/tmp.wav"
-    tts_model.tts(prompt, src_path, speaker=style, language=language)
-
-    save_path = f"{output_dir}/output.wav"
-    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,
-    )
-
-    text_hint += "Get response successfully \n"
-
-    return (
-        text_hint,
-        src_path,
-        save_path,
-    )
-
-
-description = """
-    # OpenVoice accelerated by OpenVINO:
-
-    a versatile instant voice cloning approach that requires only a short audio clip from the reference speaker to replicate their voice and generate speech in multiple languages. OpenVoice enables granular control over voice styles, including emotion, accent, rhythm, pauses, and intonation, in addition to replicating the tone color of the reference speaker. OpenVoice also achieves zero-shot cross-lingual voice cloning for languages not included in the massive-speaker training set.
-"""
-
-content = """
-<div>
-<strong>If the generated voice does not sound like the reference voice, please refer to <a href='https://github.com/myshell-ai/OpenVoice/blob/main/docs/QA.md'>this QnA</a>.</strong> <strong>For multi-lingual & cross-lingual examples, please refer to <a href='https://github.com/myshell-ai/OpenVoice/blob/main/demo_part2.ipynb'>this jupyter notebook</a>.</strong>
-This online demo mainly supports <strong>English</strong>. The <em>default</em> style also supports <strong>Chinese</strong>. But OpenVoice can adapt to any other language as long as a base speaker is provided.
-</div>
-"""
-wrapped_markdown_content = f"<div style='border: 1px solid #000; padding: 10px;'>{content}</div>"
-
-
-examples = [
-    [
-        "今天天气真好，我们一起出去吃饭吧。",
-        "default",
-        "OpenVoice/resources/demo_speaker1.mp3",
-        True,
-    ],
-    [
-        "This audio is generated by open voice with a half-performance model.",
-        "whispering",
-        "OpenVoice/resources/demo_speaker2.mp3",
-        True,
-    ],
-    [
-        "He hoped there would be stew for dinner, turnips and carrots and bruised potatoes and fat mutton pieces to be ladled out in thick, peppered, flour-fattened sauce.",
-        "sad",
-        "OpenVoice/resources/demo_speaker0.mp3",
-        True,
-    ],
-]
+    return predict
 
+OUTPUT_DIR = "output_audio"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-def get_demo(
-    output_dir,
-    tone_color_converter,
-    en_tts_model,
-    zh_tts_model,
-    en_source_default_se,
-    en_source_style_se,
-    zh_source_se,
-):
-    with gr.Blocks(analytics_enabled=False) as demo:
-        with gr.Row():
-            gr.Markdown(description)
-        with gr.Row():
-            gr.HTML(wrapped_markdown_content)
-
-        with gr.Row():
-            with gr.Column():
-                input_text_gr = gr.Textbox(
-                    label="Text Prompt",
-                    info="One or two sentences at a time is better. Up to 200 text characters.",
-                    value="He hoped there would be stew for dinner, turnips and carrots and bruised potatoes and fat mutton pieces to be ladled out in thick, peppered, flour-fattened sauce.",
-                )
-                style_gr = gr.Dropdown(
-                    label="Style",
-                    info="Select a style of output audio for the synthesised speech. (Chinese only support 'default' now)",
-                    choices=[
-                        "default",
-                        "whispering",
-                        "cheerful",
-                        "terrified",
-                        "angry",
-                        "sad",
-                        "friendly",
-                    ],
-                    max_choices=1,
-                    value="default",
-                )
-                ref_gr = gr.Audio(
-                    label="Reference Audio",
-                    type="filepath",
-                    value="OpenVoice/resources/demo_speaker2.mp3",
-                )
-                tos_gr = gr.Checkbox(
-                    label="Agree",
-                    value=False,
-                    info="I agree to the terms of the cc-by-nc-4.0 license-: https://github.com/myshell-ai/OpenVoice/blob/main/LICENSE",
-                )
-
-                tts_button = gr.Button("Send", elem_id="send-btn", visible=True)
-
-            with gr.Column():
-                out_text_gr = gr.Text(label="Info")
-                audio_orig_gr = gr.Audio(label="Synthesised Audio", autoplay=False)
-                audio_gr = gr.Audio(label="Audio with cloned voice", autoplay=True)
-                # ref_audio_gr = gr.Audio(label="Reference Audio Used")
-                predict = build_predict(
-                    output_dir,
-                    tone_color_converter,
-                    en_tts_model,
-                    zh_tts_model,
-                    en_source_default_se,
-                    en_source_style_se,
-                    zh_source_se,
-                )
-
-                gr.Examples(
-                    examples,
-                    label="Examples",
-                    inputs=[input_text_gr, style_gr, ref_gr, tos_gr],
-                    outputs=[out_text_gr, audio_gr],
-                    fn=predict,
-                    cache_examples=False,
-                )
-                tts_button.click(
-                    predict,
-                    [input_text_gr, style_gr, ref_gr, tos_gr],
-                    outputs=[out_text_gr, audio_orig_gr, audio_gr],
-                )
-    return demo
-
-demo = get_demo(
+predict_fn = build_predict(
     OUTPUT_DIR,
     tone_color_converter,
     en_base_speaker_tts,
@@ -592,13 +255,26 @@ demo = get_demo(
     en_source_default_se,
     en_source_style_se,
     zh_source_se,
+    supported_languages,
 )
-demo.queue(max_size=2)
-
-try:
-    demo.launch(debug=True, height=1000)
-except Exception:
-    demo.launch(share=True, debug=True, height=1000)
-# if you are launching remotely, specify server_name and server_port
-# demo.launch(server_name='your server name', server_port='server port in int')
-# Read more in the docs: https://gradio.app/docs/
+
+def gradio_interface():
+    input_text = gr.inputs.Textbox(lines=2, placeholder="Enter text here...")
+    reference_audio = gr.inputs.Audio(source="upload", type="file", label="Reference Audio")
+    speaker = gr.inputs.Dropdown(choices=reference_speakers, default="record_manually", label="Select Speaker")
+    noise_scale = gr.inputs.Slider(minimum=0.1, maximum=1.0, default=0.667, label="Noise Scale")
+    length_scale = gr.inputs.Slider(minimum=0.1, maximum=2.0, default=1.0, label="Length Scale")
+    noise_scale_w = gr.inputs.Slider(minimum=0.1, maximum=1.0, default=0.8, label="Noise Scale W")
+    tone_color = gr.inputs.Checkbox(default=False, label="Enable Tone Color Conversion")
+
+    gr.Interface(
+        fn=predict_fn,
+        inputs=[input_text, reference_audio, speaker, noise_scale, length_scale, noise_scale_w, tone_color],
+        outputs=gr.outputs.Audio(type="file", label="Generated Audio"),
+        title="Speech Generation and Tone Conversion",
+        description="Generate speech and convert tone using the OpenVoice model.",
+    ).launch()
+
+gradio_interface()
+
+