dskill/DiffRhythm

1from __future__ import annotations2 3import os4import random5from collections import defaultdict6from importlib.resources import files7 8import torch9from torch.nn.utils.rnn import pad_sequence10 11 12# seed everything13 14 15def seed_everything(seed=0):16    random.seed(seed)17    os.environ["PYTHONHASHSEED"] = str(seed)18    torch.manual_seed(seed)19    torch.cuda.manual_seed(seed)20    torch.cuda.manual_seed_all(seed)21    torch.backends.cudnn.deterministic = True22    torch.backends.cudnn.benchmark = False23 24 25# helpers26 27 28def exists(v):29    return v is not None30 31 32def default(v, d):33    return v if exists(v) else d34 35 36# tensor helpers37 38 39def lens_to_mask(t: int["b"], length: int | None = None) -> bool["b n"]:  # noqa: F722 F82140    if not exists(length):41        length = t.amax()42 43    seq = torch.arange(length, device=t.device)44    return seq[None, :] < t[:, None]45 46 47def mask_from_start_end_indices(seq_len: int["b"], start: int["b"], end: int["b"]):  # noqa: F722 F82148    max_seq_len = 204849    seq = torch.arange(max_seq_len, device=start.device).long()50    start_mask = seq[None, :] >= start[:, None]51    end_mask = seq[None, :] < end[:, None]52    return start_mask & end_mask53 54 55def mask_from_frac_lengths(seq_len: int["b"], frac_lengths: float["b"]):  # noqa: F722 F82156    lengths = (frac_lengths * seq_len).long()57    max_start = seq_len - lengths58 59    rand = torch.rand_like(frac_lengths)60    start = (max_start * rand).long().clamp(min=0)61    end = start + lengths62 63    return mask_from_start_end_indices(seq_len, start, end)64 65 66def maybe_masked_mean(t: float["b n d"], mask: bool["b n"] = None) -> float["b d"]:  # noqa: F72267    if not exists(mask):68        return t.mean(dim=1)69 70    t = torch.where(mask[:, :, None], t, torch.tensor(0.0, device=t.device))71    num = t.sum(dim=1)72    den = mask.float().sum(dim=1)73 74    return num / den.clamp(min=1.0)75 76 77# simple utf-8 tokenizer, since paper went character based78def list_str_to_tensor(text: list[str], padding_value=-1) -> int["b nt"]:  # noqa: F72279    list_tensors = [torch.tensor([*bytes(t, "UTF-8")]) for t in text]  # ByT5 style80    text = pad_sequence(list_tensors, padding_value=padding_value, batch_first=True)81    return text82 83 84# char tokenizer, based on custom dataset's extracted .txt file85def list_str_to_idx(86    text: list[str] | list[list[str]],87    vocab_char_map: dict[str, int],  # {char: idx}88    padding_value=-1,89) -> int["b nt"]:  # noqa: F72290    list_idx_tensors = [torch.tensor([vocab_char_map.get(c, 0) for c in t]) for t in text]  # pinyin or char style91    text = pad_sequence(list_idx_tensors, padding_value=padding_value, batch_first=True)92    return text93 94 95# Get tokenizer96 97 98def get_tokenizer(dataset_name, tokenizer: str = "pinyin"):99    """100    tokenizer   - "pinyin" do g2p for only chinese characters, need .txt vocab_file101                - "char" for char-wise tokenizer, need .txt vocab_file102                - "byte" for utf-8 tokenizer103                - "custom" if you're directly passing in a path to the vocab.txt you want to use104    vocab_size  - if use "pinyin", all available pinyin types, common alphabets (also those with accent) and symbols105                - if use "char", derived from unfiltered character & symbol counts of custom dataset106                - if use "byte", set to 256 (unicode byte range)107    """108    if tokenizer in ["pinyin", "char"]:109        tokenizer_path = os.path.join(files("diffrhythm").joinpath("../../data"), f"{dataset_name}_{tokenizer}/vocab.txt")110        with open(tokenizer_path, "r", encoding="utf-8") as f:111            vocab_char_map = {}112            for i, char in enumerate(f):113                vocab_char_map[char[:-1]] = i114        vocab_size = len(vocab_char_map)115        assert vocab_char_map[" "] == 0, "make sure space is of idx 0 in vocab.txt, cuz 0 is used for unknown char"116 117    elif tokenizer == "byte":118        vocab_char_map = None119        vocab_size = 256120 121    elif tokenizer == "custom":122        with open(dataset_name, "r", encoding="utf-8") as f:123            vocab_char_map = {}124            for i, char in enumerate(f):125                vocab_char_map[char[:-1]] = i126        vocab_size = len(vocab_char_map)127 128    return vocab_char_map, vocab_size129 130 131# convert char to pinyin132 133 134def convert_char_to_pinyin(text_list, polyphone=True):135    final_text_list = []136    god_knows_why_en_testset_contains_zh_quote = str.maketrans(137        {"“": '"', "”": '"', "‘": "'", "’": "'"}138    )  # in case librispeech (orig no-pc) test-clean139    custom_trans = str.maketrans({";": ","})  # add custom trans here, to address oov140    for text in text_list:141        char_list = []142        text = text.translate(god_knows_why_en_testset_contains_zh_quote)143        text = text.translate(custom_trans)144        for seg in jieba.cut(text):145            seg_byte_len = len(bytes(seg, "UTF-8"))146            if seg_byte_len == len(seg):  # if pure alphabets and symbols147                if char_list and seg_byte_len > 1 and char_list[-1] not in " :'\"":148                    char_list.append(" ")149                char_list.extend(seg)150            elif polyphone and seg_byte_len == 3 * len(seg):  # if pure chinese characters151                seg = lazy_pinyin(seg, style=Style.TONE3, tone_sandhi=True)152                for c in seg:153                    if c not in "。，、；：？！《》【】—…":154                        char_list.append(" ")155                    char_list.append(c)156            else:  # if mixed chinese characters, alphabets and symbols157                for c in seg:158                    if ord(c) < 256:159                        char_list.extend(c)160                    else:161                        if c not in "。，、；：？！《》【】—…":162                            char_list.append(" ")163                            char_list.extend(lazy_pinyin(c, style=Style.TONE3, tone_sandhi=True))164                        else:  # if is zh punc165                            char_list.append(c)166        final_text_list.append(char_list)167 168    return final_text_list169 170 171# filter func for dirty data with many repetitions172 173 174def repetition_found(text, length=2, tolerance=10):175    pattern_count = defaultdict(int)176    for i in range(len(text) - length + 1):177        pattern = text[i : i + length]178        pattern_count[pattern] += 1179    for pattern, count in pattern_count.items():180        if count > tolerance:181            return True182    return False183