Update app.py

app.py

@@ -1,58 +1,53 @@
 # app.py
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 from diffusers import StableDiffusionPipeline
 import torch
 
-# Load NLLB-200 translation model
+# Load translation model
 @st.cache_resource
 def load_translation_model():
-    model_name = "facebook/nllb-200-distilled-600M"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    model_name = "nithalis/tamil_translation_t5"
+    tokenizer = T5Tokenizer.from_pretrained(model_name)
+    model = T5ForConditionalGeneration.from_pretrained(model_name)
     return tokenizer, model
 
-# Load Stable Diffusion model
+# Load faster image model
 @st.cache_resource
 def load_diffusion_model():
     pipe = StableDiffusionPipeline.from_pretrained(
-        "CompVis/stable-diffusion-v1-4",
-        torch_dtype=torch.float32  # CPU only
-    )
-    pipe = pipe.to("cpu")
+        "stabilityai/stable-diffusion-2-base",
+        torch_dtype=torch.float32
+    ).to("cpu")
     return pipe
 
-# Translate Tamil to English using prompt-style method
-def translate_text(text, tokenizer, model):
-    prompt = f"Translate this from Tamil to English: {text}"
+# Translate Tamil to English
+def translate_tamil_to_english(text, tokenizer, model):
+    prompt = f"translate Tamil to English: {text}"
     inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
     with torch.no_grad():
-        translated_tokens = model.generate(**inputs, max_length=512)
-    return tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+        output = model.generate(**inputs, max_length=128)
+    return tokenizer.decode(output[0], skip_special_tokens=True)
 
-# Streamlit App
+# Main Streamlit app
 def main():
-    st.set_page_config(page_title="Tamil to Image Generator", layout="centered")
-    with open("style.css") as f:
-        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
+    st.set_page_config(page_title="தமிழ்→Image", layout="centered")
+    st.title("🌐 தமிழ் → ஆங்கிலம் → படம்")
 
-    st.title("🪔 தமிழ் உரை -> பட உருவாக்கம்")
+    user_input = st.text_area("உங்கள் உரையை தமிழ் எழுத்துப்படியாக உள்ளிடவும்:", height=150)
 
-    user_input = st.text_area("உங்கள் உரையை உள்ளிடவும் (தமிழில்)", height=150)
-
-    if st.button("படம் உருவாக்கவும்"):
+    if st.button("உருவாக்கு"):
         if not user_input.strip():
-            st.warning("தயவு செய்து தமிழ் உரையை உள்ளிடவும்.")
+            st.warning("உரை பதிவுசெய்யவும்")
             return
-
-        with st.spinner("மொழிபெயர்ப்பு மற்றும் பட உருவாக்கம் நடைபெறுகிறது..."):
-            tokenizer, model = load_translation_model()
-            translated_text = translate_text(user_input, tokenizer, model)
-            st.success(f"மொழிபெயர்ப்பு (ஆங்கிலம்): {translated_text}")
+        with st.spinner("மொழிபெயர்ப்பு மற்றும் பட உருவாக்கம்..."):
+            tok, mdl = load_translation_model()
+            translation = translate_tamil_to_english(user_input, tok, mdl)
+            st.success(f"ஆங்கில மொழிபெயர்ப்பு: {translation}")
 
             pipe = load_diffusion_model()
-            image = pipe(translated_text).images[0]
-            st.image(image, caption="உங்கள் உருவாக்கப்பட்ட படம்", use_column_width=True)
+            img = pipe(translation, num_inference_steps=25, guidance_scale=7.5).images[0]
+            st.image(img, caption="உருவாக்கப்பட்ட படம்", use_column_width=True)
 
 if __name__ == "__main__":
     main()