fix: missing module name "tools.audio" (#880)

fix #736 fix #820 fix #844

.github/workflows/checksum.yml

@@ -13,7 +13,7 @@ jobs:
 
       - name: Run RVC-Models-Downloader
         run: |
-          wget https://github.com/fumiama/RVC-Models-Downloader/releases/download/v0.2.9/rvcmd_linux_amd64.deb
+          wget https://github.com/fumiama/RVC-Models-Downloader/releases/download/v0.2.10/rvcmd_linux_amd64.deb
           sudo apt -y install ./rvcmd_linux_amd64.deb
           rm -f ./rvcmd_linux_amd64.deb
           rvcmd -notrs -w 1 -notui assets/chtts

ChatTTS/core.py

@@ -1,4 +1,5 @@
 import os
+import re
 import logging
 import tempfile
 from dataclasses import dataclass, asdict
@@ -68,13 +69,13 @@ class Chat:
         custom_path: Optional[torch.serialization.FILE_LIKE] = None,
     ) -> Optional[str]:
         if source == "local":
-            download_path = os.getcwd()
+            download_path = custom_path if custom_path is not None else os.getcwd()
             if (
                 not check_all_assets(Path(download_path), self.sha256_map, update=True)
                 or force_redownload
             ):
                 with tempfile.TemporaryDirectory() as tmp:
-                    download_all_assets(tmpdir=tmp)
+                    download_all_assets(tmpdir=tmp, homedir=download_path)
                 if not check_all_assets(
                     Path(download_path), self.sha256_map, update=False
                 ):
@@ -83,10 +84,20 @@ class Chat:
                     )
                     return None
         elif source == "huggingface":
-            hf_home = os.getenv("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
             try:
-                download_path = get_latest_modified_file(
-                    os.path.join(hf_home, "hub/models--2Noise--ChatTTS/snapshots")
+                download_path = (
+                    get_latest_modified_file(
+                        os.path.join(
+                            os.getenv(
+                                "HF_HOME", os.path.expanduser("~/.cache/huggingface")
+                            ),
+                            "hub/models--2Noise--ChatTTS/snapshots",
+                        )
+                    )
+                    if custom_path is None
+                    else get_latest_modified_file(
+                        os.path.join(custom_path, "models--2Noise--ChatTTS/snapshots")
+                    )
                 )
             except:
                 download_path = None
@@ -99,16 +110,16 @@ class Chat:
                     download_path = snapshot_download(
                         repo_id="2Noise/ChatTTS",
                         allow_patterns=["*.yaml", "*.json", "*.safetensors"],
+                        cache_dir=custom_path,
+                        force_download=force_redownload,
                     )
                 except:
                     download_path = None
-            else:
-                self.logger.log(
-                    logging.INFO, f"load latest snapshot from cache: {download_path}"
-                )
-            if download_path is None:
-                self.logger.error("download from huggingface failed.")
-                return None
+                else:
+                    self.logger.log(
+                        logging.INFO,
+                        f"load latest snapshot from cache: {download_path}",
+                    )
         elif source == "custom":
             self.logger.log(logging.INFO, f"try to load from local: {custom_path}")
             if not check_all_assets(Path(custom_path), self.sha256_map, update=False):
@@ -116,6 +127,10 @@ class Chat:
                 return None
             download_path = custom_path
 
+        if download_path is None:
+            self.logger.error("Model download failed")
+            return None
+
         return download_path
 
     def load(
@@ -199,10 +214,32 @@ class Chat:
         use_decoder=True,
         do_text_normalization=True,
         do_homophone_replacement=True,
+        split_text=True,
+        max_split_batch=4,
         params_refine_text=RefineTextParams(),
         params_infer_code=InferCodeParams(),
     ):
         self.context.set(False)
+
+        if split_text and isinstance(text, str):
+            if "\n" in text:
+                text = text.split("\n")
+            else:
+                text = re.split(r"(?<=。)|(?<=\.\s)", text)
+                nt = []
+                if isinstance(text, list):
+                    for t in text:
+                        if t:
+                            nt.append(t)
+                    text = nt
+                else:
+                    text = [text]
+            self.logger.info("split text into %d parts", len(text))
+            self.logger.debug("%s", str(text))
+
+        if len(text) == 0:
+            return []
+
         res_gen = self._infer(
             text,
             stream,
@@ -212,11 +249,21 @@ class Chat:
             use_decoder,
             do_text_normalization,
             do_homophone_replacement,
+            split_text,
+            max_split_batch,
             params_refine_text,
             params_infer_code,
         )
         if stream:
             return res_gen
+        elif not refine_text_only:
+            stripped_wavs = []
+            for wavs in res_gen:
+                for wav in wavs:
+                    stripped_wavs.append(wav[np.abs(wav) > 1e-5])
+            if split_text:
+                return [np.concatenate(stripped_wavs)]
+            return stripped_wavs
         else:
             return next(res_gen)
 
@@ -336,7 +383,7 @@ class Chat:
 
     def _infer(
         self,
-        text,
+        text: Union[List[str], str],
         stream=False,
         lang=None,
         skip_refine_text=False,
@@ -344,6 +391,8 @@ class Chat:
         use_decoder=True,
         do_text_normalization=True,
         do_homophone_replacement=True,
+        split_text=True,
+        max_split_batch=4,
         params_refine_text=RefineTextParams(),
         params_infer_code=InferCodeParams(),
     ):
@@ -376,44 +425,80 @@ class Chat:
             text = self.tokenizer.decode(text_tokens)
             refined.destroy()
             if refine_text_only:
+                if split_text and isinstance(text, list):
+                    text = "\n".join(text)
                 yield text
                 return
 
-        if stream:
-            length = 0
-            pass_batch_count = 0
-        for result in self._infer_code(
-            text,
-            stream,
-            self.device,
-            use_decoder,
-            params_infer_code,
-        ):
+        if split_text and len(text) > 1 and params_infer_code.spk_smp is None:
+            refer_text = text[0]
+            result = next(
+                self._infer_code(
+                    refer_text,
+                    False,
+                    self.device,
+                    use_decoder,
+                    params_infer_code,
+                )
+            )
             wavs = self._decode_to_wavs(
                 result.hiddens if use_decoder else result.ids,
                 use_decoder,
             )
             result.destroy()
-            if stream:
-                pass_batch_count += 1
-                if pass_batch_count <= params_infer_code.pass_first_n_batches:
-                    continue
-                a = length
-                b = a + params_infer_code.stream_speed
-                if b > wavs.shape[1]:
-                    b = wavs.shape[1]
-                new_wavs = wavs[:, a:b]
-                length = b
-                yield new_wavs
-            else:
-                yield wavs
+            assert len(wavs), 1
+            params_infer_code.spk_smp = self.sample_audio_speaker(wavs[0])
+            params_infer_code.txt_smp = refer_text
+
         if stream:
-            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]
+            length = 0
+            pass_batch_count = 0
+        if split_text:
+            n = len(text) // max_split_batch
+            if len(text) % max_split_batch:
+                n += 1
+        else:
+            n = 1
+            max_split_batch = len(text)
+        for i in range(n):
+            text_remain = text[i * max_split_batch :]
+            if len(text_remain) > max_split_batch:
+                text_remain = text_remain[:max_split_batch]
+            if split_text:
+                self.logger.info(
+                    "infer split %d~%d",
+                    i * max_split_batch,
+                    i * max_split_batch + len(text_remain),
+                )
+            for result in self._infer_code(
+                text_remain,
+                stream,
+                self.device,
+                use_decoder,
+                params_infer_code,
+            ):
+                wavs = self._decode_to_wavs(
+                    result.hiddens if use_decoder else result.ids,
+                    use_decoder,
+                )
+                result.destroy()
+                if stream:
+                    pass_batch_count += 1
+                    if pass_batch_count <= params_infer_code.pass_first_n_batches:
+                        continue
+                    a = length
+                    b = a + params_infer_code.stream_speed
+                    if b > wavs.shape[1]:
+                        b = wavs.shape[1]
+                    new_wavs = wavs[:, a:b]
+                    length = b
+                    yield new_wavs
+                else:
+                    yield wavs
+            if stream:
+                new_wavs = wavs[:, length:]
+                keep_cols = np.sum(np.abs(new_wavs) > 1e-5, axis=0) > 0
+                yield new_wavs[:][:, keep_cols]
 
     @torch.inference_mode()
     def _vocos_decode(self, spec: torch.Tensor) -> np.ndarray:

ChatTTS/utils/dl.py

@@ -151,7 +151,7 @@ def download_dns_yaml(url: str, folder: str, headers: Dict[str, str]):
         logger.get_logger().info(f"downloaded into {folder}")
 
 
-def download_all_assets(tmpdir: str, version="0.2.9"):
+def download_all_assets(tmpdir: str, homedir: str, version="0.2.10"):
     import subprocess
     import platform
 
@@ -186,7 +186,7 @@ def download_all_assets(tmpdir: str, version="0.2.9"):
         else:
             download_and_extract_tar_gz(RVCMD_URL, tmpdir)
             os.chmod(cmdfile, 0o755)
-        subprocess.run([cmdfile, "-notui", "-w", "0", "assets/chtts"])
+        subprocess.run([cmdfile, "-notui", "-w", "0", "-H", homedir, "assets/chtts"])
     except Exception:
         BASE_URL = (
             "https://gitea.seku.su/fumiama/RVC-Models-Downloader/releases/download/"
@@ -215,6 +215,8 @@ def download_all_assets(tmpdir: str, version="0.2.9"):
                 "0",
                 "-dns",
                 os.path.join(tmpdir, "dns.yaml"),
+                "-H",
+                homedir,
                 "assets/chtts",
             ]
         )

examples/api/main.py

@@ -23,7 +23,10 @@ import torch
 
 
 from pydantic import BaseModel
-
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+from tools.normalizer.en import normalizer_en_nemo_text
+from tools.normalizer.zh import normalizer_zh_tn
 
 logger = get_logger("Command")
 
@@ -35,14 +38,23 @@ async def startup_event():
     global chat
 
     chat = ChatTTS.Chat(get_logger("ChatTTS"))
+    chat.normalizer.register("en", normalizer_en_nemo_text())
+    chat.normalizer.register("zh", normalizer_zh_tn())
+
     logger.info("Initializing ChatTTS...")
-    if chat.load():
+    if chat.load(source="huggingface"):
         logger.info("Models loaded successfully.")
     else:
         logger.error("Models load failed.")
         sys.exit(1)
 
 
+@app.exception_handler(RequestValidationError)
+async def validation_exception_handler(request, exc: RequestValidationError):
+    logger.error(f"Validation error: {exc.errors()}")
+    return JSONResponse(status_code=422, content={"detail": exc.errors()})
+
+
 class ChatTTSParams(BaseModel):
     text: list[str]
     stream: bool = False
@@ -52,7 +64,7 @@ class ChatTTSParams(BaseModel):
     use_decoder: bool = True
     do_text_normalization: bool = True
     do_homophone_replacement: bool = False
-    params_refine_text: ChatTTS.Chat.RefineTextParams
+    params_refine_text: ChatTTS.Chat.RefineTextParams = None
     params_infer_code: ChatTTS.Chat.InferCodeParams
 
 

examples/web/funcs.py

@@ -4,6 +4,10 @@ from time import sleep
 
 import gradio as gr
 
+import sys
+
+sys.path.append("..")
+sys.path.append("../..")
 from tools.audio import float_to_int16, has_ffmpeg_installed, load_audio
 from tools.logger import get_logger
 
@@ -133,6 +137,7 @@ def refine_text(
     temperature,
     top_P,
     top_K,
+    split_batch,
 ):
     global chat
 
@@ -150,6 +155,7 @@ def refine_text(
             top_K=top_K,
             manual_seed=text_seed_input,
         ),
+        split_text=split_batch > 0,
     )
 
     return text[0] if isinstance(text, list) else text
@@ -165,6 +171,7 @@ def generate_audio(
     audio_seed_input,
     sample_text_input,
     sample_audio_code_input,
+    split_batch,
 ):
     global chat, has_interrupted
 
@@ -189,6 +196,8 @@ def generate_audio(
         skip_refine_text=True,
         params_infer_code=params_infer_code,
         stream=stream,
+        split_text=split_batch > 0,
+        max_split_batch=split_batch,
     )
     if stream:
         for gen in wav:

examples/web/webui.py

@@ -139,6 +139,14 @@ def main():
                 scale=1,
                 interactive=True,
             )
+            split_batch_slider = gr.Slider(
+                minimum=0,
+                maximum=100,
+                step=1,
+                value=4,
+                label="Split Batch",
+                interactive=True,
+            )
             generate_button = gr.Button(
                 "Generate", scale=2, variant="primary", interactive=True
             )
@@ -208,6 +216,7 @@ def main():
                     temperature_slider,
                     top_p_slider,
                     top_k_slider,
+                    split_batch_slider,
                 ],
                 outputs=text_output,
             ).then(
@@ -222,6 +231,7 @@ def main():
                     audio_seed_input,
                     sample_text_input,
                     sample_audio_code_input,
+                    split_batch_slider,
                 ],
                 outputs=audio_output,
             ).then(

tests/#511.py

@@ -41,6 +41,7 @@ fail = False
 wavs = chat.infer(
     texts,
     skip_refine_text=True,
+    split_text=False,
     params_infer_code=params_infer_code,
 )
 

tests/#588.py

@@ -29,6 +29,7 @@ refined = chat.infer(
     texts,
     refine_text_only=True,
     stream=False,
+    split_text=False,
     params_refine_text=ChatTTS.Chat.RefineTextParams(show_tqdm=False),
 )
 

tests/#655.py

@@ -38,6 +38,7 @@ refined_text = chat.infer(
         prompt="[oral_2][laugh_0][break_6]",
         manual_seed=12345,
     ),
+    split_text=False,
 )
 if (
     refined_text[0]

tools/audio/av.py

@@ -1,8 +1,9 @@
 from io import BufferedWriter, BytesIO
 from pathlib import Path
-from typing import Dict
+from typing import Dict, Tuple, Optional, Union, List
 
 import av
+from av.audio.frame import AudioFrame
 from av.audio.resampler import AudioResampler
 import numpy as np
 
@@ -39,41 +40,88 @@ def wav2(i: BytesIO, o: BufferedWriter, format: str):
     inp.close()
 
 
-def load_audio(file: str, sr: int) -> np.ndarray:
+def load_audio(
+    file: Union[str, BytesIO, Path],
+    sr: Optional[int] = None,
+    format: Optional[str] = None,
+    mono=True,
+) -> Union[np.ndarray, Tuple[np.ndarray, int]]:
     """
     https://github.com/fumiama/Retrieval-based-Voice-Conversion-WebUI/blob/412a9950a1e371a018c381d1bfb8579c4b0de329/infer/lib/audio.py#L39
     """
-
-    if not Path(file).exists():
+    if (isinstance(file, str) and not Path(file).exists()) or (
+        isinstance(file, Path) and not file.exists()
+    ):
         raise FileNotFoundError(f"File not found: {file}")
+    rate = 0
 
-    try:
-        container = av.open(file)
-        resampler = AudioResampler(format="fltp", layout="mono", rate=sr)
+    container = av.open(file, format=format)
+    audio_stream = next(s for s in container.streams if s.type == "audio")
+    channels = 1 if audio_stream.layout == "mono" else 2
+    container.seek(0)
+    resampler = (
+        AudioResampler(format="fltp", layout=audio_stream.layout, rate=sr)
+        if sr is not None
+        else None
+    )
 
-        # Estimated maximum total number of samples to pre-allocate the array
-        # AV stores length in microseconds by default
-        estimated_total_samples = int(container.duration * sr // 1_000_000)
-        decoded_audio = np.zeros(estimated_total_samples + 1, dtype=np.float32)
+    # Estimated maximum total number of samples to pre-allocate the array
+    # AV stores length in microseconds by default
+    estimated_total_samples = (
+        int(container.duration * sr // 1_000_000) if sr is not None else 48000
+    )
+    decoded_audio = np.zeros(
+        (
+            estimated_total_samples + 1
+            if channels == 1
+            else (channels, estimated_total_samples + 1)
+        ),
+        dtype=np.float32,
+    )
 
-        offset = 0
-        for frame in container.decode(audio=0):
-            frame.pts = None  # Clear presentation timestamp to avoid resampling issues
-            resampled_frames = resampler.resample(frame)
+    offset = 0
+
+    def process_packet(packet: List[AudioFrame]):
+        frames_data = []
+        rate = 0
+        for frame in packet:
+            # frame.pts = None  # 清除时间戳，避免重新采样问题
+            resampled_frames = (
+                resampler.resample(frame) if resampler is not None else [frame]
+            )
             for resampled_frame in resampled_frames:
-                frame_data = resampled_frame.to_ndarray()[0]
-                end_index = offset + len(frame_data)
+                frame_data = resampled_frame.to_ndarray()
+                rate = resampled_frame.rate
+                frames_data.append(frame_data)
+        return (rate, frames_data)
 
-                # Check if decoded_audio has enough space, and resize if necessary
-                if end_index > decoded_audio.shape[0]:
-                    decoded_audio = np.resize(decoded_audio, end_index + 1)
+    def frame_iter(container):
+        for p in container.demux(container.streams.audio[0]):
+            yield p.decode()
 
-                decoded_audio[offset:end_index] = frame_data
-                offset += len(frame_data)
+    for r, frames_data in map(process_packet, frame_iter(container)):
+        if not rate:
+            rate = r
+        for frame_data in frames_data:
+            end_index = offset + len(frame_data[0])
 
-        # Truncate the array to the actual size
-        decoded_audio = decoded_audio[:offset]
-    except Exception as e:
-        raise RuntimeError(f"Failed to load audio: {e}")
+            # 检查 decoded_audio 是否有足够的空间，并在必要时调整大小
+            if end_index > decoded_audio.shape[1]:
+                decoded_audio = np.resize(
+                    decoded_audio, (decoded_audio.shape[0], end_index * 4)
+                )
 
-    return decoded_audio
+            np.copyto(decoded_audio[..., offset:end_index], frame_data)
+            offset += len(frame_data[0])
+
+    container.close()
+
+    # Truncate the array to the actual size
+    decoded_audio = decoded_audio[..., :offset]
+
+    if mono and decoded_audio.shape[0] > 1:
+        decoded_audio = decoded_audio.mean(0)
+
+    if sr is not None:
+        return decoded_audio
+    return decoded_audio, rate