storytold/storyteller-ml
2
0
Fork 0
mirror of https://github.com/storytold/storyteller-ml.git synced 2026-10-28 09:08:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

F5 tts (#59)

Browse Source 
* init tts/F5-TTS

* epic hax to download ckpts

* clear default reference audio/text from inference vars

* allow blank inputs for inference-cli to just pull checkpoints without running inference

* f5-tts github workflow to publish docker images
This commit is contained in:
Madhukar Mishra
2024-10-21 18:25:28 +05:30
committed by GitHub
parent e51123956f
commit b9c516e8f0
33 changed files with 5613 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tts/F5-TTS/model/cfm.py
+279
View File
@@ -0,0 +1,279 @@
"""
ein notation:
b - batch
n - sequence
nt - text sequence
nw - raw wave length
d - dimension
"""
from __future__ import annotations
from typing import Callable
from random import random
import torch
from torch import nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence
from torchdiffeq import odeint
from einops import rearrange
from model.modules import MelSpec
from model.utils import (
default, exists,
list_str_to_idx, list_str_to_tensor,
lens_to_mask, mask_from_frac_lengths,
)
class CFM(nn.Module):
def __init__(
self,
transformer: nn.Module,
sigma = 0.,
odeint_kwargs: dict = dict(
# atol = 1e-5,
# rtol = 1e-5,
method = 'euler' # 'midpoint'
),
audio_drop_prob = 0.3,
cond_drop_prob = 0.2,
num_channels = None,
mel_spec_module: nn.Module | None = None,
mel_spec_kwargs: dict = dict(),
frac_lengths_mask: tuple[float, float] = (0.7, 1.),
vocab_char_map: dict[str: int] | None = None
):
super().__init__()
self.frac_lengths_mask = frac_lengths_mask
# mel spec
self.mel_spec = default(mel_spec_module, MelSpec(**mel_spec_kwargs))
num_channels = default(num_channels, self.mel_spec.n_mel_channels)
self.num_channels = num_channels
# classifier-free guidance
self.audio_drop_prob = audio_drop_prob
self.cond_drop_prob = cond_drop_prob
# transformer
self.transformer = transformer
dim = transformer.dim
self.dim = dim
# conditional flow related
self.sigma = sigma
# sampling related
self.odeint_kwargs = odeint_kwargs
# vocab map for tokenization
self.vocab_char_map = vocab_char_map
@property
def device(self):
return next(self.parameters()).device
@torch.no_grad()
def sample(
self,
cond: float['b n d'] | float['b nw'],
text: int['b nt'] | list[str],
duration: int | int['b'],
*,
lens: int['b'] | None = None,
steps = 32,
cfg_strength = 1.,
sway_sampling_coef = None,
seed: int | None = None,
max_duration = 4096,
vocoder: Callable[[float['b d n']], float['b nw']] | None = None,
no_ref_audio = False,
duplicate_test = False,
t_inter = 0.1,
edit_mask = None,
):
self.eval()
# raw wave
if cond.ndim == 2:
cond = self.mel_spec(cond)
cond = rearrange(cond, 'b d n -> b n d')
assert cond.shape[-1] == self.num_channels
batch, cond_seq_len, device = *cond.shape[:2], cond.device
if not exists(lens):
lens = torch.full((batch,), cond_seq_len, device = device, dtype = torch.long)
# text
if isinstance(text, list):
if exists(self.vocab_char_map):
text = list_str_to_idx(text, self.vocab_char_map).to(device)
else:
text = list_str_to_tensor(text).to(device)
assert text.shape[0] == batch
if exists(text):
text_lens = (text != -1).sum(dim = -1)
lens = torch.maximum(text_lens, lens) # make sure lengths are at least those of the text characters
# duration
cond_mask = lens_to_mask(lens)
if edit_mask is not None:
cond_mask = cond_mask & edit_mask
if isinstance(duration, int):
duration = torch.full((batch,), duration, device = device, dtype = torch.long)
duration = torch.maximum(lens + 1, duration) # just add one token so something is generated
duration = duration.clamp(max = max_duration)
max_duration = duration.amax()
# duplicate test corner for inner time step oberservation
if duplicate_test:
test_cond = F.pad(cond, (0, 0, cond_seq_len, max_duration - 2*cond_seq_len), value = 0.)
cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value = 0.)
cond_mask = F.pad(cond_mask, (0, max_duration - cond_mask.shape[-1]), value = False)
cond_mask = rearrange(cond_mask, '... -> ... 1')
step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond)) # allow direct control (cut cond audio) with lens passed in
if batch > 1:
mask = lens_to_mask(duration)
else: # save memory and speed up, as single inference need no mask currently
mask = None
# test for no ref audio
if no_ref_audio:
cond = torch.zeros_like(cond)
# neural ode
def fn(t, x):
# at each step, conditioning is fixed
# step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond))
# predict flow
pred = self.transformer(x = x, cond = step_cond, text = text, time = t, mask = mask, drop_audio_cond = False, drop_text = False)
if cfg_strength < 1e-5:
return pred
null_pred = self.transformer(x = x, cond = step_cond, text = text, time = t, mask = mask, drop_audio_cond = True, drop_text = True)
return pred + (pred - null_pred) * cfg_strength
# noise input
# to make sure batch inference result is same with different batch size, and for sure single inference
# still some difference maybe due to convolutional layers
y0 = []
for dur in duration:
if exists(seed):
torch.manual_seed(seed)
y0.append(torch.randn(dur, self.num_channels, device = self.device))
y0 = pad_sequence(y0, padding_value = 0, batch_first = True)
t_start = 0
# duplicate test corner for inner time step oberservation
if duplicate_test:
t_start = t_inter
y0 = (1 - t_start) * y0 + t_start * test_cond
steps = int(steps * (1 - t_start))
t = torch.linspace(t_start, 1, steps, device = self.device)
if sway_sampling_coef is not None:
t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
trajectory = odeint(fn, y0, t, **self.odeint_kwargs)
sampled = trajectory[-1]
out = sampled
out = torch.where(cond_mask, cond, out)
if exists(vocoder):
out = rearrange(out, 'b n d -> b d n')
out = vocoder(out)
return out, trajectory
def forward(
self,
inp: float['b n d'] | float['b nw'], # mel or raw wave
text: int['b nt'] | list[str],
*,
lens: int['b'] | None = None,
noise_scheduler: str | None = None,
):
# handle raw wave
if inp.ndim == 2:
inp = self.mel_spec(inp)
inp = rearrange(inp, 'b d n -> b n d')
assert inp.shape[-1] == self.num_channels
batch, seq_len, dtype, device, σ1 = *inp.shape[:2], inp.dtype, self.device, self.sigma
# handle text as string
if isinstance(text, list):
if exists(self.vocab_char_map):
text = list_str_to_idx(text, self.vocab_char_map).to(device)
else:
text = list_str_to_tensor(text).to(device)
assert text.shape[0] == batch
# lens and mask
if not exists(lens):
lens = torch.full((batch,), seq_len, device = device)
mask = lens_to_mask(lens, length = seq_len) # useless here, as collate_fn will pad to max length in batch
# get a random span to mask out for training conditionally
frac_lengths = torch.zeros((batch,), device = self.device).float().uniform_(*self.frac_lengths_mask)
rand_span_mask = mask_from_frac_lengths(lens, frac_lengths)
if exists(mask):
rand_span_mask &= mask
# mel is x1
x1 = inp
# x0 is gaussian noise
x0 = torch.randn_like(x1)
# time step
time = torch.rand((batch,), dtype = dtype, device = self.device)
# TODO. noise_scheduler
# sample xt (φ_t(x) in the paper)
t = rearrange(time, 'b -> b 1 1')
φ = (1 - t) * x0 + t * x1
flow = x1 - x0
# only predict what is within the random mask span for infilling
cond = torch.where(
rand_span_mask[..., None],
torch.zeros_like(x1), x1
)
# transformer and cfg training with a drop rate
drop_audio_cond = random() < self.audio_drop_prob # p_drop in voicebox paper
if random() < self.cond_drop_prob: # p_uncond in voicebox paper
drop_audio_cond = True
drop_text = True
else:
drop_text = False
# if want rigourously mask out padding, record in collate_fn in dataset.py, and pass in here
# adding mask will use more memory, thus also need to adjust batchsampler with scaled down threshold for long sequences
pred = self.transformer(x = φ, cond = cond, text = text, time = time, drop_audio_cond = drop_audio_cond, drop_text = drop_text)
# flow matching loss
loss = F.mse_loss(pred, flow, reduction = 'none')
loss = loss[rand_span_mask]
return loss.mean(), cond, pred