mirror of
https://github.com/shivammehta25/Matcha-TTS.git
synced 2026-02-04 09:49:21 +08:00
115 lines
3.4 KiB
Python
115 lines
3.4 KiB
Python
from abc import ABC
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from matcha.models.components.decoder import Decoder
|
|
from matcha.utils.pylogger import get_pylogger
|
|
|
|
log = get_pylogger(__name__)
|
|
|
|
|
|
class BASECFM(torch.nn.Module, ABC):
|
|
def __init__(
|
|
self,
|
|
n_feats,
|
|
cfm_params,
|
|
n_spks=1,
|
|
spk_emb_dim=128,
|
|
):
|
|
super().__init__()
|
|
self.n_feats = n_feats
|
|
self.n_spks = n_spks
|
|
self.spk_emb_dim = spk_emb_dim
|
|
self.solver = cfm_params.solver
|
|
if hasattr(cfm_params, "sigma_min"):
|
|
self.sigma_min = cfm_params.sigma_min
|
|
else:
|
|
self.sigma_min = 1e-4
|
|
|
|
self.estimator = None
|
|
|
|
@torch.inference_mode()
|
|
def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None):
|
|
"""Forward diffusion
|
|
|
|
Args:
|
|
z (_type_): mu + noise (we don't need this in this formulation), we will sample the noise again
|
|
mask (_type_): output_mask
|
|
mu (_type_): output of encoder
|
|
n_timesteps (_type_): number of diffusion steps
|
|
stoc (bool, optional): _description_. Defaults to False.
|
|
spks (_type_, optional): _description_. Defaults to None.
|
|
|
|
Returns:
|
|
sample: _description_
|
|
"""
|
|
z = torch.randn_like(mu) * temperature
|
|
t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device)
|
|
return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond)
|
|
|
|
def solve_euler(self, x, t_span, mu, mask, spks, cond):
|
|
"""
|
|
Fixed euler solver for ODEs.
|
|
Args:
|
|
x (_type_): _description_
|
|
t (_type_): _description_
|
|
"""
|
|
t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]
|
|
sol = []
|
|
|
|
steps = 1
|
|
while steps <= len(t_span) - 1:
|
|
dphi_dt = self.estimator(x, mask, mu, t, spks, cond)
|
|
|
|
x = x + dt * dphi_dt
|
|
t = t + dt
|
|
sol.append(x)
|
|
if steps < len(t_span) - 1:
|
|
dt = t_span[steps + 1] - t
|
|
steps += 1
|
|
|
|
return sol[-1]
|
|
|
|
def compute_loss(self, x1, mask, mu, spks=None, cond=None):
|
|
"""Computes diffusion loss
|
|
|
|
Args:
|
|
x1 (_type_): Target
|
|
mask (_type_): target mask
|
|
mu (_type_): output of encoder
|
|
spks (_type_, optional): speaker embedding. Defaults to None.
|
|
|
|
Returns:
|
|
loss: diffusion loss
|
|
y: conditional flow
|
|
"""
|
|
b, _, t = mu.shape
|
|
|
|
# random timestep
|
|
t = torch.rand([b, 1, 1], device=mu.device, dtype=mu.dtype)
|
|
# sample noise p(x_0)
|
|
z = torch.randn_like(x1)
|
|
|
|
y = (1 - (1 - self.sigma_min) * t) * z + t * x1
|
|
u = x1 - (1 - self.sigma_min) * z
|
|
|
|
loss = F.mse_loss(self.estimator(y, mask, mu, t.squeeze(), spks), u, reduction="sum") / (
|
|
torch.sum(mask) * u.shape[1]
|
|
)
|
|
return loss, y
|
|
|
|
|
|
class CFM(BASECFM):
|
|
def __init__(self, in_channels, out_channel, cfm_params, decoder_params, n_spks=1, spk_emb_dim=64):
|
|
super().__init__(
|
|
n_feats=in_channels,
|
|
cfm_params=cfm_params,
|
|
n_spks=n_spks,
|
|
spk_emb_dim=spk_emb_dim,
|
|
)
|
|
|
|
in_channels = in_channels + (spk_emb_dim if n_spks > 1 else 0)
|
|
# Just change the architecture of the estimator here
|
|
self.estimator = Decoder(in_channels=in_channels, out_channels=out_channel, **decoder_params)
|