Add CUDA Graph

.gitignore

@@ -0,0 +1,3 @@
+.*cache
+__pycache__
+pretrained_models

AR/models/structs.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Literal, Optional
+
+import torch
+
+from AR.models.t2s_model_abc import Sampler, T2SDecoderABC
+
+Tensor = torch.Tensor
+
+
+@dataclass
+class T2SResult:
+    result: List[Tensor] | None = None
+    status: Literal["Success", "Error"] = "Success"
+    exception: Optional[Exception] = None
+    traceback: Optional[str] = None
+
+
+@dataclass
+class T2SRequest:
+    x: List[torch.Tensor]
+    x_lens: Tensor
+    prompts: torch.Tensor
+    bert_feature: List[Tensor]
+    valid_length: int
+    top_k: int = 5
+    top_p: float = 1
+    early_stop_num: int = -1
+    temperature: float = 1.0
+    repetition_penalty: float = 1.35
+    use_cuda_graph: bool = False
+    debug: bool = False
+
+
+class T2SSession:
+    def __init__(self, decoder: T2SDecoderABC, request: T2SRequest, device: torch.device, dtype: torch.dtype):
+        with device:
+            self.decoder = decoder
+            self.request = request
+            self.device = device
+            self.dtype = dtype
+
+            bsz = len(request.x)
+            y_len = request.prompts.size(-1)
+            self.bsz = bsz
+            self.y_len = y_len
+
+            # Cache
+            self.kv_cache = decoder.init_cache(bsz)
+            self.sampler = Sampler(bsz, decoder.vocab_size)
+
+            # Forward args
+            self.x = request.x
+            self.x_lens = request.x_lens.to(torch.int32)
+            self.y = request.prompts
+            self.bert_feature = request.bert_feature
+
+            self.prefill_len = self.x_lens + self.y.size(1)
+
+            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_ = torch.rand((bsz, 1, decoder.embedding_dim)).to(dtype)
+            self.xy_dec_ = torch.rand((bsz, 1, decoder.embedding_dim)).to(dtype)
+
+            # EOS
+            self.completed = [False] * len(self.x)
+            self.y_results: List[Tensor] = [None] * len(self.x)  # type: ignore
+
+            self.xy_pos = decoder.embed(self.x, self.y, self.bert_feature)
+
+            attn_mask = []
+            for bs in range(bsz):
+                pos = int(self.x_lens[bs].item())
+                mask = torch.zeros(pos + y_len, pos + y_len).bool()
+                mask[:, :pos].fill_(True)
+                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)

AR/models/t2s_model_abc.py

@@ -0,0 +1,598 @@
+from __future__ import annotations
+
+import os
+from abc import ABC, abstractmethod
+from contextlib import nullcontext
+from typing import Any, Dict, List, MutableSequence, Optional, Tuple, Type
+
+import torch
+import torch._inductor.config
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.cuda.graphs import CUDAGraph
+from torch.profiler import ProfilerAction, tensorboard_trace_handler
+
+from AR.modules.embedding import (
+    SinePositionalEmbeddingNested as SinePositionalEmbedding,
+)
+from AR.modules.embedding import TokenEmbedding
+
+Tensor = torch.Tensor
+
+
+class Sampler(nn.Module):
+    def __init__(self, batch_size: int, vocab_size: int) -> None:
+        super().__init__()
+        self.batch_size = batch_size
+
+        self.logits: Tensor
+        self.samples: Tensor
+        self.register_buffer("logits", torch.zeros((batch_size, vocab_size)), persistent=False)
+        self.register_buffer("samples", torch.zeros((batch_size,), dtype=torch.int32), persistent=False)
+
+        self.__CUDAGraph: Optional[CUDAGraph] = None
+
+    def empty_cache(self):
+        self.logits.zero_()
+        self.__CUDAGraph = None
+
+    @staticmethod
+    def multinomial_sample_one_no_sync(probs_sort: Tensor):  # Does multinomial sampling without a cuda synchronization
+        q = torch.empty_like(probs_sort).exponential_(1)
+        return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.int32)
+
+    @staticmethod
+    def logits_to_probs(
+        logits: Tensor,
+        previous_tokens: Tensor,
+        temperature: float,
+        top_k: int,
+        top_p: float,
+        repetition_penalty: float,
+    ):
+        previous_tokens = previous_tokens.long()
+        score = torch.gather(logits, dim=1, index=previous_tokens)
+        score = torch.where(score < 0, score * repetition_penalty, score / repetition_penalty)
+        logits.scatter_(dim=1, index=previous_tokens, src=score)
+
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cum_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
+        sorted_indices_to_remove = cum_probs > top_p
+        sorted_indices_to_remove[:, 0] = False  # keep at least one option
+        indices_to_remove = sorted_indices_to_remove.scatter(dim=1, index=sorted_indices, src=sorted_indices_to_remove)
+        logits = logits.masked_fill(indices_to_remove, -float("Inf"))
+
+        logits = logits / max(temperature, 1e-5)
+
+        v, _ = torch.topk(logits, top_k)
+        pivot = v[:, -1].unsqueeze(-1)
+        logits = torch.where(logits < pivot, -float("Inf"), logits)
+
+        probs = torch.nn.functional.softmax(logits, dim=-1)
+        return probs
+
+    @staticmethod
+    def apply_repetition_penalty(logits: Tensor, previous_tokens: Tensor, repetition_penalty: float):
+        previous_tokens = previous_tokens.long()
+        score = torch.gather(logits, dim=1, index=previous_tokens)
+        score = torch.where(score < 0, score * repetition_penalty, score / repetition_penalty)
+        logits.scatter_(dim=1, index=previous_tokens, src=score)
+        return logits
+
+    @staticmethod
+    def logits_to_probs_cuda_graph(
+        logits: Tensor,
+        temperature: float,
+        top_k: int,
+        top_p: float,
+    ):
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cum_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
+        sorted_indices_to_remove = cum_probs > top_p
+        sorted_indices_to_remove[:, 0] = False  # keep at least one option
+        indices_to_remove = sorted_indices_to_remove.scatter(dim=1, index=sorted_indices, src=sorted_indices_to_remove)
+        logits = logits.masked_fill(indices_to_remove, -float("Inf"))
+
+        logits = logits / max(temperature, 1e-5)
+
+        v, _ = torch.topk(logits, top_k)
+        pivot = v[:, -1].unsqueeze(-1)
+        logits = torch.where(logits < pivot, -float("Inf"), logits)
+
+        probs = torch.nn.functional.softmax(logits, dim=-1)
+        return probs
+
+    def __sample(
+        self,
+        logits: Tensor,
+        previous_tokens: Tensor,
+        temperature: float,
+        top_k: int,
+        top_p: float,
+        repetition_penalty: float,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        probs = self.logits_to_probs(
+            logits=logits,
+            previous_tokens=previous_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+        )
+        idx_next = self.multinomial_sample_one_no_sync(probs)
+        return idx_next, probs
+
+    def __sample_cuda_graph(
+        self,
+        logits: Tensor,
+        temperature: float,
+        top_k: int,
+        top_p: float,
+    ):
+        probs = self.logits_to_probs_cuda_graph(
+            logits=logits,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+        )
+        idx_next = self.multinomial_sample_one_no_sync(probs)
+        return idx_next
+
+    def capture(self, temperature: float, top_k: int, top_p: float):
+        s = torch.cuda.Stream()
+        s.wait_stream(torch.cuda.current_stream())
+
+        logits = self.logits
+
+        with torch.cuda.stream(s):  # type: ignore
+            for _ in range(5):
+                self.__sample_cuda_graph(logits, temperature, top_k, top_p)
+        torch.cuda.current_stream().wait_stream(s)
+
+        self.__CUDAGraph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self.__CUDAGraph):
+            self.samples = self.__sample_cuda_graph(logits, temperature, top_k, top_p)
+        torch.cuda.synchronize()
+
+    def sample(
+        self,
+        logits: Tensor,
+        previous_tokens: Tensor,
+        temperature: float,
+        top_k: int,
+        top_p: float,
+        repetition_penalty: float,
+        use_cuda_graph=False,
+        idx=-1,
+    ) -> Tensor:
+        if use_cuda_graph and torch.cuda.is_available() and self.__CUDAGraph is None and idx > 0:
+            self.logits.copy_(logits)
+            self.capture(temperature, top_k, top_p)
+        if self.__CUDAGraph is not None:
+            self.logits.copy_(logits)
+            self.apply_repetition_penalty(self.logits, previous_tokens, repetition_penalty)
+            self.__CUDAGraph.replay()
+            samples = self.samples.clone()
+        else:
+            samples = self.__sample(logits, previous_tokens, temperature, top_k, top_p, repetition_penalty)[0]
+
+        return samples
+
+
+class KVCacheABC(ABC, nn.Module):
+    def __init__(self, *args, **kwds) -> None:
+        super().__init__()
+        self.k_cache: Tensor
+        self.v_cache: Tensor
+        self.n_head: int
+        self.head_dim: int
+        self.batch_size: int
+        self.max_seq_length: int
+
+    def empty(self):
+        self.k_cache.zero_()
+        self.v_cache.zero_()
+
+    @abstractmethod
+    def update(self, input_pos: Tensor, k_val: Tensor, v_val: Tensor, *args, **kwds) -> Tuple[Tensor, Tensor]: ...
+
+    @abstractmethod
+    def prefill_kv(self, k_val: Tensor, v_val: Tensor, bs: int) -> None: ...
+
+    def forward(self):
+        raise NotImplementedError()
+
+
+class KVCacheNHD(KVCacheABC):
+    def __init__(self, batch_size, max_seq_length, n_heads, head_dim):
+        super().__init__()
+        assert batch_size > 0
+        cache_shape = (batch_size, max_seq_length, n_heads, head_dim)
+        self.n_head = n_heads
+        self.head_dim = head_dim
+        self.batch_size = batch_size
+        self.max_seq_length = max_seq_length
+
+        self.register_buffer("k_cache", torch.zeros(size=cache_shape), persistent=False)
+        self.register_buffer("v_cache", torch.zeros(size=cache_shape), persistent=False)
+
+    def update(self, input_pos: Tensor, k_val: Tensor, v_val: Tensor):
+        # input_pos: [B, ], k_val: [B, 1, H, D]
+
+        index = (
+            (input_pos - 1)
+            .unsqueeze(-1)
+            .unsqueeze(-1)
+            .unsqueeze(-1)
+            .expand(
+                -1,
+                -1,
+                self.n_head,
+                self.head_dim,
+            )
+            .to(torch.int64)
+        )  # (bs, 1, num_head, head_dim)
+
+        k_out = self.k_cache
+        v_out = self.v_cache
+        k_out.scatter_(1, index, k_val)
+        v_out.scatter_(1, index, v_val)
+
+        return k_out, v_out
+
+    def empty(self):
+        self.k_cache.zero_()
+        self.v_cache.zero_()
+
+    def prefill_kv(self, k_val: Tensor, v_val: Tensor, bs: int):
+        # input_pos: int, k_val: [B, S, H, D]
+
+        self.k_cache[[bs], : k_val.shape[1]] = k_val
+        self.v_cache[[bs], : v_val.shape[1]] = v_val
+
+
+class KVCacheHND(KVCacheABC):
+    def __init__(self, batch_size, max_seq_length, n_heads, head_dim):
+        super().__init__()
+        assert batch_size > 0
+        cache_shape = (batch_size, n_heads, max_seq_length, head_dim)
+        self.n_head = n_heads
+        self.head_dim = head_dim
+        self.batch_size = batch_size
+        self.max_seq_length = max_seq_length
+
+        self.register_buffer("k_cache", torch.zeros(size=cache_shape), persistent=False)
+        self.register_buffer("v_cache", torch.zeros(size=cache_shape), persistent=False)
+
+    def update(self, input_pos: Tensor, k_val: Tensor, v_val: Tensor):
+        # input_pos: [B, ], k_val: [B, H, 1, D]
+
+        index = (
+            (input_pos - 1)
+            .unsqueeze(-1)
+            .unsqueeze(-1)
+            .unsqueeze(-1)
+            .expand(
+                -1,
+                self.n_head,
+                -1,
+                self.head_dim,
+            )
+            .to(torch.int64)
+        )  # (bs, num_head, 1, head_dim)
+
+        k_out = self.k_cache
+        v_out = self.v_cache
+        k_out.scatter_(2, index, k_val)
+        v_out.scatter_(2, index, v_val)
+
+        return k_out, v_out
+
+    def empty(self):
+        self.k_cache.zero_()
+        self.v_cache.zero_()
+
+    def prefill_kv(self, k_val: Tensor, v_val: Tensor, bs: int):
+        # input_pos: int, k_val: [B, S, H, D]
+
+        self.k_cache[[bs], :, : k_val.shape[1]] = k_val.transpose(1, 2)
+        self.v_cache[[bs], :, : v_val.shape[1]] = v_val.transpose(1, 2)
+
+
+class AttentionABC(ABC, nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.n_head: int
+        self.hidden_dim: int
+        self.head_dim: int
+
+        # key, query, value projections for all heads, but in a batch
+        self.in_proj: nn.Linear
+        self.out_proj: nn.Linear
+
+        self.dropout = nn.Dropout(0.1)
+
+        self._register_load_state_dict_pre_hook(self.load_hook)
+
+    def load_hook(self, state_dict: dict, prefix, *args):
+        keys_to_modify = [key for key in state_dict if "in_proj_" in key]
+        for key in keys_to_modify:
+            new_key = key.replace("in_proj_", "in_proj.")  # in_proj_ -> in_proj.
+            state_dict[new_key] = state_dict.pop(key)
+
+    @abstractmethod
+    def forward(self, x: Tensor, input_pos: Tensor, kv_cache: KVCacheABC, *args, **kwds) -> Tensor: ...
+
+    def prefill(self, x: Tensor, mask: Tensor, kv_cache: KVCacheABC) -> Tensor:
+        bsz = x.size(0)
+
+        outputs = []
+
+        for bs in range(bsz):
+            x_b = x[bs].unsqueeze(0)
+
+            q, k, v = self.in_proj.forward(x_b.unsqueeze(0)).chunk(3, dim=-1)
+
+            q = q.contiguous().view(1, -1, self.n_head, self.head_dim)
+            k = k.contiguous().view(1, -1, self.n_head, self.head_dim)
+            v = v.contiguous().view(1, -1, self.n_head, self.head_dim)
+
+            kv_cache.prefill_kv(k, v, bs)
+
+            q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v))
+
+            attn_mask = mask[bs].unsqueeze(0).unsqueeze(0).expand(1, self.n_head, -1, -1)
+
+            attn = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+
+            attn = self.dropout.forward(attn)
+
+            attn = attn.transpose(1, 2).contiguous().view(1, -1, self.hidden_dim)
+
+            output = self.out_proj.forward(attn)
+
+            outputs.append(output.squeeze(0))
+
+        return torch.nested.nested_tensor(outputs)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.linear1 = nn.Linear(dim, hidden_dim, bias=True)
+        self.linear2 = nn.Linear(hidden_dim, dim, bias=True)
+        self.dropout = nn.Dropout(0.1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.dropout.forward(self.linear2(self.dropout.forward(F.relu(self.linear1(x)))))
+
+
+class TransformerBlockABC(ABC, nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.hidden_dim: int
+        self.attention: AttentionABC
+        self.feed_forward: FeedForward
+        self.attention_norm: nn.LayerNorm
+        self.ffn_norm: nn.LayerNorm
+        self.dropout = nn.Dropout(0.1)
+
+        self._register_load_state_dict_pre_hook(self.load_hook)
+
+    def load_hook(self, state_dict: dict[str, Tensor], prefix, *args):
+        for key in list(state_dict.keys()):
+            new_key = (
+                key.replace("self_attn", "attention")
+                .replace("linear", "feed_forward.linear")
+                .replace("norm1", "attention_norm")
+                .replace("norm2", "ffn_norm")
+            )
+            state_dict[new_key] = state_dict.pop(key)
+
+    def forward(self, x: Tensor, input_pos: Tensor, kv_cache: KVCacheABC, *args, **kwds) -> Tensor:
+        h = self.attention_norm.forward(
+            x
+            + self.dropout.forward(
+                self.attention.forward(
+                    x,
+                    input_pos,
+                    kv_cache,
+                    *args,
+                    **kwds,
+                )
+            )
+        )
+        out = self.ffn_norm.forward(h + self.feed_forward.forward(h))
+        return out
+
+    def prefill(self, x: Tensor, mask: Tensor, kv_cache: KVCacheABC) -> Tensor:
+        h = self.attention_norm.forward(
+            x
+            + self.dropout.forward(
+                self.attention.prefill(
+                    x,
+                    mask,
+                    kv_cache,
+                )
+            )
+        )
+        out = self.ffn_norm.forward(h + self.feed_forward.forward(h))
+        return out
+
+
+class TransformerDecoderABC(ABC, nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+        self.hidden_dim: int
+        self.n_head: int
+        self.head_dim: int
+        self.vocab_size: int
+        self.n_layer: int
+
+        self.layers: MutableSequence[TransformerBlockABC]
+
+        self.max_seq_length: int
+        self.max_batch_size: int
+
+        self.input_pos: Tensor
+        self.xy_pos: Tensor
+        self.xy_dec: Tensor
+
+    def forward(self, input_pos: Tensor, x: Tensor, kv_caches: MutableSequence[KVCacheABC], *args, **kwds):
+        for layer, kv_cache in zip(self.layers, kv_caches):
+            x = layer.forward(x, input_pos, kv_cache, *args, **kwds)
+        return x
+
+    def prefill(self, x: Tensor, mask: Tensor, kv_caches: MutableSequence[KVCacheABC]):
+        for layer, kv_cache in zip(self.layers, kv_caches):
+            x = layer.prefill(x, mask, kv_cache)
+        return x
+
+
+class T2SDecoderABC(ABC, nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+        self.n_layer: int
+        self.hidden_dim: int
+        self.n_head: int
+
+        self.head_dim: int
+        self.embedding_dim: int
+        self.vocab_size: int
+        self.phoneme_vocab_size: int
+        self.p_dropout: float
+        self.max_seq_length: int
+        self.max_batch_size: int
+        self.EOS: int
+
+        self.bert_proj: nn.Linear
+        self.ar_text_embedding: TokenEmbedding
+        self.ar_text_position: SinePositionalEmbedding
+        self.ar_audio_embedding: TokenEmbedding
+        self.ar_audio_position: SinePositionalEmbedding
+        self.ar_predict_layer: nn.Linear
+        self.h: TransformerDecoderABC
+
+        self.kv_class: Type[KVCacheNHD] | Type[KVCacheHND]
+
+        self._register_load_state_dict_pre_hook(self.load_hook)
+
+    def load_hook(self, state_dict, prefix, *args):
+        model_keys = [key for key in state_dict if key.startswith("model.")]
+        for key in model_keys:
+            new_key = key[len("model.") :]
+            state_dict[new_key] = state_dict.pop(key)
+
+    def init_cache(self, bsz: int = 0) -> MutableSequence[KVCacheABC]:
+        bsz = bsz or self.h.max_batch_size
+        assert bsz <= self.h.max_batch_size
+        seq_lens = self.h.max_seq_length
+        device = self.bert_proj.bias.device
+        dtype = self.bert_proj.bias.dtype
+        kvclass = self.kv_class
+        return nn.ModuleList(
+            [kvclass(bsz, seq_lens, self.n_head, self.head_dim) for _ in range(self.n_layer)],
+        ).to(device, dtype)  # type: ignore
+
+    @abstractmethod
+    def embed(self, x: List[torch.Tensor], y: torch.Tensor, bert_features: List[Tensor]) -> Tensor: ...
+
+    def compile(self, *args, **kwds):
+        torch._inductor.config.triton.cudagraph_skip_dynamic_graphs = True
+        torch._inductor.config.coordinate_descent_tuning = True
+        torch._inductor.config.triton.unique_kernel_names = True
+        # Experimental features to reduce compilation times, will be on by default in future
+        torch._inductor.config.fx_graph_cache = True
+        torch._inductor.config.triton.cudagraph_trees = True
+        torch._inductor.config.triton.cudagraph_support_input_mutation = True
+        self.h.compile(fullgraph=True, mode="reduce-overhead")
+
+    def capture(self, input_pos: Tensor, x: Tensor, x_dec: Tensor, *args, **kwds) -> CUDAGraph:
+        s = torch.cuda.Stream()
+        s.wait_stream(torch.cuda.current_stream())
+
+        graph = torch.cuda.CUDAGraph()
+
+        with torch.cuda.stream(s):  # type: ignore
+            for _ in range(5):
+                self.h.forward(input_pos, x, *args, **kwds)
+        torch.cuda.current_stream().wait_stream(s)
+
+        with torch.cuda.graph(graph):
+            x_dec.copy_(self.h.forward(input_pos, x, *args, **kwds))
+        torch.cuda.synchronize()
+
+        return graph
+
+    @abstractmethod
+    def pre_forward(self, session: Any) -> Tuple[List, Dict]: ...
+
+    @abstractmethod
+    def post_forward(self, idx: int, session: Any) -> None: ...
+
+
+class TorchProfiler:
+    def __init__(self, debug: bool, log_dir: str = "./profiler") -> None:
+        self.debug = debug
+        self.log_dir = log_dir
+        self.__profiler: torch.profiler.profile
+
+        if self.debug and not os.path.exists(self.log_dir):
+            os.makedirs(self.log_dir)
+
+        self.tensorboard_handler = tensorboard_trace_handler(self.log_dir)
+
+    def profiler_callback(self, prof: torch.profiler.profile):
+        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=30))
+        print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=30))
+        self.tensorboard_handler(prof)
+
+    @staticmethod
+    def three_step_schedule(step: int) -> ProfilerAction:
+        if step == 0:
+            return ProfilerAction.NONE
+        elif step == 1:
+            return ProfilerAction.RECORD
+        elif step == 2:
+            return ProfilerAction.RECORD_AND_SAVE
+        else:
+            return ProfilerAction.NONE
+
+    def start(self):
+        if not self.debug:
+            return
+        assert self.__profiler is not None
+        self.__profiler.step()
+
+    def end(self):
+        if not self.debug:
+            return
+        assert self.__profiler is not None
+        self.__profiler.step()
+
+    def profiler(self):
+        if self.debug:
+            activities_list = [torch.profiler.ProfilerActivity.CPU]
+            if torch.cuda.is_available():
+                activities_list.append(torch.profiler.ProfilerActivity.CUDA)
+
+            self.__profiler = torch.profiler.profile(
+                activities=activities_list,
+                record_shapes=True,
+                with_stack=True,
+                with_modules=True,
+                profile_memory=True,
+                schedule=self.three_step_schedule,
+                on_trace_ready=self.profiler_callback,
+            )
+            return self.__profiler
+        else:
+            return nullcontext()
+
+    def record(self, func_name: str):
+        if self.debug:
+            return torch.profiler.record_function(func_name)
+        else:
+            return nullcontext()

AR/models/t2s_model_flash_attn.py

@@ -0,0 +1,357 @@
+import gc
+import os
+import time
+import traceback
+from typing import Dict, List, Tuple
+
+import flash_attn  # type: ignore
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+
+from AR.models.structs import T2SRequest, T2SResult, T2SSession
+from AR.models.t2s_model_abc import (
+    AttentionABC,
+    FeedForward,
+    KVCacheABC,
+    KVCacheNHD,
+    T2SDecoderABC,
+    TorchProfiler,
+    TransformerBlockABC,
+    TransformerDecoderABC,
+)
+from AR.modules.embedding import (
+    SinePositionalEmbeddingNested as SinePositionalEmbedding,
+)
+from AR.modules.embedding import TokenEmbedding
+
+Tensor = torch.Tensor
+
+
+class Attention(AttentionABC):
+    def __init__(self, n_head: int, hidden_dim: int):
+        super().__init__()
+        self.n_head = n_head
+        self.hidden_dim = hidden_dim
+        assert hidden_dim % n_head == 0
+        self.head_dim = hidden_dim // n_head
+
+        self.in_proj = nn.Linear(hidden_dim, hidden_dim * 3, bias=True)
+        self.out_proj = nn.Linear(hidden_dim, hidden_dim, bias=True)
+
+    def forward(self, x: Tensor, input_pos: Tensor, kv_cache: KVCacheABC, *args, **kwds) -> Tensor:
+        bsz, seqlen, _ = x.shape
+
+        q, k, v = self.in_proj.forward(x).chunk(3, dim=-1)
+
+        q = q.view(bsz, seqlen, self.n_head, self.head_dim)
+        k = k.view(bsz, seqlen, self.n_head, self.head_dim)
+        v = v.view(bsz, seqlen, self.n_head, self.head_dim)
+
+        attn: Tensor = flash_attn.flash_attn_with_kvcache(
+            q, kv_cache.k_cache, kv_cache.v_cache, k, v, cache_seqlens=input_pos - 1
+        )
+
+        attn = self.dropout.forward(attn)
+
+        attn = attn.view(bsz, seqlen, self.hidden_dim)
+
+        attn = self.out_proj.forward(attn)
+
+        return attn
+
+
+class TransformerBlock(TransformerBlockABC):
+    def __init__(self, n_head, ffn_dim, hidden_dim) -> None:
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.attention = Attention(n_head, hidden_dim)
+        self.feed_forward = FeedForward(hidden_dim, ffn_dim)
+        self.attention_norm = nn.LayerNorm([self.hidden_dim])
+        self.ffn_norm = nn.LayerNorm([self.hidden_dim])
+
+
+class TransformerDecoder(TransformerDecoderABC):
+    def __init__(
+        self,
+        hidden_dim,
+        n_layer,
+        n_head,
+        ffn_dim,
+        vocab_size,
+        max_seq_length,
+        max_batch_size,
+    ) -> None:
+        super().__init__()
+
+        self.hidden_dim = hidden_dim
+        self.n_head = n_head
+        assert hidden_dim % n_head == 0
+
+        self.head_dim = hidden_dim // n_head
+        self.vocab_size = vocab_size
+
+        self.n_layer = n_layer
+
+        self.layers = nn.ModuleList(  # type: ignore
+            TransformerBlock(n_head, ffn_dim, hidden_dim) for _ in range(n_layer)
+        )
+
+        self.max_seq_length: int = max_seq_length
+        self.max_batch_size: int = max_batch_size
+
+        self.setup_caches(self.max_batch_size, self.max_seq_length)
+
+    def setup_caches(self, max_batch_size=10, max_seq_length=2500):
+        self.max_seq_length = max_seq_length
+        self.max_batch_size = max_batch_size
+
+
+class T2SDecoder(T2SDecoderABC):
+    def __init__(
+        self,
+        config,
+        *args,
+        norm_first=False,
+        max_seq_length=2500,
+        max_batch_size=10,
+        **kwds,
+    ) -> None:
+        super().__init__()
+
+        hidden_dim = config["model"]["hidden_dim"]
+        embedding_dim = config["model"]["embedding_dim"]
+        n_head = config["model"]["head"]
+        n_layer = config["model"]["n_layer"]
+        vocab_size = config["model"]["vocab_size"]
+        phoneme_vocab_size = config["model"]["phoneme_vocab_size"]
+        p_dropout = config["model"]["dropout"]
+        EOS = config["model"]["EOS"]
+        ffn_dim = hidden_dim * 4
+        self.norm_first = norm_first
+
+        self.n_layer = n_layer
+        self.hidden_dim = hidden_dim
+        self.n_head = n_head
+        assert hidden_dim % n_head == 0
+
+        self.head_dim = hidden_dim // n_head
+        self.embedding_dim = embedding_dim
+        self.vocab_size = vocab_size
+        self.phoneme_vocab_size = phoneme_vocab_size
+        self.p_dropout = p_dropout
+        self.max_seq_length = max_seq_length
+        self.max_batch_size = max_batch_size
+        self.EOS = EOS
+        assert self.EOS == self.vocab_size - 1
+
+        self.bert_proj = nn.Linear(1024, self.embedding_dim)
+        self.ar_text_embedding = TokenEmbedding(self.embedding_dim, self.phoneme_vocab_size, self.p_dropout)
+        self.ar_text_position = SinePositionalEmbedding(
+            self.embedding_dim,
+            dropout=0.1,
+            scale=False,
+            alpha=True,
+            max_batch_size=max_batch_size,
+            max_seq_len=max_seq_length,
+        )
+        self.ar_audio_embedding = TokenEmbedding(self.embedding_dim, self.vocab_size, self.p_dropout)
+        self.ar_audio_position = SinePositionalEmbedding(
+            self.embedding_dim,
+            dropout=0.1,
+            scale=False,
+            alpha=True,
+            max_batch_size=max_batch_size,
+            max_seq_len=max_seq_length,
+        )
+        self.ar_predict_layer = nn.Linear(self.hidden_dim, self.vocab_size, bias=False)
+        self.h: TransformerDecoderABC = TransformerDecoder(
+            hidden_dim, n_layer, n_head, ffn_dim, vocab_size, max_seq_length, max_batch_size
+        )
+
+        self.kv_class = KVCacheNHD
+        self._register_load_state_dict_pre_hook(self.load_hook)
+
+    def embed(
+        self,
+        x: List[torch.Tensor],
+        y: torch.Tensor,
+        bert_features: List[torch.Tensor],
+    ):
+        x_nested = torch.nested.nested_tensor(x)
+        assert x_nested.size(0) <= self.max_batch_size
+        bert_features_nested = torch.nested.nested_tensor(list(map(lambda x: x.transpose(0, 1), bert_features)))
+
+        x_emb = self.ar_text_embedding.forward(x_nested)
+        bert = self.bert_proj.forward(bert_features_nested)
+        x_emb = x_emb + bert
+        x_pos = self.ar_text_position.prefill(x_emb)
+
+        y_nested = torch.nested.nested_tensor(list(y.unbind(0)))
+        y_emb = self.ar_audio_embedding.forward(y_nested)
+        y_pos = self.ar_audio_position.prefill(y_emb)
+
+        xy_pos = torch.nested.nested_tensor([torch.cat([x_pos[i], y_pos[i]]) for i in range(len(x))])
+        return xy_pos
+
+    def post_forward(self, idx: int, session: T2SSession) -> None:
+        pass
+
+    def pre_forward(self, session: T2SSession) -> Tuple[List, Dict]:
+        return list(), dict()
+
+
+class CUDAGraphRunner:
+    def __init__(
+        self,
+        decoder_model: T2SDecoderABC,
+        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}
+        self.device = device
+        self.dtype = dtype
+
+        self.decoder_path: os.PathLike
+        self.decoder_model: T2SDecoderABC = decoder_model.to(self.device, self.dtype)
+
+    def _handle_request(self, request: T2SRequest) -> List[torch.Tensor]:
+        with self.device:
+            decoder = self.decoder_model
+            session = T2SSession(decoder, request, device=self.device, dtype=self.dtype)
+
+            y = session.y
+            bsz = y.size(0)
+            t1 = 0.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, session.kv_cache)
+                        xy_dec = torch.stack([t[[-1]] for t in xy_dec.unbind()])
+                    else:
+                        if request.use_cuda_graph and session.graph is None and torch.cuda.is_available():
+                            session.xy_pos_.copy_(session.xy_pos)
+                            args, kwds = decoder.pre_forward(session)
+                            session.graph = decoder.capture(
+                                session.input_pos,
+                                session.xy_pos_,
+                                session.xy_dec_,
+                                kv_caches=session.kv_cache,
+                                *args,
+                                **kwds,
+                            )
+
+                        torch_profiler.start()
+                        with torch_profiler.record("AR"):
+                            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(
+                                    session.input_pos,
+                                    session.xy_pos,
+                                    session.kv_cache,
+                                    *args,
+                                    **kwds,
+                                )
+                    decoder.post_forward(idx, session)
+                    logits = decoder.ar_predict_layer(xy_dec[:, -1])
+                    session.input_pos.add_(1)
+
+                    if idx == 0:
+                        logits = logits[:, :-1]
+
+                    with torch_profiler.record("Sampling"):
+                        samples = session.sampler.sample(
+                            logits=logits,
+                            previous_tokens=session.y,
+                            top_k=request.top_k,
+                            top_p=request.top_p,
+                            repetition_penalty=request.repetition_penalty,
+                            temperature=request.temperature,
+                            use_cuda_graph=False,
+                            idx=idx,
+                        )
+
+                        session.y = torch.cat([session.y, samples], dim=1)
+
+                    with torch_profiler.record("EOS"):
+                        EOS_mask = (samples[:, 0] == decoder.EOS) | (torch.argmax(logits, dim=-1) == decoder.EOS)
+                        EOS_indices: List[int] = torch.where(EOS_mask)[0].tolist()
+
+                        for i in EOS_indices:
+                            if not session.completed[i]:
+                                session.y_results[i] = session.y[i, session.y_len : -1]
+                                session.completed[i] = True
+
+                        if all(session.completed):
+                            if session.y.size(1) == 0:
+                                session.y = torch.cat([session.y, torch.zeros_like(samples)], dim=1)
+                                tqdm.write("Bad Zero Prediction")
+                            else:
+                                tqdm.write(
+                                    f"T2S Decoding EOS {session.prefill_len.tolist().__str__().strip('[]')} -> \n{[i.size(0) for i in session.y_results].__str__().strip('[]')}"
+                                )
+                                tqdm.write(f"Infer Speed: {(idx - 1) / (time.perf_counter() - t1):.2f} token/s")
+                            break
+
+                        if (
+                            request.early_stop_num != -1
+                            and (session.y.size(1) - session.y_len) > request.early_stop_num
+                        ):
+                            for i in range(bsz):
+                                if not session.completed[i]:
+                                    session.y_results[i] = session.y[i, session.y_len :]
+                                    session.completed[i] = True
+                            break
+
+                    with torch_profiler.record("NextPos"):
+                        y_emb = decoder.ar_audio_embedding(session.y[:, -1:])
+                        session.xy_pos = decoder.ar_audio_position.forward(session.input_pos - session.x_lens, y_emb)
+
+                    if idx == 2:
+                        t1 = time.perf_counter()
+
+                    if idx == 51:
+                        torch_profiler.end()
+
+            match session.device.type:
+                case "cuda":
+                    torch.cuda.empty_cache()
+                case "mps":
+                    torch.mps.empty_cache()
+                case "xpu":
+                    torch.xpu.empty_cache()
+                case "mtia":
+                    torch.mtia.empty_cache()
+            gc.collect()
+
+            return session.y_results[: request.valid_length]
+
+    def generate(self, request: T2SRequest):
+        try:
+            result = self._handle_request(request)
+            t2s_result = T2SResult(result=result, status="Success")
+        except Exception as e:
+            t2s_result = T2SResult(status="Error", exception=e, traceback=traceback.format_exc())
+        return t2s_result
+
+    @staticmethod
+    def load_decoder(weights_path: os.PathLike, implement: str = "flash_attn"):
+        print(f"Loading Text2Semantic Weights from {weights_path} with {implement.replace('_', ' ').title()} Implement")
+        module_path = f"AR.models.t2s_model_{implement.lower()}"
+        cls_name = "T2SDecoder"
+        mod = __import__(module_path, fromlist=[cls_name])
+        decoder_cls: T2SDecoderABC = getattr(mod, cls_name)
+        dict_s1 = torch.load(weights_path, map_location="cpu", weights_only=False, mmap=True)
+        config = dict_s1["config"]
+        decoder: T2SDecoderABC = decoder_cls(config, max_batch_size=1)
+        state_dict = dict_s1["weight"]
+        decoder.load_state_dict(state_dict)
+        return decoder.eval()

AR/modules/embedding.py

@@ -60,14 +60,11 @@ class SinePositionalEmbedding(nn.Module):
                 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)
+            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)
+            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)
@@ -79,3 +76,68 @@ class SinePositionalEmbedding(nn.Module):
         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,
+        embedding_dim: int,
+        dropout: float = 0.0,
+        scale: bool = False,
+        alpha: bool = False,
+        max_batch_size: int = 20,
+        max_seq_len: int = 2500,
+    ):
+        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.max_batch_size = max_batch_size
+        self.max_seq_len = max_seq_len
+
+        self.reverse = False
+        self.register_buffer("pe", torch.zeros(max_batch_size, max_seq_len, embedding_dim), persistent=False)
+        self.pe: torch.Tensor
+        self.compute_pe()
+
+    def compute_pe(self):
+        """Reset the positional encodings."""
+        if self.reverse:
+            position = torch.arange(self.max_seq_len - 1, -1, -1.0, dtype=torch.float32).unsqueeze(1)
+        else:
+            position = torch.arange(self.max_seq_len, 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 = self.pe
+        pe[:, :, 0::2] = torch.sin(position * div_term)
+        pe[:, :, 1::2] = torch.cos(position * div_term)
+
+    def forward(self, input_pos: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            input_pos (Tensor): [batch_size, ]
+            x (Tensor): [batch_size, 1, embed_dim]
+
+        Returns:
+            embedded_x (Tensor): [batch_size, 1, embed_dim]
+        """
+
+        batch_size = x.shape[0]
+        pe_values = self.pe[torch.arange(batch_size), input_pos - 1]  # (batch_size, embed_dim)
+
+        return x * self.x_scale + self.alpha * pe_values.unsqueeze(1)  # (batch_size, 1, embed_dim)
+
+    def prefill(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x (Tensor): Nested Seqlen [batch_size, seq_len, embed_dim]
+
+        Returns:
+            embedded_x (Tensor): Nested Seqlen [batch_size, seq_len, embed_dim]
+        """
+
+        input_pos: torch.Tensor = torch.tensor([i.shape[0] for i in x.unbind()])
+        pe_values = torch.nested.nested_tensor([self.pe[i, : input_pos[i], :] for i in range(input_pos.size(0))])
+        return x * self.x_scale + self.alpha.item() * pe_values

inference_webui.py

@@ -1,12 +1,35 @@
 import os
-os.makedirs("pretrained_models",exist_ok=True)
+
+os.makedirs("pretrained_models", exist_ok=True)
 from huggingface_hub import snapshot_download
-snapshot_download(repo_id="lj1995/GPT-SoVITS",repo_type="model",allow_patterns="chinese*",local_dir="pretrained_models",)
-snapshot_download(repo_id="lj1995/GPT-SoVITS",repo_type="model",allow_patterns="s1v3.ckpt",local_dir="pretrained_models",)
-snapshot_download(repo_id="lj1995/GPT-SoVITS",repo_type="model",allow_patterns="sv*",local_dir="pretrained_models",)
-snapshot_download(repo_id="lj1995/GPT-SoVITS",repo_type="model",allow_patterns="v2Pro/s2Gv2ProPlus.pth",local_dir="pretrained_models",)
+
+snapshot_download(
+    repo_id="lj1995/GPT-SoVITS",
+    repo_type="model",
+    allow_patterns="chinese*",
+    local_dir="pretrained_models",
+)
+snapshot_download(
+    repo_id="lj1995/GPT-SoVITS",
+    repo_type="model",
+    allow_patterns="s1v3.ckpt",
+    local_dir="pretrained_models",
+)
+snapshot_download(
+    repo_id="lj1995/GPT-SoVITS",
+    repo_type="model",
+    allow_patterns="sv*",
+    local_dir="pretrained_models",
+)
+snapshot_download(
+    repo_id="lj1995/GPT-SoVITS",
+    repo_type="model",
+    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)
@@ -17,42 +40,47 @@ 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
-import gradio.analytics as analytics
-analytics.version_check = lambda:None
-analytics.get_local_ip_address= lambda :"127.0.0.1"##不干掉本地联不通亚马逊的get_local_ip服务器
-import nltk,torchaudio
-nltk.download('averaged_perceptron_tagger_eng')
-import LangSegment, os, re, sys, json
+
+nltk.download("averaged_perceptron_tagger_eng")
+import json
+import os
 import pdb
+import re
+import sys
+
+import LangSegment
 import spaces
 import torch
 
-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"
-)
+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(['!', '?', '…', ',', '.', '-'," "])
+punctuation = set(["!", "?", "…", ",", ".", "-", " "])
 import gradio as gr
-from transformers import AutoModelForMaskedLM, AutoTokenizer
-import numpy as np
 import librosa
+import numpy as np
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
 from feature_extractor import cnhubert
 
 cnhubert.cnhubert_base_path = cnhubert_base_path
 
+from time import time as ttime
+
+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 AR.models.t2s_lightning_module import Text2SemanticLightningModule
 from text import cleaned_text_to_sequence
 from text.cleaner import clean_text
-from time import time as ttime
-from module.mel_processing import spectrogram_torch
-from tools.my_utils import load_audio
 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
@@ -65,30 +93,30 @@ if torch.cuda.is_available():
     is_half = True  # eval(os.environ.get("is_half", "True")) and torch.cuda.is_available()
 else:
     device = "cpu"
-    is_half=False
+    is_half = False
 
 dict_language_v1 = {
-    i18n("中文"): "all_zh",#全部按中文识别
-    i18n("英文"): "en",#全部按英文识别#######不变
-    i18n("日文"): "all_ja",#全部按日文识别
-    i18n("中英混合"): "zh",#按中英混合识别####不变
-    i18n("日英混合"): "ja",#按日英混合识别####不变
-    i18n("多语种混合"): "auto",#多语种启动切分识别语种
+    i18n("中文"): "all_zh",  # 全部按中文识别
+    i18n("英文"): "en",  # 全部按英文识别#######不变
+    i18n("日文"): "all_ja",  # 全部按日文识别
+    i18n("中英混合"): "zh",  # 按中英混合识别####不变
+    i18n("日英混合"): "ja",  # 按日英混合识别####不变
+    i18n("多语种混合"): "auto",  # 多语种启动切分识别语种
 }
 dict_language_v2 = {
-    i18n("中文"): "all_zh",#全部按中文识别
-    i18n("英文"): "en",#全部按英文识别#######不变
-    i18n("日���"): "all_ja",#全部按日文识别
-    i18n("粤语"): "all_yue",#全部按中文识别
-    i18n("韩文"): "all_ko",#全部按韩文识别
-    i18n("中英混合"): "zh",#按中英混合识别####不变
-    i18n("日英混合"): "ja",#按日英混合识别####不变
-    i18n("粤英混合"): "yue",#按粤英混合识别####不变
-    i18n("韩英混合"): "ko",#按韩英混合识别####不变
-    i18n("多语种混合"): "auto",#多语种启动切分识别语种
-    i18n("多语种混合(粤语)"): "auto_yue",#多语种启动切分识别语种
+    i18n("中文"): "all_zh",  # 全部按中文识别
+    i18n("英文"): "en",  # 全部按英文识别#######不变
+    i18n("日文"): "all_ja",  # 全部按日文识别
+    i18n("粤语"): "all_yue",  # 全部按中文识别
+    i18n("韩文"): "all_ko",  # 全部按韩文识别
+    i18n("中英混合"): "zh",  # 按中英混合识别####不变
+    i18n("日英混合"): "ja",  # 按日英混合识别####不变
+    i18n("粤英混合"): "yue",  # 按粤英混合识别####不变
+    i18n("韩英混合"): "ko",  # 按韩英混合识别####不变
+    i18n("多语种混合"): "auto",  # 多语种启动切分识别语种
+    i18n("多语种混合(粤语)"): "auto_yue",  # 多语种启动切分识别语种
 }
-dict_language = dict_language_v1 if version =='v1' else dict_language_v2
+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)
@@ -149,13 +177,13 @@ else:
     ssl_model = ssl_model.to(device)
 
 
-def change_sovits_weights(sovits_path,prompt_language=None,text_language=None):
+def change_sovits_weights(sovits_path, prompt_language=None, text_language=None):
     global vq_model, hps, version, dict_language
     dict_s2 = torch.load(sovits_path, map_location="cpu")
     hps = dict_s2["config"]
     hps = DictToAttrRecursive(hps)
     hps.model.semantic_frame_rate = "25hz"
-    if dict_s2['weight']['enc_p.text_embedding.weight'].shape[0] == 322:
+    if dict_s2["weight"]["enc_p.text_embedding.weight"].shape[0] == 322:
         hps.model.version = "v1"
     else:
         hps.model.version = "v2"
@@ -165,9 +193,9 @@ def change_sovits_weights(sovits_path,prompt_language=None,text_language=None):
         hps.data.filter_length // 2 + 1,
         hps.train.segment_size // hps.data.hop_length,
         n_speakers=hps.data.n_speakers,
-        **hps.model
+        **hps.model,
     )
-    if ("pretrained" not in sovits_path):
+    if "pretrained" not in sovits_path:
         del vq_model.enc_q
     if is_half == True:
         vq_model = vq_model.half().to(device)
@@ -175,43 +203,48 @@ def change_sovits_weights(sovits_path,prompt_language=None,text_language=None):
         vq_model = vq_model.to(device)
     vq_model.eval()
     print(vq_model.load_state_dict(dict_s2["weight"], strict=False))
-    dict_language = dict_language_v1 if version =='v1' else dict_language_v2
+    dict_language = dict_language_v1 if version == "v1" else dict_language_v2
     if prompt_language is not None and text_language is not None:
         if prompt_language in list(dict_language.keys()):
-            prompt_text_update, prompt_language_update = {'__type__':'update'},  {'__type__':'update', 'value':prompt_language}
+            prompt_text_update, prompt_language_update = (
+                {"__type__": "update"},
+                {"__type__": "update", "value": prompt_language},
+            )
         else:
-            prompt_text_update = {'__type__':'update', 'value':''}
-            prompt_language_update = {'__type__':'update', 'value':i18n("中文")}
+            prompt_text_update = {"__type__": "update", "value": ""}
+            prompt_language_update = {"__type__": "update", "value": i18n("中文")}
         if text_language in list(dict_language.keys()):
-            text_update, text_language_update = {'__type__':'update'}, {'__type__':'update', 'value':text_language}
+            text_update, text_language_update = {"__type__": "update"}, {"__type__": "update", "value": text_language}
         else:
-            text_update = {'__type__':'update', 'value':''}
-            text_language_update = {'__type__':'update', 'value':i18n("中文")}
-        return  {'__type__':'update', 'choices':list(dict_language.keys())}, {'__type__':'update', 'choices':list(dict_language.keys())}, prompt_text_update, prompt_language_update, text_update, text_language_update
-
+            text_update = {"__type__": "update", "value": ""}
+            text_language_update = {"__type__": "update", "value": i18n("中文")}
+        return (
+            {"__type__": "update", "choices": list(dict_language.keys())},
+            {"__type__": "update", "choices": list(dict_language.keys())},
+            prompt_text_update,
+            prompt_language_update,
+            text_update,
+            text_language_update,
+        )
 
 
 change_sovits_weights("pretrained_models/v2Pro/s2Gv2ProPlus.pth")
 
 
 def change_gpt_weights(gpt_path):
-    global hz, max_sec, t2s_model, config
-    hz = 50
+    global t2s_model, config
     dict_s1 = torch.load(gpt_path, map_location="cpu")
     config = dict_s1["config"]
-    max_sec = config["data"]["max_sec"]
-    t2s_model = Text2SemanticLightningModule(config, "****", is_train=False)
-    t2s_model.load_state_dict(dict_s1["weight"])
-    if is_half == True:
-        t2s_model = t2s_model.half()
-    t2s_model = t2s_model.to(device)
-    t2s_model.eval()
-    total = sum([param.nelement() for param in t2s_model.parameters()])
+    t2s_model = CUDAGraphRunner(
+        CUDAGraphRunner.load_decoder(gpt_path), torch.device(device), torch.float16 if is_half else torch.float32
+    )
+    total = sum(p.numel() for p in t2s_model.decoder_model.parameters())
     print("Number of parameter: %.2fM" % (total / 1e6))
 
 
 change_gpt_weights("pretrained_models/s1v3.ckpt")
 from sv import SV
+
 sv_cn_model = SV(device, is_half)
 
 resample_transform_dict = {}
@@ -261,11 +294,14 @@ def clean_text_inf(text, language, version):
     phones = cleaned_text_to_sequence(phones, version)
     return phones, word2ph, norm_text
 
-dtype=torch.float16 if is_half == True else torch.float32
+
+dtype = torch.float16 if is_half == True else torch.float32
+
+
 def get_bert_inf(phones, word2ph, norm_text, language):
-    language=language.replace("all_","")
+    language = language.replace("all_", "")
     if language == "zh":
-        bert = get_bert_feature(norm_text, word2ph).to(device)#.to(dtype)
+        bert = get_bert_feature(norm_text, word2ph).to(device)  # .to(dtype)
     else:
         bert = torch.zeros(
             (1024, len(phones)),
@@ -275,7 +311,21 @@ def get_bert_inf(phones, word2ph, norm_text, language):
     return bert
 
 
-splits = {"，", "。", "？", "！", ",", ".", "?", "!", "~", ":", "：", "—", "…", }
+splits = {
+    "，",
+    "。",
+    "？",
+    "！",
+    ",",
+    ".",
+    "?",
+    "!",
+    "~",
+    ":",
+    "：",
+    "—",
+    "…",
+}
 
 
 def get_first(text):
@@ -283,8 +333,10 @@ def get_first(text):
     text = re.split(pattern, text)[0].strip()
     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
@@ -361,24 +413,44 @@ def merge_short_text_in_array(texts, threshold):
         if len(text) >= threshold:
             result.append(text)
             text = ""
-    if (len(text) > 0):
+    if len(text) > 0:
         if len(result) == 0:
             result.append(text)
         else:
             result[len(result) - 1] += text
     return result
 
+
 ##ref_wav_path+prompt_text+prompt_language+text(单个)+text_language+top_k+top_p+temperature
 # cache_tokens={}#暂未实现清理机制
-cache= {}
-@torch.inference_mode()
+cache = {}
+
+
 @spaces.GPU
-def get_tts_wav(ref_wav_path, prompt_text, prompt_language, text, text_language, how_to_cut=i18n("不切"), top_k=20, top_p=0.6, temperature=0.6, ref_free = False,speed=1,if_freeze=False,inp_refs=123):
+def get_tts_wav(
+    ref_wav_path,
+    prompt_text,
+    prompt_language,
+    text,
+    text_language,
+    how_to_cut=i18n("不切"),
+    top_k=20,
+    top_p=0.6,
+    temperature=0.6,
+    ref_free=False,
+    speed=1,
+    if_freeze=False,
+    inp_refs=123,
+):
     global cache
-    if ref_wav_path:pass
-    else:gr.Warning(i18n('请上传参考音频'))
-    if text:pass
-    else:gr.Warning(i18n('请填入推理文本'))
+    if ref_wav_path:
+        pass
+    else:
+        gr.Warning(i18n("请上传参考音频"))
+    if text:
+        pass
+    else:
+        gr.Warning(i18n("请填入推理文本"))
     t = []
     if prompt_text is None or len(prompt_text) == 0:
         ref_free = True
@@ -386,13 +458,14 @@ def get_tts_wav(ref_wav_path, prompt_text, prompt_language, text, text_language,
     prompt_language = dict_language[prompt_language]
     text_language = dict_language[text_language]
 
-
     if not ref_free:
         prompt_text = prompt_text.strip("\n")
-        if (prompt_text[-1] not in splits): prompt_text += "。" if prompt_language != "en" else "."
+        if prompt_text[-1] not in splits:
+            prompt_text += "。" if prompt_language != "en" else "."
         print(i18n("实际输入的参考文本:"), prompt_text)
     text = text.strip("\n")
-    if (text[0] not in splits and len(get_first(text)) < 4): text = "。" + text if text_language != "en" else "." + text
+    if text[0] not in splits and len(get_first(text)) < 4:
+        text = "。" + text if text_language != "en" else "." + text
 
     print(i18n("实际输入的目标文本:"), text)
     zero_wav = np.zeros(
@@ -402,7 +475,7 @@ def get_tts_wav(ref_wav_path, prompt_text, prompt_language, text, text_language,
     if not ref_free:
         with torch.no_grad():
             wav16k, sr = librosa.load(ref_wav_path, sr=16000)
-            if (wav16k.shape[0] > 160000 or wav16k.shape[0] < 48000):
+            if wav16k.shape[0] > 160000 or wav16k.shape[0] < 48000:
                 gr.Warning(i18n("参考音���在3~10秒范围外，请更换！"))
                 raise OSError(i18n("参考音频在3~10秒范围外，请更换！"))
             wav16k = torch.from_numpy(wav16k)
@@ -414,27 +487,23 @@ def get_tts_wav(ref_wav_path, prompt_text, prompt_language, text, text_language,
                 wav16k = wav16k.to(device)
                 zero_wav_torch = zero_wav_torch.to(device)
             wav16k = torch.cat([wav16k, zero_wav_torch])
-            ssl_content = ssl_model.model(wav16k.unsqueeze(0))[
-                "last_hidden_state"
-            ].transpose(
-                1, 2
-            )  # .float()
+            ssl_content = ssl_model.model(wav16k.unsqueeze(0))["last_hidden_state"].transpose(1, 2)  # .float()
             codes = vq_model.extract_latent(ssl_content)
             prompt_semantic = codes[0, 0]
             prompt = prompt_semantic.unsqueeze(0).to(device)
 
     t1 = ttime()
-    t.append(t1-t0)
+    t.append(t1 - t0)
 
-    if (how_to_cut == i18n("凑四句一切")):
+    if how_to_cut == i18n("凑四句一切"):
         text = cut1(text)
-    elif (how_to_cut == i18n("凑50字一切")):
+    elif how_to_cut == i18n("凑50字一切"):
         text = cut2(text)
-    elif (how_to_cut == i18n("按中文句号。切")):
+    elif how_to_cut == i18n("按中文句号。切"):
         text = cut3(text)
-    elif (how_to_cut == i18n("按英文句号.切")):
+    elif how_to_cut == i18n("按英文句号.切"):
         text = cut4(text)
-    elif (how_to_cut == i18n("按标点符号切")):
+    elif how_to_cut == i18n("按标点符号切"):
         text = cut5(text)
     while "\n\n" in text:
         text = text.replace("\n\n", "\n")
@@ -444,19 +513,20 @@ def get_tts_wav(ref_wav_path, prompt_text, prompt_language, text, text_language,
     texts = merge_short_text_in_array(texts, 5)
     audio_opt = []
     if not ref_free:
-        phones1,bert1,norm_text1=get_phones_and_bert(prompt_text, prompt_language, version)
+        phones1, bert1, norm_text1 = get_phones_and_bert(prompt_text, prompt_language, version)
 
-    for i_text,text in enumerate(texts):
+    for i_text, text in enumerate(texts):
         # 解决输入目标文本的空行导致报错的问题
-        if (len(text.strip()) == 0):
+        if len(text.strip()) == 0:
             continue
-        if (text[-1] not in splits): text += "。" if text_language != "en" else "."
+        if text[-1] not in splits:
+            text += "。" if text_language != "en" else "."
         print(i18n("实际输入的目标文本(每句):"), text)
-        phones2,bert2,norm_text2=get_phones_and_bert(text, text_language, version)
+        phones2, bert2, norm_text2 = get_phones_and_bert(text, text_language, version)
         print(i18n("前端处理后的文本(每句):"), norm_text2)
         if not ref_free:
             bert = torch.cat([bert1, bert2], 1)
-            all_phoneme_ids = torch.LongTensor(phones1+phones2).to(device).unsqueeze(0)
+            all_phoneme_ids = torch.LongTensor(phones1 + phones2).to(device).unsqueeze(0)
         else:
             bert = bert2
             all_phoneme_ids = torch.LongTensor(phones2).to(device).unsqueeze(0)
@@ -467,26 +537,33 @@ def get_tts_wav(ref_wav_path, prompt_text, prompt_language, text, text_language,
         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):pred_semantic=cache[i_text]
+        if i_text in cache and if_freeze == True:
+            pred_semantic = cache[i_text]
         else:
             with torch.no_grad():
-                pred_semantic, idx = t2s_model.model.infer_panel(
-                    all_phoneme_ids,
+                t2s_request = T2SRequest(
+                    [all_phoneme_ids.squeeze(0)],
                     all_phoneme_len,
-                    None if ref_free else prompt,
-                    bert,
-                    # prompt_phone_len=ph_offset,
+                    torch.zeros((1, 0)) if ref_free else prompt,
+                    [bert.squeeze(0)],
+                    valid_length=1,
                     top_k=top_k,
                     top_p=top_p,
                     temperature=temperature,
-                    early_stop_num=hz * max_sec,
+                    early_stop_num=1500,
+                    use_cuda_graph=True,
                 )
-                pred_semantic = pred_semantic[:, -idx:].unsqueeze(0)
-                cache[i_text]=pred_semantic
+                t2s_result = t2s_model.generate(t2s_request)
+                pred_semantic = t2s_result.result
+                if pred_semantic is None:
+                    print(t2s_result.exception)
+                    print(t2s_result.traceback)
+                    raise RuntimeError("")
+                cache[i_text] = pred_semantic
         t3 = ttime()
-        refers=[]
+        refers = []
         sv_emb = []
-        if(inp_refs):
+        if inp_refs:
             for path in inp_refs:
                 try:
                     refer, audio_tensor = get_spepc(hps, path.name, dtype, device, is_v2pro=True)
@@ -498,22 +575,28 @@ def get_tts_wav(ref_wav_path, prompt_text, prompt_language, text, text_language,
             refers, audio_tensor = get_spepc(hps, ref_wav_path, dtype, device, is_v2pro=True)
             refers = [refers]
             sv_emb = [sv_cn_model.compute_embedding3(audio_tensor)]
-        audio = vq_model.decode(
-            pred_semantic, torch.LongTensor(phones2).to(device).unsqueeze(0), refers, speed=speed, sv_emb=sv_emb
-        ).detach().cpu().numpy()[0][0]
-        max_audio=np.abs(audio).max()#简单防止16bit爆音
-        if max_audio>1:audio/=max_audio
+        audio = (
+            vq_model.decode(
+                pred_semantic[0].unsqueeze(0).unsqueeze(0),
+                torch.LongTensor(phones2).to(device).unsqueeze(0),
+                refers,
+                speed=speed,
+                sv_emb=sv_emb,
+            )
+            .detach()
+            .cpu()
+            .numpy()[0][0]
+        )
+        max_audio = np.abs(audio).max()  # 简单防止16bit爆音
+        if max_audio > 1:
+            audio /= max_audio
         audio_opt.append(audio)
         audio_opt.append(zero_wav)
         t4 = ttime()
-        t.extend([t2 - t1,t3 - t2, t4 - t3])
+        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]))
-           )
-    yield hps.data.sampling_rate, (np.concatenate(audio_opt, 0) * 32768).astype(
-        np.int16
-    )
+    print("%.3f\t%.3f\t%.3f\t%.3f" % (t[0], sum(t[1::3]), sum(t[2::3]), sum(t[3::3])))
+    yield hps.data.sampling_rate, (np.concatenate(audio_opt, 0) * 32768).astype(np.int16)
 
 
 def split(todo_text):
@@ -543,7 +626,7 @@ def cut1(inp):
     if len(split_idx) > 1:
         opts = []
         for idx in range(len(split_idx) - 1):
-            opts.append("".join(inps[split_idx[idx]: split_idx[idx + 1]]))
+            opts.append("".join(inps[split_idx[idx] : split_idx[idx + 1]]))
     else:
         opts = [inp]
     opts = [item for item in opts if not set(item).issubset(punctuation)]
@@ -579,7 +662,8 @@ def cut3(inp):
     inp = inp.strip("\n")
     opts = ["%s" % item for item in inp.strip("。").split("。")]
     opts = [item for item in opts if not set(item).issubset(punctuation)]
-    return  "\n".join(opts)
+    return "\n".join(opts)
+
 
 def cut4(inp):
     inp = inp.strip("\n")
@@ -591,13 +675,13 @@ def cut4(inp):
 # contributed by https://github.com/AI-Hobbyist/GPT-SoVITS/blob/main/GPT_SoVITS/inference_webui.py
 def cut5(inp):
     inp = inp.strip("\n")
-    punds = {',', '.', ';', '?', '!', '、', '，', '。', '？', '！', ';', '：', '…'}
+    punds = {",", ".", ";", "?", "!", "、", "，", "。", "？", "！", ";", "：", "…"}
     mergeitems = []
     items = []
 
     for i, char in enumerate(inp):
         if char in punds:
-            if char == '.' and i > 0 and i < len(inp) - 1 and inp[i - 1].isdigit() and inp[i + 1].isdigit():
+            if char == "." and i > 0 and i < len(inp) - 1 and inp[i - 1].isdigit() and inp[i + 1].isdigit():
                 items.append(char)
             else:
                 items.append(char)
@@ -615,35 +699,37 @@ def cut5(inp):
 
 def custom_sort_key(s):
     # 使用正则表达式提取字符串中的数字部分和非数字部分
-    parts = re.split('(\d+)', s)
+    parts = re.split("(\d+)", s)
     # 将数字部分转换为整数，非数字部分保持不变
     parts = [int(part) if part.isdigit() else part for part in parts]
     return parts
 
+
 def process_text(texts):
-    _text=[]
-    if all(text in [None, " ", "\n",""] for text in texts):
+    _text = []
+    if all(text in [None, " ", "\n", ""] for text in texts):
         raise ValueError(i18n("请输入有效文本"))
     for text in texts:
-        if text in  [None, " ", ""]:
+        if text in [None, " ", ""]:
             pass
         else:
             _text.append(text)
     return _text
 
 
-def html_center(text, label='p'):
+def html_center(text, label="p"):
     return f"""<div style="text-align: center; margin: 100; padding: 50;">
                 <{label} style="margin: 0; padding: 0;">{text}</{label}>
                 </div>"""
 
-def html_left(text, label='p'):
+
+def html_left(text, label="p"):
     return f"""<div style="text-align: left; margin: 0; padding: 0;">
                 <{label} style="margin: 0; padding: 0;">{text}</{label}>
                 </div>"""
 
 
-with gr.Blocks(title="GPT-SoVITS WebUI") as app:
+with gr.Blocks(title="GPT-SoVITS WebUI", analytics_enabled=False) as app:
     gr.Markdown(
         value="""# GPT-SoVITS-ProPlus Zero-shot TTS demo
 ## https://github.com/RVC-Boss/GPT-SoVITS
@@ -656,49 +742,95 @@ This demo is open source under the MIT license. The author does not have any con
 """
     )
     with gr.Group():
-        gr.Markdown(html_center(i18n("*请上传并填写参考信息"),'h3'))
+        gr.Markdown(html_center(i18n("*请上传并填写参考信息"), "h3"))
         with gr.Row():
             inp_ref = gr.Audio(label=i18n("请上传3~10秒内参考音频，超过会报错！"), type="filepath")
             with gr.Column():
-                ref_text_free = gr.Checkbox(label=i18n("开启无参考文本模式。不填参考文本亦相当于开启。"), value=False, interactive=True, show_label=True)
-                gr.Markdown(html_left(i18n("使用无参考文本模式时建议使用微调的GPT，听不清参考音频说的啥(不晓得写啥)可以开。<br>开启后无视填写的参考文本。")))
+                ref_text_free = gr.Checkbox(
+                    label=i18n("开启无参考文本模式。不填参考文本亦相当于开启。"),
+                    value=False,
+                    interactive=True,
+                    show_label=True,
+                )
+                gr.Markdown(
+                    html_left(
+                        i18n(
+                            "使用无参考文本模式时建议使用微调的GPT，听不清参考音频说的啥(不晓得写啥)可以开。<br>开启后无视填写的参考文本。"
+                        )
+                    )
+                )
                 prompt_text = gr.Textbox(label=i18n("参考音频的文本"), value="", lines=3, max_lines=3)
             prompt_language = gr.Dropdown(
                 label=i18n("参考音频的语种"), choices=list(dict_language.keys()), value=i18n("中文")
             )
-            inp_refs = gr.File(label=i18n("可选项：通过拖拽多个文件上传多个参考音频（建议同性），平均融合他们的音色。如不填写此项，音色由左侧单个参考音频控制。"),file_count="multiple")
-        gr.Markdown(html_center(i18n("*请填写需要合成的目标文本和语种模式"),'h3'))
+            inp_refs = gr.File(
+                label=i18n(
+                    "可选项：通过拖拽多个文件上传多个参考音频（建议同性），平均融合他们的音色。如不填写此项，音色由左侧单个参考音频控制。"
+                ),
+                file_count="multiple",
+            )
+        gr.Markdown(html_center(i18n("*请填写需要合成的目标文本和语种模式"), "h3"))
         with gr.Row():
             with gr.Column():
                 text = gr.Textbox(label=i18n("需要合成的文本"), value="", lines=26, max_lines=26)
             with gr.Column():
                 text_language = gr.Dropdown(
-                        label=i18n("需要合成的语种")+i18n(".限制范围越小判别效果越好。"), choices=list(dict_language.keys()), value=i18n("中文")
-                    )
+                    label=i18n("需要合成的语种") + i18n(".限制范围越小判别效果越好。"),
+                    choices=list(dict_language.keys()),
+                    value=i18n("中文"),
+                )
                 how_to_cut = gr.Dropdown(
-                        label=i18n("怎么切"),
-                        choices=[i18n("不切"), i18n("凑四句一切"), i18n("凑50字一切"), i18n("按中文句号。切"), i18n("按英文句号.切"), i18n("按标点符号切"), ],
-                        value=i18n("凑四句一切"),
-                        interactive=True
-                    )
+                    label=i18n("怎么切"),
+                    choices=[
+                        i18n("不切"),
+                        i18n("凑四句一切"),
+                        i18n("凑50字一切"),
+                        i18n("按中文句号。切"),
+                        i18n("按英文句号.切"),
+                        i18n("按标点符号切"),
+                    ],
+                    value=i18n("凑四句一切"),
+                    interactive=True,
+                )
                 gr.Markdown(value=html_center(i18n("语速调整，高为更快")))
-                if_freeze=gr.Checkbox(label=i18n("是否直接对上次合成结果调整语速和音色。防止随机性。"), value=False, interactive=True,show_label=True)
-                speed = gr.Slider(minimum=0.6,maximum=1.65,step=0.05,label=i18n("语速"),value=1,interactive=True)
+                if_freeze = gr.Checkbox(
+                    label=i18n("是否直接对上次合成结果调整语速和音色。防止随机性。"),
+                    value=False,
+                    interactive=True,
+                    show_label=True,
+                )
+                speed = gr.Slider(minimum=0.6, maximum=1.65, step=0.05, label=i18n("语速"), value=1, interactive=True)
                 gr.Markdown(html_center(i18n("GPT采样参数(无参考文本时不要太低。不懂就用默认)：")))
-                top_k = gr.Slider(minimum=1,maximum=100,step=1,label=i18n("top_k"),value=15,interactive=True)
-                top_p = gr.Slider(minimum=0,maximum=1,step=0.05,label=i18n("top_p"),value=1,interactive=True)
-                temperature = gr.Slider(minimum=0,maximum=1,step=0.05,label=i18n("temperature"),value=1,interactive=True)
+                top_k = gr.Slider(minimum=1, maximum=100, step=1, label=i18n("top_k"), value=15, interactive=True)
+                top_p = gr.Slider(minimum=0, maximum=1, step=0.05, label=i18n("top_p"), value=1, interactive=True)
+                temperature = gr.Slider(
+                    minimum=0, maximum=1, step=0.05, label=i18n("temperature"), value=1, interactive=True
+                )
         with gr.Row():
-            inference_button = gr.Button(i18n("合成语音"), variant="primary", size='lg')
+            inference_button = gr.Button(i18n("合成语音"), variant="primary", size="lg")
             output = gr.Audio(label=i18n("输出的语音"))
 
         inference_button.click(
             get_tts_wav,
-            [inp_ref, prompt_text, prompt_language, text, text_language, how_to_cut, top_k, top_p, temperature, ref_text_free,speed,if_freeze,inp_refs],
+            [
+                inp_ref,
+                prompt_text,
+                prompt_language,
+                text,
+                text_language,
+                how_to_cut,
+                top_k,
+                top_p,
+                temperature,
+                ref_text_free,
+                speed,
+                if_freeze,
+                inp_refs,
+            ],
             [output],
         )
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     app.queue().launch(
         server_name="0.0.0.0",
         inbrowser=True,

pre-requirements.txt

@@ -0,0 +1 @@
+torch==2.5.1

requirements.txt

@@ -30,10 +30,10 @@ g2pk2==0.0.3
 ko_pron==1.3
 opencc==1.1.0
 python_mecab_ko==1.3.7
-torch==2.5.1
 pydantic==2.8.2
 torchmetrics<=1.5
 nltk==3.8.1
 fast_langdetect==0.3.1
 split_lang==2.1.0
-ToJyutping==3.2.0
+ToJyutping==3.2.0
+https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.5cxx11abiTRUE-cp310-cp310-linux_x86_64.whl