| 1 | # ruff: noqa: F722 F821 |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import os |
| 6 | import random |
| 7 | from collections import defaultdict |
| 8 | from importlib.resources import files |
| 9 | |
| 10 | import rjieba |
| 11 | import torch |
| 12 | from pypinyin import Style, lazy_pinyin |
| 13 | from torch.nn.utils.rnn import pad_sequence |
| 14 | |
| 15 | |
| 16 | # seed everything |
| 17 | |
| 18 | |
| 19 | def seed_everything(seed=0): |
| 20 | random.seed(seed) |
| 21 | os.environ["PYTHONHASHSEED"] = str(seed) |
| 22 | torch.manual_seed(seed) |
| 23 | torch.cuda.manual_seed(seed) |
| 24 | torch.cuda.manual_seed_all(seed) |
| 25 | torch.backends.cudnn.deterministic = True |
| 26 | torch.backends.cudnn.benchmark = False |
| 27 | |
| 28 | |
| 29 | # helpers |
| 30 | |
| 31 | |
| 32 | def exists(v): |
| 33 | return v is not None |
| 34 | |
| 35 | |
| 36 | def default(v, d): |
| 37 | return v if exists(v) else d |
| 38 | |
| 39 | |
| 40 | def is_package_available(package_name: str) -> bool: |
| 41 | try: |
| 42 | import importlib |
| 43 | |
| 44 | package_exists = importlib.util.find_spec(package_name) is not None |
| 45 | return package_exists |
| 46 | except Exception: |
| 47 | return False |
| 48 | |
| 49 | |
| 50 | # tensor helpers |
| 51 | |
| 52 | |
| 53 | def lens_to_mask(t: int["b"], length: int | None = None) -> bool["b n"]: |
| 54 | if not exists(length): |
| 55 | length = t.amax() |
| 56 | |
| 57 | seq = torch.arange(length, device=t.device) |
| 58 | return seq[None, :] < t[:, None] |
| 59 | |
| 60 | |
| 61 | def mask_from_start_end_indices(seq_len: int["b"], start: int["b"], end: int["b"]): |
| 62 | max_seq_len = seq_len.max().item() |
| 63 | seq = torch.arange(max_seq_len, device=start.device).long() |
| 64 | start_mask = seq[None, :] >= start[:, None] |
| 65 | end_mask = seq[None, :] < end[:, None] |
| 66 | return start_mask & end_mask |
| 67 | |
| 68 | |
| 69 | def mask_from_frac_lengths(seq_len: int["b"], frac_lengths: float["b"]): |
| 70 | lengths = (frac_lengths * seq_len).long() |
| 71 | max_start = seq_len - lengths |
| 72 | |
| 73 | rand = torch.rand_like(frac_lengths) |
| 74 | start = (max_start * rand).long().clamp(min=0) |
| 75 | end = start + lengths |
| 76 | |
| 77 | return mask_from_start_end_indices(seq_len, start, end) |
| 78 | |
| 79 | |
| 80 | def maybe_masked_mean(t: float["b n d"], mask: bool["b n"] = None) -> float["b d"]: |
| 81 | if not exists(mask): |
| 82 | return t.mean(dim=1) |
| 83 | |
| 84 | t = torch.where(mask[:, :, None], t, torch.tensor(0.0, device=t.device)) |
| 85 | num = t.sum(dim=1) |
| 86 | den = mask.float().sum(dim=1) |
| 87 | |
| 88 | return num / den.clamp(min=1.0) |
| 89 | |
| 90 | |
| 91 | # simple utf-8 tokenizer, since paper went character based |
| 92 | def list_str_to_tensor(text: list[str], padding_value=-1) -> int["b nt"]: |
| 93 | list_tensors = [torch.tensor([*bytes(t, "UTF-8")]) for t in text] # ByT5 style |
| 94 | text = pad_sequence(list_tensors, padding_value=padding_value, batch_first=True) |
| 95 | return text |
| 96 | |
| 97 | |
| 98 | # char tokenizer, based on custom dataset's extracted .txt file |
| 99 | def list_str_to_idx( |
| 100 | text: list[str] | list[list[str]], |
| 101 | vocab_char_map: dict[str, int], # {char: idx} |
| 102 | padding_value=-1, |
| 103 | ) -> int["b nt"]: |
| 104 | list_idx_tensors = [torch.tensor([vocab_char_map.get(c, 0) for c in t]) for t in text] # pinyin or char style |
| 105 | text = pad_sequence(list_idx_tensors, padding_value=padding_value, batch_first=True) |
| 106 | return text |
| 107 | |
| 108 | |
| 109 | # Get tokenizer |
| 110 | |
| 111 | |
| 112 | def get_tokenizer(dataset_name, tokenizer: str = "pinyin"): |
| 113 | """ |
| 114 | tokenizer - "pinyin" do g2p for only chinese characters, need .txt vocab_file |
| 115 | - "char" for char-wise tokenizer, need .txt vocab_file |
| 116 | - "byte" for utf-8 tokenizer |
| 117 | - "custom" if you're directly passing in a path to the vocab.txt you want to use |
| 118 | vocab_size - if use "pinyin", all available pinyin types, common alphabets (also those with accent) and symbols |
| 119 | - if use "char", derived from unfiltered character & symbol counts of custom dataset |
| 120 | - if use "byte", set to 256 (unicode byte range) |
| 121 | """ |
| 122 | if tokenizer in ["pinyin", "char"]: |
| 123 | tokenizer_path = os.path.join(files("f5_tts").joinpath("../../data"), f"{dataset_name}_{tokenizer}/vocab.txt") |
| 124 | with open(tokenizer_path, "r", encoding="utf-8") as f: |
| 125 | vocab_char_map = {} |
| 126 | for i, char in enumerate(f): |
| 127 | vocab_char_map[char[:-1]] = i |
| 128 | vocab_size = len(vocab_char_map) |
| 129 | assert vocab_char_map[" "] == 0, "make sure space is of idx 0 in vocab.txt, cuz 0 is used for unknown char" |
| 130 | |
| 131 | elif tokenizer == "byte": |
| 132 | vocab_char_map = None |
| 133 | vocab_size = 256 |
| 134 | |
| 135 | elif tokenizer == "custom": |
| 136 | with open(dataset_name, "r", encoding="utf-8") as f: |
| 137 | vocab_char_map = {} |
| 138 | for i, char in enumerate(f): |
| 139 | vocab_char_map[char[:-1]] = i |
| 140 | vocab_size = len(vocab_char_map) |
| 141 | |
| 142 | return vocab_char_map, vocab_size |
| 143 | |
| 144 | |
| 145 | # convert char to pinyin |
| 146 | |
| 147 | |
| 148 | def convert_char_to_pinyin(text_list, polyphone=True): |
| 149 | final_text_list = [] |
| 150 | custom_trans = str.maketrans( |
| 151 | {";": ",", "“": '"', "”": '"', "‘": "'", "’": "'"} |
| 152 | ) # add custom trans here, to address oov |
| 153 | |
| 154 | def is_chinese(c): |
| 155 | return ( |
| 156 | "\u3100" <= c <= "\u9fff" # common chinese characters |
| 157 | ) |
| 158 | |
| 159 | for text in text_list: |
| 160 | char_list = [] |
| 161 | text = text.translate(custom_trans) |
| 162 | for seg in rjieba.cut(text): |
| 163 | seg_byte_len = len(bytes(seg, "UTF-8")) |
| 164 | if seg_byte_len == len(seg): # if pure alphabets and symbols |
| 165 | if char_list and seg_byte_len > 1 and char_list[-1] not in " :'\"": |
| 166 | char_list.append(" ") |
| 167 | char_list.extend(seg) |
| 168 | elif polyphone and seg_byte_len == 3 * len(seg): # if pure east asian characters |
| 169 | seg_ = lazy_pinyin(seg, style=Style.TONE3, tone_sandhi=True) |
| 170 | for i, c in enumerate(seg): |
| 171 | if is_chinese(c): |
| 172 | char_list.append(" ") |
| 173 | char_list.append(seg_[i]) |
| 174 | else: # if mixed characters, alphabets and symbols |
| 175 | for c in seg: |
| 176 | if ord(c) < 256: |
| 177 | char_list.extend(c) |
| 178 | elif is_chinese(c): |
| 179 | char_list.append(" ") |
| 180 | char_list.extend(lazy_pinyin(c, style=Style.TONE3, tone_sandhi=True)) |
| 181 | else: |
| 182 | char_list.append(c) |
| 183 | final_text_list.append(char_list) |
| 184 | |
| 185 | return final_text_list |
| 186 | |
| 187 | |
| 188 | # filter func for dirty data with many repetitions |
| 189 | |
| 190 | |
| 191 | def repetition_found(text, length=2, tolerance=10): |
| 192 | pattern_count = defaultdict(int) |
| 193 | for i in range(len(text) - length + 1): |
| 194 | pattern = text[i : i + length] |
| 195 | pattern_count[pattern] += 1 |
| 196 | for pattern, count in pattern_count.items(): |
| 197 | if count > tolerance: |
| 198 | return True |
| 199 | return False |
| 200 | |
| 201 | |
| 202 | # get the empirically pruned step for sampling |
| 203 | |
| 204 | |
| 205 | def get_epss_timesteps(n, device, dtype): |
| 206 | dt = 1 / 32 |
| 207 | predefined_timesteps = { |
| 208 | 5: [0, 2, 4, 8, 16, 32], |
| 209 | 6: [0, 2, 4, 6, 8, 16, 32], |
| 210 | 7: [0, 2, 4, 6, 8, 16, 24, 32], |
| 211 | 10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32], |
| 212 | 12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32], |
| 213 | 16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32], |
| 214 | } |
| 215 | t = predefined_timesteps.get(n, []) |
| 216 | if not t: |
| 217 | return torch.linspace(0, 1, n + 1, device=device, dtype=dtype) |
| 218 | return dt * torch.tensor(t, device=device, dtype=dtype) |
| 219 |