correct speaker encoder

speech/config.yaml

@@ -12,12 +12,12 @@ spk_embed_dim: 192
 qwen_pretrain_path: ''
 token_frame_rate: 25
 token_mel_ratio: 2
-
+use_speaker_encoder: True
+speaker_encoder_path: '/data/checkpoint/llm/epoch_29_step_20001.pt' 
 # stream related params
 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
@@ -51,8 +51,8 @@ llm: !new:cosyvoice.llm.llm.Qwen2LM
 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
+    max_length: 12.0
+    num_crops: 3  # Multiple crops from same utterance
     training: True
     sample_rate: !ref <sample_rate>
 
@@ -66,6 +66,9 @@ flow: !new:cosyvoice.flow.flow.CausalMaskedDiffWithXvec
     only_mask_loss: True
     token_mel_ratio: !ref <token_mel_ratio>
     pre_lookahead_len: 3
+    use_speaker_encoder: !ref <use_speaker_encoder>  # Add this
+    freeze_speaker_encoder: True  # Freeze by default for flow training
+    speaker_encoder_path: !ref <speaker_encoder_path>  # Add this
     encoder: !new:cosyvoice.transformer.upsample_encoder.UpsampleConformerEncoder
         output_size: 512
         attention_heads: 8
@@ -198,7 +201,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: 30000
+    max_frames_in_batch: 5000
 padding: !name:cosyvoice.dataset.processor.padding
     use_spk_embedding: False # change to True during sft
     use_speaker_encoder: !ref <use_speaker_encoder>

speech/cosyvoice/dataset/processor.py

@@ -103,7 +103,7 @@ def individual_file_opener(data, mode='train', tts_data={}):
                 
                 # Load embeddings if they exist
                 if os.path.exists(file_info['embedding_path']):
-                    utt_embedding = torch.load(file_info['embedding_path'])
+                    utt_embedding = torch.load(file_info['embedding_path'], weights_only=False)
                     if isinstance(utt_embedding, torch.Tensor):
                         utt_embedding = utt_embedding.tolist()
                 else:
@@ -113,7 +113,7 @@ def individual_file_opener(data, mode='train', tts_data={}):
                 
                 # Load tokens if they exist
                 if os.path.exists(file_info['token_path']):
-                    speech_token = torch.load(file_info['token_path'])
+                    speech_token = torch.load(file_info['token_path'], weights_only=False)
                     if isinstance(speech_token, torch.Tensor):
                         speech_token = speech_token.tolist()
                 else:
@@ -122,7 +122,7 @@ def individual_file_opener(data, mode='train', tts_data={}):
                 
                 # Load speaker embedding
                 if os.path.exists(file_info['spk_embedding_path']):
-                    spk_embedding = torch.load(file_info['spk_embedding_path'])
+                    spk_embedding = torch.load(file_info['spk_embedding_path'], weights_only=False)
                     if isinstance(spk_embedding, torch.Tensor):
                         spk_embedding = spk_embedding.tolist()
                 else:
@@ -686,17 +686,6 @@ def padding(data, use_spk_embedding, mode='train', gan=False, dpo=False, use_spe
                     batch['reference_mels'] = padded_reference_mels
                     batch['reference_mel_lengths'] = padded_reference_mel_lengths
                     batch['reference_mel_masks'] = reference_mel_masks
-                else:
-                    # No valid mels, create dummy
-                    batch['reference_mels'] = torch.zeros(batch_size, 1, mel_dim, 100)
-                    batch['reference_mel_lengths'] = torch.zeros(batch_size, 1, dtype=torch.int32)
-                    batch['reference_mel_masks'] = torch.zeros(batch_size, 1, 100)
-            else:
-                # No references available, create dummy for batch
-                batch_size = len(order)
-                batch['reference_mels'] = torch.zeros(batch_size, 1, 80, 100)
-                batch['reference_mel_lengths'] = torch.zeros(batch_size, 1, dtype=torch.int32)
-                batch['reference_mel_masks'] = torch.zeros(batch_size, 1, 100)
         
         if gan is True:
             # in gan train, we need pitch_feat
@@ -726,9 +715,4 @@ def padding(data, use_spk_embedding, mode='train', gan=False, dpo=False, use_spe
             batch['reject_speech_token'] = reject_speech_token
             batch['reject_speech_token_len'] = reject_speech_token_len
             
-        if use_spk_embedding is True:
-            batch["embedding"] = batch["spk_embedding"]
-        else:
-            batch["embedding"] = batch["utt_embedding"]
-            
         yield batch

speech/cosyvoice/flow/flow.py

@@ -210,6 +210,10 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
         only_mask_loss: bool = True,
         token_mel_ratio: int = 2,
         pre_lookahead_len: int = 3,
+        use_speaker_encoder: bool = False,  # Add this
+        freeze_speaker_encoder: bool = False,  # Add this
+        max_conditioning_inputs: int = 2,  # Add this
+        speaker_encoder_path: str = None,
         encoder: torch.nn.Module = None,
         decoder: torch.nn.Module = None,
         decoder_conf: Dict = {
@@ -261,6 +265,60 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
         self.input_frame_rate = input_frame_rate
         logging.info(f"input frame rate={self.input_frame_rate}")
         self.input_embedding = nn.Embedding(vocab_size, input_size)
+        self.use_speaker_encoder = use_speaker_encoder
+        # Speaker encoder setup
+        if use_speaker_encoder:
+            from cosyvoice.llm.llm import LearnableSpeakerEncoder
+            self.speaker_encoder = LearnableSpeakerEncoder(
+                mel_dim=80,
+                model_dim=512,
+                output_dim=spk_embed_dim,
+                num_blocks=6,
+                num_heads=8,
+            )
+            
+            # Load speaker encoder weights from LLM checkpoint if provided
+        if speaker_encoder_path is not None:
+            logging.info(f"Loading speaker encoder from {speaker_encoder_path}")
+            checkpoint = torch.load(speaker_encoder_path, map_location='cpu')
+            
+            # Debug: print checkpoint structure
+            print(f'Checkpoint keys: {checkpoint.keys()}')
+            
+            # Extract speaker encoder weights
+            speaker_encoder_state = {}
+            
+            # Check if checkpoint has 'state_dict' key or direct model weights
+            if 'state_dict' in checkpoint:
+                state_dict = checkpoint['state_dict']
+            else:
+                # Direct model weights (based on your save function)
+                state_dict = {k: v for k, v in checkpoint.items() if not k in ['epoch', 'step']}
+            
+            # Extract speaker encoder weights
+            for key, value in state_dict.items():
+                if 'speaker_encoder.' in key:
+                    # Remove module. prefix if exists (from DDP)
+                    new_key = key.replace('module.', '')
+                    # Remove speaker_encoder. prefix to match the local module
+                    new_key = new_key.replace('speaker_encoder.', '')
+                    speaker_encoder_state[new_key] = value
+            
+            if len(speaker_encoder_state) == 0:
+                logging.warning("No speaker encoder weights found in checkpoint!")
+                logging.warning(f"Available keys: {list(state_dict.keys())[:10]}...")  # Show first 10 keys
+            else:
+                logging.info(f"Found {len(speaker_encoder_state)} speaker encoder weights")
+                # Load the weights
+                self.speaker_encoder.load_state_dict(speaker_encoder_state, strict=True)
+                logging.info("Speaker encoder loaded successfully")
+        self.freeze_speaker_encoder = freeze_speaker_encoder
+        if freeze_speaker_encoder:
+            # Freeze speaker encoder parameters
+            for param in self.speaker_encoder.parameters():
+                param.requires_grad = False
+            logging.info("Speaker encoder frozen in flow matching")
+
         self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, output_size)
         self.encoder = encoder
         self.encoder_proj = torch.nn.Linear(self.encoder.output_size(), output_size)
@@ -271,6 +329,55 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
         print(" decoder_conf['cfm_params']: ", decoder_conf["cfm_params"])
         self.use_contrastive_fm = decoder_conf["cfm_params"]["use_contrastive_fm"]
 
+    def get_speaker_embedding(self, batch, device):
+        """Extract speaker embedding from reference mels or use provided embeddings"""
+        
+        if self.use_speaker_encoder and 'reference_mels' in batch:
+            reference_mels = batch['reference_mels'].to(device)
+            
+            # Handle multiple references
+            if reference_mels.dim() == 4:  # [B, N, C, T]
+                B, N, C, T = reference_mels.shape
+                embeddings = []
+                
+                for i in range(N):
+                    ref_mel = reference_mels[:, i, :, :]  # [B, C, T]
+                    if 'reference_mel_masks' in batch:
+                        mask = batch['reference_mel_masks'][:, i, :].unsqueeze(1).to(device)
+                    else:
+                        mask = None
+                    print('ref_mel mask: ', ref_mel.shape, mask.shape)
+                    # Apply speaker encoder
+                    with torch.set_grad_enabled(not self.freeze_speaker_encoder):
+                        emb = self.speaker_encoder(ref_mel, mask)  # [B, spk_embed_dim]
+                    embeddings.append(emb)
+                
+                # Average multiple references
+                embedding = torch.stack(embeddings, dim=1).mean(dim=1)  # [B, spk_embed_dim]
+                
+            else:  # Single reference [B, C, T]
+                if 'reference_mel_mask' in batch:
+                    mask = batch['reference_mel_mask'].unsqueeze(1).to(device)
+                else:
+                    mask = None
+                
+                with torch.set_grad_enabled(not self.freeze_speaker_encoder):
+                    embedding = self.speaker_encoder(reference_mels, mask)
+            
+            # Normalize (already normalized in speaker encoder, but just to be safe)
+            embedding = F.normalize(embedding, dim=1)
+            
+        elif 'embedding' in batch:
+            # Use provided embeddings (backward compatibility)
+            embedding = batch['embedding'].to(device)
+            embedding = F.normalize(embedding, dim=1)
+        else:
+            # No speaker conditioning
+            B = batch['speech_token'].shape[0]
+            embedding = torch.zeros(B, self.spk_embed_dim).to(device)
+            
+        return embedding
+
     def forward(
         self,
         batch: dict,
@@ -280,10 +387,11 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
         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
+        print("get speaker embedding")
+        embedding = self.get_speaker_embedding(batch, device)
 
         # xvec projection
         embedding = F.normalize(embedding, dim=1)
@@ -335,13 +443,29 @@ class CausalMaskedDiffWithXvec(torch.nn.Module):
         prompt_token_len,
         prompt_feat,
         prompt_feat_len,
-        embedding,
-        streaming,
-        finalize,
+        embedding=None,
+        reference_mels=None,
+        reference_mel_lengths=None,
+        reference_mel_masks=None,
+        streaming=False,
+        finalize=False,
     ):
         assert token.shape[0] == 1
+        
+        # Get speaker embedding
+        if self.use_speaker_encoder and reference_mels is not None:
+            batch = {
+                'reference_mels': reference_mels,
+                'reference_mel_lengths': reference_mel_lengths,
+                'reference_mel_masks': reference_mel_masks
+            }
+            embedding = self.get_speaker_embedding(batch, token.device)
+        elif embedding is not None:
+            embedding = F.normalize(embedding, dim=1)
+        else:
+            embedding = torch.zeros(1, self.spk_embed_dim).to(token.device)
+        
         # xvec projection
-        embedding = F.normalize(embedding, dim=1)
         embedding = self.spk_embed_affine_layer(embedding)
 
         # concat text and prompt_text

speech/cosyvoice/llm/llm.py

@@ -29,11 +29,11 @@ from cosyvoice.utils.file_utils import logging
 from cosyvoice.utils.mask import make_pad_mask
 from cosyvoice.transformer.attention import MultiHeadedAttention
 from cosyvoice.transformer.encoder_layer import ConformerEncoderLayer
+from cosyvoice.transformer.arch_util import AttentionBlock
 
 class LearnableSpeakerEncoder(nn.Module):
     """
-    Speaker encoder inspired by Tortoise-TTS for CosyVoice2.
-    Processes mel-spectrograms through attention blocks to extract speaker characteristics.
+    Speaker encoder using the same architecture as Tortoise-TTS ConditioningEncoder.
     """
     def __init__(
         self,
@@ -42,58 +42,60 @@ class LearnableSpeakerEncoder(nn.Module):
         output_dim: int = 192,
         num_blocks: int = 6,
         num_heads: int = 8,
-        kernel_size: int = 1,
         dropout: float = 0.0,
+        mean_pooling: bool = False,  # Tortoise uses first position by default
     ):
         super().__init__()
         
-        # Initial projection (like Tortoise's ConditioningEncoder)
-        self.init = nn.Conv1d(mel_dim, model_dim, kernel_size=kernel_size)
-                
-        self.blocks = nn.ModuleList([
-            ConformerEncoderLayer(
-                size=model_dim,
-                self_attn=MultiHeadedAttention(
-                    n_head=num_heads,
-                    n_feat=model_dim,
-                    dropout_rate=dropout,
-                ),
-                feed_forward=None,  # Can add if needed
-                feed_forward_macaron=None,
-                conv_module=None,
-                dropout_rate=dropout,
-                normalize_before=True,
-            ) for _ in range(num_blocks)
-        ])
+        # Same as Tortoise ConditioningEncoder
+        self.init = nn.Conv1d(mel_dim, model_dim, kernel_size=1)
         
-        # Output projection
-        self.output_proj = nn.Linear(model_dim, output_dim)
+        # AttentionBlock from Tortoise
+        attn = []
+        for _ in range(num_blocks):
+            attn.append(AttentionBlock(model_dim, num_heads))
+        self.attn = nn.Sequential(*attn)
+        
+        self.dim = model_dim
+        self.mean_pooling = mean_pooling
         
+        # Output projection to match CosyVoice embedding dimension
+        self.output_proj = nn.Linear(model_dim, output_dim)
+
     def forward(self, x, mask=None):
         """
         Args:
             x: mel-spectrogram [B, 80, T]
-            mask: padding mask [B, 1, T]
+            mask: padding mask [B, 1, T] (not used in Tortoise version)
         Returns:
             speaker embedding [B, output_dim]
         """
         # Initial conv
         h = self.init(x)  # [B, model_dim, T]
-        h = h.transpose(1, 2)  # [B, T, model_dim]
         
         # Apply attention blocks
-        for block in self.blocks:
-            h, _ = block(h, mask)
+        h = self.attn(h)  # [B, model_dim, T]
         
-        # Pool over time (take first position like Tortoise)
-        # Could also use mean pooling
-        output = h[:, 0, :]  # [B, model_dim]
+        # Pooling - Tortoise uses first position
+        if self.mean_pooling:
+            if mask is not None:
+                # Masked mean pooling
+                h_masked = h * mask
+                h_sum = h_masked.sum(dim=2)
+                mask_sum = mask.sum(dim=2).clamp(min=1)
+                h_pooled = h_sum / mask_sum
+            else:
+                h_pooled = h.mean(dim=2)
+        else:
+            # Use first position like Tortoise
+            h_pooled = h[:, :, 0]
         
         # Project to output dimension
-        output = self.output_proj(output)  # [B, output_dim]
+        output = self.output_proj(h_pooled)  # [B, output_dim]
         
         return F.normalize(output, p=2, dim=1)
 
+
 class TransformerLM(torch.nn.Module):
     def __init__(
             self,
@@ -163,19 +165,19 @@ class TransformerLM(torch.nn.Module):
         
         if self.use_speaker_encoder and 'reference_mels' in batch:
             reference_mels = batch['reference_mels'].to(device)
-            
+            print('reference_mels shape: ', reference_mels.shape)
             # Handle multiple references
             if reference_mels.dim() == 4:  # [B, N, T, 80]
-                B, N, T, D = reference_mels.shape
+                B, N, C, T = reference_mels.shape
                 conds = []
                 
                 for i in range(N):
-                    ref_mel = reference_mels[:, i, :, :].transpose(1, 2)  # [B, 80, T]
+                    ref_mel = reference_mels[:, i, :, :]
                     if 'reference_mel_masks' in batch:
                         mask = batch['reference_mel_masks'][:, i, :].unsqueeze(1).to(device)
                     else:
                         mask = None
-                    
+                    print('ref_mel shape: ', ref_mel.shape, 'mask: ', mask.shape)
                     cond = self.speaker_encoder(ref_mel, mask)  # [B, spk_embed_dim]
                     conds.append(cond)
                 
@@ -547,42 +549,42 @@ class Qwen2LM(TransformerLM):
         lm_target = pad_sequence(lm_target, batch_first=True, padding_value=IGNORE_ID)
         return lm_target, lm_input, lm_input_len
 
-    def forward(
-            self,
-            batch: dict,
-            device: torch.device,
-    ) -> Dict[str, Optional[torch.Tensor]]:
-        """
-        Args:
-            text: (B, L, D)
-            text_lengths: (B,)
-            audio: (B, T, N) or (B, T)
-            audio_lengths: (B,)
-        """
-        text_token = batch['text_token'].to(device)
-        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)
-
-
-        # 1. encode text_token
-        text_token_emb = self.llm.model.model.embed_tokens(text_token)
-
-        # 2. encode speech_token
-        speech_token_emb = self.speech_embedding(speech_token)
-
-        # 3. prepare llm_input/target
-        lm_target, lm_input, lm_input_len = self.prepare_lm_input_target(text_token, text_token_emb, text_token_len, speech_token, speech_token_emb, speech_token_len)
-        lm_target = lm_target.to(device)
-
-        # 4. run lm forward
-        lm_output, lm_output_mask = self.llm(lm_input, lm_input_len.to(device))
-        logits = self.llm_decoder(lm_output)
-        loss = self.criterion_ce(logits, lm_target.to(device))
-        acc = th_accuracy(logits.view(-1, self.speech_token_size + 3), lm_target, ignore_label=IGNORE_ID)
-        return {'loss': loss, 'acc': acc}
+    # def forward(
+    #         self,
+    #         batch: dict,
+    #         device: torch.device,
+    # ) -> Dict[str, Optional[torch.Tensor]]:
+    #     """
+    #     Args:
+    #         text: (B, L, D)
+    #         text_lengths: (B,)
+    #         audio: (B, T, N) or (B, T)
+    #         audio_lengths: (B,)
+    #     """
+    #     text_token = batch['text_token'].to(device)
+    #     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)
+
+
+    #     # 1. encode text_token
+    #     text_token_emb = self.llm.model.model.embed_tokens(text_token)
+
+    #     # 2. encode speech_token
+    #     speech_token_emb = self.speech_embedding(speech_token)
+
+    #     # 3. prepare llm_input/target
+    #     lm_target, lm_input, lm_input_len = self.prepare_lm_input_target(text_token, text_token_emb, text_token_len, speech_token, speech_token_emb, speech_token_len)
+    #     lm_target = lm_target.to(device)
+
+    #     # 4. run lm forward
+    #     lm_output, lm_output_mask = self.llm(lm_input, lm_input_len.to(device))
+    #     logits = self.llm_decoder(lm_output)
+    #     loss = self.criterion_ce(logits, lm_target.to(device))
+    #     acc = th_accuracy(logits.view(-1, self.speech_token_size + 3), lm_target, ignore_label=IGNORE_ID)
+    #     return {'loss': loss, 'acc': acc}
     
-    def forward_spk(
+    def forward(
             self,
             batch: dict,
             device: torch.device,

speech/cosyvoice/transformer/arch_util.py

@@ -0,0 +1,373 @@
+import os
+import functools
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+from cosyvoice.transformer.xtransformers import ContinuousTransformerWrapper, RelativePositionBias
+
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+class GroupNorm32(nn.GroupNorm):
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+
+
+def normalization(channels):
+    """
+    Make a standard normalization layer.
+
+    :param channels: number of input channels.
+    :return: an nn.Module for normalization.
+    """
+    groups = 32
+    if channels <= 16:
+        groups = 8
+    elif channels <= 64:
+        groups = 16
+    while channels % groups != 0:
+        groups = int(groups / 2)
+    assert groups > 2
+    return GroupNorm32(groups, channels)
+
+
+class QKVAttentionLegacy(nn.Module):
+    """
+    A module which performs QKV attention. Matches legacy QKVAttention + input/output heads shaping
+    """
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv, mask=None, rel_pos=None):
+        """
+        Apply QKV attention.
+
+        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
+        :return: an [N x (H * C) x T] tensor after attention.
+        """
+        bs, width, length = qkv.shape
+        assert width % (3 * self.n_heads) == 0
+        ch = width // (3 * self.n_heads)
+        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
+        scale = 1 / math.sqrt(math.sqrt(ch))
+        weight = torch.einsum(
+            "bct,bcs->bts", q * scale, k * scale
+        )  # More stable with f16 than dividing afterwards
+        if rel_pos is not None:
+            weight = rel_pos(weight.reshape(bs, self.n_heads, weight.shape[-2], weight.shape[-1])).reshape(bs * self.n_heads, weight.shape[-2], weight.shape[-1])
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        if mask is not None:
+            # The proper way to do this is to mask before the softmax using -inf, but that doesn't work properly on CPUs.
+            mask = mask.repeat(self.n_heads, 1).unsqueeze(1)
+            weight = weight * mask
+        a = torch.einsum("bts,bcs->bct", weight, v)
+
+        return a.reshape(bs, -1, length)
+
+
+class AttentionBlock(nn.Module):
+    """
+    An attention block that allows spatial positions to attend to each other.
+
+    Originally ported from here, but adapted to the N-d case.
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
+    """
+
+    def __init__(
+        self,
+        channels,
+        num_heads=1,
+        num_head_channels=-1,
+        do_checkpoint=True,
+        relative_pos_embeddings=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.do_checkpoint = do_checkpoint
+        if num_head_channels == -1:
+            self.num_heads = num_heads
+        else:
+            assert (
+                channels % num_head_channels == 0
+            ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
+            self.num_heads = channels // num_head_channels
+        self.norm = normalization(channels)
+        self.qkv = nn.Conv1d(channels, channels * 3, 1)
+        # split heads before split qkv
+        self.attention = QKVAttentionLegacy(self.num_heads)
+
+        self.proj_out = zero_module(nn.Conv1d(channels, channels, 1))
+        if relative_pos_embeddings:
+            self.relative_pos_embeddings = RelativePositionBias(scale=(channels // self.num_heads) ** .5, causal=False, heads=num_heads, num_buckets=32, max_distance=64)
+        else:
+            self.relative_pos_embeddings = None
+
+    def forward(self, x, mask=None):
+        b, c, *spatial = x.shape
+        x = x.reshape(b, c, -1)
+        qkv = self.qkv(self.norm(x))
+        h = self.attention(qkv, mask, self.relative_pos_embeddings)
+        h = self.proj_out(h)
+        return (x + h).reshape(b, c, *spatial)
+
+
+class Upsample(nn.Module):
+    """
+    An upsampling layer with an optional convolution.
+
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    """
+
+    def __init__(self, channels, use_conv, out_channels=None, factor=4):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.factor = factor
+        if use_conv:
+            ksize = 5
+            pad = 2
+            self.conv = nn.Conv1d(self.channels, self.out_channels, ksize, padding=pad)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        x = F.interpolate(x, scale_factor=self.factor, mode="nearest")
+        if self.use_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    """
+    A downsampling layer with an optional convolution.
+
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    """
+
+    def __init__(self, channels, use_conv, out_channels=None, factor=4, ksize=5, pad=2):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+
+        stride = factor
+        if use_conv:
+            self.op = nn.Conv1d(
+                self.channels, self.out_channels, ksize, stride=stride, padding=pad
+            )
+        else:
+            assert self.channels == self.out_channels
+            self.op = nn.AvgPool1d(kernel_size=stride, stride=stride)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        return self.op(x)
+
+
+class ResBlock(nn.Module):
+    def __init__(
+            self,
+            channels,
+            dropout,
+            out_channels=None,
+            use_conv=False,
+            use_scale_shift_norm=False,
+            up=False,
+            down=False,
+            kernel_size=3,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_scale_shift_norm = use_scale_shift_norm
+        padding = 1 if kernel_size == 3 else 2
+
+        self.in_layers = nn.Sequential(
+            normalization(channels),
+            nn.SiLU(),
+            nn.Conv1d(channels, self.out_channels, kernel_size, padding=padding),
+        )
+
+        self.updown = up or down
+
+        if up:
+            self.h_upd = Upsample(channels, False)
+            self.x_upd = Upsample(channels, False)
+        elif down:
+            self.h_upd = Downsample(channels, False)
+            self.x_upd = Downsample(channels, False)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        self.out_layers = nn.Sequential(
+            normalization(self.out_channels),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(
+                nn.Conv1d(self.out_channels, self.out_channels, kernel_size, padding=padding)
+            ),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = nn.Conv1d(
+                channels, self.out_channels, kernel_size, padding=padding
+            )
+        else:
+            self.skip_connection = nn.Conv1d(channels, self.out_channels, 1)
+
+    def forward(self, x):
+        if self.updown:
+            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+            h = in_rest(x)
+            h = self.h_upd(h)
+            x = self.x_upd(x)
+            h = in_conv(h)
+        else:
+            h = self.in_layers(x)
+        h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+class AudioMiniEncoder(nn.Module):
+    def __init__(self,
+                 spec_dim,
+                 embedding_dim,
+                 base_channels=128,
+                 depth=2,
+                 resnet_blocks=2,
+                 attn_blocks=4,
+                 num_attn_heads=4,
+                 dropout=0,
+                 downsample_factor=2,
+                 kernel_size=3):
+        super().__init__()
+        self.init = nn.Sequential(
+            nn.Conv1d(spec_dim, base_channels, 3, padding=1)
+        )
+        ch = base_channels
+        res = []
+        for l in range(depth):
+            for r in range(resnet_blocks):
+                res.append(ResBlock(ch, dropout, kernel_size=kernel_size))
+            res.append(Downsample(ch, use_conv=True, out_channels=ch*2, factor=downsample_factor))
+            ch *= 2
+        self.res = nn.Sequential(*res)
+        self.final = nn.Sequential(
+            normalization(ch),
+            nn.SiLU(),
+            nn.Conv1d(ch, embedding_dim, 1)
+        )
+        attn = []
+        for a in range(attn_blocks):
+            attn.append(AttentionBlock(embedding_dim, num_attn_heads,))
+        self.attn = nn.Sequential(*attn)
+        self.dim = embedding_dim
+
+    def forward(self, x):
+        h = self.init(x)
+        h = self.res(h)
+        h = self.final(h)
+        h = self.attn(h)
+        return h[:, :, 0]
+
+
+DEFAULT_MEL_NORM_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), '../data/mel_norms.pth')
+
+
+class TorchMelSpectrogram(nn.Module):
+    def __init__(self, filter_length=1024, hop_length=256, win_length=1024, n_mel_channels=80, mel_fmin=0, mel_fmax=8000,
+                 sampling_rate=22050, normalize=False, mel_norm_file=DEFAULT_MEL_NORM_FILE):
+        super().__init__()
+        # These are the default tacotron values for the MEL spectrogram.
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.n_mel_channels = n_mel_channels
+        self.mel_fmin = mel_fmin
+        self.mel_fmax = mel_fmax
+        self.sampling_rate = sampling_rate
+        self.mel_stft = torchaudio.transforms.MelSpectrogram(n_fft=self.filter_length, hop_length=self.hop_length,
+                                                             win_length=self.win_length, power=2, normalized=normalize,
+                                                             sample_rate=self.sampling_rate, f_min=self.mel_fmin,
+                                                             f_max=self.mel_fmax, n_mels=self.n_mel_channels,
+                                                             norm="slaney")
+        self.mel_norm_file = mel_norm_file
+        if self.mel_norm_file is not None:
+            self.mel_norms = torch.load(self.mel_norm_file)
+        else:
+            self.mel_norms = None
+
+    def forward(self, inp):
+        if len(inp.shape) == 3:  # Automatically squeeze out the channels dimension if it is present (assuming mono-audio)
+            inp = inp.squeeze(1)
+        assert len(inp.shape) == 2
+        if torch.backends.mps.is_available():
+            inp = inp.to('cpu')
+        self.mel_stft = self.mel_stft.to(inp.device)
+        mel = self.mel_stft(inp)
+        # Perform dynamic range compression
+        mel = torch.log(torch.clamp(mel, min=1e-5))
+        if self.mel_norms is not None:
+            self.mel_norms = self.mel_norms.to(mel.device)
+            mel = mel / self.mel_norms.unsqueeze(0).unsqueeze(-1)
+        return mel
+
+
+class CheckpointedLayer(nn.Module):
+    """
+    Wraps a module. When forward() is called, passes kwargs that require_grad through torch.checkpoint() and bypasses
+    checkpoint for all other args.
+    """
+    def __init__(self, wrap):
+        super().__init__()
+        self.wrap = wrap
+
+    def forward(self, x, *args, **kwargs):
+        for k, v in kwargs.items():
+            assert not (isinstance(v, torch.Tensor) and v.requires_grad)  # This would screw up checkpointing.
+        partial = functools.partial(self.wrap, **kwargs)
+        return partial(x, *args)
+
+
+class CheckpointedXTransformerEncoder(nn.Module):
+    """
+    Wraps a ContinuousTransformerWrapper and applies CheckpointedLayer to each layer and permutes from channels-mid
+    to channels-last that XTransformer expects.
+    """
+    def __init__(self, needs_permute=True, exit_permute=True, checkpoint=True, **xtransformer_kwargs):
+        super().__init__()
+        self.transformer = ContinuousTransformerWrapper(**xtransformer_kwargs)
+        self.needs_permute = needs_permute
+        self.exit_permute = exit_permute
+
+        if not checkpoint:
+            return
+        for i in range(len(self.transformer.attn_layers.layers)):
+            n, b, r = self.transformer.attn_layers.layers[i]
+            self.transformer.attn_layers.layers[i] = nn.ModuleList([n, CheckpointedLayer(b), r])
+
+    def forward(self, x, **kwargs):
+        if self.needs_permute:
+            x = x.permute(0,2,1)
+        h = self.transformer(x, **kwargs)
+        if self.exit_permute:
+            h = h.permute(0,2,1)
+        return h

speech/cosyvoice/transformer/xtransformers.py

@@ -0,0 +1,1248 @@
+import math
+from collections import namedtuple
+from functools import partial
+from inspect import isfunction
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch import nn, einsum
+
+DEFAULT_DIM_HEAD = 64
+
+Intermediates = namedtuple('Intermediates', [
+    'pre_softmax_attn',
+    'post_softmax_attn'
+])
+
+LayerIntermediates = namedtuple('Intermediates', [
+    'hiddens',
+    'attn_intermediates',
+    'past_key_values',
+])
+
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def cast_tuple(val, depth):
+    return val if isinstance(val, tuple) else (val,) * depth
+
+
+class always():
+    def __init__(self, val):
+        self.val = val
+
+    def __call__(self, *args, **kwargs):
+        return self.val
+
+
+class not_equals():
+    def __init__(self, val):
+        self.val = val
+
+    def __call__(self, x, *args, **kwargs):
+        return x != self.val
+
+
+class equals():
+    def __init__(self, val):
+        self.val = val
+
+    def __call__(self, x, *args, **kwargs):
+        return x == self.val
+
+
+def max_neg_value(tensor):
+    return -torch.finfo(tensor.dtype).max
+
+
+def l2norm(t):
+    return F.normalize(t, p=2, dim=-1)
+
+
+# init helpers
+
+def init_zero_(layer):
+    nn.init.constant_(layer.weight, 0.)
+    if exists(layer.bias):
+        nn.init.constant_(layer.bias, 0.)
+
+
+# keyword argument helpers
+
+def pick_and_pop(keys, d):
+    values = list(map(lambda key: d.pop(key), keys))
+    return dict(zip(keys, values))
+
+
+def group_dict_by_key(cond, d):
+    return_val = [dict(), dict()]
+    for key in d.keys():
+        match = bool(cond(key))
+        ind = int(not match)
+        return_val[ind][key] = d[key]
+    return (*return_val,)
+
+
+def string_begins_with(prefix, str):
+    return str.startswith(prefix)
+
+
+def group_by_key_prefix(prefix, d):
+    return group_dict_by_key(partial(string_begins_with, prefix), d)
+
+
+def groupby_prefix_and_trim(prefix, d):
+    kwargs_with_prefix, kwargs = group_dict_by_key(partial(string_begins_with, prefix), d)
+    kwargs_without_prefix = dict(map(lambda x: (x[0][len(prefix):], x[1]), tuple(kwargs_with_prefix.items())))
+    return kwargs_without_prefix, kwargs
+
+
+# activations
+
+class ReluSquared(nn.Module):
+    def forward(self, x):
+        return F.relu(x) ** 2
+
+
+# positional embeddings
+
+class AbsolutePositionalEmbedding(nn.Module):
+    def __init__(self, dim, max_seq_len):
+        super().__init__()
+        self.scale = dim ** -0.5
+        self.emb = nn.Embedding(max_seq_len, dim)
+
+    def forward(self, x):
+        n = torch.arange(x.shape[1], device=x.device)
+        pos_emb = self.emb(n)
+        pos_emb = rearrange(pos_emb, 'n d -> () n d')
+        return pos_emb * self.scale
+
+
+class FixedPositionalEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+
+    def forward(self, x, seq_dim=1, offset=0):
+        t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq) + offset
+        sinusoid_inp = torch.einsum('i , j -> i j', t, self.inv_freq)
+        emb = torch.cat((sinusoid_inp.sin(), sinusoid_inp.cos()), dim=-1)
+        return rearrange(emb, 'n d -> () n d')
+
+
+class RelativePositionBias(nn.Module):
+    def __init__(self, scale, causal=False, num_buckets=32, max_distance=128, heads=8):
+        super().__init__()
+        self.scale = scale
+        self.causal = causal
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.relative_attention_bias = nn.Embedding(num_buckets, heads)
+
+    @staticmethod
+    def _relative_position_bucket(relative_position, causal=True, num_buckets=32, max_distance=128):
+        ret = 0
+        n = -relative_position
+        if not causal:
+            num_buckets //= 2
+            ret += (n < 0).long() * num_buckets
+            n = torch.abs(n)
+        else:
+            n = torch.max(n, torch.zeros_like(n))
+
+        max_exact = num_buckets // 2
+        is_small = n < max_exact
+
+        val_if_large = max_exact + (
+                torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact)
+        ).long()
+        val_if_large = torch.min(val_if_large, torch.full_like(val_if_large, num_buckets - 1))
+
+        ret += torch.where(is_small, n, val_if_large)
+        return ret
+
+    def forward(self, qk_dots):
+        i, j, device = *qk_dots.shape[-2:], qk_dots.device
+        q_pos = torch.arange(i, dtype=torch.long, device=device)
+        k_pos = torch.arange(j, dtype=torch.long, device=device)
+        rel_pos = k_pos[None, :] - q_pos[:, None]
+        rp_bucket = self._relative_position_bucket(rel_pos, causal=self.causal, num_buckets=self.num_buckets,
+                                                   max_distance=self.max_distance)
+        values = self.relative_attention_bias(rp_bucket)
+        bias = rearrange(values, 'i j h -> () h i j')
+        return qk_dots + (bias * self.scale)
+
+
+class AlibiPositionalBias(nn.Module):
+    def __init__(self, heads, **kwargs):
+        super().__init__()
+        self.heads = heads
+        slopes = torch.Tensor(self._get_slopes(heads))
+        slopes = rearrange(slopes, 'h -> () h () ()')
+        self.register_buffer('slopes', slopes, persistent=False)
+        self.register_buffer('bias', None, persistent=False)
+
+    @staticmethod
+    def _get_slopes(heads):
+        def get_slopes_power_of_2(n):
+            start = (2 ** (-2 ** -(math.log2(n) - 3)))
+            ratio = start
+            return [start * ratio ** i for i in range(n)]
+
+        if math.log2(heads).is_integer():
+            return get_slopes_power_of_2(heads)
+
+        closest_power_of_2 = 2 ** math.floor(math.log2(heads))
+        return get_slopes_power_of_2(closest_power_of_2) + get_slopes_power_of_2(2 * closest_power_of_2)[0::2][
+                                                           :heads - closest_power_of_2]
+
+    def forward(self, qk_dots):
+        h, i, j, device = *qk_dots.shape[-3:], qk_dots.device
+
+        if exists(self.bias) and self.bias.shape[-1] >= j:
+            return qk_dots + self.bias[..., :j]
+
+        bias = torch.arange(j, device=device)
+        bias = rearrange(bias, 'j -> () () () j')
+        bias = bias * self.slopes
+
+        num_heads_unalibied = h - bias.shape[1]
+        bias = F.pad(bias, (0, 0, 0, 0, 0, num_heads_unalibied))
+
+        self.register_buffer('bias', bias, persistent=False)
+        return qk_dots + self.bias
+
+
+class LearnedAlibiPositionalBias(AlibiPositionalBias):
+    def __init__(self, heads, bidirectional=False):
+        super().__init__(heads)
+        los_slopes = torch.log(self.slopes)
+        self.learned_logslopes = nn.Parameter(los_slopes)
+
+        self.bidirectional = bidirectional
+        if self.bidirectional:
+            self.learned_logslopes_future = nn.Parameter(los_slopes)
+
+    def forward(self, qk_dots):
+        h, i, j, device = *qk_dots.shape[-3:], qk_dots.device
+
+        def get_slopes(param):
+            return F.pad(param.exp(), (0, 0, 0, 0, 0, h - param.shape[1]))
+
+        if exists(self.bias) and self.bias.shape[-1] >= j:
+            bias = self.bias[..., :i, :j]
+        else:
+            i_arange = torch.arange(i, device=device)
+            j_arange = torch.arange(j, device=device)
+            bias = rearrange(j_arange, 'j -> 1 1 1 j') - rearrange(i_arange, 'i -> 1 1 i 1')
+            self.register_buffer('bias', bias, persistent=False)
+
+        if self.bidirectional:
+            past_slopes = get_slopes(self.learned_logslopes)
+            future_slopes = get_slopes(self.learned_logslopes_future)
+            bias = torch.tril(bias * past_slopes) + torch.triu(bias * future_slopes)
+        else:
+            slopes = get_slopes(self.learned_logslopes)
+            bias = bias * slopes
+
+        return qk_dots + bias
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+
+    def forward(self, max_seq_len, device):
+        t = torch.arange(max_seq_len, device=device).type_as(self.inv_freq)
+        freqs = torch.einsum('i , j -> i j', t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return rearrange(emb, 'n d -> () () n d')
+
+
+def rotate_half(x):
+    x = rearrange(x, '... (j d) -> ... j d', j=2)
+    x1, x2 = x.unbind(dim=-2)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(t, freqs):
+    seq_len = t.shape[-2]
+    freqs = freqs[:, :, -seq_len:]
+    return (t * freqs.cos()) + (rotate_half(t) * freqs.sin())
+
+
+# norms
+
+class Scale(nn.Module):
+    def __init__(self, value, fn):
+        super().__init__()
+        self.value = value
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        out = self.fn(x, **kwargs)
+        scale_fn = lambda t: t * self.value
+
+        if not isinstance(out, tuple):
+            return scale_fn(out)
+
+        return (scale_fn(out[0]), *out[1:])
+
+
+class Rezero(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+        self.g = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x, **kwargs):
+        out = self.fn(x, **kwargs)
+        rezero_fn = lambda t: t * self.g
+
+        if not isinstance(out, tuple):
+            return rezero_fn(out)
+
+        return (rezero_fn(out[0]), *out[1:])
+
+
+class ScaleNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.scale = dim ** -0.5
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(1))
+
+    def forward(self, x):
+        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
+        return x / norm.clamp(min=self.eps) * self.g
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-8):
+        super().__init__()
+        self.scale = dim ** -0.5
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
+        return x / norm.clamp(min=self.eps) * self.g
+
+
+class RMSScaleShiftNorm(nn.Module):
+    def __init__(self, dim, eps=1e-8):
+        super().__init__()
+        self.scale = dim ** -0.5
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(dim))
+        self.scale_shift_process = nn.Linear(dim * 2, dim * 2)
+
+    def forward(self, x, norm_scale_shift_inp):
+        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
+        norm = x / norm.clamp(min=self.eps) * self.g
+
+        ss_emb = self.scale_shift_process(norm_scale_shift_inp)
+        scale, shift = torch.chunk(ss_emb, 2, dim=1)
+        h = norm * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        return h
+
+
+# residual and residual gates
+
+class Residual(nn.Module):
+    def __init__(self, dim, scale_residual=False):
+        super().__init__()
+        self.residual_scale = nn.Parameter(torch.ones(dim)) if scale_residual else None
+
+    def forward(self, x, residual):
+        if exists(self.residual_scale):
+            residual = residual * self.residual_scale
+
+        return x + residual
+
+
+class GRUGating(nn.Module):
+    def __init__(self, dim, scale_residual=False):
+        super().__init__()
+        self.gru = nn.GRUCell(dim, dim)
+        self.residual_scale = nn.Parameter(torch.ones(dim)) if scale_residual else None
+
+    def forward(self, x, residual):
+        if exists(self.residual_scale):
+            residual = residual * self.residual_scale
+
+        gated_output = self.gru(
+            rearrange(x, 'b n d -> (b n) d'),
+            rearrange(residual, 'b n d -> (b n) d')
+        )
+
+        return gated_output.reshape_as(x)
+
+
+# token shifting
+
+def shift(t, amount, mask=None):
+    if amount == 0:
+        return t
+
+    if exists(mask):
+        t = t.masked_fill(~mask[..., None], 0.)
+
+    return F.pad(t, (0, 0, amount, -amount), value=0.)
+
+
+class ShiftTokens(nn.Module):
+    def __init__(self, shifts, fn):
+        super().__init__()
+        self.fn = fn
+        self.shifts = tuple(shifts)
+
+    def forward(self, x, **kwargs):
+        mask = kwargs.get('mask', None)
+        shifts = self.shifts
+        segments = len(shifts)
+        feats_per_shift = x.shape[-1] // segments
+        splitted = x.split(feats_per_shift, dim=-1)
+        segments_to_shift, rest = splitted[:segments], splitted[segments:]
+        segments_to_shift = list(map(lambda args: shift(*args, mask=mask), zip(segments_to_shift, shifts)))
+        x = torch.cat((*segments_to_shift, *rest), dim=-1)
+        return self.fn(x, **kwargs)
+
+
+# feedforward
+
+class GLU(nn.Module):
+    def __init__(self, dim_in, dim_out, activation):
+        super().__init__()
+        self.act = activation
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * self.act(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(
+            self,
+            dim,
+            dim_out=None,
+            mult=4,
+            glu=False,
+            relu_squared=False,
+            post_act_ln=False,
+            dropout=0.,
+            zero_init_output=False
+    ):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = default(dim_out, dim)
+        activation = ReluSquared() if relu_squared else nn.GELU()
+
+        project_in = nn.Sequential(
+            nn.Linear(dim, inner_dim),
+            activation
+        ) if not glu else GLU(dim, inner_dim, activation)
+
+        self.net = nn.Sequential(
+            project_in,
+            nn.LayerNorm(inner_dim) if post_act_ln else nn.Identity(),
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim_out)
+        )
+
+        # init last linear layer to 0
+        if zero_init_output:
+            init_zero_(self.net[-1])
+
+    def forward(self, x):
+        return self.net(x)
+
+
+# attention.
+
+class Attention(nn.Module):
+    def __init__(
+            self,
+            dim,
+            dim_head=DEFAULT_DIM_HEAD,
+            heads=8,
+            causal=False,
+            talking_heads=False,
+            head_scale=False,
+            collab_heads=False,
+            collab_compression=.3,
+            sparse_topk=None,
+            use_entmax15=False,
+            num_mem_kv=0,
+            dropout=0.,
+            on_attn=False,
+            gate_values=False,
+            zero_init_output=False,
+            max_attend_past=None,
+            qk_norm=False,
+            scale_init_value=None,
+            rel_pos_bias=False,
+            rel_pos_num_buckets=32,
+            rel_pos_max_distance=128,
+    ):
+        super().__init__()
+        self.scale = dim_head ** -0.5
+
+        self.heads = heads
+        self.causal = causal
+        self.max_attend_past = max_attend_past
+
+        qk_dim = v_dim = dim_head * heads
+
+        # collaborative heads
+        self.collab_heads = collab_heads
+        if self.collab_heads:
+            qk_dim = int(collab_compression * qk_dim)
+            self.collab_mixing = nn.Parameter(torch.randn(heads, qk_dim))
+
+        self.to_q = nn.Linear(dim, qk_dim, bias=False)
+        self.to_k = nn.Linear(dim, qk_dim, bias=False)
+        self.to_v = nn.Linear(dim, v_dim, bias=False)
+
+        self.dropout = nn.Dropout(dropout)
+
+        # add GLU gating for aggregated values, from alphafold2
+        self.to_v_gate = None
+        if gate_values:
+            self.to_v_gate = nn.Linear(dim, v_dim)
+            nn.init.constant_(self.to_v_gate.weight, 0)
+            nn.init.constant_(self.to_v_gate.bias, 1)
+
+        # cosine sim attention
+        self.qk_norm = qk_norm
+        if qk_norm:
+            scale_init_value = default(scale_init_value,
+                                       -3)  # if not provided, initialize as though it were sequence length of 1024
+            self.scale = nn.Parameter(torch.ones(1, heads, 1, 1) * scale_init_value)
+
+        # talking heads
+        self.talking_heads = talking_heads
+        if talking_heads:
+            self.pre_softmax_proj = nn.Parameter(torch.randn(heads, heads))
+            self.post_softmax_proj = nn.Parameter(torch.randn(heads, heads))
+
+        # head scaling
+        self.head_scale = head_scale
+        if head_scale:
+            self.head_scale_params = nn.Parameter(torch.ones(1, heads, 1, 1))
+
+        # explicit topk sparse attention
+        self.sparse_topk = sparse_topk
+
+        # entmax
+        self.attn_fn = F.softmax
+
+        # add memory key / values
+        self.num_mem_kv = num_mem_kv
+        if num_mem_kv > 0:
+            self.mem_k = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
+            self.mem_v = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
+
+        # attention on attention
+        self.attn_on_attn = on_attn
+        self.to_out = nn.Sequential(nn.Linear(v_dim, dim * 2), nn.GLU()) if on_attn else nn.Linear(v_dim, dim)
+
+        self.rel_pos_bias = rel_pos_bias
+        if rel_pos_bias:
+            assert rel_pos_num_buckets <= rel_pos_max_distance, 'number of relative position buckets must be less than the relative position max distance'
+            self.rel_pos = RelativePositionBias(scale=dim_head ** 0.5, causal=causal, heads=heads,
+                                                num_buckets=rel_pos_num_buckets, max_distance=rel_pos_max_distance)
+
+        # init output projection 0
+        if zero_init_output:
+            init_zero_(self.to_out)
+
+    def forward(
+            self,
+            x,
+            context=None,
+            mask=None,
+            context_mask=None,
+            attn_mask=None,
+            sinusoidal_emb=None,
+            rotary_pos_emb=None,
+            prev_attn=None,
+            mem=None,
+            layer_past=None,
+    ):
+        b, n, _, h, talking_heads, collab_heads, head_scale, scale, device, has_context = *x.shape, self.heads, self.talking_heads, self.collab_heads, self.head_scale, self.scale, x.device, exists(
+            context)
+        kv_input = default(context, x)
+
+        q_input = x
+        k_input = kv_input
+        v_input = kv_input
+
+        if exists(mem):
+            k_input = torch.cat((mem, k_input), dim=-2)
+            v_input = torch.cat((mem, v_input), dim=-2)
+
+        if exists(sinusoidal_emb):
+            # in shortformer, the query would start at a position offset depending on the past cached memory
+            offset = k_input.shape[-2] - q_input.shape[-2]
+            q_input = q_input + sinusoidal_emb(q_input, offset=offset)
+            k_input = k_input + sinusoidal_emb(k_input)
+
+        q = self.to_q(q_input)
+        k = self.to_k(k_input)
+        v = self.to_v(v_input)
+
+        if not collab_heads:
+            q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v))
+        else:
+            q = einsum('b i d, h d -> b h i d', q, self.collab_mixing)
+            k = rearrange(k, 'b n d -> b () n d')
+            v = rearrange(v, 'b n (h d) -> b h n d', h=h)
+
+        if layer_past is not None:
+            past_key, past_value = layer_past
+            k = torch.cat([past_key, k], dim=-2)
+            v = torch.cat([past_value, v], dim=-2)
+        k_cache = k
+        v_cache = v
+
+        if exists(rotary_pos_emb) and not has_context:
+            l = rotary_pos_emb.shape[-1]
+            (ql, qr), (kl, kr), (vl, vr) = map(lambda t: (t[..., :l], t[..., l:]), (q, k, v))
+            ql, kl, vl = map(lambda t: apply_rotary_pos_emb(t, rotary_pos_emb), (ql, kl, vl))
+            q, k, v = map(lambda t: torch.cat(t, dim=-1), ((ql, qr), (kl, kr), (vl, vr)))
+
+        input_mask = None
+        if any(map(exists, (mask, context_mask))):
+            q_mask = default(mask, lambda: torch.ones((b, n), device=device).bool())
+            k_mask = q_mask if not exists(context) else context_mask
+            k_mask = default(k_mask, lambda: torch.ones((b, k.shape[-2]), device=device).bool())
+            q_mask = rearrange(q_mask, 'b i -> b () i ()')
+            k_mask = rearrange(k_mask, 'b j -> b () () j')
+            input_mask = q_mask * k_mask
+
+        if self.num_mem_kv > 0:
+            mem_k, mem_v = map(lambda t: repeat(t, 'h n d -> b h n d', b=b), (self.mem_k, self.mem_v))
+            k = torch.cat((mem_k, k), dim=-2)
+            v = torch.cat((mem_v, v), dim=-2)
+            if exists(input_mask):
+                input_mask = F.pad(input_mask, (self.num_mem_kv, 0), value=True)
+
+        if collab_heads:
+            k = k.expand(-1, h, -1, -1)
+
+        if self.qk_norm:
+            q, k = map(l2norm, (q, k))
+            scale = 1 / (self.scale.exp().clamp(min=1e-2))
+
+        dots = einsum('b h i d, b h j d -> b h i j', q, k) * scale
+        mask_value = max_neg_value(dots)
+
+        if exists(prev_attn):
+            dots = dots + prev_attn
+
+        pre_softmax_attn = dots.clone()
+
+        if talking_heads:
+            dots = einsum('b h i j, h k -> b k i j', dots, self.pre_softmax_proj).contiguous()
+
+        if self.rel_pos_bias:
+            dots = self.rel_pos(dots)
+
+        if exists(input_mask):
+            dots.masked_fill_(~input_mask, mask_value)
+            del input_mask
+
+        if exists(attn_mask):
+            assert 2 <= attn_mask.ndim <= 4, 'attention mask must have greater than 2 dimensions but less than or equal to 4'
+            if attn_mask.ndim == 2:
+                attn_mask = rearrange(attn_mask, 'i j -> () () i j')
+            elif attn_mask.ndim == 3:
+                attn_mask = rearrange(attn_mask, 'h i j -> () h i j')
+            dots.masked_fill_(~attn_mask, mask_value)
+
+        if exists(self.max_attend_past):
+            i, j = dots.shape[-2:]
+            range_q = torch.arange(j - i, j, device=device)
+            range_k = torch.arange(j, device=device)
+            dist = rearrange(range_q, 'i -> () () i ()') - rearrange(range_k, 'j -> () () () j')
+            mask = dist > self.max_attend_past
+            dots.masked_fill_(mask, mask_value)
+            del mask
+
+        if self.causal:
+            i, j = dots.shape[-2:]
+            r = torch.arange(i, device=device)
+            mask = rearrange(r, 'i -> () () i ()') < rearrange(r, 'j -> () () () j')
+            mask = F.pad(mask, (j - i, 0), value=False)
+            dots.masked_fill_(mask, mask_value)
+            del mask
+
+        if exists(self.sparse_topk) and self.sparse_topk < dots.shape[-1]:
+            top, _ = dots.topk(self.sparse_topk, dim=-1)
+            vk = top[..., -1].unsqueeze(-1).expand_as(dots)
+            mask = dots < vk
+            dots.masked_fill_(mask, mask_value)
+            del mask
+
+        attn = self.attn_fn(dots, dim=-1)
+        post_softmax_attn = attn.clone()
+
+        attn = self.dropout(attn)
+
+        if talking_heads:
+            attn = einsum('b h i j, h k -> b k i j', attn, self.post_softmax_proj).contiguous()
+
+        out = einsum('b h i j, b h j d -> b h i d', attn, v)
+
+        if head_scale:
+            out = out * self.head_scale_params
+
+        out = rearrange(out, 'b h n d -> b n (h d)')
+
+        if exists(self.to_v_gate):
+            gates = self.to_v_gate(x)
+            out = out * gates.sigmoid()
+
+        intermediates = Intermediates(
+            pre_softmax_attn=pre_softmax_attn,
+            post_softmax_attn=post_softmax_attn
+        )
+
+        return self.to_out(out), intermediates, k_cache, v_cache
+
+
+class AttentionLayers(nn.Module):
+    def __init__(
+            self,
+            dim,
+            depth,
+            heads=8,
+            causal=False,
+            cross_attend=False,
+            only_cross=False,
+            use_scalenorm=False,
+            use_rms_scaleshift_norm=False,
+            use_rmsnorm=False,
+            use_rezero=False,
+            alibi_pos_bias=False,
+            alibi_num_heads=None,
+            alibi_learned=False,
+            position_infused_attn=False,
+            rotary_pos_emb=False,
+            rotary_emb_dim=None,
+            custom_layers=None,
+            sandwich_coef=None,
+            par_ratio=None,
+            residual_attn=False,
+            cross_residual_attn=False,
+            macaron=False,
+            pre_norm=True,
+            gate_residual=False,
+            scale_residual=False,
+            shift_tokens=0,
+            sandwich_norm=False,
+            use_qk_norm_attn=False,
+            qk_norm_attn_seq_len=None,
+            zero_init_branch_output=False,
+            **kwargs
+    ):
+        super().__init__()
+        ff_kwargs, kwargs = groupby_prefix_and_trim('ff_', kwargs)
+        attn_kwargs, _ = groupby_prefix_and_trim('attn_', kwargs)
+
+        dim_head = attn_kwargs.get('dim_head', DEFAULT_DIM_HEAD)
+
+        self.dim = dim
+        self.depth = depth
+        self.layers = nn.ModuleList([])
+        self.causal = causal
+
+        rel_pos_bias = 'rel_pos_bias' in attn_kwargs
+        self.has_pos_emb = position_infused_attn or rel_pos_bias or rotary_pos_emb
+        self.pia_pos_emb = FixedPositionalEmbedding(dim) if position_infused_attn else None
+
+        rotary_emb_dim = max(default(rotary_emb_dim, dim_head // 2), 32)
+        self.rotary_pos_emb = RotaryEmbedding(rotary_emb_dim) if rotary_pos_emb else None
+
+        assert not (
+                    alibi_pos_bias and rel_pos_bias), 'you can only choose Alibi positional bias or T5 relative positional bias, not both'
+
+        if alibi_pos_bias:
+            alibi_num_heads = default(alibi_num_heads, heads)
+            assert alibi_num_heads <= heads, 'number of ALiBi heads must be less than the total number of heads'
+            alibi_pos_klass = LearnedAlibiPositionalBias if alibi_learned or not causal else AlibiPositionalBias
+            self.rel_pos = alibi_pos_klass(heads=alibi_num_heads, bidirectional=not causal)
+        else:
+            self.rel_pos = None
+
+        assert not (not pre_norm and sandwich_norm), 'sandwich norm cannot be used when not using prenorm'
+        self.pre_norm = pre_norm
+        self.sandwich_norm = sandwich_norm
+
+        self.residual_attn = residual_attn
+        self.cross_residual_attn = cross_residual_attn
+        self.cross_attend = cross_attend
+
+        norm_class = ScaleNorm if use_scalenorm else nn.LayerNorm
+        norm_class = RMSNorm if use_rmsnorm else norm_class
+        norm_class = RMSScaleShiftNorm if use_rms_scaleshift_norm else norm_class
+        norm_fn = partial(norm_class, dim)
+
+        norm_fn = nn.Identity if use_rezero else norm_fn
+        branch_fn = Rezero if use_rezero else None
+
+        if cross_attend and not only_cross:
+            default_block = ('a', 'c', 'f')
+        elif cross_attend and only_cross:
+            default_block = ('c', 'f')
+        else:
+            default_block = ('a', 'f')
+
+        if macaron:
+            default_block = ('f',) + default_block
+
+        # qk normalization
+
+        if use_qk_norm_attn:
+            attn_scale_init_value = -math.log(math.log2(qk_norm_attn_seq_len ** 2 - qk_norm_attn_seq_len)) if exists(
+                qk_norm_attn_seq_len) else None
+            attn_kwargs = {**attn_kwargs, 'qk_norm': True, 'scale_init_value': attn_scale_init_value}
+
+        # zero init
+
+        if zero_init_branch_output:
+            attn_kwargs = {**attn_kwargs, 'zero_init_output': True}
+            ff_kwargs = {**ff_kwargs, 'zero_init_output': True}
+
+        # calculate layer block order
+
+        if exists(custom_layers):
+            layer_types = custom_layers
+        elif exists(par_ratio):
+            par_depth = depth * len(default_block)
+            assert 1 < par_ratio <= par_depth, 'par ratio out of range'
+            default_block = tuple(filter(not_equals('f'), default_block))
+            par_attn = par_depth // par_ratio
+            depth_cut = par_depth * 2 // 3  # 2 / 3 attention layer cutoff suggested by PAR paper
+            par_width = (depth_cut + depth_cut // par_attn) // par_attn
+            assert len(default_block) <= par_width, 'default block is too large for par_ratio'
+            par_block = default_block + ('f',) * (par_width - len(default_block))
+            par_head = par_block * par_attn
+            layer_types = par_head + ('f',) * (par_depth - len(par_head))
+        elif exists(sandwich_coef):
+            assert sandwich_coef > 0 and sandwich_coef <= depth, 'sandwich coefficient should be less than the depth'
+            layer_types = ('a',) * sandwich_coef + default_block * (depth - sandwich_coef) + ('f',) * sandwich_coef
+        else:
+            layer_types = default_block * depth
+
+        self.layer_types = layer_types
+        self.num_attn_layers = len(list(filter(equals('a'), layer_types)))
+
+        # calculate token shifting
+
+        shift_tokens = cast_tuple(shift_tokens, len(layer_types))
+
+        # iterate and construct layers
+
+        for ind, (layer_type, layer_shift_tokens) in enumerate(zip(self.layer_types, shift_tokens)):
+            is_last_layer = ind == (len(self.layer_types) - 1)
+
+            if layer_type == 'a':
+                layer = Attention(dim, heads=heads, causal=causal, **attn_kwargs)
+            elif layer_type == 'c':
+                layer = Attention(dim, heads=heads, **attn_kwargs)
+            elif layer_type == 'f':
+                layer = FeedForward(dim, **ff_kwargs)
+                layer = layer if not macaron else Scale(0.5, layer)
+            else:
+                raise Exception(f'invalid layer type {layer_type}')
+
+            if layer_shift_tokens > 0:
+                shift_range_upper = layer_shift_tokens + 1
+                shift_range_lower = -layer_shift_tokens if not causal else 0
+                layer = ShiftTokens(range(shift_range_lower, shift_range_upper), layer)
+
+            if exists(branch_fn):
+                layer = branch_fn(layer)
+
+            residual_fn = GRUGating if gate_residual else Residual
+            residual = residual_fn(dim, scale_residual=scale_residual)
+
+            layer_uses_qk_norm = use_qk_norm_attn and layer_type in ('a', 'c')
+
+            pre_branch_norm = norm_fn() if pre_norm and not layer_uses_qk_norm else None
+            post_branch_norm = norm_fn() if sandwich_norm or layer_uses_qk_norm else None
+            post_main_norm = norm_fn() if not pre_norm and not is_last_layer else None
+
+            norms = nn.ModuleList([
+                pre_branch_norm,
+                post_branch_norm,
+                post_main_norm
+            ])
+
+            self.layers.append(nn.ModuleList([
+                norms,
+                layer,
+                residual
+            ]))
+
+    def forward(
+            self,
+            x,
+            context=None,
+            full_context=None,  # for passing a list of hidden states from an encoder
+            mask=None,
+            context_mask=None,
+            attn_mask=None,
+            mems=None,
+            return_hiddens=False,
+            norm_scale_shift_inp=None,
+            past_key_values=None,
+            expected_seq_len=None,
+    ):
+
+        assert not (self.cross_attend ^ (exists(context) or exists(
+            full_context))), 'context must be passed in if cross_attend is set to True'
+        assert context is None or full_context is None, 'only one of full_context or context can be provided'
+
+        hiddens = []
+        intermediates = []
+        prev_attn = None
+        prev_cross_attn = None
+
+        mems = mems.copy() if exists(mems) else [None] * self.num_attn_layers
+        norm_args = {}
+        if exists(norm_scale_shift_inp):
+            norm_args['norm_scale_shift_inp'] = norm_scale_shift_inp
+
+        rotary_pos_emb = None
+        if exists(self.rotary_pos_emb):
+            if not self.training and self.causal:
+                assert expected_seq_len is not None, "To decode a transformer with rotary embeddings, you must specify an `expected_seq_len`"
+            elif expected_seq_len is None:
+                expected_seq_len = 0
+            seq_len = x.shape[1]
+            if past_key_values is not None:
+                seq_len += past_key_values[0][0].shape[-2]
+            max_rotary_emb_length = max(list(map(lambda m: (m.shape[1] if exists(m) else 0) + seq_len, mems)) + [expected_seq_len])
+            rotary_pos_emb = self.rotary_pos_emb(max_rotary_emb_length, x.device)
+
+        present_key_values = []
+        cross_attn_count = 0
+        for ind, (layer_type, (norm, block, residual_fn)) in enumerate(zip(self.layer_types, self.layers)):
+            if layer_type == 'a':
+                layer_mem = mems.pop(0) if mems else None
+
+            residual = x
+
+            pre_branch_norm, post_branch_norm, post_main_norm = norm
+
+            if exists(pre_branch_norm):
+                x = pre_branch_norm(x, **norm_args)
+
+            if layer_type == 'a' or layer_type == 'c':
+                if past_key_values is not None:
+                    layer_kv = past_key_values.pop(0)
+                    layer_past = tuple(s.to(x.device) for s in layer_kv)
+                else:
+                    layer_past = None
+
+            if layer_type == 'a':
+                out, inter, k, v = block(x, None, mask, None, attn_mask, self.pia_pos_emb, rotary_pos_emb,
+                                        prev_attn, layer_mem, layer_past)
+            elif layer_type == 'c':
+                if exists(full_context):
+                    out, inter, k, v = block(x, full_context[cross_attn_count], mask, context_mask, None, None,
+                                            None, prev_attn, None, layer_past)
+                else:
+                    out, inter, k, v = block(x, context, mask, context_mask, None, None, None, prev_attn, None, layer_past)
+            elif layer_type == 'f':
+                out = block(x)
+
+            if layer_type == 'a' or layer_type == 'c' and present_key_values is not None:
+                present_key_values.append((k.detach(), v.detach()))
+
+            if exists(post_branch_norm):
+                out = post_branch_norm(out, **norm_args)
+
+            x = residual_fn(out, residual)
+
+            if layer_type in ('a', 'c'):
+                intermediates.append(inter)
+
+            if layer_type == 'a' and self.residual_attn:
+                prev_attn = inter.pre_softmax_attn
+            elif layer_type == 'c' and self.cross_residual_attn:
+                prev_cross_attn = inter.pre_softmax_attn
+
+            if exists(post_main_norm):
+                x = post_main_norm(x, **norm_args)
+
+            if layer_type == 'c':
+                cross_attn_count += 1
+
+            if layer_type == 'f':
+                hiddens.append(x)
+
+        if return_hiddens:
+            intermediates = LayerIntermediates(
+                hiddens=hiddens,
+                attn_intermediates=intermediates,
+                past_key_values=present_key_values
+            )
+
+            return x, intermediates
+
+        return x
+
+
+class Encoder(AttentionLayers):
+    def __init__(self, **kwargs):
+        assert 'causal' not in kwargs, 'cannot set causality on encoder'
+        super().__init__(causal=False, **kwargs)
+
+
+class Decoder(AttentionLayers):
+    def __init__(self, **kwargs):
+        assert 'causal' not in kwargs, 'cannot set causality on decoder'
+        super().__init__(causal=True, **kwargs)
+
+
+class CrossAttender(AttentionLayers):
+    def __init__(self, **kwargs):
+        super().__init__(cross_attend=True, only_cross=True, **kwargs)
+
+
+class ViTransformerWrapper(nn.Module):
+    def __init__(
+            self,
+            *,
+            image_size,
+            patch_size,
+            attn_layers,
+            num_classes=None,
+            dropout=0.,
+            emb_dropout=0.
+    ):
+        super().__init__()
+        assert isinstance(attn_layers, Encoder), 'attention layers must be an Encoder'
+        assert image_size % patch_size == 0, 'image dimensions must be divisible by the patch size'
+        dim = attn_layers.dim
+        num_patches = (image_size // patch_size) ** 2
+        patch_dim = 3 * patch_size ** 2
+
+        self.patch_size = patch_size
+
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+        self.patch_to_embedding = nn.Linear(patch_dim, dim)
+        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.attn_layers = attn_layers
+        self.norm = nn.LayerNorm(dim)
+        self.mlp_head = FeedForward(dim, dim_out=num_classes, dropout=dropout) if exists(num_classes) else None
+
+    def forward(
+            self,
+            img,
+            return_embeddings=False
+    ):
+        p = self.patch_size
+
+        x = rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=p, p2=p)
+        x = self.patch_to_embedding(x)
+        b, n, _ = x.shape
+
+        cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=b)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = x + self.pos_embedding[:, :(n + 1)]
+        x = self.dropout(x)
+
+        x = self.attn_layers(x)
+        x = self.norm(x)
+
+        if not exists(self.mlp_head) or return_embeddings:
+            return x
+
+        return self.mlp_head(x[:, 0])
+
+
+class TransformerWrapper(nn.Module):
+    def __init__(
+            self,
+            *,
+            num_tokens,
+            max_seq_len,
+            attn_layers,
+            emb_dim=None,
+            max_mem_len=0.,
+            shift_mem_down=0,
+            emb_dropout=0.,
+            num_memory_tokens=None,
+            tie_embedding=False,
+            use_pos_emb=True
+    ):
+        super().__init__()
+        assert isinstance(attn_layers, AttentionLayers), 'attention layers must be one of Encoder or Decoder'
+
+        dim = attn_layers.dim
+        emb_dim = default(emb_dim, dim)
+
+        self.max_seq_len = max_seq_len
+        self.max_mem_len = max_mem_len
+        self.shift_mem_down = shift_mem_down
+
+        self.token_emb = nn.Embedding(num_tokens, emb_dim)
+        self.pos_emb = AbsolutePositionalEmbedding(emb_dim, max_seq_len) if (
+                    use_pos_emb and not attn_layers.has_pos_emb) else always(0)
+        self.emb_dropout = nn.Dropout(emb_dropout)
+
+        self.project_emb = nn.Linear(emb_dim, dim) if emb_dim != dim else nn.Identity()
+        self.attn_layers = attn_layers
+        self.norm = nn.LayerNorm(dim)
+
+        self.init_()
+
+        self.to_logits = nn.Linear(dim, num_tokens) if not tie_embedding else lambda t: t @ self.token_emb.weight.t()
+
+        # memory tokens (like [cls]) from Memory Transformers paper
+        num_memory_tokens = default(num_memory_tokens, 0)
+        self.num_memory_tokens = num_memory_tokens
+        if num_memory_tokens > 0:
+            self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim))
+
+    def init_(self):
+        nn.init.kaiming_normal_(self.token_emb.weight)
+
+    def forward(
+            self,
+            x,
+            return_embeddings=False,
+            mask=None,
+            return_hiddens=False,
+            return_attn=False,
+            mems=None,
+            use_cache=False,
+            **kwargs
+    ):
+        b, n, device, num_mem = *x.shape, x.device, self.num_memory_tokens
+        x = self.token_emb(x)
+        x = x + self.pos_emb(x)
+        x = self.emb_dropout(x)
+
+        x = self.project_emb(x)
+
+        if num_mem > 0:
+            mem = repeat(self.memory_tokens, 'n d -> b n d', b=b)
+            x = torch.cat((mem, x), dim=1)
+
+            # auto-handle masking after appending memory tokens
+            if exists(mask):
+                mask = F.pad(mask, (num_mem, 0), value=True)
+
+        if self.shift_mem_down and exists(mems):
+            mems_l, mems_r = mems[:self.shift_mem_down], mems[self.shift_mem_down:]
+            mems = [*mems_r, *mems_l]
+
+        x, intermediates = self.attn_layers(x, mask=mask, mems=mems, return_hiddens=True, **kwargs)
+        x = self.norm(x)
+
+        mem, x = x[:, :num_mem], x[:, num_mem:]
+
+        out = self.to_logits(x) if not return_embeddings else x
+
+        if return_hiddens:
+            hiddens = intermediates.hiddens
+            return out, hiddens
+
+        res = [out]
+        if return_attn:
+            attn_maps = list(map(lambda t: t.post_softmax_attn, intermediates.attn_intermediates))
+            res.append(attn_maps)
+        if use_cache:
+            res.append(intermediates.past_key_values)
+
+        if len(res) > 1:
+            return tuple(res)
+        return res[0]
+
+
+class ContinuousTransformerWrapper(nn.Module):
+    def __init__(
+            self,
+            *,
+            max_seq_len,
+            attn_layers,
+            dim_in=None,
+            dim_out=None,
+            emb_dim=None,
+            emb_dropout=0.,
+            use_pos_emb=True
+    ):
+        super().__init__()
+        assert isinstance(attn_layers, AttentionLayers), 'attention layers must be one of Encoder or Decoder'
+
+        dim = attn_layers.dim
+
+        self.max_seq_len = max_seq_len
+
+        self.pos_emb = AbsolutePositionalEmbedding(dim, max_seq_len) if (
+                    use_pos_emb and not attn_layers.has_pos_emb) else always(0)
+        self.emb_dropout = nn.Dropout(emb_dropout)
+
+        self.project_in = nn.Linear(dim_in, dim) if exists(dim_in) else nn.Identity()
+
+        self.attn_layers = attn_layers
+        self.norm = nn.LayerNorm(dim)
+
+        self.project_out = nn.Linear(dim, dim_out) if exists(dim_out) else nn.Identity()
+
+    def forward(
+            self,
+            x,
+            return_embeddings=False,
+            mask=None,
+            return_attn=False,
+            mems=None,
+            use_cache=False,
+            **kwargs
+    ):
+        b, n, _, device = *x.shape, x.device
+
+        x = self.project_in(x)
+        x = x + self.pos_emb(x)
+        x = self.emb_dropout(x)
+
+        x, intermediates = self.attn_layers(x, mask=mask, mems=mems, return_hiddens=True, **kwargs)
+        x = self.norm(x)
+
+        out = self.project_out(x) if not return_embeddings else x
+
+        res = [out]
+        if return_attn:
+            attn_maps = list(map(lambda t: t.post_softmax_attn, intermediates.attn_intermediates))
+            res.append(attn_maps)
+        if use_cache:
+            res.append(intermediates.past_key_values)
+
+        if len(res) > 1:
+            return tuple(res)
+        return res[0]
+

speech/tools/download_dataset.py

@@ -0,0 +1,371 @@
+import os
+import io
+import time
+from pathlib import Path
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
+from multiprocessing import cpu_count, Queue, Process
+from threading import Lock
+from queue import Empty
+import logging
+from tqdm import tqdm
+from functools import partial
+
+import numpy as np
+import torch
+import torchaudio
+import soundfile as sf
+from datasets import load_dataset
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Global settings
+CHUNK_SIZE = 100  # Process this many samples before showing progress
+MAX_WORKERS = cpu_count() - 1  # Leave one CPU free
+PREFETCH_SIZE = 50  # How many samples to prefetch
+
+def process_single_sample(data, root_save_path):
+    """Process a single sample from the dataset"""
+    try:
+        metadata = data['json']
+        
+        # Prepare paths
+        out_wav_path = os.path.join(root_save_path, metadata['wav'].replace('/mp3', '').replace('.mp3', '.wav'))
+        os.makedirs(os.path.dirname(out_wav_path), exist_ok=True)
+        out_text_path = out_wav_path.replace('.wav', '.txt')
+        # check if the files are existsed
+        if os.path.exists(out_wav_path) and os.path.exists(out_text_path):
+            return metadata['id'], True, None
+        
+        # Save text
+        with open(out_text_path, 'w') as f:
+            f.write(metadata['text'])
+        
+        # Decode and save audio
+        raw_bytes = data['mp3']._hf_encoded['bytes']
+        audio_tensor, sample_rate = torchaudio.load(io.BytesIO(raw_bytes), format='mp3')
+        audio_array = audio_tensor.squeeze().numpy()
+        sf.write(out_wav_path, audio_array, sample_rate)
+        
+        return metadata['id'], True, None
+    except Exception as e:
+        return data.get('json', {}).get('id', 'unknown'), False, str(e)
+
+def process_batch(batch_data, root_save_path):
+    """Process a batch of samples"""
+    results = []
+    for data in batch_data:
+        result = process_single_sample(data, root_save_path)
+        results.append(result)
+    return results
+
+class ParallelDatasetProcessor:
+    """Main class for parallel processing of the dataset"""
+    
+    def __init__(self, language, root_save_path, num_workers=None):
+        self.language = language
+        self.root_save_path = root_save_path
+        self.num_workers = num_workers or MAX_WORKERS
+        self.processed_count = 0
+        self.error_count = 0
+        self.lock = Lock()
+        
+    def process_with_multiprocessing(self):
+        """Process dataset using multiprocessing (fastest for CPU-bound tasks)"""
+        logger.info(f"Starting multiprocessing with {self.num_workers} workers")
+        
+        # Load dataset
+        path = f"Emilia/{self.language.upper()}/*.tar"
+        dataset = load_dataset(
+            "amphion/Emilia-Dataset", 
+            data_files={self.language: path}, 
+            split=self.language, 
+            streaming=True
+        )
+        
+        # Create process pool
+        with ProcessPoolExecutor(max_workers=self.num_workers) as executor:
+            # Submit jobs in batches
+            futures = []
+            batch = []
+            
+            # Progress bar
+            pbar = tqdm(desc="Processing samples", unit="samples")
+            
+            for data in dataset:
+                batch.append(data)
+                
+                if len(batch) >= CHUNK_SIZE:
+                    # Submit batch for processing
+                    future = executor.submit(process_batch, batch, self.root_save_path)
+                    futures.append(future)
+                    batch = []
+                    
+                    # Process completed futures
+                    self._process_completed_futures(futures, pbar, max_pending=self.num_workers * 2)
+            
+            # Submit remaining batch
+            if batch:
+                future = executor.submit(process_batch, batch, self.root_save_path)
+                futures.append(future)
+            
+            # Wait for all remaining futures
+            for future in as_completed(futures):
+                results = future.result()
+                for sample_id, success, error in results:
+                    if success:
+                        self.processed_count += 1
+                    else:
+                        self.error_count += 1
+                        logger.error(f"Error processing {sample_id}: {error}")
+                    pbar.update(1)
+            
+            pbar.close()
+    
+    def process_with_threading(self):
+        """Process dataset using threading (good for I/O-bound tasks)"""
+        logger.info(f"Starting threading with {self.num_workers} workers")
+        
+        # Load dataset
+        path = f"Emilia/{self.language.upper()}/*.tar"
+        dataset = load_dataset(
+            "amphion/Emilia-Dataset", 
+            data_files={self.language: path}, 
+            split=self.language, 
+            streaming=True
+        )
+        
+        # Create thread pool
+        with ThreadPoolExecutor(max_workers=self.num_workers * 2) as executor:
+            futures = []
+            pbar = tqdm(desc="Processing samples", unit="samples")
+            
+            for data in dataset:
+                # Submit individual samples
+                future = executor.submit(process_single_sample, data, self.root_save_path)
+                futures.append((future, data.get('json', {}).get('id', 'unknown')))
+                
+                # Process completed futures
+                if len(futures) >= self.num_workers * 4:
+                    completed = []
+                    for i, (future, sample_id) in enumerate(futures):
+                        if future.done():
+                            try:
+                                _, success, error = future.result()
+                                if success:
+                                    self.processed_count += 1
+                                else:
+                                    self.error_count += 1
+                                    logger.error(f"Error processing {sample_id}: {error}")
+                                pbar.update(1)
+                                completed.append(i)
+                            except Exception as e:
+                                logger.error(f"Exception processing {sample_id}: {e}")
+                                self.error_count += 1
+                                pbar.update(1)
+                                completed.append(i)
+                    
+                    # Remove completed futures
+                    for i in reversed(completed):
+                        futures.pop(i)
+            
+            # Wait for remaining futures
+            for future, sample_id in futures:
+                try:
+                    _, success, error = future.result()
+                    if success:
+                        self.processed_count += 1
+                    else:
+                        self.error_count += 1
+                        logger.error(f"Error processing {sample_id}: {error}")
+                    pbar.update(1)
+                except Exception as e:
+                    logger.error(f"Exception processing {sample_id}: {e}")
+                    self.error_count += 1
+                    pbar.update(1)
+            
+            pbar.close()
+    
+    def process_with_producer_consumer(self):
+        """Process dataset using producer-consumer pattern"""
+        logger.info(f"Starting producer-consumer with {self.num_workers} workers")
+        
+        # Create queues
+        work_queue = Queue(maxsize=PREFETCH_SIZE)
+        result_queue = Queue()
+        
+        # Start producer
+        producer = Process(target=self._producer, args=(work_queue,))
+        producer.start()
+        
+        # Start consumers
+        consumers = []
+        for i in range(self.num_workers):
+            consumer = Process(target=self._consumer, args=(work_queue, result_queue, i))
+            consumer.start()
+            consumers.append(consumer)
+        
+        # Start result processor
+        result_processor = Process(target=self._result_processor, args=(result_queue,))
+        result_processor.start()
+        
+        # Wait for completion
+        producer.join()
+        
+        # Signal consumers to stop
+        for _ in range(self.num_workers):
+            work_queue.put(None)
+        
+        for consumer in consumers:
+            consumer.join()
+        
+        # Signal result processor to stop
+        result_queue.put(None)
+        result_processor.join()
+    
+    def _producer(self, work_queue):
+        """Producer process that reads from dataset"""
+        path = f"Emilia/{self.language.upper()}/*.tar"
+        dataset = load_dataset(
+            "amphion/Emilia-Dataset", 
+            data_files={self.language: path}, 
+            split=self.language, 
+            streaming=True
+        )
+        
+        for data in dataset:
+            work_queue.put(data)
+    
+    def _consumer(self, work_queue, result_queue, worker_id):
+        """Consumer process that processes samples"""
+        while True:
+            try:
+                data = work_queue.get(timeout=1)
+                if data is None:
+                    break
+                
+                result = process_single_sample(data, self.root_save_path)
+                result_queue.put(result)
+            except Empty:
+                continue
+            except Exception as e:
+                logger.error(f"Worker {worker_id} error: {e}")
+    
+    def _result_processor(self, result_queue):
+        """Process results and update progress"""
+        pbar = tqdm(desc="Processing samples", unit="samples")
+        
+        while True:
+            try:
+                result = result_queue.get(timeout=1)
+                if result is None:
+                    break
+                
+                sample_id, success, error = result
+                if success:
+                    self.processed_count += 1
+                else:
+                    self.error_count += 1
+                    logger.error(f"Error processing {sample_id}: {error}")
+                pbar.update(1)
+            except Empty:
+                continue
+        
+        pbar.close()
+    
+    def _process_completed_futures(self, futures, pbar, max_pending):
+        """Process completed futures to avoid memory buildup"""
+        while len(futures) > max_pending:
+            # Wait for at least one to complete
+            completed_futures = []
+            for i, future in enumerate(futures):
+                if future.done():
+                    results = future.result()
+                    for sample_id, success, error in results:
+                        if success:
+                            self.processed_count += 1
+                        else:
+                            self.error_count += 1
+                            logger.error(f"Error processing {sample_id}: {error}")
+                        pbar.update(1)
+                    completed_futures.append(i)
+            
+            # Remove completed futures
+            for i in reversed(completed_futures):
+                futures.pop(i)
+            
+            if not completed_futures:
+                # If nothing completed, wait a bit
+                time.sleep(0.1)
+
+def main():
+    """Main function with different processing options"""
+    language = "en"
+    root_save_path = f'/data/emilia/{language}'
+    os.makedirs(root_save_path, exist_ok=True)
+    
+    # Choose processing method
+    processor = ParallelDatasetProcessor(language, root_save_path)
+    
+    # Method 1: Multiprocessing (recommended for CPU-bound tasks)
+    logger.info("Starting parallel processing...")
+    start_time = time.time()
+    
+    # You can choose one of these methods:
+    processor.process_with_multiprocessing()  # Fastest for this use case
+    # processor.process_with_threading()      # Good for I/O heavy tasks
+    # processor.process_with_producer_consumer()  # Good for memory efficiency
+    
+    elapsed_time = time.time() - start_time
+    logger.info(f"Processing completed in {elapsed_time:.2f} seconds")
+    logger.info(f"Successfully processed: {processor.processed_count} samples")
+    logger.info(f"Errors: {processor.error_count} samples")
+
+# Simplified version for quick use
+def simple_parallel_process():
+    """Simplified parallel processing function"""
+    from joblib import Parallel, delayed
+    
+    language = "en"
+    root_save_path = f'/data/emilia/{language}'
+    os.makedirs(root_save_path, exist_ok=True)
+    
+    # Load dataset
+    path = f"Emilia/{language.upper()}/*.tar"
+    dataset = load_dataset(
+        "amphion/Emilia-Dataset", 
+        data_files={language: path}, 
+        split=language, 
+        streaming=True
+    )
+    
+    # Collect samples in batches
+    batch_size = 1000
+    batch = []
+    
+    def process_batch_simple(batch_data):
+        for data in batch_data:
+            process_single_sample(data, root_save_path)
+    
+    # Process in parallel batches
+    for i, data in enumerate(tqdm(dataset, desc="Loading samples")):
+        batch.append(data)
+        
+        if len(batch) >= batch_size:
+            # Process batch in parallel
+            Parallel(n_jobs=-1, backend='multiprocessing')(
+                delayed(process_single_sample)(sample, root_save_path) 
+                for sample in batch
+            )
+            batch = []
+    
+    # Process remaining batch
+    if batch:
+        Parallel(n_jobs=-1, backend='multiprocessing')(
+            delayed(process_single_sample)(sample, root_save_path) 
+            for sample in batch
+        )
+
+if __name__ == "__main__":
+    main()

speech/tools/download_vivoice.py

@@ -0,0 +1,210 @@
+import os
+import io
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
+from functools import partial
+import multiprocessing
+from typing import Dict, Any, Optional
+
+from datasets import load_dataset
+import soundfile as sf
+import torchaudio
+from tqdm import tqdm
+
+
+def process_single_item(data: Dict[str, Any], root_save_path: str) -> Optional[str]:
+    """
+    Process a single audio item: extract audio, convert, and save with text.
+    
+    Args:
+        data: Dictionary containing audio data and text
+        root_save_path: Root directory for saving files
+        
+    Returns:
+        Audio file path if successful, None if error
+    """
+    try:
+        # Extract audio data
+        raw_bytes = data['audio']._hf_encoded['bytes']
+        audio_name = data['audio']._hf_encoded['path']
+        # Prepare output paths
+        out_wav_path = os.path.join(root_save_path, audio_name)
+        out_text_path = out_wav_path.replace('.wav', '.txt')
+
+        if os.path.exists(out_wav_path) and os.path.exists(out_text_path):
+            # print(f'skip {out_wav_path}')
+            return out_wav_path
+        
+        text = data['text']
+        
+        # Load audio from bytes
+        bytes_io = io.BytesIO(raw_bytes)
+        audio_tensor, sample_rate = torchaudio.load(bytes_io)
+        audio_array = audio_tensor.squeeze().numpy()
+        
+        
+        
+        # Create directory if needed
+        os.makedirs(os.path.dirname(out_wav_path), exist_ok=True)
+        
+        # Save audio and text
+        sf.write(out_wav_path, audio_array, sample_rate)
+        with open(out_text_path, 'w', encoding='utf-8') as f:
+            f.write(text)
+            
+        return out_wav_path
+        
+    except Exception as e:
+        print(f"Error processing {data.get('audio', {})._hf_encoded.get('path', 'unknown')}: {e}")
+        return None
+
+
+def process_dataset_parallel(
+    dataset_name: str = "capleaf/viVoice",
+    root_save_path: str = '/mnt/nvme-temp/vivoice',
+    max_workers: Optional[int] = None,
+    batch_size: int = 100,
+    limit: Optional[int] = None,
+    use_threads: bool = False
+):
+    """
+    Process dataset in parallel with progress tracking.
+    
+    Args:
+        dataset_name: Name of the HuggingFace dataset
+        root_save_path: Root directory for saving files
+        max_workers: Maximum number of parallel workers (None for CPU count)
+        batch_size: Number of items to process in each batch
+        limit: Maximum number of items to process (None for all)
+        use_threads: Use ThreadPoolExecutor instead of ProcessPoolExecutor
+    """
+    # Load dataset
+    ds = load_dataset(dataset_name, streaming=True)
+    
+    # Set up executor
+    if max_workers is None:
+        max_workers = multiprocessing.cpu_count()
+    
+    Executor = ThreadPoolExecutor if use_threads else ProcessPoolExecutor
+    
+    # Process each split
+    for mode in ds:
+        print(f"\nProcessing '{mode}' split...")
+        
+        # Create partial function with root_save_path
+        process_func = partial(process_single_item, root_save_path=root_save_path)
+        
+        # Collect items in batches for better progress tracking
+        batch = []
+        processed_count = 0
+        
+        with Executor(max_workers=max_workers) as executor:
+            # Create progress bar
+            pbar = tqdm(desc=f"Processing {mode}", unit="files")
+            
+            for idx, data in enumerate(ds[mode]):
+                batch.append(data)
+                
+                # Process batch when full or at limit
+                if len(batch) >= batch_size or (limit and idx + 1 >= limit):
+                    # Submit batch for processing
+                    futures = [executor.submit(process_func, item) for item in batch]
+                    
+                    # Wait for completion and update progress
+                    for future in futures:
+                        result = future.result()
+                        if result:
+                            processed_count += 1
+                        pbar.update(1)
+                    
+                    batch = []
+                
+                # Stop if limit reached
+                if limit and idx + 1 >= limit:
+                    break
+            
+            # Process remaining items
+            if batch:
+                futures = [executor.submit(process_func, item) for item in batch]
+                for future in futures:
+                    result = future.result()
+                    if result:
+                        processed_count += 1
+                    pbar.update(1)
+            
+            pbar.close()
+        
+        print(f"Completed {mode}: {processed_count} files processed successfully")
+
+
+def process_dataset_streaming(
+    dataset_name: str = "capleaf/viVoice",
+    root_save_path: str = '/data/vivoice',
+    max_workers: Optional[int] = None,
+    limit: Optional[int] = None
+):
+    """
+    Alternative approach using streaming with thread pool for I/O bound operations.
+    More memory efficient for very large datasets.
+    """
+    ds = load_dataset(dataset_name, streaming=True)
+    
+    if max_workers is None:
+        max_workers = min(32, multiprocessing.cpu_count() * 4)  # More threads for I/O
+    
+    for mode in ds:
+        print(f"\nProcessing '{mode}' split...")
+        
+        process_func = partial(process_single_item, root_save_path=root_save_path)
+        
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            # Submit items as they come from the stream
+            futures = []
+            
+            for idx, data in enumerate(tqdm(ds[mode], desc=f"Submitting {mode}")):
+                future = executor.submit(process_func, data)
+                futures.append(future)
+                
+                # Limit number of pending futures to control memory usage
+                if len(futures) >= max_workers * 2:
+                    # Wait for some to complete
+                    for f in futures[:max_workers]:
+                        f.result()
+                    futures = futures[max_workers:]
+                
+                if limit and idx + 1 >= limit:
+                    break
+            
+            # Wait for remaining futures
+            for future in tqdm(futures, desc="Finishing"):
+                future.result()
+
+
+if __name__ == "__main__":
+    # Example usage - choose one approach:
+    
+    # Approach 1: Process with multiprocessing (good for CPU-bound operations)
+    process_dataset_parallel(
+        dataset_name="capleaf/viVoice",
+        root_save_path='/data/vivoice',
+        max_workers=24,  # Adjust based on your system
+        batch_size=100,
+        limit=None,  # Remove to process all data
+        use_threads=False
+    )
+    
+    # Approach 2: Process with threading (good for I/O-bound operations)
+    # process_dataset_parallel(
+    #     dataset_name="capleaf/viVoice",
+    #     root_save_path='/mnt/nvme-temp/vivoice',
+    #     max_workers=16,  # Can use more threads for I/O
+    #     batch_size=100,
+    #     limit=None,
+    #     use_threads=True
+    # )
+    
+    # Approach 3: Streaming approach (most memory efficient)
+    # process_dataset_streaming(
+    #     dataset_name="capleaf/viVoice",
+    #     root_save_path='/mnt/nvme-temp/vivoice',
+    #     max_workers=16,
+    #     limit=None