Update app.py

app.py

@@ -1,52 +1,91 @@
 import streamlit as st
-from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from diffusers import StableDiffusionPipeline
 import torch
 
 @st.cache_resource(show_spinner=False)
 def load_all_models():
-    tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
-    model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
+    # 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 other models
     text_gen = pipeline("text-generation", model="gpt2", device=-1)
     img_pipe = StableDiffusionPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-2-base", torch_dtype=torch.float32
+        "stabilityai/stable-diffusion-2-base", 
+        torch_dtype=torch.float32
     ).to("cpu")
-    return tokenizer, model, text_gen, img_pipe
+    
+    return indic_tokenizer, indic_model, text_gen, img_pipe
 
 def translate_tamil(text, tokenizer, model):
-    tokenizer.src_lang = "ta"
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    # Tokenize with batch processing
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=128
+    )
+    
+    # Generate translation
     generated = model.generate(
         **inputs,
-        forced_bos_token_id=tokenizer.get_lang_id("en"),
-        max_length=128
+        max_length=150,
+        num_beams=5,
+        early_stopping=True
     )
-    return tokenizer.decode(generated[0], skip_special_tokens=True)
+    
+    # Decode with normalization
+    return tokenizer.batch_decode(
+        generated, 
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=True
+    )[0]
 
 def main():
     st.set_page_config(page_title="Tamil→EN→Image/Text", layout="centered")
     st.title("🌐 தமிழ் → English → Creative Text & Image")
-
+    
+    # Load models once
     tokenizer, model, text_gen, img_pipe = load_all_models()
-    tamil_text = st.text_area("தமிழ் உரை:", height=150)
+    
+    # Input with Tamil placeholder
+    tamil_text = st.text_area(
+        "தமிழ் உரை:", 
+        height=150,
+        placeholder="உங்கள் உரையை இங்கே உள்ளிடவும்..."
+    )
 
-    if st.button("உருவாக்கு"):
+    if st.button("உருவாக்கு", type="primary"):
         if not tamil_text.strip():
             st.warning("தயவு செய்து உரையை உள்ளிடவும்.")
             return
 
-        with st.spinner("Translating…"):
+        with st.spinner("மொழிபெயர்க்கிறது..."):
             eng = translate_tamil(tamil_text, tokenizer, model)
-        st.success(f"🔤 Translation: {eng}")
+        
+        with st.expander("🔤 Translation", expanded=True):
+            st.success(eng)
 
-        with st.spinner("Creating creative text…"):
-            creative = text_gen(f"Write a creative description of: {eng}", max_length=60, num_return_sequences=1)[0]["generated_text"]
+        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("Generating image…"):
-            img = img_pipe(eng, num_inference_steps=25, guidance_scale=7.5).images[0]
+        with st.spinner("படத்தை உருவாக்குதல்..."):
+            img = img_pipe(
+                eng,
+                num_inference_steps=35,
+                guidance_scale=8.0
+            ).images[0]
+        
         st.image(img, caption="🎨 Generated Image")
 
 if __name__ == "__main__":
-    main()
+    main()