| 1 | import json |
| 2 | from importlib.resources import files |
| 3 | |
| 4 | import torch |
| 5 | import torch.nn.functional as F |
| 6 | import torchaudio |
| 7 | from datasets import Dataset as Dataset_ |
| 8 | from datasets import load_from_disk |
| 9 | from torch import nn |
| 10 | from torch.utils.data import Dataset, Sampler |
| 11 | from tqdm import tqdm |
| 12 | |
| 13 | from f5_tts.model.modules import MelSpec |
| 14 | from f5_tts.model.utils import default |
| 15 | |
| 16 | |
| 17 | class HFDataset(Dataset): |
| 18 | def __init__( |
| 19 | self, |
| 20 | hf_dataset: Dataset, |
| 21 | target_sample_rate=24_000, |
| 22 | n_mel_channels=100, |
| 23 | hop_length=256, |
| 24 | n_fft=1024, |
| 25 | win_length=1024, |
| 26 | mel_spec_type="vocos", |
| 27 | ): |
| 28 | self.data = hf_dataset |
| 29 | self.target_sample_rate = target_sample_rate |
| 30 | self.hop_length = hop_length |
| 31 | |
| 32 | self.mel_spectrogram = MelSpec( |
| 33 | n_fft=n_fft, |
| 34 | hop_length=hop_length, |
| 35 | win_length=win_length, |
| 36 | n_mel_channels=n_mel_channels, |
| 37 | target_sample_rate=target_sample_rate, |
| 38 | mel_spec_type=mel_spec_type, |
| 39 | ) |
| 40 | self._resamplers = {} |
| 41 | |
| 42 | def get_frame_len(self, index): |
| 43 | row = self.data[index] |
| 44 | audio = row["audio"]["array"] |
| 45 | sample_rate = row["audio"]["sampling_rate"] |
| 46 | return audio.shape[-1] / sample_rate * self.target_sample_rate / self.hop_length |
| 47 | |
| 48 | def __len__(self): |
| 49 | return len(self.data) |
| 50 | |
| 51 | def __getitem__(self, index): |
| 52 | row = self.data[index] |
| 53 | audio = row["audio"]["array"] |
| 54 | |
| 55 | sample_rate = row["audio"]["sampling_rate"] |
| 56 | duration = audio.shape[-1] / sample_rate |
| 57 | |
| 58 | if duration > 30 or duration < 0.3: |
| 59 | return self.__getitem__((index + 1) % len(self.data)) |
| 60 | |
| 61 | audio_tensor = torch.from_numpy(audio).float() |
| 62 | |
| 63 | if sample_rate != self.target_sample_rate: |
| 64 | if sample_rate not in self._resamplers: |
| 65 | self._resamplers[sample_rate] = torchaudio.transforms.Resample(sample_rate, self.target_sample_rate) |
| 66 | audio_tensor = self._resamplers[sample_rate](audio_tensor) |
| 67 | |
| 68 | audio_tensor = audio_tensor.unsqueeze(0) # 't -> 1 t') |
| 69 | |
| 70 | mel_spec = self.mel_spectrogram(audio_tensor) |
| 71 | |
| 72 | mel_spec = mel_spec.squeeze(0) # '1 d t -> d t' |
| 73 | |
| 74 | text = row["text"] |
| 75 | |
| 76 | return dict( |
| 77 | mel_spec=mel_spec, |
| 78 | text=text, |
| 79 | ) |
| 80 | |
| 81 | |
| 82 | class CustomDataset(Dataset): |
| 83 | def __init__( |
| 84 | self, |
| 85 | custom_dataset: Dataset, |
| 86 | durations=None, |
| 87 | target_sample_rate=24_000, |
| 88 | hop_length=256, |
| 89 | n_mel_channels=100, |
| 90 | n_fft=1024, |
| 91 | win_length=1024, |
| 92 | mel_spec_type="vocos", |
| 93 | preprocessed_mel=False, |
| 94 | mel_spec_module: nn.Module | None = None, |
| 95 | ): |
| 96 | self.data = custom_dataset |
| 97 | self.durations = durations |
| 98 | self.target_sample_rate = target_sample_rate |
| 99 | self.hop_length = hop_length |
| 100 | self.n_fft = n_fft |
| 101 | self.win_length = win_length |
| 102 | self.mel_spec_type = mel_spec_type |
| 103 | self.preprocessed_mel = preprocessed_mel |
| 104 | |
| 105 | if not preprocessed_mel: |
| 106 | self.mel_spectrogram = default( |
| 107 | mel_spec_module, |
| 108 | MelSpec( |
| 109 | n_fft=n_fft, |
| 110 | hop_length=hop_length, |
| 111 | win_length=win_length, |
| 112 | n_mel_channels=n_mel_channels, |
| 113 | target_sample_rate=target_sample_rate, |
| 114 | mel_spec_type=mel_spec_type, |
| 115 | ), |
| 116 | ) |
| 117 | self._resamplers = {} |
| 118 | |
| 119 | def get_frame_len(self, index): |
| 120 | if ( |
| 121 | self.durations is not None |
| 122 | ): # Please make sure the separately provided durations are correct, otherwise 99.99% OOM |
| 123 | return self.durations[index] * self.target_sample_rate / self.hop_length |
| 124 | return self.data[index]["duration"] * self.target_sample_rate / self.hop_length |
| 125 | |
| 126 | def __len__(self): |
| 127 | return len(self.data) |
| 128 | |
| 129 | def __getitem__(self, index): |
| 130 | while True: |
| 131 | row = self.data[index] |
| 132 | audio_path = row["audio_path"] |
| 133 | text = row["text"] |
| 134 | duration = row["duration"] |
| 135 | |
| 136 | # filter by given length |
| 137 | if 0.3 <= duration <= 30: |
| 138 | break # valid |
| 139 | |
| 140 | index = (index + 1) % len(self.data) |
| 141 | |
| 142 | if self.preprocessed_mel: |
| 143 | mel_spec = torch.tensor(row["mel_spec"]) |
| 144 | else: |
| 145 | audio, source_sample_rate = torchaudio.load(audio_path) |
| 146 | |
| 147 | # make sure mono input |
| 148 | if audio.shape[0] > 1: |
| 149 | audio = torch.mean(audio, dim=0, keepdim=True) |
| 150 | |
| 151 | # resample if necessary |
| 152 | if source_sample_rate != self.target_sample_rate: |
| 153 | if source_sample_rate not in self._resamplers: |
| 154 | self._resamplers[source_sample_rate] = torchaudio.transforms.Resample( |
| 155 | source_sample_rate, self.target_sample_rate |
| 156 | ) |
| 157 | audio = self._resamplers[source_sample_rate](audio) |
| 158 | |
| 159 | # to mel spectrogram |
| 160 | mel_spec = self.mel_spectrogram(audio) |
| 161 | mel_spec = mel_spec.squeeze(0) # '1 d t -> d t' |
| 162 | |
| 163 | return { |
| 164 | "mel_spec": mel_spec, |
| 165 | "text": text, |
| 166 | } |
| 167 | |
| 168 | |
| 169 | # Dynamic Batch Sampler |
| 170 | class DynamicBatchSampler(Sampler[list[int]]): |
| 171 | """Extension of Sampler that will do the following: |
| 172 | 1. Change the batch size (essentially number of sequences) |
| 173 | in a batch to ensure that the total number of frames are less |
| 174 | than a certain threshold. |
| 175 | 2. Make sure the padding efficiency in the batch is high. |
| 176 | 3. Shuffle batches each epoch while maintaining reproducibility. |
| 177 | """ |
| 178 | |
| 179 | def __init__( |
| 180 | self, sampler: Sampler[int], frames_threshold: int, max_samples=0, random_seed=None, drop_residual: bool = False |
| 181 | ): |
| 182 | self.sampler = sampler |
| 183 | self.frames_threshold = frames_threshold |
| 184 | self.max_samples = max_samples |
| 185 | self.random_seed = random_seed |
| 186 | self.epoch = 0 |
| 187 | |
| 188 | indices, batches = [], [] |
| 189 | data_source = self.sampler.data_source |
| 190 | |
| 191 | for idx in tqdm( |
| 192 | self.sampler, desc="Sorting with sampler... if slow, check whether dataset is provided with duration" |
| 193 | ): |
| 194 | indices.append((idx, data_source.get_frame_len(idx))) |
| 195 | indices.sort(key=lambda elem: elem[1]) |
| 196 | |
| 197 | batch = [] |
| 198 | batch_frames = 0 |
| 199 | for idx, frame_len in tqdm( |
| 200 | indices, desc=f"Creating dynamic batches with {frames_threshold} audio frames per gpu" |
| 201 | ): |
| 202 | if batch_frames + frame_len <= self.frames_threshold and (max_samples == 0 or len(batch) < max_samples): |
| 203 | batch.append(idx) |
| 204 | batch_frames += frame_len |
| 205 | else: |
| 206 | if len(batch) > 0: |
| 207 | batches.append(batch) |
| 208 | if frame_len <= self.frames_threshold: |
| 209 | batch = [idx] |
| 210 | batch_frames = frame_len |
| 211 | else: |
| 212 | batch = [] |
| 213 | batch_frames = 0 |
| 214 | |
| 215 | if not drop_residual and len(batch) > 0: |
| 216 | batches.append(batch) |
| 217 | |
| 218 | del indices |
| 219 | self.batches = batches |
| 220 | |
| 221 | # Ensure even batches with accelerate BatchSamplerShard cls under frame_per_batch setting |
| 222 | self.drop_last = True |
| 223 | |
| 224 | def set_epoch(self, epoch: int) -> None: |
| 225 | """Sets the epoch for this sampler.""" |
| 226 | self.epoch = epoch |
| 227 | |
| 228 | def __iter__(self): |
| 229 | # Use both random_seed and epoch for deterministic but different shuffling per epoch |
| 230 | if self.random_seed is not None: |
| 231 | g = torch.Generator() |
| 232 | g.manual_seed(self.random_seed + self.epoch) |
| 233 | # Use PyTorch's random permutation for better reproducibility across PyTorch versions |
| 234 | indices = torch.randperm(len(self.batches), generator=g).tolist() |
| 235 | batches = [self.batches[i] for i in indices] |
| 236 | else: |
| 237 | batches = self.batches |
| 238 | return iter(batches) |
| 239 | |
| 240 | def __len__(self): |
| 241 | return len(self.batches) |
| 242 | |
| 243 | |
| 244 | # Load dataset |
| 245 | |
| 246 | |
| 247 | def load_dataset( |
| 248 | dataset_name: str, |
| 249 | tokenizer: str = "pinyin", |
| 250 | dataset_type: str = "CustomDataset", |
| 251 | audio_type: str = "raw", |
| 252 | mel_spec_module: nn.Module | None = None, |
| 253 | mel_spec_kwargs: dict = dict(), |
| 254 | ) -> CustomDataset | HFDataset: |
| 255 | """ |
| 256 | dataset_type - "CustomDataset" if you want to use tokenizer name and default data path to load for train_dataset |
| 257 | - "CustomDatasetPath" if you just want to pass the full path to a preprocessed dataset without relying on tokenizer |
| 258 | """ |
| 259 | |
| 260 | print("Loading dataset ...") |
| 261 | |
| 262 | if dataset_type == "CustomDataset": |
| 263 | rel_data_path = str(files("f5_tts").joinpath(f"../../data/{dataset_name}_{tokenizer}")) |
| 264 | if audio_type == "raw": |
| 265 | try: |
| 266 | train_dataset = load_from_disk(f"{rel_data_path}/raw") |
| 267 | except: # noqa: E722 |
| 268 | train_dataset = Dataset_.from_file(f"{rel_data_path}/raw.arrow") |
| 269 | preprocessed_mel = False |
| 270 | elif audio_type == "mel": |
| 271 | train_dataset = Dataset_.from_file(f"{rel_data_path}/mel.arrow") |
| 272 | preprocessed_mel = True |
| 273 | with open(f"{rel_data_path}/duration.json", "r", encoding="utf-8") as f: |
| 274 | data_dict = json.load(f) |
| 275 | durations = data_dict["duration"] |
| 276 | train_dataset = CustomDataset( |
| 277 | train_dataset, |
| 278 | durations=durations, |
| 279 | preprocessed_mel=preprocessed_mel, |
| 280 | mel_spec_module=mel_spec_module, |
| 281 | **mel_spec_kwargs, |
| 282 | ) |
| 283 | |
| 284 | elif dataset_type == "CustomDatasetPath": |
| 285 | try: |
| 286 | train_dataset = load_from_disk(f"{dataset_name}/raw") |
| 287 | except: # noqa: E722 |
| 288 | train_dataset = Dataset_.from_file(f"{dataset_name}/raw.arrow") |
| 289 | |
| 290 | with open(f"{dataset_name}/duration.json", "r", encoding="utf-8") as f: |
| 291 | data_dict = json.load(f) |
| 292 | durations = data_dict["duration"] |
| 293 | train_dataset = CustomDataset( |
| 294 | train_dataset, durations=durations, preprocessed_mel=preprocessed_mel, **mel_spec_kwargs |
| 295 | ) |
| 296 | |
| 297 | elif dataset_type == "HFDataset": |
| 298 | print( |
| 299 | "Should manually modify the path of huggingface dataset to your need.\n" |
| 300 | + "May also the corresponding script cuz different dataset may have different format." |
| 301 | ) |
| 302 | pre, post = dataset_name.split("_") |
| 303 | train_dataset = HFDataset( |
| 304 | load_dataset(f"{pre}/{pre}", split=f"train.{post}", cache_dir=str(files("f5_tts").joinpath("../../data"))), |
| 305 | ) |
| 306 | |
| 307 | return train_dataset |
| 308 | |
| 309 | |
| 310 | # collation |
| 311 | |
| 312 | |
| 313 | def collate_fn(batch): |
| 314 | mel_specs = [item["mel_spec"].squeeze(0) for item in batch] |
| 315 | mel_lengths = torch.LongTensor([spec.shape[-1] for spec in mel_specs]) |
| 316 | max_mel_length = mel_lengths.amax() |
| 317 | |
| 318 | padded_mel_specs = [] |
| 319 | for spec in mel_specs: |
| 320 | padding = (0, max_mel_length - spec.size(-1)) |
| 321 | padded_spec = F.pad(spec, padding, value=0) |
| 322 | padded_mel_specs.append(padded_spec) |
| 323 | |
| 324 | mel_specs = torch.stack(padded_mel_specs) |
| 325 | |
| 326 | text = [item["text"] for item in batch] |
| 327 | text_lengths = torch.LongTensor([len(item) for item in text]) |
| 328 | |
| 329 | return dict( |
| 330 | mel=mel_specs, |
| 331 | mel_lengths=mel_lengths, # records for padding mask |
| 332 | text=text, |
| 333 | text_lengths=text_lengths, |
| 334 | ) |
| 335 |