Update app.py

app.py

@@ -20,7 +20,7 @@ import tempfile
 from huggingface_hub import hf_hub_download
 from sentence_transformers import SentenceTransformer
 from huggingface_hub import InferenceClient
-from transformers import VitsModel, AutoTokenizer, pipeline
+from transformers import VitsModel, AutoTokenizer, pipeline, SpeechT5HifiGan
 
 
 # ── Auth ───────────────────────────────────────────────────────
@@ -104,47 +104,52 @@ def detect_language(text):
         return "Español"
         
 # ── TTS: Parler TTS mini v1 (neutral català/spanish voice) ─────────
-print("Loading MMS TTS models...")
-tts_models, tts_tokenizers = {}, {}
-for lang_code, repo in {"en": "facebook/mms-tts-eng", "es": "facebook/mms-tts-spa", "ca": "facebook/mms-tts-cat"}.items():
-    tts_tokenizers[lang_code] = AutoTokenizer.from_pretrained(repo)
-    tts_models[lang_code] = VitsModel.from_pretrained(repo).to(device)
-    tts_models[lang_code].eval()
+#print("Loading MMS TTS models...")
+#tts_models, tts_tokenizers = {}, {}
+#for lang_code, repo in {"en": "facebook/mms-tts-eng", "es": "facebook/mms-tts-spa", "ca": "facebook/mms-tts-cat"}.items():
+#    tts_tokenizers[lang_code] = AutoTokenizer.from_pretrained(repo)
+#    tts_models[lang_code] = VitsModel.from_pretrained(repo).to(device)
+#    tts_models[lang_code].eval()
+print("Loading TTS models...")
+
+# Kokoro for English and Spanish
+from kokoro import KPipeline
+kokoro_en = KPipeline(lang_code='en')
+kokoro_es = KPipeline(lang_code='es')
+
+# Matxa (BSC) for Catalan
+tts_tokenizers, tts_models = {}, {}
+tts_tokenizers["ca"] = AutoTokenizer.from_pretrained("projecte-aina/matxa-tts-cat-multiaccent")
+tts_models["ca"] = VitsModel.from_pretrained("projecte-aina/matxa-tts-cat-multiaccent").to(device)
+tts_models["ca"].eval()
 
 def text_to_speech(text, lang="es"):
-    if not text or lang not in tts_models:
+    if not text:
         return None
     try:
-        inputs = tts_tokenizers[lang](text, return_tensors="pt").to(device)
-        with torch.no_grad():
-            audio = tts_models[lang](**inputs).waveform
-        audio_int16 = (audio.squeeze().cpu().float().numpy() * 32767).clip(-32768, 32767).astype("int16")
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
-            scipy.io.wavfile.write(f.name, rate=tts_models[lang].config.sampling_rate, data=audio_int16)
-            return f.name
+        if lang == "ca":
+            inputs = tts_tokenizers["ca"](text, return_tensors="pt").to(device)
+            with torch.no_grad():
+                audio = tts_models["ca"](**inputs).waveform
+            audio_int16 = (audio.squeeze().cpu().float().numpy() * 32767).clip(-32768, 32767).astype("int16")
+            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+                scipy.io.wavfile.write(f.name, rate=tts_models["ca"].config.sampling_rate, data=audio_int16)
+                return f.name
+        else:
+            pipeline = kokoro_en if lang == "en" else kokoro_es
+            voice = "af_heart" if lang == "en" else "ef_dora"
+            audio_chunks = []
+            for _, _, audio in pipeline(text, voice=voice):
+                audio_chunks.append(audio)
+            audio_np = np.concatenate(audio_chunks)
+            audio_int16 = (audio_np * 32767).clip(-32768, 32767).astype("int16")
+            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+                scipy.io.wavfile.write(f.name, rate=24000, data=audio_int16)
+                return f.name
     except Exception as e:
-        print(f"TTS error: {e}")
+        print(f"TTS error ({lang}): {e}")
         return None
 
-    try:
-        input_ids = tts_tokenizer(voice_desc, return_tensors="pt").input_ids.to(device)
-        prompt_ids = tts_tokenizer(text, return_tensors="pt").input_ids.to(device)
-
-        with torch.no_grad():
-            generation = tts_model.generate(
-                input_ids=input_ids,
-                prompt_input_ids=prompt_ids,
-            )
-
-        audio_array = generation.cpu().to(torch.float32).numpy().squeeze()
-
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
-            scipy.io.wavfile.write(f.name, rate=sampling_rate, data=audio_array)
-            return f.name
-
-    except Exception as e:
-        print(f"TTS error: {e}")
-        return None
 
 # ── LLM: HF Inference API + RAG ───────────────────────────────
 SYSTEM_PROMPT = """You are a warm, calm, and knowledgeable support assistant for caregivers of people with Alzheimer's disease.