Update app.py

app.py

@@ -1,8 +1,13 @@
 import streamlit as st
+import os
+from dotenv import load_dotenv
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from diffusers import StableDiffusionPipeline
 import torch
-import base64
+
+# Load environment variables
+load_dotenv()
+HF_TOKEN = os.getenv("HF_TOKEN")
 
 # Set page config
 st.set_page_config(
@@ -20,19 +25,29 @@ def load_css(file_name):
 @st.cache_resource(show_spinner=False)
 def load_all_models():
     # Load translation model
-    trans_tokenizer = AutoTokenizer.from_pretrained("ai4bharat/indictrans2-ta-en-dist-200M")
-    trans_model = AutoModelForSeq2SeqLM.from_pretrained("ai4bharat/indictrans2-ta-en-dist-200M")
-    
+    trans_tokenizer = AutoTokenizer.from_pretrained(
+        "ai4bharat/indictrans2-ta-en-dist-200M", token=HF_TOKEN
+    )
+    trans_model = AutoModelForSeq2SeqLM.from_pretrained(
+        "ai4bharat/indictrans2-ta-en-dist-200M", token=HF_TOKEN
+    )
+
     # Load text generation model
-    text_gen = pipeline("text-generation", model="gpt2", device=-1)
-    
+    text_gen = pipeline(
+        "text-generation",
+        model="gpt2",
+        device=-1,
+        token=HF_TOKEN
+    )
+
     # Load image generation model
     img_pipe = StableDiffusionPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-2-base", 
+        "stabilityai/stable-diffusion-2-base",
         torch_dtype=torch.float32,
-        safety_checker=None
+        safety_checker=None,
+        token=HF_TOKEN
     ).to("cpu")
-    
+
     return trans_tokenizer, trans_model, text_gen, img_pipe
 
 def translate_tamil(text, tokenizer, model):
@@ -43,24 +58,23 @@ def translate_tamil(text, tokenizer, model):
         truncation=True,
         max_length=128
     )
-    
+
     generated = model.generate(
         **inputs,
         max_length=150,
         num_beams=5,
         early_stopping=True
     )
-    
+
     return tokenizer.batch_decode(
-        generated, 
+        generated,
         skip_special_tokens=True,
         clean_up_tokenization_spaces=True
     )[0]
 
 def main():
     load_css("style.css")
-    
-    # Header with background
+
     st.markdown(
         """
         <div class="header">
@@ -70,37 +84,37 @@ def main():
         """,
         unsafe_allow_html=True
     )
-    
+
     tokenizer, model, text_gen, img_pipe = load_all_models()
-    
+
     with st.container():
         tamil_text = st.text_area(
-            "**தமிழ் உரை:**", 
+            "**தமிழ் உரை:**",
             height=150,
             placeholder="உங்கள் உரையை இங்கே உள்ளிடவும்...",
             key="tamil_input"
         )
-        
+
         col1, col2 = st.columns([1, 3])
         with col1:
             if st.button("**உருவாக்கு**", type="primary", use_container_width=True):
                 if not tamil_text.strip():
                     st.warning("தயவு செய்து உரையை உள்ளிடவும்.")
                     st.stop()
-                
+
                 with st.spinner("மொழிபெயர்க்கிறது..."):
                     eng = translate_tamil(tamil_text, tokenizer, model)
-                
+
                 with st.expander("**🔤 Translation**", expanded=True):
                     st.success(eng)
-                
+
                 with st.spinner("உரை உருவாக்குதல்..."):
                     creative = text_gen(
                         f"Create a creative description about: {eng}",
                         max_length=80,
                         num_return_sequences=1
                     )[0]["generated_text"]
-                
+
                 st.info("**📝 Creative Text:**")
                 st.write(creative)
 
@@ -110,8 +124,8 @@ def main():
                         num_inference_steps=40,
                         guidance_scale=8.5
                     ).images[0]
-                
+
                 st.image(img, caption="**🎨 Generated Image**", use_column_width=True)
 
 if __name__ == "__main__":
-    main()
+    main()