Cache spk weights

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+G2PWModel/

add_spk.py

@@ -0,0 +1,11 @@
+from single_inference import add_spk
+
+if __name__ == "__main__":
+    '''
+    test command:
+
+    python3 add_spk.py --spk_id 臺灣女 \
+        --speaker_prompt_audio_path "data/example.wav" \
+        --speaker_prompt_text_transcription "在密碼學中，加密是將明文資訊改變為難以讀取的密文內容，使之不可讀的方法。只有擁有解密方法的對象，經由解密過程，才能將密文還原為正常可讀的內容。"
+    '''
+    add_spk()

cache_inference.py

@@ -0,0 +1,13 @@
+from single_inference import main_cached
+
+if __name__ == "__main__":
+
+    '''
+    test command:
+
+    python3 cache_inference.py --spk_id 臺灣女 \
+        --content_to_synthesize "歡迎使用聯發創新基地 BreezyVoice 模型。" \
+        --output_path results/output.wav
+    '''
+
+    main_cached()

single_inference.py

@@ -145,6 +145,25 @@ def frontend_zero_shot_dual(self, tts_text, prompt_text, prompt_speech_16k, flow
                        'prompt_speech_feat': speech_feat, 'prompt_speech_feat_len': speech_feat_len,
                        'llm_embedding': embedding, 'flow_embedding': flow_embedding}
         return model_input
+
+    def frontend_cached(self, tts_text, spk_id):
+        tts_text_token, tts_text_token_len = self._extract_text_token(tts_text)
+        prompt_text_token = self.spk2info[spk_id]['prompt_text_token']
+        prompt_text_token_len = self.spk2info[spk_id]['prompt_text_token_len']
+        speech_feat = self.spk2info[spk_id]['speech_feat']
+        speech_feat_len = self.spk2info[spk_id]['speech_feat_len']
+        speech_token = self.spk2info[spk_id]['speech_token']
+        speech_token_len = self.spk2info[spk_id]['speech_token_len']
+        embedding = self.spk2info[spk_id]['embedding']
+
+        model_input = {'text': tts_text_token, 'text_len': tts_text_token_len,
+                       'prompt_text': prompt_text_token, 'prompt_text_len': prompt_text_token_len,
+                       'llm_prompt_speech_token': speech_token, 'llm_prompt_speech_token_len': speech_token_len,
+                       'flow_prompt_speech_token': speech_token, 'flow_prompt_speech_token_len': speech_token_len,
+                       'prompt_speech_feat': speech_feat, 'prompt_speech_feat_len': speech_feat_len,
+                       'llm_embedding': embedding, 'flow_embedding': embedding}
+
+        return model_input
 
 ####model
 class CustomCosyVoiceModel(CosyVoiceModel):
@@ -206,6 +225,7 @@ def __init__(self, model_dir):
         if not os.path.exists(model_dir):
             model_dir = snapshot_download(model_dir)
         print("model", model_dir)
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.model_dir = model_dir
 
         with open('{}/cosyvoice.yaml'.format(model_dir), 'r') as f:
@@ -223,18 +243,55 @@ def __init__(self, model_dir):
                         '{}/flow.pt'.format(model_dir),
                         '{}/hift.pt'.format(model_dir))
         del configs
-
+    
     def list_avaliable_spks(self):
         spks = list(self.frontend.spk2info.keys())
         return spks
+
+    def remove_spk(self, spk_id):
+        # 載入原始的 pt 檔案
+        model_name = f"{self.model_dir}/spk2info.pt"
+
+        if spk_id in self.frontend.spk2info:
+            del self.frontend.spk2info[spk_id]
+
+        # 儲存到新的 pt 檔案
+        torch.save(self.frontend.spk2info, model_name)
+        print(self.frontend.spk2info.keys())
+        print(f"刪除Speaker成功: {spk_id}")
+
+
+    def add_spk(self, spk_id, spk_info):
+        # 載入原始的 pt 檔案
+        model_name = f"{self.model_dir}/spk2info.pt"
+        self.frontend.spk2info[spk_id] = spk_info
+
+        # 儲存到新的 pt 檔案
+        torch.save(self.frontend.spk2info, model_name)
+        print(self.frontend.spk2info.keys())
+        print(f"新增Speaker成功，儲存為 {spk_id}")
+
+
+    def cal_spk_info(self, audio_path, prompt_text):
+        prompt_speech_16k = load_wav(audio_path, 16000)
+        ret = {}
+
+        ret['prompt_text_token'], ret['prompt_text_token_len'] = self.frontend._extract_text_token(prompt_text)
+        prompt_speech_22050 = torchaudio.transforms.Resample(orig_freq=16000, new_freq=22050)(prompt_speech_16k)
+        ret['speech_feat'], ret['speech_feat_len'] = self.frontend._extract_speech_feat(prompt_speech_22050)
+        ret['speech_token'], ret['speech_token_len'] = self.frontend._extract_speech_token(prompt_speech_16k)
+        ret['embedding'] = self.frontend._extract_spk_embedding(prompt_speech_16k)
+
+        return ret
 
     def inference_sft(self, tts_text, spk_id):
-        tts_speeches = []
-        for i in self.frontend.text_normalize(tts_text, split=True):
-            model_input = self.frontend.frontend_sft(i, spk_id)
+        for i in re.split(r'(?<=[？！。.?!])\s*', tts_text):
+            if not len(i):
+                continue
+            print("Synthesizing:",i)
+            model_input = self.frontend.frontend_cached(i, spk_id)
             model_output = self.model.inference(**model_input)
-            tts_speeches.append(model_output['tts_speech'])
-        return {'tts_speech': torch.concat(tts_speeches, dim=1)}
+            yield model_output
 
     def inference_zero_shot(self, tts_text, prompt_text, prompt_speech_16k):
         prompt_text = self.frontend.text_normalize(prompt_text, split=False)
@@ -274,6 +331,17 @@ def inference_zero_shot_no_normalize(self, tts_text, prompt_text, prompt_speech_
             model_output = self.model.inference(**model_input)
             tts_speeches.append(model_output['tts_speech'])
         return {'tts_speech': torch.concat(tts_speeches, dim=1)}
+
+    def inference_cached(self, tts_text, spk_id):
+        tts_speeches = []
+        for i in re.split(r'(?<=[？！。.?!])\s*', tts_text):
+            if not len(i):
+                continue
+            print("Synthesizing:",i)
+            model_input = self.frontend.frontend_cached(i, spk_id)
+            model_output = self.model.inference(**model_input)
+            tts_speeches.append(model_output['tts_speech'])
+        return {'tts_speech': torch.concat(tts_speeches, dim=1)}
 
 ####wav2text
 def transcribe_audio(audio_file):
@@ -363,8 +431,6 @@ def single_inference(speaker_prompt_audio_path, content_to_synthesize, output_pa
     else:
         speaker_prompt_text_transcription = transcribe_audio(speaker_prompt_audio_path)
 
-
-
     ###normalization
     speaker_prompt_text_transcription = cosyvoice.frontend.text_normalize_new(
         speaker_prompt_text_transcription, 
@@ -388,6 +454,27 @@ def single_inference(speaker_prompt_audio_path, content_to_synthesize, output_pa
     torchaudio.save(output_path, output['tts_speech'], 22050)
     print(f"Generated voice saved to {output_path}")
 
+def inference_cached(content_to_synthesize, output_path, cosyvoice, bopomofo_converter, spk_id):
+    content_to_synthesize = content_to_synthesize
+    output_path = output_path.strip()
+
+    ###normalization
+    content_to_synthesize = cosyvoice.frontend.text_normalize_new(
+        content_to_synthesize, 
+        split=False
+    )
+
+    print("Content to be synthesized before bopomofo:",content_to_synthesize)
+    content_to_synthesize_bopomo = get_bopomofo_rare(content_to_synthesize, bopomofo_converter)
+    print("Content to be synthesized:",content_to_synthesize)
+    start = time.time()
+    output = cosyvoice.inference_cached(content_to_synthesize_bopomo, spk_id)
+    end = time.time()
+    print("Elapsed time:",end - start)
+    print("Generated audio length:", output['tts_speech'].shape[1]/22050, "seconds")
+    torchaudio.save(output_path, output['tts_speech'], 22050)
+    print(f"Generated voice saved to {output_path}")
+
 def main():
     ####args
     parser = argparse.ArgumentParser(description="Run BreezyVoice text-to-speech with custom inputs")
@@ -410,6 +497,53 @@ def main():
     output_path = args.output_path.strip()
     single_inference(speaker_prompt_audio_path, content_to_synthesize, output_path, cosyvoice, bopomofo_converter, args.speaker_prompt_text_transcription)
 
+def main_cached():
+    ####args
+    parser = argparse.ArgumentParser(description="Run BreezyVoice text-to-speech with custom inputs")
+    parser.add_argument("--content_to_synthesize", type=str, required=True, help="Specifies the content that will be synthesized into speech.")
+    parser.add_argument("--output_path", type=str, required=False, default="results/output.wav", help="Specifies the name and path for the output .wav file.")
+    parser.add_argument("--model_path", type=str, required=False, default = "MediaTek-Research/BreezyVoice-300M",help="Specifies the model used for speech synthesis.")
+    parser.add_argument("--spk_id", type=str, required=False, default = "test_human",help="spk's name")
+    args = parser.parse_args()
+
+
+    cosyvoice = CustomCosyVoice(args.model_path)
+
+    bopomofo_converter = G2PWConverter()
+
+    content_to_synthesize = args.content_to_synthesize
+    spk_id = args.spk_id
+    output_path = args.output_path.strip()
+    inference_cached(content_to_synthesize, output_path, cosyvoice, bopomofo_converter, spk_id)
+
+def add_spk():
+    ####args
+    parser = argparse.ArgumentParser(description="Run BreezyVoice text-to-speech with custom inputs")
+    parser.add_argument("--speaker_prompt_audio_path", type=str, required=True, help="Specifies the path to the prompt speech audio file of the speaker.")
+    parser.add_argument("--speaker_prompt_text_transcription", type=str, required=False, help="Specifies the transcription of the speaker prompt audio (Highly Recommended, if not provided, the system will fall back to transcribing with Whisper.)")
+
+    parser.add_argument("--model_path", type=str, required=False, default = "MediaTek-Research/BreezyVoice-300M",help="Specifies the model used for speech synthesis.")
+    parser.add_argument("--spk_id", type=str, required=False, default = "test_human",help="spk's name")
+    args = parser.parse_args()
+    speaker_prompt_audio_path = args.speaker_prompt_audio_path
+
+    if args.speaker_prompt_text_transcription:
+        speaker_prompt_text_transcription = args.speaker_prompt_text_transcription
+    else:
+        speaker_prompt_text_transcription = transcribe_audio(speaker_prompt_audio_path)
+
+    cosyvoice = CustomCosyVoice(args.model_path)
+    spk_info = cosyvoice.cal_spk_info(speaker_prompt_audio_path, speaker_prompt_text_transcription)
+    cosyvoice.add_spk(args.spk_id, spk_info)
+
+def remove_spk():
+    ####args
+    parser = argparse.ArgumentParser(description="Run BreezyVoice text-to-speech with custom inputs")
+    parser.add_argument("--spk_id", type=str, required=False, default = "test_human",help="spk's name")
+    args = parser.parse_args()
+    cosyvoice = CustomCosyVoice(args.model_path)
+    cosyvoice.remove_spk(args.spk_id)
+
 if __name__ == "__main__":
     main()