Upload voice_processing.py

voice_processing.py

@@ -0,0 +1,231 @@
+import asyncio
+import datetime
+import logging
+import os
+import time
+import traceback
+import tempfile
+from concurrent.futures import ThreadPoolExecutor
+
+# import edge_tts  # Commented out as we're not using 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
+
+model_cache = {}
+
+
+# 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 voices
+# tts_voice_list = asyncio.get_event_loop().run_until_complete(edge_tts.list_voices())
+# tts_voices = ["mn-MN-BataaNeural", "mn-MN-YesuiNeural"]
+
+# RVC models directory
+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):
+    # We will not modify this function to cache models
+    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()
+
+def get_model_names():
+    return [d for d in os.listdir(model_root) if os.path.isdir(f"{model_root}/{d}")]
+
+# Initialize the global models
+hubert_model = load_hubert()
+rmvpe_model = RMVPE("rmvpe.pt", config.is_half, config.device)
+
+# voice_mapping = {
+#     "Mongolian Male": "mn-MN-BataaNeural",
+#     "Mongolian Female": "mn-MN-YesuiNeural"
+# }
+
+# Function to run async functions in a new event loop within a thread
+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(max_workers=10) as executor:
+        # futures = [executor.submit(run_async_in_thread, tts, *task) for task in tasks]  # Original line
+        futures = [executor.submit(run_async_in_thread, process_audio, *task) for task in tasks]  # New line
+        results = [future.result() for future in futures]
+    return results
+
+# Keep the original tts function but commented out
+'''
+async def tts(
+    model_name,
+    tts_text,
+    tts_voice,
+    index_rate,
+    use_uploaded_voice,
+    uploaded_voice,
+):
+    # Original TTS function code here
+    ...
+'''
+
+# New function for audio processing only
+async def process_audio(
+    model_name,
+    text_placeholder,
+    voice_placeholder,
+    index_rate,
+    use_uploaded_voice,
+    uploaded_voice,
+):
+    # Default values for parameters
+    f0_up_key = 0
+    f0_method = "rmvpe"
+    protect = 0.33
+    filter_radius = 3
+    resample_sr = 0
+    rms_mix_rate = 0.25
+
+    try:
+        if uploaded_voice is None:
+            return "No voice file uploaded.", None, None
+        
+        # Process the uploaded voice file - read the file instead of writing it
+        audio, sr = librosa.load(uploaded_voice, sr=16000, mono=True)  # Load directly from filepath
+
+        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,
+            )
+
+        # Load the model and process audio
+        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,
+            uploaded_voice,  # Use the filepath directly
+            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: npy: {times[0]}s, f0: {times[1]}s, infer: {times[2]}s"
+        print(info)
+        return (
+            info,
+            None,   
+            (tgt_sr, audio_opt),
+        )
+
+    except Exception as e:
+        traceback_info = traceback.format_exc()
+        print(traceback_info)
+        return str(e), None, None
+