feat: add instance pool for api_v3.py

GPT_SoVITS/TTS_infer_pack/TTS.py

@@ -253,9 +253,10 @@ class TTS:
         if self.configs.is_half and str(self.configs.device)!="cpu":
             self.bert_model = self.bert_model.half()
 
-    def init_vits_weights(self, weights_path: str):
+    def init_vits_weights(self, weights_path: str, save_configs: bool = True):
         print(f"Loading VITS weights from {weights_path}")
         self.configs.vits_weights_path = weights_path
+        if save_configs:
             self.configs.save_configs()
         dict_s2 = torch.load(weights_path, map_location=self.configs.device)
         hps = dict_s2["config"]
@@ -285,9 +286,10 @@ class TTS:
             self.vits_model = self.vits_model.half()
 
 
-    def init_t2s_weights(self, weights_path: str):
+    def init_t2s_weights(self, weights_path: str, save_configs: bool = True):
         print(f"Loading Text2Semantic weights from {weights_path}")
         self.configs.t2s_weights_path = weights_path
+        if save_configs:
             self.configs.save_configs()
         self.configs.hz = 50
         dict_s1 = torch.load(weights_path, map_location=self.configs.device)
@@ -334,13 +336,14 @@ class TTS:
             if self.cnhuhbert_model is not None:
                 self.cnhuhbert_model = self.cnhuhbert_model.float()
 
-    def set_device(self, device: torch.device):
+    def set_device(self, device: torch.device, save_configs: bool = True):
         '''
             To set the device for all models.
             Args:
                 device: torch.device, the device to use for all models.
         '''
         self.configs.device = device
+        if save_configs:
             self.configs.save_configs()
         if self.t2s_model is not None:
             self.t2s_model = self.t2s_model.to(device)

GPT_SoVITS/TTS_infer_pack/tts_instance_pool.py

@@ -0,0 +1,133 @@
+import threading
+from time import perf_counter
+import traceback
+from typing import Dict, Union
+
+from GPT_SoVITS.TTS_infer_pack.TTS import TTS, TTS_Config
+
+
+class TTSWrapper(TTS):
+    heat: float = 0
+    usage_count: int = 0
+    usage_counter: int = 0
+    usage_time: float = 0.0
+    first_used_time: float = 0.0
+
+    def __init__(self, configs: Union[dict, str, TTS_Config]):
+        super(TTSWrapper, self).__init__(configs)
+        self.first_used_time = perf_counter()
+
+    def __hash__(self) -> int:
+        return hash(self.first_used_time)
+
+    def run(self, *args, **kwargs):
+        self.usage_counter += 1
+        t0 = perf_counter()
+        for result in super(TTSWrapper, self).run(*args, **kwargs):
+            yield result
+        t1 = perf_counter()
+        self.usage_time += t1 - t0
+        idle_time = self.usage_time - self.first_used_time
+        self.heat = self.usage_counter / idle_time
+
+    def reset_heat(self):
+        self.heat: int = 0
+        self.usage_count: int = 0
+        self.usage_time: float = 0.0
+        self.first_used_time: float = perf_counter()
+
+
+class TTSInstancePool:
+    def __init__(self, max_size):
+        self.max_size: int = max_size
+        self.semaphore: threading.Semaphore = threading.Semaphore(max_size)
+        self.pool_lock: threading.Lock = threading.Lock()
+        self.pool: Dict[int, TTSWrapper] = dict()
+        self.current_index: int = 0
+        self.size: int = 0
+
+    def acquire(self, configs: TTS_Config):
+
+        self.semaphore.acquire()
+        try:
+            with self.pool_lock:
+                # 查询最匹配的实例
+                indexed_key = None
+                rank = []
+                for key, tts_instance in self.pool.items():
+                    if tts_instance.configs.vits_weights_path == configs.vits_weights_path \
+                            and tts_instance.configs.t2s_weights_path == configs.t2s_weights_path:
+                        indexed_key = key
+                    rank.append((tts_instance.heat, key))
+                rank.sort(key=lambda x: x[0])
+                matched_key = None if len(rank) == 0 else rank[0][1]
+
+                # 如果已有实例匹配，则直接复用
+                if indexed_key is not None:
+                    tts_instance = self._reuse_instance(indexed_key, configs)
+                    print(f"如果已有实例匹配，则直接复用: {configs.vits_weights_path} {configs.t2s_weights_path}")
+                    return tts_instance
+
+                # 如果pool未满，则创建一个新实例
+                if self.size < self.max_size:
+                    tts_instance = TTSWrapper(configs)
+                    self.size += 1
+                    print(f"如果pool未满，则创建一个新实例: {configs.vits_weights_path} {configs.t2s_weights_path}")
+                    return tts_instance
+                else:
+                    # 否则用最合适的实例进行复用
+                    tts_instance = self._reuse_instance(matched_key, configs)
+                    print(f"否则用最合适的实例进行复用: {configs.vits_weights_path} {configs.t2s_weights_path}")
+                    return tts_instance
+        except Exception as e:
+            self.semaphore.release()
+            traceback.print_exc()
+            raise e
+
+    def release(self, tts_instance: TTSWrapper):
+        assert tts_instance is not None
+        with self.pool_lock:
+            key = hash(tts_instance)
+            if key in self.pool.keys():
+                return
+            self.pool[key] = tts_instance
+        self.semaphore.release()
+
+    def clear_pool(self):
+        for i in range(self.max_size):
+            self.semaphore.acquire()
+        with self.pool_lock:
+            self.pool.clear()
+        # for i in range(self.max_size):
+        self.semaphore.release(self.max_size)
+
+    def _reuse_instance(self, instance_key: int, configs: TTS_Config) -> TTSWrapper:
+        """
+        复用已有实例
+        args:
+            instance_key: int, 已有实例的Key
+            config: TTS_Config
+        return:
+            TTS_Wrapper: 返回复用的TTS实例
+        """
+
+        # 复用已有实例
+        tts_instance = self.pool.pop(instance_key, None)
+        if tts_instance is None:
+            raise ValueError("Instance not found")
+
+        tts_instance.configs.device = configs.device
+        if tts_instance.configs.vits_weights_path != configs.vits_weights_path \
+                or tts_instance.configs.t2s_weights_path != configs.t2s_weights_path:
+            tts_instance.reset_heat()
+
+        if tts_instance.configs.vits_weights_path != configs.vits_weights_path:
+            tts_instance.init_vits_weights(configs.vits_weights_path, False)
+            tts_instance.configs.vits_weights_path = configs.vits_weights_path
+
+        if tts_instance.configs.t2s_weights_path != configs.t2s_weights_path:
+            tts_instance.init_t2s_weights(configs.t2s_weights_path, False)
+            tts_instance.configs.t2s_weights_path = configs.t2s_weights_path
+
+        tts_instance.set_device(configs.device, False)
+        return tts_instance

GPT_SoVITS/configs/voices/voice1.yaml

@@ -0,0 +1,14 @@
+custom:
+  bert_base_path: GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large
+  cnhuhbert_base_path: GPT_SoVITS/pretrained_models/chinese-hubert-base
+  device: cpu
+  is_half: false
+  t2s_weights_path: GPT_weights/liyunlong-e15.ckpt
+  vits_weights_path: SoVITS_weights/liyunlong_e8_s176.pth
+default:
+  bert_base_path: GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large
+  cnhuhbert_base_path: GPT_SoVITS/pretrained_models/chinese-hubert-base
+  device: cpu
+  is_half: false
+  t2s_weights_path: GPT_SoVITS/pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt
+  vits_weights_path: GPT_SoVITS/pretrained_models/s2G488k.pth

GPT_SoVITS/configs/voices/voice2.yaml

@@ -0,0 +1,14 @@
+custom:
+  bert_base_path: GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large
+  cnhuhbert_base_path: GPT_SoVITS/pretrained_models/chinese-hubert-base
+  device: cpu
+  is_half: false
+  t2s_weights_path: GPT_weights/jackma-e10.ckpt
+  vits_weights_path: SoVITS_weights/jackma_e8_s192.pth
+default:
+  bert_base_path: GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large
+  cnhuhbert_base_path: GPT_SoVITS/pretrained_models/chinese-hubert-base
+  device: cpu
+  is_half: false
+  t2s_weights_path: GPT_SoVITS/pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt
+  vits_weights_path: GPT_SoVITS/pretrained_models/s2G488k.pth

GPT_SoVITS/configs/voices/voice3.yaml

@@ -0,0 +1,14 @@
+custom:
+  bert_base_path: GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large
+  cnhuhbert_base_path: GPT_SoVITS/pretrained_models/chinese-hubert-base
+  device: cpu
+  is_half: false
+  t2s_weights_path: GPT_weights/stephenchow-e15.ckpt
+  vits_weights_path: SoVITS_weights/stephenchow_e8_s112.pth
+default:
+  bert_base_path: GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large
+  cnhuhbert_base_path: GPT_SoVITS/pretrained_models/chinese-hubert-base
+  device: cpu
+  is_half: false
+  t2s_weights_path: GPT_SoVITS/pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt
+  vits_weights_path: GPT_SoVITS/pretrained_models/s2G488k.pth

api_v3.py

@@ -104,6 +104,8 @@ from typing import Generator
 
 import torch
 
+from TTS_infer_pack.tts_instance_pool import TTSInstancePool, TTSWrapper
+
 now_dir = os.getcwd()
 sys.path.append(now_dir)
 sys.path.append("%s/GPT_SoVITS" % (now_dir))
@@ -119,11 +121,10 @@ from fastapi.responses import JSONResponse
 from fastapi import FastAPI
 import uvicorn
 from io import BytesIO
-from GPT_SoVITS.TTS_infer_pack.TTS import TTS, TTS_Config
+from GPT_SoVITS.TTS_infer_pack.TTS import TTS_Config
 from GPT_SoVITS.TTS_infer_pack.text_segmentation_method import get_method_names as get_cut_method_names
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
-from functools import lru_cache
 
 cut_method_names = get_cut_method_names()
 
@@ -137,6 +138,9 @@ argv = sys.argv
 
 APP = FastAPI()
 
+max_size = 10
+tts_instance_pool = TTSInstancePool(max_size)
+
 
 class TTS_Request(BaseModel):
     text: str = None
@@ -162,12 +166,6 @@ class TTS_Request(BaseModel):
     """推理时需要加载的声音模型的yaml配置文件路径，如：GPT_SoVITS/configs/tts_infer.yaml"""
 
 
-@lru_cache(maxsize=10)
-def get_tts_instance(tts_config: TTS_Config) -> TTS:
-    print(f"load tts config from {tts_config.configs_path}")
-    return TTS(tts_config)
-
-
 def pack_ogg(io_buffer: BytesIO, data: np.ndarray, rate: int):
     """modify from https://github.com/RVC-Boss/GPT-SoVITS/pull/894/files"""
     with sf.SoundFile(io_buffer, mode='w', samplerate=rate, channels=1, format='ogg') as audio_file:
@@ -318,7 +316,7 @@ async def tts_handle(req: dict):
         req["return_fragment"] = True
 
     try:
-        tts_instance = get_tts_instance(tts_config)
+        tts_instance = tts_instance_pool.acquire(tts_config)
 
         move_to_gpu(tts_instance, tts_config)
 
@@ -332,27 +330,30 @@ async def tts_handle(req: dict):
                 for sr, chunk in tts_generator:
                     yield pack_audio(BytesIO(), chunk, sr, media_type).getvalue()
                 move_to_cpu(tts_instance)
+                tts_instance_pool.release(tts_instance)
 
-            # _media_type = f"audio/{media_type}" if not (streaming_mode and media_type in ["wav", "raw"]) else f"audio/x-{media_type}"
+            # _media_type = f"audio/{media_type}" if not (streaming_mode and media_type in ["wav", "raw"]) else
+            # f"audio/x-{media_type}"
             return StreamingResponse(streaming_generator(tts_generator, media_type, ), media_type=f"audio/{media_type}")
 
         else:
             sr, audio_data = next(tts_generator)
             audio_data = pack_audio(BytesIO(), audio_data, sr, media_type).getvalue()
             move_to_cpu(tts_instance)
+            tts_instance_pool.release(tts_instance)
             return Response(audio_data, media_type=f"audio/{media_type}")
     except Exception as e:
         return JSONResponse(status_code=400, content={"message": f"tts failed", "Exception": str(e)})
 
 
-def move_to_cpu(tts):
+def move_to_cpu(tts: TTSWrapper):
     cpu_device = torch.device('cpu')
-    tts.set_device(cpu_device)
+    tts.set_device(cpu_device, False)
     print("Moved TTS models to CPU to save GPU memory.")
 
 
-def move_to_gpu(tts: TTS, tts_config: TTS_Config):
+def move_to_gpu(tts: TTSWrapper, tts_config: TTS_Config):
-    tts.set_device(tts_config.device)
+    tts.set_device(tts_config.device, False)
     print("Moved TTS models back to GPU for performance.")
 
 
@@ -422,7 +423,7 @@ async def tts_post_endpoint(request: TTS_Request):
 async def set_refer_audio(refer_audio_path: str = None, tts_infer_yaml_path: str = "GPT_SoVITS/configs/tts_infer.yaml"):
     try:
         tts_config = TTS_Config(tts_infer_yaml_path)
-        tts_instance = get_tts_instance(tts_config)
+        tts_instance = tts_instance_pool.acquire(tts_config)
         tts_instance.set_ref_audio(refer_audio_path)
     except Exception as e:
         return JSONResponse(status_code=400, content={"message": f"set refer audio failed", "Exception": str(e)})
@@ -436,7 +437,7 @@ async def set_gpt_weights(weights_path: str = None, tts_infer_yaml_path: str = "
             return JSONResponse(status_code=400, content={"message": "gpt weight path is required"})
 
         tts_config = TTS_Config(tts_infer_yaml_path)
-        tts_instance = get_tts_instance(tts_config)
+        tts_instance = tts_instance_pool.acquire(tts_config)
         tts_instance.init_t2s_weights(weights_path)
     except Exception as e:
         return JSONResponse(status_code=400, content={"message": f"change gpt weight failed", "Exception": str(e)})
@@ -451,7 +452,7 @@ async def set_sovits_weights(weights_path: str = None, tts_infer_yaml_path: str
             return JSONResponse(status_code=400, content={"message": "sovits weight path is required"})
 
         tts_config = TTS_Config(tts_infer_yaml_path)
-        tts_instance = get_tts_instance(tts_config)
+        tts_instance = tts_instance_pool.acquire(tts_config)
         tts_instance.init_vits_weights(weights_path)
     except Exception as e:
         return JSONResponse(status_code=400, content={"message": f"change sovits weight failed", "Exception": str(e)})