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

@@ -4,7 +4,7 @@ on:
 
 jobs:
   checksum:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     steps:
       - uses: actions/checkout@v4
 
@@ -13,7 +13,7 @@ jobs:
 
       - name: Run RVC-Models-Downloader
         run: |
-          wget https://github.com/fumiama/RVC-Models-Downloader/releases/download/v0.2.6/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/close-issue.yml

@@ -5,14 +5,14 @@ on:
 
 jobs:
   close-issues:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     permissions:
       issues: write
       pull-requests: write
     steps:
       - uses: actions/stale@v5
         with:
-          exempt-issue-labels: "help wanted,good first issue,documentation,following up,todo list"
+          exempt-issue-labels: "help wanted,following up,todo list,enhancement,algorithm,delayed,performance"
           days-before-issue-stale: 30
           days-before-issue-close: 15
           stale-issue-label: "stale"

.github/workflows/pull-format.yml

@@ -8,7 +8,7 @@ jobs:
   # This workflow closes invalid PR
   change-or-close-pr:
     # The type of runner that the job will run on
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     permissions: write-all
 
     # Steps represent a sequence of tasks that will be executed as part of the job
@@ -63,6 +63,14 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
 
+      - name: Create venv
+        run: python3 -m venv .venv
+
+      - name: Activate venv
+        run: |
+          . .venv/bin/activate
+          echo PATH=$PATH >> $GITHUB_ENV
+
       - name: Install Black
         run: pip install "black[jupyter]"
 

.github/workflows/push-format.yml

@@ -24,6 +24,14 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
 
+      - name: Create venv
+        run: python3 -m venv .venv
+
+      - name: Activate venv
+        run: |
+          . .venv/bin/activate
+          echo PATH=$PATH >> $GITHUB_ENV
+
       - name: Install Black
         run: pip install "black[jupyter]"
 

.github/workflows/unitest.yml

@@ -20,11 +20,20 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
+          cache: 'pip'
 
       - name: Install Dependents
         run: |
           sudo apt-get install -y portaudio19-dev python3-pyaudio
 
+      - name: Create venv
+        run: python3 -m venv .venv
+
+      - name: Activate venv
+        run: |
+          . .venv/bin/activate
+          echo PATH=$PATH >> $GITHUB_ENV
+
       - name: Test Install
         run: pip install .
 

.gitignore

@@ -160,13 +160,12 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+.idea/
 
 # MacOS System
 .DS_Store
 
 # assets and configs of ChatTTS
-
 /asset
 /config
 

ChatTTS/core.py

@@ -1,24 +1,22 @@
 import os
+import re
 import logging
 import tempfile
 from dataclasses import dataclass, asdict
-from typing import Literal, Optional, List, Tuple, Dict
+from typing import Literal, Optional, List, Tuple, Dict, Union
 from json import load
 from pathlib import Path
-import lzma
 
 import numpy as np
 import torch
-import torch.nn.functional as F
-from omegaconf import OmegaConf
 from vocos import Vocos
 from vocos.pretrained import instantiate_class
 from huggingface_hub import snapshot_download
-import pybase16384 as b14
 
 from .config import Config
-from .model import DVAE, GPT, gen_logits, Tokenizer
+from .model import DVAE, Embed, GPT, gen_logits, Tokenizer, Speaker
 from .utils import (
+    load_safetensors,
     check_all_assets,
     download_all_assets,
     select_device,
@@ -26,6 +24,7 @@ from .utils import (
     del_all,
 )
 from .utils import logger as utils_logger
+from .utils import FileLike
 
 from .norm import Normalizer
 
@@ -50,7 +49,7 @@ class Chat:
 
     def has_loaded(self, use_decoder=False):
         not_finish = False
-        check_list = ["vocos", "gpt", "tokenizer"]
+        check_list = ["vocos", "gpt", "tokenizer", "embed"]
 
         if use_decoder:
             check_list.append("decoder")
@@ -62,25 +61,22 @@ class Chat:
                 self.logger.warning(f"{module} not initialized.")
                 not_finish = True
 
-        if not not_finish:
-            self.logger.info("all models has been initialized.")
-
         return not not_finish
 
     def download_models(
         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 = 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
                 ):
@@ -89,10 +85,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
@@ -103,17 +109,18 @@ class Chat:
                 )
                 try:
                     download_path = snapshot_download(
-                        repo_id="2Noise/ChatTTS", allow_patterns=["*.pt", "*.yaml"]
+                        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):
@@ -121,17 +128,24 @@ 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(
         self,
         source: Literal["huggingface", "local", "custom"] = "local",
         force_redownload=False,
-        compile: bool = True,
-        custom_path: Optional[torch.serialization.FILE_LIKE] = None,
+        compile: bool = False,
+        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:
@@ -141,6 +155,9 @@ class Chat:
             compile=compile,
             coef=coef,
             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()
@@ -152,38 +169,17 @@ class Chat:
         self.normalizer.destroy()
         del self.normalizer
         del self.sha256_map
-        del_list = ["vocos", "gpt", "decoder", "dvae", "tokenizer"]
+        del_list = ["vocos", "gpt", "decoder", "dvae", "tokenizer", "embed"]
         for module in del_list:
             if hasattr(self, module):
                 delattr(self, module)
         self.__init__(logger)
 
     def sample_random_speaker(self) -> str:
-        return self._encode_spk_emb(self._sample_random_speaker())
+        return self.speaker.sample_random()
 
-    @staticmethod
-    @torch.no_grad()
-    def _encode_spk_emb(spk_emb: torch.Tensor) -> str:
-        arr: np.ndarray = spk_emb.to(dtype=torch.float16, device="cpu").numpy()
-        s = b14.encode_to_string(
-            lzma.compress(
-                arr.tobytes(),
-                format=lzma.FORMAT_RAW,
-                filters=[{"id": lzma.FILTER_LZMA2, "preset": 9 | lzma.PRESET_EXTREME}],
-            ),
-        )
-        del arr
-        return s
-
-    @torch.no_grad()
-    def _sample_random_speaker(self) -> torch.Tensor:
-        dim: int = self.gpt.gpt.layers[0].mlp.gate_proj.in_features
-        spk = (
-            torch.randn(dim, device=self.std.device, dtype=self.std.dtype)
-            .mul_(self.std)
-            .add_(self.mean)
-        )
-        return spk
+    def sample_audio_speaker(self, wav: Union[np.ndarray, torch.Tensor]) -> str:
+        return self.speaker.encode_prompt(self.dvae.sample_audio(wav))
 
     @dataclass(repr=False, eq=False)
     class RefineTextParams:
@@ -196,11 +192,14 @@ class Chat:
         min_new_token: int = 0
         show_tqdm: bool = True
         ensure_non_empty: bool = True
+        manual_seed: Optional[int] = None
 
     @dataclass(repr=False, eq=False)
     class InferCodeParams(RefineTextParams):
         prompt: str = "[speed_5]"
         spk_emb: Optional[str] = None
+        spk_smp: Optional[str] = None
+        txt_smp: Optional[str] = None
         temperature: float = 0.3
         repetition_penalty: float = 1.05
         max_new_token: int = 2048
@@ -218,10 +217,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,
@@ -231,11 +252,22 @@ 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 = []
+            thr = np.float32(1e-5)
+            for wavs in res_gen:
+                for wav in wavs:
+                    stripped_wavs.append(wav[np.abs(wav) > thr])
+            if split_text:
+                return [np.concatenate(stripped_wavs)]
+            return stripped_wavs
         else:
             return next(res_gen)
 
@@ -248,17 +280,22 @@ class Chat:
         vocos_ckpt_path: str = None,
         dvae_ckpt_path: str = None,
         gpt_ckpt_path: str = None,
+        embed_path: str = None,
         decoder_ckpt_path: str = None,
         tokenizer_path: str = None,
         device: Optional[torch.device] = None,
-        compile: bool = True,
+        compile: bool = False,
         coef: Optional[str] = None,
         use_flash_attn=False,
+        use_vllm=False,
+        experimental: bool = False,
+        enable_cache=True,
     ):
         if device is None:
-            device = select_device()
+            device = select_device(experimental=experimental)
             self.logger.info("use device %s", str(device))
         self.device = device
+        self.device_gpt = device if "mps" not in str(device) else torch.device("cpu")
         self.compile = compile
 
         feature_extractor = instantiate_class(
@@ -269,75 +306,81 @@ class Chat:
         vocos = (
             Vocos(feature_extractor=feature_extractor, backbone=backbone, head=head)
             .to(
-                # vocos on mps will crash, use cpu fallback
+                # Vocos on mps will crash, use cpu fallback.
+                # Plus, complex dtype used in the decode process of Vocos is not supported in torch_npu now,
+                # so we put this calculation of data on CPU instead of NPU.
                 "cpu"
-                if "mps" in str(device)
+                if "mps" in str(device) or "npu" in str(device)
                 else device
             )
             .eval()
         )
         assert vocos_ckpt_path, "vocos_ckpt_path should not be None"
-        vocos.load_state_dict(torch.load(vocos_ckpt_path, weights_only=True, mmap=True))
+        vocos.load_state_dict(load_safetensors(vocos_ckpt_path))
         self.vocos = vocos
         self.logger.log(logging.INFO, "vocos loaded.")
 
-        dvae = (
-            DVAE(
-                decoder_config=asdict(self.config.dvae.decoder),
-                encoder_config=asdict(self.config.dvae.encoder),
-                vq_config=asdict(self.config.dvae.vq),
-                dim=self.config.dvae.decoder.idim,
-                coef=coef,
-            )
-            .to(device)
-            .eval()
+        # computation of MelSpectrogram on npu is not support now, use cpu fallback.
+        dvae_device = torch.device("cpu") if "npu" in str(self.device) else device
+        dvae = DVAE(
+            decoder_config=asdict(self.config.dvae.decoder),
+            encoder_config=asdict(self.config.dvae.encoder),
+            vq_config=asdict(self.config.dvae.vq),
+            dim=self.config.dvae.decoder.idim,
+            coef=coef,
+            device=dvae_device,
         )
         coef = str(dvae)
         assert dvae_ckpt_path, "dvae_ckpt_path should not be None"
-        dvae.load_state_dict(torch.load(dvae_ckpt_path, weights_only=True, mmap=True))
-        self.dvae = dvae
+        dvae.load_pretrained(dvae_ckpt_path, dvae_device)
+        self.dvae = dvae.eval()
         self.logger.log(logging.INFO, "dvae loaded.")
 
+        embed = Embed(
+            self.config.embed.hidden_size,
+            self.config.embed.num_audio_tokens,
+            self.config.embed.num_text_tokens,
+            self.config.embed.num_vq,
+        )
+        embed.load_pretrained(embed_path, device=device)
+        self.embed = embed.to(device)
+        self.logger.log(logging.INFO, "embed loaded.")
+
         gpt = GPT(
             gpt_config=asdict(self.config.gpt),
+            embed=self.embed,
             use_flash_attn=use_flash_attn,
+            use_vllm=use_vllm,
             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_state_dict(torch.load(gpt_ckpt_path, weights_only=True, mmap=True))
+        gpt.load_pretrained(gpt_ckpt_path, embed_path, experimental=experimental)
         gpt.prepare(compile=compile and "cuda" in str(device))
         self.gpt = gpt
-        spk_stat_path = os.path.join(os.path.dirname(gpt_ckpt_path), "spk_stat.pt")
-        assert os.path.exists(spk_stat_path), f"Missing spk_stat.pt: {spk_stat_path}"
-        spk_stat: torch.Tensor = torch.load(
-            spk_stat_path,
-            weights_only=True,
-            mmap=True,
-            map_location=device,
-        )
-        self.std, self.mean = spk_stat.requires_grad_(False).chunk(2)
         self.logger.log(logging.INFO, "gpt loaded.")
 
-        decoder = (
-            DVAE(
-                decoder_config=asdict(self.config.decoder),
-                dim=self.config.decoder.idim,
-                coef=coef,
-            )
-            .to(device)
-            .eval()
+        self.speaker = Speaker(
+            self.config.gpt.hidden_size, self.config.spk_stat, device
+        )
+        self.logger.log(logging.INFO, "speaker loaded.")
+
+        decoder = DVAE(
+            decoder_config=asdict(self.config.decoder),
+            dim=self.config.decoder.idim,
+            coef=coef,
+            device=device,
         )
         coef = str(decoder)
         assert decoder_ckpt_path, "decoder_ckpt_path should not be None"
-        decoder.load_state_dict(
-            torch.load(decoder_ckpt_path, weights_only=True, mmap=True)
-        )
-        self.decoder = decoder
+        decoder.load_pretrained(decoder_ckpt_path, device)
+        self.decoder = decoder.eval()
         self.logger.log(logging.INFO, "decoder loaded.")
 
         if tokenizer_path:
-            self.tokenizer = Tokenizer(tokenizer_path, device)
+            self.tokenizer = Tokenizer(tokenizer_path)
             self.logger.log(logging.INFO, "tokenizer loaded.")
 
         self.coef = coef
@@ -346,7 +389,7 @@ class Chat:
 
     def _infer(
         self,
-        text,
+        text: Union[List[str], str],
         stream=False,
         lang=None,
         skip_refine_text=False,
@@ -354,6 +397,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(),
     ):
@@ -373,6 +418,8 @@ class Chat:
             for t in text
         ]
 
+        self.logger.debug("normed texts %s", str(text))
+
         if not skip_refine_text:
             refined = self._refine_text(
                 text,
@@ -382,50 +429,87 @@ 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):
+                    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:
-        if "mps" in str(self.device):
+        if "mps" in str(self.device) or "npu" in str(self.device):
             return self.vocos.decode(spec.cpu()).cpu().numpy()
         else:
             return self.vocos.decode(spec).cpu().numpy()
@@ -459,45 +543,6 @@ class Chat:
         del mel_specs
         return wavs
 
-    @staticmethod
-    def _decode_spk_emb(spk_emb: str) -> np.ndarray:
-        return np.frombuffer(
-            lzma.decompress(
-                b14.decode_from_string(spk_emb),
-                format=lzma.FORMAT_RAW,
-                filters=[{"id": lzma.FILTER_LZMA2, "preset": 9 | lzma.PRESET_EXTREME}],
-            ),
-            dtype=np.float16,
-        ).copy()
-
-    @torch.no_grad()
-    def _apply_spk_emb(
-        self,
-        emb: torch.Tensor,
-        spk_emb: str,
-        input_ids: torch.Tensor,
-    ):
-        n = (
-            F.normalize(
-                torch.from_numpy(
-                    self._decode_spk_emb(spk_emb),
-                ),
-                p=2.0,
-                dim=0,
-                eps=1e-12,
-            )
-            .to(self.gpt.device_gpt)
-            .unsqueeze_(0)
-            .expand(emb.size(0), -1)
-            .unsqueeze_(1)
-            .expand(emb.shape)
-        )
-        cond = (
-            input_ids.narrow(-1, 0, 1).eq(self.tokenizer.spk_emb_ids).expand(emb.shape)
-        )
-        torch.where(cond, n, emb, out=emb)
-        del cond, n
-
     @torch.no_grad()
     def _infer_code(
         self,
@@ -516,41 +561,28 @@ class Chat:
         assert len(text), "text should not be empty"
 
         if not isinstance(params.temperature, list):
-            temperature = [params.temperature] * gpt.num_vq
+            temperature = [params.temperature] * self.config.gpt.num_vq
         else:
             temperature = params.temperature
 
-        for i, t in enumerate(text):
-            text[i] = (
-                t.replace("[Stts]", "")
-                .replace("[spk_emb]", "")
-                .replace("[empty_spk]", "")
-                .strip()
-            )
-            """
-            see https://github.com/2noise/ChatTTS/issues/459
-            """
-
-        if params.prompt:
-            text = [params.prompt + i for i in text]
-
-        if params.spk_emb is not None:
-            text = [f"[Stts][spk_emb]{i}[Ptts]" for i in text]
-        else:
-            text = [f"[Stts][empty_spk]{i}[Ptts]" for i in text]
-
         input_ids, attention_mask, text_mask = self.tokenizer.encode(
-            text, self.gpt.num_vq, gpt.device_gpt
+            self.speaker.decorate_code_prompts(
+                text,
+                params.prompt,
+                params.txt_smp,
+                params.spk_emb,
+            ),
+            self.config.gpt.num_vq,
+            prompt=(
+                self.speaker.decode_prompt(params.spk_smp)
+                if params.spk_smp is not None
+                else None
+            ),
+            device=self.device_gpt,
         )
+        start_idx = input_ids.shape[-2]
 
-        emb = gpt(input_ids, text_mask)
-
-        del text_mask
-
-        if params.spk_emb is not None:
-            self._apply_spk_emb(emb, params.spk_emb, input_ids)
-
-        num_code = int(gpt.emb_code[0].num_embeddings - 1)
+        num_code = self.config.gpt.num_audio_tokens - 1
 
         logits_warpers, logits_processors = gen_logits(
             num_code=num_code,
@@ -559,6 +591,58 @@ class Chat:
             repetition_penalty=params.repetition_penalty,
         )
 
+        if gpt.is_vllm:
+            from .model.velocity import SamplingParams
+
+            sample_params = SamplingParams(
+                temperature=temperature,
+                max_new_token=params.max_new_token,
+                max_tokens=8192,
+                min_new_token=params.min_new_token,
+                logits_processors=(logits_processors, logits_warpers),
+                eos_token=num_code,
+                infer_text=False,
+                start_idx=start_idx,
+            )
+            input_ids = [i.tolist() for i in input_ids]
+
+            result = gpt.llm.generate(
+                None,
+                sample_params,
+                input_ids,
+            )
+
+            token_ids = []
+            hidden_states = []
+            for i in result:
+                token_ids.append(torch.tensor(i.outputs[0].token_ids))
+                hidden_states.append(
+                    i.outputs[0].hidden_states.to(torch.float32).to(self.device)
+                )
+
+            del text_mask, input_ids
+
+            return [
+                GPT.GenerationOutputs(
+                    ids=token_ids,
+                    hiddens=hidden_states,
+                    attentions=[],
+                ),
+            ]
+
+        emb = self.embed(input_ids, text_mask)
+
+        del text_mask
+
+        if params.spk_emb is not None:
+            self.speaker.apply(
+                emb,
+                params.spk_emb,
+                input_ids,
+                self.tokenizer.spk_emb_ids,
+                self.gpt.device_gpt,
+            )
+
         result = gpt.generate(
             emb,
             input_ids,
@@ -567,20 +651,18 @@ class Chat:
             attention_mask=attention_mask,
             max_new_token=params.max_new_token,
             min_new_token=params.min_new_token,
-            logits_warpers=logits_warpers,
-            logits_processors=logits_processors,
+            logits_processors=(*logits_processors, *logits_warpers),
             infer_text=False,
             return_hidden=return_hidden,
             stream=stream,
             show_tqdm=params.show_tqdm,
             ensure_non_empty=params.ensure_non_empty,
             stream_batch=params.stream_batch,
+            manual_seed=params.manual_seed,
             context=self.context,
         )
 
         del emb, input_ids
-        del_all(logits_warpers)
-        del_all(logits_processors)
 
         return result
 
@@ -597,10 +679,10 @@ class Chat:
         if not isinstance(text, list):
             text = [text]
 
-        text = [f"[Sbreak]{i}[Pbreak]{params.prompt}" for i in text]
-
         input_ids, attention_mask, text_mask = self.tokenizer.encode(
-            text, self.gpt.num_vq, gpt.device_gpt
+            self.speaker.decorate_text_prompts(text, params.prompt),
+            self.config.gpt.num_vq,
+            device=self.device_gpt,
         )
 
         logits_warpers, logits_processors = gen_logits(
@@ -610,7 +692,43 @@ class Chat:
             repetition_penalty=params.repetition_penalty,
         )
 
-        emb = gpt(input_ids, text_mask)
+        if gpt.is_vllm:
+            from .model.velocity import SamplingParams
+
+            sample_params = SamplingParams(
+                repetition_penalty=params.repetition_penalty,
+                temperature=params.temperature,
+                top_p=params.top_P,
+                top_k=params.top_K,
+                max_new_token=params.max_new_token,
+                max_tokens=8192,
+                min_new_token=params.min_new_token,
+                logits_processors=(logits_processors, logits_warpers),
+                eos_token=self.tokenizer.eos_token,
+                infer_text=True,
+                start_idx=input_ids.shape[-2],
+            )
+            input_ids_list = [i.tolist() for i in input_ids]
+            del input_ids
+
+            result = gpt.llm.generate(
+                None, sample_params, input_ids_list, params.show_tqdm
+            )
+            token_ids = []
+            hidden_states = []
+            for i in result:
+                token_ids.append(torch.tensor(i.outputs[0].token_ids))
+                hidden_states.append(i.outputs[0].hidden_states)
+
+            del text_mask, input_ids_list, result
+
+            return GPT.GenerationOutputs(
+                ids=token_ids,
+                hiddens=hidden_states,
+                attentions=[],
+            )
+
+        emb = self.embed(input_ids, text_mask)
 
         del text_mask
 
@@ -623,18 +741,16 @@ class Chat:
                 attention_mask=attention_mask,
                 max_new_token=params.max_new_token,
                 min_new_token=params.min_new_token,
-                logits_warpers=logits_warpers,
-                logits_processors=logits_processors,
+                logits_processors=(*logits_processors, *logits_warpers),
                 infer_text=True,
                 stream=False,
                 show_tqdm=params.show_tqdm,
                 ensure_non_empty=params.ensure_non_empty,
+                manual_seed=params.manual_seed,
                 context=self.context,
             )
         )
 
         del emb, input_ids
-        del_all(logits_warpers)
-        del_all(logits_processors)
 
         return result

ChatTTS/model/__init__.py

@@ -1,4 +1,6 @@
 from .dvae import DVAE
+from .embed import Embed
 from .gpt import GPT
 from .processors import gen_logits
+from .speaker import Speaker
 from .tokenizer import Tokenizer

ChatTTS/model/cuda/patch.py

@@ -0,0 +1,18 @@
+import torch
+
+
+class LlamaRMSNorm(torch.nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        LlamaRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight.to(hidden_states.device) * hidden_states.to(input_dtype)

ChatTTS/model/cuda/te_llama.py

@@ -6,31 +6,27 @@
 #
 # Edited by fumiama.
 
-import os
 import re
-import gc
 from contextlib import contextmanager
+from typing import Dict
 
 import transformer_engine as te
 from transformer_engine.pytorch.attention import RotaryPositionEmbedding
-from transformer_engine.pytorch.fp8 import fp8_model_init
+
+import torch
 
 import transformers
 from transformers.models.llama.modeling_llama import (
     LlamaModel,
     LlamaConfig,
 )
-from transformers.modeling_utils import (
-    _add_variant,
-    load_state_dict,
-    _load_state_dict_into_model,
-)
-from transformers.utils import WEIGHTS_INDEX_NAME
-from transformers.utils.hub import get_checkpoint_shard_files
+from transformers.modeling_utils import _load_state_dict_into_model
+
+from .patch import LlamaRMSNorm
 
 
 @contextmanager
-def replace_decoder(te_decoder_cls):
+def replace_decoder(te_decoder_cls, llama_rms_norm_cls):
     """
     Replace `LlamaDecoderLayer` with custom `TELlamaDecoderLayer`.
     """
@@ -38,12 +34,17 @@ def replace_decoder(te_decoder_cls):
         transformers.models.llama.modeling_llama.LlamaDecoderLayer
     )
     transformers.models.llama.modeling_llama.LlamaDecoderLayer = te_decoder_cls
+    original_llama_rms_norm_cls = transformers.models.llama.modeling_llama.LlamaRMSNorm
+    transformers.models.llama.modeling_llama.LlamaRMSNorm = llama_rms_norm_cls
     try:
         yield
     finally:
         transformers.models.llama.modeling_llama.LlamaDecoderLayer = (
             original_llama_decoder_cls
         )
+        transformers.models.llama.modeling_llama.LlamaRMSNorm = (
+            original_llama_rms_norm_cls
+        )
 
 
 class TELlamaDecoderLayer(te.pytorch.TransformerLayer):
@@ -103,74 +104,29 @@ class TELlamaModel:
     """
 
     def __new__(cls, config: LlamaConfig):
-        with replace_decoder(te_decoder_cls=TELlamaDecoderLayer):
+        with replace_decoder(
+            te_decoder_cls=TELlamaDecoderLayer, llama_rms_norm_cls=LlamaRMSNorm
+        ):
             model = LlamaModel(config)
         return model
 
     @classmethod
-    def from_pretrained_local(
-        cls, pretrained_model_name_or_path, *args, config, **kwargs
+    def from_state_dict(
+        cls,
+        state_dict: Dict[str, torch.Tensor],
+        config: LlamaConfig,
     ):
         """
         Custom method adapted from `from_pretrained` method in HuggingFace
         Transformers repo: https://github.com/huggingface/transformers/blob/f497f564bb76697edab09184a252fc1b1a326d1e/src/transformers/modeling_utils.py#L2579
         """
-        vanilla_model = cls(config).to(kwargs["torch_dtype"])
-        subfolder = ""
-        variant = None
-        if os.path.isfile(
-            os.path.join(
-                pretrained_model_name_or_path,
-                subfolder,
-                _add_variant("model.safetensors.index.json", variant),
-            )
-        ):
-            # Load from a sharded PyTorch checkpoint
-            archive_file = os.path.join(
-                pretrained_model_name_or_path,
-                subfolder,
-                _add_variant("model.safetensors.index.json", variant),
-            )
-            is_sharded = True
-        elif os.path.isfile(
-            os.path.join(
-                pretrained_model_name_or_path,
-                subfolder,
-                _add_variant(WEIGHTS_INDEX_NAME, variant),
-            )
-        ):
-            # Load from a sharded PyTorch checkpoint
-            archive_file = os.path.join(
-                pretrained_model_name_or_path,
-                subfolder,
-                _add_variant(WEIGHTS_INDEX_NAME, variant),
-            )
-            is_sharded = True
-        else:
-            raise AssertionError(
-                "Only sharded PyTorch ckpt format supported at the moment"
-            )
 
-        resolved_archive_file, sharded_metadata = get_checkpoint_shard_files(
-            pretrained_model_name_or_path,
-            archive_file,
-        )
+        vanilla_model = cls(config)
 
-        # If the checkpoint is not sharded, it's a trivial sharding case
-        if not is_sharded:
-            assert not isinstance(resolved_archive_file, list)
-            resolved_archive_file = [resolved_archive_file]
-
-        for shard_file in resolved_archive_file:
-            state_dict = load_state_dict(shard_file)
-            # replace_params copies parameters relevant only to TransformerEngine
-            _replace_params(state_dict, vanilla_model.state_dict(), config)
-            # _load_state_dict_into_model copies parameters other than those in TransformerEngine
-            _load_state_dict_into_model(vanilla_model, state_dict, start_prefix="")
-
-            # Force mem release. Taken from huggingface code
-            del state_dict
-            gc.collect()
+        # replace_params copies parameters relevant only to TransformerEngine
+        _replace_params(state_dict, vanilla_model.state_dict(), config)
+        # _load_state_dict_into_model copies parameters other than those in TransformerEngine
+        _load_state_dict_into_model(vanilla_model, state_dict, start_prefix="")
 
         return vanilla_model
 

ChatTTS/model/dvae.py

@@ -1,14 +1,15 @@
 import math
-from typing import List, Optional, Literal, Tuple
+from typing import List, Optional, Literal, Union
 
 import numpy as np
 import pybase16384 as b14
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 import torchaudio
 from vector_quantize_pytorch import GroupedResidualFSQ
 
+from ..utils import load_safetensors
+
 
 class ConvNeXtBlock(nn.Module):
     def __init__(
@@ -36,7 +37,7 @@ class ConvNeXtBlock(nn.Module):
         )  # pointwise/1x1 convs, implemented with linear layers
         self.act = nn.GELU()
         self.pwconv2 = nn.Linear(intermediate_dim, dim)
-        self.gamma = (
+        self.weight = (
             nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
             if layer_scale_init_value > 0
             else None
@@ -55,8 +56,8 @@ class ConvNeXtBlock(nn.Module):
         del y
         y = self.pwconv2(x)
         del x
-        if self.gamma is not None:
-            y *= self.gamma
+        if self.weight is not None:
+            y *= self.weight
         y.transpose_(1, 2)  # (B, T, C) -> (B, C, T)
 
         x = y + residual
@@ -95,17 +96,14 @@ class GFSQ(nn.Module):
         feat = self.quantizer.get_output_from_indices(x)
         return feat.transpose_(1, 2) if self.transpose else feat
 
-    def __call__(
-        self, x: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    def __call__(self, x: torch.Tensor) -> torch.Tensor:
         return super().__call__(x)
 
-    def forward(
-        self, x: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self.transpose:
             x.transpose_(1, 2)
-        feat, ind = self.quantizer(x)
+        # feat, ind = self.quantizer(x)
+        _, ind = self.quantizer(x)
         """
         ind = rearrange(
             ind, "g b t r ->b t (g r)",
@@ -113,6 +111,7 @@ class GFSQ(nn.Module):
         """
         ind = ind.permute(1, 2, 0, 3).contiguous()
         ind = ind.view(ind.size(0), ind.size(1), -1)
+        """
         embed_onehot_tmp = F.one_hot(ind.long(), self.n_ind)
         embed_onehot = embed_onehot_tmp.to(x.dtype)
         del embed_onehot_tmp
@@ -121,12 +120,12 @@ class GFSQ(nn.Module):
         torch.div(e_mean, (e_mean.sum(dim=1) + self.eps).unsqueeze(1), out=e_mean)
         perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + self.eps), dim=1))
 
-        return (
+        return 
             torch.zeros(perplexity.shape, dtype=x.dtype, device=x.device),
             feat.transpose_(1, 2) if self.transpose else feat,
             perplexity,
-            ind.transpose_(1, 2) if self.transpose else ind,
-        )
+        """
+        return ind.transpose_(1, 2) if self.transpose else ind
 
 
 class DVAEDecoder(nn.Module):
@@ -181,8 +180,10 @@ class MelSpectrogramFeatures(torch.nn.Module):
         hop_length=256,
         n_mels=100,
         padding: Literal["center", "same"] = "center",
+        device: torch.device = torch.device("cpu"),
     ):
         super().__init__()
+        self.device = device
         if padding not in ["center", "same"]:
             raise ValueError("Padding must be 'center' or 'same'.")
         self.padding = padding
@@ -199,6 +200,7 @@ class MelSpectrogramFeatures(torch.nn.Module):
         return super().__call__(audio)
 
     def forward(self, audio: torch.Tensor) -> torch.Tensor:
+        audio = audio.to(self.device)
         mel: torch.Tensor = self.mel_spec(audio)
         features = torch.log(torch.clip(mel, min=1e-5))
         return features
@@ -212,13 +214,14 @@ class DVAE(nn.Module):
         vq_config: Optional[dict] = None,
         dim=512,
         coef: Optional[str] = None,
+        device: torch.device = torch.device("cpu"),
     ):
         super().__init__()
         if coef is None:
             coef = torch.rand(100)
         else:
             coef = torch.from_numpy(
-                np.copy(np.frombuffer(b14.decode_from_string(coef), dtype=np.float32))
+                np.frombuffer(b14.decode_from_string(coef), dtype=np.float32).copy()
             )
         self.register_buffer("coef", coef.unsqueeze(0).unsqueeze_(2))
 
@@ -229,7 +232,7 @@ class DVAE(nn.Module):
                 nn.Conv1d(dim, dim, 4, 2, 1),
                 nn.GELU(),
             )
-            self.preprocessor_mel = MelSpectrogramFeatures()
+            self.preprocessor_mel = MelSpectrogramFeatures(device=device)
             self.encoder: Optional[DVAEDecoder] = DVAEDecoder(**encoder_config)
 
         self.decoder = DVAEDecoder(**decoder_config)
@@ -249,15 +252,25 @@ class DVAE(nn.Module):
     ) -> torch.Tensor:
         return super().__call__(inp, mode)
 
+    @torch.inference_mode()
+    def load_pretrained(self, filename: str, device: torch.device):
+        state_dict_tensors = load_safetensors(filename)
+        self.load_state_dict(state_dict_tensors)
+        self.to(device)
+
     @torch.inference_mode()
     def forward(
         self, inp: torch.Tensor, mode: Literal["encode", "decode"] = "decode"
     ) -> torch.Tensor:
         if mode == "encode" and hasattr(self, "encoder") and self.vq_layer is not None:
             mel = self.preprocessor_mel(inp)
-            x: torch.Tensor = self.downsample_conv(mel / self.coef)
+            x: torch.Tensor = self.downsample_conv(
+                torch.div(mel, self.coef.view(100, 1).expand(mel.shape), out=mel),
+            ).unsqueeze_(0)
+            del mel
             x = self.encoder(x)
-            ind = self.vq_layer(x)[3]
+            ind = self.vq_layer(x)
+            del x
             return ind
 
         if self.vq_layer is not None:
@@ -282,3 +295,9 @@ class DVAE(nn.Module):
         del vq_feats
 
         return torch.mul(dec_out, self.coef, out=dec_out)
+
+    @torch.inference_mode()
+    def sample_audio(self, wav: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
+        if isinstance(wav, np.ndarray):
+            wav = torch.from_numpy(wav)
+        return self(wav, "encode").squeeze_(0)

ChatTTS/model/embed.py

@@ -0,0 +1,82 @@
+import torch
+import torch.nn as nn
+from torch.nn.utils.parametrizations import weight_norm
+
+from ..utils import load_safetensors
+
+
+class Embed(nn.Module):
+    def __init__(
+        self, hidden_size: int, num_audio_tokens: int, num_text_tokens: int, num_vq=4
+    ):
+        super().__init__()
+
+        self.num_vq = num_vq
+        self.num_audio_tokens = num_audio_tokens
+
+        self.model_dim = hidden_size
+        self.emb_code = nn.ModuleList(
+            [nn.Embedding(num_audio_tokens, self.model_dim) for _ in range(num_vq)],
+        )
+        self.emb_text = nn.Embedding(num_text_tokens, self.model_dim)
+
+        self.head_text = weight_norm(
+            nn.Linear(self.model_dim, num_text_tokens, bias=False),
+            name="weight",
+        )
+        self.head_code = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Linear(self.model_dim, num_audio_tokens, bias=False),
+                    name="weight",
+                )
+                for _ in range(self.num_vq)
+            ],
+        )
+
+    @torch.inference_mode()
+    def load_pretrained(self, filename: str, device: torch.device):
+        state_dict_tensors = load_safetensors(filename)
+        self.load_state_dict(state_dict_tensors)
+        self.to(device)
+
+    def __call__(
+        self, input_ids: torch.Tensor, text_mask: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        get_emb
+        """
+        return super().__call__(input_ids, text_mask)
+
+    @torch.inference_mode()
+    def forward(self, input_ids: torch.Tensor, text_mask: torch.Tensor) -> torch.Tensor:
+        """
+        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_dev[text_mask_dev].narrow(1, 0, 1).squeeze_(1)
+        )
+
+        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)
+        ]
+        emb_code = torch.stack(emb_code, 2).sum(2)
+
+        emb = torch.zeros(
+            (input_ids_dev.shape[:-1]) + (emb_text.shape[-1],),
+            device=emb_text.device,
+            dtype=emb_text.dtype,
+        )
+        emb[text_mask_dev] = emb_text
+        emb[text_mask_inv] = emb_code.to(emb.dtype)
+
+        del emb_text, emb_code, text_mask_inv
+
+        return emb

ChatTTS/model/gpt.py

@@ -1,85 +1,107 @@
 import platform
 from dataclasses import dataclass
 import logging
-from typing import Union, List, Optional, Tuple
+from typing import Union, List, Optional, Tuple, Callable
+import gc
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.nn.utils.parametrize as P
-from torch.nn.utils.parametrizations import weight_norm
 from tqdm import tqdm
-from transformers import LlamaModel, LlamaConfig, LogitsWarper
+from transformers import LlamaModel, LlamaConfig
 from transformers.cache_utils import Cache
 from transformers.modeling_outputs import BaseModelOutputWithPast
 from transformers.utils import is_flash_attn_2_available
 
-from .processors import CustomRepetitionPenaltyLogitsProcessorRepeat
 from ..utils import del_all
+from .embed import Embed
 
 
 class GPT(nn.Module):
     def __init__(
         self,
         gpt_config: dict,
-        num_audio_tokens: int = 626,
-        num_text_tokens: int = 21178,
-        num_vq=4,
+        embed: Embed,
         use_flash_attn=False,
+        use_vllm=False,
         device=torch.device("cpu"),
+        device_gpt=torch.device("cpu"),
         logger=logging.getLogger(__name__),
+        enable_cache=True,
     ):
         super().__init__()
 
         self.logger = logger
 
         self.device = device
-        self.device_gpt = device if "mps" not in str(device) else torch.device("cpu")
+        self.device_gpt = device_gpt
 
-        self.num_vq = num_vq
-        self.num_audio_tokens = num_audio_tokens
+        self.enable_cache = enable_cache
+
+        self.generator = torch.Generator(device=device)
+
+        self.num_vq = int(gpt_config["num_vq"])
+        self.num_audio_tokens = int(gpt_config["num_audio_tokens"])
+        self.num_text_tokens = int(gpt_config["num_text_tokens"])
 
         self.use_flash_attn = use_flash_attn
+        self.is_te_llama = False
+        self.is_vllm = use_vllm
 
-        self.gpt = self._build_llama(gpt_config, self.device_gpt)
-        self.model_dim = int(self.gpt.config.hidden_size)
-        self.emb_code = nn.ModuleList(
-            [
-                nn.Embedding(
-                    num_audio_tokens,
-                    self.model_dim,
-                    device=self.device_gpt,
-                )
-                for _ in range(num_vq)
-            ],
-        )
-        self.emb_text = nn.Embedding(
-            num_text_tokens, self.model_dim, device=self.device_gpt
-        )
+        if self.is_vllm:
+            return
 
-        self.head_text = weight_norm(
-            nn.Linear(
-                self.model_dim,
-                num_text_tokens,
-                bias=False,
-                device=device,
-            ),
-            name="weight",
+        self.llama_config = self._build_llama_config(gpt_config)
+
+        self.emb_code = [ec.__call__ for ec in embed.emb_code]
+        self.emb_text = embed.emb_text.__call__
+        self.head_text = embed.head_text.__call__
+        self.head_code = [hc.__call__ for hc in embed.head_code]
+
+    def load_pretrained(
+        self, gpt_folder: str, embed_file_path: str, experimental=False
+    ):
+        if self.is_vllm and platform.system().lower() == "linux":
+
+            from .velocity import LLM
+
+            self.llm = LLM(
+                model=gpt_folder,
+                num_audio_tokens=self.num_audio_tokens,
+                num_text_tokens=self.num_text_tokens,
+                post_model_path=embed_file_path,
+            )
+            self.logger.info("vLLM model loaded")
+            return
+
+        self.gpt: LlamaModel = LlamaModel.from_pretrained(gpt_folder).to(
+            self.device_gpt
         )
-        self.head_code = nn.ModuleList(
-            [
-                weight_norm(
-                    nn.Linear(
-                        self.model_dim,
-                        num_audio_tokens,
-                        bias=False,
-                        device=device,
-                    ),
-                    name="weight",
+        del self.gpt.embed_tokens
+
+        if (
+            experimental
+            and "cuda" in str(self.device_gpt)
+            and platform.system().lower() == "linux"
+        ):  # is TELlamaModel
+            try:
+                from .cuda import TELlamaModel
+
+                self.logger.warning(
+                    "Linux with CUDA, try NVIDIA accelerated TELlamaModel because experimental is enabled"
+                )
+                state_dict = self.gpt.state_dict()
+                vanilla = TELlamaModel.from_state_dict(state_dict, self.llama_config)
+                # Force mem release. Taken from huggingface code
+                del state_dict, self.gpt
+                gc.collect()
+                self.gpt = vanilla
+                self.is_te_llama = True
+            except Exception as e:
+                self.logger.warning(
+                    f"use default LlamaModel for importing TELlamaModel error: {e}"
                 )
-                for _ in range(self.num_vq)
-            ],
-        )
 
     class Context:
         def __init__(self):
@@ -91,94 +113,38 @@ class GPT(nn.Module):
         def get(self) -> bool:
             return self._interrupt
 
-    def _build_llama(
+    def _build_llama_config(
         self,
         config: dict,
-        device: torch.device,
-    ) -> LlamaModel:
+    ) -> Tuple[LlamaModel, LlamaConfig]:
 
-        model = None
+        if self.use_flash_attn and is_flash_attn_2_available():
+            llama_config = LlamaConfig(
+                **config,
+                attn_implementation="flash_attention_2",
+            )
+            self.logger.warning(
+                "enabling flash_attention_2 may make gpt be even slower"
+            )
+        else:
+            llama_config = LlamaConfig(**config)
 
-        if "cuda" in str(device) and platform.system().lower() == "linux":
-            try:
-                from .cuda import TELlamaModel
-
-                model = TELlamaModel(LlamaConfig(**config))
-                self.logger.info("Linux with CUDA, try NVIDIA accelerated TELlamaModel")
-            except Exception as e:
-                model = None
-                self.logger.warning(
-                    f"use default LlamaModel for importing TELlamaModel error: {e}"
-                )
-
-        if model is None:
-            if self.use_flash_attn and is_flash_attn_2_available():
-                llama_config = LlamaConfig(
-                    **config,
-                    attn_implementation="flash_attention_2",
-                )
-                self.logger.warning(
-                    "enabling flash_attention_2 may make gpt be even slower"
-                )
-            else:
-                llama_config = LlamaConfig(**config)
-            model = LlamaModel(llama_config)
-        del model.embed_tokens
-
-        return model.to(device)
+        return llama_config
 
     def prepare(self, compile=False):
         if self.use_flash_attn and is_flash_attn_2_available():
             self.gpt = self.gpt.to(dtype=torch.float16)
-        if compile:
+        if compile and not self.is_te_llama and not self.is_vllm:
             try:
                 self.compile(backend="inductor", dynamic=True)
                 self.gpt.compile(backend="inductor", dynamic=True)
             except RuntimeError as e:
                 self.logger.warning(f"compile failed: {e}. fallback to normal mode.")
 
-    def __call__(
-        self, input_ids: torch.Tensor, text_mask: torch.Tensor
-    ) -> torch.Tensor:
-        """
-        get_emb
-        """
-        return super().__call__(input_ids, text_mask)
-
-    def forward(self, input_ids: torch.Tensor, text_mask: torch.Tensor) -> torch.Tensor:
-        """
-        get_emb
-        """
-
-        emb_text: torch.Tensor = self.emb_text(
-            input_ids[text_mask].narrow(1, 0, 1).squeeze_(1).to(self.device_gpt)
-        )
-
-        text_mask_inv = text_mask.logical_not().to(self.device_gpt)
-        masked_input_ids: torch.Tensor = input_ids[text_mask_inv].to(self.device_gpt)
-
-        emb_code = [
-            self.emb_code[i](masked_input_ids[:, i]) for i in range(self.num_vq)
-        ]
-        emb_code = torch.stack(emb_code, 2).sum(2)
-
-        emb = torch.zeros(
-            (input_ids.shape[:-1]) + (emb_text.shape[-1],),
-            device=emb_text.device,
-            dtype=emb_text.dtype,
-        )
-        emb[text_mask] = emb_text
-        emb[text_mask_inv] = emb_code.to(emb.dtype)
-
-        del emb_text, emb_code, text_mask_inv
-
-        return emb
-
     @dataclass(repr=False, eq=False)
     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
@@ -194,43 +160,51 @@ class GPT(nn.Module):
             if self.cache_position is not None:
                 self.cache_position = self.cache_position.to(device, dtype=dtype)
 
+    @torch.no_grad()
     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
         has_static_cache = False
         if past_key_values is None:
-            past_key_values = getattr(
-                self.gpt.layers[0].self_attn, "past_key_value", None
-            )
+            if hasattr(self.gpt.layers[0], "self_attn"):
+                past_key_values = getattr(
+                    self.gpt.layers[0].self_attn, "past_key_value", None
+                )
             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
@@ -240,8 +214,8 @@ class GPT(nn.Module):
                 attention_mask is not None
                 and attention_mask.shape[1] > input_ids.shape[1]
             ):
-                start = -(attention_mask.shape[1] - past_length)
-                input_ids = input_ids.narrow(1, start, -start)
+                start = attention_mask.shape[1] - past_length
+                input_ids = input_ids.narrow(1, -start, start)
             # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
             # input_ids based on the past_length.
             elif past_length < input_ids.shape[1]:
@@ -253,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
             ):
@@ -285,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
@@ -313,6 +287,7 @@ class GPT(nn.Module):
             del_all(self.attentions)
             del_all(self.hiddens)
 
+    @torch.no_grad()
     def _prepare_generation_outputs(
         self,
         inputs_ids: torch.Tensor,
@@ -322,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()
@@ -350,8 +330,9 @@ class GPT(nn.Module):
         attention_mask: Optional[torch.Tensor] = None,
         max_new_token=2048,
         min_new_token=0,
-        logits_warpers: List[LogitsWarper] = [],
-        logits_processors: List[CustomRepetitionPenaltyLogitsProcessorRepeat] = [],
+        logits_processors: Tuple[
+            Callable[[torch.LongTensor, torch.FloatTensor], torch.FloatTensor]
+        ] = (),
         infer_text=False,
         return_attn=False,
         return_hidden=False,
@@ -359,9 +340,12 @@ class GPT(nn.Module):
         show_tqdm=True,
         ensure_non_empty=True,
         stream_batch=24,
+        manual_seed: Optional[int] = None,
         context=Context(),
     ):
 
+        self.logger.debug("start generate")
+
         attentions: List[Optional[Tuple[torch.FloatTensor, ...]]] = []
         hiddens = []
         stream_iter = 0
@@ -371,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 = (
@@ -380,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],
@@ -388,10 +380,30 @@ 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
+        inputs_ids_buf = torch.zeros(
+            inputs_ids.size(0),
+            progress + max_new_token,
+            inputs_ids.size(2),
+            dtype=inputs_ids.dtype,
+            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)
 
         pbar: Optional[tqdm] = None
 
@@ -405,38 +417,47 @@ class GPT(nn.Module):
         past_key_values = None
 
         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=True,
             )
+            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(
@@ -449,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),
@@ -464,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()
 
@@ -475,22 +501,32 @@ class GPT(nn.Module):
                 logits = logits.permute(0, 2, 1)
                 logits = logits.reshape(-1, logits.size(2))
                 # logits_token = rearrange(inputs_ids[:, start_idx:], "b c n -> (b n) c")
-                inputs_ids_sliced = inputs_ids[:, start_idx:].permute(0, 2, 1)
+                inputs_ids_sliced = inputs_ids.narrow(
+                    1,
+                    start_idx,
+                    inputs_ids.size(1) - start_idx,
+                ).permute(0, 2, 1)
                 logits_token = inputs_ids_sliced.reshape(
                     inputs_ids_sliced.size(0) * inputs_ids_sliced.size(1),
                     -1,
                 ).to(self.device)
+                del inputs_ids_sliced
             else:
-                logits_token = inputs_ids[:, start_idx:, 0].to(self.device)
+                logits_token = (
+                    inputs_ids.narrow(
+                        1,
+                        start_idx,
+                        inputs_ids.size(1) - start_idx,
+                    )
+                    .narrow(2, 0, 1)
+                    .to(self.device)
+                )
 
             logits /= temperature
 
             for logitsProcessors in logits_processors:
                 logits = logitsProcessors(logits_token, logits)
 
-            for logitsWarpers in logits_warpers:
-                logits = logitsWarpers(logits_token, logits)
-
             del logits_token
 
             if i < min_new_token:
@@ -500,35 +536,47 @@ class GPT(nn.Module):
 
             del logits
 
-            idx_next = torch.multinomial(scores, num_samples=1).to(finish.device)
+            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:
+                idx_next = torch.multinomial(
+                    scores,
+                    num_samples=1,
+                    generator=self.generator.manual_seed(manual_seed),
+                ).to(finish.device)
 
             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)
                 finish_or = idx_next.eq(eos_token).any(1)
                 finish.logical_or_(finish_or)
                 del finish_or
-                inputs_ids_tmp = torch.cat([inputs_ids, idx_next.unsqueeze_(1)], 1)
+                inputs_ids_buf.narrow(1, progress, 1).copy_(idx_next.unsqueeze_(1))
             else:
                 finish_or = idx_next.eq(eos_token).any(1)
                 finish.logical_or_(finish_or)
                 del finish_or
-                inputs_ids_tmp = torch.cat(
-                    [
-                        inputs_ids,
-                        idx_next.unsqueeze_(-1).expand(-1, -1, self.num_vq),
-                    ],
-                    1,
+                inputs_ids_buf.narrow(1, progress, 1).copy_(
+                    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",
                     str([unexpected_idx.item() for unexpected_idx in finish.nonzero()]),
                 )
-                if ensure_non_empty:
+                if ensure_non_empty and manual_seed is None:
                     if show_tqdm:
                         pbar.close()
                     self.logger.warning("regenerate in order to ensure non-empty")
@@ -542,7 +590,7 @@ class GPT(nn.Module):
                         attention_mask_cache,
                         past_key_values,
                         idx_next,
-                        inputs_ids_tmp,
+                        inputs_ids_buf,
                     )
                     new_gen = self.generate(
                         emb,
@@ -552,7 +600,6 @@ class GPT(nn.Module):
                         attention_mask,
                         max_new_token,
                         min_new_token,
-                        logits_warpers,
                         logits_processors,
                         infer_text,
                         return_attn,
@@ -561,15 +608,19 @@ class GPT(nn.Module):
                         show_tqdm,
                         ensure_non_empty,
                         stream_batch,
+                        manual_seed,
                         context,
                     )
                     for result in new_gen:
                         yield result
+                    del inputs_ids
                 return
 
-            del inputs_ids
-            inputs_ids = inputs_ids_tmp
-            del inputs_ids_tmp, idx_next
+            self.logger.debug("start output")
+
+            del idx_next
+            progress += 1
+            inputs_ids = inputs_ids_buf.narrow(1, 0, progress)
 
             not_finished = finish.logical_not().to(end_idx.device)
             end_idx.add_(not_finished.int())
@@ -587,6 +638,8 @@ class GPT(nn.Module):
                     )
             del not_finished
 
+            self.logger.debug("finis output")
+
             if finish.all() or context.get():
                 break
 
@@ -604,7 +657,7 @@ class GPT(nn.Module):
                     f"incomplete result. hit max_new_token: {max_new_token}"
                 )
 
-        del finish
+        del finish, inputs_ids_buf
 
         yield self._prepare_generation_outputs(
             inputs_ids,

ChatTTS/model/speaker.py

@@ -0,0 +1,154 @@
+import lzma
+from typing import List, Optional, Union
+
+import pybase16384 as b14
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+class Speaker:
+    def __init__(self, dim: int, spk_cfg: str, device=torch.device("cpu")) -> None:
+        spk_stat = torch.from_numpy(
+            np.frombuffer(b14.decode_from_string(spk_cfg), dtype=np.float16).copy()
+        ).to(device=device)
+        self.std, self.mean = spk_stat.requires_grad_(False).chunk(2)
+        self.dim = dim
+
+    def sample_random(self) -> str:
+        return self._encode(self._sample_random())
+
+    @torch.inference_mode()
+    def apply(
+        self,
+        emb: torch.Tensor,
+        spk_emb: Union[str, torch.Tensor],
+        input_ids: torch.Tensor,
+        spk_emb_ids: int,
+        device: torch.device,
+        inplace: bool = True,
+    ) -> torch.Tensor:
+        if isinstance(spk_emb, str):
+            spk_emb_tensor = torch.from_numpy(self._decode(spk_emb))
+        else:
+            spk_emb_tensor = spk_emb
+        n = (
+            F.normalize(
+                spk_emb_tensor,
+                p=2.0,
+                dim=0,
+                eps=1e-12,
+            )
+            .to(device)
+            .unsqueeze_(0)
+            .expand(emb.size(0), -1)
+            .unsqueeze_(1)
+            .expand(emb.shape)
+        )
+        cond = input_ids.narrow(-1, 0, 1).eq(spk_emb_ids).expand(emb.shape)
+        out = torch.where(cond, n, emb, out=emb if inplace else None)
+        if inplace:
+            del cond, n
+        return out
+
+    @staticmethod
+    @torch.no_grad()
+    def decorate_code_prompts(
+        text: List[str],
+        prompt: str,
+        txt_smp: Optional[str],
+        spk_emb: Optional[str],
+    ) -> List[str]:
+        for i, t in enumerate(text):
+            text[i] = (
+                t.replace("[Stts]", "")
+                .replace("[spk_emb]", "")
+                .replace("[empty_spk]", "")
+                .strip()
+            )
+            """
+            see https://github.com/2noise/ChatTTS/issues/459
+            """
+
+        if prompt:
+            text = [prompt + i for i in text]
+
+        txt_smp = "" if txt_smp is None else txt_smp
+        if spk_emb is not None:
+            text = [f"[Stts][spk_emb]{txt_smp}{i}[Ptts]" for i in text]
+        else:
+            text = [f"[Stts][empty_spk]{txt_smp}{i}[Ptts]" for i in text]
+
+        return text
+
+    @staticmethod
+    @torch.no_grad()
+    def decorate_text_prompts(text: List[str], prompt: str) -> List[str]:
+        return [f"[Sbreak]{i}[Pbreak]{prompt}" for i in text]
+
+    @staticmethod
+    @torch.no_grad()
+    def encode_prompt(prompt: torch.Tensor) -> str:
+        arr: np.ndarray = prompt.cpu().numpy().astype(np.uint16)
+        shp = arr.shape
+        assert len(shp) == 2, "prompt must be a 2D tensor"
+        s = b14.encode_to_string(
+            np.array(shp, dtype="<u2").tobytes()
+            + lzma.compress(
+                arr.astype("<u2").tobytes(),
+                format=lzma.FORMAT_RAW,
+                filters=[{"id": lzma.FILTER_LZMA2, "preset": 9 | lzma.PRESET_EXTREME}],
+            ),
+        )
+        del arr
+        return s
+
+    @staticmethod
+    @torch.no_grad()
+    def decode_prompt(prompt: str) -> torch.Tensor:
+        dec = b14.decode_from_string(prompt)
+        shp = np.frombuffer(dec[:4], dtype="<u2")
+        p = np.frombuffer(
+            lzma.decompress(
+                dec[4:],
+                format=lzma.FORMAT_RAW,
+                filters=[{"id": lzma.FILTER_LZMA2, "preset": 9 | lzma.PRESET_EXTREME}],
+            ),
+            dtype="<u2",
+        ).copy()
+        del dec
+        return torch.from_numpy(p.astype(np.int32)).view(*shp)
+
+    @torch.no_grad()
+    def _sample_random(self) -> torch.Tensor:
+        spk = (
+            torch.randn(self.dim, device=self.std.device, dtype=self.std.dtype)
+            .mul_(self.std)
+            .add_(self.mean)
+        )
+        return spk
+
+    @staticmethod
+    @torch.no_grad()
+    def _encode(spk_emb: torch.Tensor) -> str:
+        arr: np.ndarray = spk_emb.to(dtype=torch.float16, device="cpu").numpy()
+        s = b14.encode_to_string(
+            lzma.compress(
+                arr.tobytes(),
+                format=lzma.FORMAT_RAW,
+                filters=[{"id": lzma.FILTER_LZMA2, "preset": 9 | lzma.PRESET_EXTREME}],
+            ),
+        )
+        del arr
+        return s
+
+    @staticmethod
+    def _decode(spk_emb: str) -> np.ndarray:
+        return np.frombuffer(
+            lzma.decompress(
+                b14.decode_from_string(spk_emb),
+                format=lzma.FORMAT_RAW,
+                filters=[{"id": lzma.FILTER_LZMA2, "preset": 9 | lzma.PRESET_EXTREME}],
+            ),
+            dtype=np.float16,
+        ).copy()

ChatTTS/model/tokenizer.py

@@ -5,21 +5,26 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false"
 https://stackoverflow.com/questions/62691279/how-to-disable-tokenizers-parallelism-true-false-warning
 """
 
-from typing import List, Tuple
+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, device: torch.device
+        self,
+        tokenizer_path: FileLike,
     ):
+        """
         tokenizer: BertTokenizerFast = torch.load(
             tokenizer_path, map_location=device, mmap=True
         )
+        # tokenizer.save_pretrained("asset/tokenizer", legacy_format=False)
+        """
+        tokenizer: BertTokenizerFast = BertTokenizerFast.from_pretrained(tokenizer_path)
         self._tokenizer = tokenizer
 
         self.len = len(tokenizer)
@@ -27,20 +32,33 @@ class Tokenizer:
         self.break_0_ids = tokenizer.convert_tokens_to_ids("[break_0]")
         self.eos_token = tokenizer.convert_tokens_to_ids("[Ebreak]")
 
-        self.decode = self._tokenizer.batch_decode
-
     @torch.inference_mode()
     def encode(
-        self, text: List[str], num_vq: int, device="cpu"
+        self,
+        text: List[str],
+        num_vq: int,
+        prompt: Optional[torch.Tensor] = None,
+        device="cpu",
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
 
         input_ids_lst = []
         attention_mask_lst = []
         max_input_ids_len = -1
         max_attention_mask_len = -1
+        prompt_size = 0
+
+        if prompt is not None:
+            assert prompt.size(0) == num_vq, "prompt dim 0 must equal to num_vq"
+            prompt_size = prompt.size(1)
+
         # 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))
@@ -52,6 +70,11 @@ class Tokenizer:
             attn_sz = attention_mask_lst[-1].size(0)
             if attn_sz > max_attention_mask_len:
                 max_attention_mask_len = attn_sz
+
+        if prompt is not None:
+            max_input_ids_len += prompt_size
+            max_attention_mask_len += prompt_size
+
         input_ids = torch.zeros(
             len(input_ids_lst),
             max_input_ids_len,
@@ -59,10 +82,15 @@ class Tokenizer:
             dtype=input_ids_lst[0].dtype,
         )
         for i in range(len(input_ids_lst)):
-            input_ids.narrow(0, i, 1).narrow(1, 0, input_ids_lst[i].size(0)).copy_(
+            input_ids.narrow(0, i, 1).narrow(
+                1,
+                max_input_ids_len - prompt_size - input_ids_lst[i].size(0),
+                input_ids_lst[i].size(0),
+            ).copy_(
                 input_ids_lst[i]
-            )
+            )  # left padding
         del_all(input_ids_lst)
+
         attention_mask = torch.zeros(
             len(attention_mask_lst),
             max_attention_mask_len,
@@ -70,12 +98,46 @@ class Tokenizer:
             dtype=attention_mask_lst[0].dtype,
         )
         for i in range(len(attention_mask_lst)):
-            attention_mask.narrow(0, i, 1).narrow(
-                1, 0, attention_mask_lst[i].size(0)
-            ).copy_(attention_mask_lst[i])
+            attn = attention_mask.narrow(0, i, 1)
+            attn.narrow(
+                1,
+                max_attention_mask_len - prompt_size - attention_mask_lst[i].size(0),
+                attention_mask_lst[i].size(0),
+            ).copy_(
+                attention_mask_lst[i]
+            )  # left padding
+            if prompt_size > 0:
+                attn.narrow(
+                    1,
+                    max_attention_mask_len - prompt_size,
+                    prompt_size,
+                ).fill_(1)
         del_all(attention_mask_lst)
 
-        text_mask = torch.ones(input_ids.shape, dtype=bool, device=device)
-        input_ids = input_ids.unsqueeze_(-1).expand(-1, -1, num_vq)
+        text_mask = attention_mask.bool()
+        new_input_ids = input_ids.unsqueeze_(-1).expand(-1, -1, num_vq).clone()
+        del input_ids
 
-        return input_ids, attention_mask, text_mask
+        if prompt_size > 0:
+            text_mask.narrow(1, max_input_ids_len - prompt_size, prompt_size).fill_(0)
+            prompt_t = prompt.t().unsqueeze_(0).expand(new_input_ids.size(0), -1, -1)
+            new_input_ids.narrow(
+                1,
+                max_input_ids_len - prompt_size,
+                prompt_size,
+            ).copy_(prompt_t)
+            del prompt_t
+
+        return new_input_ids, attention_mask, text_mask
+
+    @torch.inference_mode()
+    def decode(
+        self,
+        sequences: Union[List[int], List[List[int]]],
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: bool = None,
+        **kwargs,
+    ):
+        return self._tokenizer.batch_decode(
+            sequences, skip_special_tokens, clean_up_tokenization_spaces, **kwargs
+        )

ChatTTS/model/velocity/__init__.py

@@ -0,0 +1,2 @@
+from .llm import LLM
+from .sampling_params import SamplingParams

ChatTTS/model/velocity/block_manager.py

@@ -0,0 +1,296 @@
+"""A block manager that manages token blocks."""
+
+import enum
+from typing import Dict, List, Optional, Set, Tuple
+
+from vllm.block import PhysicalTokenBlock
+from .sequence import Sequence, SequenceGroup, SequenceStatus
+from vllm.utils import Device
+
+# Mapping: logical block number -> physical block.
+BlockTable = List[PhysicalTokenBlock]
+
+
+class BlockAllocator:
+    """Manages free physical token blocks for a device.
+
+    The allocator maintains a list of free blocks and allocates a block when
+    requested. When a block is freed, its reference count is decremented. If
+    the reference count becomes zero, the block is added back to the free list.
+    """
+
+    def __init__(
+        self,
+        device: Device,
+        block_size: int,
+        num_blocks: int,
+    ) -> None:
+        self.device = device
+        self.block_size = block_size
+        self.num_blocks = num_blocks
+
+        # Initialize the free blocks.
+        self.free_blocks: BlockTable = []
+        for i in range(num_blocks):
+            block = PhysicalTokenBlock(
+                device=device, block_number=i, block_size=block_size
+            )
+            self.free_blocks.append(block)
+
+    def allocate(self) -> PhysicalTokenBlock:
+        if not self.free_blocks:
+            raise ValueError("Out of memory! No free blocks are available.")
+        block = self.free_blocks.pop()
+        block.ref_count = 1
+        return block
+
+    def free(self, block: PhysicalTokenBlock) -> None:
+        if block.ref_count == 0:
+            raise ValueError(f"Double free! {block} is already freed.")
+        block.ref_count -= 1
+        if block.ref_count == 0:
+            self.free_blocks.append(block)
+
+    def get_num_free_blocks(self) -> int:
+        return len(self.free_blocks)
+
+
+class AllocStatus(enum.Enum):
+    """Result for BlockSpaceManager.can_allocate
+
+    1. Ok: seq_group can be allocated now.
+    2. Later: seq_group cannot be allocated.
+      The capacity of allocator is larger than seq_group required.
+    3. Never: seq_group can never be allocated.
+      The seq_group is too large to allocated in GPU.
+    """
+
+    OK = enum.auto()
+    LATER = enum.auto()
+    NEVER = enum.auto()
+
+
+class BlockSpaceManager:
+    """Manages the mapping between logical and physical token blocks."""
+
+    def __init__(
+        self,
+        block_size: int,
+        num_gpu_blocks: int,
+        num_cpu_blocks: int,
+        watermark: float = 0.01,
+        sliding_window: Optional[int] = None,
+    ) -> None:
+        self.block_size = block_size
+        self.num_total_gpu_blocks = num_gpu_blocks
+        self.num_total_cpu_blocks = num_cpu_blocks
+
+        self.block_sliding_window = None
+        if sliding_window is not None:
+            assert sliding_window % block_size == 0, (sliding_window, block_size)
+            self.block_sliding_window = sliding_window // block_size
+
+        self.watermark = watermark
+        assert watermark >= 0.0
+
+        self.watermark_blocks = int(watermark * num_gpu_blocks)
+        self.gpu_allocator = BlockAllocator(Device.GPU, block_size, num_gpu_blocks)
+        self.cpu_allocator = BlockAllocator(Device.CPU, block_size, num_cpu_blocks)
+        # Mapping: seq_id -> BlockTable.
+        self.block_tables: Dict[int, BlockTable] = {}
+
+    def can_allocate(self, seq_group: SequenceGroup) -> AllocStatus:
+        # FIXME(woosuk): Here we assume that all sequences in the group share
+        # the same prompt. This may not be true for preempted sequences.
+        seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0]
+        num_required_blocks = len(seq.logical_token_blocks)
+        if self.block_sliding_window is not None:
+            num_required_blocks = min(num_required_blocks, self.block_sliding_window)
+        num_free_gpu_blocks = self.gpu_allocator.get_num_free_blocks()
+
+        # Use watermark to avoid frequent cache eviction.
+        if self.num_total_gpu_blocks - num_required_blocks < self.watermark_blocks:
+            return AllocStatus.NEVER
+        if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks:
+            return AllocStatus.OK
+        else:
+            return AllocStatus.LATER
+
+    def allocate(self, seq_group: SequenceGroup) -> None:
+        # NOTE: Here we assume that all sequences in the group have the same
+        # prompt.
+        seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0]
+
+        # Allocate new physical token blocks that will store the prompt tokens.
+        block_table: BlockTable = []
+        for logical_idx in range(len(seq.logical_token_blocks)):
+            if (
+                self.block_sliding_window is not None
+                and logical_idx >= self.block_sliding_window
+            ):
+                block = block_table[logical_idx % self.block_sliding_window]
+            else:
+                block = self.gpu_allocator.allocate()
+            # Set the reference counts of the token blocks.
+            block.ref_count = seq_group.num_seqs()
+            block_table.append(block)
+
+        # Assign the block table for each sequence.
+        for seq in seq_group.get_seqs(status=SequenceStatus.WAITING):
+            self.block_tables[seq.seq_id] = block_table.copy()
+
+    def can_append_slot(self, seq_group: SequenceGroup) -> bool:
+        # Simple heuristic: If there is at least one free block
+        # for each sequence, we can append.
+        num_free_gpu_blocks = self.gpu_allocator.get_num_free_blocks()
+        num_seqs = seq_group.num_seqs(status=SequenceStatus.RUNNING)
+        return num_seqs <= num_free_gpu_blocks
+
+    def append_slot(self, seq: Sequence) -> Optional[Tuple[int, int]]:
+        """Allocate a physical slot for a new token."""
+        logical_blocks = seq.logical_token_blocks
+        block_table = self.block_tables[seq.seq_id]
+
+        if len(block_table) < len(logical_blocks):
+            if (
+                self.block_sliding_window
+                and len(block_table) >= self.block_sliding_window
+            ):
+                # reuse a block
+                block_table.append(
+                    block_table[len(block_table) % self.block_sliding_window]
+                )
+            else:
+                # The sequence has a new logical block.
+                # Allocate a new physical block.
+                block = self.gpu_allocator.allocate()
+                block_table.append(block)
+                return None
+
+        # We want to append the token to the last physical block.
+        last_block = block_table[-1]
+        assert last_block.device == Device.GPU
+        if last_block.ref_count == 1:
+            # Not shared with other sequences. Appendable.
+            return None
+        else:
+            # The last block is shared with other sequences.
+            # Copy on Write: Allocate a new block and copy the tokens.
+            new_block = self.gpu_allocator.allocate()
+            block_table[-1] = new_block
+            self.gpu_allocator.free(last_block)
+            return last_block.block_number, new_block.block_number
+
+    def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None:
+        # NOTE: fork does not allocate a new physical block.
+        # Thus, it is always safe from OOM.
+        src_block_table = self.block_tables[parent_seq.seq_id]
+        self.block_tables[child_seq.seq_id] = src_block_table.copy()
+        for block in src_block_table:
+            block.ref_count += 1
+
+    def _get_physical_blocks(
+        self, seq_group: SequenceGroup
+    ) -> List[PhysicalTokenBlock]:
+        # NOTE: Here, we assume that the physical blocks are only shared by
+        # the sequences in the same group.
+        blocks: Set[PhysicalTokenBlock] = set()
+        for seq in seq_group.get_seqs():
+            if seq.is_finished():
+                continue
+            blocks.update(self.block_tables[seq.seq_id])
+        return list(blocks)
+
+    def can_swap_in(self, seq_group: SequenceGroup) -> bool:
+        blocks = self._get_physical_blocks(seq_group)
+        num_swapped_seqs = seq_group.num_seqs(status=SequenceStatus.SWAPPED)
+        num_free_blocks = self.gpu_allocator.get_num_free_blocks()
+        # NOTE: Conservatively, we assume that every sequence will allocate
+        # at least one free block right after the swap-in.
+        # NOTE: This should match the logic in can_append_slot().
+        num_required_blocks = len(blocks) + num_swapped_seqs
+        return num_free_blocks - num_required_blocks >= self.watermark_blocks
+
+    def swap_in(self, seq_group: SequenceGroup) -> Dict[int, int]:
+        # CPU block -> GPU block.
+        mapping: Dict[PhysicalTokenBlock, PhysicalTokenBlock] = {}
+        for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED):
+            new_block_table: BlockTable = []
+            block_table = self.block_tables[seq.seq_id]
+
+            for cpu_block in block_table:
+                if cpu_block in mapping:
+                    gpu_block = mapping[cpu_block]
+                    gpu_block.ref_count += 1
+                else:
+                    gpu_block = self.gpu_allocator.allocate()
+                    mapping[cpu_block] = gpu_block
+                new_block_table.append(gpu_block)
+                # Free the CPU block swapped in to GPU.
+                self.cpu_allocator.free(cpu_block)
+            self.block_tables[seq.seq_id] = new_block_table
+
+        block_number_mapping = {
+            cpu_block.block_number: gpu_block.block_number
+            for cpu_block, gpu_block in mapping.items()
+        }
+        return block_number_mapping
+
+    def can_swap_out(self, seq_group: SequenceGroup) -> bool:
+        blocks = self._get_physical_blocks(seq_group)
+        return len(blocks) <= self.cpu_allocator.get_num_free_blocks()
+
+    def swap_out(self, seq_group: SequenceGroup) -> Dict[int, int]:
+        # GPU block -> CPU block.
+        mapping: Dict[PhysicalTokenBlock, PhysicalTokenBlock] = {}
+        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
+            new_block_table: BlockTable = []
+            block_table = self.block_tables[seq.seq_id]
+
+            for gpu_block in block_table:
+                if gpu_block in mapping:
+                    cpu_block = mapping[gpu_block]
+                    cpu_block.ref_count += 1
+                else:
+                    cpu_block = self.cpu_allocator.allocate()
+                    mapping[gpu_block] = cpu_block
+                new_block_table.append(cpu_block)
+                # Free the GPU block swapped out to CPU.
+                self.gpu_allocator.free(gpu_block)
+            self.block_tables[seq.seq_id] = new_block_table
+
+        block_number_mapping = {
+            gpu_block.block_number: cpu_block.block_number
+            for gpu_block, cpu_block in mapping.items()
+        }
+        return block_number_mapping
+
+    def _free_block_table(self, block_table: BlockTable) -> None:
+        for block in set(block_table):
+            if block.device == Device.GPU:
+                self.gpu_allocator.free(block)
+            else:
+                self.cpu_allocator.free(block)
+
+    def free(self, seq: Sequence) -> None:
+        if seq.seq_id not in self.block_tables:
+            # Already freed or haven't been scheduled yet.
+            return
+        block_table = self.block_tables[seq.seq_id]
+        self._free_block_table(block_table)
+        del self.block_tables[seq.seq_id]
+
+    def reset(self) -> None:
+        for block_table in self.block_tables.values():
+            self._free_block_table(block_table)
+        self.block_tables.clear()
+
+    def get_block_table(self, seq: Sequence) -> List[int]:
+        block_table = self.block_tables[seq.seq_id]
+        return [block.block_number for block in block_table]
+
+    def get_num_free_gpu_blocks(self) -> int:
+        return self.gpu_allocator.get_num_free_blocks()
+
+    def get_num_free_cpu_blocks(self) -> int:
+        return self.cpu_allocator.get_num_free_blocks()

ChatTTS/model/velocity/configs.py

@@ -0,0 +1,864 @@
+from typing import Optional, Union, Tuple
+import os
+
+import torch
+from transformers import PretrainedConfig
+
+from vllm.logger import init_logger
+from vllm.transformers_utils.config import get_config
+from vllm.utils import get_cpu_memory, is_hip
+
+import argparse
+import dataclasses
+from dataclasses import dataclass
+
+logger = init_logger(__name__)
+
+_GB = 1 << 30
+
+
+class ModelConfig:
+    """Configuration for the model.
+
+    Args:
+        model: Name or path of the huggingface model to use.
+        tokenizer: Name or path of the huggingface tokenizer to use.
+        tokenizer_mode: Tokenizer mode. "auto" will use the fast tokenizer if
+            available, and "slow" will always use the slow tokenizer.
+        trust_remote_code: Trust remote code (e.g., from HuggingFace) when
+            downloading the model and tokenizer.
+        download_dir: Directory to download and load the weights, default to the
+            default cache directory of huggingface.
+        load_format: The format of the model weights to load:
+            "auto" will try to load the weights in the safetensors format and
+                fall back to the pytorch bin format if safetensors format is
+                not available.
+            "pt" will load the weights in the pytorch bin format.
+            "safetensors" will load the weights in the safetensors format.
+            "npcache" will load the weights in pytorch format and store
+                a numpy cache to speed up the loading.
+            "dummy" will initialize the weights with random values, which is
+                mainly for profiling.
+        dtype: Data type for model weights and activations. The "auto" option
+            will use FP16 precision for FP32 and FP16 models, and BF16 precision
+            for BF16 models.
+        seed: Random seed for reproducibility.
+        revision: The specific model version to use. It can be a branch name,
+            a tag name, or a commit id. If unspecified, will use the default
+            version.
+        tokenizer_revision: The specific tokenizer version to use. It can be a
+            branch name, a tag name, or a commit id. If unspecified, will use
+            the default version.
+        max_model_len: Maximum length of a sequence (including prompt and
+            output). If None, will be derived from the model.
+        quantization: Quantization method that was used to quantize the model
+            weights. If None, we assume the model weights are not quantized.
+        enforce_eager: Whether to enforce eager execution. If True, we will
+            disable CUDA graph and always execute the model in eager mode.
+            If False, we will use CUDA graph and eager execution in hybrid.
+        max_context_len_to_capture: Maximum context len covered by CUDA graphs.
+            When a sequence has context length larger than this, we fall back
+            to eager mode.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        tokenizer: str,
+        tokenizer_mode: str,
+        trust_remote_code: bool,
+        download_dir: Optional[str],
+        load_format: str,
+        dtype: Union[str, torch.dtype],
+        seed: int,
+        revision: Optional[str] = None,
+        tokenizer_revision: Optional[str] = None,
+        max_model_len: Optional[int] = None,
+        quantization: Optional[str] = None,
+        enforce_eager: bool = False,
+        max_context_len_to_capture: Optional[int] = None,
+        num_audio_tokens: int = 1024,
+        num_text_tokens: int = 80,
+    ) -> None:
+        self.model = model
+        self.tokenizer = tokenizer
+        self.tokenizer_mode = tokenizer_mode
+        self.trust_remote_code = trust_remote_code
+        self.download_dir = download_dir
+        self.load_format = load_format
+        self.seed = seed
+        self.revision = revision
+        self.tokenizer_revision = tokenizer_revision
+        self.quantization = quantization
+        self.enforce_eager = enforce_eager
+        self.max_context_len_to_capture = max_context_len_to_capture
+        self.num_audio_tokens = num_audio_tokens
+        self.num_text_tokens = num_text_tokens
+
+        if os.environ.get("VLLM_USE_MODELSCOPE", "False").lower() == "true":
+            # download model from ModelScope hub,
+            # lazy import so that modelscope is not required for normal use.
+            from modelscope.hub.snapshot_download import (
+                snapshot_download,
+            )  # pylint: disable=C
+
+            model_path = snapshot_download(
+                model_id=model, cache_dir=download_dir, revision=revision
+            )
+            self.model = model_path
+            self.download_dir = model_path
+            self.tokenizer = model_path
+
+        self.hf_config = get_config(self.model, trust_remote_code, revision)
+        self.dtype = _get_and_verify_dtype(self.hf_config, dtype)
+        self.max_model_len = _get_and_verify_max_len(self.hf_config, max_model_len)
+        self._verify_load_format()
+        self._verify_tokenizer_mode()
+        self._verify_quantization()
+        self._verify_cuda_graph()
+
+    def _verify_load_format(self) -> None:
+        load_format = self.load_format.lower()
+        supported_load_format = ["auto", "pt", "safetensors", "npcache", "dummy"]
+        rocm_not_supported_load_format = []
+        if load_format not in supported_load_format:
+            raise ValueError(
+                f"Unknown load format: {self.load_format}. Must be one of "
+                "'auto', 'pt', 'safetensors', 'npcache', or 'dummy'."
+            )
+        if is_hip() and load_format in rocm_not_supported_load_format:
+            rocm_supported_load_format = [
+                f
+                for f in supported_load_format
+                if (f not in rocm_not_supported_load_format)
+            ]
+            raise ValueError(
+                f"load format '{load_format}' is not supported in ROCm. "
+                f"Supported load format are "
+                f"{rocm_supported_load_format}"
+            )
+
+        # TODO: Remove this check once HF updates the pt weights of Mixtral.
+        architectures = getattr(self.hf_config, "architectures", [])
+        if "MixtralForCausalLM" in architectures and load_format == "pt":
+            raise ValueError(
+                "Currently, the 'pt' format is not supported for Mixtral. "
+                "Please use the 'safetensors' format instead. "
+            )
+        self.load_format = load_format
+
+    def _verify_tokenizer_mode(self) -> None:
+        tokenizer_mode = self.tokenizer_mode.lower()
+        if tokenizer_mode not in ["auto", "slow"]:
+            raise ValueError(
+                f"Unknown tokenizer mode: {self.tokenizer_mode}. Must be "
+                "either 'auto' or 'slow'."
+            )
+        self.tokenizer_mode = tokenizer_mode
+
+    def _verify_quantization(self) -> None:
+        supported_quantization = ["awq", "gptq", "squeezellm"]
+        rocm_not_supported_quantization = ["awq"]
+        if self.quantization is not None:
+            self.quantization = self.quantization.lower()
+
+        # Parse quantization method from the HF model config, if available.
+        hf_quant_config = getattr(self.hf_config, "quantization_config", None)
+        if hf_quant_config is not None:
+            hf_quant_method = str(hf_quant_config["quant_method"]).lower()
+            if self.quantization is None:
+                self.quantization = hf_quant_method
+            elif self.quantization != hf_quant_method:
+                raise ValueError(
+                    "Quantization method specified in the model config "
+                    f"({hf_quant_method}) does not match the quantization "
+                    f"method specified in the `quantization` argument "
+                    f"({self.quantization})."
+                )
+
+        if self.quantization is not None:
+            if self.quantization not in supported_quantization:
+                raise ValueError(
+                    f"Unknown quantization method: {self.quantization}. Must "
+                    f"be one of {supported_quantization}."
+                )
+            if is_hip() and self.quantization in rocm_not_supported_quantization:
+                raise ValueError(
+                    f"{self.quantization} quantization is currently not supported "
+                    f"in ROCm."
+                )
+            logger.warning(
+                f"{self.quantization} quantization is not fully "
+                "optimized yet. The speed can be slower than "
+                "non-quantized models."
+            )
+
+    def _verify_cuda_graph(self) -> None:
+        if self.max_context_len_to_capture is None:
+            self.max_context_len_to_capture = self.max_model_len
+        self.max_context_len_to_capture = min(
+            self.max_context_len_to_capture, self.max_model_len
+        )
+
+    def verify_with_parallel_config(
+        self,
+        parallel_config: "ParallelConfig",
+    ) -> None:
+        total_num_attention_heads = self.hf_config.num_attention_heads
+        tensor_parallel_size = parallel_config.tensor_parallel_size
+        if total_num_attention_heads % tensor_parallel_size != 0:
+            raise ValueError(
+                f"Total number of attention heads ({total_num_attention_heads})"
+                " must be divisible by tensor parallel size "
+                f"({tensor_parallel_size})."
+            )
+
+        total_num_hidden_layers = self.hf_config.num_hidden_layers
+        pipeline_parallel_size = parallel_config.pipeline_parallel_size
+        if total_num_hidden_layers % pipeline_parallel_size != 0:
+            raise ValueError(
+                f"Total number of hidden layers ({total_num_hidden_layers}) "
+                "must be divisible by pipeline parallel size "
+                f"({pipeline_parallel_size})."
+            )
+
+    def get_sliding_window(self) -> Optional[int]:
+        return getattr(self.hf_config, "sliding_window", None)
+
+    def get_vocab_size(self) -> int:
+        return self.hf_config.vocab_size
+
+    def get_hidden_size(self) -> int:
+        return self.hf_config.hidden_size
+
+    def get_head_size(self) -> int:
+        # FIXME(woosuk): This may not be true for all models.
+        return self.hf_config.hidden_size // self.hf_config.num_attention_heads
+
+    def get_total_num_kv_heads(self) -> int:
+        """Returns the total number of KV heads."""
+        # For GPTBigCode & Falcon:
+        # NOTE: for falcon, when new_decoder_architecture is True, the
+        # multi_query flag is ignored and we use n_head_kv for the number of
+        # KV heads.
+        falcon_model_types = ["falcon", "RefinedWeb", "RefinedWebModel"]
+        new_decoder_arch_falcon = (
+            self.hf_config.model_type in falcon_model_types
+            and getattr(self.hf_config, "new_decoder_architecture", False)
+        )
+        if not new_decoder_arch_falcon and getattr(
+            self.hf_config, "multi_query", False
+        ):
+            # Multi-query attention, only one KV head.
+            # Currently, tensor parallelism is not supported in this case.
+            return 1
+
+        attributes = [
+            # For Falcon:
+            "n_head_kv",
+            "num_kv_heads",
+            # For LLaMA-2:
+            "num_key_value_heads",
+            # For ChatGLM:
+            "multi_query_group_num",
+        ]
+        for attr in attributes:
+            num_kv_heads = getattr(self.hf_config, attr, None)
+            if num_kv_heads is not None:
+                return num_kv_heads
+
+        # For non-grouped-query attention models, the number of KV heads is
+        # equal to the number of attention heads.
+        return self.hf_config.num_attention_heads
+
+    def get_num_kv_heads(self, parallel_config: "ParallelConfig") -> int:
+        """Returns the number of KV heads per GPU."""
+        total_num_kv_heads = self.get_total_num_kv_heads()
+        # If tensor parallelism is used, we divide the number of KV heads by
+        # the tensor parallel size. We will replicate the KV heads in the
+        # case where the number of KV heads is smaller than the tensor
+        # parallel size so each GPU has at least one KV head.
+        return max(1, total_num_kv_heads // parallel_config.tensor_parallel_size)
+
+    def get_num_layers(self, parallel_config: "ParallelConfig") -> int:
+        total_num_hidden_layers = self.hf_config.num_hidden_layers
+        return total_num_hidden_layers // parallel_config.pipeline_parallel_size
+
+
+class CacheConfig:
+    """Configuration for the KV cache.
+
+    Args:
+        block_size: Size of a cache block in number of tokens.
+        gpu_memory_utilization: Fraction of GPU memory to use for the
+            vLLM execution.
+        swap_space: Size of the CPU swap space per GPU (in GiB).
+    """
+
+    def __init__(
+        self,
+        block_size: int,
+        gpu_memory_utilization: float,
+        swap_space: int,
+        sliding_window: Optional[int] = None,
+    ) -> None:
+        self.block_size = block_size
+        self.gpu_memory_utilization = gpu_memory_utilization
+        self.swap_space_bytes = swap_space * _GB
+        self.sliding_window = sliding_window
+        self._verify_args()
+
+        # Will be set after profiling.
+        self.num_gpu_blocks = None
+        self.num_cpu_blocks = None
+
+    def _verify_args(self) -> None:
+        if self.gpu_memory_utilization > 1.0:
+            raise ValueError(
+                "GPU memory utilization must be less than 1.0. Got "
+                f"{self.gpu_memory_utilization}."
+            )
+
+    def verify_with_parallel_config(
+        self,
+        parallel_config: "ParallelConfig",
+    ) -> None:
+        total_cpu_memory = get_cpu_memory()
+        # FIXME(woosuk): Here, it is assumed that the GPUs in a tensor parallel
+        # group are in the same node. However, the GPUs may span multiple nodes.
+        num_gpus_per_node = parallel_config.tensor_parallel_size
+        cpu_memory_usage = self.swap_space_bytes * num_gpus_per_node
+
+        msg = (
+            f"{cpu_memory_usage / _GB:.2f} GiB out of "
+            f"the {total_cpu_memory / _GB:.2f} GiB total CPU memory is "
+            "allocated for the swap space."
+        )
+        if cpu_memory_usage > 0.7 * total_cpu_memory:
+            raise ValueError("Too large swap space. " + msg)
+        elif cpu_memory_usage > 0.4 * total_cpu_memory:
+            logger.warning("Possibly too large swap space. " + msg)
+
+
+class ParallelConfig:
+    """Configuration for the distributed execution.
+
+    Args:
+        pipeline_parallel_size: Number of pipeline parallel groups.
+        tensor_parallel_size: Number of tensor parallel groups.
+        worker_use_ray: Whether to use Ray for model workers. Will be set to
+            True if either pipeline_parallel_size or tensor_parallel_size is
+            greater than 1.
+    """
+
+    def __init__(
+        self,
+        pipeline_parallel_size: int,
+        tensor_parallel_size: int,
+        worker_use_ray: bool,
+        max_parallel_loading_workers: Optional[int] = None,
+    ) -> None:
+        self.pipeline_parallel_size = pipeline_parallel_size
+        self.tensor_parallel_size = tensor_parallel_size
+        self.worker_use_ray = worker_use_ray
+        self.max_parallel_loading_workers = max_parallel_loading_workers
+
+        self.world_size = pipeline_parallel_size * tensor_parallel_size
+        if self.world_size > 1:
+            self.worker_use_ray = True
+        self._verify_args()
+
+    def _verify_args(self) -> None:
+        if self.pipeline_parallel_size > 1:
+            raise NotImplementedError("Pipeline parallelism is not supported yet.")
+
+
+class SchedulerConfig:
+    """Scheduler configuration.
+
+    Args:
+        max_num_batched_tokens: Maximum number of tokens to be processed in
+            a single iteration.
+        max_num_seqs: Maximum number of sequences to be processed in a single
+            iteration.
+        max_model_len: Maximum length of a sequence (including prompt
+            and generated text).
+        max_paddings: Maximum number of paddings to be added to a batch.
+    """
+
+    def __init__(
+        self,
+        max_num_batched_tokens: Optional[int],
+        max_num_seqs: int,
+        max_model_len: int,
+        max_paddings: int,
+    ) -> None:
+        if max_num_batched_tokens is not None:
+            self.max_num_batched_tokens = max_num_batched_tokens
+        else:
+            # If max_model_len is too short, use 2048 as the default value for
+            # higher throughput.
+            self.max_num_batched_tokens = max(max_model_len, 2048)
+        self.max_num_seqs = max_num_seqs
+        self.max_model_len = max_model_len
+        self.max_paddings = max_paddings
+        self._verify_args()
+
+    def _verify_args(self) -> None:
+        if self.max_num_batched_tokens < self.max_model_len:
+            raise ValueError(
+                f"max_num_batched_tokens ({self.max_num_batched_tokens}) is "
+                f"smaller than max_model_len ({self.max_model_len}). "
+                "This effectively limits the maximum sequence length to "
+                "max_num_batched_tokens and makes vLLM reject longer "
+                "sequences. Please increase max_num_batched_tokens or "
+                "decrease max_model_len."
+            )
+        if self.max_num_batched_tokens < self.max_num_seqs:
+            raise ValueError(
+                f"max_num_batched_tokens ({self.max_num_batched_tokens}) must "
+                "be greater than or equal to max_num_seqs "
+                f"({self.max_num_seqs})."
+            )
+
+
+_STR_DTYPE_TO_TORCH_DTYPE = {
+    "half": torch.float16,
+    "float16": torch.float16,
+    "float": torch.float32,
+    "float32": torch.float32,
+    "bfloat16": torch.bfloat16,
+}
+
+_ROCM_NOT_SUPPORTED_DTYPE = ["float", "float32"]
+
+
+def _get_and_verify_dtype(
+    config: PretrainedConfig,
+    dtype: Union[str, torch.dtype],
+) -> torch.dtype:
+    # NOTE: getattr(config, "torch_dtype", torch.float32) is not correct
+    # because config.torch_dtype can be None.
+    config_dtype = getattr(config, "torch_dtype", None)
+    if config_dtype is None:
+        config_dtype = torch.float32
+
+    if isinstance(dtype, str):
+        dtype = dtype.lower()
+        if dtype == "auto":
+            if config_dtype == torch.float32:
+                # Following the common practice, we use float16 for float32
+                # models.
+                torch_dtype = torch.float16
+            else:
+                torch_dtype = config_dtype
+        else:
+            if dtype not in _STR_DTYPE_TO_TORCH_DTYPE:
+                raise ValueError(f"Unknown dtype: {dtype}")
+            torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype]
+    elif isinstance(dtype, torch.dtype):
+        torch_dtype = dtype
+    else:
+        raise ValueError(f"Unknown dtype: {dtype}")
+
+    if is_hip() and torch_dtype == torch.float32:
+        rocm_supported_dtypes = [
+            k
+            for k, v in _STR_DTYPE_TO_TORCH_DTYPE.items()
+            if (k not in _ROCM_NOT_SUPPORTED_DTYPE)
+        ]
+        raise ValueError(
+            f"dtype '{dtype}' is not supported in ROCm. "
+            f"Supported dtypes are {rocm_supported_dtypes}"
+        )
+
+    # Verify the dtype.
+    if torch_dtype != config_dtype:
+        if torch_dtype == torch.float32:
+            # Upcasting to float32 is allowed.
+            pass
+        elif config_dtype == torch.float32:
+            # Downcasting from float32 to float16 or bfloat16 is allowed.
+            pass
+        else:
+            # Casting between float16 and bfloat16 is allowed with a warning.
+            logger.warning(f"Casting {config_dtype} to {torch_dtype}.")
+
+    return torch_dtype
+
+
+def _get_and_verify_max_len(
+    hf_config: PretrainedConfig,
+    max_model_len: Optional[int],
+) -> int:
+    """Get and verify the model's maximum length."""
+    derived_max_model_len = float("inf")
+    possible_keys = [
+        # OPT
+        "max_position_embeddings",
+        # GPT-2
+        "n_positions",
+        # MPT
+        "max_seq_len",
+        # ChatGLM2
+        "seq_length",
+        # Others
+        "max_sequence_length",
+        "max_seq_length",
+        "seq_len",
+    ]
+    for key in possible_keys:
+        max_len_key = getattr(hf_config, key, None)
+        if max_len_key is not None:
+            derived_max_model_len = min(derived_max_model_len, max_len_key)
+    if derived_max_model_len == float("inf"):
+        if max_model_len is not None:
+            # If max_model_len is specified, we use it.
+            return max_model_len
+
+        default_max_len = 2048
+        logger.warning(
+            "The model's config.json does not contain any of the following "
+            "keys to determine the original maximum length of the model: "
+            f"{possible_keys}. Assuming the model's maximum length is "
+            f"{default_max_len}."
+        )
+        derived_max_model_len = default_max_len
+
+    rope_scaling = getattr(hf_config, "rope_scaling", None)
+    if rope_scaling is not None:
+        assert "factor" in rope_scaling
+        scaling_factor = rope_scaling["factor"]
+        if rope_scaling["type"] == "yarn":
+            derived_max_model_len = rope_scaling["original_max_position_embeddings"]
+        derived_max_model_len *= scaling_factor
+
+    if max_model_len is None:
+        max_model_len = derived_max_model_len
+    elif max_model_len > derived_max_model_len:
+        raise ValueError(
+            f"User-specified max_model_len ({max_model_len}) is greater than "
+            f"the derived max_model_len ({max_len_key}={derived_max_model_len}"
+            " in model's config.json). This may lead to incorrect model "
+            "outputs or CUDA errors. Make sure the value is correct and "
+            "within the model context size."
+        )
+    return int(max_model_len)
+
+
+@dataclass
+class EngineArgs:
+    """Arguments for vLLM engine."""
+
+    model: str
+    tokenizer: Optional[str] = None
+    tokenizer_mode: str = "auto"
+    trust_remote_code: bool = False
+    download_dir: Optional[str] = None
+    load_format: str = "auto"
+    dtype: str = "auto"
+    seed: int = 0
+    max_model_len: Optional[int] = None
+    worker_use_ray: bool = False
+    pipeline_parallel_size: int = 1
+    tensor_parallel_size: int = 1
+    max_parallel_loading_workers: Optional[int] = None
+    block_size: int = 16
+    swap_space: int = 4  # GiB
+    gpu_memory_utilization: float = 0.90
+    max_num_batched_tokens: Optional[int] = None
+    max_num_seqs: int = 256
+    max_paddings: int = 256
+    disable_log_stats: bool = False
+    revision: Optional[str] = None
+    tokenizer_revision: Optional[str] = None
+    quantization: Optional[str] = None
+    enforce_eager: bool = False
+    max_context_len_to_capture: int = 8192
+    num_audio_tokens: int = 1024
+    num_text_tokens: int = 80
+
+    def __post_init__(self):
+        if self.tokenizer is None:
+            self.tokenizer = self.model
+
+    @staticmethod
+    def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+        """Shared CLI arguments for vLLM engine."""
+
+        # NOTE: If you update any of the arguments below, please also
+        # make sure to update docs/source/models/engine_args.rst
+
+        # Model arguments
+        parser.add_argument(
+            "--model",
+            type=str,
+            default="facebook/opt-125m",
+            help="name or path of the huggingface model to use",
+        )
+        parser.add_argument(
+            "--tokenizer",
+            type=str,
+            default=EngineArgs.tokenizer,
+            help="name or path of the huggingface tokenizer to use",
+        )
+        parser.add_argument(
+            "--revision",
+            type=str,
+            default=None,
+            help="the specific model version to use. It can be a branch "
+            "name, a tag name, or a commit id. If unspecified, will use "
+            "the default version.",
+        )
+        parser.add_argument(
+            "--tokenizer-revision",
+            type=str,
+            default=None,
+            help="the specific tokenizer version to use. It can be a branch "
+            "name, a tag name, or a commit id. If unspecified, will use "
+            "the default version.",
+        )
+        parser.add_argument(
+            "--tokenizer-mode",
+            type=str,
+            default=EngineArgs.tokenizer_mode,
+            choices=["auto", "slow"],
+            help='tokenizer mode. "auto" will use the fast '
+            'tokenizer if available, and "slow" will '
+            "always use the slow tokenizer.",
+        )
+        parser.add_argument(
+            "--trust-remote-code",
+            action="store_true",
+            help="trust remote code from huggingface",
+        )
+        parser.add_argument(
+            "--download-dir",
+            type=str,
+            default=EngineArgs.download_dir,
+            help="directory to download and load the weights, "
+            "default to the default cache dir of "
+            "huggingface",
+        )
+        parser.add_argument(
+            "--load-format",
+            type=str,
+            default=EngineArgs.load_format,
+            choices=["auto", "pt", "safetensors", "npcache", "dummy"],
+            help="The format of the model weights to load. "
+            '"auto" will try to load the weights in the safetensors format '
+            "and fall back to the pytorch bin format if safetensors format "
+            "is not available. "
+            '"pt" will load the weights in the pytorch bin format. '
+            '"safetensors" will load the weights in the safetensors format. '
+            '"npcache" will load the weights in pytorch format and store '
+            "a numpy cache to speed up the loading. "
+            '"dummy" will initialize the weights with random values, '
+            "which is mainly for profiling.",
+        )
+        parser.add_argument(
+            "--dtype",
+            type=str,
+            default=EngineArgs.dtype,
+            choices=["auto", "half", "float16", "bfloat16", "float", "float32"],
+            help="data type for model weights and activations. "
+            'The "auto" option will use FP16 precision '
+            "for FP32 and FP16 models, and BF16 precision "
+            "for BF16 models.",
+        )
+        parser.add_argument(
+            "--max-model-len",
+            type=int,
+            default=None,
+            help="model context length. If unspecified, "
+            "will be automatically derived from the model.",
+        )
+        # Parallel arguments
+        parser.add_argument(
+            "--worker-use-ray",
+            action="store_true",
+            help="use Ray for distributed serving, will be "
+            "automatically set when using more than 1 GPU",
+        )
+        parser.add_argument(
+            "--pipeline-parallel-size",
+            "-pp",
+            type=int,
+            default=EngineArgs.pipeline_parallel_size,
+            help="number of pipeline stages",
+        )
+        parser.add_argument(
+            "--tensor-parallel-size",
+            "-tp",
+            type=int,
+            default=EngineArgs.tensor_parallel_size,
+            help="number of tensor parallel replicas",
+        )
+        parser.add_argument(
+            "--max-parallel-loading-workers",
+            type=int,
+            help="load model sequentially in multiple batches, "
+            "to avoid RAM OOM when using tensor "
+            "parallel and large models",
+        )
+        # KV cache arguments
+        parser.add_argument(
+            "--block-size",
+            type=int,
+            default=EngineArgs.block_size,
+            choices=[8, 16, 32],
+            help="token block size",
+        )
+        # TODO(woosuk): Support fine-grained seeds (e.g., seed per request).
+        parser.add_argument(
+            "--seed", type=int, default=EngineArgs.seed, help="random seed"
+        )
+        parser.add_argument(
+            "--swap-space",
+            type=int,
+            default=EngineArgs.swap_space,
+            help="CPU swap space size (GiB) per GPU",
+        )
+        parser.add_argument(
+            "--gpu-memory-utilization",
+            type=float,
+            default=EngineArgs.gpu_memory_utilization,
+            help="the fraction of GPU memory to be used for "
+            "the model executor, which can range from 0 to 1."
+            "If unspecified, will use the default value of 0.9.",
+        )
+        parser.add_argument(
+            "--max-num-batched-tokens",
+            type=int,
+            default=EngineArgs.max_num_batched_tokens,
+            help="maximum number of batched tokens per " "iteration",
+        )
+        parser.add_argument(
+            "--max-num-seqs",
+            type=int,
+            default=EngineArgs.max_num_seqs,
+            help="maximum number of sequences per iteration",
+        )
+        parser.add_argument(
+            "--max-paddings",
+            type=int,
+            default=EngineArgs.max_paddings,
+            help="maximum number of paddings in a batch",
+        )
+        parser.add_argument(
+            "--disable-log-stats",
+            action="store_true",
+            help="disable logging statistics",
+        )
+        # Quantization settings.
+        parser.add_argument(
+            "--quantization",
+            "-q",
+            type=str,
+            choices=["awq", "gptq", "squeezellm", None],
+            default=None,
+            help="Method used to quantize the weights. If "
+            "None, we first check the `quantization_config` "
+            "attribute in the model config file. If that is "
+            "None, we assume the model weights are not "
+            "quantized and use `dtype` to determine the data "
+            "type of the weights.",
+        )
+        parser.add_argument(
+            "--enforce-eager",
+            action="store_true",
+            help="Always use eager-mode PyTorch. If False, "
+            "will use eager mode and CUDA graph in hybrid "
+            "for maximal performance and flexibility.",
+        )
+        parser.add_argument(
+            "--max-context-len-to-capture",
+            type=int,
+            default=EngineArgs.max_context_len_to_capture,
+            help="maximum context length covered by CUDA "
+            "graphs. When a sequence has context length "
+            "larger than this, we fall back to eager mode.",
+        )
+        return parser
+
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace) -> "EngineArgs":
+        # Get the list of attributes of this dataclass.
+        attrs = [attr.name for attr in dataclasses.fields(cls)]
+        # Set the attributes from the parsed arguments.
+        engine_args = cls(**{attr: getattr(args, attr) for attr in attrs})
+        return engine_args
+
+    def create_engine_configs(
+        self,
+    ) -> Tuple[ModelConfig, CacheConfig, ParallelConfig, SchedulerConfig]:
+        model_config = ModelConfig(
+            self.model,
+            self.tokenizer,
+            self.tokenizer_mode,
+            self.trust_remote_code,
+            self.download_dir,
+            self.load_format,
+            self.dtype,
+            self.seed,
+            self.revision,
+            self.tokenizer_revision,
+            self.max_model_len,
+            self.quantization,
+            self.enforce_eager,
+            self.max_context_len_to_capture,
+            self.num_audio_tokens,
+            self.num_text_tokens,
+        )
+        cache_config = CacheConfig(
+            self.block_size,
+            self.gpu_memory_utilization,
+            self.swap_space,
+            model_config.get_sliding_window(),
+        )
+        parallel_config = ParallelConfig(
+            self.pipeline_parallel_size,
+            self.tensor_parallel_size,
+            self.worker_use_ray,
+            self.max_parallel_loading_workers,
+        )
+        scheduler_config = SchedulerConfig(
+            self.max_num_batched_tokens,
+            self.max_num_seqs,
+            model_config.max_model_len,
+            self.max_paddings,
+        )
+        return model_config, cache_config, parallel_config, scheduler_config
+
+
+@dataclass
+class AsyncEngineArgs(EngineArgs):
+    """Arguments for asynchronous vLLM engine."""
+
+    engine_use_ray: bool = False
+    disable_log_requests: bool = False
+    max_log_len: Optional[int] = None
+
+    @staticmethod
+    def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+        parser = EngineArgs.add_cli_args(parser)
+        parser.add_argument(
+            "--engine-use-ray",
+            action="store_true",
+            help="use Ray to start the LLM engine in a "
+            "separate process as the server process.",
+        )
+        parser.add_argument(
+            "--disable-log-requests",
+            action="store_true",
+            help="disable logging requests",
+        )
+        parser.add_argument(
+            "--max-log-len",
+            type=int,
+            default=None,
+            help="max number of prompt characters or prompt "
+            "ID numbers being printed in log. "
+            "Default: unlimited.",
+        )
+        return parser

ChatTTS/model/velocity/llama.py

@@ -0,0 +1,394 @@
+# coding=utf-8
+# Adapted from
+# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
+# Copyright 2023 The vLLM team.
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# 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
+from torch import nn
+from transformers import LlamaConfig
+
+from vllm.model_executor.input_metadata import InputMetadata
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.attention import PagedAttention
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    LinearMethodBase,
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.sampler import Sampler
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+    ParallelLMHead,
+)
+from vllm.model_executor.parallel_utils.parallel_state import (
+    get_tensor_model_parallel_world_size,
+)
+from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.model_executor.weight_utils import (
+    default_weight_loader,
+    hf_model_weights_iterator,
+)
+from vllm.sequence import SamplerOutput
+
+KVCache = Tuple[torch.Tensor, torch.Tensor]
+
+
+class LlamaMLP(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        linear_method: Optional[LinearMethodBase] = None,
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            hidden_size,
+            [intermediate_size] * 2,
+            bias=False,
+            linear_method=linear_method,
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size, hidden_size, bias=False, linear_method=linear_method
+        )
+        if hidden_act != "silu":
+            raise ValueError(
+                f"Unsupported activation: {hidden_act}. "
+                "Only silu is supported for now."
+            )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x):
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class LlamaAttention(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        rope_theta: float = 10000,
+        rope_scaling: Optional[Dict[str, Any]] = None,
+        max_position_embeddings: int = 8192,
+        linear_method: Optional[LinearMethodBase] = None,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            # Number of KV heads is greater than TP size, so we partition
+            # the KV heads across multiple tensor parallel GPUs.
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            # Number of KV heads is less than TP size, so we replicate
+            # the KV heads across multiple tensor parallel GPUs.
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = hidden_size // self.total_num_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size,
+            self.head_dim,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=False,
+            linear_method=linear_method,
+        )
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.head_dim,
+            hidden_size,
+            bias=False,
+            linear_method=linear_method,
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.head_dim,
+            max_position=max_position_embeddings,
+            base=rope_theta,
+            rope_scaling=rope_scaling,
+        )
+        self.attn = PagedAttention(
+            self.num_heads, self.head_dim, self.scaling, num_kv_heads=self.num_kv_heads
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        kv_cache: KVCache,
+        input_metadata: InputMetadata,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        k_cache, v_cache = kv_cache
+        attn_output = self.attn(q, k, v, k_cache, v_cache, input_metadata)
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class LlamaDecoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        config: LlamaConfig,
+        linear_method: Optional[LinearMethodBase] = None,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        rope_theta = getattr(config, "rope_theta", 10000)
+        rope_scaling = getattr(config, "rope_scaling", None)
+        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
+        self.self_attn = LlamaAttention(
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=config.num_key_value_heads,
+            rope_theta=rope_theta,
+            rope_scaling=rope_scaling,
+            max_position_embeddings=max_position_embeddings,
+            linear_method=linear_method,
+        )
+        self.mlp = LlamaMLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            linear_method=linear_method,
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        kv_cache: KVCache,
+        input_metadata: InputMetadata,
+        residual: Optional[torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Self Attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            kv_cache=kv_cache,
+            input_metadata=input_metadata,
+        )
+
+        # Fully Connected
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+class LlamaModel(nn.Module):
+
+    def __init__(
+        self,
+        config: LlamaConfig,
+        linear_method: Optional[LinearMethodBase] = None,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+        )
+        self.layers = nn.ModuleList(
+            [
+                LlamaDecoderLayer(config, linear_method)
+                for _ in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        input_emb: torch.Tensor,
+        positions: torch.Tensor,
+        kv_caches: List[KVCache],
+        input_metadata: InputMetadata,
+    ) -> torch.Tensor:
+        hidden_states = input_emb
+        residual = None
+        for i in range(len(self.layers)):
+            layer = self.layers[i]
+            hidden_states, residual = layer(
+                positions,
+                hidden_states,
+                kv_caches[i],
+                input_metadata,
+                residual,
+            )
+        hidden_states, _ = self.norm(hidden_states, residual)
+        return hidden_states
+
+    def load_weights(
+        self,
+        model_name_or_path: str,
+        cache_dir: Optional[str] = None,
+        load_format: str = "auto",
+        revision: Optional[str] = None,
+    ):
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+        params_dict = dict(self.named_parameters())
+        for name, loaded_weight in hf_model_weights_iterator(
+            model_name_or_path, cache_dir, load_format, revision
+        ):
+            if "rotary_emb.inv_freq" in name:
+                continue
+            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
+                # Models trained using ColossalAI may include these tensors in
+                # the checkpoint. Skip them.
+                continue
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+
+class LlamaForCausalLM(nn.Module):
+
+    def __init__(
+        self,
+        config: LlamaConfig,
+        linear_method: Optional[LinearMethodBase] = None,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.linear_method = linear_method
+        self.model = LlamaModel(config, linear_method)
+        self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
+        self.sampler = Sampler(config.vocab_size)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        kv_caches: List[KVCache],
+        input_metadata: InputMetadata,
+    ) -> torch.Tensor:
+        hidden_states = self.model(input_ids, positions, kv_caches, input_metadata)
+        return hidden_states
+
+    def sample(
+        self,
+        hidden_states: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> Optional[SamplerOutput]:
+        next_tokens = self.sampler(
+            self.lm_head.weight, hidden_states, sampling_metadata
+        )
+        return next_tokens
+
+    def load_weights(
+        self,
+        model_name_or_path: str,
+        cache_dir: Optional[str] = None,
+        load_format: str = "auto",
+        revision: Optional[str] = None,
+    ):
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+        params_dict = dict(self.named_parameters())
+        for name, loaded_weight in hf_model_weights_iterator(
+            model_name_or_path, cache_dir, load_format, revision
+        ):
+            if "rotary_emb.inv_freq" in name:
+                continue
+            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
+                # Models trained using ColossalAI may include these tensors in
+                # the checkpoint. Skip them.
+                continue
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)

ChatTTS/model/velocity/llm.py

@@ -0,0 +1,213 @@
+from typing import List, Optional, Union
+
+from tqdm import tqdm
+from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
+from vllm.utils import Counter
+
+from .configs import EngineArgs
+from .llm_engine import LLMEngine
+from .output import RequestOutput
+from .sampling_params import SamplingParams
+
+
+class LLM:
+    """An LLM for generating texts from given prompts and sampling parameters.
+
+    This class includes a tokenizer, a language model (possibly distributed
+    across multiple GPUs), and GPU memory space allocated for intermediate
+    states (aka KV cache). Given a batch of prompts and sampling parameters,
+    this class generates texts from the model, using an intelligent batching
+    mechanism and efficient memory management.
+
+    NOTE: This class is intended to be used for offline inference. For online
+    serving, use the `AsyncLLMEngine` class instead.
+    NOTE: For the comprehensive list of arguments, see `EngineArgs`.
+
+    Args:
+        model: The name or path of a HuggingFace Transformers model.
+        tokenizer: The name or path of a HuggingFace Transformers tokenizer.
+        tokenizer_mode: The tokenizer mode. "auto" will use the fast tokenizer
+            if available, and "slow" will always use the slow tokenizer.
+        trust_remote_code: Trust remote code (e.g., from HuggingFace) when
+            downloading the model and tokenizer.
+        tensor_parallel_size: The number of GPUs to use for distributed
+            execution with tensor parallelism.
+        dtype: The data type for the model weights and activations. Currently,
+            we support `float32`, `float16`, and `bfloat16`. If `auto`, we use
+            the `torch_dtype` attribute specified in the model config file.
+            However, if the `torch_dtype` in the config is `float32`, we will
+            use `float16` instead.
+        quantization: The method used to quantize the model weights. Currently,
+            we support "awq", "gptq" and "squeezellm". If None, we first check
+            the `quantization_config` attribute in the model config file. If
+            that is None, we assume the model weights are not quantized and use
+            `dtype` to determine the data type of the weights.
+        revision: The specific model version to use. It can be a branch name,
+            a tag name, or a commit id.
+        tokenizer_revision: The specific tokenizer version to use. It can be a
+            branch name, a tag name, or a commit id.
+        seed: The seed to initialize the random number generator for sampling.
+        gpu_memory_utilization: The ratio (between 0 and 1) of GPU memory to
+            reserve for the model weights, activations, and KV cache. Higher
+            values will increase the KV cache size and thus improve the model's
+            throughput. However, if the value is too high, it may cause out-of-
+            memory (OOM) errors.
+        swap_space: The size (GiB) of CPU memory per GPU to use as swap space.
+            This can be used for temporarily storing the states of the requests
+            when their `best_of` sampling parameters are larger than 1. If all
+            requests will have `best_of=1`, you can safely set this to 0.
+            Otherwise, too small values may cause out-of-memory (OOM) errors.
+        enforce_eager: Whether to enforce eager execution. If True, we will
+            disable CUDA graph and always execute the model in eager mode.
+            If False, we will use CUDA graph and eager execution in hybrid.
+        max_context_len_to_capture: Maximum context len covered by CUDA graphs.
+            When a sequence has context length larger than this, we fall back
+            to eager mode.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        tokenizer: Optional[str] = None,
+        tokenizer_mode: str = "auto",
+        trust_remote_code: bool = False,
+        tensor_parallel_size: int = 1,
+        dtype: str = "auto",
+        quantization: Optional[str] = None,
+        revision: Optional[str] = None,
+        tokenizer_revision: Optional[str] = None,
+        seed: int = 0,
+        gpu_memory_utilization: float = 0.9,
+        swap_space: int = 4,
+        enforce_eager: bool = False,
+        max_context_len_to_capture: int = 8192,
+        post_model_path: str = None,
+        num_audio_tokens: int = 0,
+        num_text_tokens: int = 0,
+        **kwargs,
+    ) -> None:
+        if "disable_log_stats" not in kwargs:
+            kwargs["disable_log_stats"] = True
+        engine_args = EngineArgs(
+            model=model,
+            tokenizer=tokenizer,
+            tokenizer_mode=tokenizer_mode,
+            trust_remote_code=trust_remote_code,
+            tensor_parallel_size=tensor_parallel_size,
+            dtype=dtype,
+            quantization=quantization,
+            revision=revision,
+            tokenizer_revision=tokenizer_revision,
+            seed=seed,
+            gpu_memory_utilization=gpu_memory_utilization,
+            swap_space=swap_space,
+            enforce_eager=enforce_eager,
+            max_context_len_to_capture=max_context_len_to_capture,
+            num_audio_tokens=num_audio_tokens,
+            num_text_tokens=num_text_tokens,
+            **kwargs,
+        )
+        self.llm_engine = LLMEngine.from_engine_args(engine_args, post_model_path)
+        self.request_counter = Counter()
+
+    def get_tokenizer(self) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:
+        return self.llm_engine.tokenizer
+
+    def set_tokenizer(
+        self,
+        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
+    ) -> None:
+        self.llm_engine.tokenizer = tokenizer
+
+    def generate(
+        self,
+        prompts: Optional[Union[str, List[str]]] = None,
+        sampling_params: Optional[SamplingParams] = None,
+        prompt_token_ids: Optional[List[List[int]]] = None,
+        use_tqdm: bool = True,
+    ) -> List[RequestOutput]:
+        """Generates the completions for the input prompts.
+
+        NOTE: This class automatically batches the given prompts, considering
+        the memory constraint. For the best performance, put all of your prompts
+        into a single list and pass it to this method.
+
+        Args:
+            prompts: A list of prompts to generate completions for.
+            sampling_params: The sampling parameters for text generation. If
+                None, we use the default sampling parameters.
+            prompt_token_ids: A list of token IDs for the prompts. If None, we
+                use the tokenizer to convert the prompts to token IDs.
+            use_tqdm: Whether to use tqdm to display the progress bar.
+
+        Returns:
+            A list of `RequestOutput` objects containing the generated
+            completions in the same order as the input prompts.
+        """
+        if prompts is None and prompt_token_ids is None:
+            raise ValueError("Either prompts or prompt_token_ids must be " "provided.")
+        if isinstance(prompts, str):
+            # Convert a single prompt to a list.
+            prompts = [prompts]
+        if (
+            prompts is not None
+            and prompt_token_ids is not None
+            and len(prompts) != len(prompt_token_ids)
+        ):
+            raise ValueError(
+                "The lengths of prompts and prompt_token_ids " "must be the same."
+            )
+        if sampling_params is None:
+            # Use default sampling params.
+            sampling_params = SamplingParams()
+
+        # Add requests to the engine.
+        num_requests = len(prompts) if prompts is not None else len(prompt_token_ids)
+        for i in range(num_requests):
+            prompt = prompts[i] if prompts is not None else None
+            token_ids = None if prompt_token_ids is None else prompt_token_ids[i]
+            self._add_request(prompt, sampling_params, token_ids)
+
+        rtns = self._run_engine(use_tqdm)
+        for i, rtn in enumerate(rtns):
+            token_ids = rtn.outputs[0].token_ids
+            for j, token_id in enumerate(token_ids):
+                if len(token_id) == 1:
+                    token_ids[j] = token_id[0]
+                else:
+                    token_ids[j] = list(token_id)
+
+        return rtns
+
+    def _add_request(
+        self,
+        prompt: Optional[str],
+        sampling_params: SamplingParams,
+        prompt_token_ids: Optional[List[int]],
+    ) -> None:
+        request_id = str(next(self.request_counter))
+        self.llm_engine.add_request(
+            request_id, prompt, sampling_params, prompt_token_ids
+        )
+
+    def _run_engine(self, use_tqdm: bool) -> List[RequestOutput]:
+        # Initialize tqdm.
+        if use_tqdm:
+            num_requests = self.llm_engine.get_num_unfinished_requests()
+            pbar = tqdm(total=num_requests, desc="Processed prompts")
+        # Run the engine.
+        outputs: List[RequestOutput] = []
+        while self.llm_engine.has_unfinished_requests():
+            step_outputs = self.llm_engine.step()
+            for output in step_outputs:
+                if output.finished:
+                    outputs.append(output)
+                    if use_tqdm:
+                        pbar.update(1)
+        if use_tqdm:
+            pbar.close()
+        # Sort the outputs by request ID.
+        # This is necessary because some requests may be finished earlier than
+        # its previous requests.
+        outputs = sorted(outputs, key=lambda x: int(x.request_id))
+        return outputs

ChatTTS/model/velocity/llm_engine.py

@@ -0,0 +1,833 @@
+import copy
+from collections import defaultdict
+import os
+import time
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union
+
+from vllm.config import CacheConfig, ModelConfig, ParallelConfig, SchedulerConfig
+from .scheduler import Scheduler, SchedulerOutputs
+from .configs import EngineArgs
+from vllm.engine.metrics import record_metrics
+from vllm.engine.ray_utils import RayWorkerVllm, initialize_cluster, ray
+from vllm.logger import init_logger
+from .output import RequestOutput
+from .sampling_params import SamplingParams
+from .sequence import (
+    SamplerOutput,
+    Sequence,
+    SequenceGroup,
+    SequenceGroupOutput,
+    SequenceOutput,
+    SequenceStatus,
+)
+from vllm.transformers_utils.tokenizer import detokenize_incrementally, get_tokenizer
+from vllm.utils import Counter, set_cuda_visible_devices, get_ip, get_open_port
+import numpy as np
+
+if ray:
+    from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+if TYPE_CHECKING:
+    from ray.util.placement_group import PlacementGroup
+
+logger = init_logger(__name__)
+
+_LOGGING_INTERVAL_SEC = 5
+
+
+class LLMEngine:
+    """An LLM engine that receives requests and generates texts.
+
+    This is the main class for the vLLM engine. It receives requests
+    from clients and generates texts from the LLM. It includes a tokenizer, a
+    language model (possibly distributed across multiple GPUs), and GPU memory
+    space allocated for intermediate states (aka KV cache). This class utilizes
+    iteration-level scheduling and efficient memory management to maximize the
+    serving throughput.
+
+    The `LLM` class wraps this class for offline batched inference and the
+    `AsyncLLMEngine` class wraps this class for online serving.
+
+    NOTE: The config arguments are derived from the `EngineArgs` class. For the
+    comprehensive list of arguments, see `EngineArgs`.
+
+    Args:
+        model_config: The configuration related to the LLM model.
+        cache_config: The configuration related to the KV cache memory
+            management.
+        parallel_config: The configuration related to distributed execution.
+        scheduler_config: The configuration related to the request scheduler.
+        placement_group: Ray placement group for distributed execution.
+            Required for distributed execution.
+        log_stats: Whether to log statistics.
+    """
+
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        cache_config: CacheConfig,
+        parallel_config: ParallelConfig,
+        scheduler_config: SchedulerConfig,
+        placement_group: Optional["PlacementGroup"],
+        post_model_path: str,
+        log_stats: bool,
+    ) -> None:
+        logger.info(
+            "Initializing an LLM engine with config: "
+            f"model={model_config.model!r}, "
+            f"tokenizer={model_config.tokenizer!r}, "
+            f"tokenizer_mode={model_config.tokenizer_mode}, "
+            f"revision={model_config.revision}, "
+            f"tokenizer_revision={model_config.tokenizer_revision}, "
+            f"trust_remote_code={model_config.trust_remote_code}, "
+            f"dtype={model_config.dtype}, "
+            f"max_seq_len={model_config.max_model_len}, "
+            f"download_dir={model_config.download_dir!r}, "
+            f"load_format={model_config.load_format}, "
+            f"tensor_parallel_size={parallel_config.tensor_parallel_size}, "
+            f"quantization={model_config.quantization}, "
+            f"enforce_eager={model_config.enforce_eager}, "
+            f"seed={model_config.seed}), "
+            f"post_model_path={post_model_path!r}"
+        )
+        # TODO(woosuk): Print more configs in debug mode.
+
+        self.model_config = model_config
+        self.cache_config = cache_config
+        self.parallel_config = parallel_config
+        self.scheduler_config = scheduler_config
+        self.log_stats = log_stats
+        self._verify_args()
+        self.post_model_path = post_model_path
+        self.seq_counter = Counter()
+
+        # Create the parallel GPU workers.
+        if self.parallel_config.worker_use_ray:
+            # Disable Ray usage stats collection.
+            ray_usage = os.environ.get("RAY_USAGE_STATS_ENABLED", "0")
+            if ray_usage != "1":
+                os.environ["RAY_USAGE_STATS_ENABLED"] = "0"
+            self._init_workers_ray(placement_group)
+        else:
+            self._init_workers()
+
+        # Profile the memory usage and initialize the cache.
+        self._init_cache()
+
+        # Create the scheduler.
+        self.scheduler = Scheduler(scheduler_config, cache_config)
+
+        # Logging.
+        self.last_logging_time = 0.0
+        # List of (timestamp, num_tokens)
+        self.num_prompt_tokens: List[Tuple[float, int]] = []
+        # List of (timestamp, num_tokens)
+        self.num_generation_tokens: List[Tuple[float, int]] = []
+
+    def _init_workers(self):
+        # Lazy import the Worker to avoid importing torch.cuda/xformers
+        # before CUDA_VISIBLE_DEVICES is set in the Worker
+        from .worker import Worker
+
+        assert (
+            self.parallel_config.world_size == 1
+        ), "Ray is required if parallel_config.world_size > 1."
+
+        self.workers: List[Worker] = []
+        distributed_init_method = f"tcp://{get_ip()}:{get_open_port()}"
+        self.driver_worker = Worker(
+            self.model_config,
+            self.parallel_config,
+            self.scheduler_config,
+            local_rank=0,
+            rank=0,
+            distributed_init_method=distributed_init_method,
+            is_driver_worker=True,
+            post_model_path=self.post_model_path,
+        )
+        self._run_workers("init_model")
+        self._run_workers("load_model")
+
+    def _init_workers_ray(self, placement_group: "PlacementGroup", **ray_remote_kwargs):
+        if self.parallel_config.tensor_parallel_size == 1:
+            num_gpus = self.cache_config.gpu_memory_utilization
+        else:
+            num_gpus = 1
+
+        self.driver_dummy_worker: RayWorkerVllm = None
+        self.workers: List[RayWorkerVllm] = []
+
+        driver_ip = get_ip()
+        for bundle_id, bundle in enumerate(placement_group.bundle_specs):
+            if not bundle.get("GPU", 0):
+                continue
+            scheduling_strategy = PlacementGroupSchedulingStrategy(
+                placement_group=placement_group,
+                placement_group_capture_child_tasks=True,
+                placement_group_bundle_index=bundle_id,
+            )
+            worker = ray.remote(
+                num_cpus=0,
+                num_gpus=num_gpus,
+                scheduling_strategy=scheduling_strategy,
+                **ray_remote_kwargs,
+            )(RayWorkerVllm).remote(self.model_config.trust_remote_code)
+
+            worker_ip = ray.get(worker.get_node_ip.remote())
+            if worker_ip == driver_ip and self.driver_dummy_worker is None:
+                # If the worker is on the same node as the driver, we use it
+                # as the resource holder for the driver process.
+                self.driver_dummy_worker = worker
+            else:
+                self.workers.append(worker)
+
+        if self.driver_dummy_worker is None:
+            raise ValueError(
+                "Ray does not allocate any GPUs on the driver node. Consider "
+                "adjusting the Ray placement group or running the driver on a "
+                "GPU node."
+            )
+
+        driver_node_id, driver_gpu_ids = ray.get(
+            self.driver_dummy_worker.get_node_and_gpu_ids.remote()
+        )
+        worker_node_and_gpu_ids = ray.get(
+            [worker.get_node_and_gpu_ids.remote() for worker in self.workers]
+        )
+
+        node_workers = defaultdict(list)
+        node_gpus = defaultdict(list)
+
+        node_workers[driver_node_id].append(0)
+        node_gpus[driver_node_id].extend(driver_gpu_ids)
+        for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids, start=1):
+            node_workers[node_id].append(i)
+            node_gpus[node_id].extend(gpu_ids)
+        for node_id, gpu_ids in node_gpus.items():
+            node_gpus[node_id] = sorted(gpu_ids)
+
+        # Set CUDA_VISIBLE_DEVICES for the driver.
+        set_cuda_visible_devices(node_gpus[driver_node_id])
+        for worker, (node_id, _) in zip(self.workers, worker_node_and_gpu_ids):
+            worker.set_cuda_visible_devices.remote(node_gpus[node_id])
+
+        distributed_init_method = f"tcp://{driver_ip}:{get_open_port()}"
+
+        # Lazy import the Worker to avoid importing torch.cuda/xformers
+        # before CUDA_VISIBLE_DEVICES is set in the Worker
+        from vllm.worker.worker import Worker
+
+        # Initialize torch distributed process group for the workers.
+        model_config = copy.deepcopy(self.model_config)
+        parallel_config = copy.deepcopy(self.parallel_config)
+        scheduler_config = copy.deepcopy(self.scheduler_config)
+
+        for rank, (worker, (node_id, _)) in enumerate(
+            zip(self.workers, worker_node_and_gpu_ids), start=1
+        ):
+            local_rank = node_workers[node_id].index(rank)
+            worker.init_worker.remote(
+                lambda rank=rank, local_rank=local_rank: Worker(
+                    model_config,
+                    parallel_config,
+                    scheduler_config,
+                    local_rank,
+                    rank,
+                    distributed_init_method,
+                )
+            )
+
+        driver_rank = 0
+        driver_local_rank = node_workers[driver_node_id].index(driver_rank)
+        self.driver_worker = Worker(
+            model_config,
+            parallel_config,
+            scheduler_config,
+            driver_local_rank,
+            driver_rank,
+            distributed_init_method,
+            is_driver_worker=True,
+        )
+
+        self._run_workers("init_model")
+        self._run_workers(
+            "load_model",
+            max_concurrent_workers=self.parallel_config.max_parallel_loading_workers,
+        )
+
+    def _verify_args(self) -> None:
+        self.model_config.verify_with_parallel_config(self.parallel_config)
+        self.cache_config.verify_with_parallel_config(self.parallel_config)
+
+    def _init_cache(self) -> None:
+        """Profiles the memory usage and initializes the KV cache."""
+        # Get the maximum number of blocks that can be allocated on GPU and CPU.
+        num_blocks = self._run_workers(
+            "profile_num_available_blocks",
+            block_size=self.cache_config.block_size,
+            gpu_memory_utilization=self.cache_config.gpu_memory_utilization,
+            cpu_swap_space=self.cache_config.swap_space_bytes,
+        )
+
+        # Since we use a shared centralized controller, we take the minimum
+        # number of blocks across all workers to make sure all the memory
+        # operators can be applied to all workers.
+        num_gpu_blocks = min(b[0] for b in num_blocks)
+        num_cpu_blocks = min(b[1] for b in num_blocks)
+        # FIXME(woosuk): Change to debug log.
+        logger.info(
+            f"# GPU blocks: {num_gpu_blocks}, " f"# CPU blocks: {num_cpu_blocks}"
+        )
+
+        if num_gpu_blocks <= 0:
+            raise ValueError(
+                "No available memory for the cache blocks. "
+                "Try increasing `gpu_memory_utilization` when "
+                "initializing the engine."
+            )
+        max_seq_len = self.cache_config.block_size * num_gpu_blocks
+        if self.model_config.max_model_len > max_seq_len:
+            raise ValueError(
+                f"The model's max seq len ({self.model_config.max_model_len}) "
+                "is larger than the maximum number of tokens that can be "
+                f"stored in KV cache ({max_seq_len}). Try increasing "
+                "`gpu_memory_utilization` or decreasing `max_model_len` when "
+                "initializing the engine."
+            )
+
+        self.cache_config.num_gpu_blocks = num_gpu_blocks
+        self.cache_config.num_cpu_blocks = num_cpu_blocks
+
+        # Initialize the cache.
+        self._run_workers("init_cache_engine", cache_config=self.cache_config)
+        # Warm up the model. This includes capturing the model into CUDA graph
+        # if enforce_eager is False.
+        self._run_workers("warm_up_model")
+
+    @classmethod
+    def from_engine_args(
+        cls, engine_args: EngineArgs, post_model_path=None
+    ) -> "LLMEngine":
+        """Creates an LLM engine from the engine arguments."""
+        # Create the engine configs.
+        engine_configs = engine_args.create_engine_configs()
+        parallel_config = engine_configs[2]
+        # Initialize the cluster.
+        placement_group = initialize_cluster(parallel_config)
+        # Create the LLM engine.
+        engine = cls(
+            *engine_configs,
+            placement_group,
+            log_stats=not engine_args.disable_log_stats,
+            post_model_path=post_model_path,
+        )
+        return engine
+
+    def add_request(
+        self,
+        request_id: str,
+        prompt: Optional[str],
+        sampling_params: SamplingParams,
+        prompt_token_ids: Optional[List[int]] = None,
+        arrival_time: Optional[float] = None,
+    ) -> None:
+        """Add a request to the engine's request pool.
+
+        The request is added to the request pool and will be processed by the
+        scheduler as `engine.step()` is called. The exact scheduling policy is
+        determined by the scheduler.
+
+        Args:
+            request_id: The unique ID of the request.
+            prompt: The prompt string. Can be None if prompt_token_ids is
+                provided.
+            sampling_params: The sampling parameters for text generation.
+            prompt_token_ids: The token IDs of the prompt. If None, we
+                use the tokenizer to convert the prompts to token IDs.
+            arrival_time: The arrival time of the request. If None, we use
+                the current monotonic time.
+        """
+        if arrival_time is None:
+            arrival_time = time.monotonic()
+
+        assert prompt_token_ids is not None, "prompt_token_ids must be provided"
+        # Create the sequences.
+        block_size = self.cache_config.block_size
+        seq_id = next(self.seq_counter)
+        seq = Sequence(seq_id, prompt, prompt_token_ids, block_size)
+
+        # Create the sequence group.
+        seq_group = SequenceGroup(request_id, [seq], sampling_params, arrival_time)
+
+        # Add the sequence group to the scheduler.
+        self.scheduler.add_seq_group(seq_group)
+
+    def abort_request(self, request_id: Union[str, Iterable[str]]) -> None:
+        """Aborts a request(s) with the given ID.
+
+        Args:
+            request_id: The ID(s) of the request to abort.
+        """
+        self.scheduler.abort_seq_group(request_id)
+
+    def get_model_config(self) -> ModelConfig:
+        """Gets the model configuration."""
+        return self.model_config
+
+    def get_num_unfinished_requests(self) -> int:
+        """Gets the number of unfinished requests."""
+        return self.scheduler.get_num_unfinished_seq_groups()
+
+    def has_unfinished_requests(self) -> bool:
+        """Returns True if there are unfinished requests."""
+        return self.scheduler.has_unfinished_seqs()
+
+    def _check_beam_search_early_stopping(
+        self,
+        early_stopping: Union[bool, str],
+        sampling_params: SamplingParams,
+        best_running_seq: Sequence,
+        current_worst_seq: Sequence,
+    ) -> bool:
+        assert sampling_params.use_beam_search
+        length_penalty = sampling_params.length_penalty
+        if early_stopping is True:
+            return True
+
+        current_worst_score = current_worst_seq.get_beam_search_score(
+            length_penalty=length_penalty, eos_token_id=self.tokenizer.eos_token_id
+        )
+        if early_stopping is False:
+            highest_attainable_score = best_running_seq.get_beam_search_score(
+                length_penalty=length_penalty, eos_token_id=self.tokenizer.eos_token_id
+            )
+        else:
+            assert early_stopping == "never"
+            if length_penalty > 0.0:
+                # If length_penalty > 0.0, beam search will prefer longer
+                # sequences. The highest attainable score calculation is
+                # based on the longest possible sequence length in this case.
+                max_possible_length = max(
+                    best_running_seq.get_prompt_len() + sampling_params.max_tokens,
+                    self.scheduler_config.max_model_len,
+                )
+                highest_attainable_score = best_running_seq.get_beam_search_score(
+                    length_penalty=length_penalty,
+                    eos_token_id=self.tokenizer.eos_token_id,
+                    seq_len=max_possible_length,
+                )
+            else:
+                # Otherwise, beam search will prefer shorter sequences. The
+                # highest attainable score calculation is based on the current
+                # sequence length.
+                highest_attainable_score = best_running_seq.get_beam_search_score(
+                    length_penalty=length_penalty,
+                    eos_token_id=self.tokenizer.eos_token_id,
+                )
+        return current_worst_score >= highest_attainable_score
+
+    def _process_sequence_group_outputs(
+        self, seq_group: SequenceGroup, outputs: SequenceGroupOutput
+    ) -> None:
+        # Process prompt logprobs
+        prompt_logprobs = outputs.prompt_logprobs
+        if prompt_logprobs is not None:
+            seq_group.prompt_logprobs = prompt_logprobs
+
+        # Process samples
+        samples = outputs.samples
+        parent_seqs = seq_group.get_seqs(status=SequenceStatus.RUNNING)
+        existing_finished_seqs = seq_group.get_finished_seqs()
+        parent_child_dict = {parent_seq.seq_id: [] for parent_seq in parent_seqs}
+        for sample in samples:
+            parent_child_dict[sample.parent_seq_id].append(sample)
+        # List of (child, parent)
+        child_seqs: List[Tuple[Sequence, Sequence]] = []
+
+        # Process the child samples for each parent sequence
+        for parent in parent_seqs:
+            child_samples: List[SequenceOutput] = parent_child_dict[parent.seq_id]
+            if len(child_samples) == 0:
+                # This parent sequence has no children samples. Remove
+                # the parent sequence from the sequence group since it will
+                # not be used in the future iterations.
+                parent.status = SequenceStatus.FINISHED_ABORTED
+                seq_group.remove(parent.seq_id)
+                self.scheduler.free_seq(parent)
+                continue
+            # Fork the parent sequence if there are multiple child samples.
+            for child_sample in child_samples[:-1]:
+                new_child_seq_id = next(self.seq_counter)
+                child = parent.fork(new_child_seq_id)
+                child.append_token_id(
+                    child_sample.output_token,
+                    child_sample.logprobs,
+                    child_sample.hidden_states,
+                    child_sample.finished,
+                )
+                child_seqs.append((child, parent))
+            # Continue the parent sequence for the last child sample.
+            # We reuse the parent sequence here to reduce redundant memory
+            # copies, especially when using non-beam search sampling methods.
+            last_child_sample = child_samples[-1]
+            parent.append_token_id(
+                last_child_sample.output_token,
+                last_child_sample.logprobs,
+                last_child_sample.hidden_states,
+                last_child_sample.finished,
+            )
+            child_seqs.append((parent, parent))
+
+        for seq, _ in child_seqs:
+            # self._decode_sequence(seq, seq_group.sampling_params)
+            self._check_stop(seq, seq_group.sampling_params)
+
+        # Non-beam search case
+        if not seq_group.sampling_params.use_beam_search:
+            # For newly created child sequences, add them to the sequence group
+            # and fork them in block manager if they are not finished.
+            for seq, parent in child_seqs:
+                if seq is not parent:
+                    seq_group.add(seq)
+                    if not seq.is_finished():
+                        self.scheduler.fork_seq(parent, seq)
+
+            # Free the finished and selected parent sequences' memory in block
+            # manager. Keep them in the sequence group as candidate output.
+            # NOTE: we need to fork the new sequences before freeing the
+            # old sequences.
+            for seq, parent in child_seqs:
+                if seq is parent and seq.is_finished():
+                    self.scheduler.free_seq(seq)
+            return
+
+        # Beam search case
+        # Select the child sequences to keep in the sequence group.
+        selected_child_seqs = []
+        unselected_child_seqs = []
+        beam_width = seq_group.sampling_params.best_of
+        length_penalty = seq_group.sampling_params.length_penalty
+
+        # Select the newly finished sequences with the highest scores
+        # to replace existing finished sequences.
+        # Tuple of (seq, parent, is_new)
+        existing_finished_seqs = [(seq, None, False) for seq in existing_finished_seqs]
+        new_finished_seqs = [
+            (seq, parent, True) for seq, parent in child_seqs if seq.is_finished()
+        ]
+        all_finished_seqs = existing_finished_seqs + new_finished_seqs
+        # Sort the finished sequences by their scores.
+        all_finished_seqs.sort(
+            key=lambda x: x[0].get_beam_search_score(
+                length_penalty=length_penalty, eos_token_id=self.tokenizer.eos_token_id
+            ),
+            reverse=True,
+        )
+        for seq, parent, is_new in all_finished_seqs[:beam_width]:
+            if is_new:
+                # A newly generated child sequence finishes and has a high
+                # score, so we will add it into the sequence group.
+                selected_child_seqs.append((seq, parent))
+        for seq, parent, is_new in all_finished_seqs[beam_width:]:
+            if is_new:
+                # A newly generated child sequence finishes but has a low
+                # score, so we will not add it into the sequence group.
+                # Additionally, if this sequence is a continuation of a
+                # parent sequence, we will need remove the parent sequence
+                # from the sequence group.
+                unselected_child_seqs.append((seq, parent))
+            else:
+                # An existing finished sequence has a low score, so we will
+                # remove it from the sequence group.
+                seq_group.remove(seq.seq_id)
+
+        # select the top beam_width sequences from the running
+        # sequences for the next iteration to continue the beam
+        # search.
+        running_child_seqs = [
+            (seq, parent) for seq, parent in child_seqs if not seq.is_finished()
+        ]
+        # Sort the running sequences by their scores.
+        running_child_seqs.sort(
+            key=lambda x: x[0].get_beam_search_score(
+                length_penalty=length_penalty, eos_token_id=self.tokenizer.eos_token_id
+            ),
+            reverse=True,
+        )
+
+        # Check if we can stop the beam search.
+        if len(running_child_seqs) == 0:
+            # No running sequences, stop the beam search.
+            stop_beam_search = True
+        elif len(all_finished_seqs) < beam_width:
+            # Not enough finished sequences, continue the beam search.
+            stop_beam_search = False
+        else:
+            # Check the early stopping criteria
+            best_running_seq = running_child_seqs[0][0]
+            current_worst_seq = all_finished_seqs[beam_width - 1][0]
+            stop_beam_search = self._check_beam_search_early_stopping(
+                seq_group.sampling_params.early_stopping,
+                seq_group.sampling_params,
+                best_running_seq,
+                current_worst_seq,
+            )
+
+        if stop_beam_search:
+            # Stop the beam search and remove all the running sequences from
+            # the sequence group.
+            unselected_child_seqs.extend(running_child_seqs)
+        else:
+            # Continue the beam search and select the top beam_width sequences
+            # to continue the beam search.
+            selected_child_seqs.extend(running_child_seqs[:beam_width])
+            # The remaining running sequences will not be used in the next
+            # iteration. Again, if these sequences are continuations of
+            # parent sequences, we will need to remove the parent sequences
+            # from the sequence group.
+            unselected_child_seqs.extend(running_child_seqs[beam_width:])
+
+        # For newly created child sequences, add them to the sequence group
+        # and fork them in block manager if they are not finished.
+        for seq, parent in selected_child_seqs:
+            if seq is not parent:
+                seq_group.add(seq)
+                if not seq.is_finished():
+                    self.scheduler.fork_seq(parent, seq)
+
+        # Free the finished and selected parent sequences' memory in block
+        # manager. Keep them in the sequence group as candidate output.
+        for seq, parent in selected_child_seqs:
+            if seq is parent and seq.is_finished():
+                self.scheduler.free_seq(seq)
+
+        # Remove the unselected parent sequences from the sequence group and
+        # free their memory in block manager.
+        for seq, parent in unselected_child_seqs:
+            if seq is parent:
+                # Remove the parent sequence if it is not selected for next
+                # iteration
+                seq_group.remove(seq.seq_id)
+                self.scheduler.free_seq(seq)
+
+    def _process_model_outputs(
+        self, output: SamplerOutput, scheduler_outputs: SchedulerOutputs
+    ) -> List[RequestOutput]:
+        # Update the scheduled sequence groups with the model outputs.
+        scheduled_seq_groups = scheduler_outputs.scheduled_seq_groups
+        for seq_group, outputs in zip(scheduled_seq_groups, output):
+            self._process_sequence_group_outputs(seq_group, outputs)
+
+        # Free the finished sequence groups.
+        self.scheduler.free_finished_seq_groups()
+
+        # Create the outputs.
+        request_outputs: List[RequestOutput] = []
+        for seq_group in scheduled_seq_groups + scheduler_outputs.ignored_seq_groups:
+            request_output = RequestOutput.from_seq_group(seq_group)
+            request_outputs.append(request_output)
+
+        if self.log_stats:
+            # Log the system stats.
+            self._log_system_stats(
+                scheduler_outputs.prompt_run, scheduler_outputs.num_batched_tokens
+            )
+        return request_outputs
+
+    def step(self) -> List[RequestOutput]:
+        """Performs one decoding iteration and returns newly generated results.
+
+        This function performs one decoding iteration of the engine. It first
+        schedules the sequences to be executed in the next iteration and the
+        token blocks to be swapped in/out/copy. Then, it executes the model
+        and updates the scheduler with the model outputs. Finally, it decodes
+        the sequences and returns the newly generated results.
+        """
+        seq_group_metadata_list, scheduler_outputs = self.scheduler.schedule()
+
+        if not scheduler_outputs.is_empty():
+            # Execute the model.
+            all_outputs = self._run_workers(
+                "execute_model",
+                driver_kwargs={
+                    "seq_group_metadata_list": seq_group_metadata_list,
+                    "blocks_to_swap_in": scheduler_outputs.blocks_to_swap_in,
+                    "blocks_to_swap_out": scheduler_outputs.blocks_to_swap_out,
+                    "blocks_to_copy": scheduler_outputs.blocks_to_copy,
+                },
+            )
+
+            # Only the driver worker returns the sampling results.
+            output = all_outputs[0]
+        else:
+            output = []
+
+        return self._process_model_outputs(output, scheduler_outputs)
+
+    def _log_system_stats(
+        self,
+        prompt_run: bool,
+        num_batched_tokens: int,
+    ) -> None:
+        now = time.monotonic()
+        # Log the number of batched input tokens.
+        if prompt_run:
+            self.num_prompt_tokens.append((now, num_batched_tokens))
+        else:
+            self.num_generation_tokens.append((now, num_batched_tokens))
+
+        should_log = now - self.last_logging_time >= _LOGGING_INTERVAL_SEC
+        if not should_log:
+            return
+
+        # Discard the old stats.
+        self.num_prompt_tokens = [
+            (t, n) for t, n in self.num_prompt_tokens if now - t < _LOGGING_INTERVAL_SEC
+        ]
+        self.num_generation_tokens = [
+            (t, n)
+            for t, n in self.num_generation_tokens
+            if now - t < _LOGGING_INTERVAL_SEC
+        ]
+
+        if len(self.num_prompt_tokens) > 1:
+            total_num_tokens = sum(n for _, n in self.num_prompt_tokens[:-1])
+            window = now - self.num_prompt_tokens[0][0]
+            avg_prompt_throughput = total_num_tokens / window
+        else:
+            avg_prompt_throughput = 0.0
+        if len(self.num_generation_tokens) > 1:
+            total_num_tokens = sum(n for _, n in self.num_generation_tokens[:-1])
+            window = now - self.num_generation_tokens[0][0]
+            avg_generation_throughput = total_num_tokens / window
+        else:
+            avg_generation_throughput = 0.0
+
+        total_num_gpu_blocks = self.cache_config.num_gpu_blocks
+        num_free_gpu_blocks = self.scheduler.block_manager.get_num_free_gpu_blocks()
+        num_used_gpu_blocks = total_num_gpu_blocks - num_free_gpu_blocks
+        gpu_cache_usage = num_used_gpu_blocks / total_num_gpu_blocks
+
+        total_num_cpu_blocks = self.cache_config.num_cpu_blocks
+        if total_num_cpu_blocks > 0:
+            num_free_cpu_blocks = self.scheduler.block_manager.get_num_free_cpu_blocks()
+            num_used_cpu_blocks = total_num_cpu_blocks - num_free_cpu_blocks
+            cpu_cache_usage = num_used_cpu_blocks / total_num_cpu_blocks
+        else:
+            cpu_cache_usage = 0.0
+
+        record_metrics(
+            avg_prompt_throughput=avg_prompt_throughput,
+            avg_generation_throughput=avg_generation_throughput,
+            scheduler_running=len(self.scheduler.running),
+            scheduler_swapped=len(self.scheduler.swapped),
+            scheduler_waiting=len(self.scheduler.waiting),
+            gpu_cache_usage=gpu_cache_usage,
+            cpu_cache_usage=cpu_cache_usage,
+        )
+
+        logger.info(
+            "Avg prompt throughput: "
+            f"{avg_prompt_throughput:.1f} tokens/s, "
+            "Avg generation throughput: "
+            f"{avg_generation_throughput:.1f} tokens/s, "
+            f"Running: {len(self.scheduler.running)} reqs, "
+            f"Swapped: {len(self.scheduler.swapped)} reqs, "
+            f"Pending: {len(self.scheduler.waiting)} reqs, "
+            f"GPU KV cache usage: {gpu_cache_usage * 100:.1f}%, "
+            f"CPU KV cache usage: {cpu_cache_usage * 100:.1f}%"
+        )
+        self.last_logging_time = now
+
+    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 = (
+            detokenize_incrementally(
+                self.tokenizer,
+                all_input_ids=seq.get_token_ids(),
+                prev_tokens=seq.tokens,
+                prefix_offset=seq.prefix_offset,
+                read_offset=seq.read_offset,
+                skip_special_tokens=prms.skip_special_tokens,
+                spaces_between_special_tokens=prms.spaces_between_special_tokens,
+            )
+        )
+        if seq.tokens is None:
+            seq.tokens = new_tokens
+        else:
+            seq.tokens.extend(new_tokens)
+        seq.prefix_offset = prefix_offset
+        seq.read_offset = read_offset
+        seq.output_text += new_output_text
+
+    def _check_stop(self, seq: Sequence, sampling_params: SamplingParams) -> None:
+        """Stop the finished sequences."""
+        for stop_str in sampling_params.stop:
+            if seq.output_text.endswith(stop_str):
+                if not sampling_params.include_stop_str_in_output:
+                    # Truncate the output text so that the stop string is
+                    # not included in the output.
+                    seq.output_text = seq.output_text[: -len(stop_str)]
+                seq.status = SequenceStatus.FINISHED_STOPPED
+                return
+        if seq.data.finished:
+            seq.status = SequenceStatus.FINISHED_STOPPED
+            return
+
+        for token_id in seq.get_last_token_id():
+            if token_id == sampling_params.eos_token:
+                seq.status = SequenceStatus.FINISHED_STOPPED
+                return
+
+        # Check if the sequence has reached max_model_len.
+        if seq.get_len() > self.scheduler_config.max_model_len:
+            seq.status = SequenceStatus.FINISHED_LENGTH_CAPPED
+            return
+
+        # Check if the sequence has reached max_tokens.
+        if seq.get_output_len() == sampling_params.max_tokens:
+            seq.status = SequenceStatus.FINISHED_LENGTH_CAPPED
+            return
+
+        # Check if the sequence has generated the EOS token.
+        if (not sampling_params.ignore_eos) and seq.get_last_token_id()[
+            0
+        ] == sampling_params.eos_token:
+            seq.status = SequenceStatus.FINISHED_STOPPED
+            return
+
+    def _run_workers(
+        self,
+        method: str,
+        *args,
+        driver_args: Optional[List[Any]] = None,
+        driver_kwargs: Optional[Dict[str, Any]] = None,
+        max_concurrent_workers: Optional[int] = None,
+        **kwargs,
+    ) -> Any:
+        """Runs the given method on all workers."""
+
+        if max_concurrent_workers:
+            raise NotImplementedError("max_concurrent_workers is not supported yet.")
+
+        # Start the ray workers first.
+        ray_worker_outputs = [
+            worker.execute_method.remote(method, *args, **kwargs)
+            for worker in self.workers
+        ]
+
+        if driver_args is None:
+            driver_args = args
+        if driver_kwargs is None:
+            driver_kwargs = kwargs
+
+        # Start the driver worker after all the ray workers.
+        driver_worker_output = getattr(self.driver_worker, method)(
+            *driver_args, **driver_kwargs
+        )
+
+        # Get the results of the ray workers.
+        if self.workers:
+            ray_worker_outputs = ray.get(ray_worker_outputs)
+
+        return [driver_worker_output] + ray_worker_outputs

ChatTTS/model/velocity/model_loader.py

@@ -0,0 +1,69 @@
+"""Utilities for selecting and loading models."""
+
+import contextlib
+
+import torch
+import torch.nn as nn
+
+from vllm.config import ModelConfig
+from vllm.model_executor.models import ModelRegistry
+from vllm.model_executor.weight_utils import get_quant_config, initialize_dummy_weights
+
+from .llama import LlamaModel
+
+
+@contextlib.contextmanager
+def _set_default_torch_dtype(dtype: torch.dtype):
+    """Sets the default torch dtype to the given dtype."""
+    old_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    yield
+    torch.set_default_dtype(old_dtype)
+
+
+def get_model(model_config: ModelConfig) -> nn.Module:
+    # Get the (maybe quantized) linear method.
+    linear_method = None
+    if model_config.quantization is not None:
+        quant_config = get_quant_config(
+            model_config.quantization,
+            model_config.model,
+            model_config.hf_config,
+            model_config.download_dir,
+        )
+        capability = torch.cuda.get_device_capability()
+        capability = capability[0] * 10 + capability[1]
+        if capability < quant_config.get_min_capability():
+            raise ValueError(
+                f"The quantization method {model_config.quantization} is not "
+                "supported for the current GPU. "
+                f"Minimum capability: {quant_config.get_min_capability()}. "
+                f"Current capability: {capability}."
+            )
+        supported_dtypes = quant_config.get_supported_act_dtypes()
+        if model_config.dtype not in supported_dtypes:
+            raise ValueError(
+                f"{model_config.dtype} is not supported for quantization "
+                f"method {model_config.quantization}. Supported dtypes: "
+                f"{supported_dtypes}"
+            )
+        linear_method = quant_config.get_linear_method()
+
+    with _set_default_torch_dtype(model_config.dtype):
+        # Create a model instance.
+        # The weights will be initialized as empty tensors.
+        with torch.device("cuda"):
+            model = LlamaModel(model_config.hf_config, linear_method)
+        if model_config.load_format == "dummy":
+            # NOTE(woosuk): For accurate performance evaluation, we assign
+            # random values to the weights.
+            initialize_dummy_weights(model)
+        else:
+            # Load the weights from the cached or downloaded files.
+            model.load_weights(
+                model_config.model,
+                model_config.download_dir,
+                model_config.load_format,
+                model_config.revision,
+            )
+    return model.eval()

ChatTTS/model/velocity/model_runner.py

@@ -0,0 +1,817 @@
+import time
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from .configs import ModelConfig, ParallelConfig, SchedulerConfig
+from vllm.logger import init_logger
+from .model_loader import get_model
+from vllm.model_executor import InputMetadata, SamplingMetadata
+from vllm.model_executor.parallel_utils.communication_op import (
+    broadcast,
+    broadcast_object_list,
+)
+from .sampling_params import SamplingParams, SamplingType
+from .sequence import (
+    SamplerOutput,
+    SequenceData,
+    SequenceGroupMetadata,
+    SequenceGroupOutput,
+    SequenceOutput,
+)
+from vllm.utils import in_wsl
+from ..embed import Embed
+from .sampler import Sampler
+from safetensors.torch import safe_open
+
+logger = init_logger(__name__)
+
+KVCache = Tuple[torch.Tensor, torch.Tensor]
+_PAD_SLOT_ID = -1
+# Capture graphs for batch size 1, 2, 4, 8, 16, 24, 32, 40, ..., 256.
+# NOTE: _get_graph_batch_size needs to be updated if this list is changed.
+_BATCH_SIZES_TO_CAPTURE = [1, 2, 4] + [8 * i for i in range(1, 33)]
+
+
+class ModelRunner:
+
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        parallel_config: ParallelConfig,
+        scheduler_config: SchedulerConfig,
+        is_driver_worker: bool = False,
+        post_model_path: str = None,
+    ):
+        self.model_config = model_config
+        self.parallel_config = parallel_config
+        self.scheduler_config = scheduler_config
+        self.is_driver_worker = is_driver_worker
+        self.post_model_path = post_model_path
+
+        # model_config can be None in tests/samplers/test_sampler.py.
+        # FIXME(woosuk): This is a hack to make the tests work. Refactor this.
+        self.sliding_window = (
+            model_config.get_sliding_window() if model_config is not None else None
+        )
+        self.model = None
+        self.block_size = None  # Set after initial profiling.
+
+        self.graph_runners: Dict[int, CUDAGraphRunner] = {}
+        self.graph_memory_pool = None  # Set during graph capture.
+
+        self.max_context_len_to_capture = (
+            self.model_config.max_context_len_to_capture
+            if self.model_config is not None
+            else 0
+        )
+        # When using CUDA graph, the input block tables must be padded to
+        # max_context_len_to_capture. However, creating the block table in
+        # Python can be expensive. To optimize this, we cache the block table
+        # in numpy and only copy the actual input content at every iteration.
+        # The shape of the cached block table will be
+        # (max batch size to capture, max context len to capture / block size).
+        self.graph_block_tables = None  # Set after initial profiling.
+        # cache in_wsl result
+        self.in_wsl = in_wsl()
+
+    def load_model(self) -> None:
+        self.model = get_model(self.model_config)
+        self.post_model = Embed(
+            self.model_config.get_hidden_size(),
+            self.model_config.num_audio_tokens,
+            self.model_config.num_text_tokens,
+        )
+        state_dict_tensors = {}
+        with safe_open(self.post_model_path, framework="pt", device=0) as f:
+            for k in f.keys():
+                state_dict_tensors[k] = f.get_tensor(k)
+        self.post_model.load_state_dict(state_dict_tensors)
+        self.post_model.to(next(self.model.parameters())).eval()
+        self.sampler = Sampler(self.post_model, self.model_config.num_audio_tokens, 4)
+
+    def set_block_size(self, block_size: int) -> None:
+        self.block_size = block_size
+
+        max_num_blocks = (
+            self.max_context_len_to_capture + block_size - 1
+        ) // block_size
+        self.graph_block_tables = np.zeros(
+            (max(_BATCH_SIZES_TO_CAPTURE), max_num_blocks), dtype=np.int32
+        )
+
+    def _prepare_prompt(
+        self,
+        seq_group_metadata_list: List[SequenceGroupMetadata],
+    ) -> Tuple[torch.Tensor, torch.Tensor, InputMetadata, List[int]]:
+        assert len(seq_group_metadata_list) > 0
+        input_tokens: List[List[int]] = []
+        input_positions: List[List[int]] = []
+        slot_mapping: List[List[int]] = []
+
+        prompt_lens: List[int] = []
+        for seq_group_metadata in seq_group_metadata_list:
+            assert seq_group_metadata.is_prompt
+            seq_ids = list(seq_group_metadata.seq_data.keys())
+            assert len(seq_ids) == 1
+            seq_id = seq_ids[0]
+
+            seq_data = seq_group_metadata.seq_data[seq_id]
+            prompt_tokens = seq_data.get_token_ids()
+            prompt_len = len(prompt_tokens)
+            prompt_lens.append(prompt_len)
+
+            input_tokens.append(prompt_tokens)
+            # NOTE(woosuk): Here we assume that the first token in the prompt
+            # is always the first token in the sequence.
+            input_positions.append(list(range(prompt_len)))
+
+            if seq_group_metadata.block_tables is None:
+                # During memory profiling, the block tables are not initialized
+                # yet. In this case, we just use a dummy slot mapping.
+                slot_mapping.append([_PAD_SLOT_ID] * prompt_len)
+                continue
+
+            # Compute the slot mapping.
+            slot_mapping.append([])
+            block_table = seq_group_metadata.block_tables[seq_id]
+            # Mask the [0, start_idx) tokens of the prompt with _PAD_SLOT_ID,
+            # where start_idx is max(0, prompt_len - sliding_window).
+            # For example, if the prompt len is 10, sliding window is 8, and
+            # block size is 4, the first two tokens are masked and the slot
+            # mapping will be [-1, -1, 2, 3, 4, 5, 6, 7, 0, 1].
+            start_idx = 0
+            if self.sliding_window is not None:
+                start_idx = max(0, prompt_len - self.sliding_window)
+            for i in range(prompt_len):
+                if i < start_idx:
+                    slot_mapping[-1].append(_PAD_SLOT_ID)
+                    continue
+
+                block_number = block_table[i // self.block_size]
+                block_offset = i % self.block_size
+                slot = block_number * self.block_size + block_offset
+                slot_mapping[-1].append(slot)
+
+        max_prompt_len = max(prompt_lens)
+        input_tokens = _make_tensor_with_pad(
+            input_tokens, max_prompt_len, pad=0, dtype=torch.long
+        )
+        input_positions = _make_tensor_with_pad(
+            input_positions, max_prompt_len, pad=0, dtype=torch.long
+        )
+        slot_mapping = _make_tensor_with_pad(
+            slot_mapping, max_prompt_len, pad=_PAD_SLOT_ID, dtype=torch.long
+        )
+
+        input_metadata = InputMetadata(
+            is_prompt=True,
+            slot_mapping=slot_mapping,
+            max_context_len=None,
+            context_lens=None,
+            block_tables=None,
+            use_cuda_graph=False,
+        )
+        return input_tokens, input_positions, input_metadata, prompt_lens
+
+    def _prepare_decode(
+        self,
+        seq_group_metadata_list: List[SequenceGroupMetadata],
+    ) -> Tuple[torch.Tensor, torch.Tensor, InputMetadata]:
+        assert len(seq_group_metadata_list) > 0
+        input_tokens: List[List[int]] = []
+        input_positions: List[List[int]] = []
+        slot_mapping: List[List[int]] = []
+        context_lens: List[int] = []
+        block_tables: List[List[int]] = []
+
+        for seq_group_metadata in seq_group_metadata_list:
+            assert not seq_group_metadata.is_prompt
+
+            seq_ids = list(seq_group_metadata.seq_data.keys())
+            for seq_id in seq_ids:
+                seq_data = seq_group_metadata.seq_data[seq_id]
+                generation_token = seq_data.get_last_token_id()
+                input_tokens.append([generation_token])
+
+                seq_len = seq_data.get_len()
+                position = seq_len - 1
+                input_positions.append([position])
+
+                context_len = (
+                    seq_len
+                    if self.sliding_window is None
+                    else min(seq_len, self.sliding_window)
+                )
+                context_lens.append(context_len)
+
+                block_table = seq_group_metadata.block_tables[seq_id]
+                block_number = block_table[position // self.block_size]
+                block_offset = position % self.block_size
+                slot = block_number * self.block_size + block_offset
+                slot_mapping.append([slot])
+
+                if self.sliding_window is not None:
+                    sliding_window_blocks = self.sliding_window // self.block_size
+                    block_table = block_table[-sliding_window_blocks:]
+                block_tables.append(block_table)
+
+        batch_size = len(input_tokens)
+        max_context_len = max(context_lens)
+        use_captured_graph = (
+            not self.model_config.enforce_eager
+            and batch_size <= _BATCH_SIZES_TO_CAPTURE[-1]
+            and max_context_len <= self.max_context_len_to_capture
+        )
+        if use_captured_graph:
+            # Pad the input tokens, positions, and slot mapping to match the
+            # batch size of the captured graph.
+            graph_batch_size = _get_graph_batch_size(batch_size)
+            assert graph_batch_size >= batch_size
+            for _ in range(graph_batch_size - batch_size):
+                input_tokens.append([])
+                input_positions.append([])
+                slot_mapping.append([])
+                context_lens.append(1)
+                block_tables.append([])
+            batch_size = graph_batch_size
+
+        input_tokens = _make_tensor_with_pad(
+            input_tokens, max_len=1, pad=0, dtype=torch.long, device="cuda"
+        )
+        input_positions = _make_tensor_with_pad(
+            input_positions, max_len=1, pad=0, dtype=torch.long, device="cuda"
+        )
+        slot_mapping = _make_tensor_with_pad(
+            slot_mapping, max_len=1, pad=_PAD_SLOT_ID, dtype=torch.long, device="cuda"
+        )
+        context_lens = torch.tensor(context_lens, dtype=torch.int, device="cuda")
+
+        if use_captured_graph:
+            # The shape of graph_block_tables is
+            # [max batch size, max context len // block size].
+            input_block_tables = self.graph_block_tables[:batch_size]
+            for i, block_table in enumerate(block_tables):
+                if block_table:
+                    input_block_tables[i, : len(block_table)] = block_table
+            block_tables = torch.tensor(input_block_tables, device="cuda")
+        else:
+            block_tables = _make_tensor_with_pad(
+                block_tables,
+                max_len=max_context_len,
+                pad=0,
+                dtype=torch.int,
+                device="cuda",
+            )
+
+        input_metadata = InputMetadata(
+            is_prompt=False,
+            slot_mapping=slot_mapping,
+            max_context_len=max_context_len,
+            context_lens=context_lens,
+            block_tables=block_tables,
+            use_cuda_graph=use_captured_graph,
+        )
+        return input_tokens, input_positions, input_metadata
+
+    def _prepare_sample(
+        self,
+        seq_group_metadata_list: List[SequenceGroupMetadata],
+        prompt_lens: List[int],
+    ) -> SamplingMetadata:
+        seq_groups: List[Tuple[List[int], SamplingParams]] = []
+        selected_token_indices: List[int] = []
+        selected_token_start_idx = 0
+        categorized_sample_indices = {t: [] for t in SamplingType}
+        categorized_sample_indices_start_idx = 0
+
+        max_prompt_len = max(prompt_lens) if prompt_lens else 1
+        for i, seq_group_metadata in enumerate(seq_group_metadata_list):
+            seq_ids = list(seq_group_metadata.seq_data.keys())
+            sampling_params = seq_group_metadata.sampling_params
+            seq_groups.append((seq_ids, sampling_params))
+
+            if seq_group_metadata.is_prompt:
+                assert len(seq_ids) == 1
+                prompt_len = prompt_lens[i]
+                if sampling_params.prompt_logprobs is not None:
+                    # NOTE: prompt token positions do not need sample, skip
+                    categorized_sample_indices_start_idx += prompt_len - 1
+
+                categorized_sample_indices[sampling_params.sampling_type].append(
+                    categorized_sample_indices_start_idx
+                )
+                categorized_sample_indices_start_idx += 1
+
+                if sampling_params.prompt_logprobs is not None:
+                    selected_token_indices.extend(
+                        range(
+                            selected_token_start_idx,
+                            selected_token_start_idx + prompt_len - 1,
+                        )
+                    )
+                selected_token_indices.append(selected_token_start_idx + prompt_len - 1)
+                selected_token_start_idx += max_prompt_len
+            else:
+                num_seqs = len(seq_ids)
+                selected_token_indices.extend(
+                    range(selected_token_start_idx, selected_token_start_idx + num_seqs)
+                )
+                selected_token_start_idx += num_seqs
+
+                categorized_sample_indices[sampling_params.sampling_type].extend(
+                    range(
+                        categorized_sample_indices_start_idx,
+                        categorized_sample_indices_start_idx + num_seqs,
+                    )
+                )
+                categorized_sample_indices_start_idx += num_seqs
+
+        selected_token_indices = _async_h2d(
+            selected_token_indices, dtype=torch.long, pin_memory=not self.in_wsl
+        )
+        categorized_sample_indices = {
+            t: _async_h2d(seq_ids, dtype=torch.int, pin_memory=not self.in_wsl)
+            for t, seq_ids in categorized_sample_indices.items()
+        }
+
+        seq_data: Dict[int, SequenceData] = {}
+        for seq_group_metadata in seq_group_metadata_list:
+            seq_data.update(seq_group_metadata.seq_data)
+
+        sampling_metadata = SamplingMetadata(
+            seq_groups=seq_groups,
+            seq_data=seq_data,
+            prompt_lens=prompt_lens,
+            selected_token_indices=selected_token_indices,
+            categorized_sample_indices=categorized_sample_indices,
+        )
+        return sampling_metadata
+
+    def prepare_input_tensors(
+        self,
+        seq_group_metadata_list: Optional[List[SequenceGroupMetadata]],
+    ) -> Tuple[torch.Tensor, torch.Tensor, InputMetadata, SamplingMetadata]:
+        if self.is_driver_worker:
+            # NOTE: We assume that all sequences in the group are all prompts or
+            # all decodes.
+            is_prompt = seq_group_metadata_list[0].is_prompt
+            # Prepare input tensors.
+            if is_prompt:
+                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(
+                    seq_group_metadata_list
+                )
+                prompt_lens = []
+            sampling_metadata = self._prepare_sample(
+                seq_group_metadata_list, prompt_lens
+            )
+
+            def get_size_or_none(x: Optional[torch.Tensor]):
+                return x.size() if x is not None else None
+
+            # Broadcast the input data. For input tensors, we first broadcast
+            # its shape and then broadcast the tensor to avoid high
+            # serialization cost.
+            py_data = {
+                "input_tokens_size": input_tokens.size(),
+                "input_positions_size": input_positions.size(),
+                "is_prompt": input_metadata.is_prompt,
+                "slot_mapping_size": get_size_or_none(input_metadata.slot_mapping),
+                "max_context_len": input_metadata.max_context_len,
+                "context_lens_size": get_size_or_none(input_metadata.context_lens),
+                "block_tables_size": get_size_or_none(input_metadata.block_tables),
+                "use_cuda_graph": input_metadata.use_cuda_graph,
+                "selected_token_indices_size": sampling_metadata.selected_token_indices.size(),
+            }
+            broadcast_object_list([py_data], src=0)
+            # TODO(zhuohan): Combine the broadcasts or set async_op=True.
+            broadcast(input_tokens, src=0)
+            broadcast(input_positions, src=0)
+            if input_metadata.slot_mapping is not None:
+                broadcast(input_metadata.slot_mapping, src=0)
+            if input_metadata.context_lens is not None:
+                broadcast(input_metadata.context_lens, src=0)
+            if input_metadata.block_tables is not None:
+                broadcast(input_metadata.block_tables, src=0)
+            broadcast(sampling_metadata.selected_token_indices, src=0)
+        else:
+            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"
+            )
+            broadcast(input_tokens, src=0)
+            input_positions = torch.empty(
+                *py_data["input_positions_size"], dtype=torch.long, device="cuda"
+            )
+            broadcast(input_positions, src=0)
+            if py_data["slot_mapping_size"] is not None:
+                slot_mapping = torch.empty(
+                    *py_data["slot_mapping_size"], dtype=torch.long, device="cuda"
+                )
+                broadcast(slot_mapping, src=0)
+            else:
+                slot_mapping = None
+            if py_data["context_lens_size"] is not None:
+                context_lens = torch.empty(
+                    *py_data["context_lens_size"], dtype=torch.int, device="cuda"
+                )
+                broadcast(context_lens, src=0)
+            else:
+                context_lens = None
+            if py_data["block_tables_size"] is not None:
+                block_tables = torch.empty(
+                    *py_data["block_tables_size"], dtype=torch.int, device="cuda"
+                )
+                broadcast(block_tables, src=0)
+            else:
+                block_tables = None
+            selected_token_indices = torch.empty(
+                *py_data["selected_token_indices_size"], dtype=torch.long, device="cuda"
+            )
+            broadcast(selected_token_indices, src=0)
+            input_metadata = InputMetadata(
+                is_prompt=py_data["is_prompt"],
+                slot_mapping=slot_mapping,
+                max_context_len=py_data["max_context_len"],
+                context_lens=context_lens,
+                block_tables=block_tables,
+                use_cuda_graph=py_data["use_cuda_graph"],
+            )
+            sampling_metadata = SamplingMetadata(
+                seq_groups=None,
+                seq_data=None,
+                prompt_lens=None,
+                selected_token_indices=selected_token_indices,
+                categorized_sample_indices=None,
+                perform_sampling=False,
+            )
+
+        return input_tokens, input_positions, input_metadata, sampling_metadata
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        seq_group_metadata_list: Optional[List[SequenceGroupMetadata]],
+        kv_caches: List[Tuple[torch.Tensor, torch.Tensor]],
+    ) -> Optional[SamplerOutput]:
+        input_tokens, input_positions, input_metadata, sampling_metadata = (
+            self.prepare_input_tensors(seq_group_metadata_list)
+        )
+        # print(sampling_metadata.seq_data)
+        seq_groups = []
+        input_tokens_history = []
+        for i, rtn in enumerate(sampling_metadata.seq_groups):
+            seq_groups.append(rtn[0][0])
+            tokens_history = sampling_metadata.seq_data[rtn[0][0]].output_token_ids
+            if len(tokens_history) >= 1:
+                if len(tokens_history[0]) == 1:
+                    tokens_history = [token[0] for token in tokens_history]
+                else:
+                    tokens_history = [list(token) for token in tokens_history]
+            input_tokens_history.append(tokens_history)
+        input_tokens_history = torch.tensor(input_tokens_history).to(
+            input_tokens.device
+        )
+        # token_ids = rtn.outputs[0].token_ids
+        # for j, token_id in enumerate(token_ids):
+        #     if len(token_id) == 1:
+        #         token_ids[j] = token_id[0]
+        #     else:
+        #         token_ids[j] = list(token_id)
+
+        # Execute the model.
+        # print("it1",input_tokens)
+        if len(input_tokens.shape) == 2:
+            input_tokens = input_tokens.unsqueeze(2).repeat(1, 1, 4)
+        if len(input_tokens_history.shape) == 2:
+            input_tokens_history = input_tokens_history.unsqueeze(2).repeat(1, 1, 4)
+        # print(input_tokens_history.shape)
+        # print("it2",input_tokens.shape)
+        text_mask = input_tokens != 0
+        text_mask = text_mask[:, :, 0]
+
+        if input_metadata.use_cuda_graph:
+            graph_batch_size = input_tokens.shape[0]
+            model_executable = self.graph_runners[graph_batch_size]
+        else:
+            model_executable = self.model
+
+        infer_text = sampling_metadata.seq_groups[0][1].infer_text
+        temperature = sampling_metadata.seq_groups[0][1].temperature
+        if not infer_text:
+            temperature = torch.tensor(temperature).to(input_tokens.device)
+        logits_processors, logits_warpers = sampling_metadata.seq_groups[0][
+            1
+        ].logits_processors
+        # print(logits_processors, logits_warpers)
+        min_new_token = sampling_metadata.seq_groups[0][1].min_new_token
+        eos_token = sampling_metadata.seq_groups[0][1].eos_token
+        start_idx = sampling_metadata.seq_groups[0][1].start_idx
+        if input_tokens.shape[-2] == 1:
+            if infer_text:
+                input_emb: torch.Tensor = self.post_model.emb_text(
+                    input_tokens[:, :, 0]
+                )
+            else:
+                code_emb = [
+                    self.post_model.emb_code[i](input_tokens[:, :, i])
+                    for i in range(self.post_model.num_vq)
+                ]
+                input_emb = torch.stack(code_emb, 3).sum(3)
+                start_idx = (
+                    input_tokens_history.shape[-2] - 1
+                    if input_tokens_history.shape[-2] > 0
+                    else 0
+                )
+        else:
+            input_emb = self.post_model(input_tokens, text_mask)
+        # print(input_emb.shape)
+        hidden_states = model_executable(
+            input_emb=input_emb,
+            positions=input_positions,
+            kv_caches=kv_caches,
+            input_metadata=input_metadata,
+        )
+        # print(hidden_states.shape)
+        # print(input_tokens)
+        B_NO_PAD = input_tokens_history.shape[0]
+        input_tokens = input_tokens[:B_NO_PAD, :, :]
+        hidden_states = hidden_states[:B_NO_PAD, :, :]
+        idx_next, logprob, finish = self.sampler.sample(
+            inputs_ids=(
+                input_tokens
+                if input_tokens_history.shape[-2] == 0
+                else input_tokens_history
+            ),
+            hidden_states=hidden_states,
+            infer_text=infer_text,
+            temperature=temperature,
+            logits_processors=logits_processors,
+            logits_warpers=logits_warpers,
+            min_new_token=min_new_token,
+            now_length=1,
+            eos_token=eos_token,
+            start_idx=start_idx,
+        )
+        # print(logprob.shape, idx_next.shape)
+        if len(logprob.shape) == 2:
+            logprob = logprob[:, None, :]
+        logprob = torch.gather(logprob, -1, idx_next.transpose(-1, -2))[:, :, 0]
+        # print("测试",idx_next.shape, logprob.shape)
+        # Sample the next token.
+        # output = self.model.sample(
+        #     hidden_states=hidden_states,
+        #     sampling_metadata=sampling_metadata,
+        # )
+        results = []
+        for i in range(idx_next.shape[0]):
+            idx_next_i = idx_next[i, 0, :].tolist()
+            logprob_i = logprob[i].tolist()
+            tmp_hidden_states = hidden_states[i]
+            if input_tokens[i].shape[-2] != 1:
+                tmp_hidden_states = tmp_hidden_states[-1:, :]
+            result = SequenceGroupOutput(
+                samples=[
+                    SequenceOutput(
+                        parent_seq_id=seq_groups[i],
+                        logprobs={tuple(idx_next_i): logprob_i},
+                        output_token=tuple(idx_next_i),
+                        hidden_states=tmp_hidden_states,
+                        finished=finish[i].item(),
+                    ),
+                ],
+                prompt_logprobs=None,
+            )
+            results.append(result)
+        # print(results)
+        # print(idx_next, idx_next.shape, logprob.shape)
+        return results
+
+    @torch.inference_mode()
+    def profile_run(self) -> None:
+        # Enable top-k sampling to reflect the accurate memory usage.
+        vocab_size = self.model_config.get_vocab_size()
+        sampling_params = SamplingParams(
+            top_p=0.99, top_k=vocab_size - 1, infer_text=True
+        )
+        max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
+        max_num_seqs = self.scheduler_config.max_num_seqs
+
+        # Profile memory usage with max_num_sequences sequences and the total
+        # number of tokens equal to max_num_batched_tokens.
+        seqs: List[SequenceGroupMetadata] = []
+        for group_id in range(max_num_seqs):
+            seq_len = max_num_batched_tokens // max_num_seqs + (
+                group_id < max_num_batched_tokens % max_num_seqs
+            )
+            seq_data = SequenceData([0] * seq_len)
+            seq = SequenceGroupMetadata(
+                request_id=str(group_id),
+                is_prompt=True,
+                seq_data={group_id: seq_data},
+                sampling_params=sampling_params,
+                block_tables=None,
+            )
+            seqs.append(seq)
+
+        # Run the model with the dummy inputs.
+        num_layers = self.model_config.get_num_layers(self.parallel_config)
+        kv_caches = [(None, None)] * num_layers
+        self.execute_model(seqs, kv_caches)
+        torch.cuda.synchronize()
+        return
+
+    @torch.inference_mode()
+    def capture_model(self, kv_caches: List[KVCache]) -> None:
+        assert not self.model_config.enforce_eager
+        logger.info(
+            "Capturing the model for CUDA graphs. This may lead to "
+            "unexpected consequences if the model is not static. To "
+            "run the model in eager mode, set 'enforce_eager=True' or "
+            "use '--enforce-eager' in the CLI."
+        )
+        logger.info(
+            "CUDA graphs can take additional 1~3 GiB memory per GPU. "
+            "If you are running out of memory, consider decreasing "
+            "`gpu_memory_utilization` or enforcing eager mode."
+        )
+        start_time = time.perf_counter()
+
+        # Prepare dummy inputs. These will be reused for all batch sizes.
+        max_batch_size = max(_BATCH_SIZES_TO_CAPTURE)
+        input_emb = torch.zeros(
+            max_batch_size,
+            1,
+            self.model_config.get_hidden_size(),
+            dtype=next(self.model.parameters()).dtype,
+        ).cuda()
+        input_positions = torch.zeros(max_batch_size, 1, dtype=torch.long).cuda()
+        slot_mapping = torch.empty(max_batch_size, 1, dtype=torch.long).cuda()
+        slot_mapping.fill_(_PAD_SLOT_ID)
+        context_lens = torch.ones(max_batch_size, dtype=torch.int32).cuda()
+        block_tables = torch.from_numpy(self.graph_block_tables).cuda()
+
+        # NOTE: Capturing the largest batch size first may help reduce the
+        # memory usage of CUDA graph.
+        for batch_size in reversed(_BATCH_SIZES_TO_CAPTURE):
+            # Create dummy input_metadata.
+            input_metadata = InputMetadata(
+                is_prompt=False,
+                slot_mapping=slot_mapping[:batch_size],
+                max_context_len=self.max_context_len_to_capture,
+                context_lens=context_lens[:batch_size],
+                block_tables=block_tables[:batch_size],
+                use_cuda_graph=True,
+            )
+
+            graph_runner = CUDAGraphRunner(self.model)
+            graph_runner.capture(
+                input_emb[:batch_size],
+                input_positions[:batch_size],
+                kv_caches,
+                input_metadata,
+                memory_pool=self.graph_memory_pool,
+            )
+            self.graph_memory_pool = graph_runner.graph.pool()
+            self.graph_runners[batch_size] = graph_runner
+
+        end_time = time.perf_counter()
+        elapsed_time = end_time - start_time
+        # This usually takes < 10 seconds.
+        logger.info(f"Graph capturing finished in {elapsed_time:.0f} secs.")
+
+
+class CUDAGraphRunner:
+
+    def __init__(self, model: nn.Module):
+        self.model = model
+        self.graph = None
+        self.input_buffers: Dict[str, torch.Tensor] = {}
+        self.output_buffers: Dict[str, torch.Tensor] = {}
+
+    def capture(
+        self,
+        input_emb: torch.Tensor,
+        positions: torch.Tensor,
+        kv_caches: List[KVCache],
+        input_metadata: InputMetadata,
+        memory_pool,
+    ) -> None:
+        assert self.graph is None
+        # Run the model once without capturing the graph.
+        # This is to make sure that the captured graph does not include the
+        # kernel launches for initial benchmarking (e.g., Triton autotune).
+        self.model(
+            input_emb,
+            positions,
+            kv_caches,
+            input_metadata,
+        )
+        torch.cuda.synchronize()
+
+        # Capture the graph.
+        self.graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(self.graph, pool=memory_pool):
+            hidden_states = self.model(
+                input_emb,
+                positions,
+                kv_caches,
+                input_metadata,
+            )
+        torch.cuda.synchronize()
+
+        # Save the input and output buffers.
+        self.input_buffers = {
+            "input_emb": input_emb,
+            "positions": positions,
+            "kv_caches": kv_caches,
+            "slot_mapping": input_metadata.slot_mapping,
+            "context_lens": input_metadata.context_lens,
+            "block_tables": input_metadata.block_tables,
+        }
+        self.output_buffers = {"hidden_states": hidden_states}
+        return
+
+    def forward(
+        self,
+        input_emb: torch.Tensor,
+        positions: torch.Tensor,
+        kv_caches: List[Tuple[torch.Tensor, torch.Tensor]],
+        input_metadata: InputMetadata,
+    ) -> torch.Tensor:
+        # KV caches are fixed tensors, so we don't need to copy them.
+        del kv_caches
+
+        # Copy the input tensors to the input buffers.
+        self.input_buffers["input_emb"].copy_(input_emb, non_blocking=True)
+        self.input_buffers["positions"].copy_(positions, non_blocking=True)
+        self.input_buffers["slot_mapping"].copy_(
+            input_metadata.slot_mapping, non_blocking=True
+        )
+        self.input_buffers["context_lens"].copy_(
+            input_metadata.context_lens, non_blocking=True
+        )
+        self.input_buffers["block_tables"].copy_(
+            input_metadata.block_tables, non_blocking=True
+        )
+
+        # Run the graph.
+        self.graph.replay()
+
+        # Return the output tensor.
+        return self.output_buffers["hidden_states"]
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+
+
+def _pad_to_max(x: List[int], max_len: int, pad: int) -> List[int]:
+    assert len(x) <= max_len
+    if len(x) == max_len:
+        return list(x)
+    return list(x) + [pad] * (max_len - len(x))
+
+
+def _make_tensor_with_pad(
+    x: List[List[int]],
+    max_len: int,
+    pad: int,
+    dtype: torch.dtype,
+    device: Union[str, torch.device] = "cuda",
+    pin_memory: bool = False,
+) -> torch.Tensor:
+    padded_x = []
+    for x_i in x:
+        pad_i = pad
+        if isinstance(x[0][0], tuple):
+            pad_i = (0,) * len(x[0][0])
+        padded_x.append(_pad_to_max(x_i, max_len, pad_i))
+
+    return torch.tensor(
+        padded_x,
+        dtype=dtype,
+        device=device,
+        pin_memory=pin_memory and str(device) == "cpu",
+    )
+
+
+def _get_graph_batch_size(batch_size: int) -> int:
+    if batch_size <= 2:
+        return batch_size
+    elif batch_size <= 4:
+        return 4
+    else:
+        return (batch_size + 7) // 8 * 8
+
+
+def _async_h2d(data: list, dtype, pin_memory):
+    t = torch.tensor(data, dtype=dtype, pin_memory=pin_memory)
+    return t.to(device="cuda", non_blocking=True)

ChatTTS/model/velocity/output.py

@@ -0,0 +1,144 @@
+from typing import List, Optional
+import torch
+
+from .sequence import (
+    PromptLogprobs,
+    SampleLogprobs,
+    SequenceGroup,
+    SequenceStatus,
+)
+
+
+class CompletionOutput:
+    """The output data of one completion output of a request.
+
+    Args:
+        index: The index of the output in the request.
+        text: The generated output text.
+        token_ids: The token IDs of the generated output text.
+        cumulative_logprob: The cumulative log probability of the generated
+            output text.
+        logprobs: The log probabilities of the top probability words at each
+            position if the logprobs are requested.
+        finish_reason: The reason why the sequence is finished.
+    """
+
+    def __init__(
+        self,
+        index: int,
+        text: str,
+        token_ids: List[int],
+        cumulative_logprob: float,
+        logprobs: Optional[SampleLogprobs],
+        finish_reason: Optional[str] = None,
+        hidden_states: Optional[torch.Tensor] = None,
+    ) -> None:
+        self.index = index
+        self.text = text
+        self.token_ids = token_ids
+        self.cumulative_logprob = cumulative_logprob
+        self.logprobs = logprobs
+        self.finish_reason = finish_reason
+        self.hidden_states = hidden_states
+
+    def finished(self) -> bool:
+        return self.finish_reason is not None
+
+    def __repr__(self) -> str:
+        return (
+            f"CompletionOutput(index={self.index}, "
+            f"text={self.text!r}, "
+            f"token_ids={self.token_ids}, "
+            f"cumulative_logprob={self.cumulative_logprob}, "
+            f"logprobs={self.logprobs}, "
+            f"finish_reason={self.finish_reason}, "
+            f"hidden_states={self.hidden_states.shape if self.hidden_states is not None else None})"
+        )
+
+
+class RequestOutput:
+    """The output data of a request to the LLM.
+
+    Args:
+        request_id: The unique ID of the request.
+        prompt: The prompt string of the request.
+        prompt_token_ids: The token IDs of the prompt.
+        prompt_logprobs: The log probabilities to return per prompt token.
+        outputs: The output sequences of the request.
+        finished: Whether the whole request is finished.
+    """
+
+    def __init__(
+        self,
+        request_id: str,
+        prompt: str,
+        prompt_token_ids: List[int],
+        prompt_logprobs: Optional[PromptLogprobs],
+        outputs: List[CompletionOutput],
+        finished: bool,
+    ) -> None:
+        self.request_id = request_id
+        self.prompt = prompt
+        self.prompt_token_ids = prompt_token_ids
+        self.prompt_logprobs = prompt_logprobs
+        self.outputs = outputs
+        self.finished = finished
+
+    @classmethod
+    def from_seq_group(cls, seq_group: SequenceGroup) -> "RequestOutput":
+        # Get the top-n sequences.
+        n = seq_group.sampling_params.n
+        seqs = seq_group.get_seqs()
+        if seq_group.sampling_params.use_beam_search:
+            sorting_key = lambda seq: seq.get_beam_search_score(
+                seq_group.sampling_params.length_penalty
+            )
+        else:
+            sorting_key = lambda seq: seq.get_cumulative_logprob()
+        sorted_seqs = sorted(seqs, key=sorting_key, reverse=True)
+        top_n_seqs = sorted_seqs[:n]
+
+        # Create the outputs.
+        outputs: List[CompletionOutput] = []
+        for seq in top_n_seqs:
+            logprobs = seq.output_logprobs
+            if seq_group.sampling_params.logprobs is None:
+                # NOTE: We need to take care of this case because the sequence
+                # always has the logprobs of the sampled tokens even if the
+                # logprobs are not requested.
+                logprobs = None
+            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,
+                finished_reason,
+                seq.data.hidden_states,
+            )
+            outputs.append(output)
+
+        # Every sequence in the sequence group should have the same prompt.
+        prompt = seq_group.prompt
+        prompt_token_ids = seq_group.prompt_token_ids
+        prompt_logprobs = seq_group.prompt_logprobs
+        finished = seq_group.is_finished()
+        return cls(
+            seq_group.request_id,
+            prompt,
+            prompt_token_ids,
+            prompt_logprobs,
+            outputs,
+            finished,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"RequestOutput(request_id={self.request_id}, "
+            f"prompt={self.prompt!r}, "
+            f"prompt_token_ids={self.prompt_token_ids}, "
+            f"prompt_logprobs={self.prompt_logprobs}, "
+            f"outputs={self.outputs}, "
+            f"finished={self.finished})"
+        )

ChatTTS/model/velocity/sampler.py

@@ -0,0 +1,120 @@
+import torch
+from torch.functional import F
+from typing import List, Callable
+
+from ..embed import Embed
+
+
+class Sampler:
+    def __init__(self, post_model: Embed, num_audio_tokens: int, num_vq: int):
+        self.post_model = post_model
+        self.device = next(self.post_model.parameters()).device
+        self.num_audio_tokens = num_audio_tokens
+        self.num_vq = num_vq
+
+    def sample(
+        self,
+        inputs_ids: torch.Tensor,
+        hidden_states: torch.Tensor,
+        infer_text: bool = False,
+        temperature: torch.Tensor = 1.0,
+        logits_processors: List[Callable] = [
+            lambda logits_token, logits: logits,
+        ],
+        logits_warpers: List[Callable] = [
+            lambda logits_token, logits: logits,
+        ],
+        min_new_token: int = 0,
+        now_length: int = 0,
+        eos_token: int = 0,
+        start_idx: int = 0,
+    ):
+        # print(inputs_ids.shape)
+        B = hidden_states.shape[0]
+
+        end_idx = torch.zeros(
+            inputs_ids.shape[0], device=inputs_ids.device, dtype=torch.long
+        )
+        finish = torch.zeros(inputs_ids.shape[0], device=inputs_ids.device).bool()
+        if not infer_text:
+            temperature = (
+                temperature.unsqueeze(0)
+                .expand(inputs_ids.shape[0], -1)
+                .contiguous()
+                .view(-1, 1)
+            )
+
+        if infer_text:
+            logits: torch.Tensor = self.post_model.head_text(hidden_states)
+        else:
+            # logits = torch.stack([self.head_code[i](hidden_states) for i in range(self.num_vq)], 3)
+            logits = torch.empty(
+                hidden_states.size(0),
+                hidden_states.size(1),
+                self.num_audio_tokens,
+                self.num_vq,
+                dtype=torch.float,
+                device=self.device,
+            )
+            for num_vq_iter in range(self.num_vq):
+                x: torch.Tensor = self.post_model.head_code[num_vq_iter](hidden_states)
+                logits[..., num_vq_iter] = x
+                del x
+
+        del hidden_states
+
+        # logits = logits[:, -1].float()
+        logits = logits.narrow(1, -1, 1).squeeze_(1).float()
+
+        if not infer_text:
+            # logits = rearrange(logits, "b c n -> (b n) c")
+            logits = logits.permute(0, 2, 1)
+            logits = logits.reshape(-1, logits.size(2))
+            # logits_token = rearrange(inputs_ids[:, start_idx:], "b c n -> (b n) c")
+            inputs_ids_sliced = inputs_ids[:, start_idx:].permute(0, 2, 1)
+            logits_token = inputs_ids_sliced.reshape(
+                inputs_ids_sliced.size(0) * inputs_ids_sliced.size(1),
+                -1,
+            ).to(self.device)
+        else:
+            logits_token = inputs_ids[:, start_idx:, 0].to(self.device)
+
+        logits /= temperature
+
+        for logitsProcessors in logits_processors:
+            logits = logitsProcessors(logits_token, logits)
+
+        for logitsWarpers in logits_warpers:
+            logits = logitsWarpers(logits_token, logits)
+
+        del logits_token
+
+        if now_length < min_new_token:
+            logits[:, eos_token] = -torch.inf
+
+        scores = F.softmax(logits, dim=-1)
+        idx_next = torch.multinomial(scores, num_samples=1).to(finish.device)
+        if not infer_text:
+            scores = scores.reshape(B, -1, scores.shape[-1])
+        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)
+            finish_or = idx_next.eq(eos_token).any(1)
+            finish.logical_or_(finish_or)
+            del finish_or
+        else:
+            finish_or = idx_next.eq(eos_token).any(1)
+            finish.logical_or_(finish_or)
+            del finish_or
+
+        del inputs_ids
+
+        not_finished = finish.logical_not().to(end_idx.device)
+
+        end_idx.add_(not_finished.int())
+        idx_next = idx_next[:, None, :]
+        return (
+            idx_next,
+            torch.log(scores),
+            finish,
+        )

ChatTTS/model/velocity/sampling_params.py

@@ -0,0 +1,296 @@
+"""Sampling parameters for text generation."""
+
+from enum import IntEnum
+from functools import cached_property
+from typing import Callable, List, Optional, Union
+
+import torch
+
+_SAMPLING_EPS = 1e-5
+
+
+class SamplingType(IntEnum):
+    GREEDY = 0
+    RANDOM = 1
+    BEAM = 2
+
+
+LogitsProcessor = Callable[[List[int], torch.Tensor], torch.Tensor]
+"""LogitsProcessor is a function that takes a list of previously generated
+tokens and a tensor of the logits for the next token, and returns a modified
+tensor of logits to sample from."""
+
+
+class SamplingParams:
+    """Sampling parameters for text generation.
+
+    Overall, we follow the sampling parameters from the OpenAI text completion
+    API (https://platform.openai.com/docs/api-reference/completions/create).
+    In addition, we support beam search, which is not supported by OpenAI.
+
+    Args:
+        n: Number of output sequences to return for the given prompt.
+        best_of: Number of output sequences that are generated from the prompt.
+            From these `best_of` sequences, the top `n` sequences are returned.
+            `best_of` must be greater than or equal to `n`. This is treated as
+            the beam width when `use_beam_search` is True. By default, `best_of`
+            is set to `n`.
+        presence_penalty: Float that penalizes new tokens based on whether they
+            appear in the generated text so far. Values > 0 encourage the model
+            to use new tokens, while values < 0 encourage the model to repeat
+            tokens.
+        frequency_penalty: Float that penalizes new tokens based on their
+            frequency in the generated text so far. Values > 0 encourage the
+            model to use new tokens, while values < 0 encourage the model to
+            repeat tokens.
+        repetition_penalty: Float that penalizes new tokens based on whether
+            they appear in the prompt and the generated text so far. Values > 1
+            encourage the model to use new tokens, while values < 1 encourage
+            the model to repeat tokens.
+        temperature: Float that controls the randomness of the sampling. Lower
+            values make the model more deterministic, while higher values make
+            the model more random. Zero means greedy sampling.
+        top_p: Float that controls the cumulative probability of the top tokens
+            to consider. Must be in (0, 1]. Set to 1 to consider all tokens.
+        top_k: Integer that controls the number of top tokens to consider. Set
+            to -1 to consider all tokens.
+        min_p: Float that represents the minimum probability for a token to be
+            considered, relative to the probability of the most likely token.
+            Must be in [0, 1]. Set to 0 to disable this.
+        use_beam_search: Whether to use beam search instead of sampling.
+        length_penalty: Float that penalizes sequences based on their length.
+            Used in beam search.
+        early_stopping: Controls the stopping condition for beam search. It
+            accepts the following values: `True`, where the generation stops as
+            soon as there are `best_of` complete candidates; `False`, where an
+            heuristic is applied and the generation stops when is it very
+            unlikely to find better candidates; `"never"`, where the beam search
+            procedure only stops when there cannot be better candidates
+            (canonical beam search algorithm).
+        stop: List of strings that stop the generation when they are generated.
+            The returned output will not contain the stop strings.
+        stop_token_ids: List of tokens that stop the generation when they are
+            generated. The returned output will contain the stop tokens unless
+            the stop tokens are special tokens.
+        include_stop_str_in_output: Whether to include the stop strings in output
+            text. Defaults to False.
+        ignore_eos: Whether to ignore the EOS token and continue generating
+            tokens after the EOS token is generated.
+        max_tokens: Maximum number of tokens to generate per output sequence.
+        logprobs: Number of log probabilities to return per output token.
+            Note that the implementation follows the OpenAI API: The return
+            result includes the log probabilities on the `logprobs` most likely
+            tokens, as well the chosen tokens. The API will always return the
+            log probability of the sampled token, so there  may be up to
+            `logprobs+1` elements in the response.
+        prompt_logprobs: Number of log probabilities to return per prompt token.
+        skip_special_tokens: Whether to skip special tokens in the output.
+        spaces_between_special_tokens: Whether to add spaces between special
+            tokens in the output.  Defaults to True.
+        logits_processors: List of functions that modify logits based on
+            previously generated tokens.
+    """
+
+    def __init__(
+        self,
+        n: int = 1,
+        best_of: Optional[int] = None,
+        presence_penalty: float = 0.0,
+        frequency_penalty: float = 0.0,
+        repetition_penalty: float = 1.0,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = -1,
+        min_p: float = 0.0,
+        use_beam_search: bool = False,
+        length_penalty: float = 1.0,
+        early_stopping: Union[bool, str] = False,
+        stop: Optional[Union[str, List[str]]] = None,
+        stop_token_ids: Optional[List[int]] = None,
+        include_stop_str_in_output: bool = False,
+        ignore_eos: bool = False,
+        max_tokens: int = 16,
+        logprobs: Optional[int] = None,
+        prompt_logprobs: Optional[int] = None,
+        skip_special_tokens: bool = True,
+        spaces_between_special_tokens: bool = True,
+        logits_processors: Optional[List[LogitsProcessor]] = (
+            [
+                lambda logits_token, logits: logits,
+            ],
+            [
+                lambda logits_token, logits: logits,
+            ],
+        ),
+        min_new_token: int = 0,
+        max_new_token: int = 8192,
+        infer_text: bool = False,
+        eos_token: int = 0,
+        spk_emb: str = None,
+        start_idx: int = 0,
+    ) -> None:
+        self.n = n
+        self.best_of = best_of if best_of is not None else n
+        self.presence_penalty = presence_penalty
+        self.frequency_penalty = frequency_penalty
+        self.repetition_penalty = repetition_penalty
+        self.temperature = temperature
+        self.top_p = top_p
+        self.top_k = top_k
+        self.min_p = min_p
+        self.use_beam_search = use_beam_search
+        self.length_penalty = length_penalty
+        self.early_stopping = early_stopping
+        self.min_new_token = min_new_token
+        self.max_new_token = max_new_token
+        self.infer_text = infer_text
+        self.eos_token = eos_token
+        self.spk_emb = spk_emb
+        self.start_idx = start_idx
+        if stop is None:
+            self.stop = []
+        elif isinstance(stop, str):
+            self.stop = [stop]
+        else:
+            self.stop = list(stop)
+        if stop_token_ids is None:
+            self.stop_token_ids = []
+        else:
+            self.stop_token_ids = list(stop_token_ids)
+        self.ignore_eos = ignore_eos
+        self.max_tokens = max_tokens
+        self.logprobs = logprobs
+        self.prompt_logprobs = prompt_logprobs
+        self.skip_special_tokens = skip_special_tokens
+        self.spaces_between_special_tokens = spaces_between_special_tokens
+        self.logits_processors = logits_processors
+        self.include_stop_str_in_output = include_stop_str_in_output
+        self._verify_args()
+        if self.use_beam_search:
+            self._verify_beam_search()
+        else:
+            self._verify_non_beam_search()
+            # if self.temperature < _SAMPLING_EPS:
+            #     # Zero temperature means greedy sampling.
+            #     self.top_p = 1.0
+            #     self.top_k = -1
+            #     self.min_p = 0.0
+            #     self._verify_greedy_sampling()
+
+    def _verify_args(self) -> None:
+        if self.n < 1:
+            raise ValueError(f"n must be at least 1, got {self.n}.")
+        if self.best_of < self.n:
+            raise ValueError(
+                f"best_of must be greater than or equal to n, "
+                f"got n={self.n} and best_of={self.best_of}."
+            )
+        if not -2.0 <= self.presence_penalty <= 2.0:
+            raise ValueError(
+                "presence_penalty must be in [-2, 2], got " f"{self.presence_penalty}."
+            )
+        if not -2.0 <= self.frequency_penalty <= 2.0:
+            raise ValueError(
+                "frequency_penalty must be in [-2, 2], got "
+                f"{self.frequency_penalty}."
+            )
+        if not 0.0 < self.repetition_penalty <= 2.0:
+            raise ValueError(
+                "repetition_penalty must be in (0, 2], got "
+                f"{self.repetition_penalty}."
+            )
+        # if self.temperature < 0.0:
+        #     raise ValueError(
+        #         f"temperature must be non-negative, got {self.temperature}.")
+        if not 0.0 < self.top_p <= 1.0:
+            raise ValueError(f"top_p must be in (0, 1], got {self.top_p}.")
+        if self.top_k < -1 or self.top_k == 0:
+            raise ValueError(
+                f"top_k must be -1 (disable), or at least 1, " f"got {self.top_k}."
+            )
+        if not 0.0 <= self.min_p <= 1.0:
+            raise ValueError("min_p must be in [0, 1], got " f"{self.min_p}.")
+        if self.max_tokens < 1:
+            raise ValueError(f"max_tokens must be at least 1, got {self.max_tokens}.")
+        if self.logprobs is not None and self.logprobs < 0:
+            raise ValueError(f"logprobs must be non-negative, got {self.logprobs}.")
+        if self.prompt_logprobs is not None and self.prompt_logprobs < 0:
+            raise ValueError(
+                f"prompt_logprobs must be non-negative, got " f"{self.prompt_logprobs}."
+            )
+
+    def _verify_beam_search(self) -> None:
+        if self.best_of == 1:
+            raise ValueError(
+                "best_of must be greater than 1 when using beam "
+                f"search. Got {self.best_of}."
+            )
+        if self.temperature > _SAMPLING_EPS:
+            raise ValueError("temperature must be 0 when using beam search.")
+        if self.top_p < 1.0 - _SAMPLING_EPS:
+            raise ValueError("top_p must be 1 when using beam search.")
+        if self.top_k != -1:
+            raise ValueError("top_k must be -1 when using beam search.")
+        if self.early_stopping not in [True, False, "never"]:
+            raise ValueError(
+                f"early_stopping must be True, False, or 'never', "
+                f"got {self.early_stopping}."
+            )
+
+    def _verify_non_beam_search(self) -> None:
+        if self.early_stopping is not False:
+            raise ValueError(
+                "early_stopping is not effective and must be "
+                "False when not using beam search."
+            )
+        if (
+            self.length_penalty < 1.0 - _SAMPLING_EPS
+            or self.length_penalty > 1.0 + _SAMPLING_EPS
+        ):
+            raise ValueError(
+                "length_penalty is not effective and must be the "
+                "default value of 1.0 when not using beam search."
+            )
+
+    def _verify_greedy_sampling(self) -> None:
+        if self.best_of > 1:
+            raise ValueError(
+                "best_of must be 1 when using greedy sampling." f"Got {self.best_of}."
+            )
+
+    @cached_property
+    def sampling_type(self) -> SamplingType:
+        if self.use_beam_search:
+            return SamplingType.BEAM
+        # if self.temperature < _SAMPLING_EPS:
+        #     return SamplingType.GREEDY
+        return SamplingType.RANDOM
+
+    def __repr__(self) -> str:
+        return (
+            f"SamplingParams(n={self.n}, "
+            f"best_of={self.best_of}, "
+            f"presence_penalty={self.presence_penalty}, "
+            f"frequency_penalty={self.frequency_penalty}, "
+            f"repetition_penalty={self.repetition_penalty}, "
+            f"temperature={self.temperature}, "
+            f"top_p={self.top_p}, "
+            f"top_k={self.top_k}, "
+            f"min_p={self.min_p}, "
+            f"use_beam_search={self.use_beam_search}, "
+            f"length_penalty={self.length_penalty}, "
+            f"early_stopping={self.early_stopping}, "
+            f"stop={self.stop}, "
+            f"stop_token_ids={self.stop_token_ids}, "
+            f"include_stop_str_in_output={self.include_stop_str_in_output}, "
+            f"ignore_eos={self.ignore_eos}, "
+            f"max_tokens={self.max_tokens}, "
+            f"logprobs={self.logprobs}, "
+            f"prompt_logprobs={self.prompt_logprobs}, "
+            f"skip_special_tokens={self.skip_special_tokens}, "
+            "spaces_between_special_tokens="
+            f"{self.spaces_between_special_tokens}), "
+            f"max_new_token={self.max_new_token}), "
+            f"min_new_token={self.min_new_token}), "
+            f"infer_text={self.infer_text})"
+        )

ChatTTS/model/velocity/scheduler.py

@@ -0,0 +1,426 @@
+import enum
+import time
+from typing import Dict, Iterable, List, Optional, Tuple, Union
+
+from vllm.config import CacheConfig, SchedulerConfig
+from .block_manager import AllocStatus, BlockSpaceManager
+from vllm.core.policy import PolicyFactory
+from vllm.logger import init_logger
+from .sequence import (
+    Sequence,
+    SequenceData,
+    SequenceGroup,
+    SequenceGroupMetadata,
+    SequenceStatus,
+)
+
+logger = init_logger(__name__)
+
+
+class PreemptionMode(enum.Enum):
+    """Preemption modes.
+
+    1. Swapping: Swap out the blocks of the preempted sequences to CPU memory
+    and swap them back in when the sequences are resumed.
+    2. Recomputation: Discard the blocks of the preempted sequences and
+    recompute them when the sequences are resumed, treating the sequences as
+    new prompts.
+    """
+
+    SWAP = enum.auto()
+    RECOMPUTE = enum.auto()
+
+
+class SchedulerOutputs:
+
+    def __init__(
+        self,
+        scheduled_seq_groups: List[SequenceGroup],
+        prompt_run: bool,
+        num_batched_tokens: int,
+        blocks_to_swap_in: Dict[int, int],
+        blocks_to_swap_out: Dict[int, int],
+        blocks_to_copy: Dict[int, List[int]],
+        ignored_seq_groups: List[SequenceGroup],
+    ) -> None:
+        self.scheduled_seq_groups = scheduled_seq_groups
+        self.prompt_run = prompt_run
+        self.num_batched_tokens = num_batched_tokens
+        self.blocks_to_swap_in = blocks_to_swap_in
+        self.blocks_to_swap_out = blocks_to_swap_out
+        self.blocks_to_copy = blocks_to_copy
+        # Swap in and swap out should never happen at the same time.
+        assert not (blocks_to_swap_in and blocks_to_swap_out)
+        self.ignored_seq_groups = ignored_seq_groups
+
+    def is_empty(self) -> bool:
+        # NOTE: We do not consider the ignored sequence groups.
+        return (
+            not self.scheduled_seq_groups
+            and not self.blocks_to_swap_in
+            and not self.blocks_to_swap_out
+            and not self.blocks_to_copy
+        )
+
+
+class Scheduler:
+
+    def __init__(
+        self,
+        scheduler_config: SchedulerConfig,
+        cache_config: CacheConfig,
+    ) -> None:
+        self.scheduler_config = scheduler_config
+        self.cache_config = cache_config
+
+        self.prompt_limit = min(
+            self.scheduler_config.max_model_len,
+            self.scheduler_config.max_num_batched_tokens,
+        )
+
+        # Instantiate the scheduling policy.
+        self.policy = PolicyFactory.get_policy(policy_name="fcfs")
+        # Create the block space manager.
+        self.block_manager = BlockSpaceManager(
+            block_size=self.cache_config.block_size,
+            num_gpu_blocks=self.cache_config.num_gpu_blocks,
+            num_cpu_blocks=self.cache_config.num_cpu_blocks,
+            sliding_window=self.cache_config.sliding_window,
+        )
+
+        # TODO(zhuohan): Use deque instead of list for better performance.
+        # Sequence groups in the WAITING state.
+        self.waiting: List[SequenceGroup] = []
+        # Sequence groups in the RUNNING state.
+        self.running: List[SequenceGroup] = []
+        # Sequence groups in the SWAPPED state.
+        self.swapped: List[SequenceGroup] = []
+
+    def add_seq_group(self, seq_group: SequenceGroup) -> None:
+        # Add sequence groups to the waiting queue.
+        self.waiting.append(seq_group)
+
+    def abort_seq_group(self, request_id: Union[str, Iterable[str]]) -> None:
+        if isinstance(request_id, str):
+            request_id = (request_id,)
+        request_ids = set(request_id)
+        for state_queue in [self.waiting, self.running, self.swapped]:
+            # We need to reverse the list as we are removing elements
+            # from it as we iterate over it. If we don't do it,
+            # indices will get messed up and we will skip over elements.
+            for seq_group in reversed(state_queue):
+                if seq_group.request_id in request_ids:
+                    # Remove the sequence group from the state queue.
+                    state_queue.remove(seq_group)
+                    for seq in seq_group.get_seqs():
+                        if seq.is_finished():
+                            continue
+                        seq.status = SequenceStatus.FINISHED_ABORTED
+                        self.free_seq(seq)
+                    request_ids.remove(seq_group.request_id)
+                    if not request_ids:
+                        return
+
+    def has_unfinished_seqs(self) -> bool:
+        return self.waiting or self.running or self.swapped
+
+    def get_num_unfinished_seq_groups(self) -> int:
+        return len(self.waiting) + len(self.running) + len(self.swapped)
+
+    def _schedule(self) -> SchedulerOutputs:
+        # 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]] = {}
+
+        # Fix the current time.
+        now = time.monotonic()
+
+        # Join waiting sequences if possible.
+        if not self.swapped:
+            ignored_seq_groups: List[SequenceGroup] = []
+            scheduled: List[SequenceGroup] = []
+            # The total number of sequences on the fly, including the
+            # requests in the generation phase.
+            num_curr_seqs = sum(
+                seq_group.get_max_num_running_seqs() for seq_group in self.running
+            )
+            seq_lens: List[int] = []
+
+            # Optimization: We do not sort the waiting queue since the preempted
+            # sequence groups are added to the front and the new sequence groups
+            # are added to the back.
+            while self.waiting:
+                seq_group = self.waiting[0]
+
+                waiting_seqs = seq_group.get_seqs(status=SequenceStatus.WAITING)
+                assert len(waiting_seqs) == 1, (
+                    "Waiting sequence group should have only one prompt " "sequence."
+                )
+                num_prompt_tokens = waiting_seqs[0].get_len()
+                if num_prompt_tokens > self.prompt_limit:
+                    logger.warning(
+                        f"Input prompt ({num_prompt_tokens} tokens) is too long"
+                        f" and exceeds limit of {self.prompt_limit}"
+                    )
+                    for seq in waiting_seqs:
+                        seq.status = SequenceStatus.FINISHED_IGNORED
+                    ignored_seq_groups.append(seq_group)
+                    self.waiting.pop(0)
+                    continue
+
+                # If the sequence group cannot be allocated, stop.
+                can_allocate = self.block_manager.can_allocate(seq_group)
+                if can_allocate == AllocStatus.LATER:
+                    break
+                elif can_allocate == AllocStatus.NEVER:
+                    logger.warning(
+                        f"Input prompt ({num_prompt_tokens} tokens) is too long"
+                        f" and exceeds the capacity of block_manager"
+                    )
+                    for seq in waiting_seqs:
+                        seq.status = SequenceStatus.FINISHED_IGNORED
+                    ignored_seq_groups.append(seq_group)
+                    self.waiting.pop(0)
+                    continue
+
+                # If the number of batched tokens exceeds the limit, stop.
+                new_seq_lens = seq_lens + [num_prompt_tokens]
+                num_batched_tokens = len(new_seq_lens) * max(new_seq_lens)
+                if num_batched_tokens > self.scheduler_config.max_num_batched_tokens:
+                    break
+
+                # The total number of sequences in the RUNNING state should not
+                # exceed the maximum number of sequences.
+                num_new_seqs = seq_group.get_max_num_running_seqs()
+                if num_curr_seqs + num_new_seqs > self.scheduler_config.max_num_seqs:
+                    break
+
+                num_paddings = num_batched_tokens - sum(new_seq_lens)
+                if num_paddings > self.scheduler_config.max_paddings:
+                    break
+                seq_lens = new_seq_lens
+
+                seq_group = self.waiting.pop(0)
+                self._allocate(seq_group)
+                self.running.append(seq_group)
+                num_curr_seqs += num_new_seqs
+                scheduled.append(seq_group)
+
+            if scheduled or ignored_seq_groups:
+                scheduler_outputs = SchedulerOutputs(
+                    scheduled_seq_groups=scheduled,
+                    prompt_run=True,
+                    num_batched_tokens=len(seq_lens) * max(seq_lens) if seq_lens else 0,
+                    blocks_to_swap_in=blocks_to_swap_in,
+                    blocks_to_swap_out=blocks_to_swap_out,
+                    blocks_to_copy=blocks_to_copy,
+                    ignored_seq_groups=ignored_seq_groups,
+                )
+                return scheduler_outputs
+
+        # NOTE(woosuk): Preemption happens only when there is no available slot
+        # to keep all the sequence groups in the RUNNING state.
+        # In this case, the policy is responsible for deciding which sequence
+        # groups to preempt.
+        self.running = self.policy.sort_by_priority(now, self.running)
+
+        # Reserve new token slots for the running sequence groups.
+        running: List[SequenceGroup] = []
+        preempted: List[SequenceGroup] = []
+        while self.running:
+            seq_group = self.running.pop(0)
+            while not self.block_manager.can_append_slot(seq_group):
+                if self.running:
+                    # Preempt the lowest-priority sequence groups.
+                    victim_seq_group = self.running.pop(-1)
+                    self._preempt(victim_seq_group, blocks_to_swap_out)
+                    preempted.append(victim_seq_group)
+                else:
+                    # No other sequence groups can be preempted.
+                    # Preempt the current sequence group.
+                    self._preempt(seq_group, blocks_to_swap_out)
+                    preempted.append(seq_group)
+                    break
+            else:
+                # Append new slots to the sequence group.
+                self._append_slot(seq_group, blocks_to_copy)
+                running.append(seq_group)
+        self.running = running
+
+        # Swap in the sequence groups in the SWAPPED state if possible.
+        self.swapped = self.policy.sort_by_priority(now, self.swapped)
+        if not preempted:
+            num_curr_seqs = sum(
+                seq_group.get_max_num_running_seqs() for seq_group in self.running
+            )
+
+            while self.swapped:
+                seq_group = self.swapped[0]
+                # If the sequence group cannot be swapped in, stop.
+                if not self.block_manager.can_swap_in(seq_group):
+                    break
+
+                # The total number of sequences in the RUNNING state should not
+                # exceed the maximum number of sequences.
+                num_new_seqs = seq_group.get_max_num_running_seqs()
+                if num_curr_seqs + num_new_seqs > self.scheduler_config.max_num_seqs:
+                    break
+
+                seq_group = self.swapped.pop(0)
+                self._swap_in(seq_group, blocks_to_swap_in)
+                self._append_slot(seq_group, blocks_to_copy)
+                num_curr_seqs += num_new_seqs
+                self.running.append(seq_group)
+
+        # Each sequence in the generation phase only takes one token slot.
+        # Therefore, the number of batched tokens is equal to the number of
+        # sequences in the RUNNING state.
+        num_batched_tokens = sum(
+            seq_group.num_seqs(status=SequenceStatus.RUNNING)
+            for seq_group in self.running
+        )
+
+        scheduler_outputs = SchedulerOutputs(
+            scheduled_seq_groups=self.running,
+            prompt_run=False,
+            num_batched_tokens=num_batched_tokens,
+            blocks_to_swap_in=blocks_to_swap_in,
+            blocks_to_swap_out=blocks_to_swap_out,
+            blocks_to_copy=blocks_to_copy,
+            ignored_seq_groups=[],
+        )
+        return scheduler_outputs
+
+    def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]:
+        # Schedule sequence groups.
+        # This function call changes the internal states of the scheduler
+        # such as self.running, self.swapped, and self.waiting.
+        scheduler_outputs = self._schedule()
+
+        # Create input data structures.
+        seq_group_metadata_list: List[SequenceGroupMetadata] = []
+        for seq_group in scheduler_outputs.scheduled_seq_groups:
+            seq_data: Dict[int, SequenceData] = {}
+            block_tables: Dict[int, List[int]] = {}
+            for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
+                seq_id = seq.seq_id
+                seq_data[seq_id] = seq.data
+                block_tables[seq_id] = self.block_manager.get_block_table(seq)
+
+            seq_group_metadata = SequenceGroupMetadata(
+                request_id=seq_group.request_id,
+                is_prompt=scheduler_outputs.prompt_run,
+                seq_data=seq_data,
+                sampling_params=seq_group.sampling_params,
+                block_tables=block_tables,
+            )
+            seq_group_metadata_list.append(seq_group_metadata)
+        return seq_group_metadata_list, scheduler_outputs
+
+    def fork_seq(self, parent_seq: Sequence, child_seq: Sequence) -> None:
+        self.block_manager.fork(parent_seq, child_seq)
+
+    def free_seq(self, seq: Sequence) -> None:
+        self.block_manager.free(seq)
+
+    def free_finished_seq_groups(self) -> None:
+        self.running = [
+            seq_group for seq_group in self.running if not seq_group.is_finished()
+        ]
+
+    def _allocate(self, seq_group: SequenceGroup) -> None:
+        self.block_manager.allocate(seq_group)
+        for seq in seq_group.get_seqs(status=SequenceStatus.WAITING):
+            seq.status = SequenceStatus.RUNNING
+
+    def _append_slot(
+        self,
+        seq_group: SequenceGroup,
+        blocks_to_copy: Dict[int, List[int]],
+    ) -> None:
+        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
+            ret = self.block_manager.append_slot(seq)
+            if ret is not None:
+                src_block, dst_block = ret
+                if src_block in blocks_to_copy:
+                    blocks_to_copy[src_block].append(dst_block)
+                else:
+                    blocks_to_copy[src_block] = [dst_block]
+
+    def _preempt(
+        self,
+        seq_group: SequenceGroup,
+        blocks_to_swap_out: Dict[int, int],
+        preemption_mode: Optional[PreemptionMode] = None,
+    ) -> None:
+        # If preemption mode is not specified, we determine the mode as follows:
+        # We use recomputation by default since it incurs lower overhead than
+        # swapping. However, when the sequence group has multiple sequences
+        # (e.g., beam search), recomputation is not currently supported. In
+        # such a case, we use swapping instead.
+        # FIXME(woosuk): This makes our scheduling policy a bit bizarre.
+        # As swapped sequences are prioritized over waiting sequences,
+        # sequence groups with multiple sequences are implicitly prioritized
+        # over sequence groups with a single sequence.
+        # TODO(woosuk): Support recomputation for sequence groups with multiple
+        # sequences. This may require a more sophisticated CUDA kernel.
+        if preemption_mode is None:
+            if seq_group.get_max_num_running_seqs() == 1:
+                preemption_mode = PreemptionMode.RECOMPUTE
+            else:
+                preemption_mode = PreemptionMode.SWAP
+        if preemption_mode == PreemptionMode.RECOMPUTE:
+            self._preempt_by_recompute(seq_group)
+        elif preemption_mode == PreemptionMode.SWAP:
+            self._preempt_by_swap(seq_group, blocks_to_swap_out)
+        else:
+            raise AssertionError("Invalid preemption mode.")
+
+    def _preempt_by_recompute(
+        self,
+        seq_group: SequenceGroup,
+    ) -> None:
+        seqs = seq_group.get_seqs(status=SequenceStatus.RUNNING)
+        assert len(seqs) == 1
+        for seq in seqs:
+            seq.status = SequenceStatus.WAITING
+            self.block_manager.free(seq)
+        # NOTE: For FCFS, we insert the preempted sequence group to the front
+        # of the waiting queue.
+        self.waiting.insert(0, seq_group)
+
+    def _preempt_by_swap(
+        self,
+        seq_group: SequenceGroup,
+        blocks_to_swap_out: Dict[int, int],
+    ) -> None:
+        self._swap_out(seq_group, blocks_to_swap_out)
+        self.swapped.append(seq_group)
+
+    def _swap_in(
+        self,
+        seq_group: SequenceGroup,
+        blocks_to_swap_in: Dict[int, int],
+    ) -> None:
+        mapping = self.block_manager.swap_in(seq_group)
+        blocks_to_swap_in.update(mapping)
+        for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED):
+            seq.status = SequenceStatus.RUNNING
+
+    def _swap_out(
+        self,
+        seq_group: SequenceGroup,
+        blocks_to_swap_out: Dict[int, int],
+    ) -> None:
+        if not self.block_manager.can_swap_out(seq_group):
+            # FIXME(woosuk): Abort the sequence group instead of aborting the
+            # entire engine.
+            raise RuntimeError(
+                "Aborted due to the lack of CPU swap space. Please increase "
+                "the swap space to avoid this error."
+            )
+        mapping = self.block_manager.swap_out(seq_group)
+        blocks_to_swap_out.update(mapping)
+        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
+            seq.status = SequenceStatus.SWAPPED

ChatTTS/model/velocity/sequence.py

@@ -0,0 +1,450 @@
+"""Sequence and its related classes."""
+
+import copy
+import enum
+from typing import Dict, List, Optional, Union
+import torch
+from vllm.block import LogicalTokenBlock
+from .sampling_params import SamplingParams
+
+PromptLogprobs = List[Optional[Dict[int, float]]]
+SampleLogprobs = List[Dict[int, float]]
+
+
+class SequenceStatus(enum.Enum):
+    """Status of a sequence."""
+
+    WAITING = enum.auto()
+    RUNNING = enum.auto()
+    SWAPPED = enum.auto()
+    FINISHED_STOPPED = enum.auto()
+    FINISHED_LENGTH_CAPPED = enum.auto()
+    FINISHED_ABORTED = enum.auto()
+    FINISHED_IGNORED = enum.auto()
+
+    @staticmethod
+    def is_finished(status: "SequenceStatus") -> bool:
+        return status in [
+            SequenceStatus.FINISHED_STOPPED,
+            SequenceStatus.FINISHED_LENGTH_CAPPED,
+            SequenceStatus.FINISHED_ABORTED,
+            SequenceStatus.FINISHED_IGNORED,
+        ]
+
+    @staticmethod
+    def get_finished_reason(status: "SequenceStatus") -> Union[str, None]:
+        if status == SequenceStatus.FINISHED_STOPPED:
+            finish_reason = "stop"
+        elif status == SequenceStatus.FINISHED_LENGTH_CAPPED:
+            finish_reason = "length"
+        elif status == SequenceStatus.FINISHED_ABORTED:
+            finish_reason = "abort"
+        elif status == SequenceStatus.FINISHED_IGNORED:
+            # The ignored sequences are the sequences whose prompt lengths
+            # are longer than the model's length cap. Therefore, the stop
+            # reason should also be "length" as in OpenAI API.
+            finish_reason = "length"
+        else:
+            finish_reason = None
+        return finish_reason
+
+
+class SequenceData:
+    """Data associated with a sequence.
+
+
+    Args:
+        prompt_token_ids: The token IDs of the prompt.
+
+    Attributes:
+        prompt_token_ids: The token IDs of the prompt.
+        output_token_ids: The token IDs of the output.
+        cumulative_logprob: The cumulative log probability of the output.
+    """
+
+    def __init__(
+        self,
+        prompt_token_ids: List[int],
+    ) -> None:
+        self.prompt_token_ids = prompt_token_ids
+        self.output_token_ids: List[int] = []
+        self.cumulative_logprob = 0.0
+        self.hidden_states: Optional[torch.Tensor] = None
+        self.finished = False
+
+    def append_token_id(self, token_id: int, logprob: float) -> None:
+        if isinstance(self.cumulative_logprob, float):
+            self.cumulative_logprob = [
+                0.0,
+            ] * len(logprob)
+        self.output_token_ids.append(token_id)
+        for i in range(len(self.cumulative_logprob)):
+            self.cumulative_logprob[i] += logprob[i]
+
+    def append_hidden_states(self, hidden_states: torch.Tensor) -> None:
+        if self.hidden_states is None:
+            self.hidden_states = hidden_states
+        else:
+            self.hidden_states = torch.cat([self.hidden_states, hidden_states], dim=0)
+
+    def get_len(self) -> int:
+        return len(self.output_token_ids) + len(self.prompt_token_ids)
+
+    def get_prompt_len(self) -> int:
+        return len(self.prompt_token_ids)
+
+    def get_output_len(self) -> int:
+        return len(self.output_token_ids)
+
+    def get_token_ids(self) -> List[int]:
+        return self.prompt_token_ids + self.output_token_ids
+
+    def get_last_token_id(self) -> int:
+        if not self.output_token_ids:
+            return self.prompt_token_ids[-1]
+        return self.output_token_ids[-1]
+
+    def __repr__(self) -> str:
+        return (
+            f"SequenceData("
+            f"prompt_token_ids={self.prompt_token_ids}, "
+            f"output_token_ids={self.output_token_ids}, "
+            f"cumulative_logprob={self.cumulative_logprob}), "
+            f"hidden_states={self.hidden_states.shape if self.hidden_states is not None else None}, "
+            f"finished={self.finished})"
+        )
+
+
+class Sequence:
+    """Stores the data, status, and block information of a sequence.
+
+    Args:
+        seq_id: The ID of the sequence.
+        prompt: The prompt of the sequence.
+        prompt_token_ids: The token IDs of the prompt.
+        block_size: The block size of the sequence. Should be the same as the
+            block size used by the block manager and cache engine.
+    """
+
+    def __init__(
+        self,
+        seq_id: int,
+        prompt: str,
+        prompt_token_ids: List[int],
+        block_size: int,
+    ) -> None:
+        self.seq_id = seq_id
+        self.prompt = prompt
+        self.block_size = block_size
+
+        self.data = SequenceData(prompt_token_ids)
+        self.output_logprobs: SampleLogprobs = []
+        self.output_text = ""
+
+        self.logical_token_blocks: List[LogicalTokenBlock] = []
+        # Initialize the logical token blocks with the prompt token ids.
+        self._append_tokens_to_blocks(prompt_token_ids)
+        self.status = SequenceStatus.WAITING
+
+        # Used for incremental detokenization
+        self.prefix_offset = 0
+        self.read_offset = 0
+        # Input + output tokens
+        self.tokens: Optional[List[str]] = None
+
+    def _append_logical_block(self) -> None:
+        block = LogicalTokenBlock(
+            block_number=len(self.logical_token_blocks),
+            block_size=self.block_size,
+        )
+        self.logical_token_blocks.append(block)
+
+    def _append_tokens_to_blocks(self, token_ids: List[int]) -> None:
+        cursor = 0
+        while cursor < len(token_ids):
+            if not self.logical_token_blocks:
+                self._append_logical_block()
+
+            last_block = self.logical_token_blocks[-1]
+            if last_block.is_full():
+                self._append_logical_block()
+                last_block = self.logical_token_blocks[-1]
+
+            num_empty_slots = last_block.get_num_empty_slots()
+            last_block.append_tokens(token_ids[cursor : cursor + num_empty_slots])
+            cursor += num_empty_slots
+
+    def append_token_id(
+        self,
+        token_id: int,
+        logprobs: Dict[int, float],
+        hidden_states: Optional[torch.Tensor] = None,
+        finished: bool = False,
+    ) -> None:
+        assert token_id in logprobs
+        self._append_tokens_to_blocks([token_id])
+        self.output_logprobs.append(logprobs)
+        self.data.append_token_id(token_id, logprobs[token_id])
+        self.data.append_hidden_states(hidden_states)
+        self.data.finished = finished
+
+    def get_len(self) -> int:
+        return self.data.get_len()
+
+    def get_prompt_len(self) -> int:
+        return self.data.get_prompt_len()
+
+    def get_output_len(self) -> int:
+        return self.data.get_output_len()
+
+    def get_token_ids(self) -> List[int]:
+        return self.data.get_token_ids()
+
+    def get_last_token_id(self) -> int:
+        return self.data.get_last_token_id()
+
+    def get_output_token_ids(self) -> List[int]:
+        return self.data.output_token_ids
+
+    def get_cumulative_logprob(self) -> float:
+        return self.data.cumulative_logprob
+
+    def get_beam_search_score(
+        self,
+        length_penalty: float = 0.0,
+        seq_len: Optional[int] = None,
+        eos_token_id: Optional[int] = None,
+    ) -> float:
+        """Calculate the beam search score with length penalty.
+
+        Adapted from
+
+        https://github.com/huggingface/transformers/blob/ccb92be23def445f2afdea94c31286f84b89eb5b/src/transformers/generation/beam_search.py#L938
+        """
+        if seq_len is None:
+            seq_len = self.get_len()
+            # NOTE: HF implementation does not count the EOS token
+            # towards the length, we align with that here for testing.
+            if eos_token_id is not None and self.get_last_token_id() == eos_token_id:
+                seq_len -= 1
+        return self.get_cumulative_logprob() / (seq_len**length_penalty)
+
+    def is_finished(self) -> bool:
+        return SequenceStatus.is_finished(self.status)
+
+    def fork(self, new_seq_id: int) -> "Sequence":
+        new_seq = copy.deepcopy(self)
+        new_seq.seq_id = new_seq_id
+        return new_seq
+
+    def __repr__(self) -> str:
+        return (
+            f"Sequence(seq_id={self.seq_id}, "
+            f"status={self.status.name}, "
+            f"num_blocks={len(self.logical_token_blocks)})"
+        )
+
+
+class SequenceGroup:
+    """A group of sequences that are generated from the same prompt.
+
+    Args:
+        request_id: The ID of the request.
+        seqs: The list of sequences.
+        sampling_params: The sampling parameters used to generate the outputs.
+        arrival_time: The arrival time of the request.
+    """
+
+    def __init__(
+        self,
+        request_id: str,
+        seqs: List[Sequence],
+        sampling_params: SamplingParams,
+        arrival_time: float,
+    ) -> None:
+        self.request_id = request_id
+        self.seqs_dict = {seq.seq_id: seq for seq in seqs}
+        self.sampling_params = sampling_params
+        self.arrival_time = arrival_time
+        self.prompt_logprobs: Optional[PromptLogprobs] = None
+
+    @property
+    def prompt(self) -> str:
+        # All sequences in the group should have the same prompt.
+        # We use the prompt of an arbitrary sequence.
+        return next(iter(self.seqs_dict.values())).prompt
+
+    @property
+    def prompt_token_ids(self) -> List[int]:
+        # All sequences in the group should have the same prompt.
+        # We use the prompt of an arbitrary sequence.
+        return next(iter(self.seqs_dict.values())).data.prompt_token_ids
+
+    def get_max_num_running_seqs(self) -> int:
+        """The maximum number of sequences running in parallel in the remaining
+        lifetime of the request."""
+        if self.sampling_params.use_beam_search:
+            # For beam search, maximally there will always be `best_of` beam
+            # candidates running in the future.
+            return self.sampling_params.best_of
+        else:
+            if self.sampling_params.best_of > self.num_seqs():
+                # At prompt stage, the sequence group is not yet filled up
+                # and only have one sequence running. However, in the
+                # generation stage, we will have `best_of` sequences running.
+                return self.sampling_params.best_of
+            # At sampling stages, return the number of actual sequences
+            # that are not finished yet.
+            return self.num_unfinished_seqs()
+
+    def get_seqs(
+        self,
+        status: Optional[SequenceStatus] = None,
+    ) -> List[Sequence]:
+        if status is None:
+            return list(self.seqs_dict.values())
+        else:
+            return [seq for seq in self.seqs_dict.values() if seq.status == status]
+
+    def get_unfinished_seqs(self) -> List[Sequence]:
+        return [seq for seq in self.seqs_dict.values() if not seq.is_finished()]
+
+    def get_finished_seqs(self) -> List[Sequence]:
+        return [seq for seq in self.seqs_dict.values() if seq.is_finished()]
+
+    def num_seqs(self, status: Optional[SequenceStatus] = None) -> int:
+        return len(self.get_seqs(status))
+
+    def num_unfinished_seqs(self) -> int:
+        return len(self.get_unfinished_seqs())
+
+    def num_finished_seqs(self) -> int:
+        return len(self.get_finished_seqs())
+
+    def find(self, seq_id: int) -> Sequence:
+        if seq_id not in self.seqs_dict:
+            raise ValueError(f"Sequence {seq_id} not found.")
+        return self.seqs_dict[seq_id]
+
+    def add(self, seq: Sequence) -> None:
+        if seq.seq_id in self.seqs_dict:
+            raise ValueError(f"Sequence {seq.seq_id} already exists.")
+        self.seqs_dict[seq.seq_id] = seq
+
+    def remove(self, seq_id: int) -> None:
+        if seq_id not in self.seqs_dict:
+            raise ValueError(f"Sequence {seq_id} not found.")
+        del self.seqs_dict[seq_id]
+
+    def is_finished(self) -> bool:
+        return all(seq.is_finished() for seq in self.get_seqs())
+
+    def __repr__(self) -> str:
+        return (
+            f"SequenceGroup(request_id={self.request_id}, "
+            f"sampling_params={self.sampling_params}, "
+            f"num_seqs={len(self.seqs_dict)})"
+        )
+
+
+class SequenceGroupMetadata:
+    """Metadata for a sequence group. Used to create `InputMetadata`.
+
+
+    Args:
+        request_id: The ID of the request.
+        is_prompt: Whether the request is at prompt stage.
+        seq_data: The sequence data. (Seq id -> sequence data)
+        sampling_params: The sampling parameters used to generate the outputs.
+        block_tables: The block tables. (Seq id -> list of physical block
+            numbers)
+    """
+
+    def __init__(
+        self,
+        request_id: str,
+        is_prompt: bool,
+        seq_data: Dict[int, SequenceData],
+        sampling_params: SamplingParams,
+        block_tables: Dict[int, List[int]],
+    ) -> None:
+        self.request_id = request_id
+        self.is_prompt = is_prompt
+        self.seq_data = seq_data
+        self.sampling_params = sampling_params
+        self.block_tables = block_tables
+
+
+class SequenceOutput:
+    """The model output associated with a sequence.
+
+    Args:
+        parent_seq_id: The ID of the parent sequence (for forking in beam
+            search).
+        output_token: The output token ID.
+        logprobs: The logprobs of the output token.
+            (Token id -> logP(x_i+1 | x_0, ..., x_i))
+    """
+
+    def __init__(
+        self,
+        parent_seq_id: int,
+        output_token: int,
+        logprobs: Dict[int, float],
+        hidden_states: Optional[torch.Tensor] = None,
+        finished: bool = False,
+    ) -> None:
+        self.parent_seq_id = parent_seq_id
+        self.output_token = output_token
+        self.logprobs = logprobs
+        self.finished = finished
+        self.hidden_states = hidden_states
+
+    def __repr__(self) -> str:
+        return (
+            f"SequenceOutput(parent_seq_id={self.parent_seq_id}, "
+            f"output_token={self.output_token}, "
+            f"logprobs={self.logprobs}),"
+            f"finished={self.finished}),"
+            f"hidden_states={self.hidden_states.shape if self.hidden_states is not None else None}"
+        )
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, SequenceOutput):
+            raise NotImplementedError()
+        return (
+            self.parent_seq_id == other.parent_seq_id
+            and self.output_token == other.output_token
+            and self.logprobs == other.logprobs
+        )
+
+
+class SequenceGroupOutput:
+    """The model output associated with a sequence group."""
+
+    def __init__(
+        self,
+        samples: List[SequenceOutput],
+        prompt_logprobs: Optional[PromptLogprobs],
+    ) -> None:
+        self.samples = samples
+        self.prompt_logprobs = prompt_logprobs
+
+    def __repr__(self) -> str:
+        return (
+            f"SequenceGroupOutput(samples={self.samples}, "
+            f"prompt_logprobs={self.prompt_logprobs})"
+        )
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, SequenceGroupOutput):
+            raise NotImplementedError()
+        return (
+            self.samples == other.samples
+            and self.prompt_logprobs == other.prompt_logprobs
+        )
+
+
+# For each sequence group, we generate a list of SequenceOutput object,
+# each of which contains one possible candidate for the next token.
+SamplerOutput = List[SequenceGroupOutput]

ChatTTS/model/velocity/worker.py

@@ -0,0 +1,251 @@
+"""A GPU worker class."""
+
+import os
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.distributed
+
+from vllm.config import CacheConfig, ModelConfig, ParallelConfig, SchedulerConfig
+from vllm.model_executor import set_random_seed
+from vllm.model_executor.parallel_utils.communication_op import broadcast_object_list
+from vllm.model_executor.parallel_utils.parallel_state import initialize_model_parallel
+from vllm.sequence import SamplerOutput, SequenceGroupMetadata
+from vllm.worker.cache_engine import CacheEngine
+
+from .model_runner import ModelRunner
+
+
+class Worker:
+    """A worker class that executes (a partition of) the model on a GPU.
+
+    Each worker is associated with a single GPU. The worker is responsible for
+    maintaining the KV cache and executing the model on the GPU. In case of
+    distributed inference, each worker is assigned a partition of the model.
+    """
+
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        parallel_config: ParallelConfig,
+        scheduler_config: SchedulerConfig,
+        local_rank: int,
+        rank: int,
+        distributed_init_method: str,
+        post_model_path: str,
+        is_driver_worker: bool = False,
+    ) -> None:
+        self.model_config = model_config
+        self.parallel_config = parallel_config
+        self.scheduler_config = scheduler_config
+        self.local_rank = local_rank
+        self.rank = rank
+        self.distributed_init_method = distributed_init_method
+        self.is_driver_worker = is_driver_worker
+        self.post_model_path = post_model_path
+
+        if self.is_driver_worker:
+            assert self.rank == 0, "The driver worker must have rank 0."
+
+        self.model_runner = ModelRunner(
+            model_config,
+            parallel_config,
+            scheduler_config,
+            is_driver_worker,
+            post_model_path,
+        )
+        # Uninitialized cache engine. Will be initialized by
+        # self.init_cache_engine().
+        self.cache_config = None
+        self.cache_engine = None
+        self.cache_events = None
+        self.gpu_cache = None
+
+    def init_model(self) -> None:
+        # torch.distributed.all_reduce does not free the input tensor until
+        # the synchronization point. This causes the memory usage to grow
+        # as the number of all_reduce calls increases. This env var disables
+        # this behavior.
+        # Related issue:
+        # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
+        os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
+
+        # This env var set by Ray causes exceptions with graph building.
+        os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
+        self.device = torch.device(f"cuda:{self.local_rank}")
+        torch.cuda.set_device(self.device)
+
+        _check_if_gpu_supports_dtype(self.model_config.dtype)
+
+        # Initialize the distributed environment.
+        _init_distributed_environment(
+            self.parallel_config, self.rank, self.distributed_init_method
+        )
+
+        # Initialize the model.
+        set_random_seed(self.model_config.seed)
+
+    def load_model(self):
+        self.model_runner.load_model()
+
+    @torch.inference_mode()
+    def profile_num_available_blocks(
+        self,
+        block_size: int,
+        gpu_memory_utilization: float,
+        cpu_swap_space: int,
+    ) -> Tuple[int, int]:
+        # Profile the memory usage of the model and get the maximum number of
+        # cache blocks that can be allocated with the remaining free memory.
+        torch.cuda.empty_cache()
+
+        # Execute a forward pass with dummy inputs to profile the memory usage
+        # of the model.
+        self.model_runner.profile_run()
+
+        # Calculate the number of blocks that can be allocated with the
+        # profiled peak memory.
+        torch.cuda.synchronize()
+        free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info()
+        peak_memory = total_gpu_memory - free_gpu_memory
+
+        cache_block_size = CacheEngine.get_cache_block_size(
+            block_size, self.model_config, self.parallel_config
+        )
+        num_gpu_blocks = int(
+            (total_gpu_memory * gpu_memory_utilization - peak_memory)
+            // cache_block_size
+        )
+        num_cpu_blocks = int(cpu_swap_space // cache_block_size)
+        num_gpu_blocks = max(num_gpu_blocks, 0)
+        num_cpu_blocks = max(num_cpu_blocks, 0)
+        torch.cuda.empty_cache()
+        return num_gpu_blocks, num_cpu_blocks
+
+    def init_cache_engine(self, cache_config: CacheConfig) -> None:
+        self.cache_config = cache_config
+        self.cache_engine = CacheEngine(
+            self.cache_config, self.model_config, self.parallel_config
+        )
+        self.cache_events = self.cache_engine.events
+        self.gpu_cache = self.cache_engine.gpu_cache
+        self.model_runner.set_block_size(self.cache_engine.block_size)
+
+    def warm_up_model(self) -> None:
+        if not self.model_config.enforce_eager:
+            self.model_runner.capture_model(self.gpu_cache)
+        # Reset the seed to ensure that the random state is not affected by
+        # the model initialization and profiling.
+        set_random_seed(self.model_config.seed)
+
+    def cache_swap(
+        self,
+        blocks_to_swap_in: Dict[int, int],
+        blocks_to_swap_out: Dict[int, int],
+        blocks_to_copy: Dict[int, List[int]],
+    ) -> None:
+        # Issue cache operations.
+        issued_cache_op = False
+        if blocks_to_swap_in:
+            self.cache_engine.swap_in(blocks_to_swap_in)
+            issued_cache_op = True
+        if blocks_to_swap_out:
+            self.cache_engine.swap_out(blocks_to_swap_out)
+            issued_cache_op = True
+        if blocks_to_copy:
+            self.cache_engine.copy(blocks_to_copy)
+            issued_cache_op = True
+
+        cache_events = self.cache_events if issued_cache_op else None
+
+        # Wait for cache operations to finish.
+        # TODO(woosuk): Profile swapping overhead and optimize if needed.
+        if cache_events is not None:
+            for event in cache_events:
+                event.wait()
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        seq_group_metadata_list: Optional[List[SequenceGroupMetadata]] = None,
+        blocks_to_swap_in: Optional[Dict[int, int]] = None,
+        blocks_to_swap_out: Optional[Dict[int, int]] = None,
+        blocks_to_copy: Optional[Dict[int, List[int]]] = None,
+    ) -> Optional[SamplerOutput]:
+        if self.is_driver_worker:
+            assert seq_group_metadata_list is not None
+            num_seq_groups = len(seq_group_metadata_list)
+            assert blocks_to_swap_in is not None
+            assert blocks_to_swap_out is not None
+            assert blocks_to_copy is not None
+            block_swapping_info = [
+                blocks_to_swap_in,
+                blocks_to_swap_out,
+                blocks_to_copy,
+            ]
+            broadcast_object_list([num_seq_groups] + block_swapping_info, src=0)
+        else:
+            # num_seq_groups, blocks_to_swap_in, blocks_to_swap_out,
+            # blocks_to_copy (4 elements)
+            recv_data = [None] * 4
+            broadcast_object_list(recv_data, src=0)
+            num_seq_groups = recv_data[0]
+            block_swapping_info = recv_data[1:]
+
+        self.cache_swap(*block_swapping_info)
+
+        # If there is no input, we don't need to execute the model.
+        if num_seq_groups == 0:
+            return {}
+
+        output = self.model_runner.execute_model(
+            seq_group_metadata_list, self.gpu_cache
+        )
+        return output
+
+
+def _init_distributed_environment(
+    parallel_config: ParallelConfig,
+    rank: int,
+    distributed_init_method: Optional[str] = None,
+) -> None:
+    """Initialize the distributed environment."""
+    if torch.distributed.is_initialized():
+        torch_world_size = torch.distributed.get_world_size()
+        if torch_world_size != parallel_config.world_size:
+            raise RuntimeError(
+                "torch.distributed is already initialized but the torch world "
+                "size does not match parallel_config.world_size "
+                f"({torch_world_size} vs. {parallel_config.world_size})."
+            )
+    elif not distributed_init_method:
+        raise ValueError(
+            "distributed_init_method must be set if torch.distributed "
+            "is not already initialized"
+        )
+    else:
+        torch.distributed.init_process_group(
+            backend="nccl",
+            world_size=parallel_config.world_size,
+            rank=rank,
+            init_method=distributed_init_method,
+        )
+
+    # A small all_reduce for warmup.
+    torch.distributed.all_reduce(torch.zeros(1).cuda())
+    initialize_model_parallel(
+        parallel_config.tensor_parallel_size, parallel_config.pipeline_parallel_size
+    )
+
+
+def _check_if_gpu_supports_dtype(torch_dtype: torch.dtype):
+    # Check if the GPU supports the dtype.
+    if torch_dtype == torch.bfloat16:
+        compute_capability = torch.cuda.get_device_capability()
+        if compute_capability[0] < 8:
+            gpu_name = torch.cuda.get_device_name()
+            raise ValueError(
+                "Bfloat16 is only supported on GPUs with compute capability "
+                f"of at least 8.0. Your {gpu_name} GPU has compute capability "
+                f"{compute_capability[0]}.{compute_capability[1]}."
+            )

ChatTTS/norm.py

@@ -10,7 +10,7 @@ import numpy as np
 from .utils import del_all
 
 
-@jit
+@jit(nopython=True)
 def _find_index(table: np.ndarray, val: np.uint16):
     for i in range(table.size):
         if table[i] == val:
@@ -18,7 +18,7 @@ def _find_index(table: np.ndarray, val: np.uint16):
     return -1
 
 
-@jit
+@jit(nopython=True)
 def _fast_replace(
     table: np.ndarray, text: bytes
 ) -> Tuple[np.ndarray, List[Tuple[str, str]]]:
@@ -34,6 +34,40 @@ def _fast_replace(
     return result, replaced_words
 
 
+@jit(nopython=True)
+def _split_tags(text: str) -> Tuple[List[str], List[str]]:
+    texts: List[str] = []
+    tags: List[str] = []
+    current_text = ""
+    current_tag = ""
+    for c in text:
+        if c == "[":
+            texts.append(current_text)
+            current_text = ""
+            current_tag = c
+        elif current_tag != "":
+            current_tag += c
+        else:
+            current_text += c
+        if c == "]":
+            tags.append(current_tag)
+            current_tag = ""
+    if current_text != "":
+        texts.append(current_text)
+    return texts, tags
+
+
+@jit(nopython=True)
+def _combine_tags(texts: List[str], tags: List[str]) -> str:
+    text = ""
+    for t in texts:
+        tg = ""
+        if len(tags) > 0:
+            tg = tags.pop(0)
+        text += t + tg
+    return text
+
+
 class Normalizer:
     def __init__(self, map_file_path: str, logger=logging.getLogger(__name__)):
         self.logger = logger
@@ -58,7 +92,7 @@ class Normalizer:
         """
         self.coding = "utf-16-le" if sys.byteorder == "little" else "utf-16-be"
         self.reject_pattern = re.compile(r"[^\u4e00-\u9fffA-Za-z，。、,\. ]")
-        self.sub_pattern = re.compile(r"\[uv_break\]|\[laugh\]|\[lbreak\]")
+        self.sub_pattern = re.compile(r"\[[\w_]+\]")
         self.chinese_char_pattern = re.compile(r"[\u4e00-\u9fff]")
         self.english_word_pattern = re.compile(r"\b[A-Za-z]+\b")
         self.character_simplifier = str.maketrans(
@@ -82,8 +116,8 @@ class Normalizer:
                 "!": ".",
                 "(": ",",
                 ")": ",",
-                "[": ",",
-                "]": ",",
+                # "[": ",",
+                # "]": ",",
                 ">": ",",
                 "<": ",",
                 "-": ",",
@@ -136,7 +170,12 @@ class Normalizer:
         if do_text_normalization:
             _lang = self._detect_language(text) if lang is None else lang
             if _lang in self.normalizers:
-                text = self.normalizers[_lang](text)
+                texts, tags = _split_tags(text)
+                self.logger.debug("split texts %s, tags %s", str(texts), str(tags))
+                texts = [self.normalizers[_lang](t) for t in texts]
+                self.logger.debug("normed texts %s", str(texts))
+                text = _combine_tags(texts, tags) if len(tags) > 0 else texts[0]
+                self.logger.debug("combined text %s", text)
             if _lang == "zh":
                 text = self._apply_half2full_map(text)
         invalid_characters = self._count_invalid_characters(text)
@@ -153,7 +192,12 @@ class Normalizer:
                 repl_res = ", ".join([f"{_[0]}->{_[1]}" for _ in replaced_words])
                 self.logger.info(f"replace homophones: {repl_res}")
         if len(invalid_characters):
-            text = self.reject_pattern.sub("", text)
+            texts, tags = _split_tags(text)
+            self.logger.debug("split texts %s, tags %s", str(texts), str(tags))
+            texts = [self.reject_pattern.sub("", t) for t in texts]
+            self.logger.debug("normed texts %s", str(texts))
+            text = _combine_tags(texts, tags) if len(tags) > 0 else texts[0]
+            self.logger.debug("combined text %s", text)
         return text
 
     def register(self, name: str, normalizer: Callable[[str], str]) -> bool:

ChatTTS/res/sha256_map.json

@@ -1,8 +1,13 @@
 {
-	"sha256_asset_Decoder_pt"   : "9964e36e840f0e3a748c5f716fe6de6490d2135a5f5155f4a642d51860e2ec38",
-	"sha256_asset_DVAE_pt"      : "613cb128adf89188c93ea5880ea0b798e66b1fe6186d0c535d99bcd87bfd6976",
-	"sha256_asset_GPT_pt"       : "d7d4ee6461ea097a2be23eb40d73fb94ad3b3d39cb64fbb50cb3357fd466cadb",
-	"sha256_asset_spk_stat_pt"  : "3228d8a4cbbf349d107a1b76d2f47820865bd3c9928c4bdfe1cefd5c7071105f",
-	"sha256_asset_tokenizer_pt" : "e911ae7c6a7c27953433f35c44227a67838fe229a1f428503bdb6cd3d1bcc69c",
-	"sha256_asset_Vocos_pt"     : "09a670eda1c08b740013679c7a90ebb7f1a97646ea7673069a6838e6b51d6c58"
+	"sha256_asset_Decoder_safetensors": "77aa55e0a977949c4733df3c6f876fa85860d3298cba63295a7bc6901729d4e0",
+	"sha256_asset_DVAE_safetensors"   : "1d0b044a8368c0513100a2eca98456b289e6be6a18b7a63be1bcaa315ea874d9",
+	"sha256_asset_Embed_safetensors"  : "2ff0be7134934155741b643b74e32fb6bf3eec41257984459b2ed60cdb4c48b0",
+	"sha256_asset_Vocos_safetensors"  : "07e5561491cce41f7f90cfdb94b2ff263ff5742c3d89339db99b17ad82cc3f44",
+
+	"sha256_asset_gpt_config_json"         : "0aaa1ecd96c49ad4f473459eb1982fa7ad79fa5de08cde2781bf6ad1f9a0c236",
+	"sha256_asset_gpt_model_safetensors"   : "cd0806fd971f52f6a22c923ec64982b305e817bcc41ca83417fcf9141b984a0f",
+
+	"sha256_asset_tokenizer_special_tokens_map_json": "bd0ac9d9bb1657996b5c5fbcaa7d80f8de530d01a283da97f89deae5b1b8d011",
+	"sha256_asset_tokenizer_tokenizer_config_json"  : "43e9d658b554fa5ee8d8e1d763349323bfef1ed7a89c0794220ab8861387d421",
+	"sha256_asset_tokenizer_tokenizer_json"         : "843838a64e121e23e774cc75874c6fe862198d9f7dd43747914633a8fd89c20e"
 }

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 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

@@ -3,7 +3,7 @@ from pathlib import Path
 import hashlib
 import requests
 from io import BytesIO
-from typing import Dict
+from typing import Dict, Tuple, Optional
 from mmap import mmap, ACCESS_READ
 
 from .log import logger
@@ -43,33 +43,81 @@ def check_model(
     return True
 
 
-def check_all_assets(base_dir: Path, sha256_map: Dict[str, str], update=False) -> bool:
-    logger.get_logger().info("checking assets...")
-    current_dir = base_dir / "asset"
-    names = [
-        "Decoder.pt",
-        "DVAE_full.pt",
-        "GPT.pt",
-        "spk_stat.pt",
-        "tokenizer.pt",
-        "Vocos.pt",
-    ]
+def check_folder(
+    base_dir: Path,
+    *innder_dirs: str,
+    names: Tuple[str],
+    sha256_map: Dict[str, str],
+    update=False,
+) -> bool:
+    key = "sha256_"
+    current_dir = base_dir
+    for d in innder_dirs:
+        current_dir /= d
+        key += f"{d}_"
+
     for model in names:
         menv = model.replace(".", "_")
-        if not check_model(
-            current_dir, model, sha256_map[f"sha256_asset_{menv}"], update
-        ):
+        if not check_model(current_dir, model, sha256_map[f"{key}{menv}"], update):
             return False
+    return True
+
+
+def check_all_assets(base_dir: Path, sha256_map: Dict[str, str], update=False) -> bool:
+    logger.get_logger().info("checking assets...")
+
+    if not check_folder(
+        base_dir,
+        "asset",
+        names=(
+            "Decoder.safetensors",
+            "DVAE.safetensors",
+            "Embed.safetensors",
+            "Vocos.safetensors",
+        ),
+        sha256_map=sha256_map,
+        update=update,
+    ):
+        return False
+
+    if not check_folder(
+        base_dir,
+        "asset",
+        "gpt",
+        names=(
+            "config.json",
+            "model.safetensors",
+        ),
+        sha256_map=sha256_map,
+        update=update,
+    ):
+        return False
+
+    if not check_folder(
+        base_dir,
+        "asset",
+        "tokenizer",
+        names=(
+            "special_tokens_map.json",
+            "tokenizer_config.json",
+            "tokenizer.json",
+        ),
+        sha256_map=sha256_map,
+        update=update,
+    ):
+        return False
 
     logger.get_logger().info("all assets are already latest.")
     return True
 
 
-def download_and_extract_tar_gz(url: str, folder: str):
+def download_and_extract_tar_gz(
+    url: str, folder: str, headers: Optional[Dict[str, str]] = None
+):
     import tarfile
 
     logger.get_logger().info(f"downloading {url}")
-    response = requests.get(url, stream=True, timeout=(5, 10))
+    response = requests.get(url, headers=headers, stream=True, timeout=(10, 3))
     with BytesIO() as out_file:
         out_file.write(response.content)
         out_file.seek(0)
@@ -79,11 +127,13 @@ def download_and_extract_tar_gz(url: str, folder: str):
         logger.get_logger().info(f"extracted into {folder}")
 
 
-def download_and_extract_zip(url: str, folder: str):
+def download_and_extract_zip(
+    url: str, folder: str, headers: Optional[Dict[str, str]] = None
+):
     import zipfile
 
     logger.get_logger().info(f"downloading {url}")
-    response = requests.get(url, stream=True, timeout=(5, 10))
+    response = requests.get(url, headers=headers, stream=True, timeout=(10, 3))
     with BytesIO() as out_file:
         out_file.write(response.content)
         out_file.seek(0)
@@ -93,15 +143,7 @@ def download_and_extract_zip(url: str, folder: str):
         logger.get_logger().info(f"extracted into {folder}")
 
 
-def download_dns_yaml(url: str, folder: str):
-    logger.get_logger().info(f"downloading {url}")
-    response = requests.get(url, stream=True, timeout=(5, 10))
-    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, version="0.2.6"):
+def download_all_assets(tmpdir: str, homedir: str, version="0.2.11"):
     import subprocess
     import platform
 
@@ -125,48 +167,15 @@ def download_all_assets(tmpdir: str, version="0.2.6"):
     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", "assets/chtts"])
-    except Exception:
-        BASE_URL = "https://raw.gitcode.com/u011570312/RVC-Models-Downloader/assets/"
-        suffix = {
-            "darwin_amd64": "987",
-            "darwin_arm64": "988",
-            "linux_386": "989",
-            "linux_amd64": "990",
-            "linux_arm64": "991",
-            "windows_386": "992",
-            "windows_amd64": "993",
-        }[f"{system_type}_{architecture}"]
-        RVCMD_URL = BASE_URL + suffix
-        download_dns_yaml(
-            "https://raw.gitcode.com/u011570312/RVC-Models-Downloader/raw/main/dns.yaml",
-            tmpdir,
-        )
-        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"),
-                "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,26 +1,38 @@
+import importlib.util
+
 import torch
 
+try:
+    import torch_npu
+except ImportError:
+    pass
+
 from .log import logger
 
 
 def select_device(min_memory=2047, experimental=False):
-    if torch.cuda.is_available():
-        selected_gpu = 0
+    has_cuda = torch.cuda.is_available()
+    if has_cuda or _is_torch_npu_available():
+        provider = torch.cuda if has_cuda else torch.npu
+        """
+        Using Ascend NPU to accelerate the process of inferencing when GPU is not found.
+        """
+        dev_idx = 0
         max_free_memory = -1
-        for i in range(torch.cuda.device_count()):
-            props = torch.cuda.get_device_properties(i)
-            free_memory = props.total_memory - torch.cuda.memory_reserved(i)
+        for i in range(provider.device_count()):
+            props = provider.get_device_properties(i)
+            free_memory = props.total_memory - provider.memory_reserved(i)
             if max_free_memory < free_memory:
-                selected_gpu = i
+                dev_idx = i
                 max_free_memory = free_memory
         free_memory_mb = max_free_memory / (1024 * 1024)
         if free_memory_mb < min_memory:
             logger.get_logger().warning(
-                f"GPU {selected_gpu} has {round(free_memory_mb, 2)} MB memory left. Switching to CPU."
+                f"{provider.device(dev_idx)} has {round(free_memory_mb, 2)} MB memory left. Switching to CPU."
             )
             device = torch.device("cpu")
         else:
-            device = torch.device(f"cuda:{selected_gpu}")
+            device = provider._get_device(dev_idx)
     elif torch.backends.mps.is_available():
         """
         Currently MPS is slower than CPU while needs more memory and core utility,
@@ -28,13 +40,35 @@ 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 found, use CPU instead")
+        logger.get_logger().warning("no GPU or NPU found, use CPU instead")
         device = torch.device("cpu")
 
     return device
+
+
+def _is_torch_npu_available():
+    try:
+        # will raise a AttributeError if torch_npu is not imported or a RuntimeError if no NPU found
+        _ = torch.npu.device_count()
+        return torch.npu.is_available()
+    except (AttributeError, RuntimeError):
+        return False

ChatTTS/utils/io.py

@@ -1,10 +1,29 @@
 import os
 import logging
-from typing import Union
+from typing import Union, IO
 from dataclasses import is_dataclass
 
+from safetensors import safe_open
+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):
+    state_dict_tensors = {}
+    with safe_open(filename, framework="pt") as f:
+        for k in f.keys():
+            state_dict_tensors[k] = f.get_tensor(k)
+    return state_dict_tensors
+
 
 def get_latest_modified_file(directory):
 

README.md

@@ -6,13 +6,13 @@
 A generative speech model for daily dialogue.
 
 [![Licence](https://img.shields.io/github/license/2noise/ChatTTS?style=for-the-badge)](https://github.com/2noise/ChatTTS/blob/main/LICENSE)
-[![PyPI](https://img.shields.io/pypi/v/ChatTTS.svg?style=for-the-badge)](https://pypi.org/project/ChatTTS)
+[![PyPI](https://img.shields.io/pypi/v/ChatTTS.svg?style=for-the-badge&color=green)](https://pypi.org/project/ChatTTS)
 
 [![Huggingface](https://img.shields.io/badge/🤗%20-Models-yellow.svg?style=for-the-badge)](https://huggingface.co/2Noise/ChatTTS)
 [![Open In Colab](https://img.shields.io/badge/Colab-F9AB00?style=for-the-badge&logo=googlecolab&color=525252)](https://colab.research.google.com/github/2noise/ChatTTS/blob/main/examples/ipynb/colab.ipynb)
 [![Discord](https://img.shields.io/badge/Discord-7289DA?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/Ud5Jxgx5yD)
 
-**English** | [**简体中文**](docs/cn/README.md) | [**日本語**](docs/jp/README.md) | [**Русский**](docs/ru/README.md) | [**Español**](docs/es/README.md) | [**Français**](docs/fr/README.md)
+**English** | [**简体中文**](docs/cn/README.md) | [**日本語**](docs/jp/README.md) | [**Русский**](docs/ru/README.md) | [**Español**](docs/es/README.md) | [**Français**](docs/fr/README.md) | [**한국어**](docs/kr/README.md)
 
 </div>
 
@@ -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.
 
@@ -45,9 +46,10 @@ ChatTTS is a text-to-speech model designed specifically for dialogue scenarios s
 - The open-source version on **[HuggingFace](https://huggingface.co/2Noise/ChatTTS)** is a 40,000 hours pre-trained model without SFT.
 
 ### Roadmap
-- [x] Open-source the 40k hour base model and spk_stats file.
+- [x] Open-source the 40k-hours-base model and spk_stats file.
 - [x] Streaming audio generation.
-- [ ] Open-source the 40k hour version with multi-emotion control.
+- [x] Open-source DVAE encoder and zero shot inferring code.
+- [ ] Multi-emotion controlling.
 - [ ] ChatTTS.cpp (new repo in `2noise` org is welcomed)
 
 ### Licenses
@@ -95,30 +97,38 @@ pip install --upgrade -r requirements.txt
 
 #### 2. Install from conda
 ```bash
-conda create -n chattts
+conda create -n chattts python=3.11
 conda activate chattts
 pip install -r requirements.txt
 ```
 
-#### Optional: Install TransformerEngine if using NVIDIA GPU (Linux only)
+#### Optional: Install vLLM (Linux only)
+```bash
+pip install safetensors vllm==0.2.7 torchaudio
+```
+
+#### Unrecommended Optional: Install TransformerEngine if using NVIDIA GPU (Linux only)
+> [!Warning]
+> DO NOT INSTALL! 
+> The adaptation of TransformerEngine is currently under development and CANNOT run properly now. 
+> Only install it on developing purpose. See more details on at #672 #676
+
 > [!Note]
 > The installation process is very slow.
 
-> [!Warning]
-> The adaptation of TransformerEngine is currently under development and CANNOT run properly now. 
-> Only install it on developing purpose.
-
 ```bash
 pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable
 ```
 
-#### Optional: Install FlashAttention-2 (mainly NVIDIA GPU)
+#### Unrecommended Optional: Install FlashAttention-2 (mainly NVIDIA GPU)
+> [!Warning]
+> DO NOT INSTALL! 
+> Currently the FlashAttention-2 will slow down the generating speed according to [this issue](https://github.com/huggingface/transformers/issues/26990). 
+> Only install it on developing purpose.
+
 > [!Note]
 > See supported devices at the [Hugging Face Doc](https://huggingface.co/docs/transformers/perf_infer_gpu_one#flashattention-2).
 
-> [!Warning]
-> Currently the FlashAttention-2 will slow down the generating speed according to [this issue](https://github.com/huggingface/transformers/issues/26990). 
-> Only install it on developing purpose.
 
 ```bash
 pip install flash-attn --no-build-isolation
@@ -171,7 +181,13 @@ texts = ["PUT YOUR 1st TEXT HERE", "PUT YOUR 2nd TEXT HERE"]
 wavs = chat.infer(texts)
 
 for i in range(len(wavs)):
-    torchaudio.save(f"basic_output{i}.wav", torch.from_numpy(wavs[i]).unsqueeze(0), 24000)
+    """
+    In some versions of torchaudio, the first line works but in other versions, so does the second line.
+    """
+    try:
+        torchaudio.save(f"basic_output{i}.wav", torch.from_numpy(wavs[i]).unsqueeze(0), 24000)
+    except:
+        torchaudio.save(f"basic_output{i}.wav", torch.from_numpy(wavs[i]), 24000)
 ```
 
 ### Advanced Usage
@@ -210,7 +226,13 @@ wavs = chat.infer(
 
 text = 'What is [uv_break]your favorite english food?[laugh][lbreak]'
 wavs = chat.infer(text, skip_refine_text=True, params_refine_text=params_refine_text,  params_infer_code=params_infer_code)
-torchaudio.save("word_level_output.wav", torch.from_numpy(wavs[0]).unsqueeze(0), 24000)
+"""
+In some versions of torchaudio, the first line works but in other versions, so does the second line.
+"""
+try:
+    torchaudio.save("word_level_output.wav", torch.from_numpy(wavs[0]).unsqueeze(0), 24000)
+except:
+    torchaudio.save("word_level_output.wav", torch.from_numpy(wavs[0]), 24000)
 ```
 
 <details open>
@@ -278,7 +300,7 @@ This is a problem that typically occurs with autoregressive models (for bark and
 In the current released model, the only token-level control units are `[laugh]`, `[uv_break]`, and `[lbreak]`. In future versions, we may open-source models with additional emotional control capabilities.
 
 ## Acknowledgements
-- [bark](https://github.com/suno-ai/bark), [XTTSv2](https://github.com/coqui-ai/TTS) and [valle](https://arxiv.org/abs/2301.02111) demostrate a remarkable TTS result by an autoregressive-style system.
+- [bark](https://github.com/suno-ai/bark), [XTTSv2](https://github.com/coqui-ai/TTS) and [valle](https://arxiv.org/abs/2301.02111) demonstrate a remarkable TTS result by an autoregressive-style system.
 - [fish-speech](https://github.com/fishaudio/fish-speech) reveals capability of GVQ as audio tokenizer for LLM modeling.
 - [vocos](https://github.com/gemelo-ai/vocos) which is used as a pretrained vocoder.
 

docs/cn/README.md

@@ -6,13 +6,13 @@
 一款适用于日常对话的生成式语音模型。
 
 [![Licence](https://img.shields.io/github/license/2noise/ChatTTS?style=for-the-badge)](https://github.com/2noise/ChatTTS/blob/main/LICENSE)
-[![PyPI](https://img.shields.io/pypi/v/ChatTTS.svg?style=for-the-badge)](https://pypi.org/project/ChatTTS)
+[![PyPI](https://img.shields.io/pypi/v/ChatTTS.svg?style=for-the-badge&color=green)](https://pypi.org/project/ChatTTS)
 
 [![Huggingface](https://img.shields.io/badge/🤗%20-Models-yellow.svg?style=for-the-badge)](https://huggingface.co/2Noise/ChatTTS)
 [![Open In Colab](https://img.shields.io/badge/Colab-F9AB00?style=for-the-badge&logo=googlecolab&color=525252)](https://colab.research.google.com/github/2noise/ChatTTS/blob/main/examples/ipynb/colab.ipynb)
 [![Discord](https://img.shields.io/badge/Discord-7289DA?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/Ud5Jxgx5yD)
 
-[**English**](../../README.md) | **简体中文** | [**日本語**](../jp/README.md) | [**Русский**](../ru/README.md) | [**Español**](../es/README.md) | [**Français**](../fr/README.md)
+[**English**](../../README.md) | **简体中文** | [**日本語**](../jp/README.md) | [**Русский**](../ru/README.md) | [**Español**](../es/README.md) | [**Français**](../fr/README.md) | [**한국어**](../kr/README.md)
 
 </div>
 
@@ -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 助手）设计的文本转语音模型。
 
@@ -52,6 +53,7 @@ ChatTTS 是一款专门为对话场景（例如 LLM 助手）设计的文本转
 
 - [x] 开源 4 万小时基础模型和 spk_stats 文件。
 - [x] 支持流式语音输出。
+- [x] 开源 DVAE 编码器和零样本推理代码
 - [ ] 开源具有多情感控制功能的 4 万小时版本。
 - [ ] ChatTTS.cpp (欢迎在 2noise 组织中新建仓库)。
 

docs/es/README.md

@@ -11,12 +11,17 @@ Un modelo de generación de voz para la conversación diaria.
 [![Open In Colab](https://img.shields.io/badge/Colab-F9AB00?style=for-the-badge&logo=googlecolab&color=525252)](https://colab.research.google.com/github/2noise/ChatTTS/blob/main/examples/ipynb/colab.ipynb)
 
 [**English**](../../README.md) | [**简体中文**](../cn/README.md) | [**日本語**](../jp/README.md) | [**Русский**](../ru/README.md) | **Español**
- | [**Français**](../fr/README.md)
+ | [**Français**](../fr/README.md) | [**한국어**](../kr/README.md)
 </div>
 
 > [!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

@@ -6,13 +6,13 @@
 Un modèle de parole génératif pour le dialogue quotidien.
 
 [![Licence](https://img.shields.io/github/license/2noise/ChatTTS?style=for-the-badge)](https://github.com/2noise/ChatTTS/blob/main/LICENSE)
-[![PyPI](https://img.shields.io/pypi/v/ChatTTS.svg?style=for-the-badge)](https://pypi.org/project/ChatTTS)
+[![PyPI](https://img.shields.io/pypi/v/ChatTTS.svg?style=for-the-badge&color=green)](https://pypi.org/project/ChatTTS)
 
 [![Huggingface](https://img.shields.io/badge/🤗%20-Models-yellow.svg?style=for-the-badge)](https://huggingface.co/2Noise/ChatTTS)
 [![Open In Colab](https://img.shields.io/badge/Colab-F9AB00?style=for-the-badge&logo=googlecolab&color=525252)](https://colab.research.google.com/github/2noise/ChatTTS/blob/main/examples/ipynb/colab.ipynb)
 [![Discord](https://img.shields.io/badge/Discord-7289DA?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/Ud5Jxgx5yD)
 
-[**English**](../../README.md) | [**简体中文**](../cn/README.md) | [**日本語**](../jp/README.md) | [**Русский**](../ru/README.md) | [**Español**](../es/README.md)| **Français** 
+[**English**](../../README.md) | [**简体中文**](../cn/README.md) | [**日本語**](../jp/README.md) | [**Русский**](../ru/README.md) | [**Español**](../es/README.md)| **Français**  | [**한국어**](../kr/README.md)
 
 </div>
 
@@ -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

@@ -4,12 +4,21 @@
 
 [![Huggingface](https://img.shields.io/badge/🤗%20-Models-yellow.svg?style=for-the-badge)](https://huggingface.co/2Noise/ChatTTS)
 
-[**English**](../../README.md) | [**简体中文**](../cn/README.md) | **日本語** | [**Русский**](../ru/README.md) | [**Español**](../es/README.md) | [**Français**](../fr/README.md)
+[**English**](../../README.md) | [**简体中文**](../cn/README.md) | **日本語** | [**Русский**](../ru/README.md) | [**Español**](../es/README.md) | [**Français**](../fr/README.md) | [**한국어**](../kr/README.md)
 
 ChatTTSは、LLMアシスタントなどの対話シナリオ用に特別に設計されたテキストから音声へのモデルです。英語と中国語の両方をサポートしています。私たちのモデルは、中国語と英語で構成される100,000時間以上でトレーニングされています。**[HuggingFace](https://huggingface.co/2Noise/ChatTTS)**でオープンソース化されているバージョンは、40,000時間の事前トレーニングモデルで、SFTは行われていません。
 
 モデルやロードマップについての正式なお問い合わせは、**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

@@ -0,0 +1,319 @@
+<div align="center">
+
+<a href="https://trendshift.io/repositories/10489" target="_blank"><img src="https://trendshift.io/api/badge/repositories/10489" alt="2noise%2FChatTTS | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
+
+# ChatTTS
+일상 대화를 위한 생성형 음성 모델입니다.
+
+[![Licence](https://img.shields.io/github/license/2noise/ChatTTS?style=for-the-badge)](https://github.com/2noise/ChatTTS/blob/main/LICENSE)
+[![PyPI](https://img.shields.io/pypi/v/ChatTTS.svg?style=for-the-badge&color=green)](https://pypi.org/project/ChatTTS)
+
+[![Huggingface](https://img.shields.io/badge/🤗%20-Models-yellow.svg?style=for-the-badge)](https://huggingface.co/2Noise/ChatTTS)
+[![Open In Colab](https://img.shields.io/badge/Colab-F9AB00?style=for-the-badge&logo=googlecolab&color=525252)](https://colab.research.google.com/github/2noise/ChatTTS/blob/main/examples/ipynb/colab.ipynb)
+[![Discord](https://img.shields.io/badge/Discord-7289DA?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/Ud5Jxgx5yD)
+
+[**English**](../../README.md) | [**简体中文**](../cn/README.md) | [**日本語**](../jp/README.md) | [**Русский**](../ru/README.md) | [**Español**](../es/README.md) | [**Français**](../fr/README.md) | **한국어**
+
+</div>
+
+> [!NOTE]
+> 이 문서는 최신 버전이 아닐 수 있습니다. [영어 문서](../../README.md)를 기준으로 작업하는 것을 권장합니다.
+
+## 프로젝트 소개
+
+> [!Note]
+> 이 저장소에는 알고리즘 구조와 간단한 예시들이 포함되어 있습니다.
+
+> [!Tip]
+> 이 프로젝트에서 파생된 프로젝트는 커뮤니티가 유지 관리하는 커뮤니티[Awesome-ChatTTS](https://github.com/libukai/Awesome-ChatTTS)를 참조하시길 바랍니다.  
+> 코드베이스의 다이어그램 시각화는 [여기](https://github.com/CodeBoarding/GeneratedOnBoardings/blob/main/ChatTTS/on_boarding.md)에서 확인할 수 있습니다.
+
+ChatTTS는 대화 기반 작업(예: LLM 어시스턴트)을 위해 설계된 텍스트-음성 변환(TTS) 모델입니다.
+
+### 지원 언어
+
+- [x] 영어
+- [x] 중국어
+- [ ] 계속 추가 예정...
+
+### 프로젝트 특징
+
+> 이 프로젝트의 내용은 **[Bilibili](https://www.bilibili.com/video/BV1zn4y1o7iV)**에서 제공되는 비디오를 참조하시길 바랍니다.
+
+1. **대화형 TTS**: ChatTTS는 대화 기반 작업에 최적화되어 자연스럽고 표현력 있는 음성 합성을 구현합니다. 다중 화자를 지원하여 상호작용적인 대화를 가능하게 합니다.
+2. **세밀한 제어**: 이 모델은 웃음, 일시 정지, 삽입어 등 세밀한 운율적 특징을 예측하고 제어할 수 있습니다.
+3. **향상된 운율**: ChatTTS는 운율 측면에서 대부분의 오픈 소스 TTS 모델을 능가하며, 추가 연구와 개발을 지원하기 위해 사전 훈련된 모델을 제공합니다.
+
+### 데이터셋 및 모델
+> [!Important]
+> 공개된 모델은 학술 목적으로만 사용 가능합니다.
+
+- 주요 모델은 100,000+ 시간의 중국어 및 영어 오디오 데이터를 사용하여 훈련되었습니다.
+- **[HuggingFace](https://huggingface.co/2Noise/ChatTTS)**에서 제공되는 오픈 소스 버전은 40,000시간의 사전 훈련된 모델로, SFT가 적용되지 않았습니다.
+
+### 로드맵
+- [x] 40,000시간 기반 모델과 spk_stats 파일 오픈 소스화.
+- [x] 스트리밍 오디오 생성.
+- [x] DVAE 인코더와 제로 샷 추론 코드 오픈 소스화.
+- [ ] 다중 감정 제어 기능.
+- [ ] ChatTTS.cpp (`2noise` 조직 내의 새로운 저장소를 환영합니다.)
+
+### 라이선스
+
+#### 코드
+코드는 `AGPLv3+` 라이선스를 따릅니다.
+
+#### 모델
+모델은 `CC BY-NC 4.0` 라이선스로 공개되었습니다. 이 모델은 교육 및 연구 목적으로만 사용되며, 상업적 또는 불법적 목적으로 사용되어서는 안 됩니다. 저자들은 정보의 정확성, 완전성, 신뢰성을 보장하지 않습니다. 이 저장소에서 사용된 정보와 데이터는 학술 및 연구 목적으로만 사용되며, 공개적으로 이용 가능한 출처에서 얻은 데이터입니다. 저자들은 데이터에 대한 소유권 또는 저작권을 주장하지 않습니다.
+
+### 면책 조항
+
+ChatTTS는 강력한 텍스트-음성 변환 시스템입니다. 그렇기에 기술을 책임감 있고 윤리적으로 사용하는 것은 아주 중요합니다. ChatTTS의 악용을 방지하기 위해 40,000시간 모델의 훈련 중 소량의 고주파 노이즈를 추가하고, 오디오 품질을 최대한 압축하여 MP3 형식으로 제공했습니다. 또한, 우리는 내부적으로 탐지 모델을 훈련했으며, 추후 이를 오픈 소스화할 계획입니다.
+
+### 연락처
+> GitHub 이슈/PR은 언제든지 환영합니다.
+
+#### 공식 문의
+모델 및 로드맵에 대한 공식적인 문의는 **open-source@2noise.com**으로 연락해 주십시오.
+
+#### 온라인 채팅
+##### 1. QQ Group (Chinese Social APP)
+- **Group 1**, 808364215
+- **Group 2**, 230696694
+- **Group 3**, 933639842
+- **Group 4**, 608667975
+
+##### 2. Discord 서버
+[이곳](https://discord.gg/Ud5Jxgx5yD)를 클릭하여 참여하십시오.
+
+## 시작하기
+### 레포지토리 클론
+```bash
+git clone https://github.com/2noise/ChatTTS
+cd ChatTTS
+```
+
+### 의존성 설치
+#### 1. 직접 설치
+```bash
+pip install --upgrade -r requirements.txt
+```
+
+#### 2. Conda에서 설치
+```bash
+conda create -n chattts
+conda activate chattts
+pip install -r requirements.txt
+```
+
+#### 선택사항: vLLM 설치 (Linux 전용)
+```bash
+pip install safetensors vllm==0.2.7 torchaudio
+```
+
+#### 권장되지 않는 선택사항: NVIDIA GPU 사용 시 TransformerEngine 설치 (Linux 전용)
+> [!Warning]
+> 설치하지 마십시오!
+> TransformerEngine의 적응 작업은 현재 개발 중이며, 아직 제대로 작동하지 않습니다.
+> 개발 목적으로만 설치하십시오. 자세한 내용은 #672 및 #676에서 확인할 수 있습니다.
+
+> [!Note]
+> 설치 과정은 매우 느립니다.
+
+```bash
+pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable
+```
+
+#### 권장되지 않는 선택사항: FlashAttention-2 설치 (주로 NVIDIA GPU)
+> [!Warning]
+> 설치하지 마십시오!
+> 현재 FlashAttention-2는 [이 이슈](https://github.com/huggingface/transformers/issues/26990)에 따르면 생성 속도를 저하시킵니다.
+> 개발 목적으로만 설치하십시오.
+
+> [!Note]
+> 지원되는 장치는 [Hugging Face 문서](https://huggingface.co/docs/transformers/perf_infer_gpu_one#flashattention-2)에서 확인할 수 있습니다.
+
+```bash
+pip install flash-attn --no-build-isolation
+```
+
+### 빠른 시작
+> 아래 명령어를 실행할 때 반드시 프로젝트 루트 디렉토리에서 실행하십시오.
+
+#### 1. WebUI 실행
+```bash
+python examples/web/webui.py
+```
+
+#### 2. 커맨드 라인에서 추론
+> 오디오는 `./output_audio_n.mp3`에 저장됩니다.
+
+```bash
+python examples/cmd/run.py "Your text 1." "Your text 2."
+```
+
+## 설치 방법
+
+1. PyPI에서 안정 버전 설치
+```bash
+pip install ChatTTS
+```
+
+2. GitHub에서 최신 버전 설치
+```bash
+pip install git+https://github.com/2noise/ChatTTS
+```
+
+3. 로컬 디렉토리에서 개발 모드로 설치
+```bash
+pip install -e .
+```
+
+### 기본 사용법
+
+```python
+import ChatTTS
+import torch
+import torchaudio
+
+chat = ChatTTS.Chat()
+chat.load(compile=False) # 성능 향상을 위해 True로 설정 가능
+
+texts = ["PUT YOUR 1st TEXT HERE", "PUT YOUR 2nd TEXT HERE"]
+
+wavs = chat.infer(texts)
+
+for i in range(len(wavs)):
+    """
+    torchaudio의 버전에 따라 첫 번째 줄이 작동할 수 있고, 다른 버전에서는 두 번째 줄이 작동할 수 있습니다.
+    """
+    try:
+        torchaudio.save(f"basic_output{i}.wav", torch.from_numpy(wavs[i]).unsqueeze(0), 24000)
+    except:
+        torchaudio.save(f"basic_output{i}.wav", torch.from_numpy(wavs[i]), 24000)
+```
+
+### Advanced Usage
+
+```python
+###################################
+# Sample a speaker from Gaussian.
+
+rand_spk = chat.sample_random_speaker()
+print(rand_spk) # save it for later timbre recovery
+
+params_infer_code = ChatTTS.Chat.InferCodeParams(
+    spk_emb = rand_spk, # add sampled speaker 
+    temperature = .3,   # using custom temperature
+    top_P = 0.7,        # top P decode
+    top_K = 20,         # top K decode
+)
+
+###################################
+# For sentence level manual control.
+
+# use oral_(0-9), laugh_(0-2), break_(0-7) 
+# to generate special token in text to synthesize.
+params_refine_text = ChatTTS.Chat.RefineTextParams(
+    prompt='[oral_2][laugh_0][break_6]',
+)
+
+wavs = chat.infer(
+    texts,
+    params_refine_text=params_refine_text,
+    params_infer_code=params_infer_code,
+)
+
+###################################
+# For word level manual control.
+
+text = 'What is [uv_break]your favorite english food?[laugh][lbreak]'
+wavs = chat.infer(text, skip_refine_text=True, params_refine_text=params_refine_text,  params_infer_code=params_infer_code)
+"""
+In some versions of torchaudio, the first line works but in other versions, so does the second line.
+"""
+try:
+    torchaudio.save("word_level_output.wav", torch.from_numpy(wavs[0]).unsqueeze(0), 24000)
+except:
+    torchaudio.save("word_level_output.wav", torch.from_numpy(wavs[0]), 24000)
+```
+
+<details open>
+  <summary><h4>Example: self introduction</h4></summary>
+
+```python
+inputs_en = """
+chat T T S is a text to speech model designed for dialogue applications. 
+[uv_break]it supports mixed language input [uv_break]and offers multi speaker 
+capabilities with precise control over prosodic elements like 
+[uv_break]laughter[uv_break][laugh], [uv_break]pauses, [uv_break]and intonation. 
+[uv_break]it delivers natural and expressive speech,[uv_break]so please
+[uv_break] use the project responsibly at your own risk.[uv_break]
+""".replace('\n', '') # English is still experimental.
+
+params_refine_text = ChatTTS.Chat.RefineTextParams(
+    prompt='[oral_2][laugh_0][break_4]',
+)
+
+audio_array_en = chat.infer(inputs_en, params_refine_text=params_refine_text)
+torchaudio.save("self_introduction_output.wav", torch.from_numpy(audio_array_en[0]), 24000)
+```
+
+<table>
+<tr>
+<td align="center">
+
+**male speaker**
+
+</td>
+<td align="center">
+
+**female speaker**
+
+</td>
+</tr>
+<tr>
+<td align="center">
+
+[male speaker](https://github.com/2noise/ChatTTS/assets/130631963/e0f51251-db7f-4d39-a0e9-3e095bb65de1)
+
+</td>
+<td align="center">
+
+[female speaker](https://github.com/2noise/ChatTTS/assets/130631963/f5dcdd01-1091-47c5-8241-c4f6aaaa8bbd)
+
+</td>
+</tr>
+</table>
+
+</details>
+
+## FAQ
+
+#### 1. VRAM이 얼마나 필요한가요? 추론 속도는 어느 정도인가요?
+30초 길이의 오디오 클립을 생성하려면 최소 4GB의 GPU 메모리가 필요합니다. 4090 GPU의 경우 초당 약 7개의 의미 토큰에 해당하는 오디오를 생성할 수 있습니다. 실시간 인자(RTF)는 약 0.3입니다.
+
+#### 2. 모델의 안정성은 불안정하며, 화자가 많은 경우 및 오디오 품질이 저하되는 이슈 존재.
+
+이는 일반적으로 autoregressive 모델(bark 및 valle 등)에서 발생하는 불가피한 문제입니다. 현재로선 여러 번 샘플링하여 적절한 결과를 찾는 것이 최선입니다.
+
+#### 3. 웃음 뿐 아니라 다른 감정도 표현할 수 있나요?
+
+현재 공개된 모델에서는 제어 가능한 토큰은 `[laugh]`, `[uv_break]`, `[lbreak]`입니다. 향후 버전의 모델에서는 추가적인 감정 제어 기능 포함하여 오픈 소스로 제공할 계획입니다.
+
+## 감사의 인사
+- [bark](https://github.com/suno-ai/bark), [XTTSv2](https://github.com/coqui-ai/TTS), [valle](https://arxiv.org/abs/2301.02111)는 autoregressive 방식의 시스템으로 뛰어난 TTS 성능을 보여주었습니다.
+- [fish-speech](https://github.com/fishaudio/fish-speech)는 LLM 모델링을 위한 오디오 토크나이저로서 GVQ의 능력을 보여주었습니다.
+- [vocos](https://github.com/gemelo-ai/vocos)는 사전 훈련된 vocoder로 사용되었습니다.
+
+## 특별 감사
+- 초기 알고리즘 실험을 위한 [wlu-audio lab](https://audio.westlake.edu.cn/)에 감사의 말씀을 전합니다.
+
+## 모든 기여자들의 노고에 감사드립니다
+[![contributors](https://contrib.rocks/image?repo=2noise/ChatTTS)](https://github.com/2noise/ChatTTS/graphs/contributors)
+
+<div align="center">
+
+  ![counter](https://counter.seku.su/cmoe?name=chattts&theme=mbs)
+
+</div>

docs/ru/README.md

@@ -4,7 +4,16 @@
 
 [![Huggingface](https://img.shields.io/badge/🤗%20-Models-yellow.svg?style=for-the-badge)](https://huggingface.co/2Noise/ChatTTS)
 
-[**English**](../../README.md) | [**简体中文**](../cn/README.md) | [**日本語**](../jp/README.md) | **Русский** | [**Español**](../es/README.md) | [**Français**](../fr/README.md)
+[**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.
 

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"
 
@@ -23,7 +22,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 +37,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,9 +63,7 @@ class ChatTTSParams(BaseModel):
     use_decoder: bool = True
     do_text_normalization: bool = True
     do_homophone_replacement: bool = False
-    audio_seed: int
-    text_seed: int
-    params_refine_text: ChatTTS.Chat.RefineTextParams
+    params_refine_text: ChatTTS.Chat.RefineTextParams = None
     params_infer_code: ChatTTS.Chat.InferCodeParams
 
 
@@ -63,13 +72,12 @@ async def generate_voice(params: ChatTTSParams):
     logger.info("Text input: %s", str(params.text))
 
     # audio seed
-    if params.audio_seed:
-        torch.manual_seed(params.audio_seed)
+    if params.params_infer_code.manual_seed is not None:
+        torch.manual_seed(params.params_infer_code.manual_seed)
         params.params_infer_code.spk_emb = chat.sample_random_speaker()
 
     # text seed for text refining
     if params.params_refine_text:
-        torch.manual_seed(params.text_seed)
         text = chat.infer(
             text=params.text, skip_refine_text=False, refine_text_only=True
         )

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/cmd/run.py

@@ -6,15 +6,17 @@ if sys.platform == "darwin":
 now_dir = os.getcwd()
 sys.path.append(now_dir)
 
-import argparse
 from typing import Optional, List
+import argparse
 
 import numpy as np
 
 import ChatTTS
 
-from tools.audio import pcm_arr_to_mp3_view
 from tools.logger import get_logger
+from tools.audio import pcm_arr_to_mp3_view
+from tools.normalizer.en import normalizer_en_nemo_text
+from tools.normalizer.zh import normalizer_zh_tn
 
 logger = get_logger("Command")
 
@@ -27,12 +29,48 @@ def save_mp3_file(wav, index):
     logger.info(f"Audio saved to {mp3_filename}")
 
 
-def main(texts: List[str], spk: Optional[str] = None, stream=False):
+def load_normalizer(chat: ChatTTS.Chat):
+    # try to load normalizer
+    try:
+        chat.normalizer.register("en", normalizer_en_nemo_text())
+    except ValueError as e:
+        logger.error(e)
+    except BaseException:
+        logger.warning("Package nemo_text_processing not found!")
+        logger.warning(
+            "Run: conda install -c conda-forge pynini=2.1.5 && pip install nemo_text_processing",
+        )
+    try:
+        chat.normalizer.register("zh", normalizer_zh_tn())
+    except ValueError as e:
+        logger.error(e)
+    except BaseException:
+        logger.warning("Package WeTextProcessing not found!")
+        logger.warning(
+            "Run: conda install -c conda-forge pynini=2.1.5 && pip install WeTextProcessing",
+        )
+
+
+def main(
+    texts: List[str],
+    spk: Optional[str] = None,
+    stream: bool = False,
+    source: str = "local",
+    custom_path: str = "",
+):
     logger.info("Text input: %s", str(texts))
 
     chat = ChatTTS.Chat(get_logger("ChatTTS"))
     logger.info("Initializing ChatTTS...")
-    if chat.load():
+    load_normalizer(chat)
+
+    is_load = False
+    if os.path.isdir(custom_path) and source == "custom":
+        is_load = chat.load(source="custom", custom_path=custom_path)
+    else:
+        is_load = chat.load(source=source)
+
+    if is_load:
         logger.info("Models loaded successfully.")
     else:
         logger.error("Models load failed.")
@@ -69,10 +107,14 @@ def main(texts: List[str], spk: Optional[str] = None, stream=False):
 
 
 if __name__ == "__main__":
+    r"""
+    python -m examples.cmd.run \
+        --source custom --custom_path ../../models/2Noise/ChatTTS 你好喲 ":)"
+    """
     logger.info("Starting ChatTTS commandline demo...")
     parser = argparse.ArgumentParser(
         description="ChatTTS Command",
-        usage='[--spk xxx] [--stream] "Your text 1." " Your text 2."',
+        usage='[--spk xxx] [--stream] [--source ***] [--custom_path XXX] "Your text 1." " Your text 2."',
     )
     parser.add_argument(
         "--spk",
@@ -85,6 +127,18 @@ if __name__ == "__main__":
         help="Use stream mode",
         action="store_true",
     )
+    parser.add_argument(
+        "--source",
+        help="source form [ huggingface(hf download), local(ckpt save to asset dir), custom(define) ]",
+        type=str,
+        default="local",
+    )
+    parser.add_argument(
+        "--custom_path",
+        help="custom defined model path(include asset ckpt dir)",
+        type=str,
+        default="",
+    )
     parser.add_argument(
         "texts",
         help="Original text",
@@ -92,5 +146,6 @@ if __name__ == "__main__":
         nargs=argparse.REMAINDER,
     )
     args = parser.parse_args()
-    main(args.texts, args.spk, args.stream)
+    logger.info(args)
+    main(args.texts, args.spk, args.stream, args.source, args.custom_path)
     logger.info("ChatTTS process finished.")

examples/cmd/stream.py

@@ -1,4 +1,3 @@
-import time
 import random
 
 import numpy as np
@@ -148,7 +147,6 @@ class ChatStreamer:
     # 流式播放接口
     def play(self, streamchat, wait=5):
         import pyaudio  # please install it manually
-        import time
 
         p = pyaudio.PyAudio()
         print(p.get_device_count())

examples/ipynb/colab.ipynb

@@ -342,6 +342,44 @@
     "Audio(wav[0], rate=24_000, autoplay=True)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Zero shot (simulate speaker)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ChatTTS.tools.audio import load_audio\n",
+    "\n",
+    "spk_smp = chat.sample_audio_speaker(load_audio(\"sample.mp3\", 24000))\n",
+    "print(spk_smp)  # save it in order to load the speaker without sample audio next time\n",
+    "\n",
+    "params_infer_code = ChatTTS.Chat.InferCodeParams(\n",
+    "    spk_smp=spk_smp,\n",
+    "    txt_smp=\"与sample.mp3内容完全一致的文本转写。\",\n",
+    ")\n",
+    "\n",
+    "wav = chat.infer(\n",
+    "    \"四川美食确实以辣闻名，但也有不辣的选择。比如甜水面、赖汤圆、蛋烘糕、叶儿粑等，这些小吃口味温和，甜而不腻，也很受欢迎。\",\n",
+    "    params_infer_code=params_infer_code,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "Audio(wav[0], rate=24_000, autoplay=True)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {

examples/ipynb/example.ipynb

@@ -268,6 +268,44 @@
     "Audio(wav[0], rate=24_000, autoplay=True)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Zero shot (simulate speaker)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from tools.audio import load_audio\n",
+    "\n",
+    "spk_smp = chat.sample_audio_speaker(load_audio(\"sample.mp3\", 24000))\n",
+    "print(spk_smp)  # save it in order to load the speaker without sample audio next time\n",
+    "\n",
+    "params_infer_code = ChatTTS.Chat.InferCodeParams(\n",
+    "    spk_smp=spk_smp,\n",
+    "    txt_smp=\"与sample.mp3内容完全一致的文本转写。\",\n",
+    ")\n",
+    "\n",
+    "wav = chat.infer(\n",
+    "    \"四川美食确实以辣闻名，但也有不辣的选择。比如甜水面、赖汤圆、蛋烘糕、叶儿粑等，这些小吃口味温和，甜而不腻，也很受欢迎。\",\n",
+    "    params_infer_code=params_infer_code,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "Audio(wav[0], rate=24_000, autoplay=True)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

examples/onnx/README.md

@@ -0,0 +1,14 @@
+# Export onnx or JIT models for deployment
+
+## Run `pip install onnx -U`.
+
+## Export GPT
+
+3. Run `python examples/onnx/exporter.py --gpt`
+
+
+## Export other models
+Run `python examples/onnx/exporter.py --decoder --vocos`
+
+## Reference
+[Run LLMs on Sophon TPU](https://github.com/sophgo/LLM-TPU)

examples/onnx/exporter.py

@@ -0,0 +1,426 @@
+import os, sys
+
+if sys.platform == "darwin":
+    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+
+from dataclasses import asdict
+import argparse
+import torch
+from tqdm import tqdm
+from ChatTTS.model.dvae import DVAE
+from ChatTTS.config import Config
+from vocos import Vocos
+from vocos.pretrained import instantiate_class
+import torch.jit as jit
+
+from gpt import GPT
+
+# disable cuda
+torch.cuda.is_available = lambda: False
+
+# add args to control which modules to export
+parser = argparse.ArgumentParser()
+parser.add_argument("--gpt", action="store_true", help="trace gpt")
+parser.add_argument("--decoder", action="store_true", help="trace decoder")
+parser.add_argument("--vocos", action="store_true", help="trace vocos")
+parser.add_argument(
+    "--pth_dir", default="./assets", type=str, help="path to the pth model directory"
+)
+parser.add_argument(
+    "--out_dir", default="./tmp", type=str, help="path to output directory"
+)
+
+args = parser.parse_args()
+chattts_config = Config()
+
+
+def export_gpt():
+    gpt_model = GPT(gpt_config=asdict(chattts_config.gpt), use_flash_attn=False).eval()
+    gpt_model.from_pretrained(asdict(chattts_config.path)["gpt_ckpt_path"])
+    gpt_model = gpt_model.eval()
+    for param in gpt_model.parameters():
+        param.requires_grad = False
+
+    config = gpt_model.gpt.config
+    layers = gpt_model.gpt.layers
+    model_norm = gpt_model.gpt.norm
+
+    NUM_OF_LAYERS = config.num_hidden_layers
+    HIDDEN_SIZE = config.hidden_size
+    NUM_ATTENTION_HEADS = config.num_attention_heads
+    NUM_KEY_VALUE_HEADS = config.num_key_value_heads
+    HEAD_DIM = HIDDEN_SIZE // NUM_ATTENTION_HEADS  # 64
+    TEXT_VOCAB_SIZE = gpt_model.emb_text.weight.shape[0]
+    AUDIO_VOCAB_SIZE = gpt_model.emb_code[0].weight.shape[0]
+    SEQ_LENGTH = 512
+
+    folder = os.path.join(args.out_dir, "gpt")
+    os.makedirs(folder, exist_ok=True)
+
+    for param in gpt_model.emb_text.parameters():
+        param.requires_grad = False
+
+    for param in gpt_model.emb_code.parameters():
+        param.requires_grad = False
+
+    for param in gpt_model.head_code.parameters():
+        param.requires_grad = False
+
+    for param in gpt_model.head_text.parameters():
+        param.requires_grad = False
+
+    class EmbeddingText(torch.nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__(*args, **kwargs)
+
+        def forward(self, input_ids):
+            return gpt_model.emb_text(input_ids)
+
+    def convert_embedding_text():
+        model = EmbeddingText()
+        input_ids = torch.tensor([range(SEQ_LENGTH)])
+
+        torch.onnx.export(
+            model,
+            (input_ids),
+            f"{folder}/embedding_text.onnx",
+            verbose=False,
+            input_names=["input_ids"],
+            output_names=["input_embed"],
+            do_constant_folding=True,
+            opset_version=15,
+        )
+
+    class EmbeddingCode(torch.nn.Module):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__(*args, **kwargs)
+
+        def forward(self, input_ids):
+            input_ids = input_ids.unsqueeze(2).expand(
+                -1, -1, gpt_model.num_vq
+            )  # for forward_first_code
+            code_emb = [
+                gpt_model.emb_code[i](input_ids[:, :, i])
+                for i in range(gpt_model.num_vq)
+            ]
+            return torch.stack(code_emb, 2).sum(2)
+
+    def convert_embedding_code():
+        model = EmbeddingCode()
+        input_ids = torch.tensor([range(SEQ_LENGTH)])
+
+        torch.onnx.export(
+            model,
+            (input_ids),
+            f"{folder}/embedding_code.onnx",
+            verbose=False,
+            input_names=["input_ids"],
+            output_names=["input_embed"],
+            do_constant_folding=True,
+            opset_version=15,
+        )
+
+    class EmbeddingCodeCache(torch.nn.Module):  # for forward_next_code
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__(*args, **kwargs)
+
+        def forward(self, input_ids):
+            code_emb = [
+                gpt_model.emb_code[i](input_ids[:, :, i])
+                for i in range(gpt_model.num_vq)
+            ]
+            return torch.stack(code_emb, 2).sum(2)
+
+    def convert_embedding_code_cache():
+        model = EmbeddingCodeCache()
+        input_ids = torch.tensor(
+            [[[416, 290, 166, 212]]]
+        )  # torch.tensor([[range(gpt_model.num_vq)]])
+        torch.onnx.export(
+            model,
+            (input_ids),
+            f"{folder}/embedding_code_cache.onnx",
+            verbose=False,
+            input_names=["input_ids"],
+            output_names=["input_embed"],
+            do_constant_folding=True,
+            opset_version=15,
+        )
+
+    class Block(torch.nn.Module):
+        def __init__(self, layer_id):
+            super().__init__()
+            self.layer_id = layer_id
+            self.layer = layers[layer_id]  # LlamaDecoderLayer
+            self.norm = model_norm
+
+        def forward(self, hidden_states, position_ids, attention_mask):
+            hidden_states, past_kv = self.layer(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                use_cache=True,
+            )
+            present_k, present_v = past_kv
+            if self.layer_id == NUM_OF_LAYERS - 1:
+                hidden_states = self.norm(hidden_states)
+            return hidden_states, present_k, present_v
+
+    def convert_block(layer_id):
+        model = Block(layer_id)
+        hidden_states = torch.randn((1, SEQ_LENGTH, HIDDEN_SIZE))
+        position_ids = torch.tensor([range(SEQ_LENGTH)], dtype=torch.long)
+        attention_mask = -1000 * torch.ones(
+            (1, 1, SEQ_LENGTH, SEQ_LENGTH), dtype=torch.float32
+        ).triu(diagonal=1)
+        model(hidden_states, position_ids, attention_mask)
+        torch.onnx.export(
+            model,
+            (hidden_states, position_ids, attention_mask),
+            f"{folder}/block_{layer_id}.onnx",
+            verbose=False,
+            input_names=["input_states", "position_ids", "attention_mask"],
+            output_names=["hidden_states", "past_k", "past_v"],
+            do_constant_folding=True,
+            opset_version=15,
+        )
+
+    class BlockCache(torch.nn.Module):
+
+        def __init__(self, layer_id):
+            super().__init__()
+            self.layer_id = layer_id
+            self.layer = layers[layer_id]
+            self.norm = model_norm
+
+        def forward(self, hidden_states, position_ids, attention_mask, past_k, past_v):
+            hidden_states, past_kv = self.layer(
+                hidden_states,
+                attention_mask,
+                position_ids=position_ids,
+                past_key_value=(past_k, past_v),
+                use_cache=True,
+            )
+            present_k, present_v = past_kv
+            if self.layer_id == NUM_OF_LAYERS - 1:
+                hidden_states = self.norm(hidden_states)
+            return hidden_states, present_k, present_v
+
+    def convert_block_cache(layer_id):
+        model = BlockCache(layer_id)
+        hidden_states = torch.randn((1, 1, HIDDEN_SIZE))
+        position_ids = torch.tensor([range(1)], dtype=torch.long)
+        attention_mask = -1000 * torch.ones(
+            (1, 1, 1, SEQ_LENGTH + 1), dtype=torch.float32
+        ).triu(diagonal=1)
+        past_k = torch.randn((1, SEQ_LENGTH, NUM_ATTENTION_HEADS, HEAD_DIM))
+        past_v = torch.randn((1, SEQ_LENGTH, NUM_ATTENTION_HEADS, HEAD_DIM))
+
+        torch.onnx.export(
+            model,
+            (hidden_states, position_ids, attention_mask, past_k, past_v),
+            f"{folder}/block_cache_{layer_id}.onnx",
+            verbose=False,
+            input_names=[
+                "input_states",
+                "position_ids",
+                "attention_mask",
+                "history_k",
+                "history_v",
+            ],
+            output_names=["hidden_states", "past_k", "past_v"],
+            do_constant_folding=True,
+            opset_version=15,
+        )
+
+    class GreedyHead(torch.nn.Module):
+
+        def __init__(self):
+            super().__init__()
+
+        def forward(self, m_logits):
+            _, token = torch.topk(m_logits.float(), 1)
+            return token
+
+    def convert_greedy_head_text():
+        model = GreedyHead()
+        m_logits = torch.randn(1, TEXT_VOCAB_SIZE)
+
+        torch.onnx.export(
+            model,
+            (m_logits),
+            f"{folder}/greedy_head_text.onnx",
+            verbose=False,
+            input_names=["m_logits"],
+            output_names=["token"],
+            do_constant_folding=True,
+            opset_version=15,
+        )
+
+    def convert_greedy_head_code():
+        model = GreedyHead()
+        m_logits = torch.randn(1, AUDIO_VOCAB_SIZE, gpt_model.num_vq)
+
+        torch.onnx.export(
+            model,
+            (m_logits),
+            f"{folder}/greedy_head_code.onnx",
+            verbose=False,
+            input_names=["m_logits"],
+            output_names=["token"],
+            do_constant_folding=True,
+            opset_version=15,
+        )
+
+    class LmHead_infer_text(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+
+        def forward(self, hidden_states):
+            m_logits = gpt_model.head_text(hidden_states)
+            return m_logits
+
+    class LmHead_infer_code(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+
+        def forward(self, hidden_states):
+            m_logits = torch.stack(
+                [
+                    gpt_model.head_code[i](hidden_states)
+                    for i in range(gpt_model.num_vq)
+                ],
+                2,
+            )
+            return m_logits
+
+    def convert_lm_head_text():
+        model = LmHead_infer_text()
+        input = torch.randn(1, HIDDEN_SIZE)
+
+        torch.onnx.export(
+            model,
+            (input),
+            f"{folder}/lm_head_text.onnx",
+            verbose=False,
+            input_names=["hidden_states"],
+            output_names=["m_logits"],
+            do_constant_folding=True,
+            opset_version=15,
+        )
+
+    def convert_lm_head_code():
+        model = LmHead_infer_code()
+        input = torch.randn(1, HIDDEN_SIZE)
+        torch.onnx.export(
+            model,
+            (input),
+            f"{folder}/lm_head_code.onnx",
+            verbose=False,
+            input_names=["hidden_states"],
+            output_names=["m_logits"],
+            do_constant_folding=True,
+            opset_version=15,
+        )
+
+    # export models
+    print(f"Convert block & block_cache")
+    for i in tqdm(range(NUM_OF_LAYERS)):
+        convert_block(i)
+        convert_block_cache(i)
+
+    print(f"Convert embedding")
+    convert_embedding_text()
+    convert_embedding_code()
+    convert_embedding_code_cache()
+
+    print(f"Convert lm_head")
+    convert_lm_head_code()
+    convert_lm_head_text()
+
+    print(f"Convert greedy_head")
+    convert_greedy_head_text()
+    convert_greedy_head_code()
+
+
+def export_decoder():
+    decoder = DVAE(
+        decoder_config=asdict(chattts_config.decoder),
+        dim=chattts_config.decoder.idim,
+    ).eval()
+    decoder.load_state_dict(
+        torch.load(
+            asdict(chattts_config.path)["decoder_ckpt_path"],
+            weights_only=True,
+            mmap=True,
+        )
+    )
+
+    for param in decoder.parameters():
+        param.requires_grad = False
+    rand_input = torch.rand([1, 768, 1024], requires_grad=False)
+
+    def mydec(_inp):
+        return decoder(_inp, mode="decode")
+
+    jitmodel = jit.trace(mydec, [rand_input])
+    jit.save(jitmodel, f"{args.out_dir}/decoder_jit.pt")
+
+
+def export_vocos():
+    feature_extractor = instantiate_class(
+        args=(), init=asdict(chattts_config.vocos.feature_extractor)
+    )
+    backbone = instantiate_class(args=(), init=asdict(chattts_config.vocos.backbone))
+    head = instantiate_class(args=(), init=asdict(chattts_config.vocos.head))
+    vocos = Vocos(
+        feature_extractor=feature_extractor, backbone=backbone, head=head
+    ).eval()
+    vocos.load_state_dict(
+        torch.load(
+            asdict(chattts_config.path)["vocos_ckpt_path"], weights_only=True, mmap=True
+        )
+    )
+
+    for param in vocos.parameters():
+        param.requires_grad = False
+    rand_input = torch.rand([1, 100, 2048], requires_grad=False)
+
+    def myvocos(_inp):
+        # return chat.vocos.decode(_inp) # TPU cannot support the istft OP, thus it has to be moved to postprocessing
+        # reference: https://github.com/gemelo-ai/vocos.git
+        x = vocos.backbone(_inp)
+        x = vocos.head.out(x).transpose(1, 2)
+        mag, p = x.chunk(2, dim=1)
+        mag = torch.exp(mag)
+        mag = torch.clip(
+            mag, max=1e2
+        )  # safeguard to prevent excessively large magnitudes
+        # wrapping happens here. These two lines produce real and imaginary value
+        x = torch.cos(p)
+        y = torch.sin(p)
+        return mag, x, y
+
+    jitmodel = jit.trace(myvocos, [rand_input])
+    torch.onnx.export(
+        jitmodel,
+        [rand_input],
+        f"{args.out_dir}/vocos_1-100-2048.onnx",
+        opset_version=12,
+        do_constant_folding=True,
+    )
+
+
+if args.gpt:
+    export_gpt()
+
+if args.decoder:
+    export_decoder()
+
+if args.vocos:
+    export_vocos()
+
+print("Done. Please check the files in", args.out_dir)

examples/onnx/gpt.py

@@ -0,0 +1,89 @@
+import logging
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+from torch.nn.utils.parametrizations import weight_norm
+
+from modeling_llama import LlamaModel, LlamaConfig
+
+
+class GPT(nn.Module):
+    def __init__(
+        self,
+        gpt_config: dict,
+        num_audio_tokens: int = 626,
+        num_text_tokens: int = 21178,
+        num_vq=4,
+        use_flash_attn=False,
+        device=torch.device("cpu"),
+        logger=logging.getLogger(__name__),
+    ):
+        super().__init__()
+
+        self.logger = logger
+
+        self.device = device
+        self.device_gpt = device if "mps" not in str(device) else torch.device("cpu")
+
+        self.num_vq = num_vq
+        self.num_audio_tokens = num_audio_tokens
+
+        self.use_flash_attn = use_flash_attn
+
+        self.gpt, self.llama_config = self._build_llama(gpt_config, self.device_gpt)
+        self.is_te_llama = False
+        self.model_dim = int(self.gpt.config.hidden_size)
+        self.emb_code = nn.ModuleList(
+            [
+                nn.Embedding(
+                    num_audio_tokens,
+                    self.model_dim,
+                    device=self.device_gpt,
+                )
+                for _ in range(num_vq)
+            ],
+        )
+        self.emb_text = nn.Embedding(
+            num_text_tokens, self.model_dim, device=self.device_gpt
+        )
+
+        self.head_text = weight_norm(
+            nn.Linear(
+                self.model_dim,
+                num_text_tokens,
+                bias=False,
+                device=device,
+            ),
+            name="weight",
+        )
+        self.head_code = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Linear(
+                        self.model_dim,
+                        num_audio_tokens,
+                        bias=False,
+                        device=device,
+                    ),
+                    name="weight",
+                )
+                for _ in range(self.num_vq)
+            ],
+        )
+
+    def from_pretrained(self, file_path: str):
+        self.load_state_dict(
+            torch.load(file_path, weights_only=True, mmap=True), strict=False
+        )
+
+    def _build_llama(
+        self,
+        config: dict,
+        device: torch.device,
+    ) -> Tuple[LlamaModel, LlamaConfig]:
+
+        llama_config = LlamaConfig(**config)
+        model = LlamaModel(llama_config)
+        del model.embed_tokens
+        return model.to(device), llama_config

examples/web/funcs.py

@@ -1,11 +1,14 @@
-import sys
 import random
 from typing import Optional
 from time import sleep
 
 import gradio as gr
 
-from tools.audio import float_to_int16, has_ffmpeg_installed
+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
 
 logger = get_logger(" WebUI ")
@@ -22,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
 
@@ -59,13 +65,26 @@ 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, compile=sys.platform != "win32")
+        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, compile=sys.platform != "win32"
+            "custom",
+            custom_path=cust_path,
+            coef=coef,
+            enable_cache=enable_cache,
+            experimental=experimental,
         )
         global custom_path
         custom_path = cust_path
@@ -101,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
@@ -110,10 +129,19 @@ 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
 
 
+def on_upload_sample_audio(sample_audio_input: Optional[str]) -> str:
+    if sample_audio_input is None:
+        return ""
+    sample_audio = load_audio(sample_audio_input, 24000)
+    spk_smp = chat.sample_audio_speaker(sample_audio)
+    del sample_audio
+    return spk_smp
+
+
 def _set_generate_buttons(generate_button, interrupt_button, is_reset=False):
     return gr.update(
         value=generate_button, visible=is_reset, interactive=is_reset
@@ -124,6 +152,10 @@ def refine_text(
     text,
     text_seed_input,
     refine_text_flag,
+    temperature,
+    top_P,
+    top_K,
+    split_batch,
 ):
     global chat
 
@@ -131,18 +163,33 @@ def refine_text(
         sleep(1)  # to skip fast answer of loading mark
         return text
 
-    with TorchSeedContext(text_seed_input):
-        text = chat.infer(
-            text,
-            skip_refine_text=False,
-            refine_text_only=True,
-        )
+    text = chat.infer(
+        text,
+        skip_refine_text=False,
+        refine_text_only=True,
+        params_refine_text=ChatTTS.Chat.RefineTextParams(
+            temperature=temperature,
+            top_P=top_P,
+            top_K=top_K,
+            manual_seed=text_seed_input,
+        ),
+        split_text=split_batch > 0,
+    )
 
     return text[0] if isinstance(text, list) else text
 
 
 def generate_audio(
-    text, temperature, top_P, top_K, spk_emb_text: str, stream, audio_seed_input
+    text,
+    temperature,
+    top_P,
+    top_K,
+    spk_emb_text: str,
+    stream,
+    audio_seed_input,
+    sample_text_input,
+    sample_audio_code_input,
+    split_batch,
 ):
     global chat, has_interrupted
 
@@ -154,23 +201,30 @@ def generate_audio(
         temperature=temperature,
         top_P=top_P,
         top_K=top_K,
+        manual_seed=audio_seed_input,
     )
 
-    with TorchSeedContext(audio_seed_input):
-        wav = chat.infer(
-            text,
-            skip_refine_text=True,
-            params_infer_code=params_infer_code,
-            stream=stream,
-        )
-        if stream:
-            for gen in wav:
-                audio = gen[0]
-                if audio is not None and len(audio) > 0:
-                    yield 24000, float_to_int16(audio).T
-                del audio
-        else:
-            yield 24000, float_to_int16(wav[0]).T
+    if sample_text_input and sample_audio_code_input:
+        params_infer_code.txt_smp = sample_text_input
+        params_infer_code.spk_smp = sample_audio_code_input
+        params_infer_code.spk_emb = None
+
+    wav = chat.infer(
+        text,
+        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:
+            audio = gen[0]
+            if audio is not None and len(audio) > 0:
+                yield 24000, float_to_int16(audio).T
+            del audio
+    else:
+        yield 24000, float_to_int16(wav[0]).T
 
 
 def interrupt_generate():

examples/web/webui.py

@@ -10,8 +10,8 @@ import argparse
 
 import gradio as gr
 
-from examples.web.funcs import *
-from examples.web.ex import ex
+from funcs import *
+from ex import ex
 
 
 def main():
@@ -21,13 +21,43 @@ def main():
         gr.Markdown("- **GitHub Repo**: https://github.com/2noise/ChatTTS")
         gr.Markdown("- **HuggingFace Repo**: https://huggingface.co/2Noise/ChatTTS")
 
-        text_input = gr.Textbox(
-            label="Input Text",
-            lines=4,
-            placeholder="Please Input Text...",
-            value=ex[0][0],
-            interactive=True,
-        )
+        with gr.Row():
+            with gr.Column(scale=2):
+                text_input = gr.Textbox(
+                    label="Input Text",
+                    lines=4,
+                    max_lines=4,
+                    placeholder="Please Input Text...",
+                    value=ex[0][0],
+                    interactive=True,
+                )
+                sample_text_input = gr.Textbox(
+                    label="Sample Text",
+                    lines=4,
+                    max_lines=4,
+                    placeholder="If Sample Audio and Sample Text are available, the Speaker Embedding will be disabled.",
+                    interactive=True,
+                )
+            with gr.Column():
+                with gr.Tab(label="Sample Audio"):
+                    sample_audio_input = gr.Audio(
+                        value=None,
+                        type="filepath",
+                        interactive=True,
+                        show_label=False,
+                        waveform_options=gr.WaveformOptions(
+                            sample_rate=24000,
+                        ),
+                        scale=1,
+                    )
+                with gr.Tab(label="Sample Audio Code"):
+                    sample_audio_code_input = gr.Textbox(
+                        lines=12,
+                        max_lines=12,
+                        show_label=False,
+                        placeholder="Paste the Code copied before after uploading Sample Audio.",
+                        interactive=True,
+                    )
 
         with gr.Row():
             refine_text_checkbox = gr.Checkbox(
@@ -72,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",
@@ -80,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,
             )
@@ -109,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
             )
@@ -123,9 +161,15 @@ def main():
         text_output = gr.Textbox(
             label="Output Text",
             interactive=False,
-            show_copy_button=True,
+            buttons=["copy"],
         )
 
+        sample_audio_input.change(
+            fn=on_upload_sample_audio,
+            inputs=sample_audio_input,
+            outputs=sample_audio_code_input,
+        ).then(fn=lambda: gr.Info("Sampled Audio Code generated at another Tab."))
+
         # 使用Gradio的回调功能来更新数值输入框
         voice_selection.change(
             fn=on_voice_change, inputs=voice_selection, outputs=audio_seed_input
@@ -169,6 +213,10 @@ def main():
                     text_input,
                     text_seed_input,
                     refine_text_checkbox,
+                    temperature_slider,
+                    top_p_slider,
+                    top_k_slider,
+                    split_batch_slider,
                 ],
                 outputs=text_output,
             ).then(
@@ -181,6 +229,9 @@ def main():
                     spk_emb_text,
                     stream_mode_checkbox,
                     audio_seed_input,
+                    sample_text_input,
+                    sample_audio_code_input,
+                    split_batch_slider,
                 ],
                 outputs=audio_output,
             ).then(
@@ -207,15 +258,17 @@ def main():
         "--server_name", type=str, default="0.0.0.0", help="server name"
     )
     parser.add_argument("--server_port", type=int, default=8080, help="server port")
-    parser.add_argument("--root_path", type=str, default=None, help="root path")
+    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(
-        "--custom_path", type=str, default=None, help="custom model path"
+        "--disable_cache", action="store_true", help="enable model cache"
     )
     parser.add_argument(
-        "--coef", type=str, default=None, help="custom dvae coefficient"
+        "--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):
@@ -232,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,6 +1,5 @@
-numpy<2.0.0
+numpy<3.0.0
 numba
-omegaconf>=2.3.0
 torch>=2.1.0
 torchaudio
 tqdm
@@ -15,3 +14,4 @@ WeTextProcessing; sys_platform == 'linux'
 nemo_text_processing; sys_platform == 'linux'
 av
 pydub
+requests

setup.py

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

tests/#511.py

@@ -12,7 +12,7 @@ import ChatTTS
 
 from tools.logger import get_logger
 
-logger = get_logger("Test #511", lv=logging.WARN)
+logger = get_logger("Test", lv=logging.WARN)
 
 chat = ChatTTS.Chat(logger)
 chat.load(compile=False, source="huggingface")  # Set to True for better performance
@@ -41,6 +41,7 @@ fail = False
 wavs = chat.infer(
     texts,
     skip_refine_text=True,
+    split_text=False,
     params_infer_code=params_infer_code,
 )
 

tests/#521.py

@@ -1,221 +0,0 @@
-import os, sys
-
-if sys.platform == "darwin":
-    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
-
-now_dir = os.getcwd()
-sys.path.append(now_dir)
-
-import logging
-import threading
-import time
-import random
-
-import numpy as np
-import ChatTTS
-
-from tools.audio import float_to_int16
-from tools.logger import get_logger
-
-fail = False
-logger = get_logger("Test #521", lv=logging.WARN)
-
-
-# 计算rms
-# nan为噪声 ！！！
-def calculate_rms(data):
-    m = np.mean(np.square(data.astype(np.int32)))
-    if m < 0:
-        logger.warning("neg RM: %f", m)
-    else:
-        logger.info("RM: %f", m)
-    return np.sqrt(m)
-
-
-# 流式声音处理器
-class AudioStreamer:
-    # 流式写入
-    @staticmethod
-    def write(waveform: np.ndarray):
-        global fail, logger
-        rms = calculate_rms(waveform)
-        if np.isnan(rms):
-            fail = True
-            logger.warning("NAN RMS found.")
-
-
-# ChatTTS流式处理
-class ChatStreamer:
-    def __init__(self, waittime_topause=50, base_block_size=8000):
-        self.streamer = AudioStreamer
-        self.accum_streamwavs = []
-        self.waittime_topause = waittime_topause
-        self.base_block_size = base_block_size
-
-    def write(self, chatstream):
-        # 已推理batch数据保存
-        def accum(accum_wavs, stream_wav):
-            n_texts = len(stream_wav)
-            if accum_wavs is None:
-                accum_wavs = [[i] for i in stream_wav]
-            else:
-                for i_text in range(n_texts):
-                    if stream_wav[i_text] is not None:
-                        accum_wavs[i_text].append(stream_wav[i_text])
-            return accum_wavs
-
-        # stream状态更新。数据量不足的stream，先存一段时间，直到拿到足够数据，监控小块数据情况
-        def update_stream(history_stream_wav, new_stream_wav, thre):
-            result_stream = []
-            randn = -1
-            if history_stream_wav is not None:
-                randn = random.random()
-                if randn > 0.1:
-                    logger.info("update_stream")
-                n_texts = len(new_stream_wav)
-                for i in range(n_texts):
-                    if new_stream_wav[i] is not None:
-                        result_stream.append(
-                            np.concatenate(
-                                [history_stream_wav[i], new_stream_wav[i]], axis=1
-                            )
-                        )
-                    else:
-                        result_stream.append(history_stream_wav[i])
-            else:
-                result_stream = [i[np.newaxis, :] for i in new_stream_wav]
-            is_keep_next = (
-                sum([i.shape[1] for i in result_stream if i is not None]) < thre
-            )
-            if randn > 0.1:
-                logger.info(
-                    "result_stream: %s %s",
-                    str(is_keep_next),
-                    str([i.shape if i is not None else None for i in result_stream]),
-                )
-            return result_stream, is_keep_next
-
-        self.finish = False
-        curr_sentence_index = 0
-        base_block_size = self.base_block_size
-        history_stream_wav = None
-        article_streamwavs = None
-        for stream_wav in chatstream:
-            n_texts = len(stream_wav)
-            n_valid_texts = len(list(filter(lambda x: x is not None, stream_wav)))
-            if n_valid_texts == 0:
-                continue
-            else:
-                block_thre = n_valid_texts * base_block_size
-                stream_wav, is_keep_next = update_stream(
-                    history_stream_wav, stream_wav, block_thre
-                )
-                # 数据量不足，先保存状态
-                if is_keep_next:
-                    history_stream_wav = stream_wav
-                    continue
-                # 数据量足够，执行写入操作
-                else:
-                    history_stream_wav = None
-                    stream_wav = [float_to_int16(i) for i in stream_wav]
-                    article_streamwavs = accum(article_streamwavs, stream_wav)
-                    # 写入当前句子
-                    if stream_wav[curr_sentence_index] is not None:
-                        if stream_wav[curr_sentence_index][0].shape[0] > 257:
-                            self.streamer.write(stream_wav[curr_sentence_index][0])
-                        # self.streamer.write(stream_wav[curr_sentence_index][0])
-                    # 当前句子已写入完成，直接写下一个句子已经推理完成的部分
-                    elif curr_sentence_index < n_texts - 1:
-                        curr_sentence_index += 1
-                        logger.info("add next sentence")
-                        finish_stream_wavs = np.concatenate(
-                            article_streamwavs[curr_sentence_index], axis=1
-                        )
-                        if finish_stream_wavs[0].shape[0] > 257:
-                            self.streamer.write(finish_stream_wavs[0])
-                        # self.streamer.write(finish_stream_wavs[0])
-                    # streamchat遍历完毕，在外层把剩余结果写入
-                    else:
-                        break
-            # 有一定概率遇到奇怪bug（一定概率遇到256维异常输出，正常是1w+维），输出全是噪声，写的快遇到的概率更高？
-            time.sleep(0.02)
-        # 本轮剩余最后一点数据写入
-        if is_keep_next:
-            if len(list(filter(lambda x: x is not None, stream_wav))) > 0:
-                stream_wav = [float_to_int16(i) for i in stream_wav]
-                if stream_wav[curr_sentence_index] is not None:
-                    if stream_wav[curr_sentence_index][0].shape[0] > 257:
-                        self.streamer.write(stream_wav[curr_sentence_index][0])
-                        article_streamwavs = accum(article_streamwavs, stream_wav)
-        # 把已经完成推理的下几轮剩余数据写入
-        for i_text in range(curr_sentence_index + 1, n_texts):
-            finish_stream_wavs = np.concatenate(article_streamwavs[i_text], axis=1)
-            if finish_stream_wavs[0].shape[0] > 257:
-                self.streamer.write(finish_stream_wavs[0])
-        self.accum_streamwavs.append(article_streamwavs)
-        self.finish = True
-
-    # 获取完整历史播放数据
-    def get_complete_speech(self):
-        complete_waveform = np.concatenate(
-            sum([sum(i_infer, []) for i_infer in self.accum_streamwavs], []), axis=1
-        )
-        return complete_waveform
-
-    # 开始音频写入。可支持多次音频写入
-    def start_writing(self, streamchat):
-        self.writer = threading.Thread(target=self.write, args=(streamchat,))
-        self.writer.start()
-
-
-chat = ChatTTS.Chat(logger)
-chat.load(compile=False, source="huggingface")  # Set to True for better performance
-
-rand_spk = chat.sample_random_speaker()
-params_infer_code = ChatTTS.Chat.InferCodeParams(
-    spk_emb=rand_spk,  # add sampled speaker
-    temperature=0.0001,  # using custom temperature
-    prompt="[speed_0]",
-    show_tqdm=False,
-)
-params_refine_text = ChatTTS.Chat.RefineTextParams(
-    prompt="[oral_2][laugh_0][break_6]",
-    show_tqdm=False,
-)
-
-# 获取ChatTTS 流式推理generator
-streamchat = chat.infer(
-    "EvilBOT : 在黑暗的角落里，杀人可能被视为一种艺术，但我要提醒你，这违背了我作为邪恶机器人的本性，因为我是无所不能的。",
-    skip_refine_text=True,
-    params_infer_code=params_infer_code,
-    stream=True,
-    params_refine_text=params_refine_text,
-)
-
-# 分别开启一个写线程和读线程，进行流式播放
-streamer = ChatStreamer()
-# 一次性生成
-streamer.write(streamchat)
-
-streamer.write(
-    chat.infer(
-        "有一个神奇的故事，传说在很远很远的未来。",
-        skip_refine_text=True,
-        params_infer_code=params_infer_code,
-        stream=True,
-    )
-)
-
-streamer.write(
-    chat.infer(
-        "有一种叫做奥特曼的物种。他是超人族的一员。",
-        skip_refine_text=True,
-        params_infer_code=params_infer_code,
-        stream=True,
-    )
-)
-
-if fail:
-    import sys
-
-    sys.exit(1)

tests/#588.py

@@ -0,0 +1,52 @@
+import os, sys
+
+if sys.platform == "darwin":
+    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+
+import logging
+import re
+
+import ChatTTS
+
+from tools.logger import get_logger
+
+logger = get_logger("Test", lv=logging.WARN)
+
+chat = ChatTTS.Chat(logger)
+chat.load(compile=False, source="huggingface")  # Set to True for better performance
+
+texts = [
+    "总结一下，AI Agent是大模型功能的扩展，让AI更接近于通用人工智能，也就是我们常说的AGI。",
+    "你真是太聪明啦。",
+]
+
+fail = False
+
+refined = chat.infer(
+    texts,
+    refine_text_only=True,
+    stream=False,
+    split_text=False,
+    params_refine_text=ChatTTS.Chat.RefineTextParams(show_tqdm=False),
+)
+
+trimre = re.compile("\\[[\w_]+\\]")
+
+
+def trim_tags(txt: str) -> str:
+    global trimre
+    return trimre.sub("", txt)
+
+
+for i, t in enumerate(refined):
+    if len(trim_tags(t)) > 4 * len(texts[i]):
+        fail = True
+        logger.warning("in: %s, out: %s", texts[i], t)
+
+if fail:
+    import sys
+
+    sys.exit(1)

tests/#655.py

@@ -0,0 +1,97 @@
+import os, sys
+
+if sys.platform == "darwin":
+    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+
+import logging
+
+import torch
+
+import ChatTTS
+
+from tools.logger import get_logger
+from tools.normalizer import normalizer_en_nemo_text
+
+logger = get_logger("Test", lv=logging.WARN)
+
+chat = ChatTTS.Chat(logger)
+chat.load(compile=False, source="huggingface")  # Set to True for better performance
+try:
+    chat.normalizer.register("en", normalizer_en_nemo_text())
+except:
+    logger.warning("Package nemo_text_processing not found!")
+
+rand_spk = chat.sample_random_speaker()
+
+
+text = ["What is [uv_break]your favorite english food?[laugh][lbreak]"]
+
+fail = False
+
+refined_text = chat.infer(
+    text,
+    refine_text_only=True,
+    params_refine_text=ChatTTS.Chat.RefineTextParams(
+        prompt="[oral_2][laugh_0][break_6]",
+        manual_seed=12345,
+    ),
+    split_text=False,
+)
+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])
+
+params = ChatTTS.Chat.InferCodeParams(
+    spk_emb=rand_spk,  # add sampled speaker
+    temperature=0.3,  # using custom temperature
+    top_P=0.7,  # top P decode
+    top_K=20,  # top K decode
+)
+input_ids, attention_mask, text_mask = chat.tokenizer.encode(
+    chat.speaker.decorate_code_prompts(
+        text,
+        params.prompt,
+        params.txt_smp,
+        params.spk_emb,
+    ),
+    chat.config.gpt.num_vq,
+    prompt=(
+        chat.speaker.decode_prompt(params.spk_smp)
+        if params.spk_smp is not None
+        else None
+    ),
+    device=chat.device_gpt,
+)
+with torch.inference_mode():
+    start_idx, end_idx = 0, torch.zeros(
+        input_ids.shape[0], device=input_ids.device, dtype=torch.long
+    ).fill_(input_ids.shape[1])
+
+    recoded_text = chat.tokenizer.decode(
+        chat.gpt._prepare_generation_outputs(
+            input_ids,
+            start_idx,
+            end_idx,
+            [],
+            [],
+            True,
+        ).ids
+    )
+
+if (
+    recoded_text[0]
+    != "[Stts] [spk_emb] [speed_5] what is [uv_break] your favorite english food? [laugh] [lbreak] [Ptts]"
+):
+    fail = True
+    logger.warning("recoded text is '%s'", refined_text)
+
+if fail:
+    import sys
+
+    sys.exit(1)

tools/audio/__init__.py

@@ -1,3 +1,4 @@
-from .pcm import pcm_arr_to_mp3_view
+from .av import load_audio
+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

@@ -1,8 +1,11 @@
 from io import BufferedWriter, BytesIO
-from typing import Dict
+from pathlib import Path
+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
 
 video_format_dict: Dict[str, str] = {
     "m4a": "mp4",
@@ -34,3 +37,90 @@ def wav2(i: BytesIO, o: BufferedWriter, format: str):
 
     out.close()
     inp.close()
+
+
+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 (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
+
+    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) 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
+
+    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()
+                rate = resampled_frame.rate
+                frames_data.append(frame_data)
+        return (rate, frames_data)
+
+    def frame_iter(container):
+        for p in container.demux(container.streams.audio[0]):
+            yield p.decode()
+
+    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])
+
+            # 检查 decoded_audio 是否有足够的空间，并在必要时调整大小
+            if end_index > decoded_audio.shape[1]:
+                decoded_audio = np.resize(
+                    decoded_audio, (decoded_audio.shape[0], end_index * 4)
+                )
+
+            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

tools/audio/np.py

@@ -4,7 +4,7 @@ import numpy as np
 from numba import jit
 
 
-@jit
+@jit(nopython=True)
 def float_to_int16(audio: np.ndarray) -> np.ndarray:
     am = int(math.ceil(float(np.abs(audio).max())) * 32768)
     am = 32767 * 32768 // am

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())

tools/checksum/tmpl.go

@@ -1,20 +1,30 @@
 package main
 
 var files = [...]string{
-	"asset/Decoder.pt",
-	"asset/DVAE_full.pt",
-	"asset/GPT.pt",
-	"asset/spk_stat.pt",
-	"asset/tokenizer.pt",
-	"asset/Vocos.pt",
+	"asset/Decoder.safetensors",
+	"asset/DVAE.safetensors",
+	"asset/Embed.safetensors",
+	"asset/Vocos.safetensors",
+
+	"asset/gpt/config.json",
+	"asset/gpt/model.safetensors",
+
+	"asset/tokenizer/special_tokens_map.json",
+	"asset/tokenizer/tokenizer_config.json",
+	"asset/tokenizer/tokenizer.json",
 }
 
 const jsontmpl = `{
-	"sha256_asset_Decoder_pt"   : "%s",
-	"sha256_asset_DVAE_full_pt"      : "%s",
-	"sha256_asset_GPT_pt"       : "%s",
-	"sha256_asset_spk_stat_pt"  : "%s",
-	"sha256_asset_tokenizer_pt" : "%s",
-	"sha256_asset_Vocos_pt"     : "%s"
+	"sha256_asset_Decoder_safetensors": "%s",
+	"sha256_asset_DVAE_safetensors"   : "%s",
+	"sha256_asset_Embed_safetensors"  : "%s",
+	"sha256_asset_Vocos_safetensors"  : "%s",
+
+	"sha256_asset_gpt_config_json"         : "%s",
+	"sha256_asset_gpt_model_safetensors"   : "%s",
+
+	"sha256_asset_tokenizer_special_tokens_map_json": "%s",
+	"sha256_asset_tokenizer_tokenizer_config_json"  : "%s",
+	"sha256_asset_tokenizer_tokenizer_json"         : "%s"
 }
 `

tools/normalizer/zh.py

@@ -4,4 +4,4 @@ from typing import Callable
 def normalizer_zh_tn() -> Callable[[str], str]:
     from tn.chinese.normalizer import Normalizer
 
-    return Normalizer().normalize
+    return Normalizer(remove_interjections=False).normalize