Update app.py

app.py

@@ -1,13 +1,15 @@
-import streamlit as st
 import os
+import streamlit as st
 from dotenv import load_dotenv
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from diffusers import StableDiffusionPipeline
 import torch
 
+# Load environment variables
+load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Set page config
+# Set Streamlit page config
 st.set_page_config(
     page_title="Tamil Creative Studio",
     page_icon="🇮🇳",
@@ -22,30 +24,27 @@ def load_css(file_name):
 
 @st.cache_resource(show_spinner=False)
 def load_all_models():
-    # Load translation model
+    # Load translation model (private)
     trans_tokenizer = AutoTokenizer.from_pretrained(
-        "ai4bharat/indictrans2-ta-en-dist-200M", token=HF_TOKEN
+        "ai4bharat/indictrans2-ta-en-dist-200M",
+        use_auth_token=HF_TOKEN
     )
     trans_model = AutoModelForSeq2SeqLM.from_pretrained(
-        "ai4bharat/indictrans2-ta-en-dist-200M", token=HF_TOKEN
+        "ai4bharat/indictrans2-ta-en-dist-200M",
+        use_auth_token=HF_TOKEN
     )
-
+    
     # Load text generation model
-    text_gen = pipeline(
-        "text-generation",
-        model="gpt2",
-        device=-1,
-        token=HF_TOKEN
-    )
-
+    text_gen = pipeline("text-generation", model="gpt2", device=-1)
+    
     # Load image generation model
     img_pipe = StableDiffusionPipeline.from_pretrained(
         "stabilityai/stable-diffusion-2-base",
+        use_auth_token=HF_TOKEN,
         torch_dtype=torch.float32,
-        safety_checker=None,
-        token=HF_TOKEN
+        safety_checker=None
     ).to("cpu")
-
+    
     return trans_tokenizer, trans_model, text_gen, img_pipe
 
 def translate_tamil(text, tokenizer, model):
@@ -56,74 +55,60 @@ 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")
-
-    st.markdown(
-        """
+    
+    # Header
+    st.markdown("""
         <div class="header">
             <h1>🌐 தமிழ் → English → Creative Studio</h1>
             <p>Translate Tamil text and generate creative content</p>
         </div>
-        """,
-        unsafe_allow_html=True
-    )
-
+    """, 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)
-
-                with st.spinner("படத்தை உருவாக்குதல்..."):
-                    img = img_pipe(
-                        eng,
-                        num_inference_steps=40,
-                        guidance_scale=8.5
-                    ).images[0]
-
-                st.image(img, caption="**🎨 Generated Image**", use_column_width=True)
+    
+    tamil_text = st.text_area("**தமிழ் உரை:**", height=150, placeholder="உங்கள் உரையை இங்கே உள்ளிடவும்...")
+    
+    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)
+
+        with st.spinner("படத்தை உருவாக்குதல்..."):
+            img = img_pipe(eng, num_inference_steps=40, guidance_scale=8.5).images[0]
+
+        st.image(img, caption="🎨 Generated Image", use_column_width=True)
 
 if __name__ == "__main__":
     main()