refacotor voice clone api

Author: suchen-sci

src/easevoice/inference/__init__.py

@@ -20,16 +20,6 @@
 logging.getLogger("torchaudio._extension").setLevel(logging.ERROR)
 
 
-@dataclasses.dataclass
-class InferenceResult:
-    """
-    Result of inference
-    """
-    items: list = dataclasses.field(default_factory=list)
-    seed: int = -1
-    error: Optional[str] = None
-
-
 @dataclasses.dataclass
 class InferenceTaskData:
     """
@@ -42,29 +32,20 @@ class InferenceTaskData:
     prompt_lang: str
     text_split_method: str
     aux_ref_audio_paths: list = dataclasses.field(default_factory=list)
-    seed = -1
-    top_k = 5
-    top_p = 1
-    temperature = 1
-    batch_size = 20
-    speed_factor = 1.0
-    ref_text_free = False
-    split_bucket = True
-    fragment_interval = 0.3
-    keep_random = True
-    parallel_infer = True
-    repetition_penalty = 1.3
-    sovits_path = ""
-    gpt_path = ""
-
-
-@dataclasses.dataclass
-class InferenceTask:
-    """
-    Task for inference
-    """
-    result_queue: multiprocessing.Queue
-    data: InferenceTaskData
+    seed: int = -1
+    top_k: int = 5
+    top_p: int = 1
+    temperature: float = 1.0
+    batch_size: int = 20
+    speed_factor: float = 1.0
+    ref_text_free: bool = False
+    split_bucket: bool = True
+    fragment_interval: float = 0.3
+    keep_random: bool = True
+    parallel_infer: bool = True
+    repetition_penalty: float = 1.3
+    sovits_path: str = ""
+    gpt_path: str = ""
 
 
 class Runner:
@@ -75,42 +56,23 @@ class Runner:
     Wait InferenceResult from the queue
     """
 
-    def __init__(self, queue: multiprocessing.Queue):
+    def __init__(self):
         tts_config = TTSConfig(os.path.join(get_base_path(), "configs", "tts_infer.yaml"))
         logger.info(f"tts config: {tts_config}")
 
         self.tts_config = tts_config
         self.tts_pipeline = TTS(tts_config)
-        self.task_queue = queue
-        self.done = False
-
-    def run(self):
-        while not self.done:
-            task: Union[InferenceTask, int] = self.task_queue.get()
-            if isinstance(task, int):
-                logger.info("Received stop signal")
-                return
-            else:
-                try:
-                    items, seed = self._inference(task)
-                    task.result_queue.put(
-                        InferenceResult(items=items, seed=seed)
-                    )
-                except Exception as e:
-                    logger.error(f"error: {e}")
-                    task.result_queue.put(InferenceResult(error=str(e)))
 
-    def _inference(self, task: InferenceTask):
+    def inference(self, data: InferenceTaskData):
         # change weight based on task
         try:
-            self.tts_pipeline.update_weights(task.data.sovits_path, task.data.gpt_path)
+            self.tts_pipeline.update_weights(data.sovits_path, data.gpt_path)
         except Exception as e:
             logger.error(f"failed to update weights: {e}")
             # change back to default weights
             self.tts_pipeline.update_weights("", "")
             raise e
 
-        data = task.data
         seed = -1 if data.keep_random else data.seed
         actual_seed = seed if seed not in [-1, "", None] else random.randrange(1 << 32)
         inputs = {

src/easevoice/inference/tts.py

@@ -678,7 +678,7 @@ def run(self, inputs: dict):
 
         if ref_audio_path in [None, ""] and \
                 ((self.prompt_cache["prompt_semantic"] is None) or (self.prompt_cache["refer_spec"] in [None, []])):
-            raise ValueError("ref_audio_path cannot be empty, when the reference audio is not set using set_ref_audio()")
+            raise ValueError("ref_audio_path cannot be empty")
 
         ###### setting reference audio and prompt text preprocessing ########
         t0 = ttime()

src/rest/rest.py

@@ -254,7 +254,7 @@ async def clone(self, request: dict):
         try:
             return self.service.clone(request)
         except Exception as e:
-            logger.error(f"failed to clone voice for {request}, err: {e}")
+            logger.error(f"failed to clone voice for {request}, err: {e}", exc_info=True)
             raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail={"error": f"failed to clone voice: {e}"})
 
     async def stop_service(self):

src/service/voice.py

@@ -1,23 +1,15 @@
 import base64
-from concurrent.futures import thread
 from enum import Enum
 import gc
 import io
 import multiprocessing as mp
-import os
-import queue
-import threading
-import time
 import numpy as np
-from scipy.io import wavfile
 import soundfile as sf
+import torch
 
-from src.api.api import ServiceNames, TaskStatus, VoiceCloneProgress
 
-
-from src.easevoice.inference import InferenceResult, InferenceTask, InferenceTaskData, Runner
+from src.easevoice.inference import InferenceTaskData, Runner
 from src.logger import logger
-from src.train import sovits
 from src.train.helper import list_train_gpts, list_train_sovits
 from src.utils.response import EaseVoiceResponse, ResponseStatus
 
@@ -35,79 +27,50 @@ class VoiceCloneService:
 
     def __init__(self):
         self.queue = mp.Queue()
-        self.runner_process = mp.Process(target=VoiceCloneService._init_runner, args=(self.queue,))
-        self.runner_process.start()
+        self.runner_process = Runner()
 
     def close(self):
         if self.runner_process is not None:
-            self.queue.put(1)
-            self.runner_process.terminate()
-            self.runner_process.join(timeout=10)
             self.runner_process = None
+            gc.collect()
+            torch.cuda.empty_cache()
 
     def get_status(self):
         if self.runner_process is None:
             return VoiceCloneStatus.COMPLETED
-        elif self.runner_process.is_alive():
-            return VoiceCloneStatus.RUNNING
-        else:
-            return VoiceCloneStatus.ERROR
-
-    @staticmethod
-    def _init_runner(queue: mp.Queue):
-        """
-        Call this method to start the runner process
-        """
-        runner = Runner(queue)
-        runner.run()
-        print("Voice clone runner process exited")
-        gc.collect()
+        return VoiceCloneStatus.RUNNING
 
     def clone(self, params: dict):
-        try:
-            data = InferenceTaskData(**params)
-            queue = mp.Queue()
-            infer_task = InferenceTask(result_queue=queue, data=data)
-            infer_task = self.update_task_path(infer_task)
-            self.queue.put(infer_task)
-            result: InferenceResult = infer_task.result_queue.get(timeout=600)
-        except Exception as e:
-            logger.error(f"failed to clone voice for {params}, error: {e}", exc_info=True)
-            result = InferenceResult(error=str(e))
+        data = InferenceTaskData(**params)
+        data = self.update_task_path(data)
+        items, seed = self.runner_process.inference(data)  # pyright: ignore
 
-            if result.error:
-                logger.error(f"failed to clone voice for {params}, error: {result.error}")
-                return EaseVoiceResponse(ResponseStatus.FAILED, result.error)
-            else:
-                try:
-                    sampling_rate = result.items[0][0]
-                    data = np.concatenate([item[1] for item in result.items])
-                    buffer = io.BytesIO()
-                    sf.write(buffer, data, sampling_rate, format="WAV")
-                    audio = base64.b64encode(buffer.getvalue()).decode("utf-8")
-                    return EaseVoiceResponse(ResponseStatus.SUCCESS, "Voice cloned successfully", {"sampling_rate": sampling_rate, "audio": audio})
-                except Exception as e:
-                    logger.error(f"failed to clone voice for {params}, error: {e}", exc_info=True)
-                    return EaseVoiceResponse(ResponseStatus.FAILED, "failed to clone voice")
+        sampling_rate = items[0][0]
+        data = np.concatenate([item[1] for item in items])
+        buffer = io.BytesIO()
+        sf.write(buffer, data, sampling_rate, format="WAV")
+        audio = base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+        return EaseVoiceResponse(ResponseStatus.SUCCESS, "Voice cloned successfully", {"sampling_rate": sampling_rate, "audio": audio})
 
-    def update_task_path(self, task: InferenceTask):
-        if task.data.gpt_path == "default":
-            task.data.gpt_path = ""
-        if task.data.sovits_path == "default":
-            task.data.sovits_path = ""
+    def update_task_path(self, data: InferenceTaskData):
+        if data.gpt_path == "default":
+            data.gpt_path = ""
+        if data.sovits_path == "default":
+            data.sovits_path = ""
 
-        if task.data.gpt_path != "":
+        if data.gpt_path != "":
             gpts = list_train_gpts()
-            if task.data.gpt_path in gpts:
-                task.data.gpt_path = gpts[task.data.gpt_path]
+            if data.gpt_path in gpts:
+                data.gpt_path = gpts[data.gpt_path]
             else:
-                logger.error(f"failed to find gpt model for {task.data.gpt_path}")
-                raise ValueError(f"failed to find gpt model for {task.data.gpt_path}")
-        if task.data.sovits_path != "":
+                logger.error(f"failed to find gpt model for {data.gpt_path}")
+                raise ValueError(f"failed to find gpt model for {data.gpt_path}")
+        if data.sovits_path != "":
             sovits = list_train_sovits()
-            if task.data.sovits_path in sovits:
-                task.data.sovits_path = sovits[task.data.sovits_path]
+            if data.sovits_path in sovits:
+                data.sovits_path = sovits[data.sovits_path]
             else:
-                logger.error(f"failed to find sovits model for {task.data.sovits_path}")
-                raise ValueError(f"failed to find sovits model for {task.data.sovits_path}")
-        return task
+                logger.error(f"failed to find sovits model for {data.sovits_path}")
+                raise ValueError(f"failed to find sovits model for {data.sovits_path}")
+        return data