feat:添加generate_long_parallel_batched

fish_speech/models/text2semantic/inference.py

@@ -4,14 +4,12 @@ import re
 import threading
 import time
 import traceback
-from copy import deepcopy
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, Literal, Optional, Tuple, Union
 
 import click
 import numpy as np
-import torch
 import torch._inductor.config
 from loguru import logger
 from tqdm import tqdm
@@ -30,7 +28,6 @@ torch._inductor.config.triton.unique_kernel_names = True
 if hasattr(torch._inductor.config, "fx_graph_cache"):
     torch._inductor.config.fx_graph_cache = True
 
-
 from torch.nn.attention import SDPBackend, sdpa_kernel
 
 from fish_speech.models.text2semantic.llama import (
@@ -50,10 +47,10 @@ RAS_HIGH_TOP_P = 0.9
 
 
 def logits_to_probs(
-    logits,
-    temperature: torch.Tensor,
-    top_p: torch.Tensor,
-    top_k: int,  # 注意: 我看到你传进来的是 int，这很关键
+        logits,
+        temperature: torch.Tensor,
+        top_p: torch.Tensor,
+        top_k: int,  # 注意: 我看到你传进来的是 int，这很关键
 ) -> torch.Tensor:
     sorted_logits, sorted_indices = torch.sort(logits, descending=True)
     cum_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
@@ -76,10 +73,10 @@ def logits_to_probs(
 
 
 def sample(
-    logits,
-    temperature: torch.Tensor,
-    top_p: torch.Tensor,
-    top_k: int,
+        logits,
+        temperature: torch.Tensor,
+        top_p: torch.Tensor,
+        top_k: int,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     probs = logits_to_probs(
         logits=logits[0, -1],
@@ -92,16 +89,16 @@ def sample(
 
 
 def decode_one_token_ar(
-    model: DualARTransformer,
-    x: torch.Tensor,
-    input_pos: torch.Tensor,
-    temperature: torch.Tensor,
-    top_p: torch.Tensor,
-    top_k: int,
-    semantic_logit_bias: torch.Tensor,
-    audio_masks: torch.Tensor,
-    audio_parts: torch.Tensor,
-    previous_tokens: Optional[torch.Tensor] = None,
+        model: DualARTransformer,
+        x: torch.Tensor,
+        input_pos: torch.Tensor,
+        temperature: torch.Tensor,
+        top_p: torch.Tensor,
+        top_k: int,
+        semantic_logit_bias: torch.Tensor,
+        audio_masks: torch.Tensor,
+        audio_parts: torch.Tensor,
+        previous_tokens: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     forward_result = model.forward_generate(
         x,
@@ -134,7 +131,7 @@ def decode_one_token_ar(
         in_window = (previous_tokens[0] == main_token_normal).any()
         # Use tensor ops (&, torch.where) instead of Python (and, if) — torch.compile requires no data-dependent branching
         is_semantic = (main_token_normal >= model.config.semantic_begin_id) & (
-            main_token_normal <= model.config.semantic_end_id
+                main_token_normal <= model.config.semantic_end_id
         )
         should_use_high = in_window & is_semantic
         main_token_normal = torch.where(
@@ -180,17 +177,17 @@ def decode_one_token_ar(
 
 
 def decode_n_tokens(
-    model: DualARTransformer,
-    cur_token: torch.Tensor,
-    input_pos: torch.Tensor,
-    num_new_tokens: int,
-    temperature: torch.Tensor,
-    top_p: torch.Tensor,
-    top_k: int,
-    semantic_logit_bias: torch.Tensor,
-    audio_masks: torch.Tensor,
-    audio_parts: torch.Tensor,
-    decode_one_token=decode_one_token_ar,
+        model: DualARTransformer,
+        cur_token: torch.Tensor,
+        input_pos: torch.Tensor,
+        num_new_tokens: int,
+        temperature: torch.Tensor,
+        top_p: torch.Tensor,
+        top_k: int,
+        semantic_logit_bias: torch.Tensor,
+        audio_masks: torch.Tensor,
+        audio_parts: torch.Tensor,
+        decode_one_token=decode_one_token_ar,
 ):
     # Rolling window for RAS (Repetition Aware Sampling)
     previous_tokens = torch.zeros(
@@ -239,15 +236,15 @@ def decode_n_tokens(
 @torch.no_grad()
 @torch.inference_mode()
 def generate(
-    *,
-    model: DualARTransformer,
-    prompt: torch.Tensor,
-    max_new_tokens: int,
-    audio_masks: torch.Tensor,
-    audio_parts: torch.Tensor,
-    decode_one_token=decode_one_token_ar,
-    num_samples: int = 1,
-    **sampling_kwargs,
+        *,
+        model: DualARTransformer,
+        prompt: torch.Tensor,
+        max_new_tokens: int,
+        audio_masks: torch.Tensor,
+        audio_parts: torch.Tensor,
+        decode_one_token=decode_one_token_ar,
+        num_samples: int = 1,
+        **sampling_kwargs,
 ):
     """
     Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.
@@ -311,7 +308,7 @@ def generate(
 
     # [MODIFIED] Use config for semantic range
     semantic_logit_bias[
-        0, 0, model.config.semantic_begin_id : model.config.semantic_end_id + 1
+        0, 0, model.config.semantic_begin_id: model.config.semantic_end_id + 1
     ] = 0.0
 
     # [MODIFIED] Use tokenizer.get_token_id (Wrapper method)
@@ -330,7 +327,7 @@ def generate(
         audio_masks,
         audio_parts,
     )
-    seq[:, T : T + 1] = first_token
+    seq[:, T: T + 1] = first_token
 
     # Recreate input_pos
     input_pos = torch.tensor([T], device=device, dtype=torch.int)
@@ -349,7 +346,7 @@ def generate(
         decode_one_token=decode_one_token,
     )
     seq = seq[:, : T + 1 + x.size(1)]
-    seq[:, T + 1 :] = x
+    seq[:, T + 1:] = x
 
     # Clean up temporary variables
     del first_token, x, prompt, empty, input_pos
@@ -483,7 +480,7 @@ def split_text_by_speaker(text: str) -> list[str]:
 
 
 def group_turns_into_batches(
-    turns: list[str], max_speakers: int = 3, max_bytes: int = 300
+        turns: list[str], max_speakers: int = 3, max_bytes: int = 300
 ) -> list[str]:
     """
     Group turns into batches based on speaker count or byte limit.
@@ -521,22 +518,22 @@ def group_turns_into_batches(
 
 
 def generate_long(
-    *,
-    model,
-    device: Union[str, torch.device],
-    decode_one_token: Callable,
-    text: str,
-    num_samples: int = 1,
-    max_new_tokens: int = 0,
-    top_p: float = 0.9,
-    top_k: int = 30,
-    repetition_penalty: float = 1.1,
-    temperature: float = 1.0,
-    compile: bool = False,
-    iterative_prompt: bool = True,
-    chunk_length: int = 512,
-    prompt_text: Optional[Union[str, list[str]]] = None,
-    prompt_tokens: Optional[Union[torch.Tensor, list[torch.Tensor]]] = None,
+        *,
+        model,
+        device: Union[str, torch.device],
+        decode_one_token: Callable,
+        text: str,
+        num_samples: int = 1,
+        max_new_tokens: int = 0,
+        top_p: float = 0.9,
+        top_k: int = 30,
+        repetition_penalty: float = 1.1,
+        temperature: float = 1.0,
+        compile: bool = False,
+        iterative_prompt: bool = True,
+        chunk_length: int = 512,
+        prompt_text: Optional[Union[str, list[str]]] = None,
+        prompt_tokens: Optional[Union[torch.Tensor, list[torch.Tensor]]] = None,
 ):
     assert 0 < top_p <= 1, "top_p must be in (0, 1]"
     assert 0 < temperature < 2, "temperature must be in (0, 2)"
@@ -770,10 +767,10 @@ class GenerateRequest:
 
 
 def launch_thread_safe_queue(
-    checkpoint_path,
-    device,
-    precision,
-    compile: bool = False,
+        checkpoint_path,
+        device,
+        precision,
+        compile: bool = False,
 ):
     input_queue = queue.Queue()
     init_event = threading.Event()
@@ -799,8 +796,8 @@ def launch_thread_safe_queue(
             response_queue = item.response_queue
 
             try:
-                for chunk in generate_long(
-                    model=model, decode_one_token=decode_one_token, **kwargs
+                for chunk in generate_long_parallel_batched(
+                        model=model, decode_one_token=decode_one_token, **kwargs
                 ):
                     response_queue.put(
                         WrappedGenerateResponse(status="success", response=chunk)
@@ -823,6 +820,241 @@ def launch_thread_safe_queue(
     return input_queue
 
 
+import torch
+from copy import deepcopy
+from typing import List, Optional
+
+
+# =========================
+# 1. Prompt 构建（从 generate_long 抽出）
+# =========================
+def build_prompt(
+        *,
+        model,
+        device,
+        text,
+        chunk_length,
+        prompt_text=None,
+        prompt_tokens=None,
+):
+    tokenizer = model.tokenizer
+    max_length = model.config.max_seq_len
+
+    use_prompt = bool(prompt_text) and bool(prompt_tokens)
+
+    if use_prompt and isinstance(prompt_text, str):
+        prompt_text = [prompt_text]
+        prompt_tokens = [prompt_tokens]
+
+    if prompt_tokens:
+        prompt_tokens = [p.cpu() for p in prompt_tokens]
+
+    # -------- system prompt --------
+    base_conversation = Conversation()
+
+    if use_prompt:
+        tagged = []
+        for i, t in enumerate(prompt_text):
+            if not re.search(r"<\|speaker:\d+\|>", t):
+                tagged.append(f"<|speaker:{i}|>{t}")
+            else:
+                tagged.append(t)
+
+        system_parts = [
+            TextPart(text="convert the provided text to speech reference to the following:\n\nText:\n", cal_loss=False, ),
+            TextPart(text="\n".join(tagged), cal_loss=False),
+            TextPart(text="\n\nSpeech:\n", cal_loss=False)
+        ]
+
+        all_codes = torch.cat(prompt_tokens, dim=1)
+        system_parts.append(VQPart(codes=all_codes, cal_loss=False))
+
+    else:
+        system_parts = [
+            TextPart(text="convert the provided text to speech", cal_loss=False)
+        ]
+
+    base_conversation.append(
+        Message(
+            role="system",
+            parts=system_parts,
+            cal_loss=False,
+            add_im_start=True,
+            add_im_end=True,
+        )
+    )
+
+    # -------- chunk text --------
+    turns = split_text_by_speaker(text)
+    batches = (
+        group_turns_into_batches(turns, max_speakers=5, max_bytes=chunk_length)
+        if turns else [text]
+    )
+
+    # 只做 chunk
+    encoded_prompts = []
+    metas = []
+
+    for batch_text in batches:
+        conv = deepcopy(base_conversation)
+
+        conv.append(
+            Message(
+                role="user",
+                parts=[TextPart(text=batch_text, cal_loss=False)],
+                cal_loss=False,
+                add_im_start=True,
+                add_im_end=True,
+            )
+        )
+
+        conv_gen = deepcopy(conv)
+        conv_gen.append(
+            Message(
+                role="assistant",
+                parts=[],
+                cal_loss=False,
+                modality="voice",
+                add_im_start=True,
+                add_im_end=False,
+            )
+        )
+
+        encoded, audio_masks, audio_parts = conv_gen.encode_for_inference(
+            tokenizer,
+            num_codebooks=model.config.num_codebooks,
+        )
+
+        if encoded.size(1) > max_length - 2048:
+            raise ValueError("Prompt too long")
+
+        encoded = encoded.to(device)
+
+        encoded_prompts.append(encoded)
+        metas.append(
+            {
+                "audio_masks": audio_masks,
+                "audio_parts": audio_parts,
+                "prompt_len": encoded.size(1),
+                "text": batch_text,
+            }
+        )
+
+    return encoded_prompts, metas
+
+
+# =========================
+# 2. Batched decode（核心）
+# =========================
+def generate_batched(
+        *,
+        model,
+        prompts: List[torch.Tensor],
+        decode_one_token,
+        max_new_tokens,
+        temperature,
+        top_p,
+        top_k,
+):
+    """
+    prompts: List[Tensor] [1, T]
+    return: List[Tensor] decoded sequences
+    """
+
+    device = prompts[0].device
+
+    max_len = max(p.size(1) for p in prompts)
+
+    padded = []
+    for p in prompts:
+        if p.size(1) < max_len:
+            pad = torch.zeros((1, max_len - p.size(1)), dtype=p.dtype, device=device)
+            p = torch.cat([p, pad], dim=1)
+        padded.append(p)
+
+    y = torch.cat(padded, dim=0)  # [B, T]
+
+    for _ in range(max_new_tokens):
+        next_token = decode_one_token(
+            model,
+            y,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            audio_masks=None,
+            audio_parts=None,
+        )  # [B, 1]
+
+        y = torch.cat([y, next_token], dim=1)
+
+    outputs = [y[i:i + 1] for i in range(y.size(0))]
+    return outputs
+
+
+# =========================
+# 3. 并行 TTS 主入口
+# =========================
+@torch.inference_mode()
+def generate_long_parallel_batched(
+        *,
+        model,
+        device,
+        decode_one_token,
+        text,
+        prompt_text=None,
+        prompt_tokens=None,
+        max_new_tokens=512,
+        chunk_length=512,
+        temperature=1.0,
+        top_p=0.9,
+        top_k=30,
+):
+    """
+    最小侵入版本：
+    chunk + batch decode + concat
+    """
+
+    # ===== 1. build prompts =====
+    encoded_prompts, metas = build_prompt(
+        model=model,
+        device=device,
+        text=text,
+        chunk_length=chunk_length,
+        prompt_text=prompt_text,
+        prompt_tokens=prompt_tokens,
+    )
+
+    # ===== 2. batched decode =====
+    outputs = generate_batched(
+        model=model,
+        prompts=encoded_prompts,
+        decode_one_token=decode_one_token,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+    )
+
+    # ===== 3. merge outputs =====
+    all_codes = []
+
+    for y, meta in zip(outputs, metas):
+        prompt_len = meta["prompt_len"]
+
+        codes = y[1:, prompt_len:-1].clone()
+
+        all_codes.append(codes)
+
+    final_codes = torch.cat(all_codes, dim=1)
+
+    return final_codes
+
+
+# ============================================
+# ===============   原始代码  =================
+# ============================================
+
+
 @click.command()
 @click.option(
     "--text",
@@ -861,24 +1093,24 @@ def launch_thread_safe_queue(
 @click.option("--chunk-length", type=int, default=300)
 @click.option("--output-dir", type=Path, default="output")
 def main(
-    text: str,
-    prompt_text: Optional[tuple[str, ...]],
-    prompt_tokens: Optional[tuple[Path, ...]],
-    prompt_audio: Optional[tuple[Path, ...]],
-    output: Optional[Path],
-    num_samples: int,
-    max_new_tokens: int,
-    top_p: float,
-    top_k: int,
-    temperature: float,
-    checkpoint_path: Path,
-    device: str,
-    compile: bool,
-    seed: int,
-    half: bool,
-    iterative_prompt: bool,
-    chunk_length: int,
-    output_dir: Path,
+        text: str,
+        prompt_text: Optional[tuple[str, ...]],
+        prompt_tokens: Optional[tuple[Path, ...]],
+        prompt_audio: Optional[tuple[Path, ...]],
+        output: Optional[Path],
+        num_samples: int,
+        max_new_tokens: int,
+        top_p: float,
+        top_k: int,
+        temperature: float,
+        checkpoint_path: Path,
+        device: str,
+        compile: bool,
+        seed: int,
+        half: bool,
+        iterative_prompt: bool,
+        chunk_length: int,
+        output_dir: Path,
 ) -> None:
     os.makedirs(output_dir, exist_ok=True)
     precision = torch.half if half else torch.bfloat16