CAMB-AI/MARS5-TTS

1import torch2import logging3 4def length_to_mask(length, offsets, max_len=None):5    """6    Convert tensor of lengths into a mask.7 8    Args:9        length (Tensor): a tensor of lengths, shape = (batch_size,)10        offsets (Tensor): a tensor of offsets, shape = (batch_size,)11        max_len (int, optional): maximum length to be considered12 13    Returns:14        mask (Tensor): a mask tensor, shape = (batch_size, max_len), 15                        True in masked positions, False otherwise.16    """17    # get the batch size18    batch_size = length.size(0)19    20    # if maximum length is not provided, then compute it from the 'length' tensor.21    if max_len is None:22        max_len = length.max().item()23    24    # Create a tensor of size `(batch_size, max_len)` filled with `True`.25    mask = torch.ones(size=(batch_size, max_len), dtype=torch.bool, device=length.device)26    27    # Create a tensor with consecutive numbers.28    range_tensor = torch.arange(max_len, device=length.device)29    30    # Expand the dim of 'length' tensor and 'offset' tensor to make it `(batch_size, max_len)`.31    # The added dimension will be used for broadcasting.32    length_exp = length.unsqueeze(-1)33    offsets_exp = offsets.unsqueeze(-1)34    35    # Create a boolean mask where `False` represents valid positions and `True` represents padding.36    mask = (range_tensor < offsets_exp) | (~(range_tensor < length_exp))37 38    return mask39 40 41def construct_padding_mask(input_tensor, pad_token):42    return (input_tensor == pad_token).cumsum(dim=1) > 0    43 44 45def nuke_weight_norm(module):46    """47    Recursively remove weight normalization from a module and its children.48 49    Args:50        module (torch.nn.Module): The module from which to remove weight normalization.51    """52    # Remove weight norm from current module if it exists53    try:54        torch.nn.utils.remove_weight_norm(module)55        logging.debug(f"Removed weight norm from {module.__class__.__name__}")56    except ValueError:57        # Ignore if the module does not have weight norm applied.58        pass59 60    # Recursively call the function on children modules61    for child in module.children():62        nuke_weight_norm(child)63