dynamic update inference weights based on request

Author: suchen-sci

src/easevoice/inference/__init__.py

@@ -54,6 +54,8 @@ class InferenceTaskData:
     keep_random = True
     parallel_infer = True
     repetition_penalty = 1.3
+    sovits_path = ""
+    gpt_path = ""
 
 
 @dataclasses.dataclass
@@ -99,6 +101,15 @@ def run(self):
                     task.result_queue.put(InferenceResult(error=str(e)))
 
     def _inference(self, task: InferenceTask):
+        # change weight based on task
+        try:
+            self.tts_pipeline.update_weights(task.data.sovits_path, task.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)

src/easevoice/inference/tts.py

@@ -194,6 +194,9 @@ def __init__(self, configs: Union[dict, str, TTSConfig]):
         self.bert_model: AutoModelForMaskedLM = None  # pyright: ignore
         self.cnhuhbert_model: CNHubert = None  # pyright: ignore
 
+        self.current_sovits_path = configs.vits_weights_path  # pyright: ignore
+        self.current_gpt_path = configs.t2s_weights_path  # pyright: ignore
+
         self._init_models()
 
         self.text_preprocessor: TextPreprocessor = \
@@ -216,6 +219,27 @@ def __init__(self, configs: Union[dict, str, TTSConfig]):
         self.stop_flag: bool = False
         self.precision: torch.dtype = torch.float16 if self.configs.is_half else torch.float32
 
+    def update_weights(self, sovits_path: str, gpt_path: str):
+        if sovits_path != self.current_sovits_path:
+            if sovits_path == "":
+                # empty sovits path, use default path
+                if self.current_sovits_path != self.configs.vits_weights_path:
+                    self.current_sovits_path = self.configs.vits_weights_path
+                    self.init_vits_weights(self.current_sovits_path)
+            else:
+                self.current_sovits_path = sovits_path
+                self.init_vits_weights(self.current_sovits_path)
+
+        if gpt_path != self.current_gpt_path:
+            if gpt_path == "":
+                # empty gpt path, use default path
+                if self.current_gpt_path != self.configs.t2s_weights_path:
+                    self.current_gpt_path = self.configs.t2s_weights_path
+                    self.init_t2s_weights(self.current_gpt_path)
+            else:
+                self.current_gpt_path = gpt_path
+                self.init_t2s_weights(self.current_gpt_path)
+
     def _init_models(self,):
         self.init_t2s_weights(self.configs.t2s_weights_path)
         self.init_vits_weights(self.configs.vits_weights_path)

tests/easevoice/tts_test.py

@@ -5,7 +5,6 @@
 
 from src.easevoice.inference import Runner
 from src.logger import logger
-from src.service.task import TaskService
 from src.service.voice import VoiceCloneService
 
 
@@ -19,8 +18,7 @@ def test_tts_runner(self):
         queue.put(1)
 
     def test_voice_clone_service(self):
-        task_service = TaskService()
-        voice_service = VoiceCloneService(task_service)
+        voice_service = VoiceCloneService()
         time.sleep(5)
         voice_service.close()