| 1 | """ |
| 2 | ein notation: |
| 3 | b - batch |
| 4 | n - sequence |
| 5 | nt - text sequence |
| 6 | nw - raw wave length |
| 7 | d - dimension |
| 8 | """ |
| 9 | # ruff: noqa: F722 F821 |
| 10 | |
| 11 | from __future__ import annotations |
| 12 | |
| 13 | from random import random |
| 14 | from typing import Callable |
| 15 | |
| 16 | import torch |
| 17 | import torch.nn.functional as F |
| 18 | from torch import nn |
| 19 | from torch.nn.utils.rnn import pad_sequence |
| 20 | from torchdiffeq import odeint |
| 21 | |
| 22 | from f5_tts.model.modules import MelSpec |
| 23 | from f5_tts.model.utils import ( |
| 24 | default, |
| 25 | exists, |
| 26 | get_epss_timesteps, |
| 27 | lens_to_mask, |
| 28 | list_str_to_idx, |
| 29 | list_str_to_tensor, |
| 30 | mask_from_frac_lengths, |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | class CFM(nn.Module): |
| 35 | def __init__( |
| 36 | self, |
| 37 | transformer: nn.Module, |
| 38 | sigma=0.0, |
| 39 | odeint_kwargs: dict = dict( |
| 40 | # atol = 1e-5, |
| 41 | # rtol = 1e-5, |
| 42 | method="euler" # 'midpoint' |
| 43 | ), |
| 44 | audio_drop_prob=0.3, |
| 45 | cond_drop_prob=0.2, |
| 46 | num_channels=None, |
| 47 | mel_spec_module: nn.Module | None = None, |
| 48 | mel_spec_kwargs: dict = dict(), |
| 49 | frac_lengths_mask: tuple[float, float] = (0.7, 1.0), |
| 50 | vocab_char_map: dict[str:int] | None = None, |
| 51 | ): |
| 52 | super().__init__() |
| 53 | |
| 54 | self.frac_lengths_mask = frac_lengths_mask |
| 55 | |
| 56 | # mel spec |
| 57 | self.mel_spec = default(mel_spec_module, MelSpec(**mel_spec_kwargs)) |
| 58 | num_channels = default(num_channels, self.mel_spec.n_mel_channels) |
| 59 | self.num_channels = num_channels |
| 60 | |
| 61 | # classifier-free guidance |
| 62 | self.audio_drop_prob = audio_drop_prob |
| 63 | self.cond_drop_prob = cond_drop_prob |
| 64 | |
| 65 | # transformer |
| 66 | self.transformer = transformer |
| 67 | dim = transformer.dim |
| 68 | self.dim = dim |
| 69 | |
| 70 | # conditional flow related |
| 71 | self.sigma = sigma |
| 72 | |
| 73 | # sampling related |
| 74 | self.odeint_kwargs = odeint_kwargs |
| 75 | |
| 76 | # vocab map for tokenization |
| 77 | self.vocab_char_map = vocab_char_map |
| 78 | |
| 79 | @property |
| 80 | def device(self): |
| 81 | return next(self.parameters()).device |
| 82 | |
| 83 | @torch.no_grad() |
| 84 | def sample( |
| 85 | self, |
| 86 | cond: float["b n d"] | float["b nw"], |
| 87 | text: int["b nt"] | list[str], |
| 88 | duration: int | int["b"], |
| 89 | *, |
| 90 | lens: int["b"] | None = None, |
| 91 | steps=32, |
| 92 | cfg_strength=1.0, |
| 93 | sway_sampling_coef=None, |
| 94 | seed: int | None = None, |
| 95 | max_duration=65536, |
| 96 | vocoder: Callable[[float["b d n"]], float["b nw"]] | None = None, |
| 97 | use_epss=True, |
| 98 | no_ref_audio=False, |
| 99 | duplicate_test=False, |
| 100 | t_inter=0.1, |
| 101 | edit_mask=None, |
| 102 | ): |
| 103 | self.eval() |
| 104 | # raw wave |
| 105 | |
| 106 | if cond.ndim == 2: |
| 107 | cond = self.mel_spec(cond) |
| 108 | cond = cond.permute(0, 2, 1) |
| 109 | assert cond.shape[-1] == self.num_channels |
| 110 | |
| 111 | cond = cond.to(next(self.parameters()).dtype) |
| 112 | |
| 113 | batch, cond_seq_len, device = *cond.shape[:2], cond.device |
| 114 | if not exists(lens): |
| 115 | lens = torch.full((batch,), cond_seq_len, device=device, dtype=torch.long) |
| 116 | |
| 117 | # text |
| 118 | |
| 119 | if isinstance(text, list): |
| 120 | if exists(self.vocab_char_map): |
| 121 | text = list_str_to_idx(text, self.vocab_char_map).to(device) |
| 122 | else: |
| 123 | text = list_str_to_tensor(text).to(device) |
| 124 | assert text.shape[0] == batch |
| 125 | |
| 126 | # duration |
| 127 | |
| 128 | cond_mask = lens_to_mask(lens) |
| 129 | if edit_mask is not None: |
| 130 | cond_mask = cond_mask & edit_mask |
| 131 | |
| 132 | if isinstance(duration, int): |
| 133 | duration = torch.full((batch,), duration, device=device, dtype=torch.long) |
| 134 | |
| 135 | duration = torch.maximum( |
| 136 | torch.maximum((text != -1).sum(dim=-1), lens) + 1, duration |
| 137 | ) # duration at least text/audio prompt length plus one token, so something is generated |
| 138 | duration = duration.clamp(max=max_duration) |
| 139 | max_duration = duration.amax() |
| 140 | |
| 141 | # duplicate test corner for inner time step oberservation |
| 142 | if duplicate_test: |
| 143 | test_cond = F.pad(cond, (0, 0, cond_seq_len, max_duration - 2 * cond_seq_len), value=0.0) |
| 144 | |
| 145 | cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value=0.0) |
| 146 | if no_ref_audio: |
| 147 | cond = torch.zeros_like(cond) |
| 148 | |
| 149 | cond_mask = F.pad(cond_mask, (0, max_duration - cond_mask.shape[-1]), value=False) |
| 150 | cond_mask = cond_mask.unsqueeze(-1) |
| 151 | step_cond = torch.where( |
| 152 | cond_mask, cond, torch.zeros_like(cond) |
| 153 | ) # allow direct control (cut cond audio) with lens passed in |
| 154 | |
| 155 | if batch > 1: |
| 156 | mask = lens_to_mask(duration) |
| 157 | else: # save memory and speed up, as single inference need no mask currently |
| 158 | mask = None |
| 159 | |
| 160 | # neural ode |
| 161 | |
| 162 | def fn(t, x): |
| 163 | # at each step, conditioning is fixed |
| 164 | # step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond)) |
| 165 | |
| 166 | # predict flow (cond) |
| 167 | if cfg_strength < 1e-5: |
| 168 | pred = self.transformer( |
| 169 | x=x, |
| 170 | cond=step_cond, |
| 171 | text=text, |
| 172 | time=t, |
| 173 | mask=mask, |
| 174 | drop_audio_cond=False, |
| 175 | drop_text=False, |
| 176 | cache=True, |
| 177 | ) |
| 178 | return pred |
| 179 | |
| 180 | # predict flow (cond and uncond), for classifier-free guidance |
| 181 | pred_cfg = self.transformer( |
| 182 | x=x, |
| 183 | cond=step_cond, |
| 184 | text=text, |
| 185 | time=t, |
| 186 | mask=mask, |
| 187 | cfg_infer=True, |
| 188 | cache=True, |
| 189 | ) |
| 190 | pred, null_pred = torch.chunk(pred_cfg, 2, dim=0) |
| 191 | return pred + (pred - null_pred) * cfg_strength |
| 192 | |
| 193 | # noise input |
| 194 | # to make sure batch inference result is same with different batch size, and for sure single inference |
| 195 | # still some difference maybe due to convolutional layers |
| 196 | y0 = [] |
| 197 | for dur in duration: |
| 198 | if exists(seed): |
| 199 | torch.manual_seed(seed) |
| 200 | y0.append(torch.randn(dur, self.num_channels, device=self.device, dtype=step_cond.dtype)) |
| 201 | y0 = pad_sequence(y0, padding_value=0, batch_first=True) |
| 202 | |
| 203 | t_start = 0 |
| 204 | |
| 205 | # duplicate test corner for inner time step oberservation |
| 206 | if duplicate_test: |
| 207 | t_start = t_inter |
| 208 | y0 = (1 - t_start) * y0 + t_start * test_cond |
| 209 | steps = int(steps * (1 - t_start)) |
| 210 | |
| 211 | if t_start == 0 and use_epss: # use Empirically Pruned Step Sampling for low NFE |
| 212 | t = get_epss_timesteps(steps, device=self.device, dtype=step_cond.dtype) |
| 213 | else: |
| 214 | t = torch.linspace(t_start, 1, steps + 1, device=self.device, dtype=step_cond.dtype) |
| 215 | if sway_sampling_coef is not None: |
| 216 | t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t) |
| 217 | |
| 218 | trajectory = odeint(fn, y0, t, **self.odeint_kwargs) |
| 219 | self.transformer.clear_cache() |
| 220 | |
| 221 | sampled = trajectory[-1] |
| 222 | out = sampled |
| 223 | out = torch.where(cond_mask, cond, out) |
| 224 | |
| 225 | if exists(vocoder): |
| 226 | out = out.permute(0, 2, 1) |
| 227 | out = vocoder(out) |
| 228 | |
| 229 | return out, trajectory |
| 230 | |
| 231 | def forward( |
| 232 | self, |
| 233 | inp: float["b n d"] | float["b nw"], # mel or raw wave |
| 234 | text: int["b nt"] | list[str], |
| 235 | *, |
| 236 | lens: int["b"] | None = None, |
| 237 | noise_scheduler: str | None = None, |
| 238 | ): |
| 239 | # handle raw wave |
| 240 | if inp.ndim == 2: |
| 241 | inp = self.mel_spec(inp) |
| 242 | inp = inp.permute(0, 2, 1) |
| 243 | assert inp.shape[-1] == self.num_channels |
| 244 | |
| 245 | batch, seq_len, dtype, device, _σ1 = *inp.shape[:2], inp.dtype, self.device, self.sigma |
| 246 | |
| 247 | # handle text as string |
| 248 | if isinstance(text, list): |
| 249 | if exists(self.vocab_char_map): |
| 250 | text = list_str_to_idx(text, self.vocab_char_map).to(device) |
| 251 | else: |
| 252 | text = list_str_to_tensor(text).to(device) |
| 253 | assert text.shape[0] == batch |
| 254 | |
| 255 | # lens and mask |
| 256 | if not exists(lens): # if lens not acquired by trainer from collate_fn |
| 257 | lens = torch.full((batch,), seq_len, device=device) |
| 258 | mask = lens_to_mask(lens, length=seq_len) |
| 259 | |
| 260 | # get a random span to mask out for training conditionally |
| 261 | frac_lengths = torch.zeros((batch,), device=self.device).float().uniform_(*self.frac_lengths_mask) |
| 262 | rand_span_mask = mask_from_frac_lengths(lens, frac_lengths) |
| 263 | |
| 264 | if exists(mask): |
| 265 | rand_span_mask &= mask |
| 266 | |
| 267 | # mel is x1 |
| 268 | x1 = inp |
| 269 | |
| 270 | # x0 is gaussian noise |
| 271 | x0 = torch.randn_like(x1) |
| 272 | |
| 273 | # time step |
| 274 | time = torch.rand((batch,), dtype=dtype, device=self.device) |
| 275 | # TODO. noise_scheduler |
| 276 | |
| 277 | # sample xt (φ_t(x) in the paper) |
| 278 | t = time.unsqueeze(-1).unsqueeze(-1) |
| 279 | φ = (1 - t) * x0 + t * x1 |
| 280 | flow = x1 - x0 |
| 281 | |
| 282 | # only predict what is within the random mask span for infilling |
| 283 | cond = torch.where(rand_span_mask[..., None], torch.zeros_like(x1), x1) |
| 284 | |
| 285 | # transformer and cfg training with a drop rate |
| 286 | drop_audio_cond = random() < self.audio_drop_prob # p_drop in voicebox paper |
| 287 | if random() < self.cond_drop_prob: # p_uncond in voicebox paper |
| 288 | drop_audio_cond = True |
| 289 | drop_text = True |
| 290 | else: |
| 291 | drop_text = False |
| 292 | |
| 293 | # apply mask will use more memory; might adjust batchsize or batchsampler long sequence threshold |
| 294 | pred = self.transformer( |
| 295 | x=φ, cond=cond, text=text, time=time, drop_audio_cond=drop_audio_cond, drop_text=drop_text, mask=mask |
| 296 | ) |
| 297 | |
| 298 | # flow matching loss |
| 299 | loss = F.mse_loss(pred, flow, reduction="none") |
| 300 | loss = loss[rand_span_mask] |
| 301 | |
| 302 | return loss.mean(), cond, pred |
| 303 |