indextts25-inference 0.1.1__py3-none-any.whl

indextts25/__init__.py

@@ -0,0 +1,5 @@
+"""IndexTTS-2.5: zero-shot speech synthesis with emotion and speed control."""
+from .api import STREAM_SEGMENT_TOKENS, Chunk, IndexTTS25
+
+__all__ = ["IndexTTS25", "Chunk", "STREAM_SEGMENT_TOKENS"]
+__version__ = "0.1.1"  # x-release-please-version

indextts25/accel/__init__.py

@@ -0,0 +1,2 @@
+from .accel_engine import AccelInferenceEngine  # noqa: F401
+from .gpt2_accel import GPT2AccelModel  # noqa: F401

indextts25/accel/accel_engine.py

@@ -0,0 +1,205 @@
+"""Accelerated decoding engine for the IndexTTS GPT (attn_backend="flash").
+
+One sequence at a time (batch 1): paged fp16 KV cache, prefill with
+flash_attn_varlen_func and every decode step captured in a CUDA graph. Samples like
+transformers with num_beams=1 (repetition_penalty, temperature, top-k, top-p); no
+beam search.
+"""
+import torch
+from torch import nn
+
+from .attention import reset_forward_context, set_forward_context
+from .kv_manager import KVCacheManager, Seq
+
+
+class AccelInferenceEngine:
+    def __init__(
+        self,
+        model,
+        lm_head,
+        num_layers: int,
+        num_heads: int,
+        head_dim: int,
+        block_size: int = 256,
+        num_blocks: int = 16,
+        use_cuda_graph: bool = True,
+    ):
+        """
+        Args:
+            model: the GPT2 with accel attention (GPT2AccelModel).
+            lm_head: final_norm + mel_head.
+            block_size / num_blocks: geometry of the paged KV pool.
+            use_cuda_graph: capture the decode step in a CUDA graph.
+        """
+        self.model = model
+        self.lm_head = lm_head
+        self.block_size = block_size
+        self.use_cuda_graph = use_cuda_graph and torch.cuda.is_available()
+        self.hidden_size = model.config.hidden_size
+        self.kv_manager = KVCacheManager(num_layers, num_heads, head_dim, block_size, num_blocks)
+        self.kv_manager.wire_kv_cache_to_model(model)
+        self.graph = None
+        self.graph_vars = None
+
+    def _capture_cuda_graph(self, mel_embedding, mel_pos_embedding):
+        # The block-table width is baked into the graph, and flash_attn_with_kvcache uses
+        # it to decide how to split the work: changing it changes the output in the last
+        # bit. It covers 2048 tokens; a longer sequence falls back to the non-graph path.
+        max_num_blocks = (2048 + self.block_size - 1) // self.block_size
+        dtype = next(self.model.parameters()).dtype
+        gv = {
+            "input_ids": torch.ones(1, dtype=torch.int64, device="cuda"),
+            "positions": torch.ones(1, dtype=torch.int64, device="cuda"),
+            "slot_mapping": torch.zeros(1, dtype=torch.int32, device="cuda"),
+            "context_lens": torch.full((1,), 2, dtype=torch.int32, device="cuda"),
+            "block_tables": torch.zeros(1, max_num_blocks, dtype=torch.int32, device="cuda"),
+            "inputs_embeds": torch.zeros(1, self.hidden_size, dtype=dtype, device="cuda"),
+            "outputs": torch.zeros(1, self.hidden_size, dtype=dtype, device="cuda"),
+        }
+
+        def step():
+            emb = mel_embedding(gv["input_ids"])
+            pos_emb = mel_pos_embedding.emb(torch.clamp(gv["positions"], min=0))
+            gv["inputs_embeds"][:] = emb + pos_emb
+            out = self.model(inputs_embeds=gv["inputs_embeds"].unsqueeze(1))
+            gv["outputs"][:] = out.squeeze(1)
+
+        set_forward_context(False, slot_mapping=gv["slot_mapping"], context_lens=gv["context_lens"],
+                            block_tables=gv["block_tables"])
+        step()  # warmup
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            step()
+        torch.cuda.synchronize()
+        reset_forward_context()
+        self.graph, self.graph_vars = graph, gv
+
+    def _decode(self, seq: Seq, mel_embedding, mel_pos_embedding):
+        """One decode step for the last token of `seq`. Returns the hidden state [1, C]."""
+        pos = len(seq) - seq.num_prompt_tokens  # mel position of the code
+        slot = seq.block_table[-1] * self.block_size + seq.last_block_num_tokens - 1
+        gv = self.graph_vars
+        if gv is not None and len(seq.block_table) <= gv["block_tables"].size(1):
+            gv["input_ids"].fill_(seq.last_token)
+            gv["positions"].fill_(pos)
+            gv["slot_mapping"].fill_(slot)
+            gv["context_lens"].fill_(len(seq))
+            if len(seq.block_table) > self._graph_blocks:  # a new block was opened
+                self._graph_blocks = len(seq.block_table)
+                gv["block_tables"][0, self._graph_blocks - 1] = seq.block_table[-1]
+            self.graph.replay()
+            return gv["outputs"]
+
+        set_forward_context(
+            False,
+            slot_mapping=torch.tensor([slot], dtype=torch.int32, device="cuda"),
+            context_lens=torch.tensor([len(seq)], dtype=torch.int32, device="cuda"),
+            block_tables=torch.tensor([seq.block_table], dtype=torch.int32, device="cuda"),
+        )
+        emb = mel_embedding(torch.tensor([seq.last_token], device="cuda"))
+        emb = emb + mel_pos_embedding.emb(torch.tensor([pos], device="cuda"))
+        out = self.model(inputs_embeds=emb.to(next(self.model.parameters()).dtype).unsqueeze(1))
+        reset_forward_context()
+        return out.squeeze(1)
+
+    def _sample(self, logits, temperature, top_k, top_p, rp, seen):
+        """One token, in transformers' order: repetition penalty on the raw logits,
+        temperature, top-k, top-p and sampling (exponential race)."""
+        if rp != 1.0:
+            logits = torch.where(seen, torch.where(logits < 0, logits * rp, logits / rp), logits)
+        if temperature <= 0:
+            tok = torch.argmax(logits, dim=-1).tolist()[0]
+        else:
+            logits = logits.float() / temperature
+            if top_k > 0:
+                kth = torch.topk(logits, min(top_k, logits.size(-1))).values[..., -1:]
+                logits = logits.masked_fill(logits < kth, float("-inf"))
+            if top_p < 1.0:
+                sl, idx = torch.sort(logits, descending=False)
+                remove = sl.softmax(-1).cumsum(-1) <= (1 - top_p)
+                remove[..., -1:] = False
+                logits = logits.masked_fill(remove.scatter(1, idx, remove), float("-inf"))
+            probs = logits.softmax(-1)
+            tok = (probs / torch.empty_like(probs).exponential_(1).clamp_min_(1e-10)).argmax(-1).tolist()[0]
+        seen[0, tok] = True
+        return tok
+
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        max_new_tokens: int,
+        temperature: float,
+        stop_tokens,
+        tts_embeddings: torch.Tensor,
+        tts_mel_embedding: nn.Module,
+        tts_text_pos_embedding: nn.Module,
+        top_k: int = 0,
+        top_p: float = 1.0,
+        repetition_penalty: float = 1.0,
+    ) -> torch.Tensor:
+        """Generate codes for one sequence.
+
+        Args:
+            input_ids: [1, n] filler ids ([1]*(n-1) + [start_mel]); the real prompt is
+                `tts_embeddings` ([1, n-1, C]: [pad][cond][text]) plus the start_mel.
+            max_new_tokens: cap on the number of generated codes.
+            stop_tokens: codes that end the sequence (included in the output).
+            tts_mel_embedding / tts_text_pos_embedding: code and position embeddings.
+
+        Returns:
+            [1, n + generated], like transformers' generate().
+        """
+        if self.use_cuda_graph and self.graph is None:
+            self._capture_cuda_graph(tts_mel_embedding, tts_text_pos_embedding)
+
+        n = input_ids.size(1)
+        vocab = tts_mel_embedding.num_embeddings
+        seen = torch.zeros(1, vocab, dtype=torch.bool, device=input_ids.device)
+        seen[0, input_ids[0]] = True  # like transformers: the prompt's filler ids count too
+        knobs = (top_k, top_p, repetition_penalty, seen)
+        seq = Seq(n, last_token=-1, block_size=self.block_size)
+        self.kv_manager.allocate(seq)
+        if self.graph_vars is not None:
+            self.graph_vars["block_tables"].fill_(-1)
+            first = seq.block_table[: self.graph_vars["block_tables"].size(1)]
+            self.graph_vars["block_tables"][0, : len(first)] = torch.tensor(first, dtype=torch.int32)
+            self._graph_blocks = len(first)
+        try:
+            # Prefill: [pad][cond][text] + start_mel at mel position 0, as in training
+            # (it used to go at n-1, the position of the last prompt token).
+            start_emb = tts_mel_embedding(input_ids[:, -1:])
+            start_pos = torch.zeros((1, 1), device="cuda", dtype=torch.long)
+            start_emb = start_emb + tts_text_pos_embedding.emb(start_pos)
+            full_embeddings = torch.cat([tts_embeddings, start_emb], dim=1)
+            model_dtype = next(self.model.parameters()).dtype
+            if full_embeddings.dtype != model_dtype:
+                full_embeddings = full_embeddings.to(model_dtype)
+
+            slots = [seq.block_table[t // self.block_size] * self.block_size + t % self.block_size
+                     for t in range(n)]
+            cu_seqlens = torch.tensor([0, n], dtype=torch.int32, device="cuda")
+            set_forward_context(True, cu_seqlens, cu_seqlens, n, n,
+                                torch.tensor(slots, dtype=torch.int32, device="cuda"))
+            hidden = self.model(inputs_embeds=full_embeddings)[:, -1, :]
+            reset_forward_context()
+            lm_dtype = next(self.lm_head.parameters()).dtype
+            if hidden.dtype != lm_dtype:
+                hidden = hidden.to(lm_dtype)
+
+            generated = []
+            token = self._sample(self.lm_head(hidden), temperature, *knobs)
+            for step in range(max_new_tokens):
+                generated.append(token)
+                # The stop token goes in the output, as in transformers: without it,
+                # infer_v2_5 warned about a max_mel_tokens cutoff on every call.
+                if token in stop_tokens or step == max_new_tokens - 1:
+                    break
+                seq.append_token(token)
+                self.kv_manager.append_to_seq(seq)
+                hidden = self._decode(seq, tts_mel_embedding, tts_text_pos_embedding)
+                token = self._sample(self.lm_head(hidden), temperature, *knobs)
+        finally:
+            self.kv_manager.free(seq)
+            reset_forward_context()
+
+        return torch.cat([input_ids, torch.tensor([generated], dtype=input_ids.dtype, device=input_ids.device)], dim=1)

indextts25/accel/attention.py

@@ -0,0 +1,148 @@
+from dataclasses import dataclass
+
+import torch
+import triton
+import triton.language as tl
+from flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache
+from torch import nn
+
+
+@dataclass
+class ForwardContext:
+    is_prefill: bool = False
+    cu_seqlens_q: torch.Tensor | None = None
+    cu_seqlens_k: torch.Tensor | None = None
+    max_seqlen_q: int = 0
+    max_seqlen_k: int = 0
+    slot_mapping: torch.Tensor | None = None
+    context_lens: torch.Tensor | None = None
+    block_tables: torch.Tensor | None = None
+
+
+_FORWARD_CONTEXT = ForwardContext()
+
+
+def get_forward_context():
+    return _FORWARD_CONTEXT
+
+
+def set_forward_context(
+    is_prefill,
+    cu_seqlens_q=None,
+    cu_seqlens_k=None,
+    max_seqlen_q=0,
+    max_seqlen_k=0,
+    slot_mapping=None,
+    context_lens=None,
+    block_tables=None,
+):
+    global _FORWARD_CONTEXT
+    _FORWARD_CONTEXT = ForwardContext(
+        is_prefill,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        slot_mapping,
+        context_lens,
+        block_tables,
+    )
+
+
+def reset_forward_context():
+    global _FORWARD_CONTEXT
+    _FORWARD_CONTEXT = ForwardContext()
+
+
+@triton.jit
+def store_kvcache_kernel(
+    key_ptr,
+    key_stride,
+    value_ptr,
+    value_stride,
+    k_cache_ptr,
+    v_cache_ptr,
+    slot_mapping_ptr,
+    D: tl.constexpr,
+):
+    BLOCK_SIZE: tl.constexpr = 2048
+    idx = tl.program_id(0)
+    slot = tl.load(slot_mapping_ptr + idx)
+    if slot == -1:
+        return
+    d_offset = 0
+    while d_offset < D:
+        cur_block_size = min(BLOCK_SIZE, D - d_offset)
+        key_offsets = idx * key_stride + d_offset + tl.arange(0, BLOCK_SIZE)
+        value_offsets = idx * value_stride + d_offset + tl.arange(0, BLOCK_SIZE)
+        cache_offsets = slot * D + d_offset + tl.arange(0, BLOCK_SIZE)
+
+        mask = tl.arange(0, BLOCK_SIZE) < cur_block_size
+        key = tl.load(key_ptr + key_offsets, mask=mask, other=0.0)
+        value = tl.load(value_ptr + value_offsets, mask=mask, other=0.0)
+        tl.store(k_cache_ptr + cache_offsets, key, mask=mask)
+        tl.store(v_cache_ptr + cache_offsets, value, mask=mask)
+
+        d_offset += BLOCK_SIZE
+
+
+def store_kvcache(
+    key: torch.Tensor,
+    value: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+):
+    N, num_heads, head_dim = key.shape
+    D = num_heads * head_dim
+    assert key.stride(-1) == 1 and value.stride(-1) == 1
+    assert key.stride(1) == head_dim and value.stride(1) == head_dim
+    assert k_cache.stride(1) == D and v_cache.stride(1) == D
+    assert slot_mapping.numel() == N
+    store_kvcache_kernel[(N,)](
+        key, key.stride(0), value, value.stride(0), k_cache, v_cache, slot_mapping, D
+    )
+
+
+class Attention(nn.Module):
+    """Attention over the paged KV cache: writes k/v into their slots and uses flash-attn
+    (varlen for prefill, with_kvcache for decode). The engine sets the context."""
+
+    def __init__(self, scale: float):
+        super().__init__()
+        self.scale = scale
+        self.k_cache = self.v_cache = torch.tensor([])
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
+        context = get_forward_context()
+        k_cache, v_cache = self.k_cache, self.v_cache
+
+        if k_cache.numel() and v_cache.numel() and context.slot_mapping is not None:
+            store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)
+
+        if context.is_prefill:
+            if context.block_tables is not None:
+                k, v = k_cache, v_cache
+            o = flash_attn_varlen_func(
+                q,
+                k,
+                v,
+                max_seqlen_q=context.max_seqlen_q,
+                cu_seqlens_q=context.cu_seqlens_q,
+                max_seqlen_k=context.max_seqlen_k,
+                cu_seqlens_k=context.cu_seqlens_k,
+                softmax_scale=self.scale,
+                causal=True,
+                block_table=context.block_tables,
+            )
+        else:
+            o = flash_attn_with_kvcache(
+                q.unsqueeze(1),
+                k_cache,
+                v_cache,
+                cache_seqlens=context.context_lens,
+                block_table=context.block_tables,
+                softmax_scale=self.scale,
+                causal=True,
+            )
+        return o

indextts25/accel/gpt2_accel.py

@@ -0,0 +1,131 @@
+import torch
+import torch.nn as nn
+from transformers.models.gpt2.modeling_gpt2 import Conv1D, GPT2Block, GPT2Model
+
+from .attention import Attention
+
+
+class GPT2AccelAttention(nn.Module):
+    def __init__(self, config, layer_idx=None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        self.split_size = self.embed_dim
+
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"`embed_dim` must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+                f" {self.num_heads})."
+            )
+
+        self.scale_attn_weights = config.scale_attn_weights
+
+        self.c_attn = Conv1D(3 * self.embed_dim, self.embed_dim)
+        self.c_proj = Conv1D(self.embed_dim, self.embed_dim)
+
+        self.attn_dropout = nn.Dropout(config.attn_pdrop)
+        self.resid_dropout = nn.Dropout(config.resid_pdrop)
+
+        scale = (self.head_dim**-0.5) if self.scale_attn_weights else 1.0
+        self.accel_attn = Attention(scale)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        layer_past=None,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        use_cache=False,
+        output_attentions=False,
+        past_key_value=None,
+        **kwargs,
+    ):
+        if encoder_hidden_states is not None:
+            raise NotImplementedError("Cross attention not supported in accel mode")
+
+        qkv = self.c_attn(hidden_states)
+        query, key, value = qkv.split(self.split_size, dim=-1)
+
+        # [B, T, H*D] -> [B, H, T, D]
+        query = self._split_heads(query, self.num_heads, self.head_dim)
+        key = self._split_heads(key, self.num_heads, self.head_dim)
+        value = self._split_heads(value, self.num_heads, self.head_dim)
+
+        # flatten to [B*T, H, D]
+        bsz, num_heads, seq_len, head_dim = query.shape
+        q_flat = query.transpose(1, 2).contiguous().view(-1, num_heads, head_dim)
+        k_flat = key.transpose(1, 2).contiguous().view(-1, num_heads, head_dim)
+        v_flat = value.transpose(1, 2).contiguous().view(-1, num_heads, head_dim)
+
+        # ensure fp16
+        if q_flat.device.type == "cuda" and q_flat.dtype != torch.float16:
+            orig_dtype = q_flat.dtype
+            q_flat = q_flat.to(torch.float16)
+            k_flat = k_flat.to(torch.float16)
+            v_flat = v_flat.to(torch.float16)
+        else:
+            orig_dtype = q_flat.dtype
+
+        o_flat = self.accel_attn(q_flat, k_flat, v_flat)  # [B*T, H, D]
+
+        if o_flat.dtype != orig_dtype:
+            o_flat = o_flat.to(orig_dtype)
+
+        # Reshape back: [B*T, H, D] -> [B, H, T, D]
+        attn_output = o_flat.view(bsz, seq_len, num_heads, head_dim).transpose(1, 2)
+
+        attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
+
+        attn_output = self.c_proj(attn_output)
+        attn_output = self.resid_dropout(attn_output)
+
+        outputs = (attn_output, None)
+        if output_attentions:
+            outputs += (None,)
+
+        return outputs
+
+    def _split_heads(self, tensor, num_heads, head_dim):
+        new_shape = tensor.size()[:-1] + (num_heads, head_dim)
+        tensor = tensor.view(new_shape)
+        return tensor.permute(0, 2, 1, 3)  # (batch, head, seq_length, head_features)
+
+    def _merge_heads(self, tensor, num_heads, head_dim):
+        tensor = tensor.permute(0, 2, 1, 3).contiguous()
+        new_shape = tensor.size()[:-2] + (num_heads * head_dim,)
+        return tensor.view(new_shape)
+
+
+class GPT2AccelBlock(GPT2Block):
+    def __init__(self, config, layer_idx=None):
+        super().__init__(config, layer_idx)
+        self.attn = GPT2AccelAttention(config, layer_idx)
+
+    def forward(self, hidden_states):
+        # Custom forward: GPT2Block's forward changed its contract across transformers
+        # versions (through 5.2 it returns a tuple; from 5.3 on, the tensor).
+        hidden_states = hidden_states + self.attn(self.ln_1(hidden_states))[0]
+        return hidden_states + self.mlp(self.ln_2(hidden_states))
+
+
+class GPT2AccelModel(GPT2Model):
+    """GPT2 with the accel attention. Only used with inputs_embeds (the codes are
+    embedded outside) and returns the final hidden state, with ln_f already applied."""
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.h = nn.ModuleList(
+            [GPT2AccelBlock(config, layer_idx=i) for i in range(config.num_hidden_layers)]
+        )
+
+    def forward(self, inputs_embeds):
+        hidden_states = inputs_embeds
+        for block in self.h:
+            hidden_states = block(hidden_states)
+        return self.ln_f(hidden_states)

indextts25/accel/kv_manager.py

@@ -0,0 +1,62 @@
+from collections import deque
+
+import torch
+
+
+class Seq:
+    """A sequence in the paged pool: how many tokens it has and which blocks hold them."""
+
+    def __init__(self, num_prompt_tokens: int, last_token: int, block_size: int = 256):
+        self.num_tokens = num_prompt_tokens
+        self.num_prompt_tokens = num_prompt_tokens
+        self.last_token = last_token
+        self.block_size = block_size
+        self.block_table: list[int] = []
+
+    def __len__(self):
+        return self.num_tokens
+
+    @property
+    def num_blocks(self):
+        return (self.num_tokens + self.block_size - 1) // self.block_size
+
+    @property
+    def last_block_num_tokens(self):
+        return self.num_tokens - (self.num_blocks - 1) * self.block_size
+
+    def append_token(self, token_id: int):
+        self.last_token = token_id
+        self.num_tokens += 1
+
+
+class KVCacheManager:
+    """fp16 KV pool, [2, layers, blocks, block_size, heads, head_dim], handed out
+    block by block. No prefix caching: in TTS the real prompt is embeddings and the ids
+    are filler, so two long prompts "matched" and overwrote each other."""
+
+    def __init__(self, num_layers: int, num_heads: int, head_dim: int, block_size: int, num_blocks: int):
+        self.block_size = block_size
+        self.free_block_ids: deque = deque(range(num_blocks))
+        self.kv_cache = torch.empty(
+            2, num_layers, num_blocks, block_size, num_heads, head_dim, dtype=torch.float16, device="cuda")
+
+    def allocate(self, seq: Seq):
+        for _ in range(seq.num_blocks):
+            seq.block_table.append(self.free_block_ids.popleft())
+
+    def append_to_seq(self, seq: Seq):
+        # The just-appended token opens a new block.
+        if len(seq) % self.block_size == 1:
+            seq.block_table.append(self.free_block_ids.popleft())
+
+    def free(self, seq: Seq):
+        self.free_block_ids.extend(reversed(seq.block_table))
+        seq.block_table.clear()
+
+    def wire_kv_cache_to_model(self, model):
+        layer_id = 0
+        for module in model.modules():
+            if hasattr(module, "k_cache") and hasattr(module, "v_cache"):
+                module.k_cache = self.kv_cache[0, layer_id]
+                module.v_cache = self.kv_cache[1, layer_id]
+                layer_id += 1