# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import copy import os import time from collections import Counter from contextlib import contextmanager import librosa import torch import torch.nn as nn import torch.nn.functional as F from lightning import LightningModule from omegaconf import DictConfig from peft import PeftModel from torch.distributed.fsdp import fully_shard from torch.distributed.tensor import Replicate, Shard from torch.distributed.tensor.parallel import ( ColwiseParallel, PrepareModuleInput, RowwiseParallel, SequenceParallel, loss_parallel, parallelize_module, ) from nemo.collections.audio.parts.utils.transforms import resample from nemo.collections.common.tokenizers import AutoTokenizer from nemo.collections.speechlm2.data.utils import get_pad_id from nemo.collections.speechlm2.modules.ear_tts_model import RVQEARTTSModel from nemo.collections.speechlm2.modules.ear_tts_vae_codec import RVQVAEModel from nemo.collections.speechlm2.parts.hf_hub import HFHubMixin from nemo.collections.speechlm2.parts.metrics.asr_bleu import ASRBLEU from nemo.collections.speechlm2.parts.metrics.asr_cer_wer import Intelligibility from nemo.collections.speechlm2.parts.metrics.results_logger import ResultsLogger from nemo.collections.speechlm2.parts.metrics.secs import SECS from nemo.collections.speechlm2.parts.optim_setup import ( configure_optimizers, configure_optimizers_exclude_norm_from_wd, is_frozen, ) from nemo.collections.speechlm2.parts.precision import fp32_precision from nemo.collections.speechlm2.parts.pretrained import ( load_checkpoint, load_pretrained_hf, set_model_dict_for_partial_init, ) from nemo.utils import logging class DuplexEARTTS(LightningModule, HFHubMixin): def __init__(self, cfg: dict) -> None: assert isinstance(cfg, dict), ( "You must pass the config to DuplexEARTTS as a Python dict to support hyperparameter serialization " f"in PTL checkpoints (we got: '{type(cfg)=}')." ) super().__init__() self.save_hyperparameters() # convert dict to config cfg = DictConfig(cfg) self.trainer_config = cfg.get("trainer", None) self.data_cfg = cfg.data self.cfg = cfg.model self.target_sample_rate = cfg.data.target_sample_rate self.source_sample_rate = cfg.data.source_sample_rate self.normalize_text = cfg.data.get("normalize_text", False) self.validation_save_path = os.path.join(cfg.exp_manager.explicit_log_dir, "validation_logs") # move back text channel by x, in inference it advance the text channel prediction by x frames self.advance_text_channel_by = self.cfg.get("advance_text_channel_by", None) # Load ForCausalLM if self.cfg.tts_config.context_hidden_size is not None: self.language_model = self._load_language_model(self.cfg) self.embed_tokens = self._load_embed_tokens(self.cfg) # delete llm because we use it only to get the embbeding tokens del self.language_model # get codec run precision self.audio_codec_run_dtype = getattr(torch, self.cfg.get("audio_codec_run_dtype", "float32"), torch.float32) # Load tokenizer tokenizer_src = self.cfg.get("tokenizer_path", None) or self.cfg.pretrained_lm_name self.tokenizer = AutoTokenizer( tokenizer_src, use_fast=True, trust_remote_code=self.cfg.get("trust_remote_code", False), bos_token=self.cfg.get("bos_token", None), eos_token=self.cfg.get("eos_token", None), pad_token=self.cfg.get("pad_token", None), ) # Note that we are using fast tokenizer # Instantiate TTS model self.tts_model = RVQEARTTSModel(DictConfig(self.cfg.tts_config), tokenizer=self.tokenizer) # Load and initialize audio codec, and bind RVQ embeddings to the TTS model setup_audio_codec(self) self._codebook_size = self.tts_model.config.codebook_size # compute samples per frame self.source_samples_per_frame = int(self.source_sample_rate * cfg.data.frame_length) # get codec silence tokens codec_silence_tokens = self.get_codec_silence_frame() self.register_buffer("codec_silence_tokens", codec_silence_tokens) # cached for quicker audio decoding self.register_buffer( "_control_codes", torch.tensor([self.speech_bos_id, self.speech_eos_id, self.speech_pad_id], device=self.device), ) self._use_fsdp = False self._use_tp = False self.audio_prompt_latents = nn.ParameterDict() def get_codec_silence_frame_last_one(self): audio = torch.zeros(1, 10 * self.target_sample_rate).float().to(self.device) audio_len = torch.tensor([audio.size(-1)]).long() audio, audio_len = self.pad_audio_to_factor(audio, audio_len, self.target_samples_per_frame) with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad(): sil_codes, sil_codes_lens = self.audio_codec.encode( audio.unsqueeze(1).to(self.audio_codec_run_dtype), audio_len ) return sil_codes[0, -1] def get_codec_silence_frame(self): # Generate long zero waveform (silence) audio = torch.zeros(1, 10 * self.target_sample_rate).float().to(self.device) audio_len = torch.tensor([audio.size(-1)]).long() audio, audio_len = self.pad_audio_to_factor(audio, audio_len, self.target_samples_per_frame) with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad(): sil_codes, _ = self.audio_codec.encode( audio.unsqueeze(1).to(self.audio_codec_run_dtype), audio_len ) # [1, T, C] sil_codes = sil_codes[0] # [T, C] # Convert each frame (C tokens) into a tuple combos = [tuple(row.tolist()) for row in sil_codes] # Count frequencies counter = Counter(combos) # Pick the most common combination most_common_combo, freq = counter.most_common(1)[0] # Return as tensor [C] return torch.tensor(most_common_combo, device=self.device, dtype=torch.long) def _load_embed_tokens(self, cfg) -> nn.Embedding: """Load token embedding layer for RVQ-EAR-TTS.""" if self.language_model: assert callable(self.language_model.get_input_embeddings) embed_tokens: nn.Embedding = self.language_model.get_input_embeddings() else: embed_tokens_state_dict = torch.load( cfg.pretrained_lm_embedding_path, map_location="cpu", weights_only=True ) # Create token embedding layer vocab_size, hidden_size = embed_tokens_state_dict["weight"].size() embed_tokens = nn.Embedding(vocab_size, hidden_size, dtype=torch.bfloat16) embed_tokens.load_state_dict(embed_tokens_state_dict) return embed_tokens def _load_language_model(self, cfg): """Load language model for RVQ-EAR-TTS.""" if cfg.pretrained_lm_name: language_model = load_pretrained_hf( self.cfg.pretrained_lm_name, pretrained_weights=True, trust_remote_code=True ).eval() else: language_model = None return language_model def restore_from_pretrained_checkpoint(self, checkpoint_path): """ Loads model weights a pretrained checkpoint file, supporting partial loading from safetensor and PyTorch formats. Args: checkpoint_path (str): Path to checkpoint file. Returns: None. The model is updated in-place. """ if checkpoint_path is not None: checkpoint_state = load_checkpoint(checkpoint_path) checkpoint_state = set_model_dict_for_partial_init(checkpoint_state, self.state_dict()) if self.cfg.get("rescale_pretrained_weights", None): checkpoint_state = rescale_state_dict( checkpoint_state, first_n_layers=self.cfg.get("rescale_first_n_layers", None) ) self.load_state_dict(checkpoint_state, strict=True) logging.info(f"Model restored from the checkpoint: {checkpoint_path} !") @property def device(self): return next(self.parameters()).device @property def speech_vocab_size(self): """Return the size of the audio codec codebook including extra speech BOS and EOS tokens.""" if self.use_local_transformer and self.local_transformer_type == "nar": # add extra token for mask return self._codebook_size + 4 return self._codebook_size + 3 @property def speech_bos_id(self) -> int: """Indicates start of utterance generation (not start of inference!).""" if self.cfg.get("custom_speech_bos_id", None): return self.cfg.get("custom_speech_bos_id") return self._codebook_size + 2 @property def speech_eos_id(self) -> int: """Indicates end of utterance generation.""" if self.cfg.get("custom_speech_eos_id", None): return self.cfg.get("custom_speech_eos_id") return self._codebook_size + 1 @property def speech_pad_id(self) -> int: """Indicates start of inference (the very first frame).""" if self.cfg.get("custom_speech_pad_id", None): return self.cfg.get("custom_speech_pad_id") return self._codebook_size @property def text_vocab_size(self): """Return the size of the text tokenizer.""" return self.tokenizer.vocab_size @property def text_bos_id(self) -> int: return self.tokenizer.bos_id @property def text_eos_id(self) -> int: return self.tokenizer.eos_id @property def text_pad_id(self) -> int: """ Text pad ID is used as a 'blank' for frames when the model is not speaking and for frames where the model is speaking but has already predicted the entire text channel's content. Example: flow: |---user---||-------assistant--------||-user-| text channel: 0000000000 1xxxxxxx0000000000000002 000000 Where 0 indicates PAD ID, 1 indicates BOS ID, 2 indacates EOS ID, and x indicates tokens corresponding to actual text """ return get_pad_id(self.tokenizer) def pad_audio_to_factor(self, audio, audio_len, samples_per_frame, downsampling_factor: int = 1): """ Zero pad the end of the audio so that we do not have a partial end frame. The output will be zero-padded to have an integer number of frames of length `samples_per_frame * downsampling_factor`. Args: audio: input time-domain signal (B, T) audio_len: valid length for each example in the batch (B,) samples_per_frame: number of samples per frame downsampling_factor: how much each frame is downsampled in later processing Returns: padded_audio: Padded time-domain signal (B, T') padded_len: Adjusted valid lengths (B,) """ with fp32_precision(): total_factor = samples_per_frame * downsampling_factor padded_len = total_factor * torch.ceil(audio_len / total_factor).int() max_len = padded_len.max().int().item() num_padding = max_len - audio.shape[1] padded_audio = F.pad(audio, (0, num_padding)) return padded_audio, padded_len def prepare_inputs(self, batch: dict): """ Prepare inputs, extracting audio tokens and padding if needed. """ # check if audios has the same batch size assert batch["source_audio"].size(0) == batch["target_audio"].size(0) assert batch["audio_prompt"].size(0) == batch["target_audio"].size(0) target_audio = batch["target_audio"] target_audio_lens = batch["target_audio_lens"] target_text_tokens = batch["target_text_tokens"] non_prompt_mask = batch["non_prompt_mask"] aligned_attention_mask = batch["aligned_attention_mask"] aligned_position_ids = batch["aligned_position_ids"] if self.training and (self.cfg.get("empty_turn_probability", 0.0) > 0): # Randomly decide whether this batch gets emptied if torch.rand(1).item() < self.cfg.empty_turn_probability: # Zero out audio target_audio = torch.zeros_like(target_audio) # Create mask for tokens we want to drop # Keep BOS and EOS, drop the rest. keep_mask = (target_text_tokens == self.text_bos_id) | (target_text_tokens == self.text_eos_id) full_dropout_mask = ~keep_mask # True = positions to replace with PAD # Replace all non-BOS/EOS with PAD target_text_tokens = torch.where( full_dropout_mask, torch.full_like(target_text_tokens, self.text_pad_id), target_text_tokens ) # extract target audio codes target_audio, target_audio_lens = self.pad_audio_to_factor( target_audio, target_audio_lens, self.target_samples_per_frame, 1 ) with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad(): target_codes, target_codes_lens = self.audio_codec.encode( target_audio.unsqueeze(1).to(self.audio_codec_run_dtype), target_audio_lens ) with fp32_precision(): target_len = target_codes.shape[1] # Pad or truncate sequence variables def pad_or_truncate(x, pad_value=0): if x.dim() == 2: # [B, T] L = x.shape[1] if L < target_len: return F.pad(x, (0, target_len - L), value=pad_value) else: return x[:, :target_len] return x # leave others for now target_text_tokens = pad_or_truncate(target_text_tokens, pad_value=self.text_pad_id) non_prompt_mask = pad_or_truncate(non_prompt_mask, pad_value=0) aligned_position_ids = pad_or_truncate(aligned_position_ids, pad_value=0) # Correct attention mask padding/truncation B, H, L1, L2 = aligned_attention_mask.shape new_len = target_len if L1 < new_len or L2 < new_len: pad_rows = new_len - L1 pad_cols = new_len - L2 aligned_attention_mask = F.pad(aligned_attention_mask, (0, pad_cols, 0, pad_rows)) elif L1 > new_len or L2 > new_len: aligned_attention_mask = aligned_attention_mask[:, :, :new_len, :new_len] # set the pad token for the first BOS frame target_codes_aligned = target_codes.clone() target_codes_aligned[:, 0] = self.speech_pad_id # set special token in the last audio prompt (it will works as a BOS token) pos = non_prompt_mask.float().argmax(dim=1) # shape: [B] row_idx = torch.arange(B, device=self.device) # set the extra self.speech_pad_id at first 1 position in non_prompt_mask target_codes_aligned[row_idx, pos] = self.speech_pad_id # EOS dropout to make the model more robust if self.training and self.cfg.get("text_eos_dropout_prob", 0.0) > 0: # Mask EOS positions eos_mask = target_text_tokens == self.text_eos_id # Random dropout only on EOS positions dropout_mask = ( torch.rand(eos_mask.sum(), device=target_text_tokens.device) < self.cfg.text_eos_dropout_prob ) # Scatter dropout decisions into [B, T] full_dropout_mask = torch.zeros_like(target_text_tokens, dtype=torch.bool) full_dropout_mask[eos_mask] = dropout_mask # Replace dropped EOS with PAD target_text_tokens = torch.where( full_dropout_mask, torch.full_like(target_text_tokens, self.text_pad_id), target_text_tokens ) if self.training and self.cfg.get("text_eos_duplicate_prob", 0.0) > 0: p = self.cfg.text_eos_duplicate_prob # [B, T] mask of EOS positions eos_mask = target_text_tokens == self.text_eos_id # Flatten EOS positions: tensor of shape [N, 2] where each row = (batch_idx, time_idx) eos_positions = eos_mask.nonzero(as_tuple=False) # [N, 2] if eos_positions.numel() > 0: N = eos_positions.shape[0] # One random decision per EOS occurrence duplicate_decision = torch.rand(N, device=target_text_tokens.device) < p # [N] # Filter only EOS tokens that will be duplicated and are not at position t=0 valid = (eos_positions[:, 1] > 0) & duplicate_decision # [N] if valid.any(): # Select only valid EOS positions valid_positions = eos_positions[valid] # [M, 2] # Indices for the token BEFORE the EOS (t-1) b_idx = valid_positions[:, 0] t_idx = valid_positions[:, 1] - 1 # Replace token before EOS with an EOS target_text_tokens[b_idx, t_idx] = self.text_eos_id # BOS dropout to make the model more robust if self.training and self.cfg.get("text_bos_dropout_prob", 0.0) > 0: # Mask BOS positions bos_mask = target_text_tokens == self.text_bos_id # Random dropout only on BOS positions dropout_mask = ( torch.rand(bos_mask.sum(), device=target_text_tokens.device) < self.cfg.text_bos_dropout_prob ) # Scatter dropout decisions into [B, T] full_dropout_mask = torch.zeros_like(target_text_tokens, dtype=torch.bool) full_dropout_mask[bos_mask] = dropout_mask # Replace dropped BOS with PAD target_text_tokens = torch.where( full_dropout_mask, torch.full_like(target_text_tokens, self.text_pad_id), target_text_tokens, ) # BOS dropout to make the model more robust if self.training and self.cfg.get("text_bos_dropout_prob", 0.0) > 0: prob = self.cfg.text_bos_dropout_prob # e.g., 0.5 # Identify all BOS positions [B, T] bos_mask = target_text_tokens == self.text_bos_id # Get indices of sequences that actually have a BOS token # We need to know *where* the BOS tokens are to drop them. # tensor of coordinates: [[batch_idx, seq_idx], ...] bos_indices = torch.nonzero(bos_mask) num_bos = bos_indices.shape[0] if num_bos > 0: # Create a random dropout decision for each BOS instance drop_decisions = torch.rand(num_bos, device=target_text_tokens.device) < prob # Ensure at least one is dropped if drop_decisions.sum() == 0: # Pick one random index from the available BOS locations to drop force_idx = torch.randint(0, num_bos, (1,), device=target_text_tokens.device) drop_decisions[force_idx] = True # 5. Apply the dropout # We need to map the decisions back to the full tensor # Create a mask of the same shape as target_text_tokens full_dropout_mask = torch.zeros_like(target_text_tokens, dtype=torch.bool) # Set True only at the specific (batch, seq) coordinates we chose to drop # bos_indices[:, 0] are batch indices, bos_indices[:, 1] are seq indices full_dropout_mask[bos_indices[:, 0], bos_indices[:, 1]] = drop_decisions # 6. Replace dropped BOS with PAD target_text_tokens = torch.where( full_dropout_mask, torch.full_like(target_text_tokens, self.text_pad_id), target_text_tokens, ) # shift text tokens subword_ids = F.pad(target_text_tokens[:, 1:], [0, 1]) # note that we are using a text mask where we are ignoring the desc + audio prompt but we are keeping 1 until the audio ends to support duplex subword_mask = F.pad(non_prompt_mask[:, 1:], [0, 1]) # detach embedding as in eartts if self.cfg.tts_config.context_hidden_size is not None: context_hidden_state = self.embed_tokens(target_text_tokens).detach() else: context_hidden_state = None if self._use_tp: tp_world_size = self.device_mesh["tensor_parallel"].size() if (remainder := (target_text_tokens.shape[1] - 1) % tp_world_size) != 0: target_text_tokens = target_text_tokens[:, :-remainder] target_codes_aligned = target_codes_aligned[:, :-remainder] target_codes_aligned = target_codes_aligned[:, :-remainder] subword_ids = subword_ids[:, :-remainder] subword_mask = subword_mask[:, :-remainder] return { "code": target_codes_aligned, "audio_mask": non_prompt_mask, # set audio_mask as non_prompt_mask to avoid the audio prompt in loss computation "attention_mask": aligned_attention_mask, "position_ids": aligned_position_ids, "subword_ids": subword_ids, "subword_mask": subword_mask, "context_hidden_state": context_hidden_state, "output_lens": target_codes_lens, "non_prompt_mask": non_prompt_mask, "target_text_tokens": target_text_tokens, } def training_step(self, batch: dict, batch_idx: int): for m in (self.tts_model,): if is_frozen(m): m.eval() inputs = self.prepare_inputs(batch) tts_output = self.tts_model( code=inputs["code"], audio_mask=inputs["audio_mask"], attention_mask=inputs["attention_mask"], position_ids=inputs["position_ids"], context_hidden_state=inputs["context_hidden_state"], subword_ids=inputs["subword_ids"], subword_mask=inputs["subword_mask"], non_prompt_mask=inputs["non_prompt_mask"], dataset_type=batch.get("dataset_type", None), ) loss_dict = {"lm_loss": tts_output.lm_loss, "c_loss": tts_output.c_loss, "k_loss": tts_output.k_loss} loss = sum(loss_dict.values()) num_frames = inputs["output_lens"].sum() B, T = inputs["code"].shape[:2] ans = { "loss": loss, "learning_rate": ( torch.as_tensor(self.trainer.optimizers[0].param_groups[0]['lr'] if self._trainer is not None else 0) ), "batch_size": B, "sequence_length": T, "num_frames": num_frames.to(torch.float32), # avoid warning "padding_ratio": num_frames / (B * T), **loss_dict, } self.log_dict(ans, on_step=True) return ans def ensures_codec_target_dtype(self) -> None: """ Ensures the audio codec is instantiated with the target dtype. This method checks whether `self.audio_codec` exists and whether its parameters match `self.audio_codec_run_dtype`. If the codec is missing or is running with the wrong dtype (e.g., due to PTL auto-downcasting), the codec is reloaded by calling `setup_audio_codec()`. Intended to be called at runtime boundaries such as: - `on_train_epoch_start` - `on_validation_epoch_start` """ if hasattr(self, "audio_codec") and next(self.audio_codec.parameters()).dtype == self.audio_codec_run_dtype: self.audio_codec.eval() return # already correct precision → no-op setup_audio_codec(self) def on_train_epoch_start(self) -> None: self.ensures_codec_target_dtype() # potentially reloads the audio codec to make sure it's in target codec precision def on_train_epoch_end(self) -> None: # log model stats to debug gradient weights issues self.log_model_stats() def log_model_stats(self): total_w_sq = 0.0 total_w_params = 0 max_abs_w = 0.0 sum_w = 0.0 total_g_sq = 0.0 total_g_params = 0 for p in self.parameters(): if not p.requires_grad: continue # ----- weights ----- w = p.detach().cpu().float() # safe offline copy total_w_sq += (w * w).sum().item() total_w_params += w.numel() max_abs_w = max(max_abs_w, w.abs().max().item()) sum_w += w.sum().item() # ----- grads (optional, disabled for speed) ----- if p.grad is not None: g = p.grad.detach().cpu().float() total_g_sq += (g * g).sum().item() total_g_params += g.numel() # L2 norms weight_l2 = (total_w_sq**0.5) if total_w_sq > 0 else 0.0 # RMS (global) weight_rms = ((total_w_sq / total_w_params) ** 0.5) if total_w_params > 0 else 0.0 # Mean weight_mean = sum_w / total_w_params if total_w_params > 0 else 0.0 # direct float logging avoids device sync penalty self.log("weights/L2", weight_l2, on_epoch=True, sync_dist=True) self.log("weights/RMS", weight_rms, on_epoch=True, sync_dist=True) self.log("weights/max_abs", max_abs_w, on_epoch=True, sync_dist=True) self.log("weights/mean", weight_mean, on_epoch=True, sync_dist=True) def on_validation_epoch_start(self) -> None: if torch.distributed.is_initialized(): self.trainer.strategy.model.require_backward_grad_sync = False self.ensures_codec_target_dtype() # potentially reloads the audio codec to make sure it's in target codec precision self.results_logger = ResultsLogger(self.validation_save_path).reset() self.asr_bleu = ASRBLEU(self.cfg.scoring_asr).reset() self.intelligibility = Intelligibility(self.cfg.scoring_asr, reuse_asr_hyps=True).reset() self.secs = SECS(self.cfg.get("scoring_se", "titanet_large")).reset() def on_validation_epoch_end(self, prefix="val") -> None: asr_bleu = self.asr_bleu.compute() for k, m in asr_bleu.items(): self.log(f"{prefix}_{k}", m.to(self.device), on_epoch=True, sync_dist=True) cer_wer = self.intelligibility.compute() for k, m in cer_wer.items(): self.log(f"{prefix}_{k}", m.to(self.device), on_epoch=True, sync_dist=True) secs = self.secs.compute() for k, m in secs.items(): self.log(f"{prefix}_{k}", m.to(self.device), on_epoch=True, sync_dist=True) self.results_logger.compute_and_save() if torch.distributed.is_initialized(): self.trainer.strategy.model.require_backward_grad_sync = True def get_teacher_force_inference_audio(self, batch, guidance_enabled=True): inputs = self.prepare_inputs(batch) tts_output = self.tts_model( code=inputs["code"], audio_mask=inputs["audio_mask"], attention_mask=inputs["attention_mask"], position_ids=inputs["position_ids"], context_hidden_state=inputs["context_hidden_state"], subword_ids=inputs["subword_ids"], subword_mask=inputs["subword_mask"], non_prompt_mask=inputs["non_prompt_mask"], generation_config=self._get_generation_config(guidance_enabled=guidance_enabled), teacher_forcing_inference=True, guidance_enabled=guidance_enabled, ) tf_audio_codes_pred = tts_output["codes"].squeeze(2) # decode audio tf_audio_codes_pred = replace_control_speech_codes( tf_audio_codes_pred, self._control_codes, self.codec_silence_tokens ) with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad(): audio_pred, audio_len = self.audio_codec.decode(tf_audio_codes_pred, inputs["output_lens"]) return audio_pred.squeeze(1), audio_len def _get_generation_config(self, guidance_enabled: bool = False): """Get default generation config for EAR-TTS.""" return { "num_iter": 8, "guidance_scale": self.cfg.get("inference_guidance_scale", 0.5) if guidance_enabled else None, "top_p_or_k": self.cfg.get("inference_top_p_or_k", 0.8), "noise_scale": self.cfg.get("inference_noise_scale", 0.8), "eos_threshold": -3.0, } @torch.inference_mode() def run_evaluation_one_batch(self, name, dataset_batch, use_dataloader_init=False): """ Runs evaluation and scoring for a single data batch, logging metrics and updating result buffers. Args: name (str): Name/id for the batch (for logging). dataset_batch (dict): Batch of data inputs, supports batched text/audio/etc. use_dataloader_init (bool, optional): If True, use dataloader initialization for prompts. Returns: None. Outputs are logged and stored in result buffers. """ results = {} inputs = self.prepare_inputs(dataset_batch) results["audio_tf"], results["audio_tf_len"] = self.get_teacher_force_inference_audio(dataset_batch) if use_dataloader_init: # cut it on prompt init_inputs = { "code": inputs["code"], "audio_mask": inputs["audio_mask"], "non_prompt_mask": inputs["non_prompt_mask"], "context_hidden_state": inputs["context_hidden_state"], "subword_ids": inputs["subword_ids"], "subword_mask": inputs["subword_mask"], } # cut init_inputs to consider only the prompt for key in init_inputs: if init_inputs[key] is not None: init_inputs[key] = torch.stack( [init_inputs[key][i, :plen] for i, plen in enumerate(dataset_batch["prompt_lens"])] ) else: sp = dataset_batch.get("system_prompts_raw") system_prompt = sp[0] if sp else None # set init inputs and get it self.set_init_inputs( speaker_audio=dataset_batch["audio_prompt"], speaker_audio_lens=dataset_batch["audio_prompt_lens"], system_prompt=system_prompt, # use the first position of the batch as system prompt ) init_inputs = self.get_init_inputs(B=inputs["subword_ids"].size(0)) # remove the prompt from the target_text_tokens to emulate S2S connected inference next_subword_ids = torch.stack( [ inputs["subword_ids"][i, plen:] # slice each element for i, plen in enumerate(dataset_batch["prompt_lens"]) ] ) results["audio"], results["audio_len"] = self.offline_inference( next_subword_ids=next_subword_ids, task=dataset_batch["task"][0], init_inputs=init_inputs, ) # remove prompt padding from the user audio as autoregressive inference does not return the prompt dataset_batch["source_audio"] = dataset_batch["source_audio"][ :, -int(next_subword_ids.size(-1) * self.source_samples_per_frame) : ] # clean prompt from the audio results["audio_tf"] = results["audio_tf"][:, -int(next_subword_ids.size(-1) * self.target_samples_per_frame) :] # remove prompt from target audio target_audio_no_prompt = dataset_batch["target_audio"][ :, -int(next_subword_ids.size(-1) * self.target_samples_per_frame) : ] target_audio_no_prompt_lens = dataset_batch["target_audio_lens"] - ( torch.tensor( dataset_batch["prompt_lens"], dtype=torch.long, device=dataset_batch["target_audio_lens"].device, ) * self.target_samples_per_frame ) with fp32_precision(): # resample is fragile to bfloat16 default dtype metric_audio_pred = results["audio"] metric_audio_pred_lens = results["audio_len"] # resample audio to the asr sampling rate metric_audio_pred = resample(metric_audio_pred, self.target_sample_rate, 16000) metric_audio_pred_lens = (metric_audio_pred_lens / self.target_sample_rate * 16000).to(torch.long) # reshape target audio without prompt target_audio_no_prompt_16khz = resample(target_audio_no_prompt, self.target_sample_rate, 16000) target_audio_no_prompt_lens_16khz = (target_audio_no_prompt_lens / self.target_sample_rate * 16000).to( torch.long ) if self.cfg.get("use_GT_transcriptions_for_metrics", True): # use target audio transcription for metrics target_asr_texts = self.asr_bleu.asr.transcribe( [ audio[:alen] for audio, alen in zip(target_audio_no_prompt_16khz, target_audio_no_prompt_lens_16khz) ], batch_size=target_audio_no_prompt_16khz.shape[0], verbose=False, ) metric_text = [asr_hyp.text for asr_hyp in target_asr_texts] else: metric_text = dataset_batch["target_texts"] asr_hyps = self.asr_bleu.update( name=name, refs=metric_text, pred_audio=metric_audio_pred, pred_audio_lens=metric_audio_pred_lens, ) self.intelligibility.update( name=name, refs=metric_text, pred_audio=metric_audio_pred, pred_audio_lens=metric_audio_pred_lens, asr_hyps=asr_hyps, ) # add ground truth intelligibility metrics self.intelligibility.update( name=name + "_gt", refs=dataset_batch["target_texts"], pred_audio=target_audio_no_prompt_16khz, pred_audio_lens=target_audio_no_prompt_lens_16khz, asr_hyps=( metric_text if self.cfg.get("use_GT_transcriptions_for_metrics", True) else None ), # reuse GT transcription ) self.secs.update( name=name, target_audio=resample(dataset_batch["target_audio"], self.target_sample_rate, 16000), target_audio_lens=(dataset_batch["target_audio_lens"] / self.target_sample_rate * 16000).to( torch.long ), pred_audio=resample(results["audio"], self.target_sample_rate, 16000), pred_audio_lens=(results["audio_len"] / self.target_sample_rate * 16000).to(torch.long), ) eou_labels = generate_multiturn_speaking_mask( next_subword_ids, bos_token_id=self.text_bos_id, eos_token_id=self.text_eos_id ) self.results_logger.update( name=name, refs=dataset_batch["target_texts"], hyps=metric_text, asr_hyps=asr_hyps, samples_id=dataset_batch['sample_id'], pred_audio=results["audio"].float(), pred_audio_tf=results["audio_tf"].float(), pre_audio_trimmed=None, reference_audio=dataset_batch["audio_prompt"].float(), target_audio=target_audio_no_prompt.float(), pred_audio_sr=self.target_sample_rate, user_audio=dataset_batch["source_audio"].float(), user_audio_sr=self.source_sample_rate, eou_pred=eou_labels, fps=self.target_fps, results=results if self.cfg.get("dump_tokens_text", False) else None, tokenizer=self.tokenizer, ) @torch.inference_mode() def validation_step(self, batch: dict, batch_idx: int): for name, dataset_batch in batch.items(): if dataset_batch is None: continue # some dataset is exhausted B = len(dataset_batch['sample_id']) # run inference for a custom speaker reference if self.cfg.get("inference_speaker_reference", None): new_dataset_batch = copy.deepcopy(dataset_batch) speaker_audio, sr = load_audio_librosa(self.cfg.inference_speaker_reference) speaker_audio = resample(speaker_audio, sr, self.target_sample_rate) speaker_audio = speaker_audio.repeat(B, 1).to(self.device) # lengths -> [B] speaker_audio_lens = torch.tensor([speaker_audio.size(1)], device=self.device).long().repeat(B) new_dataset_batch["audio_prompt"] = speaker_audio new_dataset_batch["audio_prompt_lens"] = speaker_audio_lens self.run_evaluation_one_batch(name, new_dataset_batch, use_dataloader_init=False) # run inference using dataloader speaker references else: self.run_evaluation_one_batch(name, dataset_batch, use_dataloader_init=False) def on_test_epoch_start(self) -> None: return self.on_validation_epoch_start() def on_test_epoch_end(self) -> None: return self.on_validation_epoch_end(prefix="test") def test_step(self, *args, **kwargs): return self.validation_step(*args, **kwargs) def set_audio_prompt_lantent( self, speaker_audio, speaker_audio_lens, system_prompt=None, batch_size=1, name="default_speaker", ): """ Compute and cache an audio prompt latent representation for a given speaker. The latent is stored as a registered buffer so it is saved inside checkpoints. """ # Prepare inputs self.set_init_inputs( speaker_audio=speaker_audio, speaker_audio_lens=speaker_audio_lens, system_prompt=system_prompt, ) init_inputs = self.get_init_inputs(B=batch_size) init_inputs.update( { "use_cache": True, "past_key_values": None, "guidance_enabled": False, } ) # Forward pass audio_prompt_lantent = self.tts_model(**init_inputs).audio_prompt_lantent # Detach cached = audio_prompt_lantent.detach().clone() # Store as non-trainable parameter module = nn.Parameter(cached, requires_grad=False) self.audio_prompt_latents[name] = module return cached def get_audio_prompt_lantent(self, name): """ Retrieve a cached audio prompt latent stored as a buffer. """ if name not in self.audio_prompt_latents: raise KeyError(f"Unknown audio prompt latent '{name}'. " "Call set_audio_prompt_lantent(...) first.") return self.audio_prompt_latents[name].to(self.device) def set_init_inputs(self, speaker_audio=None, speaker_audio_lens=None, system_prompt=None, speaker_name=None): """ Registers and prepares initial input buffers for text/audio prompt and context, to warm up AR inference. Args: speaker_audio (torch.Tensor): Batch of prompt audio, (B, T). speaker_audio_lens (torch.Tensor): Lengths for each sample in speaker_audio, (B,). system_prompt (str, optional): System prompt for context. speaker_name (str, optional): Speaker name. Returns: dict: Dictionary of input tensors to be passed to inference, with registered buffers. """ # compute prompt audio size and slice it with fp32_precision(): # compute the exact number of samples for the prompt duration prompt_audio_size = int( ((self.data_cfg.audio_prompt_duration * self.target_sample_rate) // self.target_samples_per_frame) * self.target_samples_per_frame ) # if a speaker name exists, use the cached latent if speaker_name is not None: speaker_audio = torch.zeros( (1, prompt_audio_size), device=self.device, dtype=torch.float32, ) speaker_audio_lens = torch.LongTensor([speaker_audio.shape[1]]).to(self.device) B, T = speaker_audio.shape device = speaker_audio.device dtype = speaker_audio.dtype # allocate result prompt_audio = torch.zeros(B, prompt_audio_size, device=device, dtype=dtype) # process each example independently for b in range(B): valid_len = min(speaker_audio_lens[b].item(), T) # handle empty if valid_len <= 0: continue # valid (non-padded) segment valid_segment = speaker_audio[b, :valid_len] if valid_len >= prompt_audio_size: # enough valid audio → crop from start (no silence) prompt_audio[b] = valid_segment[:prompt_audio_size] else: # too short → repeat and crop repeat_factor = (prompt_audio_size + valid_len - 1) // valid_len # ceil division expanded = valid_segment.repeat(repeat_factor) prompt_audio[b] = expanded[:prompt_audio_size] # add a silence in the end to smooth the transition between prompt and audio tokens prompt_audio[:, -int(self.target_samples_per_frame * 2) :] = 0 # get prompt audio size with fp32_precision(): prompt_audio_text_pad_size = int(prompt_audio_size // self.target_samples_per_frame) # create a eos token id if system_prompt is not None and self.cfg.get("use_system_prompt", None) and system_prompt != "": text_prompt = torch.as_tensor( [self.tokenizer.bos] + self.tokenizer.text_to_ids(system_prompt) + [self.tokenizer.eos], dtype=torch.long, device=self.device, ) else: text_prompt = torch.tensor([self.tokenizer.eos], dtype=torch.long, device=self.device) # create a padding tensor prompt_audio_text_pad = ( torch.ones(prompt_audio_text_pad_size, device=self.device, dtype=text_prompt.dtype) * self.text_pad_id ) prompt_audio_text_pad[-1] = self.tokenizer.eos # Prepend an initial text EOS token followed by padding tokens that match # the number of audio-prompt frames (in text-token units). target_text_tokens = torch.cat([text_prompt, prompt_audio_text_pad.to(text_prompt.dtype)]) # create pad audio for the description pad_size = text_prompt.size(-1) * self.target_samples_per_frame pad_audio = ( torch.zeros(pad_size, device=prompt_audio.device, dtype=prompt_audio.dtype) .unsqueeze(0) .repeat(prompt_audio.size(0), 1) ) # repeat to reaches the batch size target_text_tokens = target_text_tokens.unsqueeze(0).repeat(prompt_audio.size(0), 1) target_audio = torch.cat([pad_audio, prompt_audio], dim=1) # extract code codes target_audio_len = torch.tensor( [target_audio.size(-1)] * target_audio.size(0), dtype=torch.long, device=self.device ) with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad(): code, _ = self.audio_codec.encode( target_audio.unsqueeze(1).to(self.audio_codec_run_dtype), target_audio_len ) # get context hidden if self.cfg.tts_config.context_hidden_size is not None: context_hidden_state = self.embed_tokens(target_text_tokens) else: context_hidden_state = None # create masks # non_prompt_mask is all zeros, because all processed is prompt non_prompt_mask = torch.zeros_like(target_text_tokens) non_prompt_mask[:, -2:] = 1 # set last valid prompt frame as 1 to allow the addition of BOS in the right place subword_mask = torch.zeros_like( target_text_tokens ) # subword_mask is almost all zeros because on the warmup there is only the prompt subword_mask[:, -3:] = ( 1 # -3 because of the it start right after the first valid prompt token and it is shifted by 1 ) # set the pad token for the first BOS frame code[:, 0] = self.speech_pad_id # shift subword_ids subword_ids = F.pad(target_text_tokens[:, 1:], [0, 1], value=0.0) # set special token in the last audio prompt (it will works as a BOS token) pos = non_prompt_mask.float().argmax(dim=1) # shape: [B] row_idx = torch.arange(B, device=self.device) # set the extra self.speech_pad_id at first 1 position in non_prompt_mask code[row_idx, pos] = self.speech_pad_id init_inputs = { "code": code[:, :-1], "audio_mask": non_prompt_mask.bool()[ :, :-1 ], # set audio_mask as non_prompt_mask to avoid the audio prompt in loss computation "context_hidden_state": context_hidden_state[:, :-1] if context_hidden_state is not None else None, "subword_ids": subword_ids[:, :-1], "subword_mask": subword_mask.bool()[:, :-1], "non_prompt_mask": non_prompt_mask.bool()[:, :-1], } if speaker_name is not None: init_inputs["audio_prompt_lantent"] = self.get_audio_prompt_lantent(speaker_name) self._init_input_cache = {} for k, v in init_inputs.items(): if v is None: self._init_input_cache[k] = None else: # clone → removes inference tensor flag # detach → ensures no graph refs self._init_input_cache[k] = v.detach().clone() return init_inputs def get_init_inputs( self, B: int, init_inputs_names=[ "code", "audio_mask", "context_hidden_state", "subword_ids", "subword_mask", "non_prompt_mask", ], ): """ Returns a dictionary of initial inputs for inference, using registered buffers. Args: B (int): Required batch size. init_inputs_names (List[str], optional): Names of input buffers to fetch. Returns: dict: Each key is name from init_inputs_names, and value is tensor of appropriate shape (B, ...). Notes: Expands batch-1 buffers to B if necessary. """ if init_inputs_names is None: init_inputs_names = [ "code", "audio_mask", "context_hidden_state", "subword_ids", "subword_mask", "non_prompt_mask", ] # if audio_prompt_lantent available use it if self._init_input_cache.get("audio_prompt_lantent", None) is not None: init_inputs_names.append("audio_prompt_lantent") init_inputs = {} for name in init_inputs_names: buf = self._init_input_cache.get(name, None) if buf is None: init_inputs[name] = None continue # Use as-is if batch matches if buf.shape[0] == B: init_inputs[name] = buf elif buf.shape[0] >= B: init_inputs[name] = buf[:B] else: # Otherwise, assume batch=1 and expand to target B init_inputs[name] = buf[:1].expand(B, *buf.shape[1:]) return init_inputs @torch.inference_mode() def infer_codes_one_step( self, current_subword_id, prev_subword_id, current_subword_mask, prev_audio_tokens, past_key_values, guidance_enabled=True, generation_config=None, ignore_eos_flag_stop=True, ): """ Runs a single autoregressive prediction step to infer audio codec codes. Args: current_subword_id (torch.Tensor): Current text token IDs, shape (B, 1). prev_subword_id (torch.Tensor): Previous text token IDs, shape (B, 1). current_subword_mask (torch.Tensor): Current mask, shape (B, 1). prev_audio_tokens (torch.Tensor): Previously generated audio tokens, shape (B, 1, C). past_key_values: Key-value cache for transformer decoder state. guidance_enabled (bool, optional): Enables classifier-free guidance. generation_config (dict, optional): Generation hyperparameters. ignore_eos_flag_stop (bool): If True, ignore EOS flag for stopping. Returns: Tuple[torch.Tensor, Any]: - Predicted audio codec token(s), shape (B, 1, C) - Updated past_key_values for the next step. """ if self.cfg.tts_config.context_hidden_size is not None: # get context_hidden_state it is always one step behind current_subword_id # for the first step uses the last step from warmup context_hidden_state = self.embed_tokens(prev_subword_id) else: context_hidden_state = None # force silence as next token if self.cfg.get('inference_force_speech_silence_on_eos', True): silence_codes = self.codec_silence_tokens.view(1, 1, -1).expand(prev_audio_tokens.shape) prev_audio_tokens = torch.where( current_subword_id.unsqueeze(-1) == self.text_eos_id, silence_codes, # silence prev_audio_tokens, # keep original ) # get subword_ids inputs = { "code": prev_audio_tokens, "context_hidden_state": context_hidden_state, "subword_ids": current_subword_id, "subword_mask": current_subword_mask, "past_key_values": past_key_values, "use_cache": True, "guidance_enabled": guidance_enabled, "generation_config": generation_config, "ignore_eos_flag_stop": ignore_eos_flag_stop, } outputs = self.tts_model(**inputs) return outputs["codes"], outputs["past_key_values"] @torch.inference_mode() def decode_one_audio_step(self, gen_audio_codes_history, number_prev_tokens=None): """ Decodes one step of generated audio codec tokens to raw waveform. Args: gen_audio_codes_history (torch.Tensor): Audio tokens history, shape (B, T, C). number_prev_tokens (int, optional): Number of previous tokens to decode, for incremental decoding. Returns: Tuple[torch.Tensor, torch.Tensor]: - audio_pred_cur_step: Latest decoded waveform chunk, shape (B, wav_to_token_ratio). - audio_len: Lengths (number of samples), shape (B,). """ with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad(): if number_prev_tokens: gen_audio_codes_history = gen_audio_codes_history[:, -number_prev_tokens:] gen_audio_codes_history = replace_control_speech_codes( gen_audio_codes_history, self._control_codes, self.codec_silence_tokens ) gen_audio_codes_lens = torch.tensor( [gen_audio_codes_history.size(1)] * gen_audio_codes_history.size(0), device=self.device ) audio_pred, audio_len = self.audio_codec.decode(gen_audio_codes_history, gen_audio_codes_lens) # return only the current/lastest audio chunk audio_pred_cur_step = audio_pred.squeeze(1)[:, -self.audio_codec.config.wav_to_token_ratio :] audio_len[:] = self.audio_codec.config.wav_to_token_ratio return audio_pred_cur_step, audio_len @torch.inference_mode() def offline_inference( self, next_subword_ids: torch.Tensor, init_inputs: dict, task: str = "", guidance_enabled: bool = True, generation_config: dict = None, incremental_audio_decoding: bool = False, ) -> dict[str, torch.Tensor]: """ Runs offline autoregressive inference for the Duplex EAR-TTS speech decoder. This method performs **text-to-speech (TTS)** generation: given subword/text tokens and prompt-initialization states, it autoregressively generates audio codec tokens and decodes them into a waveform. Args: next_subword_ids (torch.Tensor): Conditioning subword/text token IDs for the speech decoder. Shape: (B, T_text). init_inputs (dict): Dictionary of prompt-dependent initial states produced by ``get_init_inputs()``. May include: • "code" — initial audio tokens (e.g., prompt audio) • "audio_mask" — mask for prompt audio positions • "context_hidden_state" — decoder hidden state at t = 0 • "subword_ids" — prompt text tokens • "subword_mask" — mask for prompt text • "non_prompt_mask" — mask marking positions to be generated ``get_init_inputs()`` automatically expands batch-1 buffers to batch size B. task (str, optional): Optional task identifier used to customize the prompt structure. guidance_enabled (bool, optional): Whether classifier-free guidance (CFG) is enabled. If enabled and ``generation_config`` is ``None``, guidance parameters are taken from ``_get_generation_config()``. generation_config (dict, optional): Settings controlling autoregressive generation, including sampling strategy, noise scale, refinement iterations, and EOS rules. If ``None``, defaults are taken from ``_get_generation_config(guidance_enabled)``. incremental_audio_decoding (bool, optional): If True, codec-to-waveform decoding is performed incrementally during autoregressive generation. If False, waveform decoding occurs only after all audio tokens are produced. Returns: dict[str, torch.Tensor]: Contains: • **"audio"**: Generated waveform of shape ``(B, T_audio)``, obtained via ``audio_pred.squeeze(1)``. • **"audio_len"**: Length of each generated waveform in samples, shape ``(B,)``. """ B = next_subword_ids.size(0) if generation_config is None: generation_config = self._get_generation_config(guidance_enabled) logging.info(f"Doing inference using the following config: {generation_config} !") init_inputs.update({"use_cache": True, "past_key_values": None, "guidance_enabled": guidance_enabled}) # warmup the model and generate the very first audio token outputs = self.tts_model(**init_inputs) if self.cfg.get("inference_skip_first_code_prediction_on_init", True): # use the last token on init, because we are shifthing it in the model forward, so we dont really need to compute it code = init_inputs["code"][:, -1:] else: code, _, _ = self.tts_model.generate_step(outputs.hidden_states[:, -1:], **generation_config) past_key_values = outputs["past_key_values"] # use the text tokens to stop generation max_steps = next_subword_ids.size(-1) # create variable to store the audios gen_audio_codes = torch.zeros( B, max_steps, self.tts_model.config.num_quantizers, device=self.device, dtype=torch.long ) # init subwork as all ones subword_mask = torch.ones(B, max_steps, device=self.device, dtype=torch.bool) # get first context subword_id, that is the last subword_ids from the warmup first_context_subword_id = init_inputs["subword_ids"][:, -1].unsqueeze(-1) # initialize variables used to save the output audio audio_pred = None audio_pred_len = torch.zeros(B, device=self.device, dtype=torch.long) for i in range(max_steps): step_start = time.time() # current subword id is always seem current_subword_id = next_subword_ids[:, i].unsqueeze(-1) if i == 0: prev_subword_id = first_context_subword_id else: prev_subword_id = next_subword_ids[:, i - 1].unsqueeze(-1) # create subword_mask current_subword_mask = subword_mask[:, i].unsqueeze(-1) code, past_key_values = self.infer_codes_one_step( current_subword_id=current_subword_id, prev_subword_id=prev_subword_id, current_subword_mask=current_subword_mask, prev_audio_tokens=code, past_key_values=past_key_values, guidance_enabled=guidance_enabled, generation_config=generation_config, ignore_eos_flag_stop=True, ) # cache audio tokens gen_audio_codes[:, i] = code.squeeze(1) if incremental_audio_decoding: audio_pred_i, audio_pred_i_len = self.decode_one_audio_step( gen_audio_codes[:, : i + 1], number_prev_tokens=self.cfg.get("inference_codec_decoding_prev_tokens_number", None), ) if audio_pred is None: audio_pred = audio_pred_i else: audio_pred = torch.cat([audio_pred, audio_pred_i], dim=1) audio_pred_len += audio_pred_i_len step_time = time.time() - step_start logging.info(f"Autoregressive inference step: {i} of {max_steps} take around {step_time}s") if not incremental_audio_decoding: gen_audio_codes_lens = torch.tensor( [gen_audio_codes.shape[1]] * gen_audio_codes.shape[0], dtype=torch.long ).to(self.device) # decode audio. Note that it is not necessary because the prompt is removed, so no special token should be on the output, but lets do it for safety gen_audio_codes = replace_control_speech_codes( gen_audio_codes, self._control_codes, self.codec_silence_tokens ) with ensures_target_precision(self.audio_codec_run_dtype), torch.no_grad(): audio_pred, audio_pred_len = self.audio_codec.decode(gen_audio_codes, gen_audio_codes_lens) return audio_pred.squeeze(1), audio_pred_len def backward(self, *args, **kwargs): with loss_parallel(): super().backward(*args, **kwargs) def configure_optimizers(self): # Check if we should use the custom grouping if self.cfg.get("exclude_norm_from_wd", True): return configure_optimizers_exclude_norm_from_wd(self) # Fallback to the standard NeMo behavior if the flag is False return configure_optimizers(self) @property def oomptimizer_schema(self) -> dict: """ Return a typing schema for optimal batch size calibration for various sequence lengths using OOMptimizer. """ raise NotImplementedError def configure_model(self) -> None: # TODO(pzelasko): refactor into separate module re-usable across models device_mesh = self.device_mesh if device_mesh is None: return llm = self.tts_model.backbone if isinstance(llm, PeftModel): llm = llm.base_model.model if (tp_mesh := device_mesh["tensor_parallel"]).size() > 1: self._use_tp = True plan = { "layers.0": PrepareModuleInput( input_layouts=(Replicate(),), # , None) desired_input_layouts=(Shard(1),), # , None) use_local_output=True, ), "norm": SequenceParallel(), } parallelize_module(llm, tp_mesh, plan) for transformer_block in llm.layers: plan = { "input_layernorm": SequenceParallel(), "self_attn.q_proj": ColwiseParallel(), "self_attn.k_proj": ColwiseParallel(), "self_attn.v_proj": ColwiseParallel(), "self_attn.o_proj": RowwiseParallel(output_layouts=Shard(1)), "post_attention_layernorm": SequenceParallel(), "mlp": PrepareModuleInput( input_layouts=(Shard(1),), desired_input_layouts=(Replicate(),), ), "mlp.gate_proj": ColwiseParallel(), "mlp.up_proj": ColwiseParallel(), "mlp.down_proj": RowwiseParallel(output_layouts=Shard(1)), # "pre_feedforward_layernorm": SequenceParallel(), # "post_feedforward_layernorm": SequenceParallel(), } # Adjust attention module to use the local number of heads attn_layer = transformer_block.self_attn for attr in ("num_heads", "num_key_value_heads", "hidden_size"): val = getattr(attn_layer, attr) if val % tp_mesh.size() != 0: logging.warning( f"attn_layer.{attr}={val} is not divisible by {tp_mesh.size()=}: " f"set a different tensor parallelism size to avoid errors." ) setattr(attn_layer, attr, val // tp_mesh.size()) parallelize_module(transformer_block, tp_mesh, plan) for m in ( self.tts_model.mog_head, self.tts_model.embed_subword, self.tts_model.embed_context, self.tts_model.embed_code, self.tts_model.null_emb, self.tts_model.bos_emb, self.tts_model.lm_head, ): parallelize_module( m, tp_mesh, ColwiseParallel( input_layouts=Shard(1), output_layouts=Shard(-1), use_local_output=False, ), ) if (dp_mesh := device_mesh["data_parallel"]).size() > 1: assert dp_mesh.ndim == 1 self._use_fsdp = True fsdp_config = {"mesh": dp_mesh} for idx, layer in enumerate(llm.layers): llm.layers[idx] = fully_shard(layer, **fsdp_config) for idx in range(self.tts_model._num_codebooks): self.tts_model.audio_embeddings[idx] = fully_shard(self.tts_model.audio_embeddings[idx], **fsdp_config) if self.tts_model.use_local_transformer: self.tts_model.local_transformer = fully_shard(self.tts_model.local_transformer, **fsdp_config) self.tts_model.local_transformer_in_projection = fully_shard( self.tts_model.local_transformer_in_projection, **fsdp_config ) self.embed_text_tokens = fully_shard(self.embed_text_tokens, **fsdp_config) self.tts_model.mog_head = fully_shard(self.tts_model.mog_head, **fsdp_config) self.tts_model.embed_subword = fully_shard(self.tts_model.embed_subword, **fsdp_config) self.tts_model.embed_context = fully_shard(self.tts_model.embed_context, **fsdp_config) self.tts_model.embed_code = fully_shard(self.tts_model.embed_code, **fsdp_config) self.tts_model.lm_head = fully_shard(self.tts_model.lm_head, **fsdp_config) self.llm = fully_shard(self.llm, **fsdp_config) self.tts_model = fully_shard(self.tts_model, **fsdp_config) def maybe_recreate_cached_audio_prompt_latents_structure(self, state_dict): """ Recreate audio_prompt_latents ParameterDict structure from keys present in the provided state_dict. """ for k, tensor in state_dict.items(): if "audio_prompt_latents." in k: # Extract speaker name name = k.split("audio_prompt_latents.")[1] if name not in self.audio_prompt_latents: # Create placeholder parameter with correct shape self.audio_prompt_latents[name] = nn.Parameter( torch.zeros_like(tensor), requires_grad=False, ) def load_state_dict(self, state_dict, strict: bool = True): """ Loads model weights with audio prompt latent compatibility. This method: 1. Recreates cached audio prompt latent structures if required 2. Attempts strict loading 3. Falls back to partial initialization if necessary Args: state_dict (dict): Model state dictionary. strict (bool, optional): Whether to enforce strict key matching. Returns: IncompatibleKeys: As returned by LightningModule.load_state_dict. """ # recreate audio prompt latent entries if needed self.maybe_recreate_cached_audio_prompt_latents_structure(state_dict) try: return super().load_state_dict(state_dict, strict=strict) except RuntimeError: logging.info("Error loading model state_dict !! Retrying with partial initialization!") model_dict = set_model_dict_for_partial_init(state_dict, self.state_dict()) return super().load_state_dict(model_dict, strict=False) def load_audio_librosa(path, sr=None): """ Load audio using librosa with torchaudio-like behavior. Returns: audio_tensor: torch.FloatTensor of shape [channels, time] sr: sampling rate """ # Load with librosa (preserve original sampling rate) audio, sr = librosa.load(path, sr=sr, mono=False) # Ensure shape is [channels, time] if audio.ndim == 1: # Mono: (time,) -> (1, time) audio = audio[None, :] # Convert to torch float32 (torchaudio behavior) audio_tensor = torch.from_numpy(audio).float() return audio_tensor, sr def maybe_to(x, dtype): if x is None: return None if isinstance(x, torch.Tensor) and torch.is_floating_point(x): return x.to(dtype) return x @contextmanager def ensures_target_precision(target_dtype): """ Workaround for precision related issues when training with bf16-true PyTorch Lightning precision setting. In bf16-true, PTL changes PyTorch's default dtype, which may break implicit assumptions for some models. This context manager restores default float32 precision and runs the computation in float32 autocast context. """ default_dtype = torch.get_default_dtype() torch.set_default_dtype(target_dtype) try: with torch.amp.autocast(device_type="cuda" if torch.cuda.is_available() else "cpu", dtype=target_dtype): yield finally: torch.set_default_dtype(default_dtype) def generate_multiturn_speaking_mask(input_ids: torch.Tensor, bos_token_id: int = 0, eos_token_id: int = 1): """ Efficient, batched speaking mask generator that marks 1 between and pairs. If is missing after a , mask continues to end. Handles multiple turns. Args: input_ids (torch.Tensor): LongTensor of shape (B, T) bos_token_id (int): Token ID for eos_token_id (int): Token ID for Returns: torch.Tensor: FloatTensor of shape (B, T), with 1.0 for speaking, 0.0 for silence. Note BOS is considered as speaking (1) and EOS as non speaking 0 """ device = input_ids.device bos_mask = (input_ids == bos_token_id).to(torch.int32).to(device) eos_mask = (input_ids == eos_token_id).to(torch.int32).to(device) bos_cumsum = torch.cumsum(bos_mask, dim=1) eos_cumsum = torch.cumsum(eos_mask, dim=1) speaking_mask = (bos_cumsum > eos_cumsum).to(torch.float32) return speaking_mask.long() def replace_control_speech_codes( speech_codes: torch.Tensor, control_codes: torch.Tensor, silence_tokens: torch.Tensor = None ) -> torch.Tensor: """ Replaces control codes (speech BOS, EOS, etc) in `speech_codes` with the first frame which is assumed to consist of 'valid' codes representing silence. """ if silence_tokens is not None: # Expand to [B, 1, 74] silence_tokens_expanded = silence_tokens.unsqueeze(0).unsqueeze(1).expand(speech_codes.shape[0], 1, -1) return torch.where(torch.isin(speech_codes, control_codes), silence_tokens_expanded, speech_codes) if torch.isin(speech_codes[:, :1], control_codes).any(): return torch.where( torch.isin(speech_codes, control_codes), torch.zeros_like(speech_codes[:, :1]), speech_codes ) else: return torch.where(torch.isin(speech_codes, control_codes), speech_codes[:, :1], speech_codes) def setup_audio_codec(model): """ Instantiates the RVQ audio codec and injects codec embeddings into the TTS model. This function is responsible only for: - Instantiating the codec model (`RVQVAEModel`). - Loading pretrained codec weights (if configured). - Freezing codec parameters. - Registering RVQ embeddings inside the TTS model via `set_rvq_embs`. Args: model: Model instance of DuplexEARTTS """ with ensures_target_precision(model.audio_codec_run_dtype): model.audio_codec = RVQVAEModel(DictConfig(model.cfg.codec_config)) # load pretrained codec checkpoint if model.cfg.get("pretrained_codec_model", None): checkpoint_state = load_checkpoint(model.cfg.pretrained_codec_model) checkpoint_state = set_model_dict_for_partial_init(checkpoint_state, model.audio_codec.state_dict()) model.audio_codec.load_state_dict(checkpoint_state, strict=True) for p in model.audio_codec.parameters(): p.requires_grad = False model.audio_codec.eval() model.audio_codec.to(model.device) # force codec to run in the same device as the main model assert callable(model.tts_model.set_rvq_embs) model.tts_model.set_rvq_embs(torch.stack([x.detach() for x in model.audio_codec.prvq.mus_list], 0)) model.tts_model.rvq_embs = model.tts_model.rvq_embs.to(next(model.tts_model.parameters()).dtype) # compute target fps model.target_fps = model.target_sample_rate / model.audio_codec.config.wav_to_token_ratio model.target_samples_per_frame = model.audio_codec.config.wav_to_token_ratio def rescale_state_dict(state_dict, target_std=0.02, first_n_layers=None, layer_prefix="tts_model.backbone.layers."): """ Rescale trainable weights in a state_dict for BF16/FP16 stability. Args: state_dict: PyTorch state_dict target_std: desired target std for weights first_n_layers: if not None, rescale only the first N transformer blocks layer_prefix: prefix for layer names (default: "tts_model.backbone.layers.") Returns: new_state_dict """ weight_tensors = [] # Compute which prefixes to match if first_n_layers is set prefixes_to_match = [] if first_n_layers is not None: prefixes_to_match = [f"{layer_prefix}{i}" for i in range(first_n_layers)] for name, param in state_dict.items(): if not torch.is_tensor(param): continue if "rvq_embs" in name: continue # Skip biases & 1-dim params (norm weights/gates) if param.ndim <= 1: continue # Skip layers not in the first N if first_n_layers is not None and not any(name.startswith(pfx) for pfx in prefixes_to_match): continue weight_tensors.append(param.float()) if not weight_tensors: if first_n_layers is not None: logging.info(f"No weights found for first {first_n_layers} layers with prefix '{layer_prefix}'.") else: logging.info("No weights found to rescale in state_dict.") return state_dict # Compute global std across selected weights (on CPU) cpu_weights = [p.detach().cpu() for p in weight_tensors] flat = torch.cat([p.flatten() for p in cpu_weights]) current_std = float(torch.std(flat)) scale = target_std / (current_std + 1e-8) logging.info( f"Rescaling state_dict " f"{'(first N layers)' if first_n_layers else '(all layers)'}: " f"current std = {current_std:.6f}, target = {target_std}, scale = {scale:.6f}" ) # Apply scaling new_state_dict = {} for name, param in state_dict.items(): if ( torch.is_tensor(param) and param.ndim > 1 and (first_n_layers is None or any(name.startswith(pfx) for pfx in prefixes_to_match)) ): new_state_dict[name] = param * scale else: new_state_dict[name] = param logging.info("Done: weights rescaled.") return new_state_dict