update LLM with speaker encoder

speech/config.yaml

@@ -17,6 +17,16 @@ token_mel_ratio: 2
 chunk_size: 25 # streaming inference chunk size, in token
 num_decoding_left_chunks: -1 # streaming inference flow decoder left chunk size, <0 means use all left chunks
 
+use_speaker_encoder: True  # New flag
+speaker_encoder_config:
+    mel_dim: 80
+    model_dim: 512
+    output_dim: !ref <spk_embed_dim>  # 192
+    num_blocks: 6
+    num_heads: 8
+    kernel_size: 1
+    dropout: 0.1
+    max_conditioning_inputs: 3  # Support multiple references
 # model params
 # for all class/function included in this repo, we use !<name> or !<new> for intialization, so that user may find all corresponding class/function according to one single yaml.
 # for system/third_party class/function, we do not require this.
@@ -27,6 +37,9 @@ llm: !new:cosyvoice.llm.llm.Qwen2LM
     length_normalized_loss: True
     lsm_weight: 0
     mix_ratio: [5, 15]
+    use_speaker_encoder: !ref <use_speaker_encoder>  # Add this
+    spk_embed_dim: !ref <spk_embed_dim>
+    max_conditioning_inputs: 3
     llm: !new:cosyvoice.llm.llm.Qwen2Encoder
         pretrain_path: !ref <qwen_pretrain_path>
     sampling: !name:cosyvoice.utils.common.ras_sampling
@@ -34,6 +47,14 @@ llm: !new:cosyvoice.llm.llm.Qwen2LM
         top_k: 25
         win_size: 10
         tau_r: 0.1
+#
+extract_reference_mel: !name:cosyvoice.dataset.processor.extract_reference_mel_from_speech
+    feat_extractor: !ref <feat_extractor>
+    min_length: 0.5
+    max_length: 6.0
+    num_crops: 2  # Multiple crops from same utterance
+    training: True
+    sample_rate: !ref <sample_rate>
 
 flow: !new:cosyvoice.flow.flow.CausalMaskedDiffWithXvec
     input_size: 512
@@ -180,6 +201,7 @@ batch: !name:cosyvoice.dataset.processor.batch
     max_frames_in_batch: 25000
 padding: !name:cosyvoice.dataset.processor.padding
     use_spk_embedding: False # change to True during sft
+    use_speaker_encoder: !ref <use_speaker_encoder>
 
 
 # dataset processor pipeline
@@ -189,6 +211,7 @@ data_pipeline: [
     !ref <filter>,
     !ref <resample>,
     !ref <compute_fbank>,
+    !ref <extract_reference_mel>,  # Add this for speaker encoder
     !ref <parse_embedding>,
     !ref <shuffle>,
     !ref <sort>,

speech/cosyvoice/llm/llm.py

@@ -665,6 +665,66 @@ class Qwen2LM(TransformerLM):
         rejected_logps = (rejected_logps * chosen_lm_mask).mean(dim=-1)
         return {'loss': loss, 'acc': acc, 'chosen_logps': chosen_logps, 'rejected_logps': rejected_logps}
 
+    @torch.inference_mode()
+    def inference_spk(
+            self,
+            text: torch.Tensor,
+            text_len: torch.Tensor,
+            prompt_text: torch.Tensor,
+            prompt_text_len: torch.Tensor,
+            prompt_speech_token: torch.Tensor,
+            prompt_speech_token_len: torch.Tensor,
+            embedding: torch.Tensor = None,
+            reference_mels: torch.Tensor = None,
+            reference_mel_lengths: torch.Tensor = None,
+            reference_mel_masks: torch.Tensor = None,
+            sampling: int = 25,
+            max_token_text_ratio: float = 20,
+            min_token_text_ratio: float = 2,
+            uuid: str = '',
+    ) -> Generator[torch.Tensor, None, None]:
+        device = text.device
+        text = torch.concat([prompt_text, text], dim=1)
+        text_len += prompt_text_len
+        text = self.llm.model.model.embed_tokens(text)
+        
+        # Get speaker conditioning
+        if self.use_speaker_encoder and reference_mels is not None:
+            # Use speaker encoder
+            batch = {
+                'reference_mels': reference_mels,
+                'reference_mel_lengths': reference_mel_lengths,
+                'reference_mel_masks': reference_mel_masks
+            }
+            speaker_embed = self.get_speaker_conditioning(batch, device)  # [1, 1, llm_input_size]
+        elif embedding is not None and embedding.shape[0] != 0:
+            # Use provided embeddings
+            embedding = F.normalize(embedding, dim=1)
+            speaker_embed = self.spk_embed_affine_layer(embedding)
+            speaker_embed = speaker_embed.unsqueeze(1)
+        else:
+            # No speaker conditioning
+            speaker_embed = torch.zeros(1, 1, self.llm_input_size).to(device)
+
+        # 3. concat llm_input with speaker embedding
+        sos_eos_emb = self.llm_embedding.weight[self.sos_eos].reshape(1, 1, -1)
+        task_id_emb = self.llm_embedding.weight[self.task_id].reshape(1, 1, -1)
+        if prompt_speech_token_len != 0:
+            prompt_speech_token_emb = self.speech_embedding(prompt_speech_token)
+        else:
+            prompt_speech_token_emb = torch.zeros(1, 0, self.llm_input_size, dtype=text.dtype).to(device)
+        
+        # Include speaker embedding in the sequence
+        lm_input = torch.concat([sos_eos_emb, speaker_embed, text, task_id_emb, prompt_speech_token_emb], dim=1)
+
+        # 4. cal min/max_length
+        min_len = int((text_len - prompt_text_len) * min_token_text_ratio)
+        max_len = int((text_len - prompt_text_len) * max_token_text_ratio)
+
+        # 5. step by step decode
+        for token in self.inference_wrapper(lm_input, sampling, min_len, max_len, uuid):
+            yield token
+
     @torch.inference_mode()
     def inference(
             self,