feat:添加耗时统计

fish_speech/models/text2semantic/inference.py

@@ -4,7 +4,6 @@ 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
@@ -253,6 +252,9 @@ def generate(
     """
 
     # create an empty tensor of the expected final shape and fill in the current tokens
+
+    start = time.perf_counter()
+
     T = prompt.size(1)
     prompt = prompt[None].repeat(num_samples, 1, 1)
 
@@ -274,6 +276,7 @@ def generate(
     dtype = next(
         model.parameters()
     ).dtype  # model weight dtype (bfloat16), NOT prompt dtype (int32)
+    step1 = time.perf_counter()
 
     # Critical fix: Only set up cache on first run or when necessary
     if not hasattr(model, "_cache_setup_done") or not model._cache_setup_done:
@@ -284,6 +287,7 @@ def generate(
                 dtype=next(model.parameters()).dtype,
             )
         model._cache_setup_done = True
+    step2 = time.perf_counter()
 
     codebook_dim = 1 + model.config.num_codebooks
 
@@ -292,6 +296,8 @@ def generate(
     empty = torch.empty(
         (codebook_dim, model.config.max_seq_len), dtype=prompt.dtype, device=device
     )
+    step3 = time.perf_counter()
+
     empty[:, :T] = prompt
     seq = empty
 
@@ -300,13 +306,17 @@ def generate(
     top_k_val = sampling_kwargs.get("top_k", 30)
 
     temperature = torch.tensor(temp_val, device=device, dtype=dtype)
+    step4 = time.perf_counter()
+
     top_p = torch.tensor(top_p_val, device=device, dtype=dtype)
+    step5 = time.perf_counter()
 
     # Build semantic logit bias: 0 for semantic tokens + im_end, -inf for all others
     vocab_size = model.config.vocab_size
     semantic_logit_bias = torch.full(
         (1, 1, vocab_size), float("-inf"), device=device, dtype=dtype
     )
+    step6 = time.perf_counter()
 
     # [MODIFIED] Use config for semantic range
     semantic_logit_bias[
@@ -315,9 +325,9 @@ def generate(
 
     # [MODIFIED] Use tokenizer.get_token_id (Wrapper method)
     semantic_logit_bias[0, 0, model.tokenizer.get_token_id(IM_END_TOKEN)] = 0.0
+    step7 = time.perf_counter()
 
     prefill_decode = decode_one_token_ar
-
     first_token = prefill_decode(
         model,
         prompt.view(1, codebook_dim, -1),
@@ -330,9 +340,11 @@ def generate(
         audio_parts,
     )
     seq[:, T: T + 1] = first_token
+    step8 = time.perf_counter()
 
     # Recreate input_pos
     input_pos = torch.tensor([T], device=device, dtype=torch.int)
+    step9 = time.perf_counter()
 
     x = decode_n_tokens(
         model,
@@ -349,10 +361,19 @@ def generate(
     )
     seq = seq[:, : T + 1 + x.size(1)]
     seq[:, T + 1:] = x
+    step10 = time.perf_counter()
 
     # Clean up temporary variables
     del first_token, x, prompt, empty, input_pos
 
+    step11 = time.perf_counter()
+
+    logger.info(f"elapse \n"
+                f"step1: {step1 - start}, step2: {step2 - step1}, step3: {step3 - step2}"
+                f"step4: {step4 - step3}, step5: {step5 - step4} step6: {step6 - step5}"
+                f"step7: {step7 - step6} step8: {step8 - step7} step9: {step9 - step8}"
+                f"step10: {step10 - step9} step11: {step11 - step10}")
+
     return seq
 
 
@@ -798,7 +819,7 @@ def launch_thread_safe_queue(
             response_queue = item.response_queue
 
             try:
-                for chunk in generate_long_batched(
+                for chunk in generate_long(
                         model=model, decode_one_token=decode_one_token, **kwargs
                 ):
                     response_queue.put(
@@ -822,219 +843,6 @@ def launch_thread_safe_queue(
     return input_queue
 
 
-@torch.inference_mode()
-def generate_long_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,
-):
-    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)
-
-    if use_prompt and isinstance(prompt_text, str):
-        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]
-
-    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_prompt_text = []
-        for i, t in enumerate(prompt_text):
-            if not re.search(r"<\|speaker:\d+\|>", t):
-                tagged_prompt_text.append(f"<|speaker:{i}|>{t}")
-            else:
-                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,
-            ),
-        ]
-
-        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)
-        ]
-
-    base_conversation.append(
-        Message(
-            role="system",
-            parts=system_parts,
-            cal_loss=False,
-            add_im_start=True,
-            add_im_end=True,
-        )
-    )
-
-    # =========================
-    # split batches（不动）
-    # =========================
-    turns = split_text_by_speaker(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")
-
-    # =========================
-    # worker function（核心）
-    # =========================
-    def run_one_batch(b_idx: int, b_text: str):
-        conversation = deepcopy(base_conversation)
-
-        logger.info(f"[Batch {b_idx}] start")
-
-        conversation.append(
-            Message(
-                role="user",
-                parts=[TextPart(text=b_text, cal_loss=False)],
-                cal_loss=False,
-                add_im_start=True,
-                add_im_end=True,
-            )
-        )
-
-        conversation_gen = deepcopy(conversation)
-        conversation_gen.append(
-            Message(
-                role="assistant",
-                parts=[],
-                cal_loss=False,
-                modality="voice",
-                add_im_start=True,
-                add_im_end=False,
-            )
-        )
-
-        encoded, audio_masks, audio_parts = conversation_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)
-        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,
-        )
-
-        codes = y[1:, prompt_length:-1].clone()
-
-        logger.info(f"[Batch {b_idx}] done")
-
-        return b_idx, codes, b_text
-
-    # =========================
-    # parallel execution（关键修改点）
-    # =========================
-    for sample_idx in range(num_samples):
-
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-
-        t0 = time.perf_counter()
-
-        results = {}
-
-        with ThreadPoolExecutor(max_workers=min(3, len(batches))) as executor:
-            futures = [
-                executor.submit(run_one_batch, i, b)
-                for i, b in enumerate(batches)
-            ]
-
-            for f in as_completed(futures):
-                batch_idx, codes, batch_text = f.result()
-
-                results[batch_idx] = codes
-
-                # ⭐ 保持你原来的 streaming 行为
-                yield GenerateResponse(
-                    action="sample",
-                    codes=codes,
-                    text=batch_text,
-                )
-
-        # =========================
-        # stats（不动逻辑）
-        # =========================
-        all_codes = [results[i] for i in sorted(results)]
-
-        final_latency = time.perf_counter() - t0
-
-        logger.info(f"Sample {sample_idx} done in {final_latency:.2f}s")
-
-        if torch.cuda.is_available():
-            logger.info(
-                f"GPU Memory used: {torch.cuda.max_memory_reserved() / 1e9:.2f} GB"
-            )
-
-        # cleanup
-        del results
-        import gc
-        gc.collect()
-
-        yield GenerateResponse(action="next")
-
-
 # ============================================
 # ===============   原始代码  =================
 # ============================================