add contrastive loss

speech/config.yaml

@@ -73,6 +73,10 @@ flow: !new:cosyvoice.flow.flow.CausalMaskedDiffWithXvec
                 training_cfg_rate: 0.2
                 inference_cfg_rate: 0.7
                 reg_loss_type: 'l1'
+                use_immiscible: True
+                immiscible_k: 8
+                use_contrastive_fm: True
+                contrastive_lambda: 0.05
         estimator: !new:cosyvoice.flow.decoder.CausalConditionalDecoder
             in_channels: 320
             out_channels: 80
@@ -161,6 +165,7 @@ feat_extractor: !name:matcha.utils.audio.mel_spectrogram
     center: False
 compute_fbank: !name:cosyvoice.dataset.processor.compute_fbank
     feat_extractor: !ref <feat_extractor>
+    token_mel_ratio: !ref <token_mel_ratio>
 compute_f0: !name:cosyvoice.dataset.processor.compute_f0
     sample_rate: !ref <sample_rate>
     hop_size: 480
@@ -172,7 +177,7 @@ sort: !name:cosyvoice.dataset.processor.sort
     sort_size: 500  # sort_size should be less than shuffle_size
 batch: !name:cosyvoice.dataset.processor.batch
     batch_type: 'dynamic'
-    max_frames_in_batch: 2000
+    max_frames_in_batch: 25000
 padding: !name:cosyvoice.dataset.processor.padding
     use_spk_embedding: False # change to True during sft
 
@@ -195,12 +200,12 @@ data_pipeline: [
 train_conf:
     optim: adamw
     optim_conf:
-        lr: 1e-5 # change to 1e-5 during sft
+        lr: 2e-6 # change to 1e-5 during sft
     scheduler: constantlr # change to constantlr during sft
     scheduler_conf:
         warmup_steps: 2500
     max_epoch: 200
     grad_clip: 1
     accum_grad: 1
-    log_interval: 100
+    log_interval: 5
     save_per_step: -1

speech/cosyvoice/flow/flow.py

@@ -11,9 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from ast import List
 import logging
 import random
-from typing import Dict, Optional
+from typing import Dict, Optional, Tuple
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
@@ -22,24 +23,57 @@ from cosyvoice.utils.mask import make_pad_mask
 
 
 class MaskedDiffWithXvec(torch.nn.Module):
-    def __init__(self,
-                 input_size: int = 512,
-                 output_size: int = 80,
-                 spk_embed_dim: int = 192,
-                 output_type: str = "mel",
-                 vocab_size: int = 4096,
-                 input_frame_rate: int = 50,
-                 only_mask_loss: bool = True,
-                 encoder: torch.nn.Module = None,
-                 length_regulator: torch.nn.Module = None,
-                 decoder: torch.nn.Module = None,
-                 decoder_conf: Dict = {'in_channels': 240, 'out_channel': 80, 'spk_emb_dim': 80, 'n_spks': 1,
-                                       'cfm_params': DictConfig({'sigma_min': 1e-06, 'solver': 'euler', 't_scheduler': 'cosine',
-                                                                 'training_cfg_rate': 0.2, 'inference_cfg_rate': 0.7, 'reg_loss_type': 'l1'}),
-                                       'decoder_params': {'channels': [256, 256], 'dropout': 0.0, 'attention_head_dim': 64,
-                                                          'n_blocks': 4, 'num_mid_blocks': 12, 'num_heads': 8, 'act_fn': 'gelu'}},
-                 mel_feat_conf: Dict = {'n_fft': 1024, 'num_mels': 80, 'sampling_rate': 22050,
-                                        'hop_size': 256, 'win_size': 1024, 'fmin': 0, 'fmax': 8000}):
+    def __init__(
+        self,
+        input_size: int = 512,
+        output_size: int = 80,
+        spk_embed_dim: int = 192,
+        output_type: str = "mel",
+        vocab_size: int = 4096,
+        input_frame_rate: int = 50,
+        only_mask_loss: bool = True,
+        encoder: torch.nn.Module = None,
+        length_regulator: torch.nn.Module = None,
+        decoder: torch.nn.Module = None,
+        decoder_conf: Dict = {
+            "in_channels": 240,
+            "out_channel": 80,
+            "spk_emb_dim": 80,
+            "n_spks": 1,
+            "cfm_params": DictConfig(
+                {
+                    "sigma_min": 1e-06,
+                    "solver": "euler",
+                    "t_scheduler": "cosine",
+                    "training_cfg_rate": 0.2,
+                    "inference_cfg_rate": 0.7,
+                    "reg_loss_type": "l1",
+                    "use_immiscible": True,
+                    "immiscible_k": 8,
+                    "use_contrastive_fm": False,
+                    "contrastive_lambda": 0.05
+                }
+            ),
+            "decoder_params": {
+                "channels": [256, 256],
+                "dropout": 0.0,
+                "attention_head_dim": 64,
+                "n_blocks": 4,
+                "num_mid_blocks": 12,
+                "num_heads": 8,
+                "act_fn": "gelu",
+            },
+        },
+        mel_feat_conf: Dict = {
+            "n_fft": 1024,
+            "num_mels": 80,
+            "sampling_rate": 22050,
+            "hop_size": 256,
+            "win_size": 1024,
+            "fmin": 0,
+            "fmax": 8000,
+        },
+    ):
         super().__init__()
         self.input_size = input_size
         self.output_size = output_size
@@ -58,22 +92,22 @@ class MaskedDiffWithXvec(torch.nn.Module):
         self.only_mask_loss = only_mask_loss
 
     def forward(
-            self,
-            batch: dict,
-            device: torch.device,
+        self,
+        batch: dict,
+        device: torch.device,
     ) -> Dict[str, Optional[torch.Tensor]]:
-        token = batch['speech_token'].to(device)
-        token_len = batch['speech_token_len'].to(device)
-        feat = batch['speech_feat'].to(device)
-        feat_len = batch['speech_feat_len'].to(device)
-        embedding = batch['embedding'].to(device)
+        token = batch["speech_token"].to(device)
+        token_len = batch["speech_token_len"].to(device)
+        feat = batch["speech_feat"].to(device)
+        feat_len = batch["speech_feat_len"].to(device)
+        embedding = batch["embedding"].to(device)
 
         # xvec projection
         embedding = F.normalize(embedding, dim=1)
         embedding = self.spk_embed_affine_layer(embedding)
 
         # concat text and prompt_text
-        print('token_len values: ', token_len)
+        print("token_len values: ", token_len)
         mask = (~make_pad_mask(token_len)).float().unsqueeze(-1).to(device)
         token = self.input_embedding(torch.clamp(token, min=0)) * mask
 
@@ -98,20 +132,22 @@ class MaskedDiffWithXvec(torch.nn.Module):
             mask.unsqueeze(1),
             h.transpose(1, 2).contiguous(),
             embedding,
-            cond=conds
+            cond=conds,
         )
-        return {'loss': loss}
+        return {"loss": loss}
 
     @torch.inference_mode()
-    def inference(self,
-                  token,
-                  token_len,
-                  prompt_token,
-                  prompt_token_len,
-                  prompt_feat,
-                  prompt_feat_len,
-                  embedding,
-                  flow_cache):
+    def inference(
+        self,
+        token,
+        token_len,
+        prompt_token,
+        prompt_token_len,
+        prompt_feat,
+        prompt_feat_len,
+        embedding,
+        flow_cache,
+    ):
         assert token.shape[0] == 1
         # xvec projection
         embedding = F.normalize(embedding, dim=1)
@@ -119,18 +155,31 @@ class MaskedDiffWithXvec(torch.nn.Module):
 
         # concat speech token and prompt speech token
         token_len1, token_len2 = prompt_token.shape[1], token.shape[1]
-        token, token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + token_len
+        token, token_len = (
+            torch.concat([prompt_token, token], dim=1),
+            prompt_token_len + token_len,
+        )
         mask = (~make_pad_mask(token_len)).unsqueeze(-1).to(embedding)
         token = self.input_embedding(torch.clamp(token, min=0)) * mask
 
         # text encode
         h, h_lengths = self.encoder(token, token_len)
         h = self.encoder_proj(h)
-        mel_len1, mel_len2 = prompt_feat.shape[1], int(token_len2 / self.input_frame_rate * 22050 / 256)
-        h, h_lengths = self.length_regulator.inference(h[:, :token_len1], h[:, token_len1:], mel_len1, mel_len2, self.input_frame_rate)
+        mel_len1, mel_len2 = prompt_feat.shape[1], int(
+            token_len2 / self.input_frame_rate * 22050 / 256
+        )
+        h, h_lengths = self.length_regulator.inference(
+            h[:, :token_len1],
+            h[:, token_len1:],
+            mel_len1,
+            mel_len2,
+            self.input_frame_rate,
+        )
 
         # get conditions
-        conds = torch.zeros([1, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype)
+        conds = torch.zeros(
+            [1, mel_len1 + mel_len2, self.output_size], device=token.device
+        ).to(h.dtype)
         conds[:, :mel_len1] = prompt_feat
         conds = conds.transpose(1, 2)
 
@@ -142,7 +191,7 @@ class MaskedDiffWithXvec(torch.nn.Module):
             cond=conds,
             n_timesteps=10,
             prompt_len=mel_len1,
-            cache=flow_cache
+            cache=flow_cache,
         )
         feat = feat[:, :, mel_len1:]
         assert feat.shape[2] == mel_len2
@@ -150,25 +199,58 @@ class MaskedDiffWithXvec(torch.nn.Module):
 
 
 class CausalMaskedDiffWithXvec(torch.nn.Module):
-    def __init__(self,
-                 input_size: int = 512,
-                 output_size: int = 80,
-                 spk_embed_dim: int = 192,
-                 output_type: str = "mel",
-                 vocab_size: int = 4096,
-                 input_frame_rate: int = 50,
-                 only_mask_loss: bool = True,
-                 token_mel_ratio: int = 2,
-                 pre_lookahead_len: int = 3,
-                 encoder: torch.nn.Module = None,
-                 decoder: torch.nn.Module = None,
-                 decoder_conf: Dict = {'in_channels': 240, 'out_channel': 80, 'spk_emb_dim': 80, 'n_spks': 1,
-                                       'cfm_params': DictConfig({'sigma_min': 1e-06, 'solver': 'euler', 't_scheduler': 'cosine',
-                                                                 'training_cfg_rate': 0.2, 'inference_cfg_rate': 0.7, 'reg_loss_type': 'l1'}),
-                                       'decoder_params': {'channels': [256, 256], 'dropout': 0.0, 'attention_head_dim': 64,
-                                                          'n_blocks': 4, 'num_mid_blocks': 12, 'num_heads': 8, 'act_fn': 'gelu'}},
-                 mel_feat_conf: Dict = {'n_fft': 1024, 'num_mels': 80, 'sampling_rate': 22050,
-                                        'hop_size': 256, 'win_size': 1024, 'fmin': 0, 'fmax': 8000}):
+    def __init__(
+        self,
+        input_size: int = 512,
+        output_size: int = 80,
+        spk_embed_dim: int = 192,
+        output_type: str = "mel",
+        vocab_size: int = 4096,
+        input_frame_rate: int = 50,
+        only_mask_loss: bool = True,
+        token_mel_ratio: int = 2,
+        pre_lookahead_len: int = 3,
+        encoder: torch.nn.Module = None,
+        decoder: torch.nn.Module = None,
+        decoder_conf: Dict = {
+            "in_channels": 240,
+            "out_channel": 80,
+            "spk_emb_dim": 80,
+            "n_spks": 1,
+            "cfm_params": DictConfig(
+                {
+                    "sigma_min": 1e-06,
+                    "solver": "euler",
+                    "t_scheduler": "cosine",
+                    "training_cfg_rate": 0.2,
+                    "inference_cfg_rate": 0.7,
+                    "reg_loss_type": "l1",
+                    "use_immiscible": True,
+                    "immiscible_k": 8,
+                    "use_contrastive_fm": True,
+                    "contrastive_lambda": 0.05
+                }
+            ),
+            "decoder_params": {
+                "channels": [256, 256],
+                "dropout": 0.0,
+                "attention_head_dim": 64,
+                "n_blocks": 4,
+                "num_mid_blocks": 12,
+                "num_heads": 8,
+                "act_fn": "gelu",
+            },
+        },
+        mel_feat_conf: Dict = {
+            "n_fft": 1024,
+            "num_mels": 80,
+            "sampling_rate": 22050,
+            "hop_size": 256,
+            "win_size": 1024,
+            "fmin": 0,
+            "fmax": 8000,
+        },
+    ):
         super().__init__()
         self.input_size = input_size
         self.output_size = output_size
@@ -186,32 +268,26 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
         self.only_mask_loss = only_mask_loss
         self.token_mel_ratio = token_mel_ratio
         self.pre_lookahead_len = pre_lookahead_len
+        print(" decoder_conf['cfm_params']: ", decoder_conf["cfm_params"])
+        self.use_contrastive_fm = decoder_conf["cfm_params"]["use_contrastive_fm"]
 
     def forward(
-            self,
-            batch: dict,
-            device: torch.device,
+        self,
+        batch: dict,
+        device: torch.device,
     ) -> Dict[str, Optional[torch.Tensor]]:
-        token = batch['speech_token'].to(device)
-        token_len = batch['speech_token_len'].to(device)
-        feat = batch['speech_feat'].to(device)
-        feat_len = batch['speech_feat_len'].to(device)
-        embedding = batch['embedding'].to(device)
-
-        # print('token: ', token.shape)
-        # print('token_len: ', token_len.shape)
-        # print('feat: ', feat.shape)
-        # print('feat_len: ', feat_len.shape)
-        # print('embedding: ', embedding.shape)
+        token = batch["speech_token"].to(device)
+        token_len = batch["speech_token_len"].to(device)
+        feat = batch["speech_feat"].to(device)
+        feat_len = batch["speech_feat_len"].to(device)
+        embedding = batch["embedding"].to(device)
 
         # NOTE unified training, static_chunk_size > 0 or = 0
-        streaming = False# if random.random() < 0.5 else False
+        streaming = False  # if random.random() < 0.5 else False
 
         # xvec projection
         embedding = F.normalize(embedding, dim=1)
         embedding = self.spk_embed_affine_layer(embedding)
-        # print('token_len values: ', token_len)
-        # concat text and prompt_text
         mask = (~make_pad_mask(token_len)).float().unsqueeze(-1).to(device)
         token = self.input_embedding(torch.clamp(token, min=0)) * mask
 
@@ -229,42 +305,50 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
         conds = conds.transpose(1, 2)
 
         mask = (~make_pad_mask(h_lengths.sum(dim=-1).squeeze(dim=1))).to(h)
-
-        # print('feat shape: ', feat.shape)
-        # print('mask shape: ', mask.shape)
-        # print('h shape: ', h.shape)
-        # print('embedding shape: ', embedding.shape)
-        # print('conds shape: ', conds.shape)
-        # print('streaming: ', streaming)
-
-        loss, _ = self.decoder.compute_loss(
-            feat.transpose(1, 2).contiguous(),
-            mask.unsqueeze(1),
-            h.transpose(1, 2).contiguous(),
-            embedding,
-            cond=conds,
-            streaming=streaming,
-        )
-        return {'loss': loss}
+        if not self.use_contrastive_fm:
+            loss, _ = self.decoder.compute_loss(
+                feat.transpose(1, 2).contiguous(),
+                mask.unsqueeze(1),
+                h.transpose(1, 2).contiguous(),
+                embedding,
+                cond=conds,
+                streaming=streaming,
+            )
+        else:
+            # print("use contrastive fm")
+            loss, _ = self.decoder.compute_loss_contrastive(
+                feat.transpose(1, 2).contiguous(),
+                mask.unsqueeze(1),
+                h.transpose(1, 2).contiguous(),
+                embedding,
+                cond=conds,
+                streaming=streaming,
+            )
+        return {"loss": loss}
 
     @torch.inference_mode()
-    def inference(self,
-                  token,
-                  token_len,
-                  prompt_token,
-                  prompt_token_len,
-                  prompt_feat,
-                  prompt_feat_len,
-                  embedding,
-                  streaming,
-                  finalize):
+    def inference(
+        self,
+        token,
+        token_len,
+        prompt_token,
+        prompt_token_len,
+        prompt_feat,
+        prompt_feat_len,
+        embedding,
+        streaming,
+        finalize,
+    ):
         assert token.shape[0] == 1
         # xvec projection
         embedding = F.normalize(embedding, dim=1)
         embedding = self.spk_embed_affine_layer(embedding)
 
         # concat text and prompt_text
-        token, token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + token_len
+        token, token_len = (
+            torch.concat([prompt_token, token], dim=1),
+            prompt_token_len + token_len,
+        )
         mask = (~make_pad_mask(token_len)).unsqueeze(-1).to(embedding)
         token = self.input_embedding(torch.clamp(token, min=0)) * mask
 
@@ -272,13 +356,20 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
         if finalize is True:
             h, h_lengths = self.encoder(token, token_len, streaming=streaming)
         else:
-            token, context = token[:, :-self.pre_lookahead_len], token[:, -self.pre_lookahead_len:]
-            h, h_lengths = self.encoder(token, token_len, context=context, streaming=streaming)
+            token, context = (
+                token[:, : -self.pre_lookahead_len],
+                token[:, -self.pre_lookahead_len :],
+            )
+            h, h_lengths = self.encoder(
+                token, token_len, context=context, streaming=streaming
+            )
         mel_len1, mel_len2 = prompt_feat.shape[1], h.shape[1] - prompt_feat.shape[1]
         h = self.encoder_proj(h)
 
         # get conditions
-        conds = torch.zeros([1, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype)
+        conds = torch.zeros(
+            [1, mel_len1 + mel_len2, self.output_size], device=token.device
+        ).to(h.dtype)
         conds[:, :mel_len1] = prompt_feat
         conds = conds.transpose(1, 2)
 
@@ -289,7 +380,7 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
             spks=embedding,
             cond=conds,
             n_timesteps=10,
-            streaming=streaming
+            streaming=streaming,
         )
         feat = feat[:, :, mel_len1:]
         assert feat.shape[2] == mel_len2

speech/cosyvoice/flow/flow_matching.py

@@ -34,6 +34,7 @@ class ConditionalCFM(BASECFM):
         self.estimator = estimator
         self.use_immiscible = cfm_params.use_immiscible
         self.immiscible_k = cfm_params.immiscible_k
+        self.lambda_weight = cfm_params.contrastive_lambda
 
     @torch.inference_mode()
     def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, prompt_len=0, cache=torch.zeros(1, 80, 0, 2)):
@@ -177,14 +178,6 @@ class ConditionalCFM(BASECFM):
         t = torch.rand([b, 1, 1], device=mu.device, dtype=mu.dtype)
         if self.t_scheduler == 'cosine':
             t = 1 - torch.cos(t * 0.5 * torch.pi)
-
-        
-        print(f"\n=== Immiscible Diffusion Debug ===")
-        print(f"x1 shape: {x1.shape}")
-        print(f"mu shape: {mu.shape}")
-        print(f"t shape: {t.shape}")
-        print(f"Device: {x1.device}")
-        print(f"Dtype: {x1.dtype}")
         
         # Apply immiscible diffusion with KNN
         if self.use_immiscible:
@@ -192,49 +185,87 @@ class ConditionalCFM(BASECFM):
             
             # Generate k noise samples for each data point
             z_candidates = torch.randn(b, k, d, T, device=x1.device, dtype=x1.dtype)
-            
-            print(f"z_candidates shape: {z_candidates.shape}")
-            print(f"z_candidates stats - mean: {z_candidates.mean():.4f}, std: {z_candidates.std():.4f}")
-            
-            # Flatten for distance computation
+          
             x1_flat = x1.flatten(start_dim=1).to(torch.float16)
             z_candidates_flat = z_candidates.flatten(start_dim=2).to(torch.float16)
             
-            print(f"x1_flat shape: {x1_flat.shape}")
-            print(f"z_candidates_flat shape: {z_candidates_flat.shape}")
-            
-            # Calculate distances
+          
             distances = torch.norm(x1_flat.unsqueeze(1) - z_candidates_flat, dim=2)
+           
+            min_distances, min_indices = torch.min(distances, dim=1)
             
-            print(f"distances shape: {distances.shape}")
-            print(f"distances stats - mean: {distances.mean():.4f}, std: {distances.std():.4f}")
-            print(f"distances min: {distances.min():.4f}, max: {distances.max():.4f}")
+           
+            z = torch.gather(
+                z_candidates, 
+                1, 
+                min_indices.unsqueeze(1).unsqueeze(2).unsqueeze(3).expand(-1, 1, d, T)
+            )[:, 0, :, :]
             
-            # Pick the nearest noise for each data point
-            min_distances, min_indices = torch.min(distances, dim=1)
+        else:         
+            # sample noise p(x_0)
+            z = torch.randn_like(x1)
+
+        y = (1 - (1 - self.sigma_min) * t) * z + t * x1
+        u = x1 - (1 - self.sigma_min) * z
+
+        # during training, we randomly drop condition to trade off mode coverage and sample fidelity
+        if self.training_cfg_rate > 0:
+            cfg_mask = torch.rand(b, device=x1.device) > self.training_cfg_rate
+            mu = mu * cfg_mask.view(-1, 1, 1)
+            spks = spks * cfg_mask.view(-1, 1)
+            cond = cond * cfg_mask.view(-1, 1, 1)
+       
+        pred = self.estimator(y, mask, mu, t.squeeze(), spks, cond, streaming=streaming)
+        loss = F.mse_loss(pred * mask, u * mask, reduction="sum") / (torch.sum(mask) * u.shape[1])
+        return loss, y
+
+    def compute_loss_contrastive(self, x1, mask, mu, spks=None, cond=None, streaming=False):
+        """Computes diffusion loss
+
+        Args:
+            x1 (torch.Tensor): Target
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): target mask
+                shape: (batch_size, 1, mel_timesteps)
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            spks (torch.Tensor, optional): speaker embedding. Defaults to None.
+                shape: (batch_size, spk_emb_dim)
+
+        Returns:
+            loss: conditional flow matching loss
+            y: conditional flow
+                shape: (batch_size, n_feats, mel_timesteps)
+        """
+        b, d, T = mu.shape
+
+        # random timestep
+        t = torch.rand([b, 1, 1], device=mu.device, dtype=mu.dtype)
+        if self.t_scheduler == 'cosine':
+            t = 1 - torch.cos(t * 0.5 * torch.pi)
+        
+        # Apply immiscible diffusion with KNN
+        if self.use_immiscible:
+            k = getattr(self, 'immiscible_k', 4)
             
+            # Generate k noise samples for each data point
+            z_candidates = torch.randn(b, k, d, T, device=x1.device, dtype=x1.dtype)
+          
+            x1_flat = x1.flatten(start_dim=1).to(torch.float16)
+            z_candidates_flat = z_candidates.flatten(start_dim=2).to(torch.float16)
             
-            print(f"min_indices: {min_indices[:10]}")  # First 10 indices
-            print(f"min_distances stats - mean: {min_distances.mean():.4f}, std: {min_distances.std():.4f}")
+          
+            distances = torch.norm(x1_flat.unsqueeze(1) - z_candidates_flat, dim=2)
+           
+            min_distances, min_indices = torch.min(distances, dim=1)
             
-            # Gather the selected noise samples
+           
             z = torch.gather(
                 z_candidates, 
                 1, 
                 min_indices.unsqueeze(1).unsqueeze(2).unsqueeze(3).expand(-1, 1, d, T)
             )[:, 0, :, :]
             
-            print(f"Selected z shape: {z.shape}")
-            print(f"Selected z stats - mean: {z.mean():.4f}, std: {z.std():.4f}")
-            
-            # Calculate distance reduction
-            with torch.no_grad():
-                orig_distance = distances[:, 0].mean()
-                selected_distance = min_distances.mean()
-                reduction_rate = (orig_distance - selected_distance) / orig_distance
-                print(f"Distance reduction: {reduction_rate:.3%}")
-                print(f"Original distance: {orig_distance:.4f}")
-                print(f"Selected distance: {selected_distance:.4f}")
         else:         
             # sample noise p(x_0)
             z = torch.randn_like(x1)
@@ -250,7 +281,27 @@ class ConditionalCFM(BASECFM):
             cond = cond * cfg_mask.view(-1, 1, 1)
        
         pred = self.estimator(y, mask, mu, t.squeeze(), spks, cond, streaming=streaming)
-        loss = F.mse_loss(pred * mask, u * mask, reduction="sum") / (torch.sum(mask) * u.shape[1])
+        fm_loss = F.mse_loss(pred * mask, u * mask, reduction="sum") / (torch.sum(mask) * u.shape[1])
+
+        neg_indices = torch.roll(torch.arange(b, device=x1.device), shifts=1)
+        
+        # Get negative targets from shifted indices
+        if b > 1:
+            u_neg = u[neg_indices]
+            neg_mask = mask[neg_indices]
+            
+            # Contrastive loss
+            contrastive_loss = F.mse_loss(
+                pred * neg_mask, 
+                u_neg * neg_mask,
+                reduction="sum"
+            ) / (torch.sum(neg_mask) * d)
+            print('contrastive_loss: ', contrastive_loss)
+        else:
+            contrastive_loss = torch.tensor(0.0, device=fm_loss.device)
+        
+        loss = fm_loss - self.lambda_weight * contrastive_loss
+
         return loss, y
 
 

speech/cosyvoice/utils/executor.py

@@ -33,6 +33,7 @@ class Executor:
         gan: bool = False,
         ref_model: torch.nn.Module = None,
         dpo_loss: torch.nn.Module = None,
+        use_contrastive_fm: bool = False
     ):
         self.gan = gan
         self.ref_model = ref_model
@@ -41,6 +42,7 @@ class Executor:
         self.epoch = 0
         self.rank = int(os.environ.get("RANK", 0))
         self.device = torch.device(f"cuda:{self.rank}")
+        self.use_contrastive_fm = use_contrastive_fm
 
     def train_one_epoc(
         self,
@@ -69,16 +71,20 @@ class Executor:
 
         use_ddp = info_dict["train_engine"] == "torch_ddp"
 
+
         for batch_idx, batch_dict in enumerate(train_data_loader):
             info_dict["tag"] = "TRAIN"
             info_dict["step"] = self.step
             info_dict["epoch"] = self.epoch
             info_dict["batch_idx"] = batch_idx
 
+
             if use_ddp and (batch_idx + 1) % info_dict["accum_grad"] != 0:
                 context = model.no_sync
             else:
                 context = nullcontext
+
+
             with context():
                 info_dict = batch_forward(
                     model,
@@ -88,6 +94,7 @@ class Executor:
                     ref_model=self.ref_model,
                     dpo_loss=self.dpo_loss,
                 )
+
                 info_dict = batch_backward(model, scaler, info_dict)
 
             info_dict = update_parameter_and_lr(

speech/cosyvoice/utils/train_utils.py

@@ -250,6 +250,7 @@ def batch_forward(model, batch, scaler, info_dict, ref_model=None, dpo_loss=None
 
     with autocast:
         info_dict['loss_dict'] = model(batch, device)
+        # print('infor_dict loss_dict : ', info_dict['loss_dict'])
         if ref_model is not None and dpo_loss is not None:
             chosen_logps = info_dict['loss_dict']["chosen_logps"]
             rejected_logps = info_dict['loss_dict']["rejected_logps"]