Update src/chatterbox/mtl_tts.py

src/chatterbox/mtl_tts.py

@@ -64,9 +64,7 @@ class ChatterboxMultilingualTTS:
     ENC_COND_LEN = 6 * S3_SR
     DEC_COND_LEN = 10 * S3GEN_SR
 
-    def __init__(self, t3: T3, s3gen: S3Gen, ve: VoiceEncoder,
-                 tokenizer: MTLTokenizer, device: torch.device,
-                 conds: Conditionals = None):
+    def __init__(self, t3, s3gen, ve, tokenizer, device, conds=None):
         self.sr = S3GEN_SR
         self.t3 = t3
         self.s3gen = s3gen
@@ -77,13 +75,7 @@ class ChatterboxMultilingualTTS:
         self.watermarker = perth.PerthImplicitWatermarker()
 
     @classmethod
-    def get_supported_languages(cls):
-        return SUPPORTED_LANGUAGES.copy()
-
-    @classmethod
-    def from_local(cls, ckpt_dir: Path, device: torch.device) -> "ChatterboxMultilingualTTS":
-        ckpt_dir = Path(ckpt_dir)
-
+    def from_local(cls, ckpt_dir, device=torch.device("cpu")):
         ve = VoiceEncoder()
         ve.load_state_dict(torch.load(ckpt_dir / "ve.pt", weights_only=True))
         ve.to(device).eval()
@@ -105,10 +97,10 @@ class ChatterboxMultilingualTTS:
         if (builtin_voice := ckpt_dir / "conds.pt").exists():
             conds = Conditionals.load(builtin_voice)
 
-        return cls(t3, s3gen, ve, tokenizer, device, conds=conds)
+        return cls(t3, s3gen, ve, tokenizer, device, conds)
 
     @classmethod
-    def from_pretrained(cls, device: str | torch.device | None = None) -> "ChatterboxMultilingualTTS":
+    def from_pretrained(cls, device=None):
         if device is None:
             device = torch.device("cpu")
         elif isinstance(device, str):
@@ -122,35 +114,36 @@ class ChatterboxMultilingualTTS:
             revision="main",
             allow_patterns=[
                 "ve.pt", "t3_mtl23ls_v2.safetensors", "s3gen.pt",
-                "grapheme_mtl_merged_expanded_v1.json", "conds.pt",
-                "Cangjie5_TC.json"
+                "grapheme_mtl_merged_expanded_v1.json", "conds.pt"
             ],
             token=os.getenv("HF_TOKEN"),
         ))
 
-        model = cls.from_local(ckpt_dir, device)
-
-        # Ensure all params on CPU and eval
-        model.t3.to(device).eval()
-        model.s3gen.to(device).eval()
-        model.ve.to(device).eval()
-        if model.conds:
-            for k, v in model.conds.gen.items():
-                if torch.is_tensor(v):
-                    model.conds.gen[k] = v.to(device)
-        return model
+        return cls.from_local(ckpt_dir, device)
 
     @torch.no_grad()
-    def generate(self, text: str, speaker_embedding: torch.Tensor = None) -> torch.Tensor:
+    def generate(self, text: str, speaker_embedding=None, language_id=None, **kwargs):
         """
-        Generate audio waveform (numpy array) from text.
-        CPU-compatible.
+        CPU-safe text-to-speech.
+        Accepts optional `language_id` and any other kwargs.
         """
-        text = punc_norm(text)
+        # Normalize punctuation
+        text = text.strip()
+        if not text.endswith("."):
+            text += "."
+
+        # Encode text
         token_ids = self.tokenizer.encode(text)
         token_ids = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0).to(self.device)
 
         conds = self.conds.gen if self.conds else {}
+
+        # Include language_id in conds if provided
+        if language_id is not None:
+            conds = conds.copy()
+            conds['language_id'] = language_id
+
+        # Run through T3 and S3Gen
         t3_out = self.t3(token_ids, **conds)
         audio = self.s3gen(t3_out, **conds)