voice_processing.py

import asyncio
import datetime
import logging
import os
import time
import traceback
import tempfile
from concurrent.futures import ThreadPoolExecutor

import edge_tts
import librosa
import torch
from fairseq import checkpoint_utils
import uuid

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 up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Set logging levels for other libraries
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
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 get_unique_filename(extension):
    return f"{uuid.uuid4()}.{extension}"

def model_data(model_name):
    pth_path = [
        f"{model_root}/{model_name}/{f}"
        for f in os.listdir(f"{model_root}/{model_name}")
        if f.endswith(".pth")
    ][0]
    logger.info(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)
    logger.info("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)

    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:
        logger.info("No index file found")
        index_file = ""
    else:
        index_file = index_files[0]
        logger.info(f"Index file found: {index_file}")

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

def load_hubert():
    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 get_model_names():
    model_root = "weights"  # Assuming this is where your models are stored
    return [d for d in os.listdir(model_root) if os.path.isdir(f"{model_root}/{d}")]

# Add this helper function to ensure a new event loop is created if none exists
def run_async_in_thread(fn, *args):
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    result = loop.run_until_complete(fn(*args))
    loop.close()
    return result

executor = ThreadPoolExecutor(max_workers=config.n_cpu)

def parallel_tts(tasks):
    with ThreadPoolExecutor() as executor:
        futures = [executor.submit(run_async_in_thread, tts, *task) for task in tasks]
        results = [future.result() for future in futures]
    return results

async def tts(
    model_name,
    tts_text,
    tts_voice,
    index_rate,
    use_uploaded_voice,
    uploaded_voice,
):
    try:
        # 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_output_filename = get_unique_filename("mp3")

        if use_uploaded_voice:
            if uploaded_voice is None:
                logger.error("No voice file uploaded.")
                return "No voice file uploaded.", None, None
            
            # Process the uploaded voice file
            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
                tmp_file.write(uploaded_voice)
                uploaded_file_path = tmp_file.name

            audio, sr = librosa.load(uploaded_file_path, sr=16000, mono=True)
        else:
            # EdgeTTS processing
            if limitation and len(tts_text) > 12000:
                logger.error(f"Text characters exceed limit: {len(tts_text)} characters.")
                return (
                    f"Text characters should be at most 12000 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}%"
            await 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
        logger.info(f"Audio duration: {duration}s")
        if limitation and duration >= 20000:
            logger.error(f"Audio duration exceeds limit: {duration}s")
            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:.2f}s, npy: {times[0]:.2f}s, f0: {times[1]:.2f}s, infer: {times[2]:.2f}s"
        logger.info(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."
        logger.error(info)
        return info, None, None
    except Exception as e:
        logger.exception("Error in TTS processing")
        return str(e), None, None

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

hubert_model = load_hubert()

rmvpe_model = RMVPE("rmvpe.pt", config.is_half, config.device)

# Add the optimized TTSProcessor
class TTSProcessor:
    def __init__(self, config):
        self.config = config
        self.executor = ThreadPoolExecutor(max_workers=config.n_cpu)
        self.semaphore = asyncio.Semaphore(config.max_concurrent_tts)
        self.queue = asyncio.Queue()
        self.is_processing = False

    async def tts(self, model_name, tts_text, tts_voice, index_rate, use_uploaded_voice, uploaded_voice):
        task = asyncio.create_task(self._tts_task(model_name, tts_text, tts_voice, index_rate, use_uploaded_voice, uploaded_voice))
        await self.queue.put(task)
        
        if not self.is_processing:
            asyncio.create_task(self._process_queue())

        return await task

    async def _tts_task(self, model_name, tts_text, tts_voice, index_rate, use_uploaded_voice, uploaded_voice):
        async with self.semaphore:
            return await tts(model_name, tts_text, tts_voice, index_rate, use_uploaded_voice, uploaded_voice)

    async def _process_queue(self):
        self.is_processing = True
        while not self.queue.empty():
            task = await self.queue.get()
            await task
            self.queue.task_done()
        self.is_processing = False

# Initialize the TTSProcessor
tts_processor = TTSProcessor(config)

# Update parallel_tts to use TTSProcessor
async def parallel_tts_processor(tasks):
    return await asyncio.gather(*(tts_processor.tts(*task) for task in tasks))

def parallel_tts_wrapper(tasks):
    loop = asyncio.get_event_loop()
    return loop.run_until_complete(parallel_tts_processor(tasks))

# Keep the original parallel_tts function
# def parallel_tts(tasks):
#     with ThreadPoolExecutor() as executor:
#         futures = [executor.submit(run_async_in_thread, tts, *task) for task in tasks]
#         results = [future.result() for future in futures]
#     return results