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
2noise · Jul 16, 2024 · e609093 · e609093
1 parent 3201c24
commit e609093
Showing 1 changed file with 6 additions and 1 deletion.
diff --git 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: