Update voice_processing.py

voice_processing.py

@@ -1,11 +1,18 @@
+import asyncio
+import datetime
+import logging
 import os
 import time
 import traceback
-import torch
-import numpy as np
+import tempfile
+from concurrent.futures import ThreadPoolExecutor
+
+import edge_tts
 import librosa
+import torch
 from fairseq import checkpoint_utils
-from rmvpe import RMVPE
+import uuid
+
 from config import Config
 from lib.infer_pack.models import (
     SynthesizerTrnMs256NSFsid,
@@ -13,65 +20,39 @@ from lib.infer_pack.models import (
     SynthesizerTrnMs768NSFsid,
     SynthesizerTrnMs768NSFsid_nono,
 )
+from rmvpe import RMVPE
 from vc_infer_pipeline import VC
-import uuid
-import tempfile  # Ensure this is imported
-import asyncio  # Ensure this is imported
+
+# 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()
 
-# Global models loaded once
-hubert_model = None
-rmvpe_model = None
-model_cache = {}  # Cache for RVC models
+# 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
 
-def load_hubert():
-    global hubert_model
-    if hubert_model is None:
-        print("Loading 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()
-        hubert_model.eval()
-        print("Hubert model loaded.")
-    return hubert_model
-
-def load_rmvpe():
-    global rmvpe_model
-    if rmvpe_model is None:
-        print("Loading RMVPE model...")
-        rmvpe_model = RMVPE("rmvpe.pt", config.is_half, config.device)
-        print("RMVPE model loaded.")
-    return rmvpe_model
+# 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 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}")]
-
 def model_data(model_name):
-    global model_cache
-    if model_name in model_cache:
-        # Return cached model data
-        return model_cache[model_name]
-
-    model_root = "weights"
-    pth_files = [
-        f for f in os.listdir(f"{model_root}/{model_name}") if f.endswith(".pth")
-    ]
-    if not pth_files:
-        raise FileNotFoundError(f"No .pth file found for model '{model_name}'")
-    pth_path = f"{model_root}/{model_name}/{pth_files[0]}"
-    print(f"Loading model from {pth_path}")
+    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
@@ -89,32 +70,61 @@ def model_data(model_name):
             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()
-    print(f"Model '{model_name}' loaded.")
-
     vc = VC(tgt_sr, config)
 
     index_files = [
-        f for f in os.listdir(f"{model_root}/{model_name}") if f.endswith(".index")
+        f"{model_root}/{model_name}/{f}"
+        for f in os.listdir(f"{model_root}/{model_name}")
+        if f.endswith(".index")
     ]
-    if index_files:
-        index_file = f"{model_root}/{model_name}/{index_files[0]}"
-        print(f"Index file found: {index_file}")
-    else:
+    if len(index_files) == 0:
+        print("No index file found")
         index_file = ""
-        print("No index file found.")
+    else:
+        index_file = index_files[0]
+        print(f"Index file found: {index_file}")
 
-    # Cache the loaded model data
-    model_cache[model_name] = (tgt_sr, net_g, vc, version, index_file, if_f0)
     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
+
+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,
@@ -123,58 +133,63 @@ async def tts(
     use_uploaded_voice,
     uploaded_voice,
 ):
-    try:
-        # Load models if not already loaded
-        load_hubert()
-        load_rmvpe()
-
-        # Default values for parameters used in EdgeTTS
-        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
-
-        edge_output_filename = get_unique_filename("mp3")
-        audio = None
-        sr = 16000  # Default sample rate
+    # 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
+
+    edge_output_filename = get_unique_filename("mp3")
 
+    try:
         if use_uploaded_voice:
             if uploaded_voice is None:
-                return {"error": "No voice file uploaded."}, None, None
-
+                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)
-            input_audio_path = uploaded_file_path
         else:
             # EdgeTTS processing
-            import edge_tts
+            if limitation and len(tts_text) > 12000:
+                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 = 0  # Default speech speed
             speed_str = f"+{speed}%" if speed >= 0 else f"{speed}%"
-            communicate = edge_tts.Communicate(
+            await edge_tts.Communicate(
                 tts_text, tts_voice, rate=speed_str
-            )
-            try:
-                await asyncio.wait_for(communicate.save(edge_output_filename), timeout=30)
-            except asyncio.TimeoutError:
-                return {"error": "EdgeTTS operation timed out"}, None, None
+            ).save(edge_output_filename)
             t1 = time.time()
             edge_time = t1 - t0
 
             audio, sr = librosa.load(edge_output_filename, sr=16000, mono=True)
-            input_audio_path = edge_output_filename
 
-        # Load the specified RVC model
+        # Common processing after loading the audio
+        duration = len(audio) / sr
+        print(f"Audio duration: {duration}s")
+        if limitation and duration >= 20000:
+            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)
 
-        # Set RMVPE model for pitch extraction
+        # Setup for RMVPE or other pitch extraction methods
         if f0_method == "rmvpe":
             vc.model_rmvpe = rmvpe_model
 
@@ -183,9 +198,9 @@ async def tts(
         audio_opt = vc.pipeline(
             hubert_model,
             net_g,
-            0,  # Speaker ID
+            0,
             audio,
-            input_audio_path,
+            edge_output_filename if not use_uploaded_voice else uploaded_file_path,
             times,
             f0_up_key,
             f0_method,
@@ -203,29 +218,31 @@ async def tts(
 
         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"
+        
+        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,
+            edge_output_filename if not use_uploaded_voice else None,   
             (tgt_sr, audio_opt),
         )
 
-    except asyncio.CancelledError:
-        print("TTS operation was cancelled")
-        return {"error": "Operation cancelled"}, None, None
     except EOFError:
-        info = "Output not valid. This may occur when input text and speaker do not match."
+        info = (
+            "output not valid. This may occur when input text and speaker do not match."
+        )
         print(info)
-        return {"error": info}, None, None
+        return info, None, None
     except Exception as e:
         traceback_info = traceback.format_exc()
         print(traceback_info)
-        return {"error": str(e)}, None, None
+        return str(e), None, None
 
-# Voice mapping dictionary
 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)