| 1 | import math |
| 2 | import os |
| 3 | import time |
| 4 | from functools import wraps |
| 5 | from typing import List, Optional |
| 6 | |
| 7 | import tensorrt as trt |
| 8 | import tensorrt_llm |
| 9 | import torch |
| 10 | import torch.nn as nn |
| 11 | import torch.nn.functional as F |
| 12 | from tensorrt_llm._utils import str_dtype_to_torch, trt_dtype_to_torch |
| 13 | from tensorrt_llm.logger import logger |
| 14 | from tensorrt_llm.runtime.session import Session |
| 15 | from torch.nn.utils.rnn import pad_sequence |
| 16 | |
| 17 | |
| 18 | def remove_tensor_padding(input_tensor, input_tensor_lengths=None): |
| 19 | # Audio tensor case: batch, seq_len, feature_len |
| 20 | # position_ids case: batch, seq_len |
| 21 | assert input_tensor_lengths is not None, "input_tensor_lengths must be provided for 3D input_tensor" |
| 22 | |
| 23 | # Initialize a list to collect valid sequences |
| 24 | valid_sequences = [] |
| 25 | |
| 26 | for i in range(input_tensor.shape[0]): |
| 27 | valid_length = input_tensor_lengths[i] |
| 28 | valid_sequences.append(input_tensor[i, :valid_length]) |
| 29 | |
| 30 | # Concatenate all valid sequences along the batch dimension |
| 31 | output_tensor = torch.cat(valid_sequences, dim=0).contiguous() |
| 32 | return output_tensor |
| 33 | |
| 34 | |
| 35 | class TextEmbedding(nn.Module): |
| 36 | def __init__( |
| 37 | self, text_num_embeds, text_dim, mask_padding=True, conv_layers=0, conv_mult=2, precompute_max_pos=4096 |
| 38 | ): |
| 39 | super().__init__() |
| 40 | self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim) # use 0 as filler token |
| 41 | self.mask_padding = mask_padding |
| 42 | self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, precompute_max_pos), persistent=False) |
| 43 | self.text_blocks = nn.Sequential(*[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]) |
| 44 | |
| 45 | def forward(self, text, seq_len, drop_text=False): |
| 46 | text = text + 1 |
| 47 | text = text[:, :seq_len] # curtail if character tokens are more than the mel spec tokens |
| 48 | text = F.pad(text, (0, seq_len - text.shape[1]), value=0) |
| 49 | if self.mask_padding: |
| 50 | text_mask = text == 0 |
| 51 | |
| 52 | if drop_text: # cfg for text |
| 53 | text = torch.zeros_like(text) |
| 54 | |
| 55 | text = self.text_embed(text) # b n -> b n d |
| 56 | text = text + self.freqs_cis[:seq_len, :] |
| 57 | if self.mask_padding: |
| 58 | text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0) |
| 59 | for block in self.text_blocks: |
| 60 | text = block(text) |
| 61 | text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0) |
| 62 | else: |
| 63 | text = self.text_blocks(text) |
| 64 | |
| 65 | return text |
| 66 | |
| 67 | |
| 68 | class GRN(nn.Module): |
| 69 | def __init__(self, dim): |
| 70 | super().__init__() |
| 71 | self.gamma = nn.Parameter(torch.zeros(1, 1, dim)) |
| 72 | self.beta = nn.Parameter(torch.zeros(1, 1, dim)) |
| 73 | |
| 74 | def forward(self, x): |
| 75 | Gx = torch.norm(x, p=2, dim=1, keepdim=True) |
| 76 | Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) |
| 77 | return self.gamma * (x * Nx) + self.beta + x |
| 78 | |
| 79 | |
| 80 | class ConvNeXtV2Block(nn.Module): |
| 81 | def __init__( |
| 82 | self, |
| 83 | dim: int, |
| 84 | intermediate_dim: int, |
| 85 | dilation: int = 1, |
| 86 | ): |
| 87 | super().__init__() |
| 88 | padding = (dilation * (7 - 1)) // 2 |
| 89 | self.dwconv = nn.Conv1d( |
| 90 | dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation |
| 91 | ) # depthwise conv |
| 92 | self.norm = nn.LayerNorm(dim, eps=1e-6) |
| 93 | self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers |
| 94 | self.act = nn.GELU() |
| 95 | self.grn = GRN(intermediate_dim) |
| 96 | self.pwconv2 = nn.Linear(intermediate_dim, dim) |
| 97 | |
| 98 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 99 | residual = x |
| 100 | x = x.transpose(1, 2) # b n d -> b d n |
| 101 | x = self.dwconv(x) |
| 102 | x = x.transpose(1, 2) # b d n -> b n d |
| 103 | x = self.norm(x) |
| 104 | x = self.pwconv1(x) |
| 105 | x = self.act(x) |
| 106 | x = self.grn(x) |
| 107 | x = self.pwconv2(x) |
| 108 | return residual + x |
| 109 | |
| 110 | |
| 111 | def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0): |
| 112 | # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning |
| 113 | # has some connection to NTK literature |
| 114 | # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/ |
| 115 | # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py |
| 116 | theta *= theta_rescale_factor ** (dim / (dim - 2)) |
| 117 | freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) |
| 118 | t = torch.arange(end, device=freqs.device) # type: ignore |
| 119 | freqs = torch.outer(t, freqs).float() # type: ignore |
| 120 | freqs_cos = torch.cos(freqs) # real part |
| 121 | freqs_sin = torch.sin(freqs) # imaginary part |
| 122 | return torch.cat([freqs_cos, freqs_sin], dim=-1) |
| 123 | |
| 124 | |
| 125 | def get_text_embed_dict(ckpt_path, use_ema=True): |
| 126 | ckpt_type = ckpt_path.split(".")[-1] |
| 127 | if ckpt_type == "safetensors": |
| 128 | from safetensors.torch import load_file |
| 129 | |
| 130 | checkpoint = load_file(ckpt_path) |
| 131 | else: |
| 132 | checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=True) |
| 133 | |
| 134 | if use_ema: |
| 135 | if ckpt_type == "safetensors": |
| 136 | checkpoint = {"ema_model_state_dict": checkpoint} |
| 137 | checkpoint["model_state_dict"] = { |
| 138 | k.replace("ema_model.", ""): v |
| 139 | for k, v in checkpoint["ema_model_state_dict"].items() |
| 140 | if k not in ["initted", "step"] |
| 141 | } |
| 142 | else: |
| 143 | if ckpt_type == "safetensors": |
| 144 | checkpoint = {"model_state_dict": checkpoint} |
| 145 | model_params = checkpoint["model_state_dict"] |
| 146 | |
| 147 | text_embed_dict = {} |
| 148 | for key in model_params.keys(): |
| 149 | # transformer.text_embed.text_embed.weight -> text_embed.weight |
| 150 | if "text_embed" in key: |
| 151 | text_embed_dict[key.replace("transformer.text_embed.", "")] = model_params[key] |
| 152 | return text_embed_dict |
| 153 | |
| 154 | |
| 155 | class F5TTS(object): |
| 156 | def __init__( |
| 157 | self, |
| 158 | config, |
| 159 | debug_mode=True, |
| 160 | stream: Optional[torch.cuda.Stream] = None, |
| 161 | tllm_model_dir: Optional[str] = None, |
| 162 | model_path: Optional[str] = None, |
| 163 | vocab_size: Optional[int] = None, |
| 164 | ): |
| 165 | self.dtype = config["pretrained_config"]["dtype"] |
| 166 | |
| 167 | rank = tensorrt_llm.mpi_rank() |
| 168 | world_size = config["pretrained_config"]["mapping"]["world_size"] |
| 169 | cp_size = config["pretrained_config"]["mapping"]["cp_size"] |
| 170 | tp_size = config["pretrained_config"]["mapping"]["tp_size"] |
| 171 | pp_size = config["pretrained_config"]["mapping"]["pp_size"] |
| 172 | assert pp_size == 1 |
| 173 | self.mapping = tensorrt_llm.Mapping( |
| 174 | world_size=world_size, rank=rank, cp_size=cp_size, tp_size=tp_size, pp_size=1, gpus_per_node=1 |
| 175 | ) |
| 176 | |
| 177 | local_rank = rank % self.mapping.gpus_per_node |
| 178 | self.device = torch.device(f"cuda:{local_rank}") |
| 179 | |
| 180 | torch.cuda.set_device(self.device) |
| 181 | |
| 182 | self.stream = stream |
| 183 | if self.stream is None: |
| 184 | self.stream = torch.cuda.Stream(self.device) |
| 185 | torch.cuda.set_stream(self.stream) |
| 186 | |
| 187 | engine_file = os.path.join(tllm_model_dir, f"rank{rank}.engine") |
| 188 | logger.info(f"Loading engine from {engine_file}") |
| 189 | with open(engine_file, "rb") as f: |
| 190 | engine_buffer = f.read() |
| 191 | |
| 192 | assert engine_buffer is not None |
| 193 | |
| 194 | self.session = Session.from_serialized_engine(engine_buffer) |
| 195 | |
| 196 | self.debug_mode = debug_mode |
| 197 | |
| 198 | self.inputs = {} |
| 199 | self.outputs = {} |
| 200 | self.buffer_allocated = False |
| 201 | |
| 202 | expected_tensor_names = ["noise", "cond", "time", "rope_cos", "rope_sin", "input_lengths", "denoised"] |
| 203 | |
| 204 | found_tensor_names = [self.session.engine.get_tensor_name(i) for i in range(self.session.engine.num_io_tensors)] |
| 205 | if not self.debug_mode and set(expected_tensor_names) != set(found_tensor_names): |
| 206 | logger.error( |
| 207 | f"The following expected tensors are not found: {set(expected_tensor_names).difference(set(found_tensor_names))}" |
| 208 | ) |
| 209 | logger.error( |
| 210 | f"Those tensors in engine are not expected: {set(found_tensor_names).difference(set(expected_tensor_names))}" |
| 211 | ) |
| 212 | logger.error(f"Expected tensor names: {expected_tensor_names}") |
| 213 | logger.error(f"Found tensor names: {found_tensor_names}") |
| 214 | raise RuntimeError("Tensor names in engine are not the same as expected.") |
| 215 | if self.debug_mode: |
| 216 | self.debug_tensors = list(set(found_tensor_names) - set(expected_tensor_names)) |
| 217 | |
| 218 | self.max_mel_len = 4096 |
| 219 | self.text_embedding = TextEmbedding( |
| 220 | text_num_embeds=vocab_size, |
| 221 | text_dim=config["pretrained_config"]["text_dim"], |
| 222 | mask_padding=config["pretrained_config"]["text_mask_padding"], |
| 223 | conv_layers=config["pretrained_config"]["conv_layers"], |
| 224 | precompute_max_pos=self.max_mel_len, |
| 225 | ).to(self.device) |
| 226 | self.text_embedding.load_state_dict(get_text_embed_dict(model_path), strict=True) |
| 227 | |
| 228 | self.n_mel_channels = config["pretrained_config"]["mel_dim"] |
| 229 | self.head_dim = config["pretrained_config"]["dim_head"] |
| 230 | self.base_rescale_factor = 1.0 |
| 231 | self.interpolation_factor = 1.0 |
| 232 | base = 10000.0 * self.base_rescale_factor ** (self.head_dim / (self.head_dim - 2)) |
| 233 | inv_freq = 1.0 / (base ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim)) |
| 234 | freqs = torch.outer(torch.arange(self.max_mel_len, dtype=torch.float32), inv_freq) / self.interpolation_factor |
| 235 | self.freqs = freqs.repeat_interleave(2, dim=-1).unsqueeze(0) |
| 236 | self.rope_cos = self.freqs.cos().half() |
| 237 | self.rope_sin = self.freqs.sin().half() |
| 238 | |
| 239 | self.nfe_steps = 32 |
| 240 | epss = { |
| 241 | 5: [0, 2, 4, 8, 16, 32], |
| 242 | 6: [0, 2, 4, 6, 8, 16, 32], |
| 243 | 7: [0, 2, 4, 6, 8, 16, 24, 32], |
| 244 | 10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32], |
| 245 | 12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32], |
| 246 | 16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32], |
| 247 | } |
| 248 | t = 1 / 32 * torch.tensor(epss.get(self.nfe_steps, list(range(self.nfe_steps + 1))), dtype=torch.float32) |
| 249 | time_step = 1 - torch.cos(torch.pi * t / 2) |
| 250 | delta_t = torch.diff(time_step) |
| 251 | |
| 252 | freq_embed_dim = 256 # Warning: hard coding 256 here |
| 253 | time_expand = torch.zeros((1, self.nfe_steps, freq_embed_dim), dtype=torch.float32) |
| 254 | half_dim = freq_embed_dim // 2 |
| 255 | emb_factor = math.log(10000) / (half_dim - 1) |
| 256 | emb_factor = 1000.0 * torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb_factor) |
| 257 | for i in range(self.nfe_steps): |
| 258 | emb = time_step[i] * emb_factor |
| 259 | time_expand[:, i, :] = torch.cat((emb.sin(), emb.cos()), dim=-1) |
| 260 | self.time_expand = time_expand.to(self.device) |
| 261 | self.delta_t = torch.cat((delta_t, delta_t), dim=0).contiguous().to(self.device) |
| 262 | |
| 263 | def _tensor_dtype(self, name): |
| 264 | # return torch dtype given tensor name for convenience |
| 265 | dtype = trt_dtype_to_torch(self.session.engine.get_tensor_dtype(name)) |
| 266 | return dtype |
| 267 | |
| 268 | def _setup(self, batch_size, seq_len): |
| 269 | for i in range(self.session.engine.num_io_tensors): |
| 270 | name = self.session.engine.get_tensor_name(i) |
| 271 | if self.session.engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT: |
| 272 | shape = list(self.session.engine.get_tensor_shape(name)) |
| 273 | shape[0] = batch_size |
| 274 | shape[1] = seq_len |
| 275 | self.outputs[name] = torch.empty(shape, dtype=self._tensor_dtype(name), device=self.device) |
| 276 | |
| 277 | self.buffer_allocated = True |
| 278 | |
| 279 | def cuda_stream_guard(func): |
| 280 | """Sync external stream and set current stream to the one bound to the session. Reset on exit.""" |
| 281 | |
| 282 | @wraps(func) |
| 283 | def wrapper(self, *args, **kwargs): |
| 284 | external_stream = torch.cuda.current_stream() |
| 285 | if external_stream != self.stream: |
| 286 | external_stream.synchronize() |
| 287 | torch.cuda.set_stream(self.stream) |
| 288 | ret = func(self, *args, **kwargs) |
| 289 | if external_stream != self.stream: |
| 290 | self.stream.synchronize() |
| 291 | torch.cuda.set_stream(external_stream) |
| 292 | return ret |
| 293 | |
| 294 | return wrapper |
| 295 | |
| 296 | @cuda_stream_guard |
| 297 | def forward( |
| 298 | self, |
| 299 | noise: torch.Tensor, |
| 300 | cond: torch.Tensor, |
| 301 | time_expand: torch.Tensor, |
| 302 | rope_cos: torch.Tensor, |
| 303 | rope_sin: torch.Tensor, |
| 304 | input_lengths: torch.Tensor, |
| 305 | delta_t: torch.Tensor, |
| 306 | use_perf: bool = False, |
| 307 | ): |
| 308 | if use_perf: |
| 309 | torch.cuda.nvtx.range_push("flow matching") |
| 310 | cfg_strength = 2.0 |
| 311 | batch_size = noise.shape[0] |
| 312 | half_batch = batch_size // 2 |
| 313 | noise_half = noise[:half_batch] # Store the initial half of noise |
| 314 | |
| 315 | input_type = str_dtype_to_torch(self.dtype) |
| 316 | |
| 317 | # Keep a copy of the initial tensors |
| 318 | cond = cond.to(input_type) |
| 319 | rope_cos = rope_cos.to(input_type) |
| 320 | rope_sin = rope_sin.to(input_type) |
| 321 | input_lengths = input_lengths.to(str_dtype_to_torch("int32")) |
| 322 | |
| 323 | # Instead of iteratively updating noise within a single model context, |
| 324 | # we'll do a single forward pass for each iteration with fresh context setup |
| 325 | for i in range(self.nfe_steps): |
| 326 | # Re-setup the buffers for clean execution |
| 327 | self._setup(batch_size, noise.shape[1]) |
| 328 | if not self.buffer_allocated: |
| 329 | raise RuntimeError("Buffer not allocated, please call setup first!") |
| 330 | |
| 331 | # Re-create combined noises for this iteration |
| 332 | current_noise = torch.cat([noise_half, noise_half], dim=0).to(input_type) |
| 333 | |
| 334 | # Get time step for this iteration |
| 335 | current_time = time_expand[:, i].to(input_type) |
| 336 | |
| 337 | # Create fresh input dictionary for this iteration |
| 338 | current_inputs = { |
| 339 | "noise": current_noise, |
| 340 | "cond": cond, |
| 341 | "time": current_time, |
| 342 | "rope_cos": rope_cos, |
| 343 | "rope_sin": rope_sin, |
| 344 | "input_lengths": input_lengths, |
| 345 | } |
| 346 | |
| 347 | # Update inputs and set shapes |
| 348 | self.inputs.clear() # Clear previous inputs |
| 349 | self.inputs.update(**current_inputs) |
| 350 | self.session.set_shapes(self.inputs) |
| 351 | |
| 352 | if use_perf: |
| 353 | torch.cuda.nvtx.range_push(f"execute {i}") |
| 354 | ok = self.session.run(self.inputs, self.outputs, self.stream.cuda_stream) |
| 355 | assert ok, "Failed to execute model" |
| 356 | # self.session.context.execute_async_v3(self.stream.cuda_stream) |
| 357 | if use_perf: |
| 358 | torch.cuda.nvtx.range_pop() |
| 359 | # Process results |
| 360 | t_scale = delta_t[i].unsqueeze(0).to(input_type) |
| 361 | |
| 362 | # Extract predictions |
| 363 | pred_cond = self.outputs["denoised"][:half_batch] |
| 364 | pred_uncond = self.outputs["denoised"][half_batch:] |
| 365 | |
| 366 | # Apply classifier-free guidance with safeguards |
| 367 | guidance = pred_cond + (pred_cond - pred_uncond) * cfg_strength |
| 368 | # Calculate update for noise |
| 369 | noise_half = noise_half + guidance * t_scale |
| 370 | if use_perf: |
| 371 | torch.cuda.nvtx.range_pop() |
| 372 | return noise_half |
| 373 | |
| 374 | def sample( |
| 375 | self, |
| 376 | text_pad_sequence: torch.Tensor, |
| 377 | cond_pad_sequence: torch.Tensor, |
| 378 | ref_mel_len_batch: torch.Tensor, |
| 379 | estimated_reference_target_mel_len: List[int], |
| 380 | remove_input_padding: bool = False, |
| 381 | use_perf: bool = False, |
| 382 | ): |
| 383 | if use_perf: |
| 384 | torch.cuda.nvtx.range_push("text embedding") |
| 385 | batch = text_pad_sequence.shape[0] |
| 386 | max_seq_len = cond_pad_sequence.shape[1] |
| 387 | |
| 388 | # get text_embed one by one to avoid misalignment |
| 389 | text_and_drop_embedding_list = [] |
| 390 | for i in range(batch): |
| 391 | text_embedding_i = self.text_embedding( |
| 392 | text_pad_sequence[i].unsqueeze(0).to(self.device), |
| 393 | estimated_reference_target_mel_len[i], |
| 394 | drop_text=False, |
| 395 | ) |
| 396 | text_embedding_drop_i = self.text_embedding( |
| 397 | text_pad_sequence[i].unsqueeze(0).to(self.device), |
| 398 | estimated_reference_target_mel_len[i], |
| 399 | drop_text=True, |
| 400 | ) |
| 401 | text_and_drop_embedding_list.extend([text_embedding_i[0], text_embedding_drop_i[0]]) |
| 402 | |
| 403 | # pad separately computed text_embed to form batch with max_seq_len |
| 404 | text_and_drop_embedding = pad_sequence( |
| 405 | text_and_drop_embedding_list, |
| 406 | batch_first=True, |
| 407 | padding_value=0, |
| 408 | ) |
| 409 | text_embedding = text_and_drop_embedding[0::2] |
| 410 | text_embedding_drop = text_and_drop_embedding[1::2] |
| 411 | |
| 412 | noise = torch.randn_like(cond_pad_sequence).to(self.device) |
| 413 | rope_cos = self.rope_cos[:, :max_seq_len, :].float().repeat(batch, 1, 1) |
| 414 | rope_sin = self.rope_sin[:, :max_seq_len, :].float().repeat(batch, 1, 1) |
| 415 | |
| 416 | cat_mel_text = torch.cat( |
| 417 | ( |
| 418 | cond_pad_sequence, |
| 419 | text_embedding, |
| 420 | ), |
| 421 | dim=-1, |
| 422 | ) |
| 423 | cat_mel_text_drop = torch.cat( |
| 424 | ( |
| 425 | torch.zeros((batch, max_seq_len, self.n_mel_channels), dtype=torch.float32).to(self.device), |
| 426 | text_embedding_drop, |
| 427 | ), |
| 428 | dim=-1, |
| 429 | ) |
| 430 | |
| 431 | time_expand = self.time_expand.repeat(2 * batch, 1, 1).contiguous() |
| 432 | |
| 433 | # Convert estimated_reference_target_mel_len to tensor |
| 434 | input_lengths = torch.tensor(estimated_reference_target_mel_len, dtype=torch.int32) |
| 435 | |
| 436 | # combine above along the batch dimension |
| 437 | inputs = { |
| 438 | "noise": torch.cat((noise, noise), dim=0).contiguous(), |
| 439 | "cond": torch.cat((cat_mel_text, cat_mel_text_drop), dim=0).contiguous(), |
| 440 | "time_expand": time_expand, |
| 441 | "rope_cos": torch.cat((rope_cos, rope_cos), dim=0).contiguous(), |
| 442 | "rope_sin": torch.cat((rope_sin, rope_sin), dim=0).contiguous(), |
| 443 | "input_lengths": torch.cat((input_lengths, input_lengths), dim=0).contiguous(), |
| 444 | "delta_t": self.delta_t, |
| 445 | } |
| 446 | if use_perf and remove_input_padding: |
| 447 | torch.cuda.nvtx.range_push("remove input padding") |
| 448 | if remove_input_padding: |
| 449 | max_seq_len = inputs["cond"].shape[1] |
| 450 | inputs["noise"] = remove_tensor_padding(inputs["noise"], inputs["input_lengths"]) |
| 451 | inputs["cond"] = remove_tensor_padding(inputs["cond"], inputs["input_lengths"]) |
| 452 | # for time_expand, convert from B,D to B,T,D by repeat |
| 453 | inputs["time_expand"] = inputs["time_expand"].unsqueeze(1).repeat(1, max_seq_len, 1, 1) |
| 454 | inputs["time_expand"] = remove_tensor_padding(inputs["time_expand"], inputs["input_lengths"]) |
| 455 | inputs["rope_cos"] = remove_tensor_padding(inputs["rope_cos"], inputs["input_lengths"]) |
| 456 | inputs["rope_sin"] = remove_tensor_padding(inputs["rope_sin"], inputs["input_lengths"]) |
| 457 | if use_perf and remove_input_padding: |
| 458 | torch.cuda.nvtx.range_pop() |
| 459 | for key in inputs: |
| 460 | inputs[key] = inputs[key].to(self.device) |
| 461 | if use_perf: |
| 462 | torch.cuda.nvtx.range_pop() |
| 463 | start_time = time.time() |
| 464 | denoised = self.forward(**inputs, use_perf=use_perf) |
| 465 | cost_time = time.time() - start_time |
| 466 | if use_perf and remove_input_padding: |
| 467 | torch.cuda.nvtx.range_push("remove input padding output") |
| 468 | if remove_input_padding: |
| 469 | denoised_list = [] |
| 470 | start_idx = 0 |
| 471 | for i in range(batch): |
| 472 | denoised_list.append(denoised[start_idx : start_idx + inputs["input_lengths"][i]]) |
| 473 | start_idx += inputs["input_lengths"][i] |
| 474 | if use_perf and remove_input_padding: |
| 475 | torch.cuda.nvtx.range_pop() |
| 476 | return denoised_list, cost_time |
| 477 | return denoised, cost_time |
| 478 |