Update app.py

app.py

@@ -1,114 +1,74 @@
-import os
 import streamlit as st
-from dotenv import load_dotenv
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from diffusers import StableDiffusionPipeline
 import torch
 
-
-load_dotenv()
-# Set Streamlit page config
 st.set_page_config(
     page_title="Tamil Creative Studio",
     page_icon="🇮🇳",
     layout="centered",
-    initial_sidebar_state="collapsed"
 )
 
-# Load custom CSS
-def load_css(file_name):
-    with open(file_name, "r") as f:
-        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
+def load_css():
+    st.markdown(
+        """<style>
+        .header {
+            text-align: center;
+            padding: 20px;
+            background: #f9f9f9;
+            border-radius: 10px;
+            margin-bottom: 20px;
+        }
+        .header h1 { color: #cc0000; }
+        .header p { color: #333; font-style: italic; }
+        </style>""",
+        unsafe_allow_html=True,
+    )
 
 @st.cache_resource(show_spinner=False)
 def load_all_models():
-    HF_TOKEN = os.getenv("HF_TOKEN")  # Hugging Face token from Hugging Face Spaces secret
-
-    # Load translation model (private)
-    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
+    model_id = "ai4bharat/indictrans2-indic-en-dist-200M"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
     text_gen = pipeline("text-generation", model="gpt2", device=-1)
-
-    # Load image generation model
     img_pipe = StableDiffusionPipeline.from_pretrained(
         "stabilityai/stable-diffusion-2-base",
         torch_dtype=torch.float32,
         safety_checker=None
     ).to("cpu")
-
-    return trans_tokenizer, trans_model, text_gen, img_pipe
-
+    return tokenizer, model, text_gen, img_pipe
 
 def translate_tamil(text, tokenizer, model):
-    inputs = tokenizer(
-        text,
-        return_tensors="pt",
-        padding=True,
-        truncation=True,
-        max_length=128
-    )
-    
-    generated = model.generate(
-        **inputs,
-        max_length=150,
-        num_beams=5,
-        early_stopping=True
-    )
-    
-    return tokenizer.batch_decode(
-        generated, 
-        skip_special_tokens=True,
-        clean_up_tokenization_spaces=True
-    )[0]
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
+    outs = model.generate(**inputs, max_length=150, num_beams=5, early_stopping=True)
+    return tokenizer.decode(outs[0], skip_special_tokens=True)
 
 def main():
-    load_css("style.css")
-    
-    # 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)
-    
+    load_css()
+    st.markdown(
+        '<div class="header"><h1>🌐 தமிழ் → English → Creative Studio</h1>'
+        '<p>Translate Tamil text and generate creative content</p></div>',
+        unsafe_allow_html=True
+    )
     tokenizer, model, text_gen, img_pipe = load_all_models()
-    
     tamil_text = st.text_area("**தமிழ் உரை:**", height=150, placeholder="உங்கள் உரையை இங்கே உள்ளிடவும்...")
-    
-    if st.button("**உருவாக்கு**", type="primary", use_container_width=True):
+
+    if st.button("உருவாக்கு"):
         if not tamil_text.strip():
-            st.warning("தயவு செய்து உரையை உள்ளிடவும்.")
-            st.stop()
+            st.warning("உரையை உள்ளிடவும்.")
+            return
 
         with st.spinner("மொழிபெயர்க்கிறது..."):
             eng = translate_tamil(tamil_text, tokenizer, model)
-
-        with st.expander("🔤 Translation", expanded=True):
-            st.success(eng)
+        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)
+            creative = text_gen(f"Create a creative description about: {eng}", max_length=80, num_return_sequences=1)[0]["generated_text"]
+        st.info(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)
+        st.image(img, caption="Generated Image", use_column_width=True)
 
 if __name__ == "__main__":
     main()