.

AR/models/embedding.py

@@ -33,51 +33,6 @@ class TokenEmbedding(nn.Module):
         return x
 
 
-class SinePositionalEmbedding(nn.Module):
-    def __init__(
-        self,
-        embedding_dim: int,
-        dropout: float = 0.0,
-        scale: bool = False,
-        alpha: bool = False,
-    ):
-        super().__init__()
-        self.embedding_dim = embedding_dim
-        self.x_scale = math.sqrt(embedding_dim) if scale else 1.0
-        self.alpha = nn.Parameter(torch.ones(1), requires_grad=alpha)
-        self.dropout = torch.nn.Dropout(p=dropout)
-
-        self.reverse = False
-        self.pe = None
-        self.extend_pe(torch.tensor(0.0).expand(1, 4000))
-
-    def extend_pe(self, x):
-        """Reset the positional encodings."""
-        if self.pe is not None:
-            if self.pe.size(1) >= x.size(1):
-                if self.pe.dtype != x.dtype or self.pe.device != x.device:
-                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)
-                return
-        pe = torch.zeros(x.size(1), self.embedding_dim)
-        if self.reverse:
-            position = torch.arange(x.size(1) - 1, -1, -1.0, dtype=torch.float32).unsqueeze(1)
-        else:
-            position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)
-        div_term = torch.exp(
-            torch.arange(0, self.embedding_dim, 2, dtype=torch.float32) * -(math.log(10000.0) / self.embedding_dim)
-        )
-        pe[:, 0::2] = torch.sin(position * div_term)
-        pe[:, 1::2] = torch.cos(position * div_term)
-        pe = pe.unsqueeze(0)
-        self.pe = pe.to(device=x.device, dtype=x.dtype).detach()
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        self.extend_pe(x)
-        output = x.unsqueeze(-1) if x.ndim == 2 else x
-        output = output * self.x_scale + self.alpha * self.pe[:, : x.size(1)]
-        return self.dropout(output)
-
-
 class SinePositionalEmbeddingNested(nn.Module):
     def __init__(
         self,

AR/models/structs.py

@@ -5,11 +5,11 @@ Modified From https://github.com/XXXXRT666/GPT-SoVITS
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import List, Literal, Optional
+from typing import List, Literal, MutableSequence, Optional
 
 import torch
 
-from AR.models.t2s_model_abc import Sampler, T2SDecoderABC
+from AR.models.t2s_model_abc import KVCacheABC, Sampler, T2SDecoderABC
 
 Tensor = torch.Tensor
 
@@ -53,6 +53,7 @@ class T2SSession:
             self.y_len = y_len
 
             # Cache
+            self.kv_cache: MutableSequence[KVCacheABC]
             self.sampler = Sampler(bsz, decoder.vocab_size)
 
             # Forward args
@@ -66,6 +67,11 @@ class T2SSession:
             self.input_pos = torch.zeros_like(self.prefill_len)
             self.input_pos.add_(self.prefill_len)
 
+            # CUDA Graph
+            self.graph: Optional[torch.cuda.CUDAGraph] = None
+            self.xy_pos_: Tensor
+            self.xy_dec_: Tensor
+
             # EOS
             self.completed = torch.Tensor([False] * len(self.x)).bool().to(device)
             self.y_results: List[Tensor] = [None] * len(self.x)  # type: ignore
@@ -81,3 +87,5 @@ class T2SSession:
                     mask[-y_len:, -y_len:] = ~torch.triu(torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1)
                 attn_mask.append(mask)
             self.attn_mask_nested = torch.nested.nested_tensor(attn_mask)
+
+            self.id: int = -1

AR/models/t2s_model_abc.py

@@ -5,10 +5,10 @@ Modified From https://github.com/XXXXRT666/GPT-SoVITS
 from __future__ import annotations
 
 import os
-import time
+import random
 from abc import ABC, abstractmethod
 from contextlib import nullcontext
-from typing import Any, Dict, List, MutableSequence, Optional, Tuple, Type
+from typing import Any, Dict, List, MutableSequence, Tuple, Type
 
 import torch
 import torch._inductor.config
@@ -85,6 +85,10 @@ class KVCacheABC(ABC, nn.Module):
     @abstractmethod
     def prefill_kv(self, k_val: Tensor, v_val: Tensor, bs: int) -> None: ...
 
+    def sync_cache(self, kv_cache: KVCacheABC):
+        self.k_cache.copy_(kv_cache.k_cache)
+        self.v_cache.copy_(kv_cache.v_cache)
+
     def forward(self):
         raise NotImplementedError()
 
@@ -363,6 +367,8 @@ class T2SDecoderABC(ABC, nn.Module):
 
         self.kv_class: Type[KVCacheNHD] | Type[KVCacheHND]
 
+        self.GraphCache: CUDAGraphCacheABC | None
+
         self._register_load_state_dict_pre_hook(self.load_hook)
 
     def load_hook(self, state_dict, prefix, *args):
@@ -396,6 +402,7 @@ class T2SDecoderABC(ABC, nn.Module):
         self.h.compile(fullgraph=True, mode="reduce-overhead")
 
     def capture(self, input_pos: Tensor, x: Tensor, x_dec: Tensor, *args, **kwds) -> CUDAGraph:
+        assert torch.cuda.is_available()
         s = torch.cuda.Stream()
         s.wait_stream(torch.cuda.current_stream())
 
@@ -419,6 +426,51 @@ class T2SDecoderABC(ABC, nn.Module):
     def post_forward(self, idx: int, session: Any) -> None: ...
 
 
+class CUDAGraphCacheABC(ABC):
+    def __init__(
+        self,
+        decoder: T2SDecoderABC,
+        device: torch.device = torch.device("cpu"),
+        dtype: torch.dtype = torch.float32,
+    ) -> None:
+        assert torch.cuda.is_available()
+
+        self.assigned: bool = False
+
+        self.decoder: T2SDecoderABC = decoder
+        self.kv_cache: MutableSequence[KVCacheABC] = decoder.init_cache(1)
+        self.xy_pos = torch.rand((1, 1, decoder.embedding_dim), device=device).to(dtype)
+        self.xy_dec = torch.rand((1, 1, decoder.embedding_dim), device=device).to(dtype)
+        self.input_pos = torch.tensor([10]).int().cuda()
+        self.graph: torch.cuda.CUDAGraph | None = None
+
+        self.id: int = random.randint(1, 2**32 - 1)
+
+    def assign_graph(self, session: Any):
+        if self.graph is None:
+            args, kwds = self.decoder.pre_forward(session)
+            graph = self.decoder.capture(self.input_pos, self.xy_pos, self.xy_dec, *args, **kwds)
+            self.graph = graph
+
+        if self.assigned is False:
+            self.get_cache_graph(session)
+            session.id = self.id
+            self.assigned = True
+        else:
+            self.capture_new_graph(session)
+
+    @abstractmethod
+    def release_graph(self, session: Any): ...
+
+    @abstractmethod
+    def get_cache_graph(self, session: Any):
+        pass
+
+    @abstractmethod
+    def capture_new_graph(self, session: Any):
+        pass
+
+
 class TorchProfiler:
     def __init__(self, debug: bool, log_dir: str = "./profiler") -> None:
         self.debug = debug

AR/models/t2s_model_flash_attn.py

@@ -2,13 +2,13 @@
 Modified From https://github.com/XXXXRT666/GPT-SoVITS
 """
 
+import gc
 import os
 import time
 import traceback
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Tuple
 
 import flash_attn  # type: ignore
-import gradio as gr
 import torch
 import torch.nn as nn
 from tqdm import tqdm
@@ -20,6 +20,7 @@ from AR.models.embedding import TokenEmbedding
 from AR.models.structs import T2SRequest, T2SResult, T2SSession
 from AR.models.t2s_model_abc import (
     AttentionABC,
+    CUDAGraphCacheABC,
     FeedForward,
     KVCacheABC,
     KVCacheNHD,
@@ -121,6 +122,7 @@ class T2SDecoder(T2SDecoderABC):
         max_batch_size=10,
         **kwds,
     ) -> None:
+        assert torch.cuda.is_available()
         super().__init__()
 
         hidden_dim = config["model"]["hidden_dim"]
@@ -205,6 +207,42 @@ class T2SDecoder(T2SDecoderABC):
         return list(), dict()
 
 
+class CUDAGraphCache(CUDAGraphCacheABC):
+    def __init__(
+        self,
+        decoder: T2SDecoderABC,
+        device: torch.device = torch.device("cpu"),
+        dtype: torch.dtype = torch.float32,
+    ) -> None:
+        super().__init__(decoder, device, dtype)
+
+    def release_graph(self, session: T2SSession):
+        if session.id != self.id:
+            self.assigned = False
+        else:
+            del session.graph, session.xy_pos_, session.xy_dec_, session.input_pos, session.kv_cache
+
+    def get_cache_graph(self, session: T2SSession):
+        assert self.graph
+        session.graph = self.graph
+
+        session.xy_pos_ = self.xy_pos
+        session.xy_dec_ = self.xy_dec
+        session.input_pos = self.input_pos.copy_(session.input_pos)
+
+        for cache, cache_ in zip(self.kv_cache, session.kv_cache):
+            cache.sync_cache(cache_)
+
+    def capture_new_graph(self, session: T2SSession):
+        session.xy_pos_ = self.xy_pos.clone()
+        session.xy_dec_ = self.xy_dec.clone()
+        session.input_pos = self.input_pos.clone().copy_(session.input_pos)
+
+        args, kwds = self.decoder.pre_forward(session)
+        graph = self.decoder.capture(self.input_pos, self.xy_pos, self.xy_dec, *args, **kwds)
+        session.graph = graph
+
+
 class CUDAGraphRunner:
     def __init__(
         self,
@@ -212,70 +250,51 @@ class CUDAGraphRunner:
         device: torch.device = torch.device("cpu"),
         dtype: torch.dtype = torch.float32,
     ) -> None:
-        assert device.type in {"cpu", "cuda", "mps", "xpu", "mtia"}
-        assert dtype in {torch.float16, torch.bfloat16, torch.float32}
+        assert device.type == "cuda"
         self.device = device
         self.dtype = dtype
 
-        self.decoder_path: os.PathLike
         self.decoder_model: T2SDecoderABC = decoder_model.to(self.device, self.dtype)
 
-        self.graph: Optional[torch.cuda.CUDAGraph] = None
-        self.xy_pos_ = torch.rand((1, 1, decoder_model.embedding_dim), device=device).to(dtype)
-        self.xy_dec_ = torch.rand((1, 1, decoder_model.embedding_dim), device=device).to(dtype)
-        self.kv_cache = decoder_model.init_cache(1)
-        self.input_pos = torch.tensor([10]).int().cuda()
+        self.graphcache = CUDAGraphCache(decoder_model, device, dtype)
 
     def _handle_request(self, request: T2SRequest):
         with self.device:
-            for i in self.kv_cache:
-                i.empty()
-
             decoder = self.decoder_model
             session = T2SSession(decoder, request, device=self.device, dtype=self.dtype)
-            self.input_pos.copy_(session.input_pos)
 
             t1 = 0.0
             infer_speed = 0.0
-            y = session.y
-            bsz = y.size(0)
+
             torch_profiler = TorchProfiler(request.debug)
             with torch_profiler.profiler():
                 for idx in tqdm(range(1500)):
                     if idx == 0:
-                        xy_dec = decoder.h.prefill(session.xy_pos, session.attn_mask_nested, self.kv_cache)
+                        session.kv_cache = decoder.init_cache(session.bsz)
+                        xy_dec = decoder.h.prefill(session.xy_pos, session.attn_mask_nested, session.kv_cache)
                         xy_dec = torch.stack([t[[-1]] for t in xy_dec.unbind()])
                     else:
-                        if request.use_cuda_graph and self.graph is None and torch.cuda.is_available():
-                            self.xy_pos_.copy_(session.xy_pos)
-                            args, kwds = decoder.pre_forward(session)
-                            self.graph = decoder.capture(
-                                self.input_pos,
-                                self.xy_pos_,
-                                self.xy_dec_,
-                                kv_caches=self.kv_cache,
-                                *args,
-                                **kwds,
-                            )
+                        if request.use_cuda_graph and session.graph is None and torch.cuda.is_available():
+                            self.graphcache.assign_graph(session)
 
                         with torch_profiler.record("AR"):
-                            if self.graph:
-                                self.xy_pos_.copy_(session.xy_pos)
-                                self.graph.replay()
-                                xy_dec = self.xy_dec_.clone()
+                            if session.graph:
+                                session.xy_pos_.copy_(session.xy_pos)
+                                session.graph.replay()
+                                xy_dec = session.xy_dec_.clone()
                             else:
                                 args, kwds = decoder.pre_forward(session)
                                 xy_dec = decoder.h.forward(
-                                    self.input_pos,
+                                    session.input_pos,
                                     session.xy_pos,
-                                    self.kv_cache,
+                                    session.kv_cache,
                                     *args,
                                     **kwds,
                                 )
 
                     decoder.post_forward(idx, session)
                     logits = decoder.ar_predict_layer(xy_dec[:, -1])
-                    self.input_pos.add_(1)
+                    session.input_pos.add_(1)
 
                     if idx == 0:
                         logits[:, -1] = float("-inf")
@@ -322,7 +341,7 @@ class CUDAGraphRunner:
                             request.early_stop_num != -1
                             and (session.y.size(1) - session.y_len) > request.early_stop_num
                         ) or idx == 1499:
-                            for i in range(bsz):
+                            for i in range(session.bsz):
                                 if not session.completed[i].item():
                                     session.y_results[i] = session.y[i, session.y_len :]
                                     session.completed[i] = True
@@ -330,7 +349,7 @@ class CUDAGraphRunner:
 
                     with torch_profiler.record("NextPos"):
                         y_emb = decoder.ar_audio_embedding(session.y[:, -1:])
-                        session.xy_pos = decoder.ar_audio_position.forward(self.input_pos - session.x_lens, y_emb)
+                        session.xy_pos = decoder.ar_audio_position.forward(session.input_pos - session.x_lens, y_emb)
 
                     if idx == 2:
                         torch_profiler.start()
@@ -359,8 +378,11 @@ class CUDAGraphRunner:
                     torch.xpu.empty_cache()
                 case "mtia":
                     torch.mtia.empty_cache()
+                case "cpu":
+                    gc.collect()
 
             torch_profiler.end()
+            self.graphcache.release_graph(session)
             return session.y_results[: request.valid_length], infer_speed
 
     def generate(self, request: T2SRequest):

inference_webui.py

@@ -1,7 +1,47 @@
+import logging
 import os
+import re
+import traceback
+from time import time as ttime
 
-os.makedirs("pretrained_models", exist_ok=True)
+import gradio as gr
+import gradio.themes as themes
+import librosa
+import nltk
+import numpy as np
+import spaces
+import torch
+import torchaudio
+from gradio.themes.utils import fonts
 from huggingface_hub import snapshot_download
+from transformers.models.auto.modeling_auto import AutoModelForMaskedLM
+from transformers.models.auto.tokenization_auto import AutoTokenizer
+
+from AR.models.structs import T2SRequest
+from AR.models.t2s_model_flash_attn import CUDAGraphRunner
+from feature_extractor import cnhubert
+from module.mel_processing import spectrogram_torch
+from module.models import SynthesizerTrn
+from sv import SV
+from text import chinese, cleaned_text_to_sequence
+from text.cleaner import clean_text
+from text.LangSegmenter import LangSegmenter
+from tools.i18n.i18n import I18nAuto
+
+logging.getLogger("markdown_it").setLevel(logging.ERROR)
+logging.getLogger("urllib3").setLevel(logging.ERROR)
+logging.getLogger("httpcore").setLevel(logging.ERROR)
+logging.getLogger("httpx").setLevel(logging.ERROR)
+logging.getLogger("asyncio").setLevel(logging.ERROR)
+logging.getLogger("charset_normalizer").setLevel(logging.ERROR)
+logging.getLogger("torchaudio._extension").setLevel(logging.ERROR)
+logging.getLogger("multipart.multipart").setLevel(logging.ERROR)
+logging.getLogger("python_multipart.multipart").setLevel(logging.ERROR)
+logging.getLogger("split_lang.split.splitter").setLevel(logging.ERROR)
+
+os.makedirs("pretrained_models", exist_ok=True)
+
+nltk.download("averaged_perceptron_tagger_eng")
 
 snapshot_download(
     repo_id="lj1995/GPT-SoVITS",
@@ -27,75 +67,20 @@ snapshot_download(
     allow_patterns="v2Pro/s2Gv2ProPlus.pth",
     local_dir="pretrained_models",
 )
-import logging
-import traceback
-
-logging.getLogger("markdown_it").setLevel(logging.ERROR)
-logging.getLogger("urllib3").setLevel(logging.ERROR)
-logging.getLogger("httpcore").setLevel(logging.ERROR)
-logging.getLogger("httpx").setLevel(logging.ERROR)
-logging.getLogger("asyncio").setLevel(logging.ERROR)
-logging.getLogger("charset_normalizer").setLevel(logging.ERROR)
-logging.getLogger("torchaudio._extension").setLevel(logging.ERROR)
-logging.getLogger("multipart.multipart").setLevel(logging.ERROR)
-logging.getLogger("python_multipart.multipart").setLevel(logging.ERROR)
-logging.getLogger("split_lang.split.splitter").setLevel(logging.ERROR)
-
-import nltk
-import torchaudio
-
-from text.LangSegmenter import LangSegmenter
-
-nltk.download("averaged_perceptron_tagger_eng")
-import json
-import os
-import pdb
-import re
-import sys
-import threading
-
-import LangSegment
-import spaces
-import torch
-
-lock = threading.Lock()
 
 version = "v2"  # os.environ.get("version","v2")
 cnhubert_base_path = os.environ.get("cnhubert_base_path", "pretrained_models/chinese-hubert-base")
 bert_path = os.environ.get("bert_path", "pretrained_models/chinese-roberta-wwm-ext-large")
-
-punctuation = set(["!", "?", "…", ",", ".", "-", " "])
-import gradio as gr
-import gradio.themes as themes
-import librosa
-import numpy as np
-from gradio.themes.utils import fonts
-from transformers import AutoModelForMaskedLM, AutoTokenizer
-
-from feature_extractor import cnhubert
-
 cnhubert.cnhubert_base_path = cnhubert_base_path
 
-from time import time as ttime
+punctuation = set(["!", "?", "…", ",", ".", "-", " "])
 
-from AR.models.structs import T2SRequest
-from AR.models.t2s_model_flash_attn import CUDAGraphRunner
-from module.mel_processing import spectrogram_torch
-from module.models import SynthesizerTrn
-from text import cleaned_text_to_sequence
-from text.cleaner import clean_text
-from tools.i18n.i18n import I18nAuto, scan_language_list
-from tools.my_utils import load_audio
 
-# language=os.environ.get("language","Auto")
-# language=sys.argv[-1] if sys.argv[-1] in scan_language_list() else language
 i18n = I18nAuto(language="Auto")
 
-# os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'  # 确保直接启动推理UI时也能够设置。
-
 if torch.cuda.is_available():
     device = "cuda"
-    is_half = True  # eval(os.environ.get("is_half", "True")) and torch.cuda.is_available()
+    is_half = True
 else:
     device = "cpu"
     is_half = False
@@ -125,7 +110,7 @@ dict_language = dict_language_v1 if version == "v1" else dict_language_v2
 
 tokenizer = AutoTokenizer.from_pretrained(bert_path)
 bert_model = AutoModelForMaskedLM.from_pretrained(bert_path)
-if is_half == True:
+if is_half is True:
     bert_model = bert_model.half().to(device)
 else:
     bert_model = bert_model.to(device)
@@ -176,7 +161,7 @@ class DictToAttrRecursive(dict):
 
 
 ssl_model = cnhubert.get_model()
-if is_half == True:
+if is_half is True:
     ssl_model = ssl_model.half().to(device)
 else:
     ssl_model = ssl_model.to(device)
@@ -248,7 +233,7 @@ def change_gpt_weights(gpt_path):
 
 
 change_gpt_weights("pretrained_models/s1v3.ckpt")
-from sv import SV
+
 
 sv_cn_model = SV(device, is_half)
 
@@ -288,7 +273,7 @@ def get_spepc(hps, filename, dtype, device, is_v2pro=False):
         center=False,
     )
     spec = spec.to(dtype)
-    if is_v2pro == True:
+    if is_v2pro is True:
         audio = resample(audio, sr1, 16000, device).to(dtype)
     return spec, audio
 
@@ -300,7 +285,7 @@ def clean_text_inf(text, language, version):
     return phones, word2ph, norm_text
 
 
-dtype = torch.float16 if is_half == True else torch.float32
+dtype = torch.float16 if is_half is True else torch.float32
 
 
 def get_bert_inf(phones, word2ph, norm_text, language):
@@ -310,27 +295,13 @@ def get_bert_inf(phones, word2ph, norm_text, language):
     else:
         bert = torch.zeros(
             (1024, len(phones)),
-            dtype=torch.float16 if is_half == True else torch.float32,
+            dtype=torch.float16 if is_half is True else torch.float32,
         ).to(device)
 
     return bert
 
 
-splits = {
-    "，",
-    "。",
-    "？",
-    "！",
-    ",",
-    ".",
-    "?",
-    "!",
-    "~",
-    ":",
-    "：",
-    "—",
-    "…",
-}
+splits = {"，", "。", "？", "！", ",", ".", "?", "!", "~", ":", "：", "—", "…"}
 
 
 def get_first(text):
@@ -339,9 +310,6 @@ def get_first(text):
     return text
 
 
-from text import chinese
-
-
 def get_phones_and_bert(text, language, version, final=False):
     if language in {"en", "all_zh", "all_ja", "all_ko", "all_yue"}:
         formattext = text
@@ -363,7 +331,7 @@ def get_phones_and_bert(text, language, version, final=False):
             phones, word2ph, norm_text = clean_text_inf(formattext, language, version)
             bert = torch.zeros(
                 (1024, len(phones)),
-                dtype=torch.float16 if is_half == True else torch.float32,
+                dtype=torch.float16 if is_half is True else torch.float32,
             ).to(device)
     elif language in {"zh", "ja", "ko", "yue", "auto", "auto_yue"}:
         textlist = []
@@ -475,7 +443,7 @@ def get_tts_wav(
     print(i18n("实际输入的目标文本:"), text)
     zero_wav = np.zeros(
         int(hps.data.sampling_rate * 0.3),
-        dtype=np.float16 if is_half == True else np.float32,
+        dtype=np.float16 if is_half is True else np.float32,
     )
     if not ref_free:
         with torch.no_grad():
@@ -485,7 +453,7 @@ def get_tts_wav(
                 raise OSError(i18n("参考音频在3~10秒范围外，请更换！"))
             wav16k = torch.from_numpy(wav16k)
             zero_wav_torch = torch.from_numpy(zero_wav)
-            if is_half == True:
+            if is_half is True:
                 wav16k = wav16k.half().to(device)
                 zero_wav_torch = zero_wav_torch.half().to(device)
             else:
@@ -544,10 +512,10 @@ def get_tts_wav(
         t2 = ttime()
         # cache_key="%s-%s-%s-%s-%s-%s-%s-%s"%(ref_wav_path,prompt_text,prompt_language,text,text_language,top_k,top_p,temperature)
         # print(cache.keys(),if_freeze)
-        if i_text in cache and if_freeze == True:
+        if i_text in cache and if_freeze is True:
             pred_semantic = cache[i_text]
         else:
-            with torch.no_grad(), lock:
+            with torch.no_grad():
                 t2s_request = T2SRequest(
                     [all_phoneme_ids.squeeze(0)],
                     all_phoneme_len,
@@ -564,9 +532,8 @@ def get_tts_wav(
                 t2s_result = t2s_model.generate(t2s_request)
 
                 if t2s_result.exception is not None:
-                    print(t2s_result.exception)
                     print(t2s_result.traceback)
-                    raise RuntimeError("")
+                    raise t2s_result.exception
 
                 infer_speed.append(t2s_result.infer_speed)
                 pred_semantic = t2s_result.result
@@ -608,8 +575,8 @@ def get_tts_wav(
         t.extend([t2 - t1, t3 - t2, t4 - t3])
         t1 = ttime()
     print("%.3f\t%.3f\t%.3f\t%.3f" % (t[0], sum(t[1::3]), sum(t[2::3]), sum(t[3::3])))
-    gr.Info(f"Infer Speed: {sum(infer_speed) / len(infer_speed):.2f} Token/s")
-    gr.Info("%.3f\t%.3f\t%.3f\t%.3f" % (t[0], sum(t[1::3]), sum(t[2::3]), sum(t[3::3])), duration=4)
+    gr.Info(f"{sum(infer_speed) / len(infer_speed):.2f} Token/s", title="Infer Speed")
+    gr.Info("%.3f\t%.3f\t%.3f\t%.3f" % (t[0], sum(t[1::3]), sum(t[2::3]), sum(t[3::3])), title="Time Stamps")
     yield hps.data.sampling_rate, (np.concatenate(audio_opt, 0) * 32768).astype(np.int16)
 
 
@@ -713,7 +680,7 @@ def cut5(inp):
 
 def custom_sort_key(s):
     # 使用正则表达式提取字符串中的数字部分和非数字部分
-    parts = re.split("(\d+)", s)
+    parts = re.split(r"(\d+)", s)
     # 将数字部分转换为整数，非数字部分保持不变
     parts = [int(part) if part.isdigit() else part for part in parts]
     return parts