gitea-mirror/F5-TTS
1
0
Fork 0
mirror of https://github.com/SWivid/F5-TTS.git synced 2025-12-12 07:40:43 -08:00
Code Issues Packages Projects Releases Wiki Activity

runtime trtllm: fix batch inference skipping last words in shorter sentences #1039 #1179

Browse Source
This commit is contained in:
SWivid
2025-10-24 09:12:08 +00:00
parent 6b07fb03b2
commit e67d50841e
3 changed files with 54 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
src/f5_tts/runtime/triton_trtllm/patch/f5tts/model.py
View File

@@ -4,11 +4,20 @@ import os
import sys
from collections import OrderedDict
import numpy as np
import tensorrt as trt
from tensorrt_llm._common import default_net
from ..._utils import str_dtype_to_trt
from ...functional import Tensor, concat
from ...functional import (
Tensor,
concat,
constant,
expand,
shape,
slice,
unsqueeze,
)
from ...layers import Linear
from ...module import Module, ModuleList
from ...plugin import current_all_reduce_helper
@@ -27,9 +36,9 @@ class InputEmbedding(Module):
self.proj = Linear(mel_dim * 2 + text_dim, out_dim)
self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)
def forward(self, x, cond):
def forward(self, x, cond, mask=None):
x = self.proj(concat([x, cond], dim=-1))
return self.conv_pos_embed(x) + x
return self.conv_pos_embed(x, mask=mask) + x
class F5TTS(PretrainedModel):
@@ -69,10 +78,26 @@ class F5TTS(PretrainedModel):
input_lengths,
scale=1.0,
):
if default_net().plugin_config.remove_input_padding:
mask = None
else:
N = shape(noise, 1)
B = shape(noise, 0)
seq_len_2d = concat([1, N])
max_position_embeddings = 4096
# create position ids
position_ids_buffer = constant(np.expand_dims(np.arange(max_position_embeddings).astype(np.int32), 0))
tmp_position_ids = slice(position_ids_buffer, starts=[0, 0], sizes=seq_len_2d)
tmp_position_ids = expand(tmp_position_ids, concat([B, N])) # [B, N]
tmp_input_lengths = unsqueeze(input_lengths, 1) # [B, 1]
tmp_input_lengths = expand(tmp_input_lengths, concat([B, N])) # [B, N]
mask = tmp_position_ids < tmp_input_lengths # [B, N]
mask = mask.cast("int32")
t = self.time_embed(time)
x = self.input_embed(noise, cond)
x = self.input_embed(noise, cond, mask=mask)
for block in self.transformer_blocks:
x = block(x, t, rope_cos=rope_cos, rope_sin=rope_sin, input_lengths=input_lengths, scale=scale)
x = block(x, t, rope_cos=rope_cos, rope_sin=rope_sin, input_lengths=input_lengths, scale=scale, mask=mask)
denoise = self.proj_out(self.norm_out(x, t))
denoise.mark_output("denoised", self.dtype)
return denoise

47
src/f5_tts/runtime/triton_trtllm/patch/f5tts/modules.py
View File

@@ -16,7 +16,6 @@ from ...functional import (
chunk,
concat,
constant,
expand,
expand_dims,
expand_dims_like,
expand_mask,
@@ -95,15 +94,24 @@ class ConvPositionEmbedding(Module):
self.conv1d2 = Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2)
self.mish = Mish()
def forward(self, x, mask=None): # noqa: F722
def forward(self, x, mask=None):
if default_net().plugin_config.remove_input_padding:
x = unsqueeze(x, 0)
x = permute(x, [0, 2, 1])
x = self.mish(self.conv1d2(self.mish(self.conv1d1(x))))
out = permute(x, [0, 2, 1])
if mask is not None:
mask = mask.view(concat([shape(mask, 0), 1, shape(mask, 1)])) # [B 1 N]
mask = expand_dims_like(mask, x) # [B D N]
mask = cast(mask, x.dtype)
x = permute(x, [0, 2, 1]) # [B D N]
if mask is not None:
x = self.mish(self.conv1d2(self.mish(self.conv1d1(x * mask) * mask)) * mask)
else:
x = self.mish(self.conv1d2(self.mish(self.conv1d1(x))))
x = permute(x, [0, 2, 1]) # [B N D]
if default_net().plugin_config.remove_input_padding:
out = squeeze(out, 0)
return out
x = squeeze(x, 0)
return x
class Attention(Module):
@@ -185,6 +193,7 @@ class Attention(Module):
rope_cos,
rope_sin,
input_lengths,
mask=None,
c=None, # context c
scale=1.0,
rope=None,
@@ -283,6 +292,7 @@ class AttnProcessor:
input_lengths,
scale=1.0,
rope=None,
mask=None,
) -> torch.FloatTensor:
query = attn.to_q(x)
key = attn.to_k(x)
@@ -295,20 +305,8 @@ class AttnProcessor:
inner_dim = key.shape[-1]
norm_factor = math.sqrt(attn.attention_head_size)
q_scaling = 1.0 / norm_factor
mask = None
if not default_net().plugin_config.remove_input_padding:
N = shape(x, 1)
B = shape(x, 0)
seq_len_2d = concat([1, N])
max_position_embeddings = 4096
# create position ids
position_ids_buffer = constant(np.expand_dims(np.arange(max_position_embeddings).astype(np.int32), 0))
tmp_position_ids = slice(position_ids_buffer, starts=[0, 0], sizes=seq_len_2d)
tmp_position_ids = expand(tmp_position_ids, concat([B, N])) # BxL
tmp_input_lengths = unsqueeze(input_lengths, 1) # Bx1
tmp_input_lengths = expand(tmp_input_lengths, concat([B, N])) # BxL
mask = tmp_position_ids < tmp_input_lengths # BxL
mask = mask.cast("int32")
if default_net().plugin_config.remove_input_padding:
mask = None
if default_net().plugin_config.bert_attention_plugin:
qkv = concat([query, key, value], dim=-1)
@@ -393,14 +391,15 @@ class DiTBlock(Module):
self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout)
def forward(
self, x, t, rope_cos, rope_sin, input_lengths, scale=1.0, rope=ModuleNotFoundError
self, x, t, rope_cos, rope_sin, input_lengths, scale=1.0, rope=ModuleNotFoundError, mask=None
): # x: noised input, t: time embedding
# pre-norm & modulation for attention input
norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)
# attention
# norm ----> (2,1226,1024)
attn_output = self.attn(x=norm, rope_cos=rope_cos, rope_sin=rope_sin, input_lengths=input_lengths, scale=scale)
attn_output = self.attn(
x=norm, rope_cos=rope_cos, rope_sin=rope_sin, input_lengths=input_lengths, scale=scale, mask=mask
)
# process attention output for input x
if default_net().plugin_config.remove_input_padding:
x = x + gate_msa * attn_output

2
src/f5_tts/runtime/triton_trtllm/run.sh
View File

@@ -73,7 +73,7 @@ fi
if [ $stage -le 7 ] && [ $stop_stage -ge 7 ]; then
echo "TRT-LLM: offline decoding benchmark test"
batch_size=1
batch_size=2
split_name=wenetspeech4tts
backend_type=trt
log_dir=./tests/benchmark_${model}_batch_size_${batch_size}_${split_name}_${backend_type}