feat:修改generate_long的并发逻辑

fish_speech/models/text2semantic/inference.py

@@ -4,6 +4,8 @@ import re
 import threading
 import time
 import traceback
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from copy import deepcopy
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, Literal, Optional, Tuple, Union
@@ -796,7 +798,7 @@ def launch_thread_safe_queue(
             response_queue = item.response_queue
 
             try:
-                for chunk in generate_long_parallel_batched(
+                for chunk in generate_long_batched(
                         model=model, decode_one_token=decode_one_token, **kwargs
                 ):
                     response_queue.put(
@@ -820,25 +822,38 @@ 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(
+@torch.inference_mode()
+def generate_long_batched(
         *,
         model,
-        device,
-        text,
-        chunk_length,
-        prompt_text=None,
-        prompt_tokens=None,
+        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,
 ):
-    tokenizer = model.tokenizer
-    max_length = model.config.max_seq_len
+    assert 0 < top_p <= 1, "top_p must be in (0, 1]"
+    assert 0 < temperature < 2, "temperature must be in (0, 2)"
+
+    logger.info(f"generate_long.param.device: {device}")
+    logger.info(f"generate_long.param.text: {text}")
+    logger.info(f"generate_long.param.max_new_tokens: {max_new_tokens}")
+    logger.info(f"generate_long.param.top_p: {top_p}")
+    logger.info(f"generate_long.param.top_k: {top_k}")
+    logger.info(f"generate_long.param.temperature: {temperature}")
+    logger.info(f"generate_long.param.compile: {compile}")
+    logger.info(f"generate_long.param.chunk_length: {chunk_length}")
+    logger.info(f"generate_long.param.prompt_text: {prompt_text}")
+    logger.info(f"generate_long.param.prompt_tokens: {prompt_tokens}")
 
     use_prompt = bool(prompt_text) and bool(prompt_tokens)
 
@@ -846,29 +861,41 @@ def build_prompt(
         prompt_text = [prompt_text]
         prompt_tokens = [prompt_tokens]
 
+    if use_prompt:
+        assert len(prompt_text) == len(prompt_tokens)
+
     if prompt_tokens:
         prompt_tokens = [p.cpu() for p in prompt_tokens]
 
-    # -------- system prompt --------
+    tokenizer = model.tokenizer
+    max_length = model.config.max_seq_len
+    model_size = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    # =========================
+    # build base conversation（不动）
+    # =========================
     base_conversation = Conversation()
 
     if use_prompt:
-        tagged = []
+        tagged_prompt_text = []
         for i, t in enumerate(prompt_text):
             if not re.search(r"<\|speaker:\d+\|>", t):
-                tagged.append(f"<|speaker:{i}|>{t}")
+                tagged_prompt_text.append(f"<|speaker:{i}|>{t}")
             else:
-                tagged.append(t)
+                tagged_prompt_text.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)
+            TextPart(
+                text="convert the provided text to speech reference to the following:\n\nText:\n",
+                cal_loss=False,
+            ),
         ]
 
+        system_parts.append(TextPart(text="\n".join(tagged_prompt_text), cal_loss=False))
+        system_parts.append(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)
@@ -884,32 +911,39 @@ def build_prompt(
         )
     )
 
-    # -------- chunk text --------
+    # =========================
+    # split batches（不动）
+    # =========================
     turns = split_text_by_speaker(text)
-    batches = (
-        group_turns_into_batches(turns, max_speakers=5, max_bytes=chunk_length)
-        if turns else [text]
-    )
+    if turns:
+        batches = group_turns_into_batches(
+            turns, max_speakers=5, max_bytes=chunk_length
+        )
+    else:
+        batches = [text]
+
+    logger.info(f"Split into {len(batches)} batches")
 
-    # 只做 chunk
-    encoded_prompts = []
-    metas = []
+    # =========================
+    # worker function（核心）
+    # =========================
+    def run_one_batch(b_idx: int, b_text: str):
+        conversation = deepcopy(base_conversation)
 
-    for batch_text in batches:
-        conv = deepcopy(base_conversation)
+        logger.info(f"[Batch {b_idx}] start")
 
-        conv.append(
+        conversation.append(
             Message(
                 role="user",
-                parts=[TextPart(text=batch_text, cal_loss=False)],
+                parts=[TextPart(text=b_text, cal_loss=False)],
                 cal_loss=False,
                 add_im_start=True,
                 add_im_end=True,
             )
         )
 
-        conv_gen = deepcopy(conv)
-        conv_gen.append(
+        conversation_gen = deepcopy(conversation)
+        conversation_gen.append(
             Message(
                 role="assistant",
                 parts=[],
@@ -920,138 +954,85 @@ def build_prompt(
             )
         )
 
-        encoded, audio_masks, audio_parts = conv_gen.encode_for_inference(
+        encoded, audio_masks, audio_parts = conversation_gen.encode_for_inference(
             tokenizer,
-            num_codebooks=model.config.num_codebooks,
+            num_codebooks=model.config.num_codebooks
         )
 
         if encoded.size(1) > max_length - 2048:
-            raise ValueError("Prompt too long")
+            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,
-            }
+        prompt_length = encoded.size(1)
+
+        y = generate(
+            model=model,
+            prompt=encoded,
+            max_new_tokens=max_new_tokens,
+            audio_masks=audio_masks,
+            audio_parts=audio_parts,
+            decode_one_token=decode_one_token,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
         )
 
-    return encoded_prompts, metas
+        codes = y[1:, prompt_length:-1].clone()
 
+        logger.info(f"[Batch {b_idx}] done")
 
-# =========================
-# 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
-    """
+        return b_idx, codes, b_text
 
-    device = prompts[0].device
-
-    max_len = max(p.size(1) for p in prompts)
+    # =========================
+    # parallel execution（关键修改点）
+    # =========================
+    for sample_idx in range(num_samples):
 
-    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)
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
 
-    y = torch.cat(padded, dim=0)  # [B, T]
+        t0 = time.perf_counter()
 
-    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]
+        results = {}
 
-        y = torch.cat([y, next_token], dim=1)
+        with ThreadPoolExecutor(max_workers=3) as executor:
+            futures = [
+                executor.submit(run_one_batch, i, b)
+                for i, b in enumerate(batches)
+            ]
 
-    outputs = [y[i:i + 1] for i in range(y.size(0))]
-    return outputs
+            for f in as_completed(futures):
+                batch_idx, codes, batch_text = f.result()
 
+                results[batch_idx] = codes
 
-# =========================
-# 3. 并行 TTS 主入口
-# =========================
-@torch.inference_mode()
-def generate_long_parallel_batched(
-        *,
-        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,
-):
-    """
-    最小侵入版本：
-    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,
-    )
+                # ⭐ 保持你原来的 streaming 行为
+                yield GenerateResponse(
+                    action="sample",
+                    codes=codes,
+                    text=batch_text,
+                )
 
-    # ===== 3. merge outputs =====
-    all_codes = []
+        # =========================
+        # stats（不动逻辑）
+        # =========================
+        all_codes = [results[i] for i in sorted(results)]
 
-    for y, meta in zip(outputs, metas):
-        prompt_len = meta["prompt_len"]
+        final_latency = time.perf_counter() - t0
 
-        codes = y[1:, prompt_len:-1].clone()
+        logger.info(f"Sample {sample_idx} done in {final_latency:.2f}s")
 
-        all_codes.append(codes)
+        if torch.cuda.is_available():
+            logger.info(
+                f"GPU Memory used: {torch.cuda.max_memory_reserved() / 1e9:.2f} GB"
+            )
 
-    final_codes = torch.cat(all_codes, dim=1)
+        # cleanup
+        del results
+        import gc
+        gc.collect()
 
-    return final_codes
+        yield GenerateResponse(action="next")
 
 
 # ============================================