voice_processing.py

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


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 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
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_voices():
    return list(voice_mapping.keys())

def get_model_names():
    model_root = "weights"  # Adjust this path if your models are stored elsewhere
    return [d for d in os.listdir(model_root) if os.path.isdir(f"{model_root}/{d}")]

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]
    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)

    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():
    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()

# 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

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,
):
    edge_output_filename = get_unique_filename("mp3")
    try:
        # Default values for parameters
        speed = 0
        f0_up_key = 0
        f0_method = "rmvpe"
        protect = 0.33
        filter_radius = 3
        resample_sr = 0
        rms_mix_rate = 0.25
        edge_time = 0

        if use_uploaded_voice:
            if uploaded_voice is None:
                raise ValueError("No voice file uploaded.")
            
            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:
            if limitation and len(tts_text) > 12000:
                raise ValueError(f"Text characters should be at most 12000 in this huggingface space, but got {len(tts_text)} characters.")
            
            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)
            edge_time = time.time() - t0
            audio, sr = librosa.load(edge_output_filename, sr=16000, mono=True)

        duration = len(audio) / sr
        print(f"Audio duration: {duration}s")
        if limitation and duration >= 20000:
            raise ValueError(f"Audio should be less than 20 seconds in this huggingface space, but got {duration}s.")

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

        if f0_method == "rmvpe":
            vc.model_rmvpe = rmvpe_model

        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"
        print(info)

        # Convert audio to base64
        with open(edge_output_filename, "rb") as audio_file:
            audio_base64 = base64.b64encode(audio_file.read()).decode('utf-8')

        audio_data_uri = f"data:audio/mp3;base64,{audio_base64}"

        return (
            info,
            audio_data_uri,
            (tgt_sr, audio_opt)  # Return the target sample rate and audio output
        )

    except Exception as e:
        print(f"Error in TTS task: {str(e)}")
        import traceback
        traceback.print_exc()
        if os.path.exists(edge_output_filename):
            os.remove(edge_output_filename)
        return (str(e), None, None)

voice_mapping = {
    "Mongolian Male": "mn-MN-BataaNeural",
    "Mongolian Female": "mn-MN-YesuiNeural"
    # Add more mappings as needed
}

hubert_model = load_hubert()

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

# Global semaphore to control concurrency
max_concurrent_tasks = 16  # Adjust based on server capacity
semaphore = asyncio.Semaphore(max_concurrent_tasks)

# Global ThreadPoolExecutor
executor = ThreadPoolExecutor(max_workers=max_concurrent_tasks)

class TTSProcessor:
    def __init__(self, config):
        self.config = config
        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 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()
            try:
                await task
            except asyncio.CancelledError:
                print("Task was cancelled")
            except Exception as e:
                print(f"Task failed with error: {e}")
            finally:
                self.queue.task_done()
        self.is_processing = False

# Initialize the TTSProcessor
tts_processor = TTSProcessor(config)

async def parallel_tts_processor(tasks):
    return await asyncio.gather(*(tts_processor.tts(*task) for task in tasks))

async def parallel_tts_wrapper(tasks):
    return await parallel_tts_processor(tasks)