| 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 | import threading |
| 14 | from typing import Literal |
| 15 | |
| 16 | import torch |
| 17 | import torch.nn.functional as F |
| 18 | from torch import nn |
| 19 | from x_transformers import RMSNorm |
| 20 | from x_transformers.x_transformers import RotaryEmbedding |
| 21 | |
| 22 | from f5_tts.model.modules import ( |
| 23 | Attention, |
| 24 | AttnProcessor, |
| 25 | ConvNeXtV2Block, |
| 26 | ConvPositionEmbedding, |
| 27 | FeedForward, |
| 28 | TimestepEmbedding, |
| 29 | get_pos_embed_indices, |
| 30 | precompute_freqs_cis, |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | # Text embedding |
| 35 | |
| 36 | |
| 37 | class TextEmbedding(nn.Module): |
| 38 | def __init__(self, text_num_embeds, text_dim, mask_padding=True, conv_layers=0, conv_mult=2): |
| 39 | super().__init__() |
| 40 | self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim) # use 0 as filler token |
| 41 | |
| 42 | self.mask_padding = mask_padding # mask filler and batch padding tokens or not |
| 43 | |
| 44 | if conv_layers > 0: |
| 45 | self.extra_modeling = True |
| 46 | self.precompute_max_pos = 4096 # ~44s of 24khz audio |
| 47 | self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False) |
| 48 | self.text_blocks = nn.Sequential( |
| 49 | *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)] |
| 50 | ) |
| 51 | else: |
| 52 | self.extra_modeling = False |
| 53 | |
| 54 | def forward(self, text: int["b nt"], seq_len, drop_text=False): |
| 55 | text = text + 1 # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx() |
| 56 | text = text[:, :seq_len] # curtail if character tokens are more than the mel spec tokens |
| 57 | batch, text_len = text.shape[0], text.shape[1] |
| 58 | text = F.pad(text, (0, seq_len - text_len), value=0) |
| 59 | if self.mask_padding: |
| 60 | text_mask = text == 0 |
| 61 | |
| 62 | if drop_text: # cfg for text |
| 63 | text = torch.zeros_like(text) |
| 64 | |
| 65 | text = self.text_embed(text) # b n -> b n d |
| 66 | |
| 67 | # possible extra modeling |
| 68 | if self.extra_modeling: |
| 69 | # sinus pos emb |
| 70 | batch_start = torch.zeros((batch,), dtype=torch.long) |
| 71 | pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos) |
| 72 | text_pos_embed = self.freqs_cis[pos_idx] |
| 73 | text = text + text_pos_embed |
| 74 | |
| 75 | # convnextv2 blocks |
| 76 | if self.mask_padding: |
| 77 | text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0) |
| 78 | for block in self.text_blocks: |
| 79 | text = block(text) |
| 80 | text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0) |
| 81 | else: |
| 82 | text = self.text_blocks(text) |
| 83 | |
| 84 | return text |
| 85 | |
| 86 | |
| 87 | # noised input audio and context mixing embedding |
| 88 | |
| 89 | |
| 90 | class InputEmbedding(nn.Module): |
| 91 | def __init__(self, mel_dim, text_dim, out_dim): |
| 92 | super().__init__() |
| 93 | self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim) |
| 94 | self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim) |
| 95 | |
| 96 | def forward(self, x: float["b n d"], cond: float["b n d"], text_embed: float["b n d"], drop_audio_cond=False): |
| 97 | if drop_audio_cond: # cfg for cond audio |
| 98 | cond = torch.zeros_like(cond) |
| 99 | |
| 100 | x = self.proj(torch.cat((x, cond, text_embed), dim=-1)) |
| 101 | x = self.conv_pos_embed(x) + x |
| 102 | return x |
| 103 | |
| 104 | |
| 105 | # Flat UNet Transformer backbone |
| 106 | |
| 107 | |
| 108 | class UNetT(nn.Module): |
| 109 | def __init__( |
| 110 | self, |
| 111 | *, |
| 112 | dim, |
| 113 | depth=8, |
| 114 | heads=8, |
| 115 | dim_head=64, |
| 116 | dropout=0.1, |
| 117 | ff_mult=4, |
| 118 | mel_dim=100, |
| 119 | text_num_embeds=256, |
| 120 | text_dim=None, |
| 121 | text_mask_padding=True, |
| 122 | qk_norm=None, |
| 123 | conv_layers=0, |
| 124 | pe_attn_head=None, |
| 125 | attn_backend="torch", # "torch" | "flash_attn" |
| 126 | attn_mask_enabled=False, |
| 127 | skip_connect_type: Literal["add", "concat", "none"] = "concat", |
| 128 | ): |
| 129 | super().__init__() |
| 130 | assert depth % 2 == 0, "UNet-Transformer's depth should be even." |
| 131 | |
| 132 | self.time_embed = TimestepEmbedding(dim) |
| 133 | if text_dim is None: |
| 134 | text_dim = mel_dim |
| 135 | self.text_embed = TextEmbedding( |
| 136 | text_num_embeds, text_dim, mask_padding=text_mask_padding, conv_layers=conv_layers |
| 137 | ) |
| 138 | self.input_embed = InputEmbedding(mel_dim, text_dim, dim) |
| 139 | |
| 140 | self.rotary_embed = RotaryEmbedding(dim_head) |
| 141 | |
| 142 | # transformer layers & skip connections |
| 143 | |
| 144 | self.dim = dim |
| 145 | self.skip_connect_type = skip_connect_type |
| 146 | needs_skip_proj = skip_connect_type == "concat" |
| 147 | |
| 148 | self.depth = depth |
| 149 | self.layers = nn.ModuleList([]) |
| 150 | |
| 151 | for idx in range(depth): |
| 152 | is_later_half = idx >= (depth // 2) |
| 153 | |
| 154 | attn_norm = RMSNorm(dim) |
| 155 | attn = Attention( |
| 156 | processor=AttnProcessor( |
| 157 | pe_attn_head=pe_attn_head, |
| 158 | attn_backend=attn_backend, |
| 159 | attn_mask_enabled=attn_mask_enabled, |
| 160 | ), |
| 161 | dim=dim, |
| 162 | heads=heads, |
| 163 | dim_head=dim_head, |
| 164 | dropout=dropout, |
| 165 | qk_norm=qk_norm, |
| 166 | ) |
| 167 | |
| 168 | ff_norm = RMSNorm(dim) |
| 169 | ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh") |
| 170 | |
| 171 | skip_proj = nn.Linear(dim * 2, dim, bias=False) if needs_skip_proj and is_later_half else None |
| 172 | |
| 173 | self.layers.append( |
| 174 | nn.ModuleList( |
| 175 | [ |
| 176 | skip_proj, |
| 177 | attn_norm, |
| 178 | attn, |
| 179 | ff_norm, |
| 180 | ff, |
| 181 | ] |
| 182 | ) |
| 183 | ) |
| 184 | |
| 185 | self.norm_out = RMSNorm(dim) |
| 186 | self.proj_out = nn.Linear(dim, mel_dim) |
| 187 | |
| 188 | # `_cache_local` is lazily initialized on first inference-time cache write so that |
| 189 | # training models (which never touch the cache) stay deepcopy-friendly for EMA. |
| 190 | def _get_cache_local(self): |
| 191 | cache = self.__dict__.get("_cache_local") |
| 192 | if cache is None: |
| 193 | cache = threading.local() |
| 194 | self.__dict__["_cache_local"] = cache |
| 195 | return cache |
| 196 | |
| 197 | @property |
| 198 | def text_cond(self): |
| 199 | cache = self.__dict__.get("_cache_local") |
| 200 | return getattr(cache, "text_cond", None) if cache is not None else None |
| 201 | |
| 202 | @text_cond.setter |
| 203 | def text_cond(self, value): |
| 204 | self._get_cache_local().text_cond = value |
| 205 | |
| 206 | @property |
| 207 | def text_uncond(self): |
| 208 | cache = self.__dict__.get("_cache_local") |
| 209 | return getattr(cache, "text_uncond", None) if cache is not None else None |
| 210 | |
| 211 | @text_uncond.setter |
| 212 | def text_uncond(self, value): |
| 213 | self._get_cache_local().text_uncond = value |
| 214 | |
| 215 | def get_input_embed( |
| 216 | self, |
| 217 | x, # b n d |
| 218 | cond, # b n d |
| 219 | text, # b nt |
| 220 | drop_audio_cond: bool = False, |
| 221 | drop_text: bool = False, |
| 222 | cache: bool = True, |
| 223 | ): |
| 224 | seq_len = x.shape[1] |
| 225 | if cache: |
| 226 | if drop_text: |
| 227 | if self.text_uncond is None: |
| 228 | self.text_uncond = self.text_embed(text, seq_len, drop_text=True) |
| 229 | text_embed = self.text_uncond |
| 230 | else: |
| 231 | if self.text_cond is None: |
| 232 | self.text_cond = self.text_embed(text, seq_len, drop_text=False) |
| 233 | text_embed = self.text_cond |
| 234 | else: |
| 235 | text_embed = self.text_embed(text, seq_len, drop_text=drop_text) |
| 236 | |
| 237 | x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond) |
| 238 | |
| 239 | return x |
| 240 | |
| 241 | def clear_cache(self): |
| 242 | self.text_cond, self.text_uncond = None, None |
| 243 | |
| 244 | def forward( |
| 245 | self, |
| 246 | x: float["b n d"], # nosied input audio |
| 247 | cond: float["b n d"], # masked cond audio |
| 248 | text: int["b nt"], # text |
| 249 | time: float["b"] | float[""], # time step |
| 250 | mask: bool["b n"] | None = None, |
| 251 | drop_audio_cond: bool = False, # cfg for cond audio |
| 252 | drop_text: bool = False, # cfg for text |
| 253 | cfg_infer: bool = False, # cfg inference, pack cond & uncond forward |
| 254 | cache: bool = False, |
| 255 | ): |
| 256 | batch, seq_len = x.shape[0], x.shape[1] |
| 257 | if time.ndim == 0: |
| 258 | time = time.repeat(batch) |
| 259 | |
| 260 | # t: conditioning time, c: context (text + masked cond audio), x: noised input audio |
| 261 | t = self.time_embed(time) |
| 262 | if cfg_infer: # pack cond & uncond forward: b n d -> 2b n d |
| 263 | x_cond = self.get_input_embed(x, cond, text, drop_audio_cond=False, drop_text=False, cache=cache) |
| 264 | x_uncond = self.get_input_embed(x, cond, text, drop_audio_cond=True, drop_text=True, cache=cache) |
| 265 | x = torch.cat((x_cond, x_uncond), dim=0) |
| 266 | t = torch.cat((t, t), dim=0) |
| 267 | mask = torch.cat((mask, mask), dim=0) if mask is not None else None |
| 268 | else: |
| 269 | x = self.get_input_embed(x, cond, text, drop_audio_cond=drop_audio_cond, drop_text=drop_text, cache=cache) |
| 270 | |
| 271 | # postfix time t to input x, [b n d] -> [b n+1 d] |
| 272 | x = torch.cat([t.unsqueeze(1), x], dim=1) # pack t to x |
| 273 | if mask is not None: |
| 274 | mask = F.pad(mask, (1, 0), value=1) |
| 275 | |
| 276 | rope = self.rotary_embed.forward_from_seq_len(seq_len + 1) |
| 277 | |
| 278 | # flat unet transformer |
| 279 | skip_connect_type = self.skip_connect_type |
| 280 | skips = [] |
| 281 | for idx, (maybe_skip_proj, attn_norm, attn, ff_norm, ff) in enumerate(self.layers): |
| 282 | layer = idx + 1 |
| 283 | |
| 284 | # skip connection logic |
| 285 | is_first_half = layer <= (self.depth // 2) |
| 286 | is_later_half = not is_first_half |
| 287 | |
| 288 | if is_first_half: |
| 289 | skips.append(x) |
| 290 | |
| 291 | if is_later_half: |
| 292 | skip = skips.pop() |
| 293 | if skip_connect_type == "concat": |
| 294 | x = torch.cat((x, skip), dim=-1) |
| 295 | x = maybe_skip_proj(x) |
| 296 | elif skip_connect_type == "add": |
| 297 | x = x + skip |
| 298 | |
| 299 | # attention and feedforward blocks |
| 300 | x = attn(attn_norm(x), rope=rope, mask=mask) + x |
| 301 | x = ff(ff_norm(x)) + x |
| 302 | |
| 303 | assert len(skips) == 0 |
| 304 | |
| 305 | x = self.norm_out(x)[:, 1:, :] # unpack t from x |
| 306 | |
| 307 | return self.proj_out(x) |
| 308 |