Update app.py

app.py

@@ -123,28 +123,23 @@ def load_models():
 
 
 # --- Generate Response and Audio Function ---
-@spaces.GPU # Decorate with @spaces.GPU as this function performs GPU-intensive inference
-def generate_response_and_audio(
-    message: str, # Current user message
-    history: list # Gradio Chatbot history format (list of dictionaries with 'role' and 'content')
-) -> tuple: # Returns (updated_history, audio_file_path)
-    """
-    Generates a text response from the loaded LLM and then converts it to audio
-    using the loaded TTS model.
-    """
+@spaces.GPU
+def generate_response_and_audio(message: str, history: list) -> tuple:
     global tokenizer, llm_model, tts_processor, tts_model, tts_vocoder, speaker_embeddings
 
-    # Initialize all models if not already loaded
     if tokenizer is None or llm_model is None or tts_model is None:
         load_models()
 
-    if tokenizer is None or llm_model is None: # Check LLM loading status
+    if tokenizer is None or llm_model is None:
         history.append({"role": "user", "content": message})
         history.append({"role": "assistant", "content": "Error: Chatbot LLM not loaded. Please check logs."})
         return history, None
 
+    # Initialize generated_text early
+    generated_text = ""
+
     # --- 1. Generate Text Response (LLM) ---
-    messages = history
+    messages = history.copy()
     messages.append({"role": "user", "content": message})
 
     try:
@@ -153,17 +148,15 @@ def generate_response_and_audio(
         print(f"Error applying chat template: {e}")
         input_text = ""
         for item in history:
-            if item["role"] == "user":
-                input_text += f"User: {item['content']}\n"
-            elif item["role"] == "assistant":
-                input_text += f"Assistant: {item['content']}\n"
+            input_text += f"{item['role'].capitalize()}: {item['content']}\n"
         input_text += f"User: {message}\nAssistant:"
 
+    try:
         input_ids = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True).to(llm_model.device)
         with torch.no_grad():
             output_ids = llm_model.generate(
                 input_ids["input_ids"],
-                attention_mask=input_ids["attention_mask"],  # <-- Add this line
+                attention_mask=input_ids["attention_mask"],
                 max_new_tokens=MAX_NEW_TOKENS,
                 do_sample=DO_SAMPLE,
                 temperature=TEMPERATURE,
@@ -172,6 +165,48 @@ def generate_response_and_audio(
                 pad_token_id=tokenizer.eos_token_id
             )
 
+        generated_token_ids = output_ids[0][input_ids["input_ids"].shape[-1]:]
+        generated_text = tokenizer.decode(generated_token_ids, skip_special_tokens=True).strip()
+
+    except Exception as e:
+        print(f"Error during LLM generation: {e}")
+        history.append({"role": "assistant", "content": "I encountered an error while generating a response."})
+        return history, None
+
+    # --- 2. Generate Audio from Response (TTS) ---
+    audio_path = None
+    if all([tts_processor, tts_model, tts_vocoder, speaker_embeddings]):
+        try:
+            device = llm_model.device if llm_model else 'cpu'
+            tts_model.to(device)
+            tts_vocoder.to(device)
+            speaker_embeddings = speaker_embeddings.to(device)
+
+            tts_inputs = tts_processor(
+                text=generated_text,
+                return_tensors="pt",
+                max_length=550,
+                truncation=True
+            ).to(device)
+
+            with torch.no_grad():
+                speech = tts_model.generate_speech(tts_inputs["input_ids"], speaker_embeddings, vocoder=tts_vocoder)
+
+            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_file:
+                audio_path = tmp_file.name
+                sf.write(audio_path, speech.cpu().numpy(), samplerate=16000)
+
+            print(f"Audio saved to: {audio_path}")
+
+        except Exception as e:
+            print(f"Error generating audio: {e}")
+            audio_path = None
+    else:
+        print("TTS components not fully loaded. Skipping audio generation.")
+
+    # --- 3. Update Chat History ---
+    history.append({"role": "assistant", "content": generated_text})
+    return history, audio_path
     # --- 2. Generate Audio from Response (TTS) ---
     audio_path = None
     if tts_processor and tts_model and tts_vocoder and speaker_embeddings is not None: