Update app.py

app.py

@@ -3,23 +3,20 @@ from transformers import pipeline
 from diffusers import StableDiffusionPipeline
 import torch
 
+# Cache and load models only once
 @st.cache_resource
 def load_all_models():
-    # Load Tamil ➝ English translation model (with remote code)
     translation_pipeline = pipeline(
-        "translation",
+        "text2text-generation",
         model="ai4bharat/indictrans2-indic-en-dist-200M",
         tokenizer="ai4bharat/indictrans2-indic-en-dist-200M",
-        src_lang="ta", tgt_lang="en",
-        trust_remote_code=True,  # 🔥 Must add this
         device=0 if torch.cuda.is_available() else -1,
+        trust_remote_code=True  # Required for custom tokenizer
     )
 
-    # Load image generation model
     img_pipe = StableDiffusionPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-2-1",
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        revision="fp16" if torch.cuda.is_available() else None,
+        "runwayml/stable-diffusion-v1-5",
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
     )
     img_pipe = img_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -27,26 +24,34 @@ def load_all_models():
 
 
 def main():
-    st.set_page_config(page_title="Tamil ➝ English ➝ Image", layout="centered")
-    st.title("🌐 Tamil to English to Image Generator")
+    st.title("🧠 Tamil to English Image Generator")
 
-    tamil_text = st.text_area("Enter Tamil text:", height=150)
+    tamil_text = st.text_area("✍️ Enter Tamil word or sentence:")
 
-    if st.button("Generate Image"):
-        if not tamil_text.strip():
-            st.warning("Please enter Tamil text.")
+    if st.button("Translate & Generate Image"):
+        if tamil_text.strip() == "":
+            st.warning("⚠️ Please enter some Tamil text.")
             return
 
-        with st.spinner("Loading models..."):
+        with st.spinner("⏳ Loading models..."):
             translation_pipeline, img_pipe = load_all_models()
 
-        with st.spinner("Translating Tamil to English..."):
-            translated = translation_pipeline(tamil_text)[0]["translation_text"]
-            st.success(f"🔤 Translated English: `{translated}`")
+        # Format input for IndicTrans2
+        formatted_input = f"<2ta> <2en> {tamil_text.strip()}"
+
+        try:
+            translated = translation_pipeline(formatted_input)[0]["translation_text"]
+        except Exception as e:
+            st.error(f"❌ Translation failed: {e}")
+            return
+
+        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(prompt=translated).images[0]
-            st.image(image, caption="🎨 AI-generated Image", use_column_width=True)
 
 if __name__ == "__main__":
     main()