Update rwkvtts-respark-webrwkv/tts_cli.py

rwkvtts-respark-webrwkv/tts_cli.py

@@ -10,6 +10,7 @@ import sys
 import re
 import time
 import warnings
+import logging
 from pathlib import Path
 from typing import Dict, Any, Tuple, List
 
@@ -17,6 +18,27 @@ import numpy as np
 import soundfile as sf
 import click
 
+generated_global_tokens = {}
+
+# 配置日志
+def setup_logging():
+    """设置日志配置"""
+    # 从环境变量获取日志级别，默认为WARNING
+    log_level_str = os.environ.get('LOG_LEVEL', 'INFO').upper()
+    log_level = getattr(logging, log_level_str, logging.WARNING)
+    
+    # 配置日志格式
+    logging.basicConfig(
+        level=log_level,
+        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S'
+    )
+    
+    return logging.getLogger(__name__)
+
+# 创建logger实例
+logger = setup_logging()
+
 # 抑制警告
 warnings.filterwarnings("ignore", category=UserWarning, module="numpy")
 warnings.filterwarnings("ignore", category=UserWarning, module="onnxruntime")
@@ -30,8 +52,8 @@ try:
     HAS_WEBRWKV = True
 except ImportError:
     HAS_WEBRWKV = False
-    print("❌ 错误: 需要安装 'webrwkv_py' 库")
-    print("请运行: pip install webrwkv_py")
+    logger.error("❌ 错误: 需要安装 'webrwkv_py' 库")
+    logger.error("请运行: pip install webrwkv_py")
     sys.exit(1)
 
 try:
@@ -39,8 +61,8 @@ try:
     HAS_ONNX = True
 except ImportError:
     HAS_ONNX = False
-    print("❌ 错误: 需要安装 'onnxruntime' 库")
-    print("请运行: pip install onnxruntime")
+    logger.error("❌ 错误: 需要安装 'onnxruntime' 库")
+    logger.error("请运行: pip install onnxruntime")
     sys.exit(1)
 
 try:
@@ -48,8 +70,8 @@ try:
     HAS_TRANSFORMERS = True
 except ImportError:
     HAS_TRANSFORMERS = False
-    print("❌ 错误: 需要安装 'transformers' 库")
-    print("请运行: pip install transformers")
+    logger.error("❌ 错误: 需要安装 'transformers' 库")
+    logger.error("请运行: pip install transformers")
     sys.exit(1)
 
 try:
@@ -57,8 +79,8 @@ try:
     HAS_QUESTIONARY = True
 except ImportError:
     HAS_QUESTIONARY = False
-    print("❌ 错误: 需要安装 'questionary' 库来使用交互式界面")
-    print("请运行: pip install questionary")
+    logger.warning("⚠️  警告: 无法导入 questionary 库来使用交互式界面")
+    logger.warning("请运行: pip install questionary")
     sys.exit(1)
 
 # 导入属性工具
@@ -73,7 +95,7 @@ try:
     pitch_choices = list(PITCH_MAP.keys())
     speed_choices = list(SPEED_MAP.keys())
 except ImportError:
-    print("⚠️  警告: 无法导入 properties_util，使用默认选项")
+    logger.warning("⚠️  警告: 无法导入 properties_util，使用默认选项")
     # 默认选项
     age_choices = ['child', 'teenager', 'youth-adult', 'middle-aged', 'elderly']
     gender_choices = ['female', 'male']  # 与properties_util.py保持一致
@@ -167,26 +189,26 @@ class TTSGenerator:
         self.model_path = model_path
         
         # 初始化 RefAudioUtilities 实例
-        print('🎿 开始加载音频编码器模型')
+        logger.info('🎿 开始加载音频编码器模型')
         try:
             audio_tokenizer_path = os.path.join(model_path, 'BiCodecTokenize.onnx')
             wav2vec2_path = os.path.join(model_path, 'wav2vec2-large-xlsr-53.onnx')
             from ref_audio_utilities import RefAudioUtilities
             self.ref_audio_utilities = RefAudioUtilities(audio_tokenizer_path, wav2vec2_path)
-            print('✅ 音频编码器模型加载成功')
+            logger.info('✅ 音频编码器模型加载成功')
         except Exception as e:
-            print(f'❌ 音频编码器模型加载失败: {e}')
+            logger.error(f'❌ 音频编码器模型加载失败: {e}')
             self.ref_audio_utilities = None
         
         # 缓存ONNX session
-        print('🎿 开始加载ONNX模型')
+        logger.info('🎿 开始加载ONNX模型')
         try:
             self.ort_session = ort.InferenceSession(decoder_path, 
                                                 providers=['CUDAExecutionProvider','CPUExecutionProvider'])
-            print(f"🖥️ONNX Session for generate wavform actual providers: {self.ort_session.get_providers()}")
-            print('✅ ONNX模型加载成功')
+            logger.info(f"🖥️ONNX Session for generate wavform actual providers: {self.ort_session.get_providers()}")
+            logger.info('✅ ONNX模型加载成功')
         except Exception as e:
-            print(f'❌ ONNX模型加载失败: {e}')
+            logger.error(f'❌ ONNX模型加载失败: {e}')
             raise
         
         # 生成统计信息
@@ -213,9 +235,9 @@ class TTSGenerator:
         """重置runtime状态"""
         try:
             self.runtime.reset()
-            print("🔄 Runtime状态已重置")
+            logger.info("🔄 Runtime状态已重置")
         except Exception as e:
-            print(f"⚠️  Runtime重置失败: {e}")
+            logger.warning(f"⚠️  Runtime重置失败: {e}")
     
     def generate_audio(self, params: Dict[str, Any]) -> Tuple[np.ndarray, Dict[str, Any]]:
         """生成音频"""
@@ -233,15 +255,15 @@ class TTSGenerator:
             ref_audio_path = params['ref_audio_path']
             prompt_text = params.get('prompt_text', "希望你以后能够做的，比我还好呦！")
             
-            print(f"🎯 开始生成音频 (Zero Shot 模式): {text}")
-            print(f"📊 参数: 参考音频={ref_audio_path}, 提示文本={prompt_text}")
+            logger.info(f"🎯 开始生成音频 (Zero Shot 模式): {text}")
+            logger.info(f"📊 参数: 参考音频={ref_audio_path}, 提示文本={prompt_text}")
             
             # 检测语言
             lang = detect_token_lang(text)
-            print(f"🌍 检测到语言: {lang}")
+            logger.info(f"🌍 检测到语言: {lang}")
             
             # 使用 zero shot 方法生成 tokens
-            global_tokens, semantic_tokens, global_time, global_speed, semantic_time, semantic_speed = self._generate_tokens_zeroshot(text, ref_audio_path, prompt_text)
+            global_tokens, semantic_tokens, semantic_time, semantic_speed = self._generate_tokens_zeroshot(text, ref_audio_path, prompt_text)
         else:
             # 传统模式
             age = params['age']
@@ -250,46 +272,29 @@ class TTSGenerator:
             pitch = params['pitch']
             speed = params['speed']
             
-            print(f"🎯 开始生成音频: {text}")
-            print(f"📊 参数: 年龄={age}, 性别={gender}, 情感={emotion}, 音高={pitch}, 速度={speed}")
+            logger.info(f"🎯 开始生成音频: {text}")
+            logger.info(f"📊 参数: 年龄={age}, 性别={gender}, 情感={emotion}, 音高={pitch}, 速度={speed}")
             
             # 检测语言
             lang = detect_token_lang(text)
-            print(f"🌍 检测到语言: {lang}")
+            logger.info(f"🌍 检测到语言: {lang}")
             
             # 生成global tokens和semantic tokens
-            global_tokens, semantic_tokens, global_time, global_speed, semantic_time, semantic_speed = self._generate_tokens(text, age, gender, emotion, pitch, speed)
+            generated_key = params['generated_key']
+            global_tokens, semantic_tokens, global_time, global_speed, semantic_time, semantic_speed = self._generate_tokens(text, age, gender, emotion, pitch, speed, generated_key)
         
         # 解码音频
-        print("🎵 解码音频...")
-        decode_start = time.time()
+        logger.info("🎵 解码音频...")
         
-        # 准备输入数据 - 按照tts_gui_simple.py的逻辑
-        print("🔧 准备解码器输入数据...")
-        global_tokens_array = np.array(global_tokens, dtype=np.int64).reshape(1, 1, -1)
-        semantic_tokens_array = np.array(semantic_tokens, dtype=np.int64).reshape(1, -1)
-        print(f'🎯 生成的全局token: {global_tokens}')
-        print(f'🎯 生成的语义token: {semantic_tokens}')
-        print(f'📊 解码器输入形状: global_tokens={global_tokens_array.shape}, semantic_tokens={semantic_tokens_array.shape}')
+        # 使用抽象化的音频解码函数
+        wav_data, audio_duration, decode_time, decode_speed = self._decode_audio(global_tokens, semantic_tokens)
         
-        # 使用ONNX解码器生成音频
-        print("🎵 开始ONNX解码器推理...")
-        outputs = self.ort_session.run(None, {
-                "global_tokens": global_tokens_array, 
-                "semantic_tokens": semantic_tokens_array
-            })
-        wav_data = outputs[0].reshape(-1)
-        decode_time = time.time() - decode_start
-        
-        # 计算音频时长和RTF
-        audio_duration = len(wav_data) / 16000  # 采样率16kHz
-        decode_speed = len(semantic_tokens) / decode_time if decode_time > 0 else 0
+        # 计算总耗时和RTF
         total_time = time.time() - start_time
         total_tokens = len(global_tokens) + len(semantic_tokens)
         rtf = total_time / audio_duration if audio_duration > 0 else 0
         
-        print(f"✅ 音频解码完成，时长 {audio_duration:.2f}s，耗时 {decode_time:.2f}s，速度 {decode_speed:.1f} tokens/s")
-        print(f"📊 总耗时: {total_time:.2f}s，RTF: {rtf:.2f}")
+        logger.info(f"📊 总耗时: {total_time:.2f}s，RTF: {rtf:.2f}")
         
         # 更新统计信息
         self.generation_stats['total_generations'] += 1
@@ -303,7 +308,6 @@ class TTSGenerator:
             'total_tokens': total_tokens,
             'audio_duration': audio_duration,
             'rtf': rtf,
-            'global_speed': global_speed,
             'semantic_speed': semantic_speed,
             'decode_speed': decode_speed,
             'timestamp': time.strftime('%Y-%m-%d %H:%M:%S'),
@@ -312,7 +316,7 @@ class TTSGenerator:
         
         return wav_data, self.generation_stats['last_generation']
     
-    def _generate_tokens(self, text: str, age: str, gender: str, emotion: str, pitch: str, speed: str) -> Tuple[List[int], List[int], float, float, float, float]:
+    def _generate_tokens(self, text: str, age: str, gender: str, emotion: str, pitch: str, speed: str, generated_key: str = None) -> Tuple[List[int], List[int], float, float, float, float]:
         """
         生成global tokens和semantic tokens
         
@@ -323,17 +327,17 @@ class TTSGenerator:
             emotion: 情感参数
             pitch: 音高参数
             speed: 速度参数
-            
+            generated_key: 之前生成的全局token的key
         Returns:
             Tuple: (global_tokens, semantic_tokens, global_time, global_speed, semantic_time, semantic_speed)
         """
         # 编码文本
-        print("🔤 编码文本...")
+        logger.info("🔤 编码文本...")
         tokens = self.tokenizer.encode(text)
-        print(f"✅ 文本编码完成，共 {len(tokens)} 个token")
+        logger.info(f"✅ 文本编码完成，共 {len(tokens)} 个token")
         
         # 生成全局token
-        print("🌐 生成全局token...")
+        logger.info("🌐 生成全局token...")
         global_start = time.time()
         
         # 准备输入tokens
@@ -344,7 +348,7 @@ class TTSGenerator:
         # 构建属性tokens - 使用properties_util.py
         from properties_util import convert_standard_properties_to_tokens
         properties_text = convert_standard_properties_to_tokens(age, gender, emotion, pitch, speed)
-        print(f'🔤 属性文本: {properties_text}')
+        logger.info(f'🔤 属性文本: {properties_text}')
         properties_tokens = self.tokenizer.encode(properties_text, add_special_tokens=False)
         properties_tokens = [i + 8196 + 4096 for i in properties_tokens]
         
@@ -352,36 +356,70 @@ class TTSGenerator:
         text_tokens = [i + 8196 + 4096 for i in tokens]
         
         # 组合所有tokens
-        all_idx = properties_tokens + [TTS_TAG_2] + text_tokens + [TTS_TAG_0]
-        print(f'🔢 属性token: {properties_tokens}')
-        print(f'🔢 文本token: {text_tokens}')
-        print(f'🎯 组合后的tokens: {all_idx}')
+        if generated_key is None or generated_key not in generated_global_tokens:
+            all_idx = properties_tokens + [TTS_TAG_2] + text_tokens + [TTS_TAG_0]
+        else:
+            logger.info(f"🎯 使用之前生成的全局token: {generated_key}")
+            previous_global_tokens = generated_global_tokens[generated_key]
+            global_tokens = previous_global_tokens.copy()
+            global_time = 0
+            global_speed = 0
+            logger.info(f"🎯 使用之前生成的全局token: {previous_global_tokens}")
+            previous_global_tokens = [int(i) + 8196 for i in previous_global_tokens]
+            logger.info(f"🎯 偏移后的全局token: {previous_global_tokens}")
+            all_idx = properties_tokens + [TTS_TAG_2] + text_tokens + [TTS_TAG_0] + previous_global_tokens 
+        logger.info(f'🔢 属性token: {properties_tokens}')
+        logger.info(f'🔢 文本token: {text_tokens}')
+        logger.info(f'🎯 组合后的tokens: {all_idx}')
         
         # Prefill阶段
-        print("💎 开始Prefill阶段...")
-        logits = self.runtime.predict(all_idx)
-        print(f"✅ Prefill完成，logits长度: {len(logits)}")
+        logger.info("💎 开始Prefill阶段...")
+        session = self.runtime.create_inference_session([all_idx],token_chunk_size=512)
+        step_count = 0
+        start = time.time()
+        while not session.is_complete():
+            step_count += 1
+            output = session.step()
+            if not output.batches[0].is_empty():
+                logits = output.batches[0].data
+                break
+        
+        prefill_time = time.time() - start
+        logger.info(f"✅ Prefill完成，耗时 {step_count} 步")
+        logger.info(f"✅ Prefill完成，logits长度: {len(logits)}")
+        logger.info(f"✅ Prefill完成，耗时 {prefill_time:.2f}s {len(all_idx)/prefill_time:.1f} tokens/s")
         
         # 生成全局token - 按照tts_gui_simple.py的逻辑
-        print("🌍 开始生成全局token...")
-        global_tokens_size = 32
-        global_tokens = []
-        
-        for i in range(global_tokens_size):
-            # 从logits中采样token
-            sampled_id = sample_logits(logits[0:4096], temperature=1.0, top_p=0.95, top_k=20)
-            global_tokens.append(sampled_id)
-            # 预测下一个token
-            sampled_id += 8196
-            logits = self.runtime.predict_next(sampled_id)
-        
-        global_time = time.time() - global_start
-        global_speed = global_tokens_size / global_time if global_time > 0 else 0
-        print(f"✅ 全局token生成完成，共 {len(global_tokens)} 个token，耗时 {global_time:.2f}s，速度 {global_speed:.1f} tokens/s")
-        print(f'🎯 生成的全局token: {global_tokens}')
+        if generated_key is None or generated_key not in generated_global_tokens:
+            logger.info("🌍 开始生成全局token...")
+            global_tokens_size = 32
+            global_tokens = []
+            
+            for i in range(global_tokens_size):
+                # 从logits中采样token
+                sampled_id = sample_logits(logits[0:4096], temperature=1.0, top_p=0.95, top_k=20)
+                global_tokens.append(sampled_id)
+                # 预测下一个token
+                sampled_id += 8196
+                logits = self.runtime.predict_next(sampled_id)
+            
+            global_time = time.time() - global_start
+            global_speed = global_tokens_size / global_time if global_time > 0 else 0
+            logger.info(f"✅ 全局token生成完成，共 {len(global_tokens)} 个token，耗时 {global_time:.2f}s，速度 {global_speed:.1f} tokens/s")
+            logger.info(f'🎯 生成的全局token: {global_tokens}')
+            prefix = f"{age}_{gender}_{pitch}_{emotion}_{speed}"
+            key = f"{prefix}_0"
+            if key in generated_global_tokens:
+                #found the latest index of the same key
+                latest_index = max([int(k.split('_')[-1]) for k in generated_global_tokens.keys() if k.startswith(prefix)])
+                key = f"{prefix}_{latest_index + 1}"
+            generated_global_tokens[key] = global_tokens
+            logger.info(f'🎯 生成的全局token: {generated_global_tokens[key]}, 下次可以调用generated_global_tokens[{key}]')
+
+
         
         # 生成语义token
-        print("🧠 生成语义token...")
+        logger.info("🧠 生成语义token...")
         semantic_start = time.time()
         
         # 按照tts_gui_simple.py的逻辑生成语义token
@@ -391,17 +429,78 @@ class TTSGenerator:
         for i in range(2048):  # 最大生成2048个token
             sampled_id = sample_logits(x[0:8193], temperature=1.0, top_p=0.95, top_k=80)
             if sampled_id == 8192:  # 遇到结束标记
-                print(f"🛑 语义token生成结束，遇到结束标记，共生成 {len(semantic_tokens)} 个token")
+                logger.info(f"🛑 语义token生成结束，遇到结束标记，共生成 {len(semantic_tokens)} 个token")
                 break
             semantic_tokens.append(sampled_id)
             x = self.runtime.predict_next(sampled_id)
         
         semantic_time = time.time() - semantic_start
         semantic_speed = len(semantic_tokens) / semantic_time if semantic_time > 0 else 0
-        print(f"✅ 语义token生成完成，共 {len(semantic_tokens)} 个token，耗时 {semantic_time:.2f}s，速度 {semantic_speed:.1f} tokens/s")
+        logger.info(f"✅ 语义token生成完成，共 {len(semantic_tokens)} 个token，耗时 {semantic_time:.2f}s，速度 {semantic_speed:.1f} tokens/s")
         
         return global_tokens, semantic_tokens, global_time, global_speed, semantic_time, semantic_speed
+    
+    def _generate_tokens_with_global_tokens(self, text: str, global_tokens: List[int]) -> Tuple[List[int], List[int], float, float, float, float]:
+        """
+        使用 global tokens 生成语义token
+        """
+        # 编码文本
+        logger.info("🔤 编码文本...")
+        text_tokens = self.tokenizer.encode(text, add_special_tokens=False)
+        text_tokens = [i + 8196 + 4096 for i in text_tokens]
+        logger.info(f"✅ 文本编码完成，共 {len(text_tokens)} 个token")
+        global_tokens = [int(i) + 8196 for i in global_tokens]
+        logger.info(f'🎯 参考音频 global_tokens: {global_tokens}')
+        start = time.time()
+        
+        # 准备输入tokens
+        TTS_TAG_0 = 8193
+        TTS_TAG_1 = 8194
+        TTS_TAG_2 = 8195
+        
+        # 组合所有tokens
+        all_idx = [TTS_TAG_2] + text_tokens + [TTS_TAG_0] + global_tokens + [TTS_TAG_1] 
+        logger.info(f'🎯 组合后的tokens: {all_idx}')
+        
+        # Prefill阶段
+        logger.info("💎 开始Prefill阶段...")
+        session = self.runtime.create_inference_session([all_idx],token_chunk_size=512)
+        step_count = 0
+        while not session.is_complete():
+            step_count += 1
+            output = session.step()
+            if not output.batches[0].is_empty():
+                logits = output.batches[0].data[0]
+                break
+        logger.info(f"✅ Prefill完成，耗时 {step_count} 步")
+        logger.info(f"✅ Prefill完成，速度 {step_count/output.time:.1f} tokens/s")
+        logger.info(f"✅ Prefill完成，logits长度: {len(logits)}")
+        prefill_time = time.time() - start
+        prefill_speed = len(all_idx) / prefill_time if prefill_time > 0 else 0
+        logger.info(f"✅ Prefill完成，耗时 {prefill_time:.2f}s，速度 {prefill_speed:.1f} tokens/s")
+        
+        # 生成语义token
+        logger.info("🧠 生成语义token...")
+        semantic_start = time.time()
 
+        # 从当前logits开始生成语义token
+        x = logits
+        semantic_tokens = []
+        
+        for i in range(2048):  # 最大生成2048个token
+            sampled_id = sample_logits(x[0:8193], temperature=1.0, top_p=0.95, top_k=80)
+            if sampled_id == 8192:  # 遇到结束标记
+                logger.info(f"🛑 语义token生成结束，遇到结束标记，共生成 {len(semantic_tokens)} 个token")
+                break
+            semantic_tokens.append(sampled_id)
+            x = self.runtime.predict_next(sampled_id)
+        
+        semantic_time = time.time() - semantic_start
+        semantic_speed = len(semantic_tokens) / semantic_time if semantic_time > 0 else 0
+        logger.info(f"✅ 语义token生成完成，共 {len(semantic_tokens)} 个token，耗时 {semantic_time:.2f}s，速度 {semantic_speed:.1f} tokens/s")
+        
+        return global_tokens, semantic_tokens, prefill_time, prefill_speed, semantic_time, semantic_speed
+    
     def _generate_tokens_zeroshot(self, text: str, ref_audio_path: str, prompt_text: str = "希望你以后能够做的，比我还好呦！") -> Tuple[List[int], List[int], float, float, float, float]:
         """
         使用 zero shot 方式生成global tokens和semantic tokens
@@ -418,26 +517,24 @@ class TTSGenerator:
             raise RuntimeError("RefAudioUtilities 未初始化，无法使用 zero shot 模式")
         
         # 编码文本
-        print("🔤 编码文本...")
+        logger.info("🔤 编码文本...")
         text_tokens = self.tokenizer.encode(prompt_text + text, add_special_tokens=False)
         text_tokens = [i + 8196 + 4096 for i in text_tokens]
-        print(f"✅ 文本编码完成，共 {len(text_tokens)} 个token")
+        logger.info(f"✅ 文本编码完成，共 {len(text_tokens)} 个token")
         
         # 从参考音频获取 global tokens 和 semantic tokens
-        print("🎵 处理参考音频...")
+        logger.info("🎵 处理参考音频...")
         global_tokens, prompt_semantic_tokens = self.ref_audio_utilities.tokenize(ref_audio_path)
-        print(f"✅ 参考音频处理完成")
+        logger.info(f"✅ 参考音频处理完成")
         
         # 直接使用flatten()展平数组并转换为Python一维数组
         global_tokens = [int(i) + 8196 for i in global_tokens.flatten()]
         prompt_semantic_tokens = [int(i) for i in prompt_semantic_tokens.flatten()]
         
-        print(f'🎯 参考音频 global_tokens: {global_tokens}')
-        print(f'🎯 参考音频 semantic_tokens: {prompt_semantic_tokens}')
+        logger.info(f'🎯 参考音频 global_tokens: {global_tokens}')
+        logger.info(f'🎯 参考音频 semantic_tokens: {prompt_semantic_tokens}')
+        
         
-        # 生成全局token
-        print("🌐 生成全局token...")
-        global_start = time.time()
         
         # 准备输入tokens
         TTS_TAG_0 = 8193
@@ -446,19 +543,27 @@ class TTSGenerator:
         
         # 组合所有tokens
         all_idx = [TTS_TAG_2] + text_tokens + [TTS_TAG_0] + global_tokens + [TTS_TAG_1] + prompt_semantic_tokens
-        print(f'🎯 组合后的tokens: {all_idx}')
+        logger.info(f'🎯 组合后的tokens: {all_idx}')
         
         # Prefill阶段
-        print("💎 开始Prefill阶段...")
-        logits = self.runtime.predict(all_idx)
-        print(f"✅ Prefill完成，logits长度: {len(logits)}")
-        
-        global_time = time.time() - global_start
-        global_speed = len(global_tokens) / global_time if global_time > 0 else 0
-        print(f"✅ 全局token处理完成，共 {len(global_tokens)} 个token，耗时 {global_time:.2f}s，速度 {global_speed:.1f} tokens/s")
+        logger.info("💎 开始Prefill阶段...")
+        session = self.runtime.create_inference_session([all_idx],token_chunk_size=512)
+        step_count = 0
+        start = time.time()
+        while not session.is_complete():
+            step_count += 1
+            output = session.step()
+            if not output.batches[0].is_empty():
+                logits = output.batches[0].data
+                break
+        prefill_time = time.time() - start
+        logger.info(f"✅ Prefill完成，logits长度: {len(logits)}")
+        logger.info(f"✅ Prefill完成，耗时 {step_count} 步")
+        logger.info(f"✅ Prefill完成，耗时 {prefill_time:.2f}s {len(all_idx)/prefill_time:.1f} tokens/s")
+
         
         # 生成语义token
-        print("🧠 生成语义token...")
+        logger.info("🧠 生成语义token...")
         semantic_start = time.time()
         
         # 从当前logits开始生成语义token
@@ -468,52 +573,110 @@ class TTSGenerator:
         for i in range(2048):  # 最大生成2048个token
             sampled_id = sample_logits(x[0:8193], temperature=1.0, top_p=0.95, top_k=80)
             if sampled_id == 8192:  # 遇到结束标记
-                print(f"🛑 语义token生成结束，遇到结束标记，共生成 {len(semantic_tokens)} 个token")
+                logger.info(f"🛑 语义token生成结束，遇到结束标记，共生成 {len(semantic_tokens)} 个token")
                 break
             semantic_tokens.append(sampled_id)
             x = self.runtime.predict_next(sampled_id)
         
         semantic_time = time.time() - semantic_start
         semantic_speed = len(semantic_tokens) / semantic_time if semantic_time > 0 else 0
-        print(f"✅ 语义token生成完成，共 {len(semantic_tokens)} 个token，耗时 {semantic_time:.2f}s，速度 {semantic_speed:.1f} tokens/s")
+        logger.info(f"✅ 语义token生成完成，共 {len(semantic_tokens)} 个token，耗时 {semantic_time:.2f}s，速度 {semantic_speed:.1f} tokens/s")
         
         global_tokens = [i - 8196 for i in global_tokens]
-        return global_tokens, semantic_tokens, global_time, global_speed, semantic_time, semantic_speed
+        return global_tokens, semantic_tokens, semantic_time, semantic_speed
+
+    def _decode_audio(self, global_tokens: List[int], semantic_tokens: List[int]) -> Tuple[np.ndarray, float, float, float]:
+        """
+        解码音频的核心函数
+        
+        Args:
+            global_tokens: 全局tokens列表
+            semantic_tokens: 语义tokens列表
+            
+        Returns:
+            Tuple: (wav_data, audio_duration, decode_time, decode_speed)
+        """
+        # 开始计时
+        decode_start = time.time()
+        
+        # 准备输入数据
+        logger.info("🔧 准备解码器输入数据...")
+        global_tokens_array = np.array(global_tokens, dtype=np.int64).reshape(1, 1, -1)
+        semantic_tokens_array = np.array(semantic_tokens, dtype=np.int64).reshape(1, -1)
+        logger.info(f'🎯 生成的全局token: {global_tokens}')
+        logger.info(f'🎯 生成的语义token: {semantic_tokens}')
+        logger.info(f'📊 解码器输入形状: global_tokens={global_tokens_array.shape}, semantic_tokens={semantic_tokens_array.shape}')
+        
+        # 使用ONNX解码器生成音频
+        logger.info("🎵 开始ONNX解码器推理...")
+        outputs = self.ort_session.run(None, {
+                "global_tokens": global_tokens_array, 
+                "semantic_tokens": semantic_tokens_array
+            })
+        wav_data = outputs[0].reshape(-1)
+        decode_time = time.time() - decode_start
+        
+        # 计算音频时长和解码速度
+        audio_duration = len(wav_data) / 16000  # 采样率16kHz
+        decode_speed = len(semantic_tokens) / decode_time if decode_time > 0 else 0
+        
+        logger.info(f"✅ 音频解码完成，时长 {audio_duration:.2f}s，耗时 {decode_time:.2f}s，速度 {decode_speed:.1f} tokens/s")
+        
+        return wav_data, audio_duration, decode_time, decode_speed
+
+    def _save_audio(self, wav_data: np.ndarray, output_path: str, sample_rate: int = 16000) -> bool:
+        """
+        保存音频文件
+        
+        Args:
+            wav_data: 音频数据
+            output_path: 输出文件路径
+            sample_rate: 采样率，默认16kHz
+            
+        Returns:
+            bool: 保存是否成功
+        """
+        try:
+            sf.write(output_path, wav_data, sample_rate)
+            logger.info(f"💾 音频保存成功: {output_path}")
+            return True
+        except Exception as e:
+            logger.error(f"❌ 音频保存失败: {e}")
+            return False
 
 def display_stats(stats: Dict[str, Any]):
     """显示生成统计信息"""
-    print("\n" + "="*60)
-    print("📊 生成统计信息")
-    print("="*60)
+    logger.info("\n" + "="*60)
+    logger.info("📊 生成统计信息")
+    logger.info("="*60)
     
     if stats['text']:
-        print(f"🎯 生成参数: {stats['params']}")
-        print(f"📝 文本: {stats['text']}")
-        print(f"⏱️  总耗时: {stats['total_time']:.2f}s")
-        print(f"🎵 音频时长: {stats['audio_duration']:.2f}s")
-        print(f"📈 RTF: {stats['rtf']:.2f}")
-        print(f"🔢 总token数: {stats['total_tokens']}")
-        print(f"🌐 全局token速度: {stats['global_speed']:.1f} tokens/s")
-        print(f"🧠 语义token速度: {stats['semantic_speed']:.1f} tokens/s")
-        print(f"🎵 解码速度: {stats['decode_speed']:.1f} tokens/s")
-        print(f"🕐 时间: {stats['timestamp']}")
+        logger.info(f"🎯 生成参数: {stats['params']}")
+        logger.info(f"📝 文本: {stats['text']}")
+        logger.info(f"⏱️  总耗时: {stats['total_time']:.2f}s")
+        logger.info(f"🎵 音频时长: {stats['audio_duration']:.2f}s")
+        logger.info(f"📈 RTF: {stats['rtf']:.2f}")
+        logger.info(f"🔢 总token数: {stats['total_tokens']}")
+        logger.info(f"🧠 语义token速度: {stats['semantic_speed']:.1f} tokens/s")
+        logger.info(f"🎵 解码速度: {stats['decode_speed']:.1f} tokens/s")
+        logger.info(f"🕐 时间: {stats['timestamp']}")
         if stats['output_path']:
-            print(f"💾 保存路径: {stats['output_path']}")
+            logger.info(f"💾 保存路径: {stats['output_path']}")
     else:
-        print("暂无生成记录")
+        logger.info("暂无生成记录")
     
-    print("="*60)
+    logger.info("="*60)
 
 def interactive_parameter_selection(generator: TTSGenerator):
     """交互式参数选择界面"""
-    print("\n🎮 进入交互式配置界面")
-    print("💡 使用方向键选择，回车确认，Ctrl+C退出")
+    logger.info("\n🎮 进入交互式配置界面")
+    logger.info("💡 使用方向键选择，回车确认，Ctrl+C退出")
     
     while True:
         try:
-            print("\n" + "="*60)
-            print("🎵 RWKV TTS 参数配置")
-            print("="*60)
+            logger.info("\n" + "="*60)
+            logger.info("🎵 RWKV TTS 参数配置")
+            logger.info("="*60)
             
             # 选择生成模式
             generation_mode = questionary.select(
@@ -596,16 +759,18 @@ def interactive_parameter_selection(generator: TTSGenerator):
                         output_path = get_unique_filename(output_dir, text)
                         
                         # 保存音频
-                        sf.write(output_path, wav_data, 16000)
-                        stats['output_path'] = output_path
+                        if generator._save_audio(wav_data, output_path, 16000):
+                            stats['output_path'] = output_path
+                        else:
+                            logger.warning("⚠️ 音频保存失败，但生成统计已更新")
                         
-                        print(f"✅ 音频生成成功，保存至: {output_path}")
+                        logger.info(f"✅ 音频生成成功，保存至: {output_path}")
                         stats['生成参数'] = f'参考音频={ref_audio_path}, 提示文本={prompt_text}'
                         # 显示统计信息
                         display_stats(stats)
                         
                     except Exception as e:
-                        print(f"❌ 生成失败: {e}")
+                        logger.error(f"❌ 生成失败: {e}")
                         import traceback
                         traceback.print_exc()
             else:
@@ -659,6 +824,20 @@ def interactive_parameter_selection(generator: TTSGenerator):
                 
                 if speed is None:
                     break
+                prefix = f"{age}_{gender}"
+                list_of_generated_keys = []
+                for generated_key in generated_global_tokens.keys():
+                    if generated_key.startswith(prefix):
+                        list_of_generated_keys.append(generated_key)
+                if len(list_of_generated_keys) > 0:
+                    list_of_generated_keys.append("None")
+                    generated_key = questionary.select(
+                        "🎯 是否使用之前生成的全局token?",
+                        choices=list_of_generated_keys,
+                        default="None"
+                    ).ask()
+                else:
+                    generated_key = None
                 
              
                 # 确认生成
@@ -666,7 +845,8 @@ def interactive_parameter_selection(generator: TTSGenerator):
                     f"🚀 确认生成音频?\n"
                     f"文本: {text}\n"
                     f"参数: 年龄={age}, 性别={gender}, 情感={emotion}, 音高={pitch}, 速度={speed}\n"
-                    f"输出目录: {output_dir}",
+                    f"输出目录: {output_dir}\n"
+                    f"是否使用之前生成的全局token: {generated_key is not None}",
                     default=True
                 ).ask()
                 
@@ -680,7 +860,8 @@ def interactive_parameter_selection(generator: TTSGenerator):
                         'emotion': emotion,
                         'pitch': pitch,
                         'speed': speed,
-                        'output_dir': output_dir
+                        'output_dir': output_dir,
+                        'generated_key': generated_key
                     }
                     
                     # 生成音频
@@ -691,16 +872,18 @@ def interactive_parameter_selection(generator: TTSGenerator):
                         output_path = get_unique_filename(output_dir, text)
                         
                         # 保存音频
-                        sf.write(output_path, wav_data, 16000)
-                        stats['output_path'] = output_path
+                        if generator._save_audio(wav_data, output_path, 16000):
+                            stats['output_path'] = output_path
+                        else:
+                            logger.warning("⚠️ 音频保存失败，但生成统计已更新")
                         
-                        print(f"✅ 音频生成成功，保存至: {output_path}")
+                        logger.info(f"✅ 音频生成成功，保存至: {output_path}")
                         stats['生成参数'] = f'年龄={age}, 性别={gender}, 情感={emotion}, 音高={pitch}, 速度={speed}'
                         # 显示统计信息
                         display_stats(stats)
                         
                     except Exception as e:
-                        print(f"❌ 生成失败: {e}")
+                        logger.error(f"❌ 生成失败: {e}")
                         import traceback
                         traceback.print_exc()
             
@@ -714,61 +897,57 @@ def interactive_parameter_selection(generator: TTSGenerator):
                 break
                 
         except KeyboardInterrupt:
-            print("\n👋 用户中断，退出程序")
+            logger.info("\n👋 用户中断，退出程序")
             break
         except Exception as e:
-            print(f"❌ 发生错误: {e}")
+            logger.error(f"❌ 发生错误: {e}")
             import traceback
             traceback.print_exc()
             break
     
-    print("👋 感谢使用 RWKV TTS!")
+    logger.info("👋 感谢使用 RWKV TTS!")
 
 @click.command()
 @click.option('--model_path', required=True, help='RWKV模型路径')
 def main(model_path):
     """RWKV TTS 主程序"""
-    print("🚀 欢迎使用 RWKV TTS 交互式音频生成工具!")
+    logger.info("🚀 欢迎使用 RWKV TTS 交互式音频生成工具!")
     
     # 检查模型文件
     if not os.path.exists(model_path):
-        print(f"❌ 错误: 模型路径不存在: {model_path}")
+        logger.error(f"❌ 错误: 模型路径不存在: {model_path}")
         return
     
     # 自动构建解码器路径
     decoder_path = os.path.join(model_path, "BiCodecDetokenize.onnx")
-    print(f"🔍 自动设置解码器路径: {decoder_path}")
+    logger.info(f"🔍 自动设置解码器路径: {decoder_path}")
     
     # 检查模型目录中的文件
-    print(f"🔍 检查模型目录: {model_path}")
+    logger.info(f"🔍 检查模型目录: {model_path}")
     try:
         model_files = os.listdir(model_path)
-        print(f"📁 模型目录中的文件:")
+        logger.info(f"📁 模型目录中的文件:")
         for file in model_files:
             file_path = os.path.join(model_path, file)
             if os.path.isfile(file_path):
                 size = os.path.getsize(file_path)
-                print(f"   📄 {file} ({size:,} bytes)")
+                logger.info(f"   📄 {file} ({size:,} bytes)")
             else:
-                print(f"   📁 {file}/")
+                logger.info(f"   📁 {file}/")
     except Exception as e:
-        print(f"⚠️  无法列出模型目录内容: {e}")
+        logger.warning(f"⚠️  无法列出模型目录内容: {e}")
     
     if not os.path.exists(decoder_path):
-        print(f"❌ 错误: 解码器路径不存在: {decoder_path}")
+        logger.error(f"❌ 错误: 解码器路径不存在: {decoder_path}")
         return
     
     # 选择设备
-    print("\n💎 选择设备 💎")
+    logger.info("\n💎 选择设备 💎")
     try:
         devices = webrwkv_py.get_available_adapters_py()
-    except AttributeError:
-        # 如果新API不存在，尝试旧API
-        try:
-            devices = webrwkv_py.get_available_devices()
-        except AttributeError:
-            print("❌ 无法获取可用设备列表")
-            return
+    except Exception as e:
+        logger.error(f"❌ 无法获取可用设备列表: {e}")
+        return
     
     for i, device in enumerate(devices):
         print(f"{i}: {device}")
@@ -777,16 +956,16 @@ def main(model_path):
     try:
         device_idx = int(device_choice)
         if device_idx < 0 or device_idx >= len(devices):
-            print("❌ 无效的设备选择")
+            logger.error("❌ 无效的设备选择")
             return
         device = devices[device_idx]
-        print(f"✅ 选择设备: {device}")
+        logger.info(f"✅ 选择设备: {device}")
     except ValueError:
-        print("❌ 无效的设备选择")
+        logger.error("❌ 无效的设备选择")
         return
     
     # 加载模型
-    print("\n💎 加载模型 💎")
+    logger.info("\n💎 加载模型 💎")
     try:
         # 尝试多种可能的模型文件名
         possible_model_files = [
@@ -798,55 +977,56 @@ def main(model_path):
             test_path = os.path.join(model_path, model_file)
             if os.path.exists(test_path):
                 webrwkv_model_path = test_path
-                print(f"✅ 找到模型文件: {model_file}")
+                logger.info(f"✅ 找到模型文件: {model_file}")
                 break
         
         if webrwkv_model_path is None:
-            print(f"❌ 未找到模型文件")
-            print(f"💡 请检查模型目录 {model_path} 中是否包含以下文件之一:")
+            logger.error(f"❌ 未找到模型文件")
+            logger.info(f"💡 请检查模型目录 {model_path} 中是否包含以下文件之一:")
             for model_file in possible_model_files:
-                print(f"   - {model_file}")
+                logger.info(f"   - {model_file}")
             return
         
-        print(f"🔍 尝试加载模型文件: {webrwkv_model_path}")
+        logger.info(f"🔍 尝试加载模型文件: {webrwkv_model_path}")
         
         # 尝试新的API
         model = webrwkv_py.Model(webrwkv_model_path, 'fp32', device_idx)
-        print(f"✅ 模型加载成功: {webrwkv_model_path}")
+        logger.info(f"✅ 模型加载成功: {webrwkv_model_path}")
     except Exception as e:
-        print(f"❌ 模型加载失败: {e}")
-        print(f"💡 请检查:")
-        print(f"   1. 模型文件路径是否正确: {webrwkv_model_path}")
-        print(f"   2. 模型文件是否完整")
-        print(f"   3. 设备索引是否正确: {device_idx}")
-        print(f"   4. 模型文件格式是否支持")
+        logger.error(f"❌ 模型加载失败: {e}")
+        logger.info(f"💡 请检查:")
+        logger.info(f"   1. 模型文件路径是否正确: {webrwkv_model_path}")
+        logger.info(f"   2. 模型文件是否完整")
+        logger.info(f"   3. 设备索引是否正确: {device_idx}")
+        logger.info(f"   4. 模型文件格式是否支持")
         return
     
     # 创建runtime
-    print("\n💎 创建 runtime 💎")
+    logger.info("\n💎 创建 runtime 💎")
     try:
         runtime = model.create_thread_runtime()
-        print("✅ runtime 创建成功")
+        logger.info("✅ runtime 创建成功")
     except Exception as e:
-        print(f"❌ runtime 创建失败: {e}")
+        logger.error(f"❌ runtime 创建失败: {e}")
         return
     
     # 加载tokenizer
-    print("\n💎 加载 tokenizer 💎")
+    logger.info("\n💎 加载 tokenizer 💎")
     try:
         tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-        print(f"✅ tokenizer 加载成功: {model_path}")
+        logger.info(f"✅ tokenizer 加载成功: {model_path}")
     except Exception as e:
-        print(f"❌ tokenizer 加载失败: {e}")
-        print(f"💡 请检查模型目录 {model_path} 中是否包含正确的tokenizer文件")
+        logger.error(f"❌ tokenizer 加载失败: {e}")
+        logger.info(f"💡 请检查模型目录 {model_path} 中是否包含正确的tokenizer文件")
         return
     
     # 创建TTS生成器
     generator = TTSGenerator(runtime, tokenizer, decoder_path, device, model_path)
     
     # 启动交互式界面
-    print("\n🎯 启动交互式配置界面...")
+    logger.info("\n🎯 启动交互式配置界面...")
     interactive_parameter_selection(generator)
 
 if __name__ == "__main__":
     main()
+