Update LLMwithvoice.py

LLMwithvoice.py

@@ -1,10 +1,23 @@
 import requests
-from IPython.display import Audio
+import torch
+import soundfile as sf
+from transformers import AutoTokenizer
+from parler_tts import ParlerTTSForConditionalGeneration
 
 # Hugging Face API URL for Roberta model
 API_URL_ROBERTA = "https://api-inference.huggingface.co/models/deepset/roberta-base-squad2"
-# Hugging Face API URL for text-to-speech model
-API_URL_TTS = "https://api-inference.huggingface.co/models/suno/bark"
+
+device = "cpu"
+if torch.cuda.is_available():
+    device = "cuda:0"
+if torch.backends.mps.is_available():
+    device = "mps"
+if torch.xpu.is_available():
+    device = "xpu"
+torch_dtype = torch.float16 if device != "cpu" else torch.float32
+
+model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler_tts_mini_v0.1").to(device, dtype=torch_dtype)
+tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler_tts_mini_v0.1")
 
 def query_roberta(api_token, payload):
     headers = {"Authorization": f"Bearer {api_token}"}
@@ -14,10 +27,15 @@ def query_roberta(api_token, payload):
     except ValueError:
         return {"error": "Invalid JSON response"}
 
-def query_tts(api_token, payload):
-    headers = {"Authorization": f"Bearer {api_token}"}
-    response = requests.post(API_URL_TTS, headers=headers, json=payload)
-    return response.content
+def generate_speech(prompt, description):
+    input_ids = tokenizer(description, return_tensors="pt").input_ids.to(device)
+    prompt_input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+
+    generation = model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids).to(torch.float32)
+    audio_arr = generation.cpu().numpy().squeeze()
+    audio_path = "/mnt/data/parler_tts_out.wav"
+    sf.write(audio_path, audio_arr, model.config.sampling_rate)
+    return audio_path
 
 def chat_with_roberta(api_token, question, context):
     payload = {
@@ -35,12 +53,8 @@ def chat_with_roberta(api_token, question, context):
         except (IndexError, KeyError):
             return f"Unexpected response structure: {response}"
 
-def generate_speech(api_token, text):
-    payload = {"inputs": text}
-    audio_bytes = query_tts(api_token, payload)
-    return audio_bytes
-
 def gradio_interface(api_token, context, question):
     answer = chat_with_roberta(api_token, question, context)
-    audio_bytes = generate_speech(api_token, answer)
-    return answer, audio_bytes
+    description = "A female speaker with a slightly low-pitched voice delivers her words quite expressively, in a very confined sounding environment with clear audio quality. She speaks very fast."
+    audio_path = generate_speech(answer, description)
+    return answer, audio_path