| 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 | |
| 15 | import torch |
| 16 | import torch.nn.functional as F |
| 17 | from torch import nn |
| 18 | from x_transformers.x_transformers import RotaryEmbedding |
| 19 | |
| 20 | from f5_tts.model.modules import ( |
| 21 | AdaLayerNorm_Final, |
| 22 | ConvNeXtV2Block, |
| 23 | ConvPositionEmbedding, |
| 24 | DiTBlock, |
| 25 | TimestepEmbedding, |
| 26 | precompute_freqs_cis, |
| 27 | ) |
| 28 | |
| 29 | |
| 30 | # Text embedding |
| 31 | |
| 32 | |
| 33 | class TextEmbedding(nn.Module): |
| 34 | def __init__( |
| 35 | self, text_num_embeds, text_dim, mask_padding=True, average_upsampling=False, conv_layers=0, conv_mult=2 |
| 36 | ): |
| 37 | super().__init__() |
| 38 | self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim) # use 0 as filler token |
| 39 | |
| 40 | self.mask_padding = mask_padding # mask filler and batch padding tokens or not |
| 41 | self.average_upsampling = average_upsampling # zipvoice-style text late average upsampling (after text encoder) |
| 42 | if average_upsampling: |
| 43 | assert mask_padding, "text_embedding_average_upsampling requires text_mask_padding to be True" |
| 44 | |
| 45 | if conv_layers > 0: |
| 46 | self.extra_modeling = True |
| 47 | self.precompute_max_pos = 8192 # 8192 is ~87.38s of 24khz audio; 4096 is ~43.69s of 24khz audio |
| 48 | self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False) |
| 49 | self.text_blocks = nn.Sequential( |
| 50 | *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)] |
| 51 | ) |
| 52 | else: |
| 53 | self.extra_modeling = False |
| 54 | |
| 55 | def average_upsample_text_by_mask(self, text, text_mask, target_lens): |
| 56 | batch, max_seq_len, text_dim = text.shape |
| 57 | text_lens = text_mask.sum(dim=1) # [batch] |
| 58 | |
| 59 | upsampled_text = torch.zeros_like(text) |
| 60 | |
| 61 | for i in range(batch): |
| 62 | text_len = int(text_lens[i].item()) |
| 63 | audio_len = int(target_lens[i].item()) |
| 64 | |
| 65 | if text_len == 0 or audio_len <= 0: |
| 66 | continue |
| 67 | |
| 68 | valid_ind = torch.where(text_mask[i])[0] |
| 69 | valid_data = text[i, valid_ind, :] # [text_len, text_dim] |
| 70 | |
| 71 | base_repeat = audio_len // text_len |
| 72 | remainder = audio_len % text_len |
| 73 | |
| 74 | indices = [] |
| 75 | for j in range(text_len): |
| 76 | repeat_count = base_repeat + (1 if j >= text_len - remainder else 0) |
| 77 | indices.extend([j] * repeat_count) |
| 78 | |
| 79 | indices = torch.tensor(indices[:audio_len], device=text.device, dtype=torch.long) |
| 80 | upsampled = valid_data[indices] # [audio_len, text_dim] |
| 81 | |
| 82 | upsampled_text[i, :audio_len, :] = upsampled |
| 83 | |
| 84 | return upsampled_text |
| 85 | |
| 86 | def forward(self, text: int["b nt"], seq_len, drop_text=False): |
| 87 | text = text + 1 # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx() |
| 88 | valid_pos_mask = None |
| 89 | if torch.is_tensor(seq_len): |
| 90 | seq_len = seq_len.to(device=text.device, dtype=torch.long) |
| 91 | max_seq_len = int(seq_len.max().item()) |
| 92 | else: |
| 93 | max_seq_len = int(seq_len) |
| 94 | |
| 95 | text = text[:, :max_seq_len] # curtail if character tokens are more than the mel spec tokens |
| 96 | text = F.pad(text, (0, max_seq_len - text.shape[1]), value=0) |
| 97 | |
| 98 | if torch.is_tensor(seq_len): |
| 99 | seq_pos = torch.arange(max_seq_len, device=text.device).unsqueeze(0) |
| 100 | valid_pos_mask = seq_pos < seq_len.unsqueeze(1) |
| 101 | text = text.masked_fill(~valid_pos_mask, 0) |
| 102 | |
| 103 | if self.mask_padding: |
| 104 | text_mask = text == 0 |
| 105 | |
| 106 | if drop_text: # cfg for text |
| 107 | text = torch.zeros_like(text) |
| 108 | |
| 109 | text = self.text_embed(text) # b n -> b n d |
| 110 | if valid_pos_mask is not None: |
| 111 | # Keep short-sample tail strictly zero (equivalent to per-sample pad_sequence(..., 0)). |
| 112 | text = text.masked_fill(~valid_pos_mask.unsqueeze(-1), 0.0) |
| 113 | |
| 114 | # possible extra modeling |
| 115 | if self.extra_modeling: |
| 116 | # sinus pos emb; for variable seq lengths, only add positions within each sample's valid range. |
| 117 | freqs = self.freqs_cis[:max_seq_len, :] |
| 118 | if valid_pos_mask is not None: |
| 119 | freqs = freqs.unsqueeze(0) * valid_pos_mask.unsqueeze(-1).to(freqs.dtype) |
| 120 | text = text + freqs |
| 121 | |
| 122 | # convnextv2 blocks |
| 123 | if self.mask_padding: |
| 124 | text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0) |
| 125 | for block in self.text_blocks: |
| 126 | text = block(text) |
| 127 | text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0) |
| 128 | else: |
| 129 | text = self.text_blocks(text) |
| 130 | |
| 131 | if self.average_upsampling: |
| 132 | if torch.is_tensor(seq_len): |
| 133 | target_lens = seq_len.to(device=text.device, dtype=torch.long) |
| 134 | else: |
| 135 | target_lens = torch.full((text.shape[0],), int(seq_len), device=text.device, dtype=torch.long) |
| 136 | |
| 137 | text = self.average_upsample_text_by_mask(text, ~text_mask, target_lens) |
| 138 | |
| 139 | return text |
| 140 | |
| 141 | |
| 142 | # noised input audio and context mixing embedding |
| 143 | |
| 144 | |
| 145 | class InputEmbedding(nn.Module): |
| 146 | def __init__(self, mel_dim, text_dim, out_dim): |
| 147 | super().__init__() |
| 148 | self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim) |
| 149 | self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim) |
| 150 | |
| 151 | def forward( |
| 152 | self, |
| 153 | x: float["b n d"], |
| 154 | cond: float["b n d"], |
| 155 | text_embed: float["b n d"], |
| 156 | drop_audio_cond=False, |
| 157 | audio_mask: bool["b n"] | None = None, |
| 158 | ): |
| 159 | if drop_audio_cond: # cfg for cond audio |
| 160 | cond = torch.zeros_like(cond) |
| 161 | |
| 162 | x = self.proj(torch.cat((x, cond, text_embed), dim=-1)) |
| 163 | x = self.conv_pos_embed(x, mask=audio_mask) + x |
| 164 | return x |
| 165 | |
| 166 | |
| 167 | # Transformer backbone using DiT blocks |
| 168 | |
| 169 | |
| 170 | class DiT(nn.Module): |
| 171 | def __init__( |
| 172 | self, |
| 173 | *, |
| 174 | dim, |
| 175 | depth=8, |
| 176 | heads=8, |
| 177 | dim_head=64, |
| 178 | dropout=0.1, |
| 179 | ff_mult=4, |
| 180 | mel_dim=100, |
| 181 | text_num_embeds=256, |
| 182 | text_dim=None, |
| 183 | text_mask_padding=True, |
| 184 | text_embedding_average_upsampling=False, |
| 185 | qk_norm=None, |
| 186 | conv_layers=0, |
| 187 | pe_attn_head=None, |
| 188 | attn_backend="torch", # "torch" | "flash_attn" |
| 189 | attn_mask_enabled=False, |
| 190 | long_skip_connection=False, |
| 191 | checkpoint_activations=False, |
| 192 | ): |
| 193 | super().__init__() |
| 194 | |
| 195 | self.time_embed = TimestepEmbedding(dim) |
| 196 | if text_dim is None: |
| 197 | text_dim = mel_dim |
| 198 | self.text_embed = TextEmbedding( |
| 199 | text_num_embeds, |
| 200 | text_dim, |
| 201 | mask_padding=text_mask_padding, |
| 202 | average_upsampling=text_embedding_average_upsampling, |
| 203 | conv_layers=conv_layers, |
| 204 | ) |
| 205 | self.input_embed = InputEmbedding(mel_dim, text_dim, dim) |
| 206 | |
| 207 | self.rotary_embed = RotaryEmbedding(dim_head) |
| 208 | |
| 209 | self.dim = dim |
| 210 | self.depth = depth |
| 211 | |
| 212 | self.transformer_blocks = nn.ModuleList( |
| 213 | [ |
| 214 | DiTBlock( |
| 215 | dim=dim, |
| 216 | heads=heads, |
| 217 | dim_head=dim_head, |
| 218 | ff_mult=ff_mult, |
| 219 | dropout=dropout, |
| 220 | qk_norm=qk_norm, |
| 221 | pe_attn_head=pe_attn_head, |
| 222 | attn_backend=attn_backend, |
| 223 | attn_mask_enabled=attn_mask_enabled, |
| 224 | ) |
| 225 | for _ in range(depth) |
| 226 | ] |
| 227 | ) |
| 228 | self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None |
| 229 | |
| 230 | self.norm_out = AdaLayerNorm_Final(dim) # final modulation |
| 231 | self.proj_out = nn.Linear(dim, mel_dim) |
| 232 | |
| 233 | self.checkpoint_activations = checkpoint_activations |
| 234 | |
| 235 | self.initialize_weights() |
| 236 | |
| 237 | # `_cache_local` is lazily initialized on first inference-time cache write so that |
| 238 | # training models (which never touch the cache) stay deepcopy-friendly for EMA. |
| 239 | def _get_cache_local(self): |
| 240 | cache = self.__dict__.get("_cache_local") |
| 241 | if cache is None: |
| 242 | cache = threading.local() |
| 243 | self.__dict__["_cache_local"] = cache |
| 244 | return cache |
| 245 | |
| 246 | @property |
| 247 | def text_cond(self): |
| 248 | cache = self.__dict__.get("_cache_local") |
| 249 | return getattr(cache, "text_cond", None) if cache is not None else None |
| 250 | |
| 251 | @text_cond.setter |
| 252 | def text_cond(self, value): |
| 253 | self._get_cache_local().text_cond = value |
| 254 | |
| 255 | @property |
| 256 | def text_uncond(self): |
| 257 | cache = self.__dict__.get("_cache_local") |
| 258 | return getattr(cache, "text_uncond", None) if cache is not None else None |
| 259 | |
| 260 | @text_uncond.setter |
| 261 | def text_uncond(self, value): |
| 262 | self._get_cache_local().text_uncond = value |
| 263 | |
| 264 | def initialize_weights(self): |
| 265 | # Zero-out AdaLN layers in DiT blocks: |
| 266 | for block in self.transformer_blocks: |
| 267 | nn.init.constant_(block.attn_norm.linear.weight, 0) |
| 268 | nn.init.constant_(block.attn_norm.linear.bias, 0) |
| 269 | |
| 270 | # Zero-out output layers: |
| 271 | nn.init.constant_(self.norm_out.linear.weight, 0) |
| 272 | nn.init.constant_(self.norm_out.linear.bias, 0) |
| 273 | nn.init.constant_(self.proj_out.weight, 0) |
| 274 | nn.init.constant_(self.proj_out.bias, 0) |
| 275 | |
| 276 | def ckpt_wrapper(self, module): |
| 277 | # https://github.com/chuanyangjin/fast-DiT/blob/main/models.py |
| 278 | def ckpt_forward(*inputs): |
| 279 | outputs = module(*inputs) |
| 280 | return outputs |
| 281 | |
| 282 | return ckpt_forward |
| 283 | |
| 284 | def get_input_embed( |
| 285 | self, |
| 286 | x, # b n d |
| 287 | cond, # b n d |
| 288 | text, # b nt |
| 289 | drop_audio_cond: bool = False, |
| 290 | drop_text: bool = False, |
| 291 | cache: bool = True, |
| 292 | audio_mask: bool["b n"] | None = None, |
| 293 | ): |
| 294 | if self.text_uncond is None or self.text_cond is None or not cache: |
| 295 | if audio_mask is None: |
| 296 | seq_len = x.shape[1] |
| 297 | else: |
| 298 | seq_len = audio_mask.sum(dim=1) # per-sample valid speech length |
| 299 | text_embed = self.text_embed(text, seq_len=seq_len, drop_text=drop_text) |
| 300 | if cache: |
| 301 | if drop_text: |
| 302 | self.text_uncond = text_embed |
| 303 | else: |
| 304 | self.text_cond = text_embed |
| 305 | |
| 306 | if cache: |
| 307 | if drop_text: |
| 308 | text_embed = self.text_uncond |
| 309 | else: |
| 310 | text_embed = self.text_cond |
| 311 | |
| 312 | x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond, audio_mask=audio_mask) |
| 313 | |
| 314 | return x |
| 315 | |
| 316 | def clear_cache(self): |
| 317 | self.text_cond, self.text_uncond = None, None |
| 318 | |
| 319 | def forward( |
| 320 | self, |
| 321 | x: float["b n d"], # nosied input audio |
| 322 | cond: float["b n d"], # masked cond audio |
| 323 | text: int["b nt"], # text |
| 324 | time: float["b"] | float[""], # time step |
| 325 | mask: bool["b n"] | None = None, |
| 326 | drop_audio_cond: bool = False, # cfg for cond audio |
| 327 | drop_text: bool = False, # cfg for text |
| 328 | cfg_infer: bool = False, # cfg inference, pack cond & uncond forward |
| 329 | cache: bool = False, |
| 330 | ): |
| 331 | batch, seq_len = x.shape[0], x.shape[1] |
| 332 | if time.ndim == 0: |
| 333 | time = time.repeat(batch) |
| 334 | |
| 335 | # t: conditioning time, text: text, x: noised audio + cond audio + text |
| 336 | t = self.time_embed(time) |
| 337 | if cfg_infer: # pack cond & uncond forward: b n d -> 2b n d |
| 338 | x_cond = self.get_input_embed( |
| 339 | x, cond, text, drop_audio_cond=False, drop_text=False, cache=cache, audio_mask=mask |
| 340 | ) |
| 341 | x_uncond = self.get_input_embed( |
| 342 | x, cond, text, drop_audio_cond=True, drop_text=True, cache=cache, audio_mask=mask |
| 343 | ) |
| 344 | x = torch.cat((x_cond, x_uncond), dim=0) |
| 345 | t = torch.cat((t, t), dim=0) |
| 346 | mask = torch.cat((mask, mask), dim=0) if mask is not None else None |
| 347 | else: |
| 348 | x = self.get_input_embed( |
| 349 | x, cond, text, drop_audio_cond=drop_audio_cond, drop_text=drop_text, cache=cache, audio_mask=mask |
| 350 | ) |
| 351 | |
| 352 | rope = self.rotary_embed.forward_from_seq_len(seq_len) |
| 353 | |
| 354 | if self.long_skip_connection is not None: |
| 355 | residual = x |
| 356 | |
| 357 | for block in self.transformer_blocks: |
| 358 | if self.checkpoint_activations: |
| 359 | # https://pytorch.org/docs/stable/checkpoint.html#torch.utils.checkpoint.checkpoint |
| 360 | x = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(block), x, t, mask, rope, use_reentrant=False) |
| 361 | else: |
| 362 | x = block(x, t, mask=mask, rope=rope) |
| 363 | |
| 364 | if self.long_skip_connection is not None: |
| 365 | x = self.long_skip_connection(torch.cat((x, residual), dim=-1)) |
| 366 | |
| 367 | x = self.norm_out(x, t) |
| 368 | output = self.proj_out(x) |
| 369 | |
| 370 | return output |
| 371 |