Upload 13 files

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+EchoX-Vocoder/g_00500000 filter=lfs diff=lfs merge=lfs -text
+show_case/2.wav filter=lfs diff=lfs merge=lfs -text
+show_case/Translate_de_audio_prompt.wav filter=lfs diff=lfs merge=lfs -text

ACLlama_el_s2s.py

@@ -0,0 +1,565 @@
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import CrossEntropyLoss, CTCLoss
+
+
+from transformers import AutoConfig, AutoModelForCausalLM, \
+                         LlamaConfig, LlamaModel, LlamaForCausalLM
+from transformers.trainer_pt_utils import LabelSmoother
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers import (
+    WhisperProcessor,
+    WhisperModel,
+)
+from T2ULlama_CR_online import T2ULlamaForCausalLM
+
+IGNORE_TOKEN_ID = LabelSmoother.ignore_index
+
+class ACLlamaConfig(LlamaConfig):
+    model_type = "ACLlama"
+    
+def load_whisper(audio_tower_name, device="cuda"):
+    model = WhisperModel.from_pretrained(
+            audio_tower_name,torch_dtype=torch.float16,low_cpu_mem_usage=True).to(device)
+    model.config.forced_decoder_ids = None
+    return model
+
+class LookBackModule(nn.Module):
+    def __init__(self, cfg: LlamaConfig):
+        super().__init__()
+        self.encoder_attn = nn.MultiheadAttention(
+            cfg.hidden_size,
+            cfg.num_attention_heads,
+            dropout=0.1,
+            batch_first=True
+        )
+        self.atten_layer_norm = nn.LayerNorm(cfg.hidden_size)
+
+
+    def forward(self, x, wav_feature, bf_shrink_padding_mask):
+
+        residual = x
+        x, _ = self.encoder_attn(
+            query=x,
+            key=wav_feature,
+            value=wav_feature,
+            key_padding_mask=bf_shrink_padding_mask,
+            #attn_mask=padding_mask,
+        )
+        x += residual
+        x = self.atten_layer_norm(x)
+        return x
+
+class ACLlamaModel(LlamaModel):
+    config_class = ACLlamaConfig
+
+    def __init__(self, config: LlamaConfig):
+        super(ACLlamaModel, self).__init__(config)
+
+        if hasattr(config, "audio_tower"):
+            self.audio_tower = [load_whisper(config.audio_tower)]
+
+        if hasattr(config, "adapter_size"):
+
+            self.mm_projector1 = nn.Linear(config.adapter_size*2 , config.hidden_size)
+            asr_encoder_layer = nn.TransformerEncoderLayer(
+                d_model=config.hidden_size,
+                nhead=config.num_attention_heads,
+                dim_feedforward=config.hidden_size*2,
+                dropout=0.1,
+                norm_first=True
+            )
+            self.lbm =  LookBackModule(config)
+            self.out_norm = nn.LayerNorm(config.hidden_size)
+            self.audio_feature_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+            self.asr_transformer_encoder = nn.TransformerEncoder(asr_encoder_layer, num_layers=1)
+            self.mask_tensor=(torch.ones([1, 2048])>0)
+            self.length=-1
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        audios: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+
+        # HACK: replace back original embeddings for LLaAA pretraining
+        orig_embeds_params = getattr(self, 'orig_embeds_params', None)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        audio_tower = getattr(self, 'audio_tower', None)
+        if audio_tower is not None and (input_ids.shape[1] != 1 or self.training) and audios is not None:
+            audio_tower = audio_tower[0]  # HACK: for FSDP
+            audio_list=[]
+            
+            audio_config = audio_tower.config
+            for audio in audios:
+                with torch.no_grad():
+                    audio_feature = audio_tower.encoder(audio).last_hidden_state
+           
+                audio_feature = audio_feature.view(audio_feature.shape[0], audio_feature.shape[1]//2, 2 * audio_feature.shape[2])
+                audio_feature = self.mm_projector1(audio_feature)
+                audio_feature = self.asr_transformer_encoder(audio_feature)
+                audio_feature = self.out_norm(audio_feature)
+                audio_list.append(audio_feature)
+
+            audio_features = torch.stack(audio_list, dim=0)
+            batch = audio_features.shape[0]
+            audio_turn = audio_features.shape[1]
+            audio_features = audio_features.view((batch * audio_turn,)+audio_features.shape[2:])
+ 
+            predict_logits = self.audio_feature_head(audio_features)
+
+            new_input_embeds = []
+            label_shift = []
+            speech_pos = []
+            label_extend = -1
+            new_input_ids = []
+            tokens = predict_logits.argmax(dim=-1)
+            shrink_mask = tokens.roll(1) != tokens
+            shrink_mask[:,0] = True
+                
+            lengths = shrink_mask.long().sum(-1)
+            shrink_2d = audio_features[shrink_mask]
+            #num_patches = audio_features.shape[1]
+            num_patches = audio_config.audio_patch_size
+            l_index=0
+            shrink_features_raw = []
+            for v, audio_feature, mask in zip(lengths, audio_features, ~shrink_mask):
+                shrink_feature = shrink_2d[l_index:l_index+v]
+                shrink_feature = self.lbm(shrink_feature, audio_feature, bf_shrink_padding_mask=mask)
+                shrink_features_raw.append(shrink_feature)
+                l_index += v
+
+            shrink_features = []
+            for i in range(0, len(shrink_features_raw), audio_turn):
+                shrink_features.append(shrink_features_raw[i:i+audio_turn])  
+            if self.training: 
+                maxn_length = lengths.view(batch,audio_turn).sum(-1).max()
+                label_extend = maxn_length - num_patches * audio_turn
+                old_seq_length = inputs_embeds.shape[1]
+                for cur_input_ids, cur_input_embeds, cur_shrink_features in zip(input_ids, inputs_embeds, shrink_features):
+                    pad_ids = torch.full(size=(maxn_length,), fill_value=audio_config.llm_pad_token_id, dtype=torch.long).to(attention_mask.device)
+                    pad_embeds = self.embed_tokens(pad_ids)
+                    audio_start_token_pos_all = torch.where(cur_input_ids == audio_config.audio_patch_token)[0]
+                    #print(cur_input_embeds.shape,cur_input_ids.shape)
+                    inner_label_shift = []
+                    inner_speech_pos = []
+                    for audio_start_token_pos, shrink_feature in reversed(list(zip(audio_start_token_pos_all, cur_shrink_features))): #zip(audio_start_token_pos_all, cur_shrink_features):
+                        cur_speech_length = shrink_feature.shape[0]
+ 
+                        cur_input_ids = torch.cat((cur_input_ids[:audio_start_token_pos], 
+                          cur_input_ids[audio_start_token_pos: audio_start_token_pos+1].repeat(cur_speech_length), 
+                          cur_input_ids[audio_start_token_pos + num_patches:]), dim=0)
+                        cur_input_embeds = torch.cat((
+                          cur_input_embeds[:audio_start_token_pos],
+                          shrink_feature,
+                          cur_input_embeds[audio_start_token_pos + num_patches:]), dim=0)
+                        inner_label_shift.insert(0, cur_speech_length - num_patches)
+                        inner_speech_pos.insert(0, audio_start_token_pos)
+
+                    label_shift = label_shift + inner_label_shift
+                    speech_pos = speech_pos + inner_speech_pos
+
+                    cur_new_input_embeds = torch.cat((cur_input_embeds, pad_embeds[:old_seq_length + label_extend - cur_input_embeds.shape[0]]),dim=0)
+                    cur_new_input_ids = torch.cat((cur_input_ids, pad_ids[:old_seq_length + label_extend - cur_input_ids.shape[0]]),dim=0)
+                    new_input_embeds.append(cur_new_input_embeds)
+                    new_input_ids.append(cur_new_input_ids)
+                    
+                input_ids = torch.stack(new_input_ids, dim=0)
+                attention_mask=input_ids.ne(audio_config.llm_pad_token_id)
+                inputs_embeds = torch.stack(new_input_embeds, dim=0)
+
+                batch_label_shift = []
+                batch_speech_pos=[]
+                for i in range(0, len(label_shift), audio_turn):
+                    batch_label_shift.append(label_shift[i:i+audio_turn])
+                    batch_speech_pos.append(speech_pos[i:i+audio_turn])
+            else:
+                # Inference mode with batch_size=1
+                assert input_ids.shape[0] == 1, "This implementation only supports batch_size=1 during inference"
+                
+                # Get all audio token positions in this sample
+                audio_start_token_positions = torch.where(input_ids[0] == audio_config.audio_patch_token)[0]
+                
+                # Initialize with original embeddings
+                current_embeds = inputs_embeds[0]  # [seq_len, embed_dim]
+                current_ids = input_ids[0]         # [seq_len]
+                
+                # Process each audio token position sequentially
+                position_shift = 0  # Track position changes due to expansions
+                
+                # Ensure shrink_features is properly formatted
+                if isinstance(shrink_features[0], list):
+                    # If it's a list of lists (batch_size=1 but multiple turns), flatten it
+                    shrink_features = [item for sublist in shrink_features for item in sublist]
+                
+                for pos_idx, audio_pos in enumerate(audio_start_token_positions):
+                    adjusted_pos = audio_pos + position_shift
+                    
+                    # Get corresponding shrink feature (ensure it's a tensor)
+                    shrink_feature = shrink_features[pos_idx]
+                    if isinstance(shrink_feature, list):
+                        shrink_feature = torch.stack(shrink_feature, dim=0)
+                    
+                    v = shrink_feature.shape[0]  # Now this should work
+                    # print('len: ', v)
+                    
+                    # Expand the input ids and embeddings
+                    current_ids = torch.cat([
+                        current_ids[:adjusted_pos],
+                        current_ids[adjusted_pos:adjusted_pos+1].repeat(v),
+                        current_ids[adjusted_pos + num_patches:]
+                    ], dim=0)
+                    
+                    current_embeds = torch.cat([
+                        current_embeds[:adjusted_pos],
+                        shrink_feature,
+                        current_embeds[adjusted_pos + num_patches:]
+                    ], dim=0)
+                    
+                    # Update position shift for next iteration
+                    position_shift += (v - num_patches)
+                
+                # Update the tensors (unsqueeze to restore batch dim)
+                input_ids = current_ids.unsqueeze(0)          # [1, new_seq_len]
+                inputs_embeds = current_embeds.unsqueeze(0)   # [1, new_seq_len, embed_dim]
+                attention_mask = input_ids.ne(audio_config.llm_pad_token_id)
+                
+                # Update inference state tracking
+                if not hasattr(self, 'mask_tensor'):
+                    # Initialize with current attention mask
+                    self.mask_tensor = attention_mask.clone()
+                    self.length = attention_mask.shape[1]
+                else:
+                    # Ensure mask tensor is on correct device
+                    self.mask_tensor = self.mask_tensor.to(attention_mask.device)
+                    
+                    # Expand mask tensor if needed
+                    if self.mask_tensor.shape[1] < attention_mask.shape[1]:
+                        new_mask = torch.zeros(1, attention_mask.shape[1], 
+                                            dtype=torch.bool,
+                                            device=attention_mask.device)
+                        new_mask[0, :self.mask_tensor.shape[1]] = self.mask_tensor
+                        self.mask_tensor = new_mask
+                    
+                    # Update mask tensor
+                    self.mask_tensor[0, :attention_mask.shape[1]] = attention_mask[0]
+                    self.length = attention_mask.shape[1]
+        
+        attention_mask=self.mask_tensor[:,:self.length]
+        self.length+=1
+        
+        return_state=super(ACLlamaModel, self).forward(
+            input_ids=None, attention_mask=attention_mask, past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds, use_cache=use_cache,
+            output_attentions=output_attentions, output_hidden_states=output_hidden_states,
+            return_dict=return_dict
+        )
+        if self.training and audios is not None:
+            return_state["audio_features"] =  predict_logits
+            return_state["label_shift"] = batch_label_shift
+            return_state["label_extend"] = label_extend
+            return_state["speech_pos"] = batch_speech_pos
+        #return_state = {"audio_features":predict_logits}
+        return return_state 
+
+
+class ACLlamaForCausalLM(LlamaForCausalLM):
+    config_class = ACLlamaConfig
+
+    def __init__(self, config):
+        super(LlamaForCausalLM, self).__init__(config)
+        self.model = ACLlamaModel(config)
+
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # t2u by kkq
+        if hasattr(config, "unit_output"):
+            self.unit_output = config.unit_output
+            self.unit_translator = T2ULlamaForCausalLM(config, self.lm_head.weight)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_model(self):
+        return self.model
+    
+    def get_unit_translator(self):
+        return self.unit_translator
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        t2u_input_ids: Optional[torch.LongTensor] = None,
+        t2u_labels: Optional[torch.LongTensor] = None,
+        t2u_attention_mask: Optional[torch.Tensor] = None,
+        unit_targets: Optional[torch.Tensor] = None,
+        sub_lengths: Optional[torch.Tensor] = None,
+        asr_targets: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        audios: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        do_task: str = None,
+        assistant_after_audio_shifts: Optional[torch.Tensor] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        
+        # t2u by kkq
+        # pretrain(t2u only) finetune(s2t&e2u)
+        do_task = do_task if do_task != None  else getattr(self, 'unit_output', None)
+
+        outputs = None
+        hidden_states = None
+        new_shift_labels = None
+        if do_task != "pretrain":
+            outputs = self.model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                audios=audios
+            )
+        
+
+            hidden_states = outputs[0]
+            logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None and do_task != "pretrain" and do_task != "finetune_kd":
+            if asr_targets is not None:
+                asr_logits = outputs["audio_features"]
+                asr_targets = asr_targets.view(asr_targets.shape[0] * asr_targets.shape[1], asr_targets.shape[2])
+                mask_asr_targets = (asr_targets != IGNORE_TOKEN_ID)
+                target_lengths = mask_asr_targets.sum(1)
+                input_lengths = torch.full(size=(asr_logits.shape[0],), fill_value=asr_logits.shape[1], dtype=torch.long)
+
+                loss_ctc = CTCLoss()
+
+                log_probs = F.log_softmax(asr_logits, dim=-1).transpose(0, 1)
+                #print(asr_targets.shape)
+                #print(input_lengths, target_lengths)
+
+                with torch.backends.cudnn.flags(enabled=False):
+                    loss_asr = F.ctc_loss(
+                        log_probs,
+                        asr_targets,
+                        input_lengths,
+                        target_lengths,
+                        blank=self.model.audio_tower[0].config.audio_patch_token,
+                        reduction='mean',
+                        zero_infinity=True,
+                    )
+            else:
+                loss_asr=0
+
+            shift_labels = labels
+            if "label_shift" in outputs.keys() and len(outputs["label_shift"]) >0:
+                if outputs["label_extend"] != -1:
+                    new_shift_labels = torch.full(size=(shift_labels.shape[0], outputs["label_extend"]+shift_labels.shape[1]), fill_value=IGNORE_TOKEN_ID, dtype=torch.long).to(shift_labels.device)
+                    for batch in range(len(outputs["label_shift"])):
+                        it_lable_shift = outputs["label_shift"][batch]
+                        it_speech_pos = outputs["speech_pos"][batch]
+                        prefix = 0
+                        for i in range(len(it_lable_shift)):
+                            if i == len(it_lable_shift) - 1:
+                                length = shift_labels.shape[1] - it_speech_pos[i] #len(shift_labels[batch]) - it_speech_pos[i]
+                            else:
+                                length =  it_speech_pos[i + 1] - it_speech_pos[i]
+                            prefix += it_lable_shift[i]
+                            new_shift_labels[batch][it_speech_pos[i] + prefix: it_speech_pos[i] + length + prefix]= shift_labels[batch][it_speech_pos[i]:it_speech_pos[i]+length]
+                    shift_labels = new_shift_labels
+                else:
+                    raise NotImplementedError
+
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = shift_labels[..., 1:].contiguous()
+            #print(shift_labels[:,:50])
+
+            #print(shift_labels[:,:150])
+            loss_fct = CrossEntropyLoss()
+            # Flatten the tokens
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+                    
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+            loss = loss + 0.3 * loss_asr 
+
+        t2u_output = None
+        if do_task != None and do_task != "skip":
+            if do_task == "finetune_kd":
+                text_start_index = []
+                for batch in range(len(outputs["label_shift"])):
+                    text_start_index.append(outputs["speech_pos"][batch][0] + outputs["label_shift"][batch][0]+assistant_after_audio_shifts[batch])
+
+                t2u_embeds_output = self.unit_translator.insert_text_embedding(
+                    input_ids=t2u_input_ids,
+                    attention_mask=t2u_attention_mask,
+                    inputs_embeds=None,
+                    labels=t2u_labels,
+                    text_labels=labels,
+                    shift_text_labels=new_shift_labels,
+                    shift_text_hidden_states=hidden_states,
+                    unit_targets=unit_targets,
+                    sub_lengths=sub_lengths,
+                    text_start_index=text_start_index,
+                    do_task=do_task,
+                )
+
+                vae_loss, t2u_inputs_embeds, unit_targets, t2u_attention_mask = t2u_embeds_output
+
+                t2u_output = self.unit_translator(
+                    input_ids=None,
+                    attention_mask=t2u_attention_mask,
+                    past_key_values=past_key_values,
+                    inputs_embeds=t2u_inputs_embeds,
+                    use_cache=use_cache,
+                    labels=unit_targets,
+                    output_attentions=output_attentions,
+                    output_hidden_states=output_hidden_states,
+                    return_dict=return_dict,
+                )
+            else:
+                t2u_embeds_output = self.unit_translator.insert_text_embedding(
+                    input_ids=t2u_input_ids,
+                    attention_mask=t2u_attention_mask,
+                    inputs_embeds=None,
+                    labels=t2u_labels,
+                    text_labels=labels,
+                    shift_text_labels=new_shift_labels,
+                    shift_text_hidden_states=hidden_states,
+                    do_task=do_task,
+                )
+                vae_loss, t2u_inputs_embeds = t2u_embeds_output
+
+                t2u_output = self.unit_translator(
+                    input_ids=None,
+                    attention_mask=t2u_attention_mask,
+                    past_key_values=past_key_values,
+                    inputs_embeds=t2u_inputs_embeds,
+                    use_cache=use_cache,
+                    labels=t2u_labels,
+                    output_attentions=output_attentions,
+                    output_hidden_states=output_hidden_states,
+                    return_dict=return_dict,
+                )
+            t2u_loss = t2u_output[0]
+            # print(do_task, t2u_loss, vae_loss)
+            if vae_loss != None:
+                target_scale = t2u_loss.item() * 0.2
+                vae_loss_weight = target_scale / vae_loss.item() if vae_loss > target_scale else 1.0
+                t2u_loss = t2u_loss + vae_loss_weight * vae_loss
+                #print(vae_loss)
+
+            if loss != None:            # S2T + T2U loss
+                # ignore LLM loss
+                # t2u_output["loss"] = t2u_loss
+                # return t2u_output
+                # original version
+                assert do_task in ["finetune"]
+                if loss.item() < 1.0:   # 1.7
+                    loss = 0.2 * loss + t2u_loss * 2.0
+                else:
+                    loss = loss + t2u_loss
+            else:
+                assert do_task in ["pretrain", "finetune_kd"]
+                t2u_output["loss"] = t2u_loss
+                return t2u_output
+        
+        #return CausalLMOutputWithPast(
+        #    loss=loss,
+        #    logits=outputs["audio_features"],
+        #)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        **kwargs,
+    ):
+        # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
+        # Exception 1: when passing input_embeds, input_ids may be missing entries
+        # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
+        if past_key_values is not None:
+            if inputs_embeds is not None:  # Exception 1
+                input_ids = input_ids[:, -cache_position.shape[0] :]
+            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
+                input_ids = input_ids[:, cache_position]
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and cache_position[0] == 0:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids.contiguous()}  # `contiguous()` needed for compilation use cases
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "use_cache": use_cache,
+                "attention_mask": attention_mask,
+            }
+        )
+        model_inputs.update({"audios": kwargs["audios"]} if "audios" in kwargs.keys() else {})
+        model_inputs.update({"do_task": kwargs["do_task"]} if "do_task" in kwargs.keys() else {})
+        model_inputs.update({"return_dict": kwargs["return_dict_in_generate"]} if "return_dict_in_generate" in kwargs.keys() else {})
+        return model_inputs
+
+AutoConfig.register("ACLlama", ACLlamaConfig)
+AutoModelForCausalLM.register(ACLlamaConfig, ACLlamaForCausalLM)

EchoX-Vocoder/checkpoint_last.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:abb7b49cc59bbf058719cdae2252069dce1bb7b73362a0ec5273670ed7a6d4cc
+size 389348172

EchoX-Vocoder/config.json

@@ -0,0 +1,53 @@
+{
+    "input_wavs_dir": "/private/home/adampolyak/datasets/LJ/LJSpeech-1.1/wavs_16khz_padded",
+    "input_training_file": "/large_experiments/ust/annl/datasets/tts/LJSpeech/filelist/mhubert_vp_en_es_fr_it3_400k/lj_train_layer11_hubert1000_filelist.txt",
+    "input_validation_file": "/large_experiments/ust/annl/datasets/tts/LJSpeech/filelist/mhubert_vp_en_es_fr_it3_400k/lj_dev_layer11_hubert1000_filelist.txt",
+
+    "resblock": "1",
+    "num_gpus": 0,
+    "batch_size": 16,
+    "learning_rate": 0.0002,
+    "adam_b1": 0.8,
+    "adam_b2": 0.99,
+    "lr_decay": 0.999,
+    "seed": 1234,
+
+    "upsample_rates": [5,4,4,2,2],
+    "upsample_kernel_sizes": [11,8,8,4,4],
+    "upsample_initial_channel": 512,
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+    "num_embeddings": 1000,
+    "embedding_dim": 128,
+    "model_in_dim": 128,
+
+    "segment_size": 8960,
+    "code_hop_size": 320,
+    "f0": false,
+    "num_mels": 80,
+    "num_freq": 1025,
+    "n_fft": 1024,
+    "hop_size": 256,
+    "win_size": 1024,
+
+    "dur_prediction_weight": 1.0,
+    "dur_predictor_params": {
+        "encoder_embed_dim": 128,
+        "var_pred_hidden_dim": 128,
+        "var_pred_kernel_size": 3,
+        "var_pred_dropout": 0.5
+    },
+
+    "sampling_rate": 16000,
+
+    "fmin": 0,
+    "fmax": 8000,
+    "fmax_for_loss": null,
+
+    "num_workers": 4,
+
+    "dist_config": {
+        "dist_backend": "nccl",
+        "dist_url": "env://"
+    }
+}

EchoX-Vocoder/g_00500000

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d1f7188b95b06304bc05e524fddf93c7fe682fdd93acff022685663a5e26b97
+size 54051213

EchoX-Vocoder/spm_1k.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d95d5585291329feaf35d3cb39fe5181e4987549097a9daa36f468dab9e82556
+size 254653

Echox_copy_stream.py

@@ -0,0 +1,426 @@
+import sys
+from ACLlama_el_s2s import ACLlamaForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoConfig, WhisperProcessor
+from peft import PeftModel, PeftConfig
+import json
+from tqdm import tqdm
+import torch
+import re
+import os
+torch.backends.cudnn.benchmark = False
+import librosa
+from text_to_speech import *
+import torch.nn.functional as F
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+from transformers import logging as hf_logging
+hf_logging.set_verbosity_error()
+from huggingface_hub import hf_hub_download
+from typing import Dict, Optional, List
+import tempfile
+import select
+from copy import deepcopy
+from typing import Generator, Tuple
+
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+
+def load_model(args, device):
+    quantization_config = None
+    hf_token = os.getenv("HF_TOKEN")
+    
+    # load based model
+    model = ACLlamaForCausalLM.from_pretrained(
+        args.base_model_path,
+        device_map=None,
+        torch_dtype=torch.float16,
+        quantization_config=quantization_config,
+        token=hf_token,
+    ).eval().to(device)
+    for module in model.model.audio_tower:
+        module = module.to(device)
+
+    if args.peft_model_id:
+        lora_config = PeftConfig.from_pretrained(args.peft_model_id)
+        torch.cuda.empty_cache()
+        model = PeftModel.from_pretrained(model, args.peft_model_id, config=lora_config).to(
+            dtype=torch.float16, device=device
+        )
+        model = model.merge_and_unload()
+
+    model.eval()
+
+    # load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(args.base_model_path, token=hf_token)
+
+    audio_config = model.get_model().audio_tower[0].config
+    audio_config.audio_patch_token = tokenizer.get_vocab()["<audio_patch>"]
+    audio_config.llm_pad_token_id = tokenizer.pad_token_id
+    audio_config.audio_patch_size = args.audio_token_len
+
+
+    # whisper processor
+    audio_processor = WhisperProcessor.from_pretrained(args.audio_tower, torch_dtype=torch.float16)
+
+    # t2u
+    unit_translator = model.get_unit_translator().eval()
+    return model, audio_processor, tokenizer, unit_translator
+
+def load_speech_model(device):
+    vocoder = "./EchoX-Vocoder/g_00500000"
+    vocoder_cfg = "./EchoX-Vocoder/config.json"
+    voc_cfg = get_vocoder_config(vocoder, vocoder_cfg)
+    vocoder = load_units_vocoder(voc_cfg, device)
+    return vocoder, voc_cfg
+
+# def load_speech_model(device):
+#     hf_token = os.getenv("HF_TOKEN")
+    
+#     vocoder_repo_id = "FreedomIntelligence/EchoX-Vocoder"
+    
+#     cache_path = './hf_cache'
+#     vocoder_path = hf_hub_download(repo_id=vocoder_repo_id, filename="g_00500000", token=hf_token, cache_dir=cache_path)
+#     vocoder_cfg_path = hf_hub_download(repo_id=vocoder_repo_id, filename="config.json", token=hf_token, cache_dir=cache_path)
+
+#     voc_cfg = get_vocoder_config(vocoder_path, vocoder_cfg_path)
+#     vocoder = load_units_vocoder(voc_cfg, device)
+#     return vocoder, voc_cfg
+
+class EchoxAssistant():
+    def __init__(self):
+        class BasicSetting:
+            def __init__(self):
+                self.device = "cuda:0"
+                self.sampling_rate = 16000
+                self.audio_token_len = 1  # 1500 = 300 token x 5 compress
+                self.stop = "</s>"
+                self.base_model_path = "FreedomIntelligence/EchoX-8B"
+                self.peft_model_id = None
+                self.audio_tower = "openai/whisper-large-v3"
+        self.args = BasicSetting()
+        self.device = "cuda"
+        self.vocoder, self.voc_cfg= load_speech_model(self.device)
+        self.model, self.audio_processor, self.tokenizer, self.unit_translator = load_model(self.args, self.device)
+        self.audio_executor = ThreadPoolExecutor(max_workers=2)
+        # self.specAug = SpecAugmentTransform()
+        # special_token
+        DEFAULT_AUDIO_PATCH_TOKEN = "<audio_patch>"
+        audio_placeholder = DEFAULT_AUDIO_PATCH_TOKEN * self.args.audio_token_len
+        audio_placeholder = "\n"+audio_placeholder
+        self.audio_placeholder_ids = self.tokenizer(audio_placeholder).input_ids
+
+        self.begin_of_text_id = self.tokenizer.get_vocab()["<|begin_of_text|>"]
+        self.start_header_id = self.tokenizer.get_vocab()["<|start_header_id|>"]
+        self.end_header_id = self.tokenizer.get_vocab()["<|end_header_id|>"]
+        self.eot_id = self.tokenizer.get_vocab()["<|eot_id|>"]
+        self.nl_tokens = self.tokenizer('\n').input_ids
+        self._system = self.tokenizer('system').input_ids
+        self._user = self.tokenizer('user').input_ids
+        self._assistant = self.tokenizer('assistant').input_ids
+        self._speaker = self.tokenizer('speaker').input_ids
+
+        self.max_len = 1024
+        self.unit_max_len = 2048
+        self.system_message = "You are a helpful language and speech assistant. You are able to understand the speech content that the user provides, and assist the user with a variety of tasks using natural language."
+    
+    def _generate_audio_segment(self, segment_hidden_states):
+        try:
+            audio_units = self._generate_audio_units_from_hidden_states(segment_hidden_states)
+            if audio_units:
+                audio_float32 = self.generate_with_speech_model([list(map(int, audio_units.split(" ")))])
+                audio_int16 = (audio_float32 * 32767).astype(np.int16)
+
+                print(f"Generated audio segment in background: {len(audio_units.split())} units")
+                return (16000, audio_int16)
+            return None
+        except Exception as e:
+            print(f"Background audio generation error: {e}")
+            return None
+
+    def gen_model_inputs(
+            self,
+            sources,
+            tokenizer,
+            max_len,
+            system_message,
+            audio_placeholder_ids, begin_of_text_id, start_header_id, end_header_id, eot_id, nl_tokens, _system, _user, _assistant,
+        ) -> dict:
+        # max_len 512
+
+        # Apply prompt templates
+        input_ids, audio_paths = [], []
+        audio_path = []
+        
+        for source in sources:
+            input_id = []
+            system = [begin_of_text_id] + [start_header_id] + _system + [end_header_id] + nl_tokens + tokenizer(system_message).input_ids + [eot_id]
+            input_id += system
+
+            for j, item in enumerate(source["conversations"]):
+                role = item["from"]
+                value = item["value"]
+                _audio_path = None
+
+                if role == 'user':
+                    if "audio" in item.keys():
+                        _input_id = [start_header_id] + _user + [end_header_id] + audio_placeholder_ids + tokenizer(value).input_ids + [eot_id]
+                        _audio_path = item["audio"]
+                    else:
+                        _input_id = [start_header_id] + _user + [end_header_id] + tokenizer(value).input_ids + [eot_id]
+                
+                elif role == 'assistant':
+                    _input_id = [start_header_id] + _assistant + [end_header_id] + nl_tokens + tokenizer(value).input_ids + [eot_id]
+                
+                else:
+                    raise NotImplementedError
+                input_id += _input_id
+
+                if _audio_path:
+                    audio_path.append(_audio_path)
+            assistant_input_id = [start_header_id] + _assistant + [end_header_id] + nl_tokens
+            input_id += assistant_input_id
+
+            audio_num = int(input_id.count(audio_placeholder_ids[-1]) / self.args.audio_token_len)
+            assert len(audio_path) == audio_num
+            if len(input_id) >= max_len:
+                print(f"[WARNING] Your Input Length More Than {max_len}")
+            input_ids.append(input_id[:max_len])
+            audio_paths.append(audio_path)
+        input_ids = torch.tensor(input_ids, dtype=torch.int)
+        return dict(
+            input_ids=input_ids,
+            audio_paths=audio_paths,
+            attention_mask=input_ids.ne(tokenizer.pad_token_id),
+        )
+
+    def get_unit_result(self, ret):
+        # print(ret)
+        self.unit_translator.generation_config.pad_token_id = self.tokenizer.eos_token_id
+        input_ids = ret["input_ids"]
+        ret["input_ids"] = None
+        model_outputs = self.unit_translator.generate(
+            **ret,
+            max_new_tokens=2048,
+            eos_token_id=self.tokenizer.eos_token_id,
+        )
+        # print(model_outputs, model_outputs.shape)
+        output_ids = model_outputs
+        unit_output = self.tokenizer.batch_decode(output_ids)[0]
+        if "▁" in unit_output:
+            unit_output = ''.join(re.findall(r"<\|unit_(.*?)\|>", unit_output))
+            
+        units = re.findall(r'\d+', unit_output)
+
+        #TODO grid of unk unit
+        new_units = []
+        for unit in units:
+            if int(unit) < 1000:
+                new_units.append(unit)
+
+        units = ' '.join(new_units)
+        return units
+    
+
+    def _inference(
+        self,
+        prompt,
+        **kwargs,
+    ):
+        audio_paths = []
+        response = []
+        for item in prompt:
+            for conv in item["conversations"]:
+                if "audio" in conv:
+                    audio_paths.append(conv["audio"])
+
+        model_inputs = self.gen_model_inputs(
+            prompt,
+            self.tokenizer,
+            self.max_len,
+            self.system_message,
+            self.audio_placeholder_ids, self.begin_of_text_id, self.start_header_id, self.end_header_id, self.eot_id, self.nl_tokens, self._system, self._user, self._assistant)
+        
+        audio_list = []
+        if audio_paths and audio_paths[0] is not None:
+            for audio_path in audio_paths:
+                # print("read audio file name: ", audio_path)
+                audio, _ = librosa.load(audio_path, sr=self.args.sampling_rate)
+                audio_feat = self.audio_processor(audio, sampling_rate=self.args.sampling_rate, return_tensors="pt").input_features
+                audio_list.append(audio_feat)
+            audio_feats = torch.stack(audio_list, dim=0)
+            audio_feats = audio_feats.to(dtype=torch.float16).to(self.device)
+        
+        if not audio_list:
+            ret = dict(
+                    input_ids=model_inputs["input_ids"].to(self.device),
+                    attention_mask=model_inputs["attention_mask"].to(self.device),
+                )
+        else:
+            ret = dict(
+                    input_ids=model_inputs["input_ids"].to(self.device),
+                    attention_mask=model_inputs["attention_mask"].to(self.device),
+                    audios=audio_feats,
+                )    
+        
+        self.model.generation_config.pad_token_id = self.tokenizer.eos_token_id
+        #print(self.model.lm_head.weight.shape)
+
+        dot_input_ids = self.tokenizer(".", return_tensors="pt").input_ids.to(self.device) # 形状: (1, 2), 值: [[128000, 13]]
+        period_token_id = dot_input_ids[0, -1]
+        period_lm_head_embedding = self.model.lm_head.weight[period_token_id]
+
+        input_ids = ret["input_ids"]
+        attention_mask = ret["attention_mask"]
+        input_token_len = input_ids.shape[1]
+
+        max_new_tokens = kwargs.get('max_new_tokens', 512)
+        temperature = kwargs.get('temperature', 0.2)
+        top_p = kwargs.get('top_p', 0.9)
+        do_sample = kwargs.get('do_sample', True)
+        
+        current_text = ""
+        accumulated_hidden_states = []
+        accumulated_tokens = []
+        similarity_scores = []
+        segment_start_idx = 0
+        
+        current_input_ids = input_ids
+        current_attention_mask = attention_mask
+        past_key_values = None
+        
+        audio_futures = []
+        segmentation_latency = 5
+
+        with torch.no_grad():
+            for step in range(max_new_tokens):
+                while audio_futures and audio_futures[0].done():
+                    completed_future = audio_futures.pop(0)
+                    audio_data = completed_future.result()
+                    if audio_data:
+                        yield None, audio_data
+                
+                if current_input_ids is None:
+                    break
+
+                model_kwargs = {
+                    "input_ids": current_input_ids,
+                    "attention_mask": current_attention_mask,
+                    "past_key_values": past_key_values,
+                    "use_cache": True,
+                    "output_hidden_states": True,
+                    "do_task": "skip"
+                }
+                
+                if step == 0 and "audios" in ret:
+                    model_kwargs["audios"] = ret["audios"]
+                
+                outputs = self.model(**model_kwargs)
+                
+                logits = outputs.logits
+                hidden_states = outputs.hidden_states[-1]
+                past_key_values = outputs.past_key_values
+                
+                next_token_logits = logits[:, -1, :]  # [batch_size, vocab_size]
+                
+                if do_sample:
+                    next_token_logits = next_token_logits / temperature
+                    
+                    sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+                    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                    
+                    sorted_indices_to_remove = cumulative_probs > top_p
+                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                    sorted_indices_to_remove[..., 0] = 0
+                    
+                    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+                    next_token_logits[indices_to_remove] = float('-inf')
+                    
+                    probs = F.softmax(next_token_logits, dim=-1)
+                    next_token = torch.multinomial(probs, num_samples=1)
+                else:
+                    next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+                
+                if next_token.item() == self.tokenizer.eos_token_id:
+                    current_input_ids = None
+                    continue
+
+                accumulated_tokens.append(next_token.item())
+                last_hidden_state = hidden_states[0, -1]  # [hidden_dim]
+                accumulated_hidden_states.append(last_hidden_state)
+                
+                similarity = F.cosine_similarity(last_hidden_state, period_lm_head_embedding, dim=0).item()
+                similarity_scores.append(similarity)
+                
+                token_text = self.tokenizer.decode([next_token.item()], skip_special_tokens=True)
+                current_text += token_text
+                
+                yield current_text, None
+                
+                current_idx = len(similarity_scores) - 1
+                check_idx = current_idx - segmentation_latency
+                if check_idx >= 0:
+                    similarity_at_check = similarity_scores[check_idx]
+                    is_peak = self._is_local_maximum(similarity_scores, check_idx, window=segmentation_latency)
+                    should_segment = (is_peak and
+                        check_idx - segment_start_idx >= 50) or (
+                        is_peak and
+                        similarity_at_check > 0.1 and
+                        check_idx - segment_start_idx >= 20
+                    )
+                    
+                    if should_segment:
+                        segment_end_idx = check_idx + 1
+                        print(f"Segmenting at step {segment_end_idx-1}, similarity={similarity_at_check:.4f}. Submitting to background audio generation.")
+                        
+                        segment_hidden_states = torch.stack(
+                            accumulated_hidden_states[segment_start_idx:segment_end_idx], dim=0
+                        ).unsqueeze(0)
+                        
+                        future = self.audio_executor.submit(self._generate_audio_segment, segment_hidden_states)
+                        audio_futures.append(future)
+                        
+                        segment_start_idx = segment_end_idx
+                
+                current_input_ids = next_token
+                current_attention_mask = torch.ones_like(next_token)
+            
+            if segment_start_idx < len(accumulated_hidden_states):
+                print(f"Processing final segment from {segment_start_idx} to {len(accumulated_hidden_states)}")
+                segment_hidden_states = torch.stack(
+                    accumulated_hidden_states[segment_start_idx:], dim=0
+                ).unsqueeze(0)
+                future = self.audio_executor.submit(self._generate_audio_segment, segment_hidden_states)
+                audio_futures.append(future)
+
+            for future in audio_futures:
+                audio_data = future.result()
+                if audio_data:
+                    yield None, audio_data
+
+    def _is_local_maximum(self, scores, idx, window=5):
+        start = max(0, idx - window)
+        end = min(len(scores), idx + window + 1)
+        local_scores = scores[start:end]
+        return scores[idx] == max(local_scores)
+
+    def _generate_audio_units_from_hidden_states(self, hidden_states):
+        try:
+            _, adapted_inputs_embeds = self.unit_translator.insert_text_embedding(
+                inputs_embeds=hidden_states,
+                do_task="skip",
+            )
+            
+            attention_mask = torch.ones(adapted_inputs_embeds.shape[:2]).to(self.device)
+            ret = dict(
+                input_ids=None,
+                inputs_embeds=adapted_inputs_embeds,
+                attention_mask=attention_mask,
+            )
+            
+            return self.get_unit_result(ret)
+        except Exception as e:
+            print(f"Error generating audio units: {e}")
+            return None
+        
+    def generate_with_speech_model(self, units):
+        wav = gen_wav(self.vocoder, self.voc_cfg, units, self.device)
+        return wav

T2ULlama_CR_online.py

@@ -0,0 +1,412 @@
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import CrossEntropyLoss, CTCLoss
+import transformers
+
+from transformers import AutoConfig, AutoModelForCausalLM, \
+                         LlamaConfig, LlamaModel, LlamaForCausalLM
+from transformers.trainer_pt_utils import LabelSmoother
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers import (
+    WhisperProcessor,
+    WhisperModel,
+)
+
+IGNORE_TOKEN_ID = LabelSmoother.ignore_index
+
+
+def padding_tensor(tensor, length, dim=0, pad=False):
+
+    if length == 0:
+        return tensor
+    
+    assert length > 0, f"Wrong padding length: {length}"
+
+    shape = list(tensor.shape)
+    assert dim < len(shape), f"dim {dim} out of shape {shape}"
+    shape[dim] = length
+    padding_tensor = torch.cat(
+        (
+            tensor,
+            torch.full(tuple(shape), pad, dtype=tensor.dtype, device=tensor.device)
+        ), 
+        dim=dim
+    )
+    return padding_tensor
+
+
+class T2ULlamaConfig(LlamaConfig):
+    model_type = "T2ULlama"
+    
+class T2ULlamaForCausalLM(LlamaForCausalLM):
+    config_class = T2ULlamaConfig
+
+    def __init__(self, config, embedding_weight=None):
+        
+        self.current_step = 0
+        self.log = {}
+
+        super(LlamaForCausalLM, self).__init__(config)
+        self.config = config
+        self.training_stage = config.unit_output
+        self.pad_token_id = 128009
+
+        llama_config = T2ULlamaConfig(**config.to_dict(), 
+                                        batch_first=True, 
+                                        norm_first=True
+                                    )
+        llama_config.architectures = ["T2ULlamaForCausalLM"]
+        llama_config.pad_token_id = self.pad_token_id
+        llama_config.vocab_size += llama_config.unit_vocab_size
+        #######################################################
+        llama_config.unit_model = "medium"
+        llama_config.max_position_embeddings = 2048     # 1024 1536 2048       # origin 1024 reduced 512
+        #######################################################
+        if hasattr(llama_config, "unit_model"):
+            if llama_config.unit_model == "large":
+                llama_config.num_hidden_layers = 2
+                # llama_config.hidden_size = 4096
+                # llama_config.num_attention_heads = 32
+                # llama_config.intermediate_size = 14336
+                # llama_config.head_dim = llama_config.hidden_size // llama_config.num_attention_heads
+
+            elif llama_config.unit_model == "tiny":
+                llama_config.num_hidden_layers = 4
+                llama_config.hidden_size = 512
+                llama_config.num_attention_heads = 8
+                llama_config.intermediate_size = 2048
+                llama_config.head_dim = llama_config.hidden_size // llama_config.num_attention_heads
+            else:
+                llama_config.num_hidden_layers = 8
+                llama_config.hidden_size = 768
+                llama_config.num_attention_heads = 12
+                llama_config.num_key_value_heads = 12
+                llama_config.intermediate_size = 2048
+                llama_config.head_dim = llama_config.hidden_size // llama_config.num_attention_heads
+        else:
+            llama_config.num_hidden_layers = 6
+            llama_config.hidden_size = 512
+            llama_config.num_attention_heads = 8
+            llama_config.intermediate_size = 2048
+            llama_config.head_dim = llama_config.hidden_size // llama_config.num_attention_heads
+        # print(llama_config)
+        
+        self.model = LlamaModel(llama_config)
+        # share embedding 0501 by kkq
+        self.model.embed_tokens = nn.Embedding(num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, padding_idx=self.pad_token_id)   # redefine
+        self.unit_embedding = nn.Linear(config.hidden_size, llama_config.unit_vocab_size, bias=False) 
+        self.adapter = nn.Linear(config.hidden_size, llama_config.hidden_size, bias = True) 
+        self.lm_head = nn.Linear(llama_config.hidden_size, llama_config.vocab_size, bias=False)
+
+        if self.training_stage == "pretrain":
+            pass
+        elif self.training_stage == "finetune" or self.training_stage == "finetune_kd" or self.training_stage == "finetune_kd_online":
+            self.aligner_MLP = nn.Sequential(
+                nn.Linear(config.hidden_size, config.intermediate_size),
+                nn.ReLU(),
+                nn.Dropout(0.1),
+                nn.Linear(config.intermediate_size, config.hidden_size),
+            )
+            torch.nn.init.ones_(self.aligner_MLP[0].weight)
+            torch.nn.init.zeros_(self.aligner_MLP[0].bias)
+            torch.nn.init.ones_(self.aligner_MLP[3].weight)
+            torch.nn.init.zeros_(self.aligner_MLP[3].bias)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_model(self):
+        return self.model
+
+    def insert_text_embedding(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        text_labels: Optional[torch.LongTensor] = None,
+        shift_text_labels: Optional[torch.LongTensor] = None,
+        shift_text_hidden_states: Optional[torch.FloatTensor] = None,
+        unit_targets: Optional[torch.LongTensor] = None,
+        sub_lengths: Optional[torch.LongTensor] = None,
+        text_start_index: Optional[torch.LongTensor] = None,
+        do_task: str = None,
+        **kwargs: dict,
+    ):  
+
+        if inputs_embeds == None:
+            # share embedding 0501 by kkq
+            embed_tokens_weight = torch.cat(
+                [
+                    self.model.embed_tokens.weight.detach(), self.unit_embedding.weight
+                ],
+                dim = 0,
+            )
+            # print(embed_tokens_weight, embed_tokens_weight.shape)
+            inputs_embeds = F.embedding(input_ids, embed_tokens_weight, padding_idx=self.pad_token_id)
+        
+        emb_loss = None
+        if do_task == "pretrain":
+            if self.training:
+                if hasattr(self, "embedding_dropout"):
+                    emb_origin_mask = text_labels != -100
+                    origin_padding_length = labels.shape[-1] - emb_origin_mask.shape[-1]
+                    extend_emb_origin_mask = padding_tensor(emb_origin_mask, origin_padding_length, 1, False)
+                    extend_emb_origin_mask = ~extend_emb_origin_mask.unsqueeze(-1).expand_as(inputs_embeds)
+
+                    # Π-Model + noise
+                    log_var = self.perturb(inputs_embeds)
+                    perturbed_inputs_embeds_2 = inputs_embeds + torch.randn_like(inputs_embeds) * (torch.exp(0.5 * log_var) + 1e-6)
+                    # Π-Model + dropout
+                    perturbed_inputs_embeds_1 = self.embedding_dropout(inputs_embeds)
+                    perturbed_inputs_embeds_2 = self.embedding_dropout(perturbed_inputs_embeds_2)
+                    perturbed_inputs_embeds_1 = torch.where(extend_emb_origin_mask, inputs_embeds, perturbed_inputs_embeds_1)
+                    perturbed_inputs_embeds_2 = torch.where(extend_emb_origin_mask, inputs_embeds, perturbed_inputs_embeds_2)
+
+                    inputs_embeds = torch.cat(
+                        (perturbed_inputs_embeds_1, perturbed_inputs_embeds_2),
+                        dim=0,
+                    )
+
+                    kl_loss = -0.5 * (1 + log_var - log_var.exp()).mean(dim=-1).sum(dim=-1).mean()
+                    contrastive_loss = (1 - F.cosine_similarity(perturbed_inputs_embeds_1, perturbed_inputs_embeds_2, dim=-1)).sum(dim=-1).mean()
+                    emb_loss = kl_loss + contrastive_loss
+
+                    if kl_loss.device == torch.device("cuda:0"):
+                        self.log["kl_loss"] = kl_loss.item()
+                        self.log["std"] = torch.exp(0.5 * log_var).mean().item()
+                        self.log["contrastive_loss"] = contrastive_loss.item()
+
+            pass
+        elif do_task == "finetune":
+            inputs_embeds = inputs_embeds.detach()
+            inputs_embeds_refer = inputs_embeds.clone().detach()
+            shift_text_hidden_states = self.aligner_MLP(shift_text_hidden_states)
+            emb_origin_mask = text_labels != -100    # get output text pos
+            emb_shift_mask = shift_text_labels != -100
+
+            origin_padding_length = labels.shape[-1] - emb_origin_mask.shape[-1]
+            shift_padding_length = labels.shape[-1] - emb_shift_mask.shape[-1]
+            
+            extend_emb_origin_mask = padding_tensor(emb_origin_mask, origin_padding_length, 1, False)
+            extend_emb_shift_mask = padding_tensor(emb_shift_mask, shift_padding_length, 1, False)
+            extend_shift_text_hidden_states = padding_tensor(shift_text_hidden_states, shift_padding_length, 1, 1e-9)
+            # check
+            extend_text_labels = padding_tensor(text_labels, origin_padding_length, 1, -100)
+            extend_shift_text_labels = padding_tensor(shift_text_labels, shift_padding_length, 1, -100)
+            
+            assert torch.equal(
+                extend_text_labels[extend_emb_origin_mask], 
+                extend_shift_text_labels[extend_emb_shift_mask]
+            ), "{}\n{}\n{}\n{}".format(labels, extend_emb_origin_mask, extend_shift_text_labels, extend_emb_shift_mask)
+            
+            inputs_embeds[extend_emb_origin_mask.unsqueeze(-1).expand_as(inputs_embeds)] = \
+                extend_shift_text_hidden_states[extend_emb_shift_mask.unsqueeze(-1).expand_as(extend_shift_text_hidden_states)].to(dtype=inputs_embeds.dtype)
+            
+            if self.training:
+                contrastive_loss = (1 - F.cosine_similarity(inputs_embeds, inputs_embeds_refer, dim=-1)).sum(-1).mean()
+                emb_loss = contrastive_loss
+                if emb_loss.device == torch.device("cuda:0"):
+                    self.log["contrastive_loss"] = contrastive_loss.item()
+                pass
+        elif do_task == "finetune_kd" :
+            #inputs_embeds = inputs_embeds.detach()
+            #inputs_embeds_refer = inputs_embeds.clone().detach()
+            #print(text_labels)
+            #print(sub_lengths.sum())
+            emb_origin_mask = text_labels != -100
+
+            fetch_lables_list = [] 
+            for batch in range(emb_origin_mask.shape[0]):
+                fetch_lables_list.append(text_labels[batch][emb_origin_mask[batch]])
+            shift_text_hidden_states = self.aligner_MLP(shift_text_hidden_states)
+
+            #split the shift_text_hidden_states
+            #[128006, 128000, 78191, 128007, 128000, 198, 128000]
+            maxn_length = sub_lengths.max() + 8
+            pad_ids = torch.full(size=(sub_lengths.shape[0], sub_lengths.shape[1], maxn_length), fill_value=self.pad_token_id, dtype=torch.long).to(shift_text_hidden_states.device)
+
+            pad_text_ids = torch.full(size=(sub_lengths.shape[0], sub_lengths.shape[1], maxn_length), fill_value=self.pad_token_id, dtype=torch.long).to(shift_text_hidden_states.device)
+
+            atten_mask = pad_ids.ne(self.pad_token_id)
+            #target_mask_part1 = pad_ids.ne(self.pad_token_id)
+            shift_text_hidden_states_slice = F.embedding(pad_ids, embed_tokens_weight, padding_idx=self.pad_token_id)
+            
+            #print(shift_text_hidden_states_slice.shape,shift_text_hidden_states.shape)
+            for batch in range(sub_lengths.shape[0]):
+                cot=0
+                start_index = text_start_index[batch]
+                for index, sub_length in enumerate(sub_lengths[batch]):
+                    if sub_length==-1:
+                        break
+                    #print(shift_text_hidden_states_slice[batch][index][:sub_length].shape, shift_text_hidden_states[batch][cot:cot+sub_length].shape)
+                    eos_id = torch.IntTensor([128009]).to(inputs_embeds.device)
+                    eos = self.model.embed_tokens(eos_id)
+                    if index == 0:
+                        text_prefix_ids = torch.IntTensor([128006, 128000, 65576, 128007, 128000, 198]).to(inputs_embeds.device)
+                        preifx_embed = self.model.embed_tokens(text_prefix_ids)
+                        pad_text_ids[batch][index][:sub_length+7] = torch.cat([text_prefix_ids, fetch_lables_list[batch][cot:cot+sub_length], eos_id],dim=0)
+                        atten_mask[batch][index][:sub_length+7]=True
+                    else:
+                        text_prefix_ids = torch.IntTensor([128006, 128000, 65576, 128007, 128000, 198, 12800]).to(inputs_embeds.device)
+                        preifx_embed = self.model.embed_tokens(text_prefix_ids)
+                        pad_text_ids[batch][index][:sub_length+8] = torch.cat([text_prefix_ids, fetch_lables_list[batch][cot:cot+sub_length], eos_id], dim=0)
+                        atten_mask[batch][index][:sub_length+8]=True
+
+                    new_shift_text_hidden_states = torch.cat([preifx_embed, shift_text_hidden_states[batch][start_index+cot:start_index+cot+sub_length], eos], dim = 0)
+                    shift_text_hidden_states_slice[batch][index][:new_shift_text_hidden_states.shape[0]] = new_shift_text_hidden_states
+
+                    cot+=sub_length
+            shift_text_hidden_states_slice = shift_text_hidden_states_slice.reshape(shift_text_hidden_states_slice.shape[0]*shift_text_hidden_states_slice.shape[1],shift_text_hidden_states_slice.shape[2],shift_text_hidden_states_slice.shape[3])    
+
+
+            padding_unit_targets = unit_targets.clone()
+            padding_unit_targets = torch.where(padding_unit_targets == IGNORE_TOKEN_ID, self.pad_token_id, padding_unit_targets)
+            target_mask_part = padding_unit_targets.ne(self.pad_token_id)
+            atten_mask = torch.cat([atten_mask, target_mask_part], dim = -1)
+            atten_mask = atten_mask.reshape(atten_mask.shape[0]*atten_mask.shape[1],atten_mask.shape[2])
+
+            pad_text_ids = pad_text_ids.reshape(pad_text_ids.shape[0]*pad_text_ids.shape[1],pad_text_ids.shape[2])
+            shift_text_embeddings = F.embedding(pad_text_ids, embed_tokens_weight, padding_idx=self.pad_token_id)
+
+            unit_target_slice = F.embedding(padding_unit_targets, embed_tokens_weight, padding_idx=self.pad_token_id)
+            # unit_target_slice = F.embedding(unit_targets, embed_tokens_weight, padding_idx=self.pad_token_id)             
+            unit_target_slice = unit_target_slice.reshape(unit_target_slice.shape[0]*unit_target_slice.shape[1],unit_target_slice.shape[2],unit_target_slice.shape[3])
+
+            inputs_embeds = torch.cat([shift_text_hidden_states_slice, unit_target_slice], dim = 1)
+
+            ignore_ids = torch.full(size=(sub_lengths.shape[0], sub_lengths.shape[1], maxn_length), fill_value=IGNORE_TOKEN_ID, dtype=torch.long).to(shift_text_hidden_states.device)
+            unit_targets = torch.cat([ignore_ids,unit_targets],dim=-1) 
+            unit_targets = unit_targets.reshape(unit_targets.shape[0]*unit_targets.shape[1],unit_targets.shape[2])
+
+            if self.training:
+                #print(shift_text_hidden_states_slice.shape, shift_text_embeddings.shape)
+                contrastive_loss = (1 - F.cosine_similarity(shift_text_hidden_states_slice, shift_text_embeddings, dim=-1)).sum(-1).mean()
+                emb_loss = contrastive_loss
+                if emb_loss.device == torch.device("cuda:0"):
+                    self.log["contrastive_loss"] = contrastive_loss.item()
+
+        elif do_task == "finetune_kd_online":
+            shift_text_hidden_states = self.aligner_MLP(shift_text_hidden_states)
+            gold_inputs_embeds = inputs_embeds.clone()
+            for batch in range (inputs_embeds.shape[0]):
+                start_index = text_start_index[batch]
+                for slice_index in range (inputs_embeds.shape[1]):
+                    sub_length= sub_lengths[batch][slice_index]
+                    inputs_embeds[batch][slice_index][7:7+sub_length] = shift_text_hidden_states[batch][start_index+1:start_index+1+sub_length]
+                    start_index += sub_length
+            if self.training:
+                #print(shift_text_hidden_states_slice.shape, shift_text_embeddings.shape)
+                contrastive_loss = ((1 - F.cosine_similarity(inputs_embeds, gold_inputs_embeds, dim=-1)) * attention_mask).sum(-1).mean()
+                emb_loss = contrastive_loss
+                if emb_loss.device == torch.device("cuda:0"):
+                    self.log["contrastive_loss"] = contrastive_loss.item()
+            unit_embeds = F.embedding(unit_targets, embed_tokens_weight, padding_idx=self.pad_token_id)
+
+            inputs_embeds = torch.cat([inputs_embeds,unit_embeds], dim=2)
+        else:
+            inputs_embeds = self.aligner_MLP(inputs_embeds)
+            #[start_header_id] + _speaker + [end_header_id] + nl_tokens only for batch one!
+            units_ids = torch.IntTensor([[128009, 128006, 128000, 65576, 128007, 128000, 198]]).to(inputs_embeds.device)
+            units_prefix = self.model.embed_tokens(units_ids)
+            text_ids = torch.IntTensor([[128006, 128000, 65576, 128007, 128000, 198, 12800]]).to(inputs_embeds.device)
+            text_prefix = self.model.embed_tokens(text_ids)
+            inputs_embeds = torch.cat([text_prefix, inputs_embeds, units_prefix], dim = 1)
+
+        inputs_embeds = self.adapter(inputs_embeds)
+        if do_task == "finetune_kd":
+            return (emb_loss, inputs_embeds, unit_targets, atten_mask,)
+        else:
+            return (emb_loss, inputs_embeds)
+    
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if inputs_embeds == None:
+            # inputs_embeds = self.model.embed_tokens(input_ids)
+            # share embedding 0501 by kkq
+            embed_tokens_weight = torch.cat(
+                [
+                    self.model.embed_tokens.weight.detach(), self.unit_embedding.weight
+                ],
+                dim = 0,
+            )
+            # print(embed_tokens_weight, embed_tokens_weight.shape)
+            inputs_embeds = F.embedding(input_ids, embed_tokens_weight, padding_idx=self.pad_token_id)
+            inputs_embeds = self.adapter(inputs_embeds)
+        
+        outputs = self.model(
+            input_ids=None,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        cr_loss = None
+        if labels != None:
+            shift_labels = labels
+
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = shift_labels[..., 1:].contiguous()
+
+            loss_fct = CrossEntropyLoss()
+        
+            shift_logits = shift_logits.view(-1, (self.config.vocab_size + self.config.unit_vocab_size))
+            shift_labels = shift_labels.view(-1)
+            shift_labels = shift_labels.to(shift_logits.device)
+            
+            loss = loss_fct(shift_logits, shift_labels)
+                
+            if loss.device == torch.device("cuda:0"):
+                self.log["unit_loss"] = loss.item()
+
+            if cr_loss != None:
+                target_scale = loss.item() * 0.2
+                cr_loss_weight = target_scale / cr_loss.item() if cr_loss > target_scale else 1.0
+                loss = loss + cr_loss_weight * cr_loss
+
+            if loss.device == torch.device("cuda:0") and (self.current_step - 10) % 100 == 0:
+                print(self.log, loss.device)
+        
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+AutoConfig.register("T2ULlama", T2ULlamaConfig)
+AutoModelForCausalLM.register(T2ULlamaConfig, T2ULlamaForCausalLM)

show_case/2.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03198985b12bd892d05b8ae9b2e6c8303b15b0a570eea9647b2e314332340711
+size 165742

show_case/Translate_de_audio_prompt.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aee2ba64475fb597ec0a68da4e1a552fc4463bca0e9393d9f170e37f563d7792
+size 288164

text_to_speech.py

@@ -0,0 +1,326 @@
+from fairseq.dataclass.configs import FairseqConfig
+from fairseq import utils
+from fairseq.models.text_to_speech.vocoder import CodeHiFiGANVocoder
+from fairseq import checkpoint_utils, options, tasks, utils
+from fairseq.distributed import utils as distributed_utils
+import torch
+import json
+from tqdm import tqdm
+import random
+import soundfile as sf
+import numpy as np
+import ast
+import time
+import math
+from fairseq.dataclass.utils import convert_namespace_to_omegaconf
+from fairseq.token_generation_constraints import pack_constraints, unpack_constraints
+from fairseq_cli.generate import get_symbols_to_strip_from_output
+from collections import namedtuple
+import sys
+from argparse import Namespace
+import argparse
+import sentencepiece as spm
+import re
+
+Batch = namedtuple("Batch", "ids src_tokens src_lengths constraints")
+Translation = namedtuple("Translation", "src_str hypos pos_scores alignments")
+
+def make_batches(lines, cfg, task, max_positions, encode_fn):
+    def encode_fn_target(x):
+        return encode_fn(x)
+
+    if cfg.generation.constraints:
+        # Strip (tab-delimited) contraints, if present, from input lines,
+        # store them in batch_constraints
+        batch_constraints = [list() for _ in lines]
+        for i, line in enumerate(lines):
+            if "\t" in line:
+                lines[i], *batch_constraints[i] = line.split("\t")
+
+        # Convert each List[str] to List[Tensor]
+        for i, constraint_list in enumerate(batch_constraints):
+            batch_constraints[i] = [
+                task.target_dictionary.encode_line(
+                    encode_fn_target(constraint),
+                    append_eos=False,
+                    add_if_not_exist=False,
+                )
+                for constraint in constraint_list
+            ]
+
+    if cfg.generation.constraints:
+        constraints_tensor = pack_constraints(batch_constraints)
+    else:
+        constraints_tensor = None
+
+    tokens, lengths = task.get_interactive_tokens_and_lengths(lines, encode_fn)
+
+    itr = task.get_batch_iterator(
+        dataset=task.build_dataset_for_inference(
+            tokens, lengths, constraints=constraints_tensor
+        ),
+        max_tokens=cfg.dataset.max_tokens,
+        max_sentences=cfg.dataset.batch_size,
+        max_positions=max_positions,
+        ignore_invalid_inputs=cfg.dataset.skip_invalid_size_inputs_valid_test,
+    ).next_epoch_itr(shuffle=False)
+    for batch in itr:
+        ids = batch["id"]
+        src_tokens = batch["net_input"]["src_tokens"]
+        src_lengths = batch["net_input"]["src_lengths"]
+        constraints = batch.get("constraints", None)
+
+        yield Batch(
+            ids=ids,
+            src_tokens=src_tokens,
+            src_lengths=src_lengths,
+            constraints=constraints,
+        )
+
+def tokenize(inputs, sp):
+    text = re.sub(r'[^\w\s]', '', inputs.lower())
+    inputs = ' '.join(sp.EncodeAsPieces(text))
+    # print(inputs)
+    return inputs
+
+def get_t2u_config(model, beam=5):
+
+    sys.argv = [
+        "fairseq-interactive",
+        "libri_t2u", 
+        "--path", model, 
+        "--gen-subset", "valid", 
+        "--max-len-b", "1024", 
+        "--max-source-positions", "500", 
+        "--max-target-positions", "1024", 
+        "--beam", str(beam), 
+        "--results-path", "decode" 
+    ]
+
+    parser = options.get_interactive_generation_parser()
+    args = options.parse_args_and_arch(parser)
+    # distributed_utils.call_main(convert_namespace_to_omegaconf(args), load_text2units_model)
+    return convert_namespace_to_omegaconf(args)
+
+def load_text2units_model(cfg: FairseqConfig, device):
+
+    if isinstance(cfg, Namespace):
+        cfg = convert_namespace_to_omegaconf(cfg)
+
+    utils.import_user_module(cfg.common)
+    if cfg.interactive.buffer_size < 1:
+        cfg.interactive.buffer_size = 1
+    if cfg.dataset.max_tokens is None and cfg.dataset.batch_size is None:
+        cfg.dataset.batch_size = 1
+
+    assert (
+        not cfg.generation.sampling or cfg.generation.nbest == cfg.generation.beam
+    ), "--sampling requires --nbest to be equal to --beam"
+    assert (
+        not cfg.dataset.batch_size
+        or cfg.dataset.batch_size <= cfg.interactive.buffer_size
+    ), "--batch-size cannot be larger than --buffer-size"
+
+    # Fix seed for stochastic decoding
+    if cfg.common.seed is not None and not cfg.generation.no_seed_provided:
+        np.random.seed(cfg.common.seed)
+        utils.set_torch_seed(cfg.common.seed)
+
+    use_cuda = torch.cuda.is_available() and not cfg.common.cpu
+
+    # Setup task, e.g., translation
+    task = tasks.setup_task(cfg.task)
+
+    # Load ensemble
+    overrides = ast.literal_eval(cfg.common_eval.model_overrides)
+    models, _model_args = checkpoint_utils.load_model_ensemble(
+        utils.split_paths(cfg.common_eval.path),
+        arg_overrides=overrides,
+        task=task,
+        suffix=cfg.checkpoint.checkpoint_suffix,
+        strict=(cfg.checkpoint.checkpoint_shard_count == 1),
+        num_shards=cfg.checkpoint.checkpoint_shard_count,
+    )
+
+    # Set dictionaries
+    src_dict = task.source_dictionary
+    tgt_dict = task.target_dictionary
+
+    # Optimize ensemble for generation
+    for model in models:
+        if model is None:
+            continue
+        if cfg.common.fp16:
+            model.half()
+        if use_cuda and not cfg.distributed_training.pipeline_model_parallel:
+            model.cuda()
+        model.prepare_for_inference_(cfg)
+
+    # Initialize generator
+    generator = task.build_generator(models, cfg.generation)
+
+    # Handle tokenization and BPE
+    tokenizer = task.build_tokenizer(cfg.tokenizer)
+    bpe = task.build_bpe(cfg.bpe)
+
+    return {
+        "models": models, 
+        "generator": generator, 
+        "tokenizer": tokenizer, 
+        "bpe": bpe, 
+        "task": task, 
+        "src_dict": src_dict,
+        "tgt_dict": tgt_dict,
+        "use_cuda": use_cuda
+    }
+
+def gen_units(model, cfg, inputs):
+    inputs = [inputs]
+    
+    models = model['models']
+    generator = model['generator']
+    tokenizer = model['tokenizer']
+    bpe = model['bpe']
+    task = model['task']
+    src_dict = model['src_dict']
+    tgt_dict = model['tgt_dict']
+    use_cuda = model['use_cuda']
+
+    def encode_fn(x):
+        if tokenizer is not None:
+            x = tokenizer.encode(x)
+        if bpe is not None:
+            x = bpe.encode(x)
+        return x
+
+    def decode_fn(x):
+        if bpe is not None:
+            x = bpe.decode(x)
+        if tokenizer is not None:
+            x = tokenizer.decode(x)
+        return x
+
+    align_dict = utils.load_align_dict(cfg.generation.replace_unk)
+
+    max_positions = utils.resolve_max_positions(
+        task.max_positions(), *[model.max_positions() for model in models]
+    )
+
+    start_id = 0
+    results = []
+    for batch in make_batches(inputs, cfg, task, max_positions, encode_fn):
+        print("[INFO_DEBUG]", batch)
+        bsz = batch.src_tokens.size(0)
+        src_tokens = batch.src_tokens
+        src_lengths = batch.src_lengths
+        constraints = batch.constraints
+        if use_cuda:
+            src_tokens = src_tokens.cuda()
+            src_lengths = src_lengths.cuda()
+            if constraints is not None:
+                constraints = constraints.cuda()
+
+        sample = {
+            "net_input": {
+                "src_tokens": src_tokens,
+                "src_lengths": src_lengths,
+            },
+        }
+        translate_start_time = time.time()
+        translations = task.inference_step(
+            generator, models, sample, constraints=constraints
+        )
+        translate_time = time.time() - translate_start_time
+        list_constraints = [[] for _ in range(bsz)]
+        if cfg.generation.constraints:
+            list_constraints = [unpack_constraints(c) for c in constraints]
+        for i, (id, hypos) in enumerate(zip(batch.ids.tolist(), translations)):
+            src_tokens_i = utils.strip_pad(src_tokens[i], tgt_dict.pad())
+            constraints = list_constraints[i]
+            results.append(
+                (
+                    start_id + id,
+                    src_tokens_i,
+                    hypos,
+                    {
+                        "constraints": constraints,
+                        "time": translate_time / len(translations),
+                    },
+                )
+            )
+
+    # print(results)
+
+    units = []
+    for id_, _, hypos, info in sorted(results, key=lambda x: x[0]):
+        print("W-{}\t{:.3f}\tseconds".format(id_, info["time"]))
+
+        # Process top predictions
+        for hypo in hypos[: min(len(hypos), cfg.generation.nbest)]:
+            hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
+                hypo_tokens=hypo["tokens"].int().cpu(),
+                src_str="",
+                alignment=hypo["alignment"],
+                align_dict=align_dict,
+                tgt_dict=tgt_dict,
+                remove_bpe=cfg.common_eval.post_process,
+                extra_symbols_to_ignore=get_symbols_to_strip_from_output(generator),
+            )
+
+        units.append(list(map(int, hypo_str.split(' '))))
+
+    return units
+
+def get_vocoder_config(vocoder, config):
+    
+    args = argparse.Namespace(
+        vocoder=vocoder,
+        vocoder_cfg=config,
+        dur_prediction=True,  
+        speaker_id=1,
+        cpu=False  
+    )
+
+    return args
+
+def load_units_vocoder(args, device):
+    with open(args.vocoder_cfg) as f:
+        vocoder_cfg = json.load(f)
+    vocoder = CodeHiFiGANVocoder(args.vocoder, vocoder_cfg).to(device)
+
+    multispkr = vocoder.model.multispkr
+    if multispkr:
+        num_speakers = vocoder_cfg.get(
+            "num_speakers", 200
+        )  # following the default in codehifigan to set to 200
+        assert (
+            args.speaker_id < num_speakers
+        ), f"invalid --speaker-id ({args.speaker_id}) with total #speakers = {num_speakers}"
+
+    return vocoder, num_speakers if multispkr else 1, 'cuda' in device
+
+def gen_wav(vocoder, args, data, device):
+    vocoder, num_speakers, use_cuda = vocoder
+    res = []
+    for i, d in enumerate(data): # tqdm is removed for cleaner streaming
+        x = {
+            "code": torch.LongTensor(d).view(1, -1).to(device),
+        }
+        suffix = ""
+
+        multispkr = vocoder.model.multispkr
+        if multispkr:
+            spk = (
+                random.randint(0, num_speakers - 1)
+                if args.speaker_id == -1
+                else args.speaker_id
+            )
+            suffix = f"_spk{spk}"
+            x["spkr"] = torch.LongTensor([spk]).view(1, 1)
+
+        x = utils.move_to_cuda(x) if use_cuda else x
+        wav = vocoder(x, args.dur_prediction).detach().cpu().numpy()
+        
+        res.append(wav)
+
+    return res[0]