Update app.py

app.py

@@ -2,24 +2,40 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from diffusers import StableDiffusionPipeline
 import torch
+import base64
+
+# Set 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)
 
 @st.cache_resource(show_spinner=False)
 def load_all_models():
-    # Load IndicTrans2 model for Tamil→English
-    indic_tokenizer = AutoTokenizer.from_pretrained("ai4bharat/indictrans2-ta-en-dist-200M")
-    indic_model = AutoModelForSeq2SeqLM.from_pretrained("ai4bharat/indictrans2-ta-en-dist-200M")
+    # 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")
     
-    # Load other models
+    # Load text generation model
     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
+        torch_dtype=torch.float32,
+        safety_checker=None
     ).to("cpu")
     
-    return indic_tokenizer, indic_model, text_gen, img_pipe
+    return trans_tokenizer, trans_model, text_gen, img_pipe
 
 def translate_tamil(text, tokenizer, model):
-    # Tokenize with batch processing
     inputs = tokenizer(
         text,
         return_tensors="pt",
@@ -28,7 +44,6 @@ def translate_tamil(text, tokenizer, model):
         max_length=128
     )
     
-    # Generate translation
     generated = model.generate(
         **inputs,
         max_length=150,
@@ -36,7 +51,6 @@ def translate_tamil(text, tokenizer, model):
         early_stopping=True
     )
     
-    # Decode with normalization
     return tokenizer.batch_decode(
         generated, 
         skip_special_tokens=True,
@@ -44,48 +58,60 @@ def translate_tamil(text, tokenizer, model):
     )[0]
 
 def main():
-    st.set_page_config(page_title="Tamil→EN→Image/Text", layout="centered")
-    st.title("🌐 தமிழ் → English → Creative Text & Image")
+    load_css("style.css")
+    
+    # Header with background
+    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 models once
     tokenizer, model, text_gen, img_pipe = load_all_models()
     
-    # Input with Tamil placeholder
-    tamil_text = st.text_area(
-        "தமிழ் உரை:", 
-        height=150,
-        placeholder="உங்கள் உரையை இங்கே உள்ளிடவும்..."
-    )
-
-    if st.button("உருவாக்கு", type="primary"):
-        if not tamil_text.strip():
-            st.warning("தயவு செய்து உரையை உள்ளிடவும்.")
-            return
-
-        with st.spinner("மொழிபெயர்க்கிறது..."):
-            eng = translate_tamil(tamil_text, tokenizer, model)
+    with st.container():
+        tamil_text = st.text_area(
+            "**தமிழ் உரை:**", 
+            height=150,
+            placeholder="உங்கள் உரையை இங்கே உள்ளிடவும்...",
+            key="tamil_input"
+        )
         
-        with st.expander("🔤 Translation", expanded=True):
-            st.success(eng)
+        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("உரை உருவாக்குதல்..."):
-            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=35,
-                guidance_scale=8.0
-            ).images[0]
-        
-        st.image(img, caption="🎨 Generated Image")
+                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()