New model changes

src/chatterbox/models/t3/t3.py

@@ -10,7 +10,11 @@ import torch
 import torch.nn.functional as F
 from torch import nn, Tensor
 from transformers import LlamaModel, LlamaConfig
-from transformers.generation.logits_process import TopPLogitsWarper, RepetitionPenaltyLogitsProcessor, MinPLogitsWarper
+from transformers.generation.logits_process import (
+    TopPLogitsWarper,
+    RepetitionPenaltyLogitsProcessor,
+    MinPLogitsWarper,
+)
 
 from .modules.learned_pos_emb import LearnedPositionEmbeddings
 
@@ -27,8 +31,12 @@ logger = logging.getLogger(__name__)
 
 def _ensure_BOT_EOT(text_tokens: Tensor, hp):
     B = text_tokens.size(0)
-    assert (text_tokens == hp.start_text_token).int().sum() >= B, "missing start_text_token"
-    assert (text_tokens == hp.stop_text_token).int().sum() >= B, "missing stop_text_token"
+    assert (
+        text_tokens == hp.start_text_token
+    ).int().sum() >= B, "missing start_text_token"
+    assert (
+        text_tokens == hp.stop_text_token
+    ).int().sum() >= B, "missing stop_text_token"
 
 
 class T3(nn.Module):
@@ -43,7 +51,9 @@ class T3(nn.Module):
 
     def __init__(self, hp=None):
         if hp is None:
-            hp = T3Config.english_only()  # Default to English-only config for backward compatibility
+            hp = (
+                T3Config.english_only()
+            )  # Default to English-only config for backward compatibility
         super().__init__()
         self.hp = hp
         self.cfg = LlamaConfig(**LLAMA_CONFIGS[hp.llama_config_name])
@@ -65,8 +75,12 @@ class T3(nn.Module):
             self.speech_pos_emb = LearnedPositionEmbeddings(max_mel_seq_len, self.dim)
 
         # logit projection
-        self.text_head = nn.Linear(self.cfg.hidden_size, hp.text_tokens_dict_size, bias=False)
-        self.speech_head = nn.Linear(self.cfg.hidden_size, hp.speech_tokens_dict_size, bias=False)
+        self.text_head = nn.Linear(
+            self.cfg.hidden_size, hp.text_tokens_dict_size, bias=False
+        )
+        self.speech_head = nn.Linear(
+            self.cfg.hidden_size, hp.speech_tokens_dict_size, bias=False
+        )
         self.compiled = False
 
     @property
@@ -77,9 +91,13 @@ class T3(nn.Module):
         """
         Token cond data needs to be embedded, so that needs to be here instead of in `T3CondEnc`.
         """
-        if t3_cond.cond_prompt_speech_tokens is not None and t3_cond.cond_prompt_speech_emb is None:
-            t3_cond.cond_prompt_speech_emb = self.speech_emb(t3_cond.cond_prompt_speech_tokens) + \
-                self.speech_pos_emb(t3_cond.cond_prompt_speech_tokens)
+        if (
+            t3_cond.cond_prompt_speech_tokens is not None
+            and t3_cond.cond_prompt_speech_emb is None
+        ):
+            t3_cond.cond_prompt_speech_emb = self.speech_emb(
+                t3_cond.cond_prompt_speech_tokens
+            ) + self.speech_pos_emb(t3_cond.cond_prompt_speech_tokens)
         return self.cond_enc(t3_cond)  # (B, len_cond, dim)
 
     def prepare_input_embeds(
@@ -103,13 +121,15 @@ class T3(nn.Module):
         len_cond = cond_emb.size(1)
 
         if cond_emb.size(0) != text_emb.size(0):
-             cond_emb = cond_emb.expand(text_emb.size(0), -1, -1)
+            cond_emb = cond_emb.expand(text_emb.size(0), -1, -1)
 
         # concat
-        embeds = torch.stack([
-            torch.cat((ce, te, se))
-            for ce, te, se in zip(cond_emb, text_emb, speech_emb)
-        ])  # (B, length, dim)
+        embeds = torch.stack(
+            [
+                torch.cat((ce, te, se))
+                for ce, te, se in zip(cond_emb, text_emb, speech_emb)
+            ]
+        )  # (B, length, dim)
         return embeds, len_cond
 
     def forward(
@@ -140,7 +160,9 @@ class T3(nn.Module):
             return_dict=True,
             use_cache=(not training),
         )
-        hidden_states = tfmr_out.hidden_states[-1]  # final tfmr layer output, (B, seq, dim)
+        hidden_states = tfmr_out.hidden_states[
+            -1
+        ]  # final tfmr layer output, (B, seq, dim)
 
         # post-processing: splice out text and speech parts of hidden states
         len_text = text_tokens.size(1)
@@ -154,8 +176,8 @@ class T3(nn.Module):
             text_end = len_cond + ttl[i].item()
             speech_start = len_cond + text_tokens.size(1)
             speech_end = speech_start + stl[i].item()
-            text_latents[i, :ttl[i]] = hidden_states[i, len_cond:text_end]
-            speech_latents[i, :stl[i]] = hidden_states[i, speech_start:speech_end]
+            text_latents[i, : ttl[i]] = hidden_states[i, len_cond:text_end]
+            speech_latents[i, : stl[i]] = hidden_states[i, speech_start:speech_end]
 
         # logit projection
         text_logits = self.text_head(text_latents)
@@ -173,17 +195,21 @@ class T3(nn.Module):
         self,
         *,
         t3_cond: T3Cond,
-        text_tokens: torch.LongTensor,
-        text_token_lens: torch.LongTensor,
-        speech_tokens: torch.LongTensor,
-        speech_token_lens: torch.LongTensor,
+        text_tokens: torch.LongTensor,  # (B, S_text_padded), includes BOS & EOS
+        text_token_lens: torch.LongTensor,  # (B,), actual lengths including BOS & EOS
+        speech_tokens: torch.LongTensor,  # (B, S_speech_padded), includes BOS & EOS
+        speech_token_lens: torch.LongTensor,  # (B,), actual lengths including BOS & EOS
+        labels_text: torch.LongTensor,  # (B, S_text_padded-1), already masked with –100
+        labels_speech: torch.LongTensor,  # (B, S_speech_padded-1), already masked with –100
     ):
-        "training method"
-        len_text = text_tokens.size(1)
-        len_speech = speech_tokens.size(1)
-        assert len_text == text_token_lens.max()
-        assert len_speech == speech_token_lens.max()
+        """
+        Compute text and speech cross-entropy using pre-masked labels from the collator.
+        Assumes:
+        - labels_text[t] corresponds to predicting text_tokens[:, 1:] with –100 where ignored
+        - labels_speech[t] corresponds to predicting speech_tokens[:, 1:] with –100 where ignored
+        """
 
+        # 1) Run model to get logits
         out = self.forward(
             t3_cond=t3_cond,
             text_tokens=text_tokens,
@@ -191,19 +217,42 @@ class T3(nn.Module):
             speech_tokens=speech_tokens,
             speech_token_lens=speech_token_lens,
             training=True,
-        )  # (B, seq, vocab_size)
-
-        # Calc CCE losses
-        IGNORE_ID = -100
+        )
+        # out.text_logits: (B, S_text_padded, V_text)
+        # out.speech_logits: (B, S_speech_padded, V_speech)
         device = out.text_logits.device
-        mask_text = torch.arange(len_text, device=device)[None] >= text_token_lens[:, None]  # (B, len_text)
-        mask_speech = torch.arange(len_speech, device=device)[None] >= speech_token_lens[:, None]  # (B, len_speech)
-        masked_text = text_tokens.masked_fill(mask_text, IGNORE_ID)
-        masked_speech = speech_tokens.masked_fill(mask_speech, IGNORE_ID)
-        loss_text = F.cross_entropy(out.text_logits, masked_text, ignore_index=IGNORE_ID)
-        loss_speech = F.cross_entropy(out.speech_logits, masked_speech, ignore_index=IGNORE_ID)
+        IGNORE_ID = -100
 
-        return loss_text, loss_speech
+        # --- Text Loss (use labels_text directly) ---
+        # Align logits: predict t₁..EOS from inputs [BOS, t₁..]
+        logits_for_text = out.text_logits[
+            :, :-1, :
+        ].contiguous()  # (B, S_text_padded-1, V_text)
+        # labels_text already has shape (B, S_text_padded-1) with –100 where masked
+        if logits_for_text.size(1) == 0:
+            loss_text = torch.tensor(0.0, device=device, requires_grad=self.training)
+        else:
+            loss_text = F.cross_entropy(
+                logits_for_text.transpose(1, 2),  # (B, V_text, S_text_padded-1)
+                labels_text,  # (B, S_text_padded-1), ignore_index=–100
+                ignore_index=IGNORE_ID,
+            )
+
+        # --- Speech Loss (use labels_speech directly) ---
+        logits_for_speech = out.speech_logits[
+            :, :-1, :
+        ].contiguous()  # (B, S_speech_padded-1, V_speech)
+        # labels_speech already has shape (B, S_speech_padded-1) with –100 where masked
+        if logits_for_speech.size(1) == 0:
+            loss_speech = torch.tensor(0.0, device=device, requires_grad=self.training)
+        else:
+            loss_speech = F.cross_entropy(
+                logits_for_speech.transpose(1, 2),  # (B, V_speech, S_speech_padded-1)
+                labels_speech,  # (B, S_speech_padded-1), ignore_index=–100
+                ignore_index=IGNORE_ID,
+            )
+
+        return loss_text, loss_speech, out.speech_logits
 
     @torch.inference_mode()
     def inference(
@@ -211,11 +260,9 @@ class T3(nn.Module):
         *,
         t3_cond: T3Cond,
         text_tokens: Tensor,
-        initial_speech_tokens: Optional[Tensor]=None,
-
+        initial_speech_tokens: Optional[Tensor] = None,
         # misc conditioning
-        prepend_prompt_speech_tokens: Optional[Tensor]=None,
-
+        prepend_prompt_speech_tokens: Optional[Tensor] = None,
         # HF generate args
         num_return_sequences=1,
         max_new_tokens=None,
@@ -235,11 +282,15 @@ class T3(nn.Module):
         # Validate / sanitize inputs
         assert prepend_prompt_speech_tokens is None, "not implemented"
         _ensure_BOT_EOT(text_tokens, self.hp)
-        text_tokens = torch.atleast_2d(text_tokens).to(dtype=torch.long, device=self.device)
+        text_tokens = torch.atleast_2d(text_tokens).to(
+            dtype=torch.long, device=self.device
+        )
 
         # Default initial speech to a single start-of-speech token
         if initial_speech_tokens is None:
-            initial_speech_tokens = self.hp.start_speech_token * torch.ones_like(text_tokens[:, :1])
+            initial_speech_tokens = self.hp.start_speech_token * torch.ones_like(
+                text_tokens[:, :1]
+            )
 
         # Prepare custom input embeds
         embeds, len_cond = self.prepare_input_embeds(
@@ -264,7 +315,7 @@ class T3(nn.Module):
                     self.tfmr,
                     None,
                     text_tokens_slice=(len_cond, len_cond + text_tokens.size(-1)),
-                    alignment_layer_idx=9, # TODO: hparam or something?
+                    alignment_layer_idx=9,  # TODO: hparam or something?
                     eos_idx=self.hp.stop_speech_token,
                 )
                 assert alignment_stream_analyzer.eos_idx == self.hp.stop_speech_token
@@ -298,7 +349,9 @@ class T3(nn.Module):
 
         device = embeds.device
 
-        bos_token = torch.tensor([[self.hp.start_speech_token]], dtype=torch.long, device=device)
+        bos_token = torch.tensor(
+            [[self.hp.start_speech_token]], dtype=torch.long, device=device
+        )
         bos_embed = self.speech_emb(bos_token)  # shape: (B, 1, embed_dim)
         bos_embed = bos_embed + self.speech_pos_emb.get_fixed_embedding(0)
 
@@ -316,7 +369,9 @@ class T3(nn.Module):
         top_p_warper = TopPLogitsWarper(top_p=top_p)
         min_p_warper = MinPLogitsWarper(min_p=min_p)
         top_p_warper = TopPLogitsWarper(top_p=top_p)
-        repetition_penalty_processor = RepetitionPenaltyLogitsProcessor(penalty=float(repetition_penalty))
+        repetition_penalty_processor = RepetitionPenaltyLogitsProcessor(
+            penalty=float(repetition_penalty)
+        )
 
         # ---- Initial Forward Pass (no kv_cache yet) ----
         output = self.patched_model(
@@ -332,29 +387,33 @@ class T3(nn.Module):
 
         # ---- Generation Loop using kv_cache ----
         for i in tqdm(range(max_new_tokens), desc="Sampling", dynamic_ncols=True):
-            logits_step = output.logits[:, -1, :]                
+            logits_step = output.logits[:, -1, :]
             # CFG combine  → (1, V)
-            cond   = logits_step[0:1, :]
+            cond = logits_step[0:1, :]
             uncond = logits_step[1:2, :]
             cfg = torch.as_tensor(cfg_weight, device=cond.device, dtype=cond.dtype)
             logits = cond + cfg * (cond - uncond)
-            
+
             # Apply alignment stream analyzer integrity checks
             if self.patched_model.alignment_stream_analyzer is not None:
-                if logits.dim() == 1:            # guard in case something upstream squeezed
-                    logits = logits.unsqueeze(0) # (1, V)
+                if logits.dim() == 1:  # guard in case something upstream squeezed
+                    logits = logits.unsqueeze(0)  # (1, V)
                 # Pass the last generated token for repetition tracking
-                last_token = generated_ids[0, -1].item() if len(generated_ids[0]) > 0 else None
-                logits = self.patched_model.alignment_stream_analyzer.step(logits, next_token=last_token)  # (1, V)
+                last_token = (
+                    generated_ids[0, -1].item() if len(generated_ids[0]) > 0 else None
+                )
+                logits = self.patched_model.alignment_stream_analyzer.step(
+                    logits, next_token=last_token
+                )  # (1, V)
 
             # Apply repetition penalty
-            ids_for_proc = generated_ids[:1, ...]   # batch = 1
+            ids_for_proc = generated_ids[:1, ...]  # batch = 1
             logits = repetition_penalty_processor(ids_for_proc, logits)  # expects (B,V)
-            
+
             # Apply temperature scaling.
             if temperature != 1.0:
                 logits = logits / temperature
-                
+
             # Apply min_p and top_p filtering
             logits = min_p_warper(ids_for_proc, logits)
             logits = top_p_warper(ids_for_proc, logits)
@@ -373,7 +432,9 @@ class T3(nn.Module):
 
             # Get embedding for the new token.
             next_token_embed = self.speech_emb(next_token)
-            next_token_embed = next_token_embed + self.speech_pos_emb.get_fixed_embedding(i + 1)
+            next_token_embed = (
+                next_token_embed + self.speech_pos_emb.get_fixed_embedding(i + 1)
+            )
 
             #  For CFG
             next_token_embed = torch.cat([next_token_embed, next_token_embed])

src/chatterbox/mtl_tts.py

@@ -22,36 +22,36 @@ REPO_ID = "ResembleAI/chatterbox"
 
 # Supported languages for the multilingual model
 SUPPORTED_LANGUAGES = {
-  "ar": "Arabic",
-  "da": "Danish",
-  "de": "German",
-  "el": "Greek",
-  "en": "English",
-  "es": "Spanish",
-  "fi": "Finnish",
-  "fr": "French",
-  "he": "Hebrew",
-  "hi": "Hindi",
-  "it": "Italian",
-  "ja": "Japanese",
-  "ko": "Korean",
-  "ms": "Malay",
-  "nl": "Dutch",
-  "no": "Norwegian",
-  "pl": "Polish",
-  "pt": "Portuguese",
-  "ru": "Russian",
-  "sv": "Swedish",
-  "sw": "Swahili",
-  "tr": "Turkish",
-  "zh": "Chinese",
+    "ar": "Arabic",
+    "da": "Danish",
+    "de": "German",
+    "el": "Greek",
+    "en": "English",
+    "es": "Spanish",
+    "fi": "Finnish",
+    "fr": "French",
+    "he": "Hebrew",
+    "hi": "Hindi",
+    "it": "Italian",
+    "ja": "Japanese",
+    "ko": "Korean",
+    "ms": "Malay",
+    "nl": "Dutch",
+    "no": "Norwegian",
+    "pl": "Polish",
+    "pt": "Portuguese",
+    "ru": "Russian",
+    "sv": "Swedish",
+    "sw": "Swahili",
+    "tr": "Turkish",
+    "zh": "Chinese",
 }
 
 
 def punc_norm(text: str) -> str:
     """
-        Quick cleanup func for punctuation from LLMs or
-        containing chars not seen often in the dataset
+    Quick cleanup func for punctuation from LLMs or
+    containing chars not seen often in the dataset
     """
     if len(text) == 0:
         return "You need to add some text for me to talk."
@@ -73,8 +73,8 @@ def punc_norm(text: str) -> str:
         ("—", "-"),
         ("–", "-"),
         (" ,", ","),
-        ("“", "\""),
-        ("”", "\""),
+        ("“", '"'),
+        ("”", '"'),
         ("‘", "'"),
         ("’", "'"),
     ]
@@ -83,7 +83,7 @@ def punc_norm(text: str) -> str:
 
     # Add full stop if no ending punc
     text = text.rstrip(" ")
-    sentence_enders = {".", "!", "?", "-", ",","、","，","。","？","！"}
+    sentence_enders = {".", "!", "?", "-", ",", "、", "，", "。", "？", "！"}
     if not any(text.endswith(p) for p in sentence_enders):
         text += "."
 
@@ -107,6 +107,7 @@ class Conditionals:
         - prompt_feat_len
         - embedding
     """
+
     t3: T3Cond
     gen: dict
 
@@ -118,16 +119,13 @@ class Conditionals:
         return self
 
     def save(self, fpath: Path):
-        arg_dict = dict(
-            t3=self.t3.__dict__,
-            gen=self.gen
-        )
+        arg_dict = dict(t3=self.t3.__dict__, gen=self.gen)
         torch.save(arg_dict, fpath)
 
     @classmethod
     def load(cls, fpath, map_location="cpu"):
         kwargs = torch.load(fpath, map_location=map_location, weights_only=True)
-        return cls(T3Cond(**kwargs['t3']), kwargs['gen'])
+        return cls(T3Cond(**kwargs["t3"]), kwargs["gen"])
 
 
 class ChatterboxMultilingualTTS:
@@ -158,13 +156,11 @@ class ChatterboxMultilingualTTS:
         return SUPPORTED_LANGUAGES.copy()
 
     @classmethod
-    def from_local(cls, ckpt_dir, device) -> 'ChatterboxMultilingualTTS':
+    def from_local(cls, ckpt_dir, device) -> "ChatterboxMultilingualTTS":
         ckpt_dir = Path(ckpt_dir)
 
         ve = VoiceEncoder()
-        ve.load_state_dict(
-            torch.load(ckpt_dir / "ve.pt", weights_only=True)
-        )
+        ve.load_state_dict(torch.load(ckpt_dir / "ve.pt", weights_only=True))
         ve.to(device).eval()
 
         t3 = T3(T3Config.multilingual())
@@ -175,14 +171,10 @@ class ChatterboxMultilingualTTS:
         t3.to(device).eval()
 
         s3gen = S3Gen()
-        s3gen.load_state_dict(
-            torch.load(ckpt_dir / "s3gen.pt", weights_only=True)
-        )
+        s3gen.load_state_dict(torch.load(ckpt_dir / "s3gen.pt", weights_only=True))
         s3gen.to(device).eval()
 
-        tokenizer = MTLTokenizer(
-            str(ckpt_dir / "grapheme_mtl_merged_expanded_v1.json")
-        )
+        tokenizer = MTLTokenizer(str(ckpt_dir / "grapheme_mtl_merged_expanded_v1.json"))
 
         conds = None
         if (builtin_voice := ckpt_dir / "conds.pt").exists():
@@ -191,36 +183,94 @@ class ChatterboxMultilingualTTS:
         return cls(t3, s3gen, ve, tokenizer, device, conds=conds)
 
     @classmethod
-    def from_pretrained(cls, device: torch.device) -> 'ChatterboxMultilingualTTS':
+    def from_pretrained(cls, device: torch.device) -> "ChatterboxMultilingualTTS":
         ckpt_dir = Path(
             snapshot_download(
                 repo_id=REPO_ID,
                 repo_type="model",
-                revision="main", 
-                allow_patterns=["ve.pt", "t3_mtl23ls_v2.safetensors", "s3gen.pt", "grapheme_mtl_merged_expanded_v1.json", "conds.pt", "Cangjie5_TC.json"],
+                revision="main",
+                allow_patterns=[
+                    "ve.pt",
+                    "t3_mtl23ls_v2.safetensors",
+                    "s3gen.pt",
+                    "grapheme_mtl_merged_expanded_v1.json",
+                    "conds.pt",
+                    "Cangjie5_TC.json",
+                ],
                 token=os.getenv("HF_TOKEN"),
             )
         )
         return cls.from_local(ckpt_dir, device)
-    
+
+    @classmethod
+    def from_checkpoint(
+        cls, save_dir, device: torch.device
+    ) -> "ChatterboxMultilingualTTS":
+        ckpt_dir = Path(
+            snapshot_download(
+                repo_id=REPO_ID,
+                repo_type="model",
+                revision="main",
+                allow_patterns=[
+                    "ve.pt",
+                    "t3_mtl23ls_v2.safetensors",
+                    "s3gen.pt",
+                    "grapheme_mtl_merged_expanded_v1.json",
+                    "conds.pt",
+                    "Cangjie5_TC.json",
+                ],
+                token=os.getenv("HF_TOKEN"),
+            )
+        )
+        ckpt_dir = Path(ckpt_dir)
+
+        ve = VoiceEncoder()
+        ve.load_state_dict(torch.load(ckpt_dir / "ve.pt", weights_only=True))
+        ve.to(device).eval()
+
+        t3 = T3(T3Config.multilingual())
+        t3_state = load_safetensors(save_dir + "t3_mtl23ls_v2.safetensors")
+        if "model" in t3_state.keys():
+            t3_state = t3_state["model"][0]
+        t3.load_state_dict(t3_state)
+        t3.to(device).eval()
+
+        s3gen = S3Gen()
+        s3gen.load_state_dict(torch.load(ckpt_dir / "s3gen.pt", weights_only=True))
+        s3gen.to(device).eval()
+
+        tokenizer = MTLTokenizer(str(ckpt_dir / "grapheme_mtl_merged_expanded_v1.json"))
+
+        conds = Conditionals.load(save_dir + "conds.pt").to(device)
+
+        return cls(t3, s3gen, ve, tokenizer, device, conds=conds)
+
     def prepare_conditionals(self, wav_fpath, exaggeration=0.5):
         ## Load reference wav
         s3gen_ref_wav, _sr = librosa.load(wav_fpath, sr=S3GEN_SR)
 
         ref_16k_wav = librosa.resample(s3gen_ref_wav, orig_sr=S3GEN_SR, target_sr=S3_SR)
 
-        s3gen_ref_wav = s3gen_ref_wav[:self.DEC_COND_LEN]
-        s3gen_ref_dict = self.s3gen.embed_ref(s3gen_ref_wav, S3GEN_SR, device=self.device)
+        s3gen_ref_wav = s3gen_ref_wav[: self.DEC_COND_LEN]
+        s3gen_ref_dict = self.s3gen.embed_ref(
+            s3gen_ref_wav, S3GEN_SR, device=self.device
+        )
 
         # Speech cond prompt tokens
         t3_cond_prompt_tokens = None
         if plen := self.t3.hp.speech_cond_prompt_len:
             s3_tokzr = self.s3gen.tokenizer
-            t3_cond_prompt_tokens, _ = s3_tokzr.forward([ref_16k_wav[:self.ENC_COND_LEN]], max_len=plen)
-            t3_cond_prompt_tokens = torch.atleast_2d(t3_cond_prompt_tokens).to(self.device)
+            t3_cond_prompt_tokens, _ = s3_tokzr.forward(
+                [ref_16k_wav[: self.ENC_COND_LEN]], max_len=plen
+            )
+            t3_cond_prompt_tokens = torch.atleast_2d(t3_cond_prompt_tokens).to(
+                self.device
+            )
 
         # Voice-encoder speaker embedding
-        ve_embed = torch.from_numpy(self.ve.embeds_from_wavs([ref_16k_wav], sample_rate=S3_SR))
+        ve_embed = torch.from_numpy(
+            self.ve.embeds_from_wavs([ref_16k_wav], sample_rate=S3_SR)
+        )
         ve_embed = ve_embed.mean(axis=0, keepdim=True).to(self.device)
 
         t3_cond = T3Cond(
@@ -249,11 +299,13 @@ class ChatterboxMultilingualTTS:
                 f"Unsupported language_id '{language_id}'. "
                 f"Supported languages: {supported_langs}"
             )
-        
+
         if audio_prompt_path:
             self.prepare_conditionals(audio_prompt_path, exaggeration=exaggeration)
         else:
-            assert self.conds is not None, "Please `prepare_conditionals` first or specify `audio_prompt_path`"
+            assert (
+                self.conds is not None
+            ), "Please `prepare_conditionals` first or specify `audio_prompt_path`"
 
         # Update exaggeration if needed
         if float(exaggeration) != float(self.conds.t3.emotion_adv[0, 0, 0].item()):
@@ -266,8 +318,12 @@ class ChatterboxMultilingualTTS:
 
         # Norm and tokenize text
         text = punc_norm(text)
-        text_tokens = self.tokenizer.text_to_tokens(text, language_id=language_id.lower() if language_id else None).to(self.device)
-        text_tokens = torch.cat([text_tokens, text_tokens], dim=0)  # Need two seqs for CFG
+        text_tokens = self.tokenizer.text_to_tokens(
+            text, language_id=language_id.lower() if language_id else None
+        ).to(self.device)
+        text_tokens = torch.cat(
+            [text_tokens, text_tokens], dim=0
+        )  # Need two seqs for CFG
 
         sot = self.t3.hp.start_text_token
         eot = self.t3.hp.stop_text_token
@@ -297,5 +353,5 @@ class ChatterboxMultilingualTTS:
                 ref_dict=self.conds.gen,
             )
             wav = wav.squeeze(0).detach().cpu().numpy()
-            watermarked_wav = self.watermarker.apply_watermark(wav, sample_rate=self.sr)
-        return torch.from_numpy(watermarked_wav).unsqueeze(0)
+            # wav = self.watermarker.apply_watermark(wav, sample_rate=self.sr)
+        return torch.from_numpy(wav).unsqueeze(0)