|
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
|
|
|
