| """ |
| Modified From https://github.com/XXXXRT666/GPT-SoVITS |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Literal, MutableSequence, Optional, Protocol |
|
|
| import torch |
|
|
| from .sample_funcs import SampleProtocol, sample_naive |
|
|
| Tensor = torch.Tensor |
|
|
|
|
| @dataclass |
| class T2SResult: |
| result: list[Tensor] | None = None |
| infer_speed: tuple[float, float] = (0.0, 0.0) |
| 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 KVCacheProtocol(Protocol): |
| k_cache: Tensor |
| v_cache: Tensor |
|
|
| def __init__(self, batch_size: int, max_seq_length: int, n_heads: int, head_dim: int) -> None: ... |
|
|
| def empty(self) -> None: ... |
|
|
| def update(self, input_pos: Tensor, k_val: Tensor, v_val: Tensor, *args, **kwds) -> tuple[Tensor, Tensor]: ... |
|
|
| def prefill_kv(self, k_val: Tensor, v_val: Tensor) -> None: ... |
|
|
| def sync_cache(self, kv_cache: KVCacheProtocol) -> None: ... |
|
|
|
|
| class T2SDecoderProtocol(Protocol): |
| max_seq_length: int |
| EOS: int |
| n_head: int |
|
|
| @property |
| def device(self) -> torch.device: ... |
|
|
| def embed(self, x: list[Tensor], y: Tensor, bert_features: list[Tensor]) -> Tensor: ... |
|
|
|
|
| class T2SEngineProtocol(Protocol): |
| def _handle_request(self, request: T2SRequest) -> tuple[list[Tensor], float, float]: ... |
|
|
| def generate(self, request: T2SRequest) -> T2SResult: ... |
|
|
|
|
| class T2SSession: |
| def __init__( |
| self, |
| decoder: T2SDecoderProtocol, |
| request: T2SRequest, |
| sapmle_func: type[SampleProtocol] = sample_naive, |
| device: torch.device = torch.device("cpu"), |
| dtype: torch.dtype = torch.float32, |
| ): |
| 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 |
| request.prompts = request.prompts.to(device, torch.int32) |
|
|
| |
| self.kv_cache: MutableSequence[KVCacheProtocol] |
| self.sample = sapmle_func() |
|
|
| |
| self.x = [i.to(device) for i in request.x] |
| self.x_lens = request.x_lens.to(torch.int32) |
| self.y = torch.zeros((bsz, decoder.max_seq_length)).to(torch.int32) |
| self.y[:, : request.prompts.shape[-1]] = request.prompts |
| self.bert_feature = [i.to(device, dtype) for i in request.bert_feature] |
|
|
| self.prefill_len = self.x_lens + request.prompts.size(1) |
|
|
| self.input_pos = torch.zeros_like(self.prefill_len) |
| self.input_pos.add_(self.prefill_len) |
|
|
| |
| self.stream: Optional[torch.cuda.Stream] = None |
| self.graph: Optional[torch.cuda.CUDAGraph] = None |
| self.xy_pos_: Tensor |
| self.xy_dec_: Tensor |
|
|
| |
| self.completed = torch.Tensor([False] * len(self.x)).bool().to(device) |
| self.y_results: list[Tensor] = [None] * len(self.x) |
|
|
| self.xy_pos = decoder.embed(self.x, request.prompts, self.bert_feature) |
|
|
| max_len = int(self.prefill_len.max().item()) |
| attn_mask = torch.zeros(size=(bsz, max_len, max_len), dtype=torch.bool) |
|
|
| for bs in range(bsz): |
| pos = int(self.x_lens[bs]) |
| seq_len = pos + y_len |
|
|
| attn_mask[bs, :seq_len, :pos] = True |
|
|
| ar_mask = ~torch.triu( |
| input=torch.ones( |
| size=( |
| y_len, |
| y_len, |
| ), |
| dtype=torch.bool, |
| ), |
| diagonal=1, |
| ) |
| attn_mask[bs, pos:seq_len, pos:seq_len] = ar_mask |
|
|
| self.attn_mask = attn_mask |
| self.attn_mask = attn_mask.unsqueeze(0).expand(-1, decoder.n_head, -1, -1) |
|
|
| self.id: int = -1 |
|
|
| |
| self.cu_seqlens_q: Tensor |
| self.cu_seqlens_kv: Tensor |
|
|
