update app.py and requirements

app.py

@@ -1,20 +1,17 @@
 import os
 import traceback
-import spaces
 import gradio as gr
 import numpy as np
-import pyrootutils
 import torch
+import spaces
 from loguru import logger
-from transformers import AutoTokenizer
-from vllm import LLM, SamplingParams, TokensPrompt
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from funasr_onnx import Paraformer
 from huggingface_hub import snapshot_download
 
 from tools.wer import compute_wers
 
 os.environ["EINX_FILTER_TRACEBACK"] = "false"
-os.environ["VLLM_USE_V1"] = "0"
 
 from i18n import i18n
 from text.chn_text_norm.text import Text as ChnNormedText
@@ -31,54 +28,27 @@ PARAFORMER_REPO_ID = "funasr/Paraformer-large"
 # logo
 LOGO_URL = "https://raw.githubusercontent.com/ASLP-lab/VoiceSculptor/main/assets/logo.png"
 
-
+# ===== Global cache =====
 model = None
 codec_model = None
 asr_model = None
 tokenizer = None
-
-@spaces.GPU
-def load_models():
-    global model, codec_model, asr_model, tokenizer
-    
-    # 只有当模型为空时才加载
-    if tokenizer is None:
-        tokenizer = AutoTokenizer.from_pretrained(LLASA_MODEL_ID)
-        
-    if model is None:
-        logger.info("🚀 Loading vLLM model on GPU...")
-        model = LLM(
-            model=LLASA_MODEL_ID,
-            gpu_memory_utilization=0.90, 
-            max_model_len=2048, 
-            enable_prefix_caching=True, 
-            dtype='auto',
-            quantization=None,
-            enforce_eager=False,
-            kv_cache_dtype='auto'
-        )
-        
-    if codec_model is None:
-        logger.info("🚀 Loading XCodec2...")
-        codec_model = XCodec2Model.from_pretrained(XCODEC_MODEL_ID).eval().to("cuda")
-        
-    if asr_model is None:
-        logger.info("🚀 Loading Paraformer...")
-        paraformer_dir = snapshot_download(repo_id=PARAFORMER_REPO_ID, local_dir="checkpoints/Paraformer-large")
-        asr_model = Paraformer(paraformer_dir, batch_size=5, quantize=True)
+device= 'cuda'
 
 def normalize_text_final(user_input: str) -> str:
     return ChnNormedText(raw_text=user_input).normalize()
 
 
-def extract_speech_ids(speech_tokens_str):
+def extract_speech_ids(token_strs: list[str]) -> list[int]:
+    """把 tokenizer 输出的 token 字符串列表中形如 <|s_123|> 的 token 提取成 int id"""
     speech_ids = []
-    for token_str in speech_tokens_str:
-        if token_str.startswith("<|s_") and token_str.endswith("|>"):
-            num_str = token_str[4:-2]
-            speech_ids.append(int(num_str))
-        else:
-            logger.warning(f"Unexpected token: {token_str}")
+    for t in token_strs:
+        if t.startswith("<|s_") and t.endswith("|>"):
+            num_str = t[4:-2]
+            try:
+                speech_ids.append(int(num_str))
+            except Exception:
+                logger.warning(f"Bad speech token: {t}")
     return speech_ids
 
 
@@ -97,7 +67,6 @@ def get_asr(asr_model: Paraformer, wav_list: list[np.ndarray]) -> list[str]:
             else:
                 texts.append(preds[0] if len(preds) > 0 else "")
 
-        # 容错：batch 返回数量不一致 -> fallback
         if len(texts) != len(wav_list):
             logger.warning(f"[ASR] batch返回数量不一致: got {len(texts)} expect {len(wav_list)}，fallback逐条补齐")
             texts = []
@@ -136,17 +105,71 @@ def get_asr(asr_model: Paraformer, wav_list: list[np.ndarray]) -> list[str]:
                 texts.append("")
         return texts
 
+
+def _safe_load_tokenizer(model_id: str) -> AutoTokenizer:
+    try:
+        tok = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    except TypeError:
+        tok = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, use_fast=False)
+
+    if tok.pad_token_id is None:
+        if tok.eos_token_id is not None:
+            tok.pad_token = tok.eos_token
+    return tok
+
+
+def _safe_load_lm(model_id: str, device: str) -> AutoModelForCausalLM:
+    m = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        trust_remote_code=True,
+    )
+    m.eval().to(device)
+    return m
+
 @spaces.GPU
-def inference_batch(
-    model: LLM,
-    codec_model: XCodec2Model,
-    device: str,
-    tokenizer: AutoTokenizer,
+def load_models(force_device: str | None = None):
+    """本地：加载并缓存模型（无 spaces/ZeroGPU）"""
+    global model, codec_model, asr_model, tokenizer
+
+    logger.info(f"Using device: {device}")
+
+    if tokenizer is None:
+        logger.info("Loading tokenizer...")
+        tokenizer = _safe_load_tokenizer(LLASA_MODEL_ID)
+
+    if model is None:
+        logger.info("Loading AutoModelForCausalLM...")
+        model = _safe_load_lm(LLASA_MODEL_ID, device=device)
+
+    if codec_model is None:
+        logger.info("Loading XCodec2...")
+        codec_model = XCodec2Model.from_pretrained(XCODEC_MODEL_ID).eval().to(device)
+
+    if asr_model is None:
+        logger.info("Loading Paraformer (funasr_onnx)...")
+        paraformer_dir = snapshot_download(
+            repo_id=PARAFORMER_REPO_ID,
+            local_dir="checkpoints/Paraformer-large",
+            local_dir_use_symlinks=False,
+        )
+        asr_model = Paraformer(paraformer_dir, batch_size=5, quantize=True)
+
+    logger.info("✅ All models loaded.")
+
+load_models()
+
+@torch.inference_mode()
+def inference_batch_transformers(
+    lm: AutoModelForCausalLM,
+    codec: XCodec2Model,
+    tok: AutoTokenizer,
     refined_text: str,
     instruct_text: str,
     control_tags: str,
     batch_size: int = 5,
+    max_new_tokens: int = 2048,
 ) -> list[tuple[int, np.ndarray]]:
+
     refined_text_norm = normalize_text_final(refined_text)
     instruct_text_norm = normalize_text_final(instruct_text)
 
@@ -162,61 +185,53 @@ def inference_batch(
         {"role": "assistant", "content": "<|SPEECH_GENERATION_START|>"},
     ]
 
-    with torch.no_grad():
-        input_ids = tokenizer.apply_chat_template(
-            chat,
-            tokenize=True,
-            return_tensors="pt",
-            continue_final_message=True,
-        ).to(device)
-
-        speech_end_id = tokenizer.convert_tokens_to_ids("<|SPEECH_GENERATION_END|>")
-        prompt_ids = input_ids.squeeze(0).tolist()
-        prompts = [TokensPrompt(prompt_token_ids=prompt_ids) for _ in range(batch_size)]
-
-        base_seed = int.from_bytes(os.urandom(4), "little")
-
-        try:
-            sampling_params_list = [
-                SamplingParams(
-                    temperature=0.9,
-                    top_p=0.95,
-                    top_k=15,
-                    max_tokens=2048,
-                    repetition_penalty=1.05,
-                    stop_token_ids=[speech_end_id],
-                    seed=base_seed + i,
-                )
-                for i in range(batch_size)
-            ]
-            outputs = model.generate(prompts=prompts, sampling_params=sampling_params_list)
-        except TypeError:
-            logger.warning("[vLLM] 当前版本不支持 SamplingParams(seed=...)，将不带 seed 生成")
-            sampling_params = SamplingParams(
-                temperature=0.9,
-                top_p=0.95,
-                top_k=15,
-                max_tokens=2048,
-                repetition_penalty=1.05,
-                stop_token_ids=[speech_end_id],
-            )
-            outputs = model.generate(prompts=prompts, sampling_params=sampling_params)
-
-        audios: list[tuple[int, np.ndarray]] = []
-        for out in outputs:
-            token_ids = out.outputs[0].token_ids
-            if len(token_ids) > 0 and token_ids[-1] == speech_end_id:
-                token_ids = token_ids[:-1]
-
-            speech_tokens = tokenizer.batch_decode(token_ids, skip_special_tokens=True)
-            speech_tokens = extract_speech_ids(speech_tokens)
-
-            speech_tokens_t = torch.tensor(speech_tokens, device=device).unsqueeze(0).unsqueeze(0)
-            wav = codec_model.decode_code(speech_tokens_t)
-            wav = wav.squeeze(0).squeeze(0).detach().cpu().numpy().astype(np.float32)
-            audios.append((16000, wav))
-
-        return audios
+    input_ids_1 = tok.apply_chat_template(
+        chat,
+        tokenize=True,
+        return_tensors="pt",
+        continue_final_message=True,
+    ).to(device)
+
+    speech_end_id = tok.convert_tokens_to_ids("<|SPEECH_GENERATION_END|>")
+    pad_id = tok.pad_token_id if tok.pad_token_id is not None else (tok.eos_token_id or speech_end_id)
+
+    outputs = lm.generate(
+        input_ids=input_ids_1,
+        do_sample=True,
+        top_p=0.95,
+        temperature=0.9,
+        top_k=15,
+        repetition_penalty=1.05,
+        max_new_tokens=max_new_tokens,
+        eos_token_id=speech_end_id,
+        pad_token_id=pad_id,
+        num_return_sequences=batch_size, 
+        use_cache=True,
+    )
+
+    prompt_len = input_ids_1.shape[1]
+    audios: list[tuple[int, np.ndarray]] = []
+
+    for i in range(outputs.shape[0]):
+        gen_ids = outputs[i, prompt_len:].tolist()
+        if len(gen_ids) > 0 and gen_ids[-1] == speech_end_id:
+            gen_ids = gen_ids[:-1]
+
+        token_strs = tok.convert_ids_to_tokens(gen_ids, skip_special_tokens=False)
+        speech_ids = extract_speech_ids(token_strs)
+
+        if len(speech_ids) == 0:
+            logger.warning("[TTS] No speech tokens extracted, outputting silence.")
+            audios.append((16000, np.zeros((16000,), dtype=np.float32)))
+            continue
+
+        speech_tokens_t = torch.tensor(speech_ids, device=device).unsqueeze(0).unsqueeze(0)
+        wav = codec.decode_code(speech_tokens_t)
+        wav = wav.squeeze(0).squeeze(0).detach().cpu().numpy().astype(np.float32)
+        audios.append((16000, wav))
+
+    return audios
+
 
 def build_control_tags(age, gender, pitch, pitch_var, volume, speed, emo):
     tag_map = {
@@ -230,7 +245,8 @@ def build_control_tags(age, gender, pitch, pitch_var, volume, speed, emo):
         "音量较小": "<|音量较小|>", "音量很小": "<|音量很小|>",
         "语速很快": "<|语速很快|>", "语速较快": "<|语速较快|>", "语速中等": "<|语速中等|>",
         "语速较慢": "<|语速较慢|>", "语速很慢": "<|语速很慢|>",
-        "开心": "<|开心|>", "生气": "<|生气|>", "难过": "<|难过|>", "惊讶": "<|惊讶|>", "厌恶": "<|厌恶|>", "害怕": "<|害怕|>",
+        "开心": "<|开心|>", "生气": "<|生气|>", "难过": "<|难过|>", "惊讶": "<|惊讶|>",
+        "厌恶": "<|厌恶|>", "害怕": "<|害怕|>",
     }
     tags = []
     for v in [gender, age, speed, volume, pitch, pitch_var, emo]:
@@ -238,21 +254,23 @@ def build_control_tags(age, gender, pitch, pitch_var, volume, speed, emo):
             tags.append(tag_map[v])
     return "".join(tags)
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=240)
 def inference_select_best3(refined_text, instruct_text, age, gender, pitch, pitch_var, volume, speed, emo):
-    load_models()
+
     control_tags = build_control_tags(age, gender, pitch, pitch_var, volume, speed, emo)
+
     try:
-        audios5 = inference_batch(
-            model=model,
-            codec_model=codec_model,
-            device='cuda',
-            tokenizer=tokenizer,
+        audios5 = inference_batch_transformers(
+            lm=model,
+            codec=codec_model,
+            tok=tokenizer,
             refined_text=refined_text,
             instruct_text=instruct_text,
             control_tags=control_tags,
             batch_size=5,
+            max_new_tokens=2048,
         )
+
         wav_list = [wav for (_, wav) in audios5]
         asr_texts = get_asr(asr_model, wav_list)
 
@@ -264,14 +282,15 @@ def inference_select_best3(refined_text, instruct_text, age, gender, pitch, pitc
             logger.info(f"[ASR/WER] idx={i} wer={w:.4f}  gt='{refined_text_norm}'  asr='{hyp}'")
 
         best_idx = np.argsort(np.array(wers))[:3].tolist()
-        logger.info(f"[ASR/WER] best_idx={best_idx} best_wers={[float(wers[i]) for i in best_idx]}")
         best3 = [audios5[i] for i in best_idx]
         return best3[0], best3[1], best3[2]
+
     except Exception as e:
         logger.error(f"推理/ASR/WER 失败: {e}", exc_info=True)
         logger.error("错误详细信息:\n" + traceback.format_exc())
         return None, None, None
 
+
 def build_app():
 
     INSTRUCT_TEMPLATES = {

requirements.txt

@@ -1,7 +1,7 @@
 gradio
 torch
 transformers
-vllm
+spaces
 funasr-onnx
 huggingface_hub
 jiwer