change from chaterbox to gwen-tts

src/anki_hsk_creator/utility.py

@@ -12,10 +12,10 @@ import argostranslate.package
 import argostranslate.translate
 import torch
 from cedict_utils.cedict import CedictEntry, CedictParser
-from chatterbox.mtl_tts import ChatterboxMultilingualTTS
+from qwen_tts import Qwen3TTSModel
 
 # Local
-from .constants import CCCEDICT_PATH, INPUT, LANGUAGES, OUTPUT, RESOURCES
+from .constants import CCCEDICT_PATH, GWEN_TTS, INPUT, LANGUAGES, OUTPUT, RESOURCES
 
 # Static Clases
 
@@ -50,10 +50,12 @@ class TRANS:
                     if out_package.to_code == to_code:
                         if in_package.to_code == out_package.from_code:
                             print(
-                                f"Check in_package {in_package.from_code} {in_package.to_code}"
+                                f"Check in_package {in_package.from_code}"
+                                f"{in_package.to_code}"
                             )
                             print(
-                                f"Check out_package {out_package.from_code} {out_package.to_code}"
+                                f"Check out_package {out_package.from_code}"
+                                f"{out_package.to_code}"
                             )
                             packages_to_install.append(in_package)
                             packages_to_install.append(out_package)
@@ -137,7 +139,11 @@ class TTS:
 
     MODEL = None
     DEVICE = None
-    DEFAULTS = {"cfg_weight": 0.2, "exaggeration": 0.8}
+    DEFAULTS = {
+        "language": "Chinese",
+        "speaker": "Uncle_Fu",
+        "instruct": "语速缓慢而审慎，每个音节的语调都拿捏得恰到好处，宛如教授在指导新生。",
+    }
 
     @staticmethod
     def create_tts():
@@ -151,10 +157,18 @@ class TTS:
             else:
                 TTS.DEVICE = "cpu"
         if TTS.MODEL is None:
-            TTS.MODEL = ChatterboxMultilingualTTS.from_pretrained(
-                device=TTS.DEVICE, t3_model="v3"
+            TTS.MODEL = Qwen3TTSModel.from_pretrained(
+                GWEN_TTS,
+                device_map=TTS.DEVICE,
+                dtype=torch.bfloat16,
+                # attn_implementation="flash_attention_2",
             )
 
+    @staticmethod
+    def generate(text: str):
+        """Generates a Waw using the defaulst values"""
+        return TTS.MODEL.generate_custom_voice(text=text, **TTS.DEFAULTS)
+
 
 # Clases