447 lines
20 KiB
Python
447 lines
20 KiB
Python
import math
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
import warnings
|
|
import random
|
|
|
|
try:
|
|
import librosa # type: ignore
|
|
except Exception: # pragma: no cover
|
|
librosa = None # Fallback: user must install librosa when using local audio paths
|
|
|
|
try:
|
|
import resampy # type: ignore
|
|
except Exception: # pragma: no cover
|
|
resampy = None
|
|
|
|
|
|
def _resample_if_needed(wav: np.ndarray, orig_sr: int, target_sr: int) -> np.ndarray:
|
|
if orig_sr == target_sr:
|
|
return wav.astype(np.float32, copy=False)
|
|
if resampy is not None:
|
|
return resampy.resample(wav.astype(np.float32), orig_sr, target_sr)
|
|
if librosa is not None:
|
|
return librosa.resample(y=wav.astype(np.float32), orig_sr=orig_sr, target_sr=target_sr)
|
|
warnings.warn(
|
|
"No resampler available; treating audio as target_sr without resampling. Install resampy or librosa.",
|
|
RuntimeWarning,
|
|
)
|
|
return wav.astype(np.float32, copy=False)
|
|
|
|
|
|
# Lightweight HF-style dataset wrapper (optional). Trainer can also pass raw HF datasets directly.
|
|
class VibeVoiceDataset:
|
|
def __init__(
|
|
self,
|
|
dataset: Any,
|
|
text_column: str = "text",
|
|
audio_column: str = "audio",
|
|
voice_prompts_column: Optional[str] = "voice_prompts",
|
|
) -> None:
|
|
self.dataset = dataset
|
|
self.text_column = text_column
|
|
self.audio_column = audio_column
|
|
self.voice_prompts_column = voice_prompts_column
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.dataset)
|
|
|
|
def __getitem__(self, idx: int) -> Dict[str, Any]:
|
|
item = self.dataset[idx]
|
|
data: Dict[str, Any] = {}
|
|
data["text"] = item[self.text_column]
|
|
data["audio"] = item[self.audio_column]
|
|
|
|
user_provided_prompt = None
|
|
if self.voice_prompts_column and self.voice_prompts_column in item:
|
|
user_provided_prompt = item[self.voice_prompts_column]
|
|
|
|
if user_provided_prompt:
|
|
# A prompt was provided in the dataset, so we use it.
|
|
if not isinstance(user_provided_prompt, list):
|
|
data["voice_prompts"] = [user_provided_prompt]
|
|
else:
|
|
data["voice_prompts"] = user_provided_prompt
|
|
else:
|
|
# FALLBACK: No prompt provided, so we auto-generate one from the target audio.
|
|
try:
|
|
target_sr = 24000
|
|
wav_array = _load_audio_to_24k(item[self.audio_column], target_sr=target_sr)
|
|
audio_len_seconds = len(wav_array) / target_sr
|
|
|
|
min_len_sec = min(5.0, audio_len_seconds / 4.0)
|
|
max_len_sec = min(15.0, audio_len_seconds / 2.0)
|
|
|
|
if min_len_sec > max_len_sec:
|
|
min_len_sec = max_len_sec
|
|
max_len_sec = min(max_len_sec, audio_len_seconds)
|
|
|
|
if max_len_sec > 0.1:
|
|
prompt_len_sec = random.uniform(min_len_sec, max_len_sec)
|
|
prompt_len_samples = int(prompt_len_sec * target_sr)
|
|
|
|
max_start_sample = len(wav_array) - prompt_len_samples
|
|
start_sample = random.randint(0, max_start_sample)
|
|
|
|
prompt_crop = wav_array[start_sample : start_sample + prompt_len_samples]
|
|
|
|
data["voice_prompts"] = [prompt_crop]
|
|
else:
|
|
data["voice_prompts"] = None
|
|
|
|
except Exception as e:
|
|
warnings.warn(f"Could not create voice prompt for item {idx}: {e}")
|
|
data["voice_prompts"] = None
|
|
return data
|
|
|
|
|
|
|
|
def _apply_silence_with_crossfade(
|
|
wav: np.ndarray,
|
|
*,
|
|
sample_rate: int,
|
|
pre_silence_sec: float = 0.25,
|
|
pre_crossfade_sec: float = 0.25,
|
|
post_crossfade_sec: float = 0.25,
|
|
post_silence_sec: float = 0.75,
|
|
) -> np.ndarray:
|
|
"""Pad audio with leading/trailing silence and apply crossfades.
|
|
|
|
Structure: [pre_silence][pre_crossfade][audio_body][post_crossfade][post_silence]
|
|
Crossfades blend the audio with silence linearly to avoid hard edges.
|
|
"""
|
|
|
|
wav = np.asarray(wav, dtype=np.float32).reshape(-1)
|
|
|
|
start_sil_samples = int(round(pre_silence_sec * sample_rate))
|
|
end_sil_samples = int(round(post_silence_sec * sample_rate))
|
|
pre_crossfade_samples = int(round(pre_crossfade_sec * sample_rate))
|
|
post_crossfade_samples = int(round(post_crossfade_sec * sample_rate))
|
|
|
|
total_len = wav.shape[0]
|
|
if total_len == 0:
|
|
pieces: List[np.ndarray] = []
|
|
if start_sil_samples > 0:
|
|
pieces.append(np.zeros(start_sil_samples, dtype=np.float32))
|
|
if end_sil_samples > 0:
|
|
pieces.append(np.zeros(end_sil_samples, dtype=np.float32))
|
|
return np.concatenate(pieces) if pieces else wav
|
|
|
|
start_len = min(pre_crossfade_samples, total_len)
|
|
remaining_after_start = max(total_len - start_len, 0)
|
|
end_len = min(post_crossfade_samples, remaining_after_start)
|
|
middle_end_idx = total_len - end_len
|
|
|
|
start_segment = wav[:start_len]
|
|
middle_segment = wav[start_len:middle_end_idx]
|
|
end_segment = wav[middle_end_idx:]
|
|
|
|
def _linear_fade(num_samples: int, start: float, end: float) -> np.ndarray:
|
|
if num_samples <= 0:
|
|
return np.zeros((0,), dtype=np.float32)
|
|
return np.linspace(start, end, num_samples, endpoint=True, dtype=np.float32)
|
|
|
|
start_crossfade = start_segment * _linear_fade(start_len, 0.0, 1.0)
|
|
end_crossfade = end_segment * _linear_fade(end_segment.shape[0], 1.0, 0.0)
|
|
|
|
pieces: List[np.ndarray] = []
|
|
if start_sil_samples > 0:
|
|
pieces.append(np.zeros(start_sil_samples, dtype=np.float32))
|
|
if start_crossfade.size > 0:
|
|
pieces.append(start_crossfade.astype(np.float32, copy=False))
|
|
if middle_segment.size > 0:
|
|
pieces.append(middle_segment.astype(np.float32, copy=False))
|
|
if end_crossfade.size > 0:
|
|
pieces.append(end_crossfade.astype(np.float32, copy=False))
|
|
if end_sil_samples > 0:
|
|
pieces.append(np.zeros(end_sil_samples, dtype=np.float32))
|
|
|
|
return np.concatenate(pieces)
|
|
|
|
|
|
def _load_audio_to_24k(
|
|
audio: Union[str, np.ndarray, torch.Tensor, Dict[str, Any]],
|
|
*,
|
|
target_sr: int = 24000,
|
|
augment_with_silence: bool = False,
|
|
) -> np.ndarray:
|
|
if isinstance(audio, np.ndarray):
|
|
wav_out = audio.astype(np.float32)
|
|
elif isinstance(audio, torch.Tensor):
|
|
wav_out = audio.detach().cpu().float().numpy()
|
|
elif isinstance(audio, str):
|
|
if librosa is None:
|
|
raise RuntimeError("librosa is required to load audio file paths. Please pip install librosa.")
|
|
wav, sr = librosa.load(audio, sr=None, mono=True)
|
|
wav_out = _resample_if_needed(wav, int(sr), target_sr)
|
|
elif isinstance(audio, dict) and "array" in audio and "sampling_rate" in audio:
|
|
arr = np.asarray(audio["array"], dtype=np.float32)
|
|
sr = int(audio["sampling_rate"])
|
|
wav_out = _resample_if_needed(arr, sr, target_sr)
|
|
else:
|
|
raise ValueError(f"Unsupported audio type: {type(audio)}")
|
|
|
|
wav_out = np.asarray(wav_out, dtype=np.float32)
|
|
|
|
if augment_with_silence:
|
|
wav_out = _apply_silence_with_crossfade(wav_out, sample_rate=target_sr)
|
|
|
|
return wav_out
|
|
|
|
|
|
@dataclass
|
|
class VibeVoiceCollator:
|
|
processor: Any # VibeVoiceProcessor
|
|
max_length: Optional[int] = None
|
|
speech_compress_ratio: int = 3200
|
|
semantic_vae_dim: int = 128
|
|
compute_semantics: bool = False
|
|
debug_checks: bool = False
|
|
|
|
text_field: str = "text"
|
|
audio_field: str = "audio"
|
|
voice_prompts_field: str = "voice_prompts"
|
|
voice_prompt_drop_rate: float = 0.0
|
|
|
|
def __call__(self, features: Sequence[Dict[str, Any]]) -> Dict[str, Any]:
|
|
batch_size = len(features)
|
|
|
|
sample_input_ids: List[List[int]] = []
|
|
sample_attention_masks: List[List[int]] = []
|
|
sample_acoustic_input_masks: List[List[bool]] = []
|
|
sample_acoustic_loss_masks: List[List[bool]] = []
|
|
|
|
all_speech_waveforms: List[np.ndarray] = []
|
|
all_speech_latent_lengths: List[int] = []
|
|
per_segment_is_target: List[bool] = []
|
|
|
|
for ex in features:
|
|
text: str = ex.get(self.text_field, "")
|
|
voice_prompts: Optional[List[Union[str, np.ndarray, torch.Tensor]]] = ex.get(self.voice_prompts_field)
|
|
target_audio: Union[str, np.ndarray, torch.Tensor, Dict[str, Any]] = ex.get(self.audio_field)
|
|
|
|
# Clamp drop rate for safety
|
|
_drop_rate = self.voice_prompt_drop_rate
|
|
if _drop_rate < 0.0:
|
|
_drop_rate = 0.0
|
|
elif _drop_rate > 1.0:
|
|
_drop_rate = 1.0
|
|
|
|
proc = self.processor(
|
|
text=[text],
|
|
voice_samples=[voice_prompts] if voice_prompts is not None and random.random() >= _drop_rate else None,
|
|
padding=False,
|
|
truncation=False,
|
|
max_length=self.max_length,
|
|
return_tensors="pt",
|
|
)
|
|
|
|
ids = proc["input_ids"][0].tolist()
|
|
attn = proc.get("attention_mask", torch.ones_like(proc["input_ids"]))[0].tolist()
|
|
speech_input_mask = proc.get("speech_input_mask")
|
|
if speech_input_mask is None:
|
|
speech_input_mask = torch.zeros_like(proc["input_ids"], dtype=torch.bool)
|
|
speech_input_mask_list = speech_input_mask[0].tolist()
|
|
|
|
wav_target = _load_audio_to_24k(target_audio, target_sr=24000, augment_with_silence=True)
|
|
# Prefer exact frame count from acoustic tokenizer if available; fallback to compress ratio
|
|
target_latent_len = None
|
|
try:
|
|
acoustic_tok = getattr(self.processor, "acoustic_tokenizer", None)
|
|
if acoustic_tok is not None and hasattr(acoustic_tok, "encode"):
|
|
enc_out = acoustic_tok.encode(wav_target)
|
|
# Normalize various possible return formats to get time dimension
|
|
T = None
|
|
try:
|
|
# Direct array-like with shape (T, D) or (T,)
|
|
if hasattr(enc_out, "shape") and len(getattr(enc_out, "shape", [])) >= 1:
|
|
T = int(enc_out.shape[0])
|
|
else:
|
|
# Nested lists/tuples or ModelOutput-like
|
|
cand = enc_out
|
|
# Drill down a couple of levels safely
|
|
for _ in range(2):
|
|
if isinstance(cand, (list, tuple)) and len(cand) > 0:
|
|
cand = cand[0]
|
|
if hasattr(cand, "shape") and len(getattr(cand, "shape", [])) >= 1:
|
|
T = int(cand.shape[0])
|
|
except Exception:
|
|
T = None
|
|
if T is not None and T > 0:
|
|
target_latent_len = T
|
|
except Exception:
|
|
target_latent_len = None
|
|
if target_latent_len is None:
|
|
target_latent_len = max(1, int(math.ceil(len(wav_target) / float(self.speech_compress_ratio))))
|
|
|
|
speech_diff_id = self.processor.tokenizer.speech_diffusion_id
|
|
target_placeholders = [speech_diff_id] * target_latent_len
|
|
|
|
ids_extended = ids + target_placeholders
|
|
attn_extended = attn + [1] * target_latent_len
|
|
|
|
acoustic_input_mask = speech_input_mask_list + [True] * target_latent_len
|
|
acoustic_loss_mask = ([False] * len(speech_input_mask_list)) + [True] * target_latent_len
|
|
|
|
speech_end_id = self.processor.tokenizer.speech_end_id
|
|
ids_extended.append(speech_end_id)
|
|
attn_extended.append(1)
|
|
acoustic_input_mask.append(False)
|
|
acoustic_loss_mask.append(False)
|
|
|
|
# Ensure text decoding sees an explicit end-of-sequence token after speech output.
|
|
eos_token_id = getattr(self.processor.tokenizer, "eos_id", None)
|
|
if eos_token_id is None:
|
|
eos_token_id = getattr(self.processor.tokenizer, "eos_token_id", None)
|
|
if eos_token_id is not None and eos_token_id >= 0:
|
|
ids_extended.append(eos_token_id)
|
|
attn_extended.append(1)
|
|
acoustic_input_mask.append(False)
|
|
acoustic_loss_mask.append(False)
|
|
|
|
if self.max_length is not None and len(ids_extended) > self.max_length:
|
|
cut = len(ids_extended) - int(self.max_length)
|
|
leading_non_acoustic = 0
|
|
for v in acoustic_input_mask:
|
|
if v:
|
|
break
|
|
leading_non_acoustic += 1
|
|
if cut > leading_non_acoustic:
|
|
raise ValueError(
|
|
f"--max_length={self.max_length} would truncate into acoustic tokens. "
|
|
f"Needed cut={cut}, but only {leading_non_acoustic} leading non-acoustic tokens available. "
|
|
"Increase max_length or shorten text/voice-prompt preamble."
|
|
)
|
|
ids_extended = ids_extended[cut:]
|
|
attn_extended = attn_extended[cut:]
|
|
acoustic_input_mask = acoustic_input_mask[cut:]
|
|
acoustic_loss_mask = acoustic_loss_mask[cut:]
|
|
|
|
sample_input_ids.append(ids_extended)
|
|
sample_attention_masks.append(attn_extended)
|
|
sample_acoustic_input_masks.append(acoustic_input_mask)
|
|
sample_acoustic_loss_masks.append(acoustic_loss_mask)
|
|
|
|
voice_speeches = []
|
|
voice_latent_lengths = []
|
|
if proc.get("speech_tensors") is not None:
|
|
voice_np = proc["speech_tensors"].cpu().numpy()
|
|
voice_masks = proc["speech_masks"].cpu().numpy().astype(bool)
|
|
for seg_idx in range(voice_np.shape[0]):
|
|
voice_speeches.append(voice_np[seg_idx])
|
|
voice_latent_lengths.append(int(voice_masks[seg_idx].sum()))
|
|
|
|
all_speech_waveforms.extend(voice_speeches)
|
|
all_speech_latent_lengths.extend(voice_latent_lengths)
|
|
per_segment_is_target.extend([False] * len(voice_speeches))
|
|
|
|
all_speech_waveforms.append(wav_target)
|
|
all_speech_latent_lengths.append(target_latent_len)
|
|
per_segment_is_target.append(True)
|
|
|
|
max_seq_len = max(len(x) for x in sample_input_ids)
|
|
padded_input_ids = []
|
|
padded_attention_masks = []
|
|
padded_acoustic_input_masks = []
|
|
padded_acoustic_loss_masks = []
|
|
tok = self.processor.tokenizer
|
|
pad_token_id = getattr(tok, "pad_token_id", None)
|
|
if pad_token_id is None or pad_token_id < 0:
|
|
pad_token_id = getattr(tok, "eos_token_id", None)
|
|
if pad_token_id is None or pad_token_id < 0:
|
|
raise ValueError(
|
|
"Tokenizer has no pad_token_id or eos_token_id; please set one or pass a valid pad id."
|
|
)
|
|
for ids, attn, ain_mask, aloss_mask in zip(
|
|
sample_input_ids, sample_attention_masks, sample_acoustic_input_masks, sample_acoustic_loss_masks
|
|
):
|
|
pad_len = max_seq_len - len(ids)
|
|
padded_input_ids.append(ids + [pad_token_id] * pad_len)
|
|
padded_attention_masks.append(attn + [0] * pad_len)
|
|
padded_acoustic_input_masks.append(ain_mask + [False] * pad_len)
|
|
padded_acoustic_loss_masks.append(aloss_mask + [False] * pad_len)
|
|
|
|
input_ids_tensor = torch.tensor(padded_input_ids, dtype=torch.long)
|
|
attention_mask_tensor = torch.tensor(padded_attention_masks, dtype=torch.long)
|
|
acoustic_input_mask_tensor = torch.tensor(padded_acoustic_input_masks, dtype=torch.bool)
|
|
acoustic_loss_mask_tensor = torch.tensor(padded_acoustic_loss_masks, dtype=torch.bool)
|
|
|
|
if all_speech_waveforms:
|
|
max_wave_len = max(w.shape[0] for w in all_speech_waveforms)
|
|
padded_speeches = np.zeros((len(all_speech_waveforms), max_wave_len), dtype=np.float32)
|
|
for i, w in enumerate(all_speech_waveforms):
|
|
L = w.shape[0]
|
|
padded_speeches[i, :L] = w
|
|
|
|
max_latent_len = max(all_speech_latent_lengths) if all_speech_latent_lengths else 1
|
|
speech_masks_np = np.zeros((len(all_speech_waveforms), max_latent_len), dtype=np.bool_)
|
|
for i, L_lat in enumerate(all_speech_latent_lengths):
|
|
speech_masks_np[i, :L_lat] = True
|
|
|
|
speech_tensors_tensor = torch.tensor(padded_speeches, dtype=torch.float32)
|
|
speech_masks_tensor = torch.tensor(speech_masks_np, dtype=torch.bool)
|
|
|
|
speeches_loss_input_np = np.zeros_like(speech_masks_np, dtype=np.bool_)
|
|
for i, is_target in enumerate(per_segment_is_target):
|
|
if is_target:
|
|
speeches_loss_input_np[i] = speech_masks_np[i]
|
|
speeches_loss_input_tensor = torch.tensor(speeches_loss_input_np, dtype=torch.bool)
|
|
|
|
# Semantic features
|
|
if self.compute_semantics and hasattr(self.processor, "semantic_tokenizer") and self.processor.semantic_tokenizer is not None:
|
|
sem_feats: List[np.ndarray] = []
|
|
for w in all_speech_waveforms:
|
|
try:
|
|
# Expect [T, D] where T ≈ ceil(len(w)/compress_ratio)
|
|
sem = self.processor.semantic_tokenizer.encode(w)
|
|
sem = np.asarray(sem, dtype=np.float32)
|
|
except Exception:
|
|
sem = np.zeros((0, self.semantic_vae_dim), dtype=np.float32)
|
|
if sem.ndim != 2:
|
|
raise RuntimeError(f"Semantic tokenizer returned unexpected shape {sem.shape}. Expect [T, D].")
|
|
L = sem.shape[0]
|
|
D = sem.shape[1]
|
|
if D != self.semantic_vae_dim:
|
|
if D < self.semantic_vae_dim:
|
|
pad_d = np.zeros((L, self.semantic_vae_dim - D), dtype=np.float32)
|
|
sem = np.concatenate([sem, pad_d], axis=1)
|
|
else:
|
|
sem = sem[:, : self.semantic_vae_dim]
|
|
if L < max_latent_len:
|
|
pad = np.zeros((max_latent_len - L, self.semantic_vae_dim), dtype=np.float32)
|
|
sem = np.concatenate([sem, pad], axis=0)
|
|
elif L > max_latent_len:
|
|
sem = sem[:max_latent_len]
|
|
sem_feats.append(sem.astype(np.float32))
|
|
speech_semantic_tensors = torch.tensor(np.stack(sem_feats, axis=0), dtype=torch.float32)
|
|
else:
|
|
# Semantic tokenizer unavailable while semantics are required for training.
|
|
# Raise to avoid silently degrading alignment with zeroed features.
|
|
raise RuntimeError(
|
|
"Semantic features are required but could not be computed. "
|
|
"Ensure processor.semantic_tokenizer is available or precompute and provide features."
|
|
)
|
|
else:
|
|
speech_tensors_tensor = None
|
|
speech_masks_tensor = None
|
|
speeches_loss_input_tensor = None
|
|
speech_semantic_tensors = None # No segments in batch
|
|
|
|
if self.debug_checks:
|
|
assert (input_ids_tensor >= 0).all(), "input_ids contains negative indices"
|
|
if speech_tensors_tensor is not None:
|
|
assert speech_tensors_tensor.dim() == 2, "Expected speech_tensors 2D [segments, samples]"
|
|
|
|
return {
|
|
"input_ids": input_ids_tensor,
|
|
"attention_mask": attention_mask_tensor,
|
|
"speech_tensors": speech_tensors_tensor,
|
|
"speech_masks": speech_masks_tensor,
|
|
"speech_semantic_tensors": speech_semantic_tensors,
|
|
"acoustic_input_mask": acoustic_input_mask_tensor,
|
|
"acoustic_loss_mask": acoustic_loss_mask_tensor,
|
|
"speeches_loss_input": speeches_loss_input_tensor,
|
|
} |