updating classes used for translation

app.py

@@ -1,17 +1,20 @@
+import os
 import streamlit as st
 from diffusers import StableDiffusionPipeline
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import MBart50TokenizerFast, MBartForConditionalGeneration
 
 DIFFUSION_MODEL_ID = "runwayml/stable-diffusion-v1-5"
 TRANSLATION_MODEL_ID = "Narrativa/mbart-large-50-finetuned-opus-pt-en-translation"  # noqa
+DEVICE_NAME = os.getenv("DEVICE_NAME", "cuda")
 
 
 def load_translation_models(translation_model_id):
-    tokenizer = AutoTokenizer.from_pretrained(
+    tokenizer = MBart50TokenizerFast.from_pretrained(
         translation_model_id,
         use_auth_token=True
     )
-    text_model = AutoModelForSeq2SeqLM.from_pretrained(
+    tokenizer.src_lang = 'pt_XX'
+    text_model = MBartForConditionalGeneration.from_pretrained(
         translation_model_id,
         use_auth_token=True
     )
@@ -24,7 +27,7 @@ def pipeline_generate(diffusion_model_id):
         diffusion_model_id,
         use_auth_token=True
     )
-    pipe = pipe.to("mps")
+    pipe = pipe.to(DEVICE_NAME)
 
     # Recommended if your computer has < 64 GB of RAM
     pipe.enable_attention_slicing()
@@ -39,7 +42,6 @@ def translate(prompt, tokenizer, text_model):
         num_beams=8, early_stopping=True
     )
     en_prompt = tokenizer.batch_decode(en_tokens, skip_special_tokens=True)
-    print(f"translation: [PT] {prompt} -> [EN] {en_prompt[0]}")
 
     return en_prompt[0]