adding token access

app.py

@@ -1,24 +1,34 @@
 import os
+import torch
 import streamlit as st
 from diffusers import StableDiffusionPipeline
 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")
+DEVICE_NAME = os.getenv("DEVICE_NAME", "cpu")
+HUGGING_FACE_TOKEN = os.getenv("HUGGING_FACE_TOKEN")
 
 
 def load_translation_models(translation_model_id):
-    tokenizer = MBart50TokenizerFast.from_pretrained(translation_model_id)
+    tokenizer = MBart50TokenizerFast.from_pretrained(
+        translation_model_id,
+        use_auth_token=HUGGING_FACE_TOKEN
+    )
     tokenizer.src_lang = 'pt_XX'
-    text_model = \
-        MBartForConditionalGeneration.from_pretrained(translation_model_id)
+    text_model = MBartForConditionalGeneration.from_pretrained(
+        translation_model_id,
+        use_auth_token=HUGGING_FACE_TOKEN
+    )
 
     return tokenizer, text_model
 
 
 def pipeline_generate(diffusion_model_id):
-    pipe = StableDiffusionPipeline.from_pretrained(diffusion_model_id)
+    pipe = StableDiffusionPipeline.from_pretrained(
+        diffusion_model_id, torch_dtype=torch.float16, revision="fp16",
+        use_auth_token=HUGGING_FACE_TOKEN
+    )
     pipe = pipe.to(DEVICE_NAME)
 
     # Recommended if your computer has < 64 GB of RAM