Update VoxCPM/inference_lora.py

VoxCPM/inference_lora.py

@@ -7,7 +7,7 @@ from voxcpm.training.config import load_yaml_config
 import argparse
 import torch
 import os
-
+import re
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--lora_ckpt", type=str, required=True)
@@ -34,7 +34,10 @@ def main():
         training=False,
         lora_config=lora_cfg,
     )
-
+        
+    from src.voxcpm.utils.text_normalize import TextNormalizer
+    text_normalizer = TextNormalizer()
+    
     # 3. 加载 LoRA 权重（在 compile 后也能正常工作）
     ckpt_dir = Path(args.lora_ckpt)
     if not ckpt_dir.exists():
@@ -49,8 +52,11 @@ def main():
     print(f"\n[3/3] 开始推理...")
     if args.text:
         with torch.inference_mode():
+            target_text = args.text.replace("\n", " ")
+            target_text = re.sub(r'\s+', ' ', target_text)
+            target_text = text_normalizer.normalize(target_text)
             wav = model.generate(
-            target_text=args.text,
+            target_text=target_text,
             cfg_value=args.cfg_value,             # LM guidance on LocDiT, higher for better adherence to the prompt, but maybe worse
             inference_timesteps=args.inference_timesteps,   # LocDiT inference timesteps, higher for better result, lower for fast speed
             retry_badcase=True,        # enable retrying mode for some bad cases (unstoppable)
@@ -73,8 +79,11 @@ def main():
                 texts.append((wav_id, text))
         for wav_id, text in texts:
             with torch.inference_mode():
+                target_text = text.replace("\n", " ")
+                target_text = re.sub(r'\s+', ' ', target_text)
+                target_text = text_normalizer.normalize(target_text)
                 wav = model.generate(
-                target_text=text,
+                target_text=target_text,
                 cfg_value=args.cfg_value,             # LM guidance on LocDiT, higher for better adherence to the prompt, but maybe worse
                 inference_timesteps=args.inference_timesteps,   # LocDiT inference timesteps, higher for better result, lower for fast speed
                 retry_badcase=True,        # enable retrying mode for some bad cases (unstoppable)