Update inference demo

README.md

@@ -91,16 +91,16 @@ python -m spacy download en_core_web_sm
 
 ```bash
 中文：
-python kokoro_ax.py --text "致力于打造世界领先的人工智能感知与边缘计算芯片。" --lang zh --voice checkpoints/voices/zf_xiaoyi.pt --output output_zh.wav -d models -f 0.3
+python demo_kokoro_ax.py --text "致力于打造世界领先的人工智能感知与边缘计算芯片。" --lang z --voice checkpoints/voices/zf_xiaoyi.pt --output output_zh.wav -d models -f 0.3
 英文：
-python kokoro_ax.py --text "The sky above the port was the color of television, tuned to a dead channel." --lang en --voice checkpoints/voices/af_heart.pt --output output_en.wav -d models -f 0.3
+python demo_kokoro_ax.py --text "The sky above the port was the color of television, tuned to a dead channel." --lang a --voice checkpoints/voices/af_heart.pt --output output_en.wav -d models -f 0.3
 日文：
-python kokoro_ax.py --text "「もしおれがただ偶然、そしてこうしようというつもりでなくここに立っているのなら、ちょっとばかり絶望するところだな」と、そんなことが彼の頭に思い浮かんだ。" --lang ja --voice checkpoints/voices/jm_kumo.pt --output output_jp.wav -d models -f 0.3
+python demo_kokoro_ax.py --text "「もしおれがただ偶然、そしてこうしようというつもりでなくここに立っているのなら、ちょっとばかり絶望するところだな」と、そんなことが彼の頭に思い浮かんだ。" --lang j --voice checkpoints/voices/jm_kumo.pt --output output_jp.wav -d models -f 0.3
 ```
 
 注：AX630C平台推理，更改模型文件目录即可，将 -d models 改成 -d models_620E 如下：
 ```
-python kokoro_ax.py --text "The sky above the port was the color of television, tuned to a dead channel." --lang en --voice checkpoints/voices/af_heart.pt --output output_en.wav -d models_620E -f 0.3
+python kokoro_ax.py --text "The sky above the port was the color of television, tuned to a dead channel." --lang a --voice checkpoints/voices/af_heart.pt --output output_en.wav -d models_620E -f 0.3
 ```
 
 参数说明：
@@ -109,7 +109,7 @@ python kokoro_ax.py --text "The sky above the port was the color of television,
 | --axmodel_dir, -d    | 模型文件目录（默认 models）   |
 | --voice, -v          | 声纹文件路径（必填）         |
 | --text, -t           | 合成文本（支持多语言）       |
-| --lang, -l           | 语言代码（如 en, zh, ja, ...）,目前只测了中英日  |
+| --lang, -l           | 语言代码（如 a, z, j, ...）,目前只测了中英日  |
 | --config, -c         | 配置文件路径                 |
 | --output, -o         | 输出 wav 文件名              |
 | --fade_out, -f       | 音频结尾淡出时长（秒），减少音频末尾有杂音  |

demo_kokoro_ax.py

@@ -0,0 +1,321 @@
+import argparse
+import os
+import numpy as np
+import soundfile as sf
+import time
+import torch
+from typing import Tuple
+#from onnxruntime import InferenceSession
+from axengine import InferenceSession
+
+from inference_utils import (
+    audio_numpy_concat, load_vocab_from_config, init_g2p,
+    split_input_ids_semantic, generate_input_ids_from_text,
+    process_and_merge_sentences, run_batch_inference, apply_fade_out,
+    SAMPLE_RATE, DEFAULT_SPEED, DEFAULT_PAUSE, DEFAULT_FADE_OUT
+)
+
+
+class InferenceEngine:
+    """推理"""
+    
+    # 常量
+    FIXED_SEQ_LEN = 96
+    N_FFT = 20
+    HOP_LENGTH = 5
+    DOUBLE_INPUT_THRESHOLD = 32  # 输入长度小于此值时复制一倍,适配短文本
+    
+    def __init__(self, axmodel_dir: str):
+        self.axmodel_dir = axmodel_dir
+        
+        # 加载模型
+        model_files = {
+            'model1': "kokoro_part1_96.axmodel",
+            'model2': "kokoro_part2_96.axmodel",
+            'model3': "kokoro_part3_96.axmodel",
+            'model4': "model4_har_sim.onnx"
+        }
+        
+        providers = ['CPUExecutionProvider']
+        self.session1 = InferenceSession(os.path.join(axmodel_dir, model_files['model1']))
+        self.session2 = InferenceSession(os.path.join(axmodel_dir, model_files['model2']))
+        self.session3 = InferenceSession(os.path.join(axmodel_dir, model_files['model3']))
+        import onnxruntime as ort
+        self.session4 = ort.InferenceSession(os.path.join(axmodel_dir, model_files['model4']), providers=providers)
+        
+        # 统计
+        self.model1_time = 0.0
+        self.model2_time = 0.0
+        self.model3_time = 0.0
+        self.har_time = 0.0
+        self.inference_count = 0
+    
+    def _compute_external_preprocessing(self, input_ids: np.ndarray, actual_len: int = None) -> Tuple[np.ndarray, np.ndarray]:
+        """计算输入预处理：长度和mask"""
+        if actual_len is None:
+            actual_len = self.FIXED_SEQ_LEN
+        input_lengths = np.full((input_ids.shape[0],), actual_len, dtype=np.int64)
+        text_mask = np.arange(self.FIXED_SEQ_LEN)[np.newaxis, :] >= input_lengths[:, np.newaxis]
+        return input_lengths, text_mask
+    
+    def _compute_har_onnx(self, F0_pred: np.ndarray) -> np.ndarray:
+        """使用ONNX模型计算谐波"""
+        return self.session4.run(None, {'F0_pred': F0_pred.astype(np.float32)})[0]
+    
+    def _postprocess_x_to_audio(self, x: np.ndarray) -> np.ndarray:
+        """将频谱转换为音频波形"""
+        spec_part = x[:, :self.N_FFT//2+1, :]
+        phase_part = x[:, self.N_FFT//2+1:, :]
+        
+        spec = np.exp(spec_part)
+        phase = np.sin(phase_part)
+        
+        spec_torch = torch.from_numpy(spec).float()
+        phase_torch = torch.from_numpy(phase).float()
+        cos_part = torch.sqrt(1.0 - phase_torch.pow(2).clamp(0, 1))
+        
+        real = spec_torch * cos_part
+        imag = spec_torch * phase_torch
+        complex_spec = torch.complex(real, imag)
+        
+        audio = torch.istft(
+            complex_spec, n_fft=self.N_FFT, hop_length=self.HOP_LENGTH,
+            win_length=self.N_FFT, window=torch.hann_window(self.N_FFT),
+            center=True, return_complex=False
+        )
+        return audio.squeeze().numpy()
+    
+    def _prepare_input_ids(self, input_ids: np.ndarray, actual_len: int) -> Tuple[np.ndarray, int, bool]:
+        """准备输入ID，对短输入进行复制处理"""
+        is_doubled = False
+        original_actual_len = actual_len
+        
+        if actual_len <= self.DOUBLE_INPUT_THRESHOLD:
+            is_doubled = True
+            valid_content = input_ids[:, :actual_len]
+            input_ids_doubled = np.concatenate([valid_content, valid_content], axis=1)
+            
+            padding_len = self.FIXED_SEQ_LEN - input_ids_doubled.shape[1]
+            if padding_len > 0:
+                input_ids = np.concatenate([input_ids_doubled, np.zeros((1, padding_len), dtype=input_ids.dtype)], axis=1)
+            else:
+                input_ids = input_ids_doubled[:, :self.FIXED_SEQ_LEN]
+            
+            actual_len = min(original_actual_len * 2, self.FIXED_SEQ_LEN)
+        
+        return input_ids, actual_len, is_doubled
+    
+    def inference_single_chunk(
+        self,
+        input_ids: np.ndarray,
+        ref_s: np.ndarray,
+        actual_len: int,
+        speed: float
+    ) -> Tuple[np.ndarray, int, int]:
+        """单个chunk的推理"""
+        self.inference_count += 1
+        
+        input_ids, actual_len, is_doubled = self._prepare_input_ids(input_ids, actual_len)
+        
+        input_lengths, text_mask = self._compute_external_preprocessing(input_ids, actual_len=actual_len)
+        
+        # Model1: 预测duration
+        t1 = time.time()
+        outputs1 = self.session1.run(None, {'input_ids': input_ids.astype(np.int32), 'ref_s': ref_s, 'text_mask': text_mask.astype(np.uint8)})
+        self.model1_time += time.time() - t1
+        duration, d = outputs1
+        
+        # 处理duration并对齐
+        pred_dur, total_frames = self._process_duration(duration, actual_len, speed)
+        pred_aln_trg = self._create_alignment_matrix(pred_dur, total_frames)
+        
+        # Model2: 预测F0和ASR特征
+        d_transposed = np.transpose(d, (0, 2, 1))
+        en = d_transposed @ pred_aln_trg
+        
+        t2 = time.time()
+        outputs2 = self.session2.run(None, {
+            'en': en.astype(np.float32),
+            'ref_s': ref_s,
+            'input_ids': input_ids.astype(np.int32),
+            'text_mask': text_mask.astype(np.float32),
+            'pred_aln_trg': pred_aln_trg.astype(np.float32)
+        })
+        self.model2_time += time.time() - t2
+        F0_pred, N_pred, asr = outputs2
+        
+        # Model4: 计算谐波
+        t_har = time.time()
+        har = self._compute_har_onnx(F0_pred)
+        self.har_time += time.time() - t_har
+        
+        # Model3: 解码生成频谱
+        t3 = time.time()
+        outputs3 = self.session3.run(None, {
+            'asr': asr, 'F0_pred': F0_pred, 'N_pred': N_pred, 'ref_s': ref_s, 'har': har
+        })
+        self.model3_time += time.time() - t3
+        x = outputs3[0]
+        
+        # 转换为音频
+        audio = self._postprocess_x_to_audio(x)
+        actual_content_frames = pred_dur[:actual_len].sum()
+        
+        # 如果输入被复制了，截取前一半音频
+        if is_doubled:
+            audio = audio[:len(audio) // 2]
+            actual_content_frames = actual_content_frames // 2
+            total_frames = total_frames // 2
+        
+        return audio, actual_content_frames, total_frames
+    
+    def _process_duration(self, duration: np.ndarray, actual_len: int, speed: float) -> Tuple[np.ndarray, int]:
+        """处理duration预测，调整到固定帧数"""
+        duration_processed = 1.0 / (1.0 + np.exp(-duration))
+        duration_processed = duration_processed.sum(axis=-1) / speed
+        pred_dur_original = np.round(duration_processed).clip(min=1).astype(np.int64).squeeze()
+        
+        # 分离实际内容和padding
+        pred_dur_actual = pred_dur_original[:actual_len]
+        pred_dur_padding = np.zeros(self.FIXED_SEQ_LEN - actual_len, dtype=np.int64)
+        pred_dur = np.concatenate([pred_dur_actual, pred_dur_padding])
+        
+        # 调整实际内容部分，只处理长度超出情况
+        fixed_total_frames = self.FIXED_SEQ_LEN * 2
+        diff = fixed_total_frames - pred_dur[:actual_len].sum()
+        
+        if diff < 0:
+            # 减少帧数
+            indices = np.argsort(pred_dur[:actual_len])[::-1]
+            decreased = 0
+            for idx in indices:
+                if pred_dur[idx] > 1 and decreased < abs(diff):
+                    pred_dur[idx] -= 1
+                    decreased += 1
+                if decreased >= abs(diff):
+                    break
+        
+        # 将剩余帧数分配到padding部分
+        remaining_frames = fixed_total_frames - pred_dur[:actual_len].sum()
+        padding_len = self.FIXED_SEQ_LEN - actual_len
+        if remaining_frames > 0 and padding_len > 0:
+            frames_per_padding = remaining_frames // padding_len
+            remainder = remaining_frames % padding_len
+            pred_dur[actual_len:] = frames_per_padding
+            if remainder > 0:
+                pred_dur[actual_len:actual_len+remainder] += 1
+        
+        total_frames = pred_dur.sum()
+        return pred_dur, total_frames
+    
+    def _create_alignment_matrix(self, pred_dur: np.ndarray, total_frames: int) -> np.ndarray:
+        """创建对齐矩阵"""
+        indices = np.repeat(np.arange(self.FIXED_SEQ_LEN), pred_dur)
+        pred_aln_trg = np.zeros((self.FIXED_SEQ_LEN, total_frames), dtype=np.float32)
+        if len(indices) > 0:
+            pred_aln_trg[indices, np.arange(total_frames)] = 1.0
+        return pred_aln_trg[np.newaxis, ...]
+    
+    def _trim_audio_by_content(self, audio: np.ndarray, actual_content_frames: int, 
+                              total_frames: int, actual_len: int) -> np.ndarray:
+        """根据实际内容比例裁剪音频"""
+        padding_len = self.FIXED_SEQ_LEN - actual_len
+        if padding_len > 0:
+            content_ratio = actual_content_frames / total_frames
+            audio_len_to_keep = int(len(audio) * content_ratio)
+            return audio[:audio_len_to_keep]
+        return audio
+    
+    
+    def inference(
+        self,
+        input_ids: np.ndarray,
+        ref_s: np.ndarray,
+        phonemes: str,
+        vocab: dict,
+        speed: float = DEFAULT_SPEED,
+        fade_out_duration: float = DEFAULT_FADE_OUT
+    ) -> np.ndarray:
+        """推理生成音频"""
+        chunks = split_input_ids_semantic(input_ids, self.FIXED_SEQ_LEN)
+        fade_samples = int(SAMPLE_RATE * fade_out_duration) if fade_out_duration > 0 else 0
+        
+        # if len(chunks) == 1:
+        #     # 单个chunk
+        audio, actual_content_frames, total_frames = self.inference_single_chunk(
+            chunks[0]['input_ids'], ref_s, chunks[0]['actual_len'], speed
+        )
+        audio_trimmed = self._trim_audio_by_content(
+            audio, actual_content_frames, total_frames, chunks[0]['actual_len']
+        )
+        if fade_samples > 0:
+            audio_trimmed = apply_fade_out(audio_trimmed, fade_samples)
+        return audio_trimmed
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--axmodel_dir", "-d", type=str, default="onnx")
+    parser.add_argument("--text", "-t", type=str, default="The sky above the port was the color of television, tuned to a dead channel.")
+    parser.add_argument("--lang", "-l", type=str, default='a')
+    parser.add_argument("--voice", "-v", type=str, required=True)
+    parser.add_argument("--config", "-c", type=str, default="checkpoints/config.json")
+    parser.add_argument("--output", "-o", type=str, default="output.wav")
+    parser.add_argument("--fade_out", "-f", type=float, default=0.3)
+    parser.add_argument("--max_len", "-m", type=int, default=96)
+    args = parser.parse_args()
+    
+    SPEED = 1.0 #速度暂时固定
+    PAUSE = 0.0 #不加长停顿
+    
+    start_total = time.time()
+    
+    # 初始化
+    print("初始化模型...")
+    engine = InferenceEngine(args.axmodel_dir)
+    vocab = load_vocab_from_config(args.config)
+    g2p, g2p_type = init_g2p(args.lang)
+    
+    # Warmup g2p模型
+    warmup_text = "你好" if args.lang == 'z' else "hello"
+    _, _ = generate_input_ids_from_text(warmup_text, g2p=g2p, g2p_type=g2p_type, vocab=vocab)
+    init_time = time.time() - start_total
+    print(f"初始化完成: {init_time:.3f}s\n")
+    
+    # 文本处理
+    time_cost = time.time()
+    merged_groups = process_and_merge_sentences(
+        args.text, args.lang, g2p, g2p_type, vocab, max_merge_len=args.max_len
+    )
+    
+    # pdb.set_trace()
+    # 推理
+    audio_list = run_batch_inference(
+        engine, merged_groups, args.voice, vocab, 
+        speed=SPEED, fade_out_duration=args.fade_out
+    )
+    
+    inference_time = time.time() - time_cost
+    
+    # 音频保存
+    final_audio = audio_numpy_concat(audio_list, sr=SAMPLE_RATE, speed=SPEED, pause_duration=PAUSE)
+    sf.write(args.output, final_audio, SAMPLE_RATE)
+    
+    audio_duration = len(final_audio) / SAMPLE_RATE
+    print("\n" + "="*60)
+    print(f"输出: {args.output} | 时长: {audio_duration:.2f}s")
+    print("="*60)
+    print(f"初始化: {init_time:.3f}s")
+    print(f"音频推理: {inference_time:.3f}s (共{engine.inference_count}次)")
+    if engine.inference_count > 0:
+        print(f"  ├─ Model1: {engine.model1_time:.3f}s (平均{engine.model1_time/engine.inference_count*1000:.1f}ms)")
+        print(f"  ├─ Model2: {engine.model2_time:.3f}s (平均{engine.model2_time/engine.inference_count*1000:.1f}ms)")
+        print(f"  ├─ Model3: {engine.model3_time:.3f}s (平均{engine.model3_time/engine.inference_count*1000:.1f}ms)")
+        print(f"  └─ Model4 onnx: {engine.har_time:.3f}s (平均{engine.har_time/engine.inference_count*1000:.1f}ms)")
+    
+    print(f"\n rtf:{inference_time/audio_duration:.3f}")
+    print("="*60)
+
+
+if __name__ == "__main__":
+    main()

inference_utils.py

@@ -160,7 +160,10 @@ def phonemes_to_input_ids(phonemes: str, vocab: Dict[str, int], debug: bool = Fa
 
 
 def load_voice_embedding(voice_path: str, phoneme_len: Optional[int] = None) -> np.ndarray:
-    voice_path = voice_path.replace('checkpoints/voices', 'checkpoints/voices_npy').replace('.pt', '.npy')
+    if "checkpoints/voices_npy" in voice_path:
+        pass
+    else:
+        voice_path = voice_path.replace('checkpoints/voices', 'checkpoints/voices_npy').replace('.pt', '.npy')
     pack = np.load(voice_path).reshape(510,1,256)
     if phoneme_len is not None:
         ref_s = pack[phoneme_len:phoneme_len+1]