| 1 | from __future__ import annotations |
| 2 | |
| 3 | import os |
| 4 | import sys |
| 5 | from collections import OrderedDict |
| 6 | |
| 7 | import numpy as np |
| 8 | import tensorrt as trt |
| 9 | from tensorrt_llm._common import default_net |
| 10 | |
| 11 | from ..._utils import str_dtype_to_trt |
| 12 | from ...functional import ( |
| 13 | Tensor, |
| 14 | concat, |
| 15 | constant, |
| 16 | expand, |
| 17 | shape, |
| 18 | slice, |
| 19 | unsqueeze, |
| 20 | ) |
| 21 | from ...layers import Linear |
| 22 | from ...module import Module, ModuleList |
| 23 | from ...plugin import current_all_reduce_helper |
| 24 | from ..modeling_utils import PretrainedConfig, PretrainedModel |
| 25 | from .modules import AdaLayerNormZero_Final, ConvPositionEmbedding, DiTBlock, TimestepEmbedding |
| 26 | |
| 27 | |
| 28 | current_file_path = os.path.abspath(__file__) |
| 29 | parent_dir = os.path.dirname(current_file_path) |
| 30 | sys.path.append(parent_dir) |
| 31 | |
| 32 | |
| 33 | class InputEmbedding(Module): |
| 34 | def __init__(self, mel_dim, text_dim, out_dim): |
| 35 | super().__init__() |
| 36 | self.proj = Linear(mel_dim * 2 + text_dim, out_dim) |
| 37 | self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim) |
| 38 | |
| 39 | def forward(self, x, cond, mask=None): |
| 40 | x = self.proj(concat([x, cond], dim=-1)) |
| 41 | return self.conv_pos_embed(x, mask=mask) + x |
| 42 | |
| 43 | |
| 44 | class F5TTS(PretrainedModel): |
| 45 | def __init__(self, config: PretrainedConfig): |
| 46 | super().__init__(config) |
| 47 | self.dtype = str_dtype_to_trt(config.dtype) |
| 48 | |
| 49 | self.time_embed = TimestepEmbedding(config.hidden_size) |
| 50 | self.input_embed = InputEmbedding(config.mel_dim, config.text_dim, config.hidden_size) |
| 51 | |
| 52 | self.dim = config.hidden_size |
| 53 | self.depth = config.num_hidden_layers |
| 54 | self.transformer_blocks = ModuleList( |
| 55 | [ |
| 56 | DiTBlock( |
| 57 | dim=self.dim, |
| 58 | heads=config.num_attention_heads, |
| 59 | dim_head=config.dim_head, |
| 60 | ff_mult=config.ff_mult, |
| 61 | dropout=config.dropout, |
| 62 | pe_attn_head=config.pe_attn_head, |
| 63 | ) |
| 64 | for _ in range(self.depth) |
| 65 | ] |
| 66 | ) |
| 67 | |
| 68 | self.norm_out = AdaLayerNormZero_Final(config.hidden_size) # final modulation |
| 69 | self.proj_out = Linear(config.hidden_size, config.mel_dim) |
| 70 | |
| 71 | def forward( |
| 72 | self, |
| 73 | noise, # nosied input audio |
| 74 | cond, # masked cond audio |
| 75 | time, # time step |
| 76 | rope_cos, |
| 77 | rope_sin, |
| 78 | input_lengths, |
| 79 | scale=1.0, |
| 80 | ): |
| 81 | if default_net().plugin_config.remove_input_padding: |
| 82 | mask = None |
| 83 | else: |
| 84 | N = shape(noise, 1) |
| 85 | B = shape(noise, 0) |
| 86 | seq_len_2d = concat([1, N]) |
| 87 | max_position_embeddings = 4096 |
| 88 | # create position ids |
| 89 | position_ids_buffer = constant(np.expand_dims(np.arange(max_position_embeddings).astype(np.int32), 0)) |
| 90 | tmp_position_ids = slice(position_ids_buffer, starts=[0, 0], sizes=seq_len_2d) |
| 91 | tmp_position_ids = expand(tmp_position_ids, concat([B, N])) # [B, N] |
| 92 | tmp_input_lengths = unsqueeze(input_lengths, 1) # [B, 1] |
| 93 | tmp_input_lengths = expand(tmp_input_lengths, concat([B, N])) # [B, N] |
| 94 | mask = tmp_position_ids < tmp_input_lengths # [B, N] |
| 95 | mask = mask.cast("int32") |
| 96 | |
| 97 | t = self.time_embed(time) |
| 98 | x = self.input_embed(noise, cond, mask=mask) |
| 99 | for block in self.transformer_blocks: |
| 100 | x = block(x, t, rope_cos=rope_cos, rope_sin=rope_sin, input_lengths=input_lengths, scale=scale, mask=mask) |
| 101 | denoise = self.proj_out(self.norm_out(x, t)) |
| 102 | denoise.mark_output("denoised", self.dtype) |
| 103 | return denoise |
| 104 | |
| 105 | def prepare_inputs(self, **kwargs): |
| 106 | max_batch_size = kwargs["max_batch_size"] |
| 107 | batch_size_range = [2, 2, max_batch_size] |
| 108 | mel_size = self.config.mel_dim |
| 109 | max_seq_len = 3000 # 4096 |
| 110 | num_frames_range = [mel_size * 2, max_seq_len * 2, max_seq_len * max_batch_size] |
| 111 | concat_feature_dim = mel_size + self.config.text_dim |
| 112 | freq_embed_dim = 256 # Warning: hard coding 256 here |
| 113 | head_dim = self.config.dim_head |
| 114 | mapping = self.config.mapping |
| 115 | if mapping.tp_size > 1: |
| 116 | current_all_reduce_helper().set_workspace_tensor(mapping, 1) |
| 117 | if default_net().plugin_config.remove_input_padding: |
| 118 | noise = Tensor( |
| 119 | name="noise", |
| 120 | dtype=self.dtype, |
| 121 | shape=[-1, mel_size], |
| 122 | dim_range=OrderedDict( |
| 123 | [ |
| 124 | ("num_frames", [num_frames_range]), |
| 125 | ("n_mels", [mel_size]), |
| 126 | ] |
| 127 | ), |
| 128 | ) |
| 129 | cond = Tensor( |
| 130 | name="cond", |
| 131 | dtype=self.dtype, |
| 132 | shape=[-1, concat_feature_dim], |
| 133 | dim_range=OrderedDict( |
| 134 | [ |
| 135 | ("num_frames", [num_frames_range]), |
| 136 | ("embeded_length", [concat_feature_dim]), |
| 137 | ] |
| 138 | ), |
| 139 | ) |
| 140 | time = Tensor( |
| 141 | name="time", |
| 142 | dtype=self.dtype, |
| 143 | shape=[-1, freq_embed_dim], |
| 144 | dim_range=OrderedDict( |
| 145 | [ |
| 146 | ("num_frames", [num_frames_range]), |
| 147 | ("freq_dim", [freq_embed_dim]), |
| 148 | ] |
| 149 | ), |
| 150 | ) |
| 151 | rope_cos = Tensor( |
| 152 | name="rope_cos", |
| 153 | dtype=self.dtype, |
| 154 | shape=[-1, head_dim], |
| 155 | dim_range=OrderedDict( |
| 156 | [ |
| 157 | ("num_frames", [num_frames_range]), |
| 158 | ("head_dim", [head_dim]), |
| 159 | ] |
| 160 | ), |
| 161 | ) |
| 162 | rope_sin = Tensor( |
| 163 | name="rope_sin", |
| 164 | dtype=self.dtype, |
| 165 | shape=[-1, head_dim], |
| 166 | dim_range=OrderedDict( |
| 167 | [ |
| 168 | ("num_frames", [num_frames_range]), |
| 169 | ("head_dim", [head_dim]), |
| 170 | ] |
| 171 | ), |
| 172 | ) |
| 173 | |
| 174 | else: |
| 175 | noise = Tensor( |
| 176 | name="noise", |
| 177 | dtype=self.dtype, |
| 178 | shape=[-1, -1, mel_size], |
| 179 | dim_range=OrderedDict( |
| 180 | [ |
| 181 | ("batch_size", [batch_size_range]), |
| 182 | ("max_duratuion", [[100, max_seq_len // 2, max_seq_len]]), |
| 183 | ("n_mels", [mel_size]), |
| 184 | ] |
| 185 | ), |
| 186 | ) |
| 187 | cond = Tensor( |
| 188 | name="cond", |
| 189 | dtype=self.dtype, |
| 190 | shape=[-1, -1, concat_feature_dim], |
| 191 | dim_range=OrderedDict( |
| 192 | [ |
| 193 | ("batch_size", [batch_size_range]), |
| 194 | ("max_duratuion", [[100, max_seq_len // 2, max_seq_len]]), |
| 195 | ("embeded_length", [concat_feature_dim]), |
| 196 | ] |
| 197 | ), |
| 198 | ) |
| 199 | time = Tensor( |
| 200 | name="time", |
| 201 | dtype=self.dtype, |
| 202 | shape=[-1, freq_embed_dim], |
| 203 | dim_range=OrderedDict( |
| 204 | [ |
| 205 | ("batch_size", [batch_size_range]), |
| 206 | ("freq_dim", [freq_embed_dim]), |
| 207 | ] |
| 208 | ), |
| 209 | ) |
| 210 | rope_cos = Tensor( |
| 211 | name="rope_cos", |
| 212 | dtype=self.dtype, |
| 213 | shape=[-1, -1, head_dim], |
| 214 | dim_range=OrderedDict( |
| 215 | [ |
| 216 | ("batch_size", [batch_size_range]), |
| 217 | ("max_duratuion", [[100, max_seq_len // 2, max_seq_len]]), |
| 218 | ("head_dim", [head_dim]), |
| 219 | ] |
| 220 | ), |
| 221 | ) |
| 222 | rope_sin = Tensor( |
| 223 | name="rope_sin", |
| 224 | dtype=self.dtype, |
| 225 | shape=[-1, -1, head_dim], |
| 226 | dim_range=OrderedDict( |
| 227 | [ |
| 228 | ("batch_size", [batch_size_range]), |
| 229 | ("max_duratuion", [[100, max_seq_len // 2, max_seq_len]]), |
| 230 | ("head_dim", [head_dim]), |
| 231 | ] |
| 232 | ), |
| 233 | ) |
| 234 | input_lengths = Tensor( |
| 235 | name="input_lengths", |
| 236 | dtype=trt.int32, |
| 237 | shape=[-1], |
| 238 | dim_range=OrderedDict([("batch_size", [batch_size_range])]), |
| 239 | ) |
| 240 | return { |
| 241 | "noise": noise, |
| 242 | "cond": cond, |
| 243 | "time": time, |
| 244 | "rope_cos": rope_cos, |
| 245 | "rope_sin": rope_sin, |
| 246 | "input_lengths": input_lengths, |
| 247 | } |
| 248 |