a2.py

import asyncio
import datetime
import logging
import os
import time
import traceback

import edge_tts
import gradio as gr
import librosa
import torch
from fairseq import checkpoint_utils

from config import Config
from lib.infer_pack.models import (
    SynthesizerTrnMs256NSFsid,
    SynthesizerTrnMs256NSFsid_nono,
    SynthesizerTrnMs768NSFsid,
    SynthesizerTrnMs768NSFsid_nono,
)
from rmvpe import RMVPE
from vc_infer_pipeline import VC

# Set logging levels
logging.getLogger("fairseq").setLevel(logging.WARNING)
logging.getLogger("numba").setLevel(logging.WARNING)
logging.getLogger("markdown_it").setLevel(logging.WARNING)
logging.getLogger("urllib3").setLevel(logging.WARNING)
logging.getLogger("matplotlib").setLevel(logging.WARNING)

limitation = os.getenv("SYSTEM") == "spaces"

config = Config()

# Edge TTS
edge_output_filename = "edge_output.mp3"
tts_voice_list = asyncio.get_event_loop().run_until_complete(edge_tts.list_voices())
tts_voices = ["mn-MN-BataaNeural", "mn-MN-YesuiNeural"]  # Specific voices

# RVC models
model_root = "weights"
models = [d for d in os.listdir(model_root) if os.path.isdir(f"{model_root}/{d}")]
models.sort()


def model_data(model_name):
    # global n_spk, tgt_sr, net_g, vc, cpt, version, index_file
    pth_path = [
        f"{model_root}/{model_name}/{f}"
        for f in os.listdir(f"{model_root}/{model_name}")
        if f.endswith(".pth")
    ][0]
    print(f"Loading {pth_path}")
    cpt = torch.load(pth_path, map_location="cpu")
    tgt_sr = cpt["config"][-1]
    cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0]  # n_spk
    if_f0 = cpt.get("f0", 1)
    version = cpt.get("version", "v1")
    if version == "v1":
        if if_f0 == 1:
            net_g = SynthesizerTrnMs256NSFsid(*cpt["config"], is_half=config.is_half)
        else:
            net_g = SynthesizerTrnMs256NSFsid_nono(*cpt["config"])
    elif version == "v2":
        if if_f0 == 1:
            net_g = SynthesizerTrnMs768NSFsid(*cpt["config"], is_half=config.is_half)
        else:
            net_g = SynthesizerTrnMs768NSFsid_nono(*cpt["config"])
    else:
        raise ValueError("Unknown version")
    del net_g.enc_q
    net_g.load_state_dict(cpt["weight"], strict=False)
    print("Model loaded")
    net_g.eval().to(config.device)
    if config.is_half:
        net_g = net_g.half()
    else:
        net_g = net_g.float()
    vc = VC(tgt_sr, config)
    # n_spk = cpt["config"][-3]

    index_files = [
        f"{model_root}/{model_name}/{f}"
        for f in os.listdir(f"{model_root}/{model_name}")
        if f.endswith(".index")
    ]
    if len(index_files) == 0:
        print("No index file found")
        index_file = ""
    else:
        index_file = index_files[0]
        print(f"Index file found: {index_file}")

    return tgt_sr, net_g, vc, version, index_file, if_f0


def load_hubert():
    # global hubert_model
    models, _, _ = checkpoint_utils.load_model_ensemble_and_task(
        ["hubert_base.pt"],
        suffix="",
    )
    hubert_model = models[0]
    hubert_model = hubert_model.to(config.device)
    if config.is_half:
        hubert_model = hubert_model.half()
    else:
        hubert_model = hubert_model.float()
    return hubert_model.eval()


def tts(
    model_name,
    tts_text,
    tts_voice,
    index_rate,
    use_uploaded_voice,
    uploaded_voice,
):
    # Default values for parameters used in EdgeTTS
    speed = 0  # Default speech speed
    f0_up_key = 0  # Default pitch adjustment
    f0_method = "rmvpe"  # Default pitch extraction method
    protect = 0.33  # Default protect value
    filter_radius = 3
    resample_sr = 0
    rms_mix_rate = 0.25
    edge_time = 0  # Initialize edge_time

    try:
        if use_uploaded_voice:
            if uploaded_voice is None:
                return "No voice file uploaded.", None, None
            
            # Process the uploaded voice file
            uploaded_file_path = uploaded_voice.name
            audio, sr = librosa.load(uploaded_file_path, sr=16000, mono=True)
        else:
            # EdgeTTS processing
            if limitation and len(tts_text) > 4000:
                return (
                    f"Text characters should be at most 280 in this huggingface space, but got {len(tts_text)} characters.",
                    None,
                    None,
                )
            
            # Invoke Edge TTS
            t0 = time.time()
            speed_str = f"+{speed}%" if speed >= 0 else f"{speed}%"
            asyncio.run(
                edge_tts.Communicate(
                    tts_text, tts_voice, rate=speed_str
                ).save(edge_output_filename)
            )
            t1 = time.time()
            edge_time = t1 - t0

            audio, sr = librosa.load(edge_output_filename, sr=16000, mono=True)

        # Common processing after loading the audio
        duration = len(audio) / sr
        print(f"Audio duration: {duration}s")
        if limitation and duration >= 20:
            return (
                f"Audio should be less than 20 seconds in this huggingface space, but got {duration}s.",
                None,
                None,
            )

        f0_up_key = int(f0_up_key)
        tgt_sr, net_g, vc, version, index_file, if_f0 = model_data(model_name)

        # Setup for RMVPE or other pitch extraction methods
        if f0_method == "rmvpe":
            vc.model_rmvpe = rmvpe_model

        # Perform voice conversion pipeline
        times = [0, 0, 0]
        audio_opt = vc.pipeline(
            hubert_model,
            net_g,
            0,
            audio,
            edge_output_filename if not use_uploaded_voice else uploaded_file_path,
            times,
            f0_up_key,
            f0_method,
            index_file,
            index_rate,
            if_f0,
            filter_radius,
            tgt_sr,
            resample_sr,
            rms_mix_rate,
            version,
            protect,
            None,
        )

        if tgt_sr != resample_sr and resample_sr >= 16000:
            tgt_sr = resample_sr
        
        info = f"Success. Time: tts: {edge_time}s, npy: {times[0]}s, f0: {times[1]}s, infer: {times[2]}s"
        print(info)
        return (
            info,
            edge_output_filename if not use_uploaded_voice else None,
            (tgt_sr, audio_opt),
        )

    except EOFError:
        info = (
            "output not valid. This may occur when input text and speaker do not match."
        )
        print(info)
        return info, None, None
    except Exception as e:
        traceback_info = traceback.format_exc()
        print(traceback_info)
        return str(e), None, None

voice_mapping = {
    "Mongolian Male": "mn-MN-BataaNeural",
    "Mongolian Female": "mn-MN-YesuiNeural"
}



print("Loading hubert model...")
hubert_model = load_hubert()
print("Hubert model loaded.")

print("Loading rmvpe model...")
rmvpe_model = RMVPE("rmvpe.pt", config.is_half, config.device)
print("rmvpe model loaded.")



app = gr.Blocks()

with app:
    with gr.Row():
        model_name = gr.Dropdown(label="Model", choices=models, value=models[0])
        # Update the Dropdown for voice selection
        tts_voice = gr.Dropdown(label="Speaker", choices=list(voice_mapping.keys()), value="Mongolian Male")
        slang_rate = gr.Slider(minimum=0, maximum=1, label="Slang rate", value=0.75, interactive=True)
    with gr.Row():
        use_uploaded_voice = gr.Checkbox(label="Use uploaded voice file", value=False)
        voice_upload = gr.File(label="Upload voice file")
        tts_text = gr.Textbox(label="Input Text", value="Текстыг оруулна уу.")
        but0 = gr.Button("Convert", variant="primary")
    with gr.Row():
        info_text = gr.Textbox(label="Output info")
        edge_tts_output = gr.Audio(label="Voice", type="filepath")
        tts_output = gr.Audio(label="Result")

    # Modify the button click event to use the mapping
    def on_convert_click(model_name, tts_text, selected_voice, slang_rate, use_uploaded_voice, voice_upload):
        # Use the mapping to get the technical name
        technical_voice_name = voice_mapping[selected_voice]
        # Call your tts function with the technical voice name
        return tts(model_name, tts_text, technical_voice_name, slang_rate, use_uploaded_voice, voice_upload)

    but0.click(
        on_convert_click,
        inputs=[model_name, tts_text, tts_voice, slang_rate, use_uploaded_voice, voice_upload],
        outputs=[info_text, edge_tts_output, tts_output],
    )

# Launch the app
app.launch(share=True)