optimize(tools): drop unsafe convert

ChatTTS/core.py

@@ -194,7 +194,6 @@ class Chat:
         min_new_token: int = 0
         show_tqdm: bool = True
         ensure_non_empty: bool = True
-        stream_batch: int = 24
 
     @dataclass(repr=False, eq=False)
     class InferCodeParams(RefineTextParams):
@@ -203,6 +202,7 @@ class Chat:
         temperature: float = 0.3
         repetition_penalty: float = 1.05
         max_new_token: int = 2048
+        stream_batch: int = 24
 
     def infer(
         self,
@@ -374,7 +374,7 @@ class Chat:
                 yield text
                 return
 
-        length = [0 for _ in range(len(text))]
+        length = np.zeros(len(text), dtype=np.uint16)
         for result in self._infer_code(
             text,
             stream,
@@ -382,9 +382,11 @@ class Chat:
             use_decoder,
             params_infer_code,
         ):
-            wav = self._decode_to_wavs(result, length, use_decoder)
+            wavs = self._decode_to_wavs(
+                result, length, use_decoder,
+            )
             result.destroy()
-            yield wav
+            yield wavs
 
     @torch.inference_mode()
     def _vocos_decode(self, spec: torch.Tensor) -> np.ndarray:
@@ -395,12 +397,15 @@ class Chat:
 
     @torch.inference_mode()
     def _decode_to_wavs(
-        self, result: GPT.GenerationOutputs, start_seeks: List[int], use_decoder: bool
+        self,
+        result: GPT.GenerationOutputs,
+        start_seeks: np.ndarray,
+        use_decoder: bool,
     ):
         x = result.hiddens if use_decoder else result.ids
         wavs: List[Optional[np.ndarray]] = []
         for i, chunk_data in enumerate(x):
-            start_seek = start_seeks[i]
+            start_seek: int = start_seeks[i]
             length = len(chunk_data)
             if length <= start_seek:
                 wavs.append(None)
@@ -585,7 +590,6 @@ class Chat:
                 stream=False,
                 show_tqdm=params.show_tqdm,
                 ensure_non_empty=params.ensure_non_empty,
-                stream_batch=params.stream_batch,
                 context=self.context,
             )
         )

ChatTTS/model/dvae.py

@@ -195,6 +195,9 @@ class DVAE(nn.Module):
             self.coef.cpu().numpy().astype(np.float32).tobytes()
         )
 
+    def __call__(self, inp: torch.Tensor) -> torch.Tensor:
+        return super().__call__(inp)
+
     @torch.inference_mode()
     def forward(self, inp: torch.Tensor) -> torch.Tensor:
         if self.vq_layer is not None:
@@ -216,4 +219,6 @@ class DVAE(nn.Module):
             ),
         )
 
+        del vq_feats
+
         return torch.mul(dec_out, self.coef, out=dec_out)

ChatTTS/model/gpt.py

@@ -509,6 +509,8 @@ class GPT(nn.Module):
 
             idx_next = torch.multinomial(scores, num_samples=1).to(finish.device)
 
+            del scores
+
             if not infer_text:
                 # idx_next = rearrange(idx_next, "(b n) 1 -> b n", n=self.num_vq)
                 idx_next = idx_next.view(-1, self.num_vq)

examples/web/funcs.py

@@ -6,7 +6,7 @@ from time import sleep
 import gradio as gr
 import numpy as np
 
-from tools.audio import unsafe_float_to_int16, has_ffmpeg_installed
+from tools.audio import float_to_int16, has_ffmpeg_installed
 from tools.logger import get_logger
 
 logger = get_logger(" WebUI ")
@@ -168,10 +168,10 @@ def generate_audio(
             for gen in wav:
                 audio = gen[0]
                 if audio is not None and len(audio) > 0:
-                    yield 24000, unsafe_float_to_int16(audio).T
+                    yield 24000, float_to_int16(audio).T
                 del audio
         else:
-            yield 24000, unsafe_float_to_int16(wav[0]).T
+            yield 24000, float_to_int16(wav[0]).T
 
 
 def interrupt_generate():

tools/audio/__init__.py

@@ -1,3 +1,3 @@
 from .mp3 import wav_arr_to_mp3_view
 from .ffmpeg import has_ffmpeg_installed
-from .np import unsafe_float_to_int16, batch_unsafe_float_to_int16
+from .np import float_to_int16, batch_unsafe_float_to_int16

tools/audio/mp3.py

@@ -3,7 +3,7 @@ from io import BytesIO
 
 import numpy as np
 
-from .np import unsafe_float_to_int16
+from .np import float_to_int16
 from .av import wav2
 
 
@@ -13,7 +13,7 @@ def wav_arr_to_mp3_view(wav: np.ndarray):
         wf.setnchannels(1)  # Mono channel
         wf.setsampwidth(2)  # Sample width in bytes
         wf.setframerate(24000)  # Sample rate in Hz
-        wf.writeframes(unsafe_float_to_int16(wav))
+        wf.writeframes(float_to_int16(wav))
     buf.seek(0, 0)
     buf2 = BytesIO()
     wav2(buf, buf2, "mp3")

tools/audio/np.py

@@ -1,17 +1,14 @@
 import numpy as np
 from numba import jit
 
-
 @jit
-def unsafe_float_to_int16(audio: np.ndarray) -> np.ndarray:
+def float_to_int16(audio: np.ndarray) -> np.ndarray:
     """
     This function will destroy audio, use only once.
     """
     am = np.abs(audio).max() * 32768
     am = 32767 * 32768 / am
-    np.multiply(audio, am, audio)
-    audio16 = audio.astype(np.int16)
-    return audio16
+    return np.multiply(audio, am).astype(np.int16)
 
 
 @jit