| 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 | from torch import nn |
| 17 | from x_transformers.x_transformers import RotaryEmbedding |
| 18 | |
| 19 | from f5_tts.model.modules import ( |
| 20 | AdaLayerNorm_Final, |
| 21 | ConvPositionEmbedding, |
| 22 | MMDiTBlock, |
| 23 | TimestepEmbedding, |
| 24 | get_pos_embed_indices, |
| 25 | precompute_freqs_cis, |
| 26 | ) |
| 27 | |
| 28 | |
| 29 | # text embedding |
| 30 | |
| 31 | |
| 32 | class TextEmbedding(nn.Module): |
| 33 | def __init__(self, out_dim, text_num_embeds, mask_padding=True): |
| 34 | super().__init__() |
| 35 | self.text_embed = nn.Embedding(text_num_embeds + 1, out_dim) # will use 0 as filler token |
| 36 | |
| 37 | self.mask_padding = mask_padding # mask filler and batch padding tokens or not |
| 38 | |
| 39 | self.precompute_max_pos = 1024 |
| 40 | self.register_buffer("freqs_cis", precompute_freqs_cis(out_dim, self.precompute_max_pos), persistent=False) |
| 41 | |
| 42 | def forward(self, text: int["b nt"], drop_text=False) -> int["b nt d"]: |
| 43 | text = text + 1 # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx() |
| 44 | if self.mask_padding: |
| 45 | text_mask = text == 0 |
| 46 | |
| 47 | if drop_text: # cfg for text |
| 48 | text = torch.zeros_like(text) |
| 49 | |
| 50 | text = self.text_embed(text) # b nt -> b nt d |
| 51 | |
| 52 | # sinus pos emb |
| 53 | batch_start = torch.zeros((text.shape[0],), dtype=torch.long) |
| 54 | batch_text_len = text.shape[1] |
| 55 | pos_idx = get_pos_embed_indices(batch_start, batch_text_len, max_pos=self.precompute_max_pos) |
| 56 | text_pos_embed = self.freqs_cis[pos_idx] |
| 57 | |
| 58 | text = text + text_pos_embed |
| 59 | |
| 60 | if self.mask_padding: |
| 61 | text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0) |
| 62 | |
| 63 | return text |
| 64 | |
| 65 | |
| 66 | # noised input & masked cond audio embedding |
| 67 | |
| 68 | |
| 69 | class AudioEmbedding(nn.Module): |
| 70 | def __init__(self, in_dim, out_dim): |
| 71 | super().__init__() |
| 72 | self.linear = nn.Linear(2 * in_dim, out_dim) |
| 73 | self.conv_pos_embed = ConvPositionEmbedding(out_dim) |
| 74 | |
| 75 | def forward(self, x: float["b n d"], cond: float["b n d"], drop_audio_cond=False): |
| 76 | if drop_audio_cond: |
| 77 | cond = torch.zeros_like(cond) |
| 78 | x = torch.cat((x, cond), dim=-1) |
| 79 | x = self.linear(x) |
| 80 | x = self.conv_pos_embed(x) + x |
| 81 | return x |
| 82 | |
| 83 | |
| 84 | # Transformer backbone using MM-DiT blocks |
| 85 | |
| 86 | |
| 87 | class MMDiT(nn.Module): |
| 88 | def __init__( |
| 89 | self, |
| 90 | *, |
| 91 | dim, |
| 92 | depth=8, |
| 93 | heads=8, |
| 94 | dim_head=64, |
| 95 | dropout=0.1, |
| 96 | ff_mult=4, |
| 97 | mel_dim=100, |
| 98 | text_num_embeds=256, |
| 99 | text_mask_padding=True, |
| 100 | qk_norm=None, |
| 101 | checkpoint_activations=False, |
| 102 | attn_backend="torch", |
| 103 | attn_mask_enabled=False, |
| 104 | ): |
| 105 | super().__init__() |
| 106 | |
| 107 | self.time_embed = TimestepEmbedding(dim) |
| 108 | self.text_embed = TextEmbedding(dim, text_num_embeds, mask_padding=text_mask_padding) |
| 109 | self.audio_embed = AudioEmbedding(mel_dim, dim) |
| 110 | |
| 111 | self.rotary_embed = RotaryEmbedding(dim_head) |
| 112 | |
| 113 | self.dim = dim |
| 114 | self.depth = depth |
| 115 | |
| 116 | self.transformer_blocks = nn.ModuleList( |
| 117 | [ |
| 118 | MMDiTBlock( |
| 119 | dim=dim, |
| 120 | heads=heads, |
| 121 | dim_head=dim_head, |
| 122 | dropout=dropout, |
| 123 | ff_mult=ff_mult, |
| 124 | context_pre_only=i == depth - 1, |
| 125 | qk_norm=qk_norm, |
| 126 | attn_backend=attn_backend, |
| 127 | attn_mask_enabled=attn_mask_enabled, |
| 128 | ) |
| 129 | for i in range(depth) |
| 130 | ] |
| 131 | ) |
| 132 | self.norm_out = AdaLayerNorm_Final(dim) # final modulation |
| 133 | self.proj_out = nn.Linear(dim, mel_dim) |
| 134 | |
| 135 | self.checkpoint_activations = checkpoint_activations |
| 136 | |
| 137 | self.initialize_weights() |
| 138 | |
| 139 | # `_cache_local` is lazily initialized on first inference-time cache write so that |
| 140 | # training models (which never touch the cache) stay deepcopy-friendly for EMA. |
| 141 | def _get_cache_local(self): |
| 142 | cache = self.__dict__.get("_cache_local") |
| 143 | if cache is None: |
| 144 | cache = threading.local() |
| 145 | self.__dict__["_cache_local"] = cache |
| 146 | return cache |
| 147 | |
| 148 | @property |
| 149 | def text_cond(self): |
| 150 | cache = self.__dict__.get("_cache_local") |
| 151 | return getattr(cache, "text_cond", None) if cache is not None else None |
| 152 | |
| 153 | @text_cond.setter |
| 154 | def text_cond(self, value): |
| 155 | self._get_cache_local().text_cond = value |
| 156 | |
| 157 | @property |
| 158 | def text_uncond(self): |
| 159 | cache = self.__dict__.get("_cache_local") |
| 160 | return getattr(cache, "text_uncond", None) if cache is not None else None |
| 161 | |
| 162 | @text_uncond.setter |
| 163 | def text_uncond(self, value): |
| 164 | self._get_cache_local().text_uncond = value |
| 165 | |
| 166 | def initialize_weights(self): |
| 167 | # Zero-out AdaLN layers in MMDiT blocks: |
| 168 | for block in self.transformer_blocks: |
| 169 | nn.init.constant_(block.attn_norm_x.linear.weight, 0) |
| 170 | nn.init.constant_(block.attn_norm_x.linear.bias, 0) |
| 171 | nn.init.constant_(block.attn_norm_c.linear.weight, 0) |
| 172 | nn.init.constant_(block.attn_norm_c.linear.bias, 0) |
| 173 | |
| 174 | # Zero-out output layers: |
| 175 | nn.init.constant_(self.norm_out.linear.weight, 0) |
| 176 | nn.init.constant_(self.norm_out.linear.bias, 0) |
| 177 | nn.init.constant_(self.proj_out.weight, 0) |
| 178 | nn.init.constant_(self.proj_out.bias, 0) |
| 179 | |
| 180 | def ckpt_wrapper(self, module): |
| 181 | def ckpt_forward(*inputs): |
| 182 | outputs = module(*inputs) |
| 183 | return outputs |
| 184 | |
| 185 | return ckpt_forward |
| 186 | |
| 187 | def get_input_embed( |
| 188 | self, |
| 189 | x, # b n d |
| 190 | cond, # b n d |
| 191 | text, # b nt |
| 192 | drop_audio_cond: bool = False, |
| 193 | drop_text: bool = False, |
| 194 | cache: bool = True, |
| 195 | ): |
| 196 | if cache: |
| 197 | if drop_text: |
| 198 | if self.text_uncond is None: |
| 199 | self.text_uncond = self.text_embed(text, drop_text=True) |
| 200 | c = self.text_uncond |
| 201 | else: |
| 202 | if self.text_cond is None: |
| 203 | self.text_cond = self.text_embed(text, drop_text=False) |
| 204 | c = self.text_cond |
| 205 | else: |
| 206 | c = self.text_embed(text, drop_text=drop_text) |
| 207 | x = self.audio_embed(x, cond, drop_audio_cond=drop_audio_cond) |
| 208 | |
| 209 | return x, c |
| 210 | |
| 211 | def clear_cache(self): |
| 212 | self.text_cond, self.text_uncond = None, None |
| 213 | |
| 214 | def forward( |
| 215 | self, |
| 216 | x: float["b n d"], # nosied input audio |
| 217 | cond: float["b n d"], # masked cond audio |
| 218 | text: int["b nt"], # text |
| 219 | time: float["b"] | float[""], # time step |
| 220 | mask: bool["b n"] | None = None, |
| 221 | drop_audio_cond: bool = False, # cfg for cond audio |
| 222 | drop_text: bool = False, # cfg for text |
| 223 | cfg_infer: bool = False, # cfg inference, pack cond & uncond forward |
| 224 | cache: bool = False, |
| 225 | ): |
| 226 | batch = x.shape[0] |
| 227 | if time.ndim == 0: |
| 228 | time = time.repeat(batch) |
| 229 | |
| 230 | # t: conditioning (time), c: context (text + masked cond audio), x: noised input audio |
| 231 | t = self.time_embed(time) |
| 232 | c_mask = (text + 1) != 0 # True = valid, False = padding (-1 tokens) |
| 233 | if cfg_infer: # pack cond & uncond forward: b n d -> 2b n d |
| 234 | x_cond, c_cond = self.get_input_embed(x, cond, text, drop_audio_cond=False, drop_text=False, cache=cache) |
| 235 | x_uncond, c_uncond = self.get_input_embed(x, cond, text, drop_audio_cond=True, drop_text=True, cache=cache) |
| 236 | x = torch.cat((x_cond, x_uncond), dim=0) |
| 237 | c = torch.cat((c_cond, c_uncond), dim=0) |
| 238 | t = torch.cat((t, t), dim=0) |
| 239 | mask = torch.cat((mask, mask), dim=0) if mask is not None else None |
| 240 | c_mask = torch.cat((c_mask, c_mask), dim=0) |
| 241 | else: |
| 242 | x, c = self.get_input_embed( |
| 243 | x, cond, text, drop_audio_cond=drop_audio_cond, drop_text=drop_text, cache=cache |
| 244 | ) |
| 245 | |
| 246 | seq_len = x.shape[1] |
| 247 | text_len = text.shape[1] |
| 248 | rope_audio = self.rotary_embed.forward_from_seq_len(seq_len) |
| 249 | rope_text = self.rotary_embed.forward_from_seq_len(text_len) |
| 250 | |
| 251 | for block in self.transformer_blocks: |
| 252 | if self.checkpoint_activations: |
| 253 | c, x = torch.utils.checkpoint.checkpoint( |
| 254 | self.ckpt_wrapper(block), x, c, t, mask, rope_audio, rope_text, c_mask, use_reentrant=False |
| 255 | ) |
| 256 | else: |
| 257 | c, x = block(x, c, t, mask=mask, rope=rope_audio, c_rope=rope_text, c_mask=c_mask) |
| 258 | |
| 259 | x = self.norm_out(x, t) |
| 260 | output = self.proj_out(x) |
| 261 | |
| 262 | return output |
| 263 |