fix: add retry mechanism to prevent EOS early termination sentence dropping

genie_tts/Core/Inference.py

@@ -115,8 +115,13 @@ class GENIE:
             first_stage_decoder: ort.InferenceSession,
             stage_decoder: ort.InferenceSession,
     ) -> Optional[np.ndarray]:
-        """在CPU上运行T2S模型"""
-        # Encoder
+        """在CPU上运行T2S模型，带重试机制防止 EOS 过早终止"""
+        
+        # 动态阈值：最小期望 tokens 数量（参考 AstraTTS）
+        min_expected_tokens = max(8, text_seq.shape[-1] * 2)
+        max_retries = 5
+        
+        # Encoder 只需运行一次
         x, prompts = encoder.run(
             None,
             {
@@ -127,30 +132,53 @@ class GENIE:
                 "ssl_content": ssl_content,
             },
         )
-
-        # First Stage Decoder
-        y, y_emb, *present_key_values = first_stage_decoder.run(
-            None, {"x": x, "prompts": prompts}
-        )
-
-        # Stage Decoder
+        
         input_names: List[str] = [inp.name for inp in stage_decoder.get_inputs()]
-        idx: int = 0
-        for idx in range(0, 500):
+        best_y = None
+        best_idx = 0
+        
+        for retry in range(max_retries):
             if self.stop_event.is_set():
                 return None
-            input_feed = {
-                name: data
-                for name, data in zip(input_names, [y, y_emb, *present_key_values])
-            }
-            outputs = stage_decoder.run(None, input_feed)
-            y, y_emb, stop_condition_tensor, *present_key_values = outputs
-
-            if stop_condition_tensor:
-                break
-
-        y[0, -1] = 0
-        return np.expand_dims(y[:, -idx:], axis=0)
+            
+            # First Stage Decoder（每次重试都重新运行以获取新的随机采样状态）
+            y, y_emb, *present_key_values = first_stage_decoder.run(
+                None, {"x": x, "prompts": prompts}
+            )
+            
+            # Stage Decoder Loop
+            idx: int = 0
+            for idx in range(0, 500):
+                if self.stop_event.is_set():
+                    return None
+                input_feed = {
+                    name: data
+                    for name, data in zip(input_names, [y, y_emb, *present_key_values])
+                }
+                outputs = stage_decoder.run(None, input_feed)
+                y, y_emb, stop_condition_tensor, *present_key_values = outputs
+
+                if stop_condition_tensor:
+                    break
+            
+            # 保存最佳结果（tokens 数量最多的）
+            if idx > best_idx:
+                best_idx = idx
+                best_y = y.copy()
+            
+            # 验证生成数量是否达到预期
+            if idx >= min_expected_tokens:
+                break  # 成功，退出重试循环
+            
+            # 否则继续重试
+        
+        # 使用最佳结果
+        if best_y is None:
+            best_y = y
+            best_idx = idx
+        
+        best_y[0, -1] = 0
+        return np.expand_dims(best_y[:, -best_idx:], axis=0)
 
 
 tts_client: GENIE = GENIE()