From e6090938c380749564714c3187e84f45accbb763 Mon Sep 17 00:00:00 2001
From: zhipeng <5310853+Ox0400@users.noreply.github.com>
Date: Tue, 16 Jul 2024 20:20:06 +0800
Subject: [PATCH] fix: stream mode, last chunk of audio file size (#564)

* fix stream mode, last chunk of audio file size

* Update core.py

* Update core.py

* Update core.py
---
 ChatTTS/core.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/ChatTTS/core.py b/ChatTTS/core.py
index 4ab901659..a6195acb6 100644
--- a/ChatTTS/core.py
+++ b/ChatTTS/core.py
@@ -415,7 +415,12 @@ def _infer(
             else:
                 yield wavs
         if stream:
-            yield wavs[:, length:]
+            new_wavs = wavs[:, length:]
+            # Identify rows with non-zero elements using np.any
+            # keep_rows = np.any(array != 0, axis=1)
+            keep_cols = np.sum(new_wavs != 0, axis=0) > 0
+            # Filter both rows and columns using slicing
+            yield new_wavs[:][:, keep_cols]
 
     @torch.inference_mode()
     def _vocos_decode(self, spec: torch.Tensor) -> np.ndarray: