support check available model api

Author: suchen-sci

src/rest/rest.py

@@ -22,6 +22,7 @@
 from src.service.train import TrainGPTService, TrainSovitsService
 from src.service.voice import VoiceCloneService
 from src.train.gpt import GPTTrainParams
+from src.train.helper import list_train_gpts, list_train_sovits
 from src.train.sovits import SovitsTrainParams
 from src.utils.response import EaseVoiceResponse
 
@@ -211,8 +212,14 @@ def _register_routes(self):
         self.router.post("/voiceclone/start")(self.start_service)
         self.router.post("/voiceclone/clone")(self.clone)
         self.router.get("/voiceclone/stop")(self.stop_service)
+        self.router.get("/voiceclone/models")(self.get_available_models)
         self.router.get("/voiceclone/status")(self.get_status)
 
+    async def get_available_models(self):
+        gpts = ["default"].extend(list_train_gpts().keys())
+        sovits = ["default"].extend(list_train_sovits().keys())
+        return {"gpts": gpts, "sovits": sovits}
+
     async def get_status(self):
         if self.service is None:
             raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail={"error": "voice clone service is not started"})

src/service/voice.py

@@ -17,13 +17,17 @@
 
 from src.easevoice.inference import InferenceResult, InferenceTask, 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
 
+
 class VoiceCloneStatus(Enum):
     RUNNING = "Running"
     COMPLETED = "Completed"
     ERROR = "Error"
 
+
 class VoiceCloneService:
     """
     VoiceService is a long run service that listens for voice clone tasks and processes them.
@@ -40,7 +44,7 @@ def close(self):
             self.runner_process.terminate()
             self.runner_process.join(timeout=10)
             self.runner_process = None
-    
+
     def get_status(self):
         if self.runner_process is None:
             return VoiceCloneStatus.COMPLETED
@@ -64,6 +68,7 @@ def clone(self, params: dict):
             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:
@@ -84,3 +89,25 @@ def clone(self, params: dict):
                 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")
+
+    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 = ""
+
+        if task.data.gpt_path != "":
+            gpts = list_train_gpts()
+            if task.data.gpt_path in gpts:
+                task.data.gpt_path = gpts[task.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 != "":
+            sovits = list_train_sovits()
+            if task.data.sovits_path in sovits:
+                task.data.sovits_path = sovits[task.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

src/train/gpt.py

@@ -19,7 +19,7 @@
 
 from src.easevoice.soundstorm.auto_reg.data.data_module import Text2SemanticDataModule
 from src.easevoice.soundstorm.auto_reg.models.t2s_lightning_module import Text2SemanticLightningModule
-from src.train.helper import get_gpt_train_dir, train_logs_path, train_ckpt_path
+from src.train.helper import get_gpt_train_dir, train_logs_path
 from src.utils.config import gpt_config_path, cfg, gpt_pretrained_model_path, semantic_output, text_output_name
 
 
@@ -102,7 +102,7 @@ def __init__(self, params: GPTTrainParams):
         self.train_input_dir = params.train_input_dir
         self.train_output = get_gpt_train_dir(params.output_model_name)
         self.train_logs_output = os.path.join(self.train_output, train_logs_path)
-        self.train_ckpts_output = os.path.join(self.train_output, train_ckpt_path)
+        self.train_ckpts_output = os.path.join(self.train_output, "ckpt")
         os.makedirs(self.train_output, exist_ok=True)
         os.makedirs(self.train_logs_output, exist_ok=True)
         os.makedirs(self.train_ckpts_output, exist_ok=True)

src/train/helper.py

@@ -1,11 +1,11 @@
 
 import datetime
 import os
+from pathlib import Path
 from typing import Optional
 from src.utils import config
 
 train_logs_path = "logs"
-train_ckpt_path = "ckpt"
 
 
 def generate_random_name():
@@ -25,8 +25,22 @@ def get_sovits_train_dir(name: Optional[str]):
 
 
 def list_train_gpts():
-    return os.listdir(config.all_gpt_train_output_dir)
+    all_gpts = Path(config.all_gpt_train_output_dir)
+    res = {}
+    for dir in all_gpts.iterdir():
+        if dir.is_dir():
+            for file in dir.iterdir():
+                if file.is_file() and file.name.endswith(".ckpt"):
+                    res[f"{dir.name}/{file.name}"] = str(file)
+    return res
 
 
 def list_train_sovits():
-    return os.listdir(config.all_sovits_train_output_dir)
+    all_sovits = Path(config.all_sovits_train_output_dir)
+    res = {}
+    for dir in all_sovits.iterdir():
+        if dir.is_dir():
+            for file in dir.iterdir():
+                if file.is_file() and file.name.endswith(".pth"):
+                    res[f"{dir.name}/{file.name}"] = str(file)
+    return res

src/train/sovits.py

@@ -4,7 +4,7 @@
 from typing import Any, List, Tuple
 import torch.distributed as dist
 import os
-from src.train.helper import get_sovits_train_dir, train_ckpt_path, train_logs_path
+from src.train.helper import get_sovits_train_dir, train_logs_path
 from src.utils import config
 from src.utils import helper
 from src.utils.helper import load_json
@@ -134,10 +134,9 @@ def _update_hparams(self, hps: TrainConfig, params: SovitsTrainParams):
         hps.data.exp_dir = params.train_input_dir
         hps.train.output_dir = get_sovits_train_dir(params.output_model_name)
         hps.train.train_logs_dir = os.path.join(hps.train.output_dir, train_logs_path)
-        hps.train.save_weight_dir = os.path.join(hps.train.output_dir, train_ckpt_path)
+        hps.train.save_weight_dir = hps.train.output_dir
         os.makedirs(hps.train.output_dir, exist_ok=True)
         os.makedirs(hps.train.train_logs_dir, exist_ok=True)
-        os.makedirs(hps.train.save_weight_dir, exist_ok=True)
 
         # set pretrained model path
         if params.pretrained_s2G == "":
@@ -570,7 +569,7 @@ def _train_and_evaluate(
                     hps.train.learning_rate,
                     epoch,
                     os.path.join(
-                        hps.train.save_weight_dir, f"sovits_G_epoch{epoch}_step{self.step}.pth"
+                        hps.train.train_logs_dir, f"sovits_G_epoch{epoch}_step{self.step}.pth"
                     ),
                 )
                 ckpt.save_checkpoint(
@@ -579,7 +578,7 @@ def _train_and_evaluate(
                     hps.train.learning_rate,
                     epoch,
                     os.path.join(
-                        hps.train.save_weight_dir, f"sovits_D_epoch{epoch}_step{self.step}.pth"
+                        hps.train.train_logs_dir, f"sovits_D_epoch{epoch}_step{self.step}.pth"
                     ),
                 )
             else:
@@ -589,7 +588,7 @@ def _train_and_evaluate(
                     hps.train.learning_rate,
                     epoch,
                     os.path.join(
-                        hps.train.save_weight_dir, "sovits_G_latest.pth"
+                        hps.train.train_logs_dir, "sovits_G_latest.pth"
                     ),
                 )
                 ckpt.save_checkpoint(
@@ -598,7 +597,7 @@ def _train_and_evaluate(
                     hps.train.learning_rate,
                     epoch,
                     os.path.join(
-                        hps.train.save_weight_dir, "sovits_D_latest.pth"
+                        hps.train.train_logs_dir, "sovits_D_latest.pth"
                     ),
                 )
             if rank == 0 and hps.train.if_save_every_weights == True:

tests/train_test.py

@@ -9,7 +9,7 @@
 
 class TestTrain(unittest.TestCase):
     gpt_service = TrainGPTService(gpt_params=GPTTrainParams(
-        output_path="./output",
+        train_input_dir="./output",
         output_model_name="test",
     ))
     sovits_service = TrainSovitsService(sovits_params=SovitsTrainParams(