xinference 1.0.1__py3-none-any.whl → 1.2.1__py3-none-any.whl

xinference/thirdparty/f5_tts/train/train.py

@@ -0,0 +1,75 @@
+# training script.
+
+import os
+from importlib.resources import files
+
+import hydra
+
+from f5_tts.model import CFM, DiT, Trainer, UNetT
+from f5_tts.model.dataset import load_dataset
+from f5_tts.model.utils import get_tokenizer
+
+os.chdir(str(files("f5_tts").joinpath("../..")))  # change working directory to root of project (local editable)
+
+
+@hydra.main(version_base="1.3", config_path=str(files("f5_tts").joinpath("configs")), config_name=None)
+def main(cfg):
+    tokenizer = cfg.model.tokenizer
+    mel_spec_type = cfg.model.mel_spec.mel_spec_type
+    exp_name = f"{cfg.model.name}_{mel_spec_type}_{cfg.model.tokenizer}_{cfg.datasets.name}"
+
+    # set text tokenizer
+    if tokenizer != "custom":
+        tokenizer_path = cfg.datasets.name
+    else:
+        tokenizer_path = cfg.model.tokenizer_path
+    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)
+
+    # set model
+    if "F5TTS" in cfg.model.name:
+        model_cls = DiT
+    elif "E2TTS" in cfg.model.name:
+        model_cls = UNetT
+    wandb_resume_id = None
+
+    model = CFM(
+        transformer=model_cls(**cfg.model.arch, text_num_embeds=vocab_size, mel_dim=cfg.model.mel_spec.n_mel_channels),
+        mel_spec_kwargs=cfg.model.mel_spec,
+        vocab_char_map=vocab_char_map,
+    )
+
+    # init trainer
+    trainer = Trainer(
+        model,
+        epochs=cfg.optim.epochs,
+        learning_rate=cfg.optim.learning_rate,
+        num_warmup_updates=cfg.optim.num_warmup_updates,
+        save_per_updates=cfg.ckpts.save_per_updates,
+        checkpoint_path=str(files("f5_tts").joinpath(f"../../{cfg.ckpts.save_dir}")),
+        batch_size=cfg.datasets.batch_size_per_gpu,
+        batch_size_type=cfg.datasets.batch_size_type,
+        max_samples=cfg.datasets.max_samples,
+        grad_accumulation_steps=cfg.optim.grad_accumulation_steps,
+        max_grad_norm=cfg.optim.max_grad_norm,
+        logger=cfg.ckpts.logger,
+        wandb_project="CFM-TTS",
+        wandb_run_name=exp_name,
+        wandb_resume_id=wandb_resume_id,
+        last_per_steps=cfg.ckpts.last_per_steps,
+        log_samples=True,
+        bnb_optimizer=cfg.optim.bnb_optimizer,
+        mel_spec_type=mel_spec_type,
+        is_local_vocoder=cfg.model.vocoder.is_local,
+        local_vocoder_path=cfg.model.vocoder.local_path,
+    )
+
+    train_dataset = load_dataset(cfg.datasets.name, tokenizer, mel_spec_kwargs=cfg.model.mel_spec)
+    trainer.train(
+        train_dataset,
+        num_workers=cfg.datasets.num_workers,
+        resumable_with_seed=666,  # seed for shuffling dataset
+    )
+
+
+if __name__ == "__main__":
+    main()

xinference/thirdparty/fish_speech/fish_speech/conversation.py

@@ -2,41 +2,10 @@ from dataclasses import dataclass, field
 from typing import Literal
 
 import torch
-from transformers import AutoTokenizer, PretrainedConfig, PreTrainedTokenizerFast
-
-IM_START_TOKEN = "<|im_start|>"
-IM_END_TOKEN = "<|im_end|>"
-SEMANTIC_TOKEN = "<|semantic|>"
-MEL_TOKEN = "<|mel|>"
-PHONEME_START_TOKEN = "<|phoneme_start|>"
-PHONEME_END_TOKEN = "<|phoneme_end|>"
-ALL_SPECIAL_TOKENS = [
-    IM_START_TOKEN,
-    IM_END_TOKEN,
-    SEMANTIC_TOKEN,
-    MEL_TOKEN,
-    PHONEME_START_TOKEN,
-    PHONEME_END_TOKEN,
-]
-
-CODEBOOK_PAD_TOKEN_ID = 0
-
-
-class FishTokenizerConfig(PretrainedConfig):
-    share_codebook_embeddings: bool = True
-    codebook_size: int = 1024
-    num_codebooks: int = 8
 
+from .tokenizer import MODALITY_TOKENS, FishTokenizer
 
-class FishTokenizerFast(PreTrainedTokenizerFast):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.share_codebook_embeddings = kwargs.pop("share_codebook_embeddings", True)
-        self.codebook_size = kwargs.pop("codebook_size", 1024)
-        self.num_codebooks = kwargs.pop("num_codebooks", 8)
-
-
-AutoTokenizer.register(FishTokenizerConfig, fast_tokenizer_class=FishTokenizerFast)
+CODEBOOK_PAD_TOKEN_ID = 0
 
 
 @dataclass(kw_only=True)
@@ -54,77 +23,72 @@ class TextPart(BasePart):
     text: str
 
 
-@dataclass(kw_only=True)
-class MelPart(BasePart):
-    mels: torch.Tensor
-
-
 @dataclass(kw_only=True)
 class EncodedMessage:
     tokens: torch.Tensor
     labels: torch.Tensor
+    vq_mask_tokens: torch.Tensor | None = None
+    vq_mask_labels: torch.Tensor | None = None
     vq_parts: list[torch.Tensor]
-    mel_parts: list[torch.Tensor]
     vq_require_losses: torch.Tensor | None = None
 
 
 @dataclass(kw_only=True)
 class Message:
     role: Literal["system", "user", "assistant"]
-    parts: list[VQPart | TextPart | MelPart] = field(default_factory=list)
+    parts: list[VQPart | TextPart] = field(default_factory=list)
     add_im_start: bool = True
     add_im_end: bool = True
     cal_loss: bool = False
+    modality: Literal["text", "voice", "interleave"] | None = None
 
     # By default, ignore the loss of the auto-generated im_start token
     ignore_im_start_loss: bool = True
 
     def encode(
         self: "Message",
-        tokenizer: AutoTokenizer,
+        tokenizer: FishTokenizer,
     ) -> EncodedMessage:
         all_tokens = []
         all_labels = []
 
         # Multi-modal tokens
         vq_parts = []
-        mel_parts = []
-
-        semantic_id, mel_id = tokenizer.convert_tokens_to_ids(
-            [SEMANTIC_TOKEN, MEL_TOKEN]
-        )
+        vq_masks = []
 
         parts = self.parts.copy()
         if self.add_im_start:
-            parts.insert(0, TextPart(text=f"<|im_start|>{self.role}\n"))
+            modality_token = MODALITY_TOKENS[self.modality] if self.modality else ""
+            parts.insert(0, TextPart(text=f"<|im_start|>{self.role}\n{modality_token}"))
 
         if self.add_im_end:
             parts.append(TextPart(text="<|im_end|>"))
 
         for part in parts:
             if isinstance(part, TextPart):
-                tokens = tokenizer.encode(
-                    part.text,
-                    add_special_tokens=False,
-                    truncation=False,
-                    return_tensors="pt",
-                ).int()[0]
+                tokens = torch.tensor(
+                    tokenizer.encode(part.text),
+                    dtype=torch.int,
+                )
             elif isinstance(part, VQPart):
-                tokens = torch.zeros(part.codes.shape[1], dtype=torch.int) + semantic_id
-                codes = part.codes.clone() + 1
-
-                if getattr(tokenizer, "share_codebook_embeddings", True) is False:
-                    for i in range(len(codes)):
-                        codes[i] += tokenizer.codebook_size * i
-
-                vq_parts.append(codes)
-            elif isinstance(part, MelPart):
-                tokens = torch.zeros(part.mels.shape[1], dtype=torch.int) + mel_id
-                mel_parts.append(part.mels)
+                curr_codes = part.codes.clone()
+                tokens = torch.tensor(
+                    [
+                        tokenizer.semantic_id_to_token_id[i.item()]
+                        for i in curr_codes[0].int()
+                    ],
+                    dtype=torch.int,
+                )
+                vq_parts.append(curr_codes)
             else:
                 raise ValueError(f"Unsupported part type: {type(part)}")
 
             all_tokens.append(tokens)
+            if isinstance(part, VQPart):
+                vq_masks.append(torch.ones_like(tokens, dtype=torch.bool))
+            else:
+                vq_masks.append(torch.zeros_like(tokens, dtype=torch.bool))
+
             if self.cal_loss:
                 all_labels.append(tokens.clone())
             else:
@@ -132,7 +96,9 @@ class Message:
 
         tokens = torch.cat(all_tokens, dim=0)
         labels = torch.cat(all_labels, dim=0)
-        assert tokens.shape == labels.shape
+        vq_masks = torch.cat(vq_masks, dim=0)
+
+        assert tokens.shape == labels.shape == vq_masks.shape
 
         if self.ignore_im_start_loss and self.add_im_start:
             labels[: len(all_tokens[0])] = -100
@@ -141,7 +107,8 @@ class Message:
             tokens=tokens,
             labels=labels,
             vq_parts=vq_parts,
-            mel_parts=mel_parts,
+            vq_mask_tokens=vq_masks,
+            vq_mask_labels=vq_masks,
         )
 
 
@@ -149,17 +116,23 @@ class Message:
 class Conversation:
     messages: list[Message]
 
+    def __init__(self: "Conversation", messages: list[Message] | None = None):
+        self.messages = messages or []
+
     def encode(
         self: "Conversation",
-        tokenizer: AutoTokenizer,
+        tokenizer: FishTokenizer,
         add_shift: bool = True,
+        ignore_loss_tokens: list[str] = [],
     ) -> EncodedMessage:
         # Build the input_ids and labels
         tokens = []
         labels = []
         vq_parts = []
-        mel_parts = []
+        vq_mask_tokens = []
+        vq_mask_labels = []
         vq_require_losses = []
+        ignore_loss_token_ids = [tokenizer.get_token_id(i) for i in ignore_loss_tokens]
 
         for message in self.messages:
             encoded = message.encode(
@@ -168,16 +141,25 @@ class Conversation:
             tokens.append(encoded.tokens)
             labels.append(encoded.labels)
             vq_parts.extend(encoded.vq_parts)
-            mel_parts.extend(encoded.mel_parts)
+            vq_mask_tokens.append(encoded.vq_mask_tokens)
+            vq_mask_labels.append(encoded.vq_mask_labels)
             vq_require_losses.extend([message.cal_loss] * len(encoded.vq_parts))
 
         tokens = torch.cat(tokens, dim=0)
         labels = torch.cat(labels, dim=0)
+        vq_mask_tokens = torch.cat(vq_mask_tokens, dim=0)
+        vq_mask_labels = torch.cat(vq_mask_labels, dim=0)
         vq_require_losses = torch.tensor(vq_require_losses, dtype=torch.bool)
 
         if add_shift:
             tokens = tokens[:-1]
             labels = labels[1:]
+            vq_mask_tokens = vq_mask_tokens[:-1]
+            vq_mask_labels = vq_mask_labels[1:]
+
+        for i in ignore_loss_token_ids:
+            assert i != -100 and i is not None
+            labels[labels == i] = -100
 
         assert tokens.dtype in [
             torch.int,
@@ -188,15 +170,18 @@ class Conversation:
             tokens=tokens,
             labels=labels,
             vq_parts=vq_parts,
-            mel_parts=mel_parts,
+            vq_mask_tokens=vq_mask_tokens,
+            vq_mask_labels=vq_mask_labels,
             vq_require_losses=vq_require_losses,
         )
 
     def encode_for_inference(
         self: "Conversation",
-        tokenizer: AutoTokenizer,
+        tokenizer: FishTokenizer,
         num_codebooks: int,
     ) -> EncodedMessage:
+        # self.visualize(tokenizer)
+
         encoded = self.encode(tokenizer, add_shift=False)
         tokens = encoded.tokens
         values = torch.zeros((num_codebooks + 1, len(tokens)), dtype=torch.int)
@@ -205,24 +190,47 @@ class Conversation:
         if encoded.vq_parts is None or len(encoded.vq_parts) == 0:
             return values
 
-        semantic_id, mel_id = tokenizer.convert_tokens_to_ids(
-            [SEMANTIC_TOKEN, MEL_TOKEN]
-        )
         vq_parts = encoded.vq_parts
+        vq_parts = [part.to(values.device) for part in vq_parts]
         vq_parts = torch.cat(vq_parts, dim=1)
-        values[1:, tokens == semantic_id] = vq_parts
+        values[0, encoded.vq_mask_tokens] = vq_parts[0] + tokenizer.semantic_begin_id
+        values[1:, encoded.vq_mask_tokens] = vq_parts
+
         return values
 
-    def visualize(self: "Conversation", tokenizer: AutoTokenizer):
-        encoded = self.encode(tokenizer, add_shift=False)
+    def visualize(
+        self: "Conversation",
+        tokenizer: FishTokenizer,
+        ignore_loss_tokens: list[str] = [],
+    ):
+        encoded = self.encode(
+            tokenizer, add_shift=False, ignore_loss_tokens=ignore_loss_tokens
+        )
 
-        print_in_blue = lambda x: print("\033[94m" + x + "\033[0m", end="")
-        print_in_green = lambda x: print("\033[92m" + x + "\033[0m", end="")
+        # Colors for alternating tokens
+        colors = {
+            "blue": "\033[94m",  # Light blue
+            "cyan": "\033[96m",  # Cyan
+            "green": "\033[92m",  # Light green
+            "dark_green": "\033[32m",  # Dark green
+        }
+        blue_idx = 0
+        green_idx = 0
+
+        def print_in_blue(x):
+            nonlocal blue_idx
+            color = colors["blue"] if blue_idx % 2 == 0 else colors["cyan"]
+            print(f"{color}{x}\033[0m", end="")
+            blue_idx += 1
+
+        def print_in_green(x):
+            nonlocal green_idx
+            color = colors["green"] if green_idx % 2 == 0 else colors["dark_green"]
+            print(f"{color}{x}\033[0m", end="")
+            green_idx += 1
 
         for tok, lab in zip(encoded.tokens, encoded.labels):
-            val = tokenizer.decode(tok, skip_special_tokens=False)
-            if val == "\n":
-                val = "\\n\n"
+            val = tokenizer.decode([tok])
 
             if lab == -100:
                 print_in_green(val)
@@ -231,6 +239,9 @@ class Conversation:
 
         print()
 
+    def append(self: "Conversation", message: Message):
+        self.messages.append(message)
+
 
 if __name__ == "__main__":
     message0 = Message(
@@ -248,7 +259,7 @@ if __name__ == "__main__":
         cal_loss=True,
     )
     conversation = Conversation([message0, message1])
-    tokenizer = AutoTokenizer.from_pretrained("checkpoints/Qwen2-1.5B-Instruct")
+    tokenizer = FishTokenizer.from_pretrained("checkpoints/Qwen2-1.5B-Instruct")
     conversation.visualize(tokenizer)
 
     encoded = conversation.encode(tokenizer)

xinference/thirdparty/fish_speech/fish_speech/models/text2semantic/llama.py

@@ -16,7 +16,7 @@ from torch.nn.attention import SDPBackend, sdpa_kernel
 from torch.utils.checkpoint import checkpoint
 from transformers import AutoTokenizer
 
-from fish_speech.conversation import SEMANTIC_TOKEN
+from fish_speech.tokenizer import SEMANTIC_TOKENS, FishTokenizer
 from fish_speech.utils import RankedLogger
 
 from .lora import LoraConfig, setup_lora
@@ -61,6 +61,7 @@ class BaseModelArgs:
     # Dummy vars
     is_reward_model: bool = False
     share_codebook_embeddings: bool = True
+    scale_codebook_embeddings: bool = False
 
     def __post_init__(self):
         if self.n_local_heads == -1:
@@ -164,13 +165,17 @@ class BaseTransformerForwardResult:
 
 class BaseTransformer(nn.Module):
     def __init__(
-        self, config: BaseModelArgs, tokenizer: AutoTokenizer, init_weights: bool = True
+        self,
+        config: BaseModelArgs,
+        tokenizer: FishTokenizer | AutoTokenizer,
+        init_weights: bool = True,
     ) -> None:
         super().__init__()
         self.config = config
         self.tokenizer = tokenizer
-
-        self.semantic_token_id = tokenizer.convert_tokens_to_ids(SEMANTIC_TOKEN)
+        self.semantic_token_ids = [
+            tokenizer.get_token_id(SEMANTIC_TOKEN) for SEMANTIC_TOKEN in SEMANTIC_TOKENS
+        ]
 
         # Slow transformer
         self.embeddings = nn.Embedding(
@@ -245,8 +250,10 @@ class BaseTransformer(nn.Module):
         vocab_embeds = [self.embeddings(x[:, 0])]
         for i in range(self.config.num_codebooks):
             emb = self.codebook_embeddings(x[:, i + 1] + i * self.config.codebook_size)
-            emb[x[:, 0] != self.semantic_token_id] = 0
-            vocab_embeds.append(emb)
+            semantic_token_ids_tensor = torch.tensor(
+                self.semantic_token_ids, device=x.device
+            )
+            emb[~torch.isin(x[:, 0], semantic_token_ids_tensor)] = 0
 
         x = torch.stack(vocab_embeds, dim=3)
         x = x.sum(dim=3)
@@ -294,20 +301,45 @@ class BaseTransformer(nn.Module):
 
     def forward_generate(
         self,
-        x: Tensor,
+        inp: Tensor,
         input_pos: Optional[Tensor] = None,
+        vq_masks: Optional[Tensor] = None,  # this is not used in fact
         return_all: bool = False,
     ) -> BaseTransformerForwardResult:
         # This is used for generation, optimized for torch compile
-        assert (
-            self.max_seq_len != -1 and self.max_batch_size != -1
-        ), "Please call setup_caches before forward_generate"
+        # assert (
+        #     self.max_seq_len != -1 and self.max_batch_size != -1
+        # ), "Please call setup_caches before forward_generate"
 
-        x = self.embed(x)
+        embeds = []
+        for i in range(self.config.num_codebooks):
+            if self.config.share_codebook_embeddings:
+                _tokens = inp[:, i + 1] + i * self.config.codebook_size
+            else:
+                _tokens = inp[:, i + 1]
 
-        mask = self.causal_mask[
-            None, None, input_pos, : self.max_seq_len
-        ]  # (B, N, Q, K)
+            emb = self.codebook_embeddings(_tokens)
+            embeds.append(emb)
+
+        vq_embeds_sum = torch.stack(embeds, dim=1).sum(dim=1)
+        # if self.config.use_codebook_mlp:
+        #     vq_embeds_sum = vq_embeds_sum / self.config.num_codebooks
+        #     vq_embeds_sum = self.codebook_mlp(vq_embeds_sum)
+
+        vq_masks = (inp[:, 0] >= self.tokenizer.semantic_begin_id) & (
+            inp[:, 0] <= self.tokenizer.semantic_end_id
+        )
+
+        vq_embeds_sum[~vq_masks] = 0
+        x = self.embeddings(inp[:, 0]) + vq_embeds_sum
+
+        if input_pos is None:
+            input_pos = torch.arange(inp.shape[-1], device=x.device)
+            max_seq_len = inp.shape[-1]
+        else:
+            max_seq_len = self.max_seq_len
+
+        mask = self.causal_mask[None, None, input_pos, :max_seq_len]  # (B, N, Q, K)
         freqs_cis = self.freqs_cis[input_pos]
 
         for layer in self.layers:
@@ -320,7 +352,9 @@ class BaseTransformer(nn.Module):
         # We got slow_out here
         slow_out = self.norm(x)
 
-        if self.config.tie_word_embeddings:
+        if self.config.is_reward_model:
+            token_logits = self.score_output(slow_out)
+        elif self.config.tie_word_embeddings:
             token_logits = F.linear(slow_out, self.embeddings.weight)
         else:
             token_logits = self.output(slow_out)
@@ -348,6 +382,7 @@ class BaseTransformer(nn.Module):
         max_length: int | None = None,
         lora_config: LoraConfig | None = None,
         rope_base: int | None = None,
+        is_agent: bool = False,
     ) -> "BaseTransformer":
         config = BaseModelArgs.from_pretrained(str(path))
         if max_length is not None:
@@ -366,7 +401,12 @@ class BaseTransformer(nn.Module):
             case _:
                 raise ValueError(f"Unknown model type: {config.model_type}")
 
-        tokenizer = AutoTokenizer.from_pretrained(str(path))
+        if is_agent:
+            tokenizer = AutoTokenizer.from_pretrained(str(path))
+        else:
+            tokenizer_path = str(path) + "/tokenizer.tiktoken"
+            tokenizer = FishTokenizer(tokenizer_path)
+
         log.info(f"Loading model from {path}, config: {config}")
         model = model_cls(config, tokenizer=tokenizer)
 
@@ -452,7 +492,7 @@ class BaseTransformer(nn.Module):
 
 
 class NaiveTransformer(BaseTransformer):
-    def __init__(self, config: NaiveModelArgs, tokenizer: AutoTokenizer) -> None:
+    def __init__(self, config: NaiveModelArgs, tokenizer: FishTokenizer) -> None:
         super().__init__(config, init_weights=False, tokenizer=tokenizer)
 
         self.codebook_norm = RMSNorm(config.dim, eps=config.norm_eps)
@@ -498,7 +538,7 @@ class NaiveTransformer(BaseTransformer):
 
 
 class DualARTransformer(BaseTransformer):
-    def __init__(self, config: NaiveModelArgs, tokenizer: AutoTokenizer) -> None:
+    def __init__(self, config: NaiveModelArgs, tokenizer: FishTokenizer) -> None:
         super().__init__(config, init_weights=False, tokenizer=tokenizer)
 
         # Project to fast dim if needed
@@ -654,9 +694,12 @@ class DualARTransformer(BaseTransformer):
         return codebook_logits
 
     def forward_generate(
-        self, x: Tensor, input_pos: Optional[Tensor] = None
+        self,
+        x: Tensor,
+        input_pos: Optional[Tensor] = None,
+        vq_masks: Optional[Tensor] = None,
     ) -> TransformerForwardResult:
-        x = super().forward_generate(x, input_pos)
+        x = super().forward_generate(x, input_pos, vq_masks)
         x.hidden_states = self.fast_project_in(x.hidden_states)
         return x
 

xinference/thirdparty/fish_speech/fish_speech/text/clean.py

@@ -1,33 +1,8 @@
 import re
 
 SYMBOLS_MAPPING = {
-    "\n": "",
-    "…": ".",
-    "“": "'",
-    "”": "'",
     "‘": "'",
     "’": "'",
-    "【": "",
-    "】": "",
-    "[": "",
-    "]": "",
-    "（": "",
-    "）": "",
-    "(": "",
-    ")": "",
-    "・": "",
-    "·": "",
-    "「": "'",
-    "」": "'",
-    "《": "'",
-    "》": "'",
-    "—": "",
-    "～": "",
-    "~": "",
-    "：": ",",
-    "；": ",",
-    ";": ",",
-    ":": ",",
 }
 
 REPLACE_SYMBOL_REGEX = re.compile(
@@ -57,6 +32,6 @@ def clean_text(text):
     text = EMOJI_REGEX.sub(r"", text)
 
     # Remove continuous periods (...) and commas (,,,)
-    text = re.sub(r"[.,]{2,}", lambda m: m.group()[0], text)
+    text = re.sub(r"[,]{2,}", lambda m: m.group()[0], text)
 
     return text

xinference/thirdparty/fish_speech/fish_speech/text/spliter.py

@@ -4,7 +4,7 @@ import string
 from fish_speech.text.clean import clean_text
 
 
-def utf_8_len(text):
+def utf_8_len(text: str):
     return len(text.encode("utf-8"))