Update app.py

app.py

@@ -1,57 +1,56 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 from diffusers import StableDiffusionPipeline
 import torch
 
-# Cache and load models only once
+# Cache models for faster loading
 @st.cache_resource
 def load_all_models():
+    # Translation model
+    translation_model = AutoModelForSeq2SeqLM.from_pretrained(
+        "ai4bharat/indictrans2-indic-en-dist-200M", trust_remote_code=True
+    )
+    translation_tokenizer = AutoTokenizer.from_pretrained(
+        "ai4bharat/indictrans2-indic-en-dist-200M", trust_remote_code=True
+    )
     translation_pipeline = pipeline(
-        "text2text-generation",
-        model="ai4bharat/indictrans2-indic-en-dist-200M",
-        tokenizer="ai4bharat/indictrans2-indic-en-dist-200M",
-        device=0 if torch.cuda.is_available() else -1,
-        trust_remote_code=True  # Required for custom tokenizer
+        "text2text-generation", model=translation_model, tokenizer=translation_tokenizer
     )
 
+    # Image generation model (Stable Diffusion)
     img_pipe = StableDiffusionPipeline.from_pretrained(
-        "runwayml/stable-diffusion-v1-5",
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+        "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
     )
     img_pipe = img_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
-
+    
     return translation_pipeline, img_pipe
 
-
 def main():
-    st.title("🧠 Tamil to English Image Generator")
+    st.title("📘 Tamil to English Translator & Image Generator")
 
-    tamil_text = st.text_area("✍️ Enter Tamil word or sentence:")
+    tamil_text = st.text_area("📝 Enter Tamil text (word or sentence)", height=100)
 
-    if st.button("Translate & Generate Image"):
-        if tamil_text.strip() == "":
-            st.warning("⚠️ Please enter some Tamil text.")
+    if st.button("🔄 Translate & Generate Image"):
+        if not tamil_text.strip():
+            st.warning("Please enter some Tamil text.")
             return
 
-        with st.spinner("⏳ Loading models..."):
+        try:
             translation_pipeline, img_pipe = load_all_models()
 
-        # Format input for IndicTrans2
-        formatted_input = f"<2ta> <2en> {tamil_text.strip()}"
+            # Prepare translation input
+            formatted_input = "<2en><|ta|>" + tamil_text.strip()
+            translated = translation_pipeline(formatted_input, max_length=256)[0]["generated_text"]
 
-        try:
-            translated = translation_pipeline(formatted_input)[0]["translation_text"]
-        except Exception as e:
-            st.error(f"❌ Translation failed: {e}")
-            return
+            st.success("✅ English Translation:")
+            st.write(translated)
 
-        st.success("✅ Translation Successful!")
-        st.markdown(f"**🗣️ English Translation:** {translated}")
-
-        with st.spinner("🎨 Generating image..."):
-            image = img_pipe(translated).images[0]
-            st.image(image, caption="🖼️ AI-Generated Image", use_column_width=True)
+            with st.spinner("🖼️ Generating image..."):
+                image = img_pipe(translated).images[0]
+                st.image(image, caption="🖼️ Generated from English text")
 
+        except Exception as e:
+            st.error(f"❌ Error: {str(e)}")
 
 if __name__ == "__main__":
     main()