update task service, update vocie clone service

Author: suchen-sci

models/pretrained/.gitkeep

@@ -1 +1 @@
-Please download model from 
+Please download model from https://huggingface.co/lj1995/GPT-SoVITS

src/api/api.py

@@ -18,13 +18,19 @@ class TaskStatus:
     FAILED = "FAILED"
 
 
-class Progress(object):
+class ServiceNames:
+    AUDIO = "audio"
+    VOICE_CLONE = "voice_clone"
+
+
+class Progress(BaseModel):
     """Progress of a task."""
     status: str = TaskStatus.PENDING
     current_step: str = ""
     total_steps: int = 0
     completed_steps: int = 0
     current_step_progress: int = 0
+    message: str = ""
 
 
 class AudioTaskProgressInitial(Progress):
@@ -36,6 +42,16 @@ class AudioTaskProgressInitial(Progress):
     current_step_progress: int = 0
 
 
+class VoiceCloneProgress(Progress):
+    """Progress of a task."""
+    status: str = TaskStatus.PENDING
+    current_step: str = ""
+    total_steps: int = 1
+    completed_steps: int = 0
+    current_step_progress: int = 0
+    message: str = ""
+
+
 # Task models
 class Task(BaseModel):
     """Task model."""

src/easevoice/inference/__init__.py

@@ -22,13 +22,19 @@
 
 @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:
+    """
+    Data for inference
+    """
     text: str
     text_lang: str
     ref_audio_path: str
@@ -52,6 +58,9 @@ class InferenceTaskData:
 
 @dataclasses.dataclass
 class InferenceTask:
+    """
+    Task for inference
+    """
     result_queue: multiprocessing.Queue
     data: InferenceTaskData
 

src/rest/rest.py

@@ -2,7 +2,9 @@
 from src.service.task import TaskService
 from src.service.file import FileService
 from src.service.audio import AudioService
+from src.service.voice import VoiceCloneService
 from src.api.api import (
+    Progress,
     Task,
     CreateTaskResponse,
     UpdateTaskRequest,
@@ -196,6 +198,7 @@ async def delete_files(self, request: DeleteFilesRequest):
 app = FastAPI()
 
 task_service = TaskService()
+voice_service = VoiceCloneService(task_service=task_service)
 task_api = TaskAPI(task_service)
 app.include_router(task_api.router, prefix="/apis/v1")
 

src/service/task.py

@@ -2,7 +2,7 @@
 import uuid
 import json
 from datetime import datetime, timezone
-from typing import List, Optional
+from typing import Callable, Dict, List, Optional
 from src.api.api import Task, Progress, TaskStatus
 
 
@@ -22,10 +22,22 @@ def __init__(self, base_dir: Optional[str] = None):
         self.base_dir = base_dir
         os.makedirs(self.base_dir, exist_ok=True)
 
+        tasks = self.get_tasks()
+        self._tasks: Dict[str, Task] = {}
+        for task in tasks:
+            self._tasks[task.taskID] = task
+
     def _task_metadata_path(self, task_id: str) -> str:
         """Get the path to the task metadata file."""
         return os.path.join(self.base_dir, task_id, ".metadata.json")
 
+    def filter_tasks(self, fn: Callable[[Task], bool]) -> List[Task]:
+        """
+        Filter task using a function. 
+        For example, to get all pending tasks for a service.
+        """
+        return sorted(list(filter(fn, self._tasks.values())), key=lambda t: t.createdAt)
+
     def create_task(self, service_name: str, args: dict) -> Task:
         """Create a new task."""
         task_id = str(uuid.uuid4())
@@ -43,10 +55,12 @@ def create_task(self, service_name: str, args: dict) -> Task:
             args=args,
             progress=Progress(),
         )
+        self._tasks[task_id] = task
         return task
 
     def submit_task(self, task: Task):
         self._save_task_metadata(task)
+        self._tasks[task.taskID] = task
 
     def get_tasks(self) -> List[Task]:
         """Get all tasks."""
@@ -57,6 +71,7 @@ def get_tasks(self) -> List[Task]:
                 try:
                     task = self._load_task_metadata(task_id)
                     tasks.append(task)
+                    self._tasks[task_id] = task
                 except FileNotFoundError:
                     pass  # Skip invalid tasks
         return tasks
@@ -66,11 +81,13 @@ def update_task(self, task_id: str, name: str) -> Task:
         task = self._load_task_metadata(task_id)
         task.name = name
         self._save_task_metadata(task)
+        self._tasks[task.taskID] = task
         return task
 
     def delete_task(self, task_id: str):
         """Delete a task."""
         task = self._load_task_metadata(task_id)
+        self._tasks.pop(task_id)
         self._delete_directory(task.homePath)
 
     def _save_task_metadata(self, task: Task):

src/service/voice.py

@@ -1,57 +1,91 @@
+from concurrent.futures import thread
 import gc
 import multiprocessing as mp
+import os
+import queue
 import threading
+import time
+import numpy as np
+from scipy.io import wavfile
 
-from torch import mul
-
+from src.api.api import ServiceNames, TaskStatus, VoiceCloneProgress
+from src.service.task import TaskService
 from ..utils.response import EaseVoiceResponse, ResponseStatus
 
 
 from ..easevoice.inference import InferenceResult, InferenceTask, InferenceTaskData, Runner
 from ..logger import logger
+from src.easevoice import inference
+
+
+class VoiceCloneService:
+    """
+    VoiceService is a long run service that listens for voice clone tasks and processes them.
+    """
 
+    def __init__(self, task_service: TaskService):
+        self.task_service = task_service
 
-class VoiceService:
-    def __init__(self):
         self.queue = mp.Queue()
-        self.runner_process = mp.Process(target=VoiceService.init_runner, args=(self.queue,))
+        self.runner_process = mp.Process(target=VoiceCloneService._init_runner, args=(self.queue,))
         self.runner_process.start()
-        self.locker = threading.Lock()
+
+        self._run_tasks = threading.Thread(target=self._run)
+        self._run_tasks.start()
 
     @staticmethod
-    def init_runner(queue: mp.Queue):
+    def _init_runner(queue: mp.Queue):
         """
         Call this method to start the runner process
         """
         runner = Runner(queue)
         runner.run()
         gc.collect()
 
-    def clone(self, input: dict):
-        ok = self.locker.acquire(timeout=5)
-        if not ok:
-            return EaseVoiceResponse(ResponseStatus.FAILED, "There is another task running, please try again later")
-
-        try:
-            data = InferenceTaskData(**input)
-            queue = mp.Queue()
-            task = InferenceTask(result_queue=queue, data=data)
-            self.queue.put(task)
-            result: InferenceResult = task.result_queue.get(timeout=600)
-        except Exception as e:
-            logger.error(f"failed to clone voice for {input}, error: {e}", exc_info=True)
-            result = InferenceResult(error=str(e))
-
-        finally:
-            self.locker.release()
-
-        if result.error:
-            return EaseVoiceResponse(ResponseStatus.FAILED, result.error)
-
-        return EaseVoiceResponse(
-            ResponseStatus.SUCCESS,
-            "Cloned voice successfully",
-            {
-                "items": result.items,
-                "seed": result.seed
-            })
+    def _run(self):
+        while True:
+            tasks = self.task_service.filter_tasks(lambda t: t.service_name == ServiceNames.VOICE_CLONE and t.progress.status == TaskStatus.PENDING)
+            if len(tasks) == 0:
+                logger.debug("No pending tasks found for voice clone")
+            else:
+                task = tasks[0]
+                logger.info(f"Processing task {task.taskID}, args: {task.args}")
+
+                task.progress.status = TaskStatus.IN_PROGRESS
+                self.task_service.submit_task(task)
+
+                try:
+                    data = InferenceTaskData(**task.args)
+                    queue = mp.Queue()
+                    infer_task = InferenceTask(result_queue=queue, data=data)
+                    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 {task.args}, error: {e}", exc_info=True)
+                    result = InferenceResult(error=str(e))
+
+                if result.error:
+                    progress = VoiceCloneProgress(status=TaskStatus.FAILED, current_step="Failed", total_steps=1, completed_steps=1, current_step_progress=100, message=result.error)
+                    task.progress = progress
+                    self.task_service.submit_task(task)
+                    logger.error(f"failed to clone voice for {task.args}, error: {result.error}")
+                else:
+                    try:
+                        sampling_rate = result.items[0][0]
+                        audio = np.concatenate([item[1] for item in result.items])
+                        output_file = os.path.join(task.homePath, "output.wav")
+                        wavfile.write(output_file, sampling_rate, audio)
+
+                        progress = VoiceCloneProgress(status=TaskStatus.COMPLETED, current_step="Completed", total_steps=1, completed_steps=1, current_step_progress=100)
+                        task.progress = progress
+                        self.task_service.submit_task(task)
+                        logger.info(f"Successfully cloned voice for {task.args}")
+
+                    except Exception as e:
+                        logger.error(f"failed to clone voice for {task.args}, error: {e}", exc_info=True)
+                        progress = VoiceCloneProgress(status=TaskStatus.FAILED, current_step="Failed", total_steps=1, completed_steps=1, current_step_progress=100, message=str(e))
+                        task.progress = progress
+                        self.task_service.submit_task(task)
+                        logger.error(f"failed to clone voice for {task.args}, error: {e}")
+
+            time.sleep(1)