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)