Use original S3Gen checkpoint

chatterbox/src/chatterbox/tts.py

@@ -24,7 +24,6 @@ from .models.t3.modules.cond_enc import T3Cond
 REPO_ID = "ResembleAI/chatterbox"
 BASE_REPO_ID = "ResembleAI/chatterbox"
 T3_FILENAME = "t3_mtl23ls_v3.safetensors"
-S3GEN_FILENAME = "s3gen_v3.safetensors"
 TOKENIZER_FILENAME = "grapheme_mtl_merged_expanded_v1.json"
 CANGJIE_FILENAME = "Cangjie5_TC.json"
 T3_TEXT_VOCAB_SIZE = 2454
@@ -239,8 +238,7 @@ class ChatterboxTTS:
 
         s3gen = S3Gen()
         s3gen.load_state_dict(
-            load_safetensors(ckpt_dir / S3GEN_FILENAME),
-            strict=False,
+            torch.load(ckpt_dir / "s3gen.pt", weights_only=True, map_location="cpu")
         )
         s3gen.to(device).eval()
 
@@ -258,7 +256,7 @@ class ChatterboxTTS:
     @classmethod
     def from_pretrained(cls, device: torch.device) -> 'ChatterboxTTS':
         token = os.getenv("HF_TOKEN")
-        base_files = ["ve.pt", S3GEN_FILENAME, TOKENIZER_FILENAME, CANGJIE_FILENAME, "conds.pt"]
+        base_files = ["ve.pt", "s3gen.pt", TOKENIZER_FILENAME, CANGJIE_FILENAME, "conds.pt"]
         base_dir = Path(
             snapshot_download(
                 repo_id=BASE_REPO_ID,