Update app.py

app.py

@@ -2,6 +2,7 @@ from transformers import MBartForConditionalGeneration, MBart50Tokenizer, AutoMo
 import gradio as gr
 import torch
 from diffusers import FluxPipeline
+import os
 
 # Load the translation model and tokenizer
 model_name = "facebook/mbart-large-50-many-to-one-mmt"
@@ -16,8 +17,15 @@ text_model = AutoModelForCausalLM.from_pretrained(text_generation_model_name)
 # Create a pipeline for text generation using the selected model
 text_generator = pipeline("text-generation", model=text_model, tokenizer=text_tokenizer)
 
-# Set up the new FluxPipeline for the text-to-image model
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+# Get the Hugging Face API token from environment variables
+hf_token = os.getenv("HF_TOKEN")
+
+# Authenticate and set up the new FluxPipeline for the text-to-image model
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    use_auth_token=hf_token,  # Use the token for authentication
+    torch_dtype=torch.bfloat16
+)
 pipe.enable_model_cpu_offload()  # Enable CPU offloading to save GPU memory if needed
 
 # Function to generate an image using the new FluxPipeline model