changes for gwen tts

2026-06-26 19:11:41 +08:00
parent 5896f0db66
commit 777f4b4de1
6 changed files with 22 additions and 15 deletions
--- a/src/anki_hsk_creator/utility.py
+++ b/src/anki_hsk_creator/utility.py
@@ -151,7 +151,7 @@ class TTS:
        if TTS.DEVICE is None:
            # Automatically detect the best available device
            if torch.cuda.is_available():
-                TTS.DEVICE = "cuda"
+                TTS.DEVICE = "cuda:0"
            elif torch.backends.mps.is_available():
                TTS.DEVICE = "mps"
            else:
@@ -161,7 +161,7 @@ class TTS:
                GWEN_TTS,
                device_map=TTS.DEVICE,
                dtype=torch.bfloat16,
-                # attn_implementation="flash_attention_2",
+                attn_implementation="flash_attention_2",
            )

    @staticmethod