Update app.py

app.py

@@ -1,18 +1,16 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from diffusers import StableDiffusionPipeline
 import torch
 
-# Load models only once
 @st.cache_resource
 def load_all_models():
     # Load translation model
-    trans_model_id = "ai4bharat/indictrans2-indic-en-dist-200M"
-    tokenizer = AutoTokenizer.from_pretrained(trans_model_id, trust_remote_code=True)
-    model = AutoModelForSeq2SeqLM.from_pretrained(trans_model_id, trust_remote_code=True)
-    translation_pipeline = pipeline("translation", model=model, tokenizer=tokenizer)
+    model_id = "ai4bharat/indictrans2-indic-en-dist-200M"
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_id, trust_remote_code=True)
 
-    # Load image generation model (Stable Diffusion 2.1)
+    # Load Stable Diffusion image generator
     img_pipe = StableDiffusionPipeline.from_pretrained(
         "stabilityai/stable-diffusion-2-1",
         torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
@@ -20,35 +18,33 @@ def load_all_models():
     )
     img_pipe = img_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-    return tokenizer, model, translation_pipeline, img_pipe
+    return tokenizer, model, img_pipe
 
-# Streamlit UI
 def main():
-    st.set_page_config(page_title="Tamil to English to Image Generator", layout="centered")
-    st.title("📸 Tamil → English → AI Image Generator")
-    st.markdown("Translate Tamil text to English and generate an image from it!")
+    st.set_page_config(page_title="Tamil to English to Image", layout="centered")
+    st.title("📸 Tamil ➝ English ➝ AI Image Generator")
 
-    # Load models
-    with st.spinner("Loading models..."):
-        tokenizer, model, translation_pipeline, img_pipe = load_all_models()
+    tamil_text = st.text_area("Enter Tamil text:", height=150)
 
-    # Input
-    tamil_text = st.text_area("Enter Tamil text here:", height=150)
-    
     if st.button("Generate Image"):
-        if tamil_text.strip() == "":
+        if not tamil_text.strip():
             st.warning("Please enter some Tamil text.")
             return
 
-        # Step 1: Translate Tamil to English
-        with st.spinner("Translating to English..."):
-            translated = translation_pipeline(tamil_text, src_lang="ta", tgt_lang="en")[0]["translation_text"]
+        with st.spinner("Loading models..."):
+            tokenizer, model, img_pipe = load_all_models()
+
+        with st.spinner("Translating Tamil to English..."):
+            # Prepare special format: "<2en> <tamil sentence>"
+            formatted_input = f"<2en> {tamil_text.strip()}"
+            inputs = tokenizer(formatted_input, return_tensors="pt")
+            output_ids = model.generate(**inputs)
+            translated = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
             st.success(f"🔤 English Translation: `{translated}`")
 
-        # Step 2: Generate image
         with st.spinner("Generating image..."):
             image = img_pipe(prompt=translated).images[0]
-            st.image(image, caption="Generated Image", use_column_width=True)
+            st.image(image, caption="🖼️ AI-generated Image", use_column_width=True)
 
 if __name__ == "__main__":
     main()