changed model translation

app.py

@@ -4,7 +4,7 @@ import gradio as gr
 from transformers import pipeline
 #from googletrans import Translator
 import os
-
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
 
 if torch.cuda.is_available():
@@ -27,14 +27,14 @@ import base64
 with open("Iso_Logotipo_Ceibal.png", "rb") as image_file:
     encoded_image = base64.b64encode(image_file.read()).decode()
 
+CKPT = "facebook/nllb-200-distilled-600M"
 
+model = AutoModelForSeq2SeqLM.from_pretrained(CKPT)
+tokenizer = AutoTokenizer.from_pretrained(CKPT)
 def generate_sound(text,steps,audio_length,negative_prompt): 
     print(text)
-    # text=translate_text(text)
-    text = translate_text(text)
+    text=translate_text(text)
     negative_prompt = translate_text(negative_prompt)
-    #translator = Translator()
-    #text=translator.translate(text, src='es',dest="en").text
     print(text)
     waveforms = pipe(text, 
                      num_inference_steps=steps, 
@@ -42,15 +42,21 @@ def generate_sound(text,steps,audio_length,negative_prompt):
                      negative_prompt = negative_prompt).audios
     rate =16000
     return rate, waveforms[0]
-    #return gr.make_waveform((rate, waveforms[0]))
-
-es_en_translator = pipeline("translation",model = "Helsinki-NLP/opus-mt-es-en")
-
 
 def translate_text(text):
-    text = es_en_translator(text)[0].get("translation_text")
-    return text
+    translation_pipeline = pipeline("translation",
+                                    model=model,
+                                    tokenizer=tokenizer,
+                                    src_lang="spa_Latn",
+                                    tgt_lang="eng_Latn",
+                                    max_length=400,
+                                    device=device)
 
+    result = translation_pipeline(text)
+    return result[0]['translation_text']
+# def translate_text(text):
+#     text = es_en_translator(text)[0].get("translation_text")
+#     return text
 with gr.Blocks() as demo:
     gr.Markdown("""
     <center>
@@ -90,4 +96,5 @@ with gr.Blocks() as demo:
     btn.click(fn=generate_sound, inputs=[prompt,steps,audio_len,negative_prompt], outputs=[output])  #steps,guidance,width,height]
 
 gr.close_all()
+demo.queue()
 demo.launch()