use spk_embedding when sft

hexisyztem · Jul 10, 2024 · 0fd15bb · 0fd15bb
1 parent a723ea3
commit 0fd15bb
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 2 deletions.
diff --git a/cosyvoice/flow/flow.py b/cosyvoice/flow/flow.py
@@ -60,7 +60,7 @@ def forward(
         token_len = batch['speech_token_len'].to(device)
         feat = batch['speech_feat'].to(device)
         feat_len = batch['speech_feat_len'].to(device)
-        embedding = batch['utt_embedding'].to(device)
+        embedding = batch['embedding'].to(device)
 
         # xvec projection
         embedding = F.normalize(embedding, dim=1)

diff --git a/cosyvoice/llm/llm.py b/cosyvoice/llm/llm.py
@@ -97,7 +97,7 @@ def forward(
         text_token_len = batch['text_token_len'].to(device)
         speech_token = batch['speech_token'].to(device)
         speech_token_len = batch['speech_token_len'].to(device)
-        embedding = batch['utt_embedding'].to(device)
+        embedding = batch['embedding'].to(device)
 
         # 1. prepare llm_target
         lm_target = [torch.tensor([IGNORE_ID] * (2 + text_token_len[i]) + speech_token[i, :speech_token_len[i]].tolist() + [self.speech_token_size]) for i in range(text_token.size(0))]

diff --git a/cosyvoice/utils/executor.py b/cosyvoice/utils/executor.py
@@ -52,6 +52,10 @@ def train_one_epoc(self, model, optimizer, scheduler, train_data_loader, cv_data
                 info_dict["batch_idx"] = batch_idx
                 if cosyvoice_join(group_join, info_dict):
                     break
+                if info_dict["use_spk_embedding"] is True:
+                    batch_dict["embedding"] = batch_dict["spk_embedding"]
+                else:
+                    batch_dict["embedding"] = batch_dict["utt_embedding"]
 
                 # Disable gradient synchronizations across DDP processes.
                 # Within this context, gradients will be accumulated on module

diff --git a/examples/libritts/cosyvoice/conf/cosyvoice.fromscratch.yaml b/examples/libritts/cosyvoice/conf/cosyvoice.fromscratch.yaml
@@ -190,6 +190,7 @@ train_conf:
     scheduler: warmuplr
     scheduler_conf:
         warmup_steps: 25000
+    use_spk_embedding: False # change to True during sft
     max_epoch: 200
     grad_clip: 5
     accum_grad: 2

diff --git a/examples/libritts/cosyvoice/conf/cosyvoice.yaml b/examples/libritts/cosyvoice/conf/cosyvoice.yaml
@@ -190,6 +190,7 @@ train_conf:
     scheduler: warmuplr # change to constantlr during sft
     scheduler_conf:
         warmup_steps: 2500
+    use_spk_embedding: False # change to True during sft
     max_epoch: 200
     grad_clip: 5
     accum_grad: 2