chore(format): run black on dev (#994)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

.gitattributes

@@ -1,2 +1,3 @@
 # ignore jupyter notebooks in the language bar on github
 **/*.ipynb linguist-vendored
+*.ipynb

.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.10/rvcmd_linux_amd64.deb
+          wget https://github.com/fumiama/RVC-Models-Downloader/releases/download/v0.2.11/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

.github/workflows/unitest.yml

@@ -20,6 +20,7 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
+          cache: 'pip'
 
       - name: Install Dependents
         run: |

ChatTTS/core.py

@@ -24,6 +24,7 @@ from .utils import (
     del_all,
 )
 from .utils import logger as utils_logger
+from .utils import FileLike
 
 from .norm import Normalizer
 
@@ -66,7 +67,7 @@ class Chat:
         self,
         source: Literal["huggingface", "local", "custom"] = "local",
         force_redownload=False,
-        custom_path: Optional[torch.serialization.FILE_LIKE] = None,
+        custom_path: Optional[FileLike] = None,
     ) -> Optional[str]:
         if source == "local":
             download_path = custom_path if custom_path is not None else os.getcwd()
@@ -138,12 +139,13 @@ class Chat:
         source: Literal["huggingface", "local", "custom"] = "local",
         force_redownload=False,
         compile: bool = False,
-        custom_path: Optional[torch.serialization.FILE_LIKE] = None,
+        custom_path: Optional[FileLike] = None,
         device: Optional[torch.device] = None,
-        coef: Optional[torch.Tensor] = None,
+        coef: Optional[str] = None,
         use_flash_attn=False,
         use_vllm=False,
         experimental: bool = False,
+        enable_cache=True,
     ) -> bool:
         download_path = self.download_models(source, force_redownload, custom_path)
         if download_path is None:
@@ -155,6 +157,7 @@ class Chat:
             use_flash_attn=use_flash_attn,
             use_vllm=use_vllm,
             experimental=experimental,
+            enable_cache=enable_cache,
             **{
                 k: os.path.join(download_path, v)
                 for k, v in asdict(self.config.path).items()
@@ -258,9 +261,10 @@ class Chat:
             return res_gen
         elif not refine_text_only:
             stripped_wavs = []
+            thr = np.float32(1e-5)
             for wavs in res_gen:
                 for wav in wavs:
-                    stripped_wavs.append(wav[np.abs(wav) > 1e-5])
+                    stripped_wavs.append(wav[np.abs(wav) > thr])
             if split_text:
                 return [np.concatenate(stripped_wavs)]
             return stripped_wavs
@@ -285,6 +289,7 @@ class Chat:
         use_flash_attn=False,
         use_vllm=False,
         experimental: bool = False,
+        enable_cache=True,
     ):
         if device is None:
             device = select_device(experimental=experimental)
@@ -349,6 +354,7 @@ class Chat:
             device=device,
             device_gpt=self.device_gpt,
             logger=self.logger,
+            enable_cache=enable_cache,
         ).eval()
         assert gpt_ckpt_path, "gpt_ckpt_path should not be None"
         gpt.load_pretrained(gpt_ckpt_path, embed_path, experimental=experimental)
@@ -423,6 +429,7 @@ class Chat:
             text_tokens = refined.ids
             text_tokens = [i[i.less(self.tokenizer.break_0_ids)] for i in text_tokens]
             text = self.tokenizer.decode(text_tokens)
+            self.logger.debug("refined texts %s", str(text))
             refined.destroy()
             if refine_text_only:
                 if split_text and isinstance(text, list):

ChatTTS/model/embed.py

@@ -54,12 +54,15 @@ class Embed(nn.Module):
         get_emb
         """
         device = next(self.parameters()).device
+        input_ids_dev = input_ids.to(device)
+        text_mask_dev = text_mask.to(device)
+
         emb_text: torch.Tensor = self.emb_text(
-            input_ids[text_mask].narrow(1, 0, 1).squeeze_(1).to(device)
+            input_ids_dev[text_mask_dev].narrow(1, 0, 1).squeeze_(1)
         )
 
-        text_mask_inv = text_mask.logical_not().to(device)
-        masked_input_ids: torch.Tensor = input_ids[text_mask_inv].to(device)
+        text_mask_inv = text_mask_dev.logical_not()
+        masked_input_ids: torch.Tensor = input_ids_dev[text_mask_inv]
 
         emb_code = [
             self.emb_code[i](masked_input_ids[:, i]) for i in range(self.num_vq)
@@ -67,11 +70,11 @@ class Embed(nn.Module):
         emb_code = torch.stack(emb_code, 2).sum(2)
 
         emb = torch.zeros(
-            (input_ids.shape[:-1]) + (emb_text.shape[-1],),
+            (input_ids_dev.shape[:-1]) + (emb_text.shape[-1],),
             device=emb_text.device,
             dtype=emb_text.dtype,
         )
-        emb[text_mask] = emb_text
+        emb[text_mask_dev] = emb_text
         emb[text_mask_inv] = emb_code.to(emb.dtype)
 
         del emb_text, emb_code, text_mask_inv

ChatTTS/model/gpt.py

@@ -28,6 +28,7 @@ class GPT(nn.Module):
         device=torch.device("cpu"),
         device_gpt=torch.device("cpu"),
         logger=logging.getLogger(__name__),
+        enable_cache=True,
     ):
         super().__init__()
 
@@ -36,6 +37,8 @@ class GPT(nn.Module):
         self.device = device
         self.device_gpt = device_gpt
 
+        self.enable_cache = enable_cache
+
         self.generator = torch.Generator(device=device)
 
         self.num_vq = int(gpt_config["num_vq"])
@@ -142,7 +145,6 @@ class GPT(nn.Module):
     class _GenerationInputs:
         position_ids: torch.Tensor
         cache_position: torch.Tensor
-        use_cache: bool
         input_ids: Optional[torch.Tensor] = None
         past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
         attention_mask: Optional[torch.Tensor] = None
@@ -162,12 +164,11 @@ class GPT(nn.Module):
     def _prepare_generation_inputs(
         self,
         input_ids: torch.Tensor,
-        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        past_key_values: Optional[Union[Tuple[Tuple[torch.FloatTensor]], Cache]] = None,
         attention_mask: Optional[torch.Tensor] = None,
         inputs_embeds: Optional[torch.Tensor] = None,
         cache_position: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
-        use_cache=True,
     ) -> _GenerationInputs:
         # With static cache, the `past_key_values` is None
         # TODO joao: standardize interface for the different Cache classes and remove of this if
@@ -180,23 +181,30 @@ class GPT(nn.Module):
             has_static_cache = past_key_values is not None
 
         past_length = 0
+        max_cache_length = None
+        cache_length = 0
         if past_key_values is not None:
             if isinstance(past_key_values, Cache):
-                past_length = (
-                    int(cache_position[0])
-                    if cache_position is not None
-                    else past_key_values.get_seq_length()
-                )
-                max_cache_length = past_key_values.get_max_length()
-                cache_length = (
-                    past_length
-                    if max_cache_length is None
-                    else min(max_cache_length, past_length)
-                )
+                if past_key_values.layers and len(past_key_values.layers):
+                    past_length = (
+                        int(cache_position[0])
+                        if cache_position is not None
+                        else past_key_values.get_seq_length()
+                    )
+                    try:
+                        max_cache_length = past_key_values.get_max_cache_shape()
+                    except:
+                        max_cache_length = (
+                            past_key_values.get_max_length()
+                        )  # deprecated in transformers 4.48
+                    cache_length = (
+                        past_length
+                        if max_cache_length is None
+                        else min(max_cache_length, past_length)
+                    )
             # TODO joao: remove this `else` after `generate` prioritizes `Cache` objects
             else:
                 cache_length = past_length = past_key_values[0][0].shape[2]
-                max_cache_length = None
 
             # Keep only the unprocessed tokens:
             # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
@@ -219,6 +227,7 @@ class GPT(nn.Module):
             # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
             if (
                 max_cache_length is not None
+                and max_cache_length > 0
                 and attention_mask is not None
                 and cache_length + input_ids.shape[1] > max_cache_length
             ):
@@ -251,7 +260,6 @@ class GPT(nn.Module):
         model_inputs = self._GenerationInputs(
             position_ids=position_ids,
             cache_position=cache_position,
-            use_cache=use_cache,
         )
 
         # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
@@ -289,22 +297,27 @@ class GPT(nn.Module):
         hiddens: List[torch.Tensor],
         infer_text: bool,
     ) -> GenerationOutputs:
-        inputs_ids = [
-            inputs_ids[idx].narrow(0, start_idx, i) for idx, i in enumerate(end_idx)
+        end_idx_int = end_idx.int()
+
+        inputs_ids_lst = [
+            inputs_ids[idx].narrow(0, start_idx, int(i))
+            for idx, i in enumerate(end_idx_int)
         ]
         if infer_text:
-            inputs_ids = [i.narrow(1, 0, 1).squeeze_(1) for i in inputs_ids]
+            inputs_ids_lst = [i.narrow(1, 0, 1).squeeze_(1) for i in inputs_ids_lst]
 
+        hiddens_lst = []
         if len(hiddens) > 0:
-            hiddens = torch.stack(hiddens, 1)
-            hiddens = [
-                hiddens[idx].narrow(0, 0, i) for idx, i in enumerate(end_idx.int())
+            hiddens_lst = torch.stack(hiddens, 1)
+            hiddens_lst = [
+                hiddens_lst[idx].narrow(0, 0, int(i))
+                for idx, i in enumerate(end_idx_int)
             ]
 
         return self.GenerationOutputs(
-            ids=inputs_ids,
+            ids=inputs_ids_lst,
             attentions=attentions,
-            hiddens=hiddens,
+            hiddens=hiddens_lst,
         )
 
     @torch.no_grad()
@@ -331,6 +344,8 @@ class GPT(nn.Module):
         context=Context(),
     ):
 
+        self.logger.debug("start generate")
+
         attentions: List[Optional[Tuple[torch.FloatTensor, ...]]] = []
         hiddens = []
         stream_iter = 0
@@ -340,6 +355,10 @@ class GPT(nn.Module):
         )
         finish = torch.zeros(inputs_ids.shape[0], device=inputs_ids.device).bool()
 
+        self.logger.debug(
+            f"set start_idx: {start_idx}, end_idx and finish with all zeros, len {inputs_ids.shape[0]}"
+        )
+
         old_temperature = temperature
 
         temperature = (
@@ -349,6 +368,10 @@ class GPT(nn.Module):
             .view(-1, 1)
         )
 
+        self.logger.debug(
+            f"expand temperature from shape {old_temperature.shape} to {temperature.shape}"
+        )
+
         attention_mask_cache = torch.ones(
             (
                 inputs_ids.shape[0],
@@ -357,10 +380,14 @@ class GPT(nn.Module):
             dtype=torch.bool,
             device=inputs_ids.device,
         )
+        self.logger.debug(
+            f"init attention_mask_cache with shape {attention_mask_cache.shape}"
+        )
         if attention_mask is not None:
             attention_mask_cache.narrow(1, 0, attention_mask.shape[1]).copy_(
                 attention_mask
             )
+            self.logger.debug(f"copy attention_mask with shape {attention_mask.shape}")
 
         progress = inputs_ids.size(1)
         # pre-allocate inputs_ids
@@ -372,6 +399,9 @@ class GPT(nn.Module):
             device=inputs_ids.device,
         )
         inputs_ids_buf.narrow(1, 0, progress).copy_(inputs_ids)
+        self.logger.debug(
+            f"expand inputs_ids buf from shape {inputs_ids.shape} to {inputs_ids_buf.shape}"
+        )
         del inputs_ids
         inputs_ids = inputs_ids_buf.narrow(1, 0, progress)
 
@@ -388,38 +418,46 @@ class GPT(nn.Module):
 
         for i in range(max_new_token):
 
+            self.logger.debug("start _prepare_generation_inputs")
             model_input = self._prepare_generation_inputs(
                 inputs_ids,
                 past_key_values,
                 attention_mask_cache.narrow(1, 0, inputs_ids.shape[1]),
-                use_cache=not self.is_te_llama,
             )
+            self.logger.debug("finis _prepare_generation_inputs")
 
             if i > 0:
                 del emb
                 inputs_ids_emb = model_input.input_ids.to(self.device_gpt)
                 if infer_text:
+                    self.logger.debug("start emb_text")
                     emb: torch.Tensor = self.emb_text(inputs_ids_emb[:, :, 0])
+                    self.logger.debug("finis emb_text")
                 else:
+                    self.logger.debug("start code_emb")
                     code_emb = [
-                        self.emb_code[i](inputs_ids_emb[:, :, i])
+                        self.emb_code[i](inputs_ids_emb[:, :, i]).to(self.device)
                         for i in range(self.num_vq)
                     ]
                     emb = torch.stack(code_emb, 3).sum(3)
+                    self.logger.debug("finis code_emb")
                 del inputs_ids_emb, model_input.input_ids
             model_input.inputs_embeds = emb
 
+            self.logger.debug(f"move model_input to device_gpt: {str(self.device_gpt)}")
             model_input.to(self.device_gpt, self.gpt.dtype)
 
+            self.logger.debug("start gpt...")
             outputs: BaseModelOutputWithPast = self.gpt(
                 attention_mask=model_input.attention_mask,
                 position_ids=model_input.position_ids,
                 past_key_values=model_input.past_key_values,
                 inputs_embeds=model_input.inputs_embeds,
-                use_cache=model_input.use_cache,
+                use_cache=not self.is_te_llama and self.enable_cache,
                 output_attentions=return_attn,
                 cache_position=model_input.cache_position,
             )
+            self.logger.debug("finis gpt")
             del_all(model_input)
             attentions.append(outputs.attentions)
             hidden_states = outputs.last_hidden_state.to(
@@ -432,8 +470,11 @@ class GPT(nn.Module):
 
             with P.cached():
                 if infer_text:
+                    self.logger.debug("start head_text")
                     logits: torch.Tensor = self.head_text(hidden_states)
+                    self.logger.debug("finis head_text")
                 else:
+                    self.logger.debug("start head_code")
                     # logits = torch.stack([self.head_code[i](hidden_states) for i in range(self.num_vq)], 3)
                     logits = torch.empty(
                         hidden_states.size(0),
@@ -447,9 +488,11 @@ class GPT(nn.Module):
                         x: torch.Tensor = self.head_code[num_vq_iter](hidden_states)
                         logits[..., num_vq_iter] = x
                         del x
+                    self.logger.debug("finis head_code")
 
             del hidden_states
 
+            self.logger.debug("start logits")
             # logits = logits[:, -1].float()
             logits = logits.narrow(1, -1, 1).squeeze_(1).float()
 
@@ -493,6 +536,9 @@ class GPT(nn.Module):
 
             del logits
 
+            self.logger.debug("finis logits")
+
+            self.logger.debug("start seed")
             if manual_seed is None:
                 idx_next = torch.multinomial(scores, num_samples=1).to(finish.device)
             else:
@@ -504,6 +550,10 @@ class GPT(nn.Module):
 
             del scores
 
+            self.logger.debug("finis seed")
+
+            self.logger.debug("start finish")
+
             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)
@@ -519,6 +569,8 @@ class GPT(nn.Module):
                     idx_next.unsqueeze_(-1).expand(-1, -1, self.num_vq),
                 )
 
+            self.logger.debug("finis finish")
+
             if i == 0 and finish.any():
                 self.logger.warning(
                     "unexpected end at index %s",
@@ -564,6 +616,8 @@ class GPT(nn.Module):
                     del inputs_ids
                 return
 
+            self.logger.debug("start output")
+
             del idx_next
             progress += 1
             inputs_ids = inputs_ids_buf.narrow(1, 0, progress)
@@ -584,6 +638,8 @@ class GPT(nn.Module):
                     )
             del not_finished
 
+            self.logger.debug("finis output")
+
             if finish.all() or context.get():
                 break
 

ChatTTS/model/tokenizer.py

@@ -10,13 +10,13 @@ from typing import List, Tuple, Optional, Union
 import torch
 from transformers import BertTokenizerFast
 
-from ..utils import del_all
+from ..utils import del_all, FileLike
 
 
 class Tokenizer:
     def __init__(
         self,
-        tokenizer_path: torch.serialization.FILE_LIKE,
+        tokenizer_path: FileLike,
     ):
         """
         tokenizer: BertTokenizerFast = torch.load(
@@ -53,7 +53,12 @@ class Tokenizer:
 
         # avoid random speaker embedding of tokenizer in the other dims
         for t in text:
-            x = self._tokenizer.encode_plus(
+            encode_plus = (
+                self._tokenizer.encode_plus
+                if hasattr(self._tokenizer, "encode_plus")
+                else self._tokenizer._encode_plus
+            )
+            x = encode_plus(
                 t, return_tensors="pt", add_special_tokens=False, padding=True
             )
             input_ids_lst.append(x["input_ids"].squeeze_(0))
@@ -125,7 +130,7 @@ class Tokenizer:
 
         return new_input_ids, attention_mask, text_mask
 
-    @torch.inference_mode
+    @torch.inference_mode()
     def decode(
         self,
         sequences: Union[List[int], List[List[int]]],

ChatTTS/model/velocity/block_manager.py

@@ -156,7 +156,7 @@ class BlockSpaceManager:
                 self.block_sliding_window
                 and len(block_table) >= self.block_sliding_window
             ):
-                # re-use a block
+                # reuse a block
                 block_table.append(
                     block_table[len(block_table) % self.block_sliding_window]
                 )

ChatTTS/model/velocity/configs.py

@@ -12,7 +12,6 @@ import argparse
 import dataclasses
 from dataclasses import dataclass
 
-
 logger = init_logger(__name__)
 
 _GB = 1 << 30

ChatTTS/model/velocity/llama.py

@@ -21,6 +21,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Inference-only LLaMA model compatible with HuggingFace weights."""
+
 from typing import Any, Dict, List, Optional, Tuple
 
 import torch

ChatTTS/model/velocity/llm_engine.py

@@ -741,7 +741,7 @@ class LLMEngine:
 
     def _decode_sequence(self, seq: Sequence, prms: SamplingParams) -> None:
         """Decodes the new token for a sequence."""
-        (new_tokens, new_output_text, prefix_offset, read_offset) = (
+        new_tokens, new_output_text, prefix_offset, read_offset = (
             detokenize_incrementally(
                 self.tokenizer,
                 all_input_ids=seq.get_token_ids(),

ChatTTS/model/velocity/model_runner.py

@@ -360,11 +360,11 @@ class ModelRunner:
             is_prompt = seq_group_metadata_list[0].is_prompt
             # Prepare input tensors.
             if is_prompt:
-                (input_tokens, input_positions, input_metadata, prompt_lens) = (
+                input_tokens, input_positions, input_metadata, prompt_lens = (
                     self._prepare_prompt(seq_group_metadata_list)
                 )
             else:
-                (input_tokens, input_positions, input_metadata) = self._prepare_decode(
+                input_tokens, input_positions, input_metadata = self._prepare_decode(
                     seq_group_metadata_list
                 )
                 prompt_lens = []
@@ -401,9 +401,9 @@ class ModelRunner:
                 broadcast(input_metadata.block_tables, src=0)
             broadcast(sampling_metadata.selected_token_indices, src=0)
         else:
-            receving_list = [None]
-            broadcast_object_list(receving_list, src=0)
-            py_data = receving_list[0]
+            receiving_list = [None]
+            broadcast_object_list(receiving_list, src=0)
+            py_data = receiving_list[0]
             input_tokens = torch.empty(
                 *py_data["input_tokens_size"], dtype=torch.long, device="cuda"
             )
@@ -505,9 +505,9 @@ class ModelRunner:
             model_executable = self.model
 
         infer_text = sampling_metadata.seq_groups[0][1].infer_text
-        temperture = sampling_metadata.seq_groups[0][1].temperature
+        temperature = sampling_metadata.seq_groups[0][1].temperature
         if not infer_text:
-            temperture = torch.tensor(temperture).to(input_tokens.device)
+            temperature = torch.tensor(temperature).to(input_tokens.device)
         logits_processors, logits_warpers = sampling_metadata.seq_groups[0][
             1
         ].logits_processors
@@ -553,7 +553,7 @@ class ModelRunner:
             ),
             hidden_states=hidden_states,
             infer_text=infer_text,
-            temperature=temperture,
+            temperature=temperature,
             logits_processors=logits_processors,
             logits_warpers=logits_warpers,
             min_new_token=min_new_token,

ChatTTS/model/velocity/output.py

@@ -107,14 +107,14 @@ class RequestOutput:
                 # always has the logprobs of the sampled tokens even if the
                 # logprobs are not requested.
                 logprobs = None
-            finshed_reason = SequenceStatus.get_finished_reason(seq.status)
+            finished_reason = SequenceStatus.get_finished_reason(seq.status)
             output = CompletionOutput(
                 seqs.index(seq),
                 seq.output_text,
                 seq.get_output_token_ids(),
                 seq.get_cumulative_logprob(),
                 logprobs,
-                finshed_reason,
+                finished_reason,
                 seq.data.hidden_states,
             )
             outputs.append(output)

ChatTTS/model/velocity/scheduler.py

@@ -128,7 +128,7 @@ class Scheduler:
         return len(self.waiting) + len(self.running) + len(self.swapped)
 
     def _schedule(self) -> SchedulerOutputs:
-        # Blocks that need to be swaped or copied before model execution.
+        # Blocks that need to be swapped or copied before model execution.
         blocks_to_swap_in: Dict[int, int] = {}
         blocks_to_swap_out: Dict[int, int] = {}
         blocks_to_copy: Dict[int, List[int]] = {}

ChatTTS/utils/__init__.py

@@ -1,4 +1,4 @@
 from .dl import check_all_assets, download_all_assets
 from .gpu import select_device
-from .io import load_safetensors, get_latest_modified_file, del_all
+from .io import load_safetensors, get_latest_modified_file, del_all, FileLike
 from .log import logger

ChatTTS/utils/dl.py

@@ -143,15 +143,7 @@ def download_and_extract_zip(
         logger.get_logger().info(f"extracted into {folder}")
 
 
-def download_dns_yaml(url: str, folder: str, headers: Dict[str, str]):
-    logger.get_logger().info(f"downloading {url}")
-    response = requests.get(url, headers=headers, stream=True, timeout=(100, 3))
-    with open(os.path.join(folder, "dns.yaml"), "wb") as out_file:
-        out_file.write(response.content)
-        logger.get_logger().info(f"downloaded into {folder}")
-
-
-def download_all_assets(tmpdir: str, homedir: str, version="0.2.10"):
+def download_all_assets(tmpdir: str, homedir: str, version="0.2.11"):
     import subprocess
     import platform
 
@@ -175,48 +167,15 @@ def download_all_assets(tmpdir: str, homedir: str, version="0.2.10"):
     if not architecture:
         logger.get_logger().error(f"architecture {architecture} is not supported")
         exit(1)
-    try:
-        BASE_URL = "https://github.com/fumiama/RVC-Models-Downloader/releases/download/"
-        suffix = "zip" if is_win else "tar.gz"
-        RVCMD_URL = BASE_URL + f"v{version}/rvcmd_{system_type}_{architecture}.{suffix}"
-        cmdfile = os.path.join(tmpdir, "rvcmd")
-        if is_win:
-            download_and_extract_zip(RVCMD_URL, tmpdir)
-            cmdfile += ".exe"
-        else:
-            download_and_extract_tar_gz(RVCMD_URL, tmpdir)
-            os.chmod(cmdfile, 0o755)
-        subprocess.run([cmdfile, "-notui", "-w", "0", "-H", homedir, "assets/chtts"])
-    except Exception:
-        BASE_URL = (
-            "https://gitea.seku.su/fumiama/RVC-Models-Downloader/releases/download/"
-        )
-        suffix = "zip" if is_win else "tar.gz"
-        RVCMD_URL = BASE_URL + f"v{version}/rvcmd_{system_type}_{architecture}.{suffix}"
-        download_dns_yaml(
-            "https://gitea.seku.su/fumiama/RVC-Models-Downloader/raw/branch/main/dns.yaml",
-            tmpdir,
-            headers={
-                "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36 Edg/128.0.0.0"
-            },
-        )
-        cmdfile = os.path.join(tmpdir, "rvcmd")
-        if is_win:
-            download_and_extract_zip(RVCMD_URL, tmpdir)
-            cmdfile += ".exe"
-        else:
-            download_and_extract_tar_gz(RVCMD_URL, tmpdir)
-            os.chmod(cmdfile, 0o755)
-        subprocess.run(
-            [
-                cmdfile,
-                "-notui",
-                "-w",
-                "0",
-                "-dns",
-                os.path.join(tmpdir, "dns.yaml"),
-                "-H",
-                homedir,
-                "assets/chtts",
-            ]
-        )
+
+    BASE_URL = "https://github.com/fumiama/RVC-Models-Downloader/releases/download/"
+    suffix = "zip" if is_win else "tar.gz"
+    RVCMD_URL = BASE_URL + f"v{version}/rvcmd_{system_type}_{architecture}.{suffix}"
+    cmdfile = os.path.join(tmpdir, "rvcmd")
+    if is_win:
+        download_and_extract_zip(RVCMD_URL, tmpdir)
+        cmdfile += ".exe"
+    else:
+        download_and_extract_tar_gz(RVCMD_URL, tmpdir)
+        os.chmod(cmdfile, 0o755)
+    subprocess.run([cmdfile, "-notui", "-w", "0", "-H", homedir, "assets/chtts"])

ChatTTS/utils/gpu.py

@@ -1,3 +1,5 @@
+import importlib.util
+
 import torch
 
 try:
@@ -38,11 +40,24 @@ def select_device(min_memory=2047, experimental=False):
         """
         if experimental:
             # For Apple M1/M2 chips with Metal Performance Shaders
-            logger.get_logger().warning("experimantal: found apple GPU, using MPS.")
+            logger.get_logger().warning("experimental: found apple GPU, using MPS.")
             device = torch.device("mps")
         else:
             logger.get_logger().info("found Apple GPU, but use CPU.")
             device = torch.device("cpu")
+    elif importlib.util.find_spec("torch_directml") is not None:
+        """
+        Currently DML is under developing and may output wrong result,
+        so only enable this for experimental use.
+        """
+        if experimental:
+            logger.get_logger().warning("experimental: using DML.")
+            import torch_directml
+
+            device = torch_directml.device(torch_directml.default_device())
+        else:
+            logger.get_logger().info("found DML, but use CPU.")
+            device = torch.device("cpu")
     else:
         logger.get_logger().warning("no GPU or NPU found, use CPU instead")
         device = torch.device("cpu")

ChatTTS/utils/io.py

@@ -1,6 +1,6 @@
 import os
 import logging
-from typing import Union
+from typing import Union, IO
 from dataclasses import is_dataclass
 
 from safetensors import safe_open
@@ -8,6 +8,13 @@ import torch
 
 from .log import logger
 
+if hasattr(torch.serialization, "FILE_LIKE"):
+    FileLike = torch.serialization.FILE_LIKE
+elif hasattr(torch.types, "FILE_LIKE"):
+    FileLike = torch.types.FileLike
+else:
+    FileLike = Union[str, os.PathLike, IO[bytes]]
+
 
 @torch.inference_mode()
 def load_safetensors(filename: str):

README.md

@@ -21,7 +21,8 @@ A generative speech model for daily dialogue.
 > This repo contains the algorithm infrastructure and some simple examples.
 
 > [!Tip]
-> For the extended end-user products, please refer to the index repo [Awesome-ChatTTS](https://github.com/libukai/Awesome-ChatTTS/tree/en) maintained by the community.
+> For the extended end-user products, please refer to the index repo [Awesome-ChatTTS](https://github.com/libukai/Awesome-ChatTTS/tree/en) maintained by the community.  
+> You can find a diagram visualization of the codebase [here](https://github.com/CodeBoarding/GeneratedOnBoardings/blob/main/ChatTTS/on_boarding.md).
 
 ChatTTS is a text-to-speech model designed specifically for dialogue scenarios such as LLM assistant.
 

docs/cn/README.md

@@ -25,7 +25,8 @@
 > 这个仓库包含算法架构和一些简单的示例。
 
 > [!Tip]
-> 由本仓库衍生出的用户端产品，请参见由社区维护的索引仓库  [Awesome-ChatTTS](https://github.com/libukai/Awesome-ChatTTS)。
+> 由本仓库衍生出的用户端产品，请参见由社区维护的索引仓库  [Awesome-ChatTTS](https://github.com/libukai/Awesome-ChatTTS)。  
+> 您可以在[这里](https://github.com/CodeBoarding/GeneratedOnBoardings/blob/main/ChatTTS/on_boarding.md)查看代码库的图解。
 
 ChatTTS 是一款专门为对话场景（例如 LLM 助手）设计的文本转语音模型。
 

docs/es/README.md

@@ -17,6 +17,11 @@ Un modelo de generación de voz para la conversación diaria.
 > [!NOTE]
 > Atención, es posible que esta versión no sea la última. Por favor, consulte la versión en inglés para conocer todo el contenido.
 
+> [!Tip]
+> Para los productos finales ampliados, consulta el repositorio índice [Awesome-ChatTTS](https://github.com/libukai/Awesome-ChatTTS/tree/en) mantenido por la comunidad.  
+> Puedes encontrar una visualización en forma de diagrama del código [aquí](https://github.com/CodeBoarding/GeneratedOnBoardings/blob/main/ChatTTS/on_boarding.md).
+
+
 ## Introducción
 
 ChatTTS es un modelo de texto a voz diseñado específicamente para escenarios conversacionales como LLM assistant.

docs/fr/README.md

@@ -21,7 +21,8 @@ Un modèle de parole génératif pour le dialogue quotidien.
 > Ce dépôt contient l'infrastructure de l'algorithme et quelques exemples simples.
 
 > [!Tip]
-> Pour les produits finaux étendus pour les utilisateurs, veuillez consulter le dépôt index [Awesome-ChatTTS](https://github.com/libukai/Awesome-ChatTTS/tree/en) maintenu par la communauté.
+> Pour les produits finaux étendus pour les utilisateurs, veuillez consulter le dépôt index [Awesome-ChatTTS](https://github.com/libukai/Awesome-ChatTTS/tree/en) maintenu par la communauté.  
+> Vous pouvez consulter un diagramme du code [ici](https://github.com/CodeBoarding/GeneratedOnBoardings/blob/main/ChatTTS/on_boarding.md).
 
 ChatTTS est un modèle de synthèse vocale conçu spécifiquement pour les scénarios de dialogue tels que les assistants LLM.
 

docs/jp/README.md

@@ -10,6 +10,15 @@ ChatTTSは、LLMアシスタントなどの対話シナリオ用に特別に設
 
 モデルやロードマップについての正式なお問い合わせは、**open-source@2noise.com**までご連絡ください。QQグループ：808364215に参加してディスカッションすることもできます。GitHubでの問題提起も歓迎します。
 
+## はじめに
+> [!Note]
+> このリポジトリにはアルゴリズムのインフラといくつかの簡単な例が含まれています。
+
+> [!Tip]
+> エンドユーザー向けに拡張された製品については、コミュニティによって管理されているインデックスリポジトリ [Awesome-ChatTTS](https://github.com/libukai/Awesome-ChatTTS/tree/en) を参照してください。  
+> コードベースの図解は[こちら](https://github.com/CodeBoarding/GeneratedOnBoardings/blob/main/ChatTTS/on_boarding.md)でご覧いただけます。
+
+
 ---
 ## ハイライト
 1. **会話型TTS**: ChatTTSは対話ベースのタスクに最適化されており、自然で表現豊かな音声合成を実現します。複数の話者をサポートし、対話型の会話を容易にします。

docs/kr/README.md

@@ -25,7 +25,8 @@
 > 이 저장소에는 알고리즘 구조와 간단한 예시들이 포함되어 있습니다.
 
 > [!Tip]
-> 이 프로젝트에서 파생된 프로젝트는 커뮤니티가 유지 관리하는 커뮤니티[Awesome-ChatTTS](https://github.com/libukai/Awesome-ChatTTS)를 참조하시길 바랍니다.
+> 이 프로젝트에서 파생된 프로젝트는 커뮤니티가 유지 관리하는 커뮤니티[Awesome-ChatTTS](https://github.com/libukai/Awesome-ChatTTS)를 참조하시길 바랍니다.  
+> 코드베이스의 다이어그램 시각화는 [여기](https://github.com/CodeBoarding/GeneratedOnBoardings/blob/main/ChatTTS/on_boarding.md)에서 확인할 수 있습니다.
 
 ChatTTS는 대화 기반 작업(예: LLM 어시스턴트)을 위해 설계된 텍스트-음성 변환(TTS) 모델입니다.
 

docs/ru/README.md

@@ -6,6 +6,15 @@
 
 [**English**](../../README.md) | [**简体中文**](../cn/README.md) | [**日本語**](../jp/README.md) | **Русский** | [**Español**](../es/README.md) | [**Français**](../fr/README.md) | [**한국어**](../kr/README.md)
 
+## Введение
+> [!Note]
+> Этот репозиторий содержит инфраструктуру алгоритма и некоторые простые примеры.
+
+> [!Tip]
+> Для полнофункциональных пользовательских продуктов обратитесь к индексному репозиторию [Awesome-ChatTTS](https://github.com/libukai/Awesome-ChatTTS/tree/en), поддерживаемому сообществом.  
+> Схематичную визуализацию кодовой базы можно найти [здесь](https://github.com/CodeBoarding/GeneratedOnBoardings/blob/main/ChatTTS/on_boarding.md).
+
+
 ChatTTS - это модель преобразования текста в речь, специально разработанная для диалоговых сценариев, таких как помощник LLM. Она поддерживает как английский, так и китайский языки. Наша модель обучена на более чем 100 000 часах английского и китайского языков. Открытая версия на **[HuggingFace](https://huggingface.co/2Noise/ChatTTS)** - это предварительно обученная модель с 40 000 часами без SFT.
 
 Для официальных запросов о модели и плане развития, пожалуйста, свяжитесь с нами по адресу **open-source@2noise.com**. Вы можете присоединиться к нашей группе QQ: 808364215 для обсуждения. Добавление вопросов на GitHub также приветствуется.

examples/api/README.md

@@ -14,6 +14,11 @@ pip install -r examples/api/requirements.txt
 fastapi dev examples/api/main.py --host 0.0.0.0 --port 8000
 ```
 
+## Run openAI_API server
+
+```
+fastapi dev examples/api/openai_api.py --host 0.0.0.0 --port 8000
+```
 ## Generate audio using requests
 
 ```

examples/api/main.py

@@ -6,7 +6,6 @@ import zipfile
 from fastapi import FastAPI
 from fastapi.responses import StreamingResponse
 
-
 if sys.platform == "darwin":
     os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
 

examples/api/openai_api.py

@@ -0,0 +1,285 @@
+"""
+openai_api.py
+This module implements a FastAPI-based text-to-speech API compatible with OpenAI's interface specification.
+
+Main features and improvements:
+- Use app.state to manage global state, ensuring thread safety
+- Add exception handling and unified error responses to improve stability
+- Support multiple voice options and audio formats for greater flexibility
+- Add input validation to ensure the validity of request parameters
+- Support additional OpenAI TTS parameters (e.g., speed) for richer functionality
+- Implement health check endpoint for easy service status monitoring
+- Use asyncio.Lock to manage model access, improving concurrency performance
+- Load and manage speaker embedding files to support personalized speech synthesis
+"""
+
+import io
+import os
+import sys
+import asyncio
+import time
+from typing import Optional, Dict
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse, JSONResponse
+from pydantic import BaseModel, Field
+import torch
+
+# Cross-platform compatibility settings
+if sys.platform == "darwin":
+    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+# Set working directory and add to system path
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+
+# Import necessary modules
+import ChatTTS
+from tools.audio import pcm_arr_to_mp3_view, pcm_arr_to_ogg_view, pcm_arr_to_wav_view
+from tools.logger import get_logger
+from tools.normalizer.en import normalizer_en_nemo_text
+from tools.normalizer.zh import normalizer_zh_tn
+
+# Initialize logger
+logger = get_logger("Command")
+
+# Initialize FastAPI application
+app = FastAPI()
+
+# Voice mapping table
+# Download stable voices:
+# ModelScope Community: https://modelscope.cn/studios/ttwwwaa/ChatTTS_Speaker
+# HuggingFace: https://huggingface.co/spaces/taa/ChatTTS_Speaker
+VOICE_MAP = {
+    "default": "1528.pt",
+    "alloy": "1384.pt",
+    "echo": "2443.pt",
+}
+
+# Allowed audio formats
+ALLOWED_FORMATS = {"mp3", "wav", "ogg"}
+
+
+@app.on_event("startup")
+async def startup_event():
+    """Load ChatTTS model and default speaker embedding when the application starts"""
+    # Initialize ChatTTS and async lock
+    app.state.chat = ChatTTS.Chat(get_logger("ChatTTS"))
+    app.state.model_lock = asyncio.Lock()  # Use async lock instead of thread lock
+
+    # Register text normalizers
+    app.state.chat.normalizer.register("en", normalizer_en_nemo_text())
+    app.state.chat.normalizer.register("zh", normalizer_zh_tn())
+
+    logger.info("Initializing ChatTTS...")
+    if app.state.chat.load(source="huggingface"):
+        logger.info("Model loaded successfully.")
+    else:
+        logger.error("Model loading failed, exiting application.")
+        raise RuntimeError("Failed to load ChatTTS model")
+
+    # Load default speaker embedding
+    # Preload all supported speaker embeddings into memory at startup to avoid repeated loading during runtime
+    app.state.spk_emb_map = {}
+    for voice, spk_path in VOICE_MAP.items():
+        if os.path.exists(spk_path):
+            app.state.spk_emb_map[voice] = torch.load(
+                spk_path, map_location=torch.device("cpu")
+            )
+            logger.info(f"Preloading speaker embedding: {voice} -> {spk_path}")
+        else:
+            logger.warning(f"Speaker embedding not found: {spk_path}, skipping preload")
+    app.state.spk_emb = app.state.spk_emb_map.get("default")  # Default embedding
+
+
+# Request parameter whitelist
+ALLOWED_PARAMS = {
+    "model",
+    "input",
+    "voice",
+    "response_format",
+    "speed",
+    "stream",
+    "output_format",
+}
+
+
+class OpenAITTSRequest(BaseModel):
+    """OpenAI TTS request data model"""
+
+    model: str = Field(..., description="Speech synthesis model, fixed as 'tts-1'")
+    input: str = Field(
+        ..., description="Text content to synthesize", max_length=2048
+    )  # Length limit
+    voice: Optional[str] = Field(
+        "default", description="Voice selection, supports: default, alloy, echo"
+    )
+    response_format: Optional[str] = Field(
+        "mp3", description="Audio format: mp3, wav, ogg"
+    )
+    speed: Optional[float] = Field(
+        1.0, ge=0.5, le=2.0, description="Speed, range 0.5-2.0"
+    )
+    stream: Optional[bool] = Field(False, description="Whether to stream")
+    output_format: Optional[str] = "mp3"  # Optional formats: mp3, wav, ogg
+    extra_params: Dict[str, Optional[str]] = Field(
+        default_factory=dict, description="Unsupported extra parameters"
+    )
+
+    @classmethod
+    def validate_request(cls, request_data: Dict):
+        """Filter unsupported request parameters and unify model value to 'tts-1'"""
+        request_data["model"] = "tts-1"  # Unify model value
+        unsupported_params = set(request_data.keys()) - ALLOWED_PARAMS
+        if unsupported_params:
+            logger.warning(f"Ignoring unsupported parameters: {unsupported_params}")
+        return {key: request_data[key] for key in ALLOWED_PARAMS if key in request_data}
+
+
+# Unified error response
+@app.exception_handler(Exception)
+async def custom_exception_handler(request, exc):
+    """Custom exception handler"""
+    logger.error(f"Error: {str(exc)}")
+    return JSONResponse(
+        status_code=getattr(exc, "status_code", 500),
+        content={"error": {"message": str(exc), "type": exc.__class__.__name__}},
+    )
+
+
+@app.post("/v1/audio/speech")
+async def generate_voice(request_data: Dict):
+    """Handle speech synthesis request"""
+    request_data = OpenAITTSRequest.validate_request(request_data)
+    request = OpenAITTSRequest(**request_data)
+
+    logger.info(
+        f"Received request: text={request.input}..., voice={request.voice}, stream={request.stream}"
+    )
+
+    # Validate audio format
+    if request.response_format not in ALLOWED_FORMATS:
+        raise HTTPException(
+            400,
+            detail=f"Unsupported audio format: {request.response_format}, supported formats: {', '.join(ALLOWED_FORMATS)}",
+        )
+
+    # Load speaker embedding for the specified voice
+    spk_emb = app.state.spk_emb_map.get(request.voice, app.state.spk_emb)
+
+    # Inference parameters
+    params_infer_main = {
+        "text": [request.input],
+        "stream": request.stream,
+        "lang": None,
+        "skip_refine_text": True,  # Do not use text refinement
+        "refine_text_only": False,
+        "use_decoder": True,
+        "audio_seed": 12345678,
+        # "text_seed": 87654321,  # Random seed for text processing, used to control text refinement
+        "do_text_normalization": True,  # Perform text normalization
+        "do_homophone_replacement": True,  # Perform homophone replacement
+    }
+
+    # Inference code parameters
+    params_infer_code = app.state.chat.InferCodeParams(
+        # prompt=f"[speed_{int(request.speed * 10)}]",  # Convert to format supported by ChatTTS
+        prompt="[speed_5]",
+        top_P=0.5,
+        top_K=10,
+        temperature=0.1,
+        repetition_penalty=1.1,
+        max_new_token=2048,
+        min_new_token=0,
+        show_tqdm=True,
+        ensure_non_empty=True,
+        manual_seed=42,
+        spk_emb=spk_emb,
+        spk_smp=None,
+        txt_smp=None,
+        stream_batch=24,
+        stream_speed=12000,
+        pass_first_n_batches=2,
+    )
+
+    try:
+        async with app.state.model_lock:
+            wavs = app.state.chat.infer(
+                text=params_infer_main["text"],
+                stream=params_infer_main["stream"],
+                lang=params_infer_main["lang"],
+                skip_refine_text=params_infer_main["skip_refine_text"],
+                use_decoder=params_infer_main["use_decoder"],
+                do_text_normalization=params_infer_main["do_text_normalization"],
+                do_homophone_replacement=params_infer_main["do_homophone_replacement"],
+                # params_refine_text = params_refine_text,
+                params_infer_code=params_infer_code,
+            )
+    except Exception as e:
+        raise HTTPException(500, detail=f"Speech synthesis failed: {str(e)}")
+
+    def generate_wav_header(sample_rate=24000, bits_per_sample=16, channels=1):
+        """Generate WAV file header (without data length)"""
+        header = bytearray()
+        header.extend(b"RIFF")
+        header.extend(b"\xff\xff\xff\xff")  # File size unknown
+        header.extend(b"WAVEfmt ")
+        header.extend((16).to_bytes(4, "little"))  # fmt chunk size
+        header.extend((1).to_bytes(2, "little"))  # PCM format
+        header.extend((channels).to_bytes(2, "little"))  # Channels
+        header.extend((sample_rate).to_bytes(4, "little"))  # Sample rate
+        byte_rate = sample_rate * channels * bits_per_sample // 8
+        header.extend((byte_rate).to_bytes(4, "little"))  # Byte rate
+        block_align = channels * bits_per_sample // 8
+        header.extend((block_align).to_bytes(2, "little"))  # Block align
+        header.extend((bits_per_sample).to_bytes(2, "little"))  # Bits per sample
+        header.extend(b"data")
+        header.extend(b"\xff\xff\xff\xff")  # Data size unknown
+        return bytes(header)
+
+    # Handle audio output format
+    def convert_audio(wav, format):
+        """Convert audio format"""
+        if format == "mp3":
+            return pcm_arr_to_mp3_view(wav)
+        elif format == "wav":
+            return pcm_arr_to_wav_view(
+                wav, include_header=False
+            )  # No header in streaming
+        elif format == "ogg":
+            return pcm_arr_to_ogg_view(wav)
+        return pcm_arr_to_mp3_view(wav)
+
+    # Return streaming audio data
+    if request.stream:
+        first_chunk = True
+
+        async def audio_stream():
+            nonlocal first_chunk
+            for wav in wavs:
+                if request.response_format == "wav" and first_chunk:
+                    yield generate_wav_header()  # Send WAV header
+                    first_chunk = False
+                yield convert_audio(wav, request.response_format)
+
+        media_type = "audio/wav" if request.response_format == "wav" else "audio/mpeg"
+        return StreamingResponse(audio_stream(), media_type=media_type)
+
+    # Return audio file directly
+    if request.response_format == "wav":
+        music_data = pcm_arr_to_wav_view(wavs[0])
+    else:
+        music_data = convert_audio(wavs[0], request.response_format)
+
+    return StreamingResponse(
+        io.BytesIO(music_data),
+        media_type="audio/mpeg",
+        headers={
+            "Content-Disposition": f"attachment; filename=output.{request.response_format}"
+        },
+    )
+
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {"status": "healthy", "model_loaded": bool(app.state.chat)}

examples/api/postScript.py

@@ -0,0 +1,242 @@
+import argparse
+import datetime
+import os
+import zipfile
+from io import BytesIO
+
+import requests
+
+chattts_service_host = os.environ.get("CHATTTS_SERVICE_HOST", "127.0.0.1")
+chattts_service_port = os.environ.get("CHATTTS_SERVICE_PORT", "9900")
+
+CHATTTS_URL = f"http://{chattts_service_host}:{chattts_service_port}/generate_voice"
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="HTTP client for ChatTTS service")
+    parser.add_argument(
+        "--text", type=str, nargs="+", required=True, help="Text to synthesize"
+    )
+    parser.add_argument(
+        "--audio_seed", type=int, required=True, help="Audio generation seed"
+    )
+    parser.add_argument(
+        "--text_seed", type=int, required=True, help="Text generation seed"
+    )
+    parser.add_argument(
+        "--stream", type=bool, default=False, help="Enable/disable streaming"
+    )
+    parser.add_argument("--lang", type=str, default=None, help="Language code for text")
+    parser.add_argument(
+        "--skip_refine_text", type=bool, default=True, help="Skip text refinement"
+    )
+    parser.add_argument(
+        "--refine_text_only", type=bool, default=False, help="Only refine text"
+    )
+    parser.add_argument(
+        "--use_decoder", type=bool, default=True, help="Use decoder during inference"
+    )
+    parser.add_argument(
+        "--do_text_normalization",
+        type=bool,
+        default=True,
+        help="Enable text normalization",
+    )
+    parser.add_argument(
+        "--do_homophone_replacement",
+        type=bool,
+        default=False,
+        help="Enable homophone replacement",
+    )
+    parser.add_argument(
+        "--tgt",
+        type=str,
+        default="./output",
+        help="Target directory to save output files",
+    )
+    parser.add_argument(
+        "--filename",
+        type=str,
+        default="test.mp3",
+        help="Target directory to save output files",
+    )
+
+    # Refinement text parameters
+    parser.add_argument(
+        "--refine_prompt", type=str, default="", help="Prompt for text refinement"
+    )
+    parser.add_argument(
+        "--refine_top_P",
+        type=float,
+        default=0.7,
+        help="Top P value for text refinement",
+    )
+    parser.add_argument(
+        "--refine_top_K", type=int, default=20, help="Top K value for text refinement"
+    )
+    parser.add_argument(
+        "--refine_temperature",
+        type=float,
+        default=0.7,
+        help="Temperature for text refinement",
+    )
+    parser.add_argument(
+        "--refine_repetition_penalty",
+        type=float,
+        default=1.0,
+        help="Repetition penalty for text refinement",
+    )
+    parser.add_argument(
+        "--refine_max_new_token",
+        type=int,
+        default=384,
+        help="Max new tokens for text refinement",
+    )
+    parser.add_argument(
+        "--refine_min_new_token",
+        type=int,
+        default=0,
+        help="Min new tokens for text refinement",
+    )
+    parser.add_argument(
+        "--refine_show_tqdm",
+        type=bool,
+        default=True,
+        help="Show progress bar for text refinement",
+    )
+    parser.add_argument(
+        "--refine_ensure_non_empty",
+        type=bool,
+        default=True,
+        help="Ensure non-empty output",
+    )
+    parser.add_argument(
+        "--refine_stream_batch",
+        type=int,
+        default=24,
+        help="Stream batch size for refinement",
+    )
+
+    # Infer code parameters
+    parser.add_argument(
+        "--infer_prompt", type=str, default="[speed_5]", help="Prompt for inference"
+    )
+    parser.add_argument(
+        "--infer_top_P", type=float, default=0.1, help="Top P value for inference"
+    )
+    parser.add_argument(
+        "--infer_top_K", type=int, default=20, help="Top K value for inference"
+    )
+    parser.add_argument(
+        "--infer_temperature", type=float, default=0.3, help="Temperature for inference"
+    )
+    parser.add_argument(
+        "--infer_repetition_penalty",
+        type=float,
+        default=1.05,
+        help="Repetition penalty for inference",
+    )
+    parser.add_argument(
+        "--infer_max_new_token",
+        type=int,
+        default=2048,
+        help="Max new tokens for inference",
+    )
+    parser.add_argument(
+        "--infer_min_new_token",
+        type=int,
+        default=0,
+        help="Min new tokens for inference",
+    )
+    parser.add_argument(
+        "--infer_show_tqdm",
+        type=bool,
+        default=True,
+        help="Show progress bar for inference",
+    )
+    parser.add_argument(
+        "--infer_ensure_non_empty",
+        type=bool,
+        default=True,
+        help="Ensure non-empty output",
+    )
+    parser.add_argument(
+        "--infer_stream_batch",
+        type=bool,
+        default=True,
+        help="Stream batch for inference",
+    )
+    parser.add_argument(
+        "--infer_spk_emb",
+        type=str,
+        default=None,
+        help="Speaker embedding for inference",
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_arguments()
+
+    # Main infer params
+    body = {
+        "text": args.text,
+        "stream": args.stream,
+        "lang": args.lang,
+        "filename": args.filename,
+        "skip_refine_text": args.skip_refine_text,
+        "refine_text_only": args.refine_text_only,
+        "use_decoder": args.use_decoder,
+        "audio_seed": args.audio_seed,
+        "text_seed": args.text_seed,
+        "do_text_normalization": args.do_text_normalization,
+        "do_homophone_replacement": args.do_homophone_replacement,
+    }
+    # Refinement text parameters
+    params_refine_text = {
+        "prompt": args.refine_prompt,
+        "top_P": args.refine_top_P,
+        "top_K": args.refine_top_K,
+        "temperature": args.refine_temperature,
+        "repetition_penalty": args.refine_repetition_penalty,
+        "max_new_token": args.refine_max_new_token,
+        "min_new_token": args.refine_min_new_token,
+        "show_tqdm": args.refine_show_tqdm,
+        "ensure_non_empty": args.refine_ensure_non_empty,
+        "stream_batch": args.refine_stream_batch,
+    }
+    body["params_refine_text"] = params_refine_text
+
+    # Infer code parameters
+    params_infer_code = {
+        "prompt": args.infer_prompt,
+        "top_P": args.infer_top_P,
+        "top_K": args.infer_top_K,
+        "temperature": args.infer_temperature,
+        "repetition_penalty": args.infer_repetition_penalty,
+        "max_new_token": args.infer_max_new_token,
+        "min_new_token": args.infer_min_new_token,
+        "show_tqdm": args.infer_show_tqdm,
+        "ensure_non_empty": args.infer_ensure_non_empty,
+        "stream_batch": args.infer_stream_batch,
+        "spk_emb": args.infer_spk_emb,
+    }
+    body["params_infer_code"] = params_infer_code
+
+    try:
+        response = requests.post(CHATTTS_URL, json=body)
+        response.raise_for_status()
+        with zipfile.ZipFile(BytesIO(response.content), "r") as zip_ref:
+            tgt = args.tgt
+            # filename=args.filename
+            os.makedirs(tgt, exist_ok=True)
+            zip_ref.extractall(tgt)
+            print(f"Extracted files:{tgt}/{filename}")
+            # print(tgt)
+    except requests.exceptions.RequestException as e:
+        print(f"Request Error: {e}")
+
+
+if __name__ == "__main__":
+    main()

examples/onnx/modeling_llama.py

@@ -21,6 +21,7 @@
 PyTorch LLaMA model.
 Copied from https://github.com/sophgo/LLM-TPU/blob/main/models/Llama2/compile/files/llama-2-7b-chat-hf/modeling_llama.py
 """
+
 import math
 from typing import List, Optional, Tuple, Union
 
@@ -45,7 +46,6 @@ from transformers.utils import (
 )
 from transformers.models.llama.configuration_llama import LlamaConfig
 
-
 logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "LlamaConfig"

examples/web/funcs.py

@@ -25,6 +25,9 @@ custom_path: Optional[str] = None
 has_interrupted = False
 is_in_generate = False
 
+enable_cache = True
+experimental = False
+
 seed_min = 1
 seed_max = 4294967295
 
@@ -62,12 +65,27 @@ def on_audio_seed_change(audio_seed_input):
     return rand_spk
 
 
+def set_params(en_cache, exp):
+    global enable_cache, experimental
+
+    enable_cache = en_cache
+    experimental = exp
+
+
 def load_chat(cust_path: Optional[str], coef: Optional[str]) -> bool:
+    global enable_cache, experimental
+
     if cust_path == None:
-        ret = chat.load(coef=coef)
+        ret = chat.load(coef=coef, enable_cache=enable_cache, experimental=experimental)
     else:
         logger.info("local model path: %s", cust_path)
-        ret = chat.load("custom", custom_path=cust_path, coef=coef)
+        ret = chat.load(
+            "custom",
+            custom_path=cust_path,
+            coef=coef,
+            enable_cache=enable_cache,
+            experimental=experimental,
+        )
         global custom_path
         custom_path = cust_path
     if ret:
@@ -102,7 +120,7 @@ def reload_chat(coef: Optional[str]) -> str:
     chat.unload()
     gr.Info("Model unloaded.")
     if len(coef) != 230:
-        gr.Warning("Ingore invalid DVAE coefficient.")
+        gr.Warning("Ignore invalid DVAE coefficient.")
         coef = None
     try:
         global custom_path
@@ -111,7 +129,7 @@ def reload_chat(coef: Optional[str]) -> str:
         raise gr.Error(str(e))
     if not ret:
         raise gr.Error("Unable to load model.")
-    gr.Info("Reload succeess.")
+    gr.Info("Reload success.")
     return chat.coef
 
 

examples/web/webui.py

@@ -102,7 +102,7 @@ def main():
                 minimum=seed_min,
                 maximum=seed_max,
             )
-            generate_audio_seed = gr.Button("\U0001F3B2", interactive=True)
+            generate_audio_seed = gr.Button("\U0001f3b2", interactive=True)
             text_seed_input = gr.Number(
                 value=ex[0][5],
                 label="Text Seed",
@@ -110,20 +110,20 @@ def main():
                 minimum=seed_min,
                 maximum=seed_max,
             )
-            generate_text_seed = gr.Button("\U0001F3B2", interactive=True)
+            generate_text_seed = gr.Button("\U0001f3b2", interactive=True)
 
         with gr.Row():
             spk_emb_text = gr.Textbox(
                 label="Speaker Embedding",
                 max_lines=3,
-                show_copy_button=True,
+                buttons=["copy"],
                 interactive=True,
                 scale=2,
             )
             dvae_coef_text = gr.Textbox(
                 label="DVAE Coefficient",
                 max_lines=3,
-                show_copy_button=True,
+                buttons=["copy"],
                 interactive=True,
                 scale=2,
             )
@@ -161,7 +161,7 @@ def main():
         text_output = gr.Textbox(
             label="Output Text",
             interactive=False,
-            show_copy_button=True,
+            buttons=["copy"],
         )
 
         sample_audio_input.change(
@@ -261,8 +261,14 @@ def main():
     parser.add_argument("--root_path", type=str, help="root path")
     parser.add_argument("--custom_path", type=str, help="custom model path")
     parser.add_argument("--coef", type=str, help="custom dvae coefficient")
+    parser.add_argument(
+        "--disable_cache", action="store_true", help="enable model cache"
+    )
+    parser.add_argument(
+        "--experimental", action="store_true", help="enable model cache"
+    )
     args = parser.parse_args()
-
+    set_params(not args.disable_cache, args.experimental)
     logger.info("loading ChatTTS model...")
 
     if load_chat(args.custom_path, args.coef):
@@ -279,7 +285,7 @@ def main():
         server_port=args.server_port,
         root_path=args.root_path,
         inbrowser=True,
-        show_api=False,
+        footer_links=["api", "gradio", "settings"],
     )
 
 

requirements.txt

@@ -1,4 +1,4 @@
-numpy<2.0.0
+numpy<3.0.0
 numba
 torch>=2.1.0
 torchaudio
@@ -14,3 +14,4 @@ WeTextProcessing; sys_platform == 'linux'
 nemo_text_processing; sys_platform == 'linux'
 av
 pydub
+requests

setup.py

@@ -20,7 +20,7 @@ setup(
     license="AGPLv3+",
     install_requires=[
         "numba",
-        "numpy<2.0.0",
+        "numpy<3.0.0",
         "pybase16384",
         "torch>=2.1.0",
         "torchaudio",

tests/#655.py

@@ -40,10 +40,10 @@ refined_text = chat.infer(
     ),
     split_text=False,
 )
-if (
-    refined_text[0]
-    != "what is [uv_break] your favorite english [uv_break] food [laugh] like [lbreak]"
-):
+if refined_text[0] not in [
+    "what is [uv_break] your favorite english [uv_break] food [laugh] like [lbreak]",
+    "like what is [uv_break] your favorite english food [laugh] [lbreak]",
+]:
     fail = True
     logger.warning("refined text is '%s'", refined_text[0])
 

tools/audio/__init__.py

@@ -1,4 +1,4 @@
 from .av import load_audio
-from .pcm import pcm_arr_to_mp3_view
+from .pcm import pcm_arr_to_mp3_view, pcm_arr_to_ogg_view, pcm_arr_to_wav_view
 from .ffmpeg import has_ffmpeg_installed
 from .np import float_to_int16

tools/audio/av.py

@@ -7,7 +7,6 @@ from av.audio.frame import AudioFrame
 from av.audio.resampler import AudioResampler
 import numpy as np
 
-
 video_format_dict: Dict[str, str] = {
     "m4a": "mp4",
 }

tools/audio/pcm.py

@@ -1,21 +1,91 @@
 import wave
 from io import BytesIO
-
 import numpy as np
-
 from .np import float_to_int16
 from .av import wav2
 
 
-def pcm_arr_to_mp3_view(wav: np.ndarray):
+def _pcm_to_wav_buffer(wav: np.ndarray, sample_rate: int = 24000) -> BytesIO:
+    """
+    Convert PCM audio data to a WAV format byte stream (internal utility function).
+
+    :param wav: PCM data, NumPy array, typically in float32 format.
+    :param sample_rate: Sample rate (in Hz), defaults to 24000.
+    :return: WAV format byte stream, stored in a BytesIO object.
+    """
+    # Create an in-memory byte stream buffer
     buf = BytesIO()
+
+    # Open a WAV file stream in write mode
     with wave.open(buf, "wb") as wf:
-        wf.setnchannels(1)  # Mono channel
-        wf.setsampwidth(2)  # Sample width in bytes
-        wf.setframerate(24000)  # Sample rate in Hz
+        # Set number of channels to 1 (mono)
+        wf.setnchannels(1)
+        # Set sample width to 2 bytes (16-bit)
+        wf.setsampwidth(2)
+        # Set sample rate
+        wf.setframerate(sample_rate)
+        # Convert PCM to 16-bit integer and write
         wf.writeframes(float_to_int16(wav))
+
+    # Reset buffer pointer to the beginning
     buf.seek(0, 0)
+    return buf
+
+
+def pcm_arr_to_mp3_view(wav: np.ndarray, sample_rate: int = 24000) -> memoryview:
+    """
+    Convert PCM audio data to MP3 format.
+
+    :param wav: PCM data, NumPy array, typically in float32 format.
+    :param sample_rate: Sample rate (in Hz), defaults to 24000.
+    :return: MP3 format byte data, returned as a memoryview.
+    """
+    # Get WAV format byte stream
+    buf = _pcm_to_wav_buffer(wav, sample_rate)
+
+    # Create output buffer
     buf2 = BytesIO()
+    # Convert WAV data to MP3
     wav2(buf, buf2, "mp3")
-    buf.seek(0, 0)
+    # Return MP3 data
     return buf2.getbuffer()
+
+
+def pcm_arr_to_ogg_view(wav: np.ndarray, sample_rate: int = 24000) -> memoryview:
+    """
+    Convert PCM audio data to OGG format (using Vorbis encoding).
+
+    :param wav: PCM data, NumPy array, typically in float32 format.
+    :param sample_rate: Sample rate (in Hz), defaults to 24000.
+    :return: OGG format byte data, returned as a memoryview.
+    """
+    # Get WAV format byte stream
+    buf = _pcm_to_wav_buffer(wav, sample_rate)
+
+    # Create output buffer
+    buf2 = BytesIO()
+    # Convert WAV data to OGG
+    wav2(buf, buf2, "ogg")
+    # Return OGG data
+    return buf2.getbuffer()
+
+
+def pcm_arr_to_wav_view(
+    wav: np.ndarray, sample_rate: int = 24000, include_header: bool = True
+) -> memoryview:
+    """
+    Convert PCM audio data to WAV format, with an option to include header.
+
+    :param wav: PCM data, NumPy array, typically in float32 format.
+    :param sample_rate: Sample rate (in Hz), defaults to 24000.
+    :param include_header: Whether to include WAV header, defaults to True.
+    :return: WAV format or raw PCM byte data, returned as a memoryview.
+    """
+    if include_header:
+        # Get complete WAV byte stream
+        buf = _pcm_to_wav_buffer(wav, sample_rate)
+        return buf.getbuffer()
+    else:
+        # Return only converted 16-bit PCM data
+        pcm_data = float_to_int16(wav)
+        return memoryview(pcm_data.tobytes())