Update app.py

app.py

@@ -7,6 +7,29 @@ import requests
 from genai_chat_ai import AI,create_chat_session
 api_key = os.environ.get("Id_mode_vits") 
 headers = {"Authorization": f"Bearer {api_key}"}
+from transformers import pipeline
+from transformers import AutoTokenizer,VitsModel
+import torch
+models= {}
+tokenizer = AutoTokenizer.from_pretrained("asg2024/vits-ar-sa-huba",token=api_key)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+def  get_model(name_model):
+    global models
+    if name_model in not  models:
+        models[name_model]=VitsModel.from_pretrained(name_model,token=api_key).to(device)
+    return models[name_model]
+
+def  genrate_speech(text,name_model):
+    inputs=tokenizer(text,return_tensors="pt")
+    model=get_model(name_model) 
+    with torch.no_grad():
+         wav=model(
+             input_ids= input_ids.input_ids.to(device),
+             attention_mask=input_ids.attention_mask.to(device),
+             speaker_id=0
+             ).waveform.cpu().numpy().reshape(-1)
+    return model.config.sampling_rate,wav
+     
 
 def remove_extra_spaces(text):
   """
@@ -69,9 +92,9 @@ with gr.Blocks() as demo:  # Use gr.Blocks to wrap the entire interface
             API_URL = f"https://api-inference.huggingface.co/models/{model_choice}"
             text_answer = get_answer_ai(text)
             text_answer = remove_extra_spaces(text_answer)
-            data_ai = query(text_answer, API_URL)
+            data_ai = genrate_speech(text_answer,model_choice)#query(text_answer, API_URL)
             if generate_user_audio:  # Generate user audio if needed
-                data_user = query(text, API_URL)
+                data_user =genrate_speech(text_answer,model_choice)# query(text, API_URL)
                 return data_user, data_ai, text_answer
             else:
                 return  data_ai  # Return None for user_audio