Update app.py

app.py

@@ -6,7 +6,9 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 from deepface import DeepFace
 import time
-from transformers import AutoProcessor, DiaForConditionalGeneration
+from kokoro import KPipeline
+from IPython.display import display, Audio
+import soundfile as sf
 from sentence_transformers import SentenceTransformer
 import numpy as np
 import chromadb
@@ -44,10 +46,7 @@ print("Model loaded successfully!")
 
 
 print("Loading Dia TTS Model...")
-tts_device = "cuda:0" if torch.cuda.is_available() else "cpu"
-tts_model = "nari-labs/Dia-1.6B-0626"
-tts_processor = AutoProcessor.from_pretrained(tts_model)
-tts_model = DiaForConditionalGeneration.from_pretrained(tts_model, torch_dtype=torch.float16).to(tts_device)
+tts_pipeline = KPipeline(lang_code='b')
 print("Dia TTS Model loaded successfully!")
 
 print("Initializing Memory Components...")
@@ -479,50 +478,31 @@ You are Kenko, a compassionate mental health therapist. Provide empathetic, help
         return f"I'm sorry, I'm having trouble processing your message right now. Error: {str(e)}"
 
 def generate_tts(text):
-    try:
-        text = text[:600]
-
-        print(f"[TTS] Generating speech for {len(text)} chars: '{text[:50]}...'")
-
-        inputs = tts_processor(text=text, return_tensors="pt", padding=True)
-        inputs = {k: v.to(tts_device) for k, v in inputs.items()}
-
-        print(f"[TTS] Inputs prepared, generating audio codes...")
-
-        with torch.no_grad():
-            generated_ids = tts_model.generate(**inputs, max_length=2500)
+  try:
+    text = text[:600]
 
-        print(f"[TTS] Audio codes generated, shape: {generated_ids.shape}")
-        print(f"[TTS] Decoding codes to waveform...")
+    generator = tts_pipeline(
+      text, voice='af_heart',
+      speed=1, split_pattern=r'\n+'
+    )
 
-        audio_values = tts_processor.batch_decode(generated_ids, return_tensors="pt")
+    audio_chunks = []
+    for gs, ps, audio in generator:
+      audio_chunks.append(audio)
 
-        if isinstance(audio_values, dict) and 'audio_values' in audio_values:
-            audio_arr = audio_values['audio_values'][0].cpu().numpy()
-        elif isinstance(audio_values, torch.Tensor):
-            audio_arr = audio_values[0].cpu().numpy()
-        elif isinstance(audio_values, list):
-            audio_arr = np.array(audio_values[0])
-        else:
-            audio_arr = np.array(audio_values).squeeze()
-
-        audio_arr = audio_arr.astype(np.float32)
-
-        sample_rate = 44100
+    if not audio_chunks:
+      print("TTS generation failed")
+      return None
 
-        print(f"[TTS] Audio decoded: {len(audio_arr)} samples at {sample_rate}Hz = {len(audio_arr)/sample_rate:.2f} seconds")
+    audio_array = np.concatenate(audio_chunks, axis=0)
+    audio_array = audio_array.astype(np.float32)
+    sample_rate = 24000
 
-        if len(audio_arr) == 0:
-            print("Decoded audio is empty!")
-            return None
+    return (sample_rate, audio_array)
 
-        return (sample_rate, audio_arr)
-
-    except Exception as e:
-        print(f"TTS generation error: {str(e)}")
-        import traceback
-        traceback.print_exc()
-        return None
+  except Exception as e:
+    print(f"TTS generation error: {str(e)}")
+    return None
 
 css = """
 .gradio-container {