wehub-resource-sync
ba4be087d5
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
CICD NeMo / cicd-test-container-build (push) Has been cancelled
CICD NeMo / cicd-import-tests (push) Has been cancelled
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
CICD NeMo / cicd-main-speech (push) Has been cancelled
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
CICD NeMo / Coverage (e2e) (push) Has been cancelled
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
1768 lines
74 KiB
Python
1768 lines
74 KiB
Python
# 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 <bos> and <eos> pairs.
|
|
If <eos> is missing after a <bos>, mask continues to end. Handles multiple turns.
|
|
|
|
Args:
|
|
input_ids (torch.Tensor): LongTensor of shape (B, T)
|
|
bos_token_id (int): Token ID for <bos>
|
|
eos_token_id (int): Token ID for <eos>
|
|
|
|
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
|