|
from dataclasses import dataclass |
|
from pathlib import Path |
|
from typing import Optional |
|
|
|
import librosa |
|
import numpy as np |
|
import torch |
|
from lightning import LightningDataModule |
|
from torch.utils.data import DataLoader, Dataset |
|
|
|
from fish_speech.utils import RankedLogger |
|
|
|
logger = RankedLogger(__name__, rank_zero_only=False) |
|
|
|
|
|
class VQGANDataset(Dataset): |
|
def __init__( |
|
self, |
|
filelist: str, |
|
sample_rate: int = 32000, |
|
hop_length: int = 640, |
|
slice_frames: Optional[int] = None, |
|
): |
|
super().__init__() |
|
|
|
filelist = Path(filelist) |
|
root = filelist.parent |
|
|
|
self.files = [ |
|
root / line.strip() |
|
for line in filelist.read_text(encoding="utf-8").splitlines() |
|
if line.strip() |
|
] |
|
self.sample_rate = sample_rate |
|
self.hop_length = hop_length |
|
self.slice_frames = slice_frames |
|
|
|
def __len__(self): |
|
return len(self.files) |
|
|
|
def get_item(self, idx): |
|
file = self.files[idx] |
|
|
|
audio, _ = librosa.load(file, sr=self.sample_rate, mono=True) |
|
|
|
|
|
if ( |
|
self.slice_frames is not None |
|
and audio.shape[0] > self.slice_frames * self.hop_length |
|
): |
|
start = np.random.randint( |
|
0, audio.shape[0] - self.slice_frames * self.hop_length |
|
) |
|
audio = audio[start : start + self.slice_frames * self.hop_length] |
|
|
|
if len(audio) == 0: |
|
return None |
|
|
|
max_value = np.abs(audio).max() |
|
if max_value > 1.0: |
|
audio = audio / max_value |
|
|
|
return { |
|
"audio": torch.from_numpy(audio), |
|
} |
|
|
|
def __getitem__(self, idx): |
|
try: |
|
return self.get_item(idx) |
|
except Exception as e: |
|
import traceback |
|
|
|
traceback.print_exc() |
|
logger.error(f"Error loading {self.files[idx]}: {e}") |
|
return None |
|
|
|
|
|
@dataclass |
|
class VQGANCollator: |
|
def __call__(self, batch): |
|
batch = [x for x in batch if x is not None] |
|
|
|
audio_lengths = torch.tensor([len(x["audio"]) for x in batch]) |
|
audio_maxlen = audio_lengths.max() |
|
|
|
|
|
audios = [] |
|
for x in batch: |
|
audios.append( |
|
torch.nn.functional.pad(x["audio"], (0, audio_maxlen - len(x["audio"]))) |
|
) |
|
|
|
return { |
|
"audios": torch.stack(audios), |
|
"audio_lengths": audio_lengths, |
|
} |
|
|
|
|
|
class VQGANDataModule(LightningDataModule): |
|
def __init__( |
|
self, |
|
train_dataset: VQGANDataset, |
|
val_dataset: VQGANDataset, |
|
batch_size: int = 32, |
|
num_workers: int = 4, |
|
val_batch_size: Optional[int] = None, |
|
): |
|
super().__init__() |
|
|
|
self.train_dataset = train_dataset |
|
self.val_dataset = val_dataset |
|
self.batch_size = batch_size |
|
self.val_batch_size = val_batch_size or batch_size |
|
self.num_workers = num_workers |
|
|
|
def train_dataloader(self): |
|
return DataLoader( |
|
self.train_dataset, |
|
batch_size=self.batch_size, |
|
collate_fn=VQGANCollator(), |
|
num_workers=self.num_workers, |
|
shuffle=True, |
|
persistent_workers=True, |
|
) |
|
|
|
def val_dataloader(self): |
|
return DataLoader( |
|
self.val_dataset, |
|
batch_size=self.val_batch_size, |
|
collate_fn=VQGANCollator(), |
|
num_workers=self.num_workers, |
|
persistent_workers=True, |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
dataset = VQGANDataset("data/LibriTTS_R/vq_train_filelist.txt") |
|
dataloader = DataLoader( |
|
dataset, batch_size=4, shuffle=False, collate_fn=VQGANCollator() |
|
) |
|
|
|
for batch in dataloader: |
|
print(batch["audios"].shape) |
|
print(batch["features"].shape) |
|
print(batch["audio_lengths"]) |
|
print(batch["feature_lengths"]) |
|
break |
|
|
