Update app.py

app.py

@@ -1,61 +1,60 @@
 # app.py
 import streamlit as st
-from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer, pipeline
-from diffusers import DiffusionPipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from diffusers import StableDiffusionPipeline
 import torch
 
-@st.cache_resource(show_spinner=False)
-def load_all_models():
+@st.cache_resource
+def load_models():
     # Load translation model
-    model_name = "facebook/m2m100_418M"
-    tokenizer = M2M100Tokenizer.from_pretrained(model_name)
-    model = M2M100ForConditionalGeneration.from_pretrained(model_name)
+    translation_model_name = "Helsinki-NLP/opus-mt-ta-en"
+    translator_tokenizer = AutoTokenizer.from_pretrained(translation_model_name)
+    translator_model = AutoModelForSeq2SeqLM.from_pretrained(translation_model_name)
 
-    # Load creative text model (smaller GPT-2)
-    textgen = pipeline("text-generation", model="gpt2", device=-1)
+    # Load text generation model
+    textgen = pipeline("text-generation", model="gpt2")
 
-    # Load lightweight image generation pipeline
-    img_pipe = DiffusionPipeline.from_pretrained(
-        "stabilityai/sdxl-lite", torch_dtype=torch.float32
-    ).to("cpu")
+    # Load stable diffusion (light version)
+    image_model_name = "CompVis/stable-diffusion-v1-4"
+    image_pipe = StableDiffusionPipeline.from_pretrained(
+        image_model_name,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+    )
+    if torch.cuda.is_available():
+        image_pipe = image_pipe.to("cuda")
 
-    return tokenizer, model, textgen, img_pipe
+    return translator_tokenizer, translator_model, textgen, image_pipe
 
 def translate(text, tokenizer, model):
-    tokenizer.src_lang = "ta"
-    inputs = tokenizer(text, return_tensors="pt")
-    output = model.generate(inputs["input_ids"], forced_bos_token_id=tokenizer.get_lang_id("en"), max_length=100)
-    return tokenizer.decode(output[0], skip_special_tokens=True)
-
-def generate_text(prompt, pipe):
-    output = pipe(prompt, max_length=60, do_sample=True)[0]
-    return output["generated_text"]
+    inputs = tokenizer.encode(text, return_tensors="pt", truncation=True, max_length=512)
+    outputs = model.generate(inputs, max_length=512, num_beams=4)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 def main():
-    st.set_page_config(page_title="Tamil to English → Creative → Image", layout="centered")
-    st.title("🌐 தமிழ் ➝ English ➝ Creative Text + Image")
-
-    tokenizer, model, textgen, img_pipe = load_all_models()
+    st.set_page_config(page_title="Tamil Text to English + Image + Text", layout="centered")
+    st.title("🌐 தமிழ் உரையை ஆங்கிலம் + படம் + உரை ஆக மாற்றவும்")
 
-    tamil_text = st.text_area("தமிழ் உரையை உள்ளிடவும்:", height=130)
+    tamil_input = st.text_area("தமிழ் உரையை உள்ளிடவும்", height=150)
 
-    if st.button("உருவாக்கு"):
-        if not tamil_text.strip():
-            st.warning("தயவுசெய்து உரையை உள்ளிடவும்.")
+    if st.button("உரையை மொழிபெயர்த்து உருவாக்கவும்"):
+        if not tamil_input.strip():
+            st.warning("உரையை உள்ளிடவும்.")
             return
 
-        with st.spinner("மொழிபெயர்ப்பு..."):
-            english_text = translate(tamil_text, tokenizer, model)
-        st.success(f"🔁 Translated: {english_text}")
+        with st.spinner("மாடல்கள் ஏற்றப்படுகிறது..."):
+            tokenizer, model, textgen, image_pipe = load_models()
 
-        with st.spinner("உரையாக்கம்..."):
-            creative_text = generate_text(english_text, textgen)
-        st.info("📝 Creative Output:")
-        st.write(creative_text)
+        with st.spinner("மொழிபெயர்ப்பு நடைபெறுகிறது..."):
+            english_text = translate(tamil_input, tokenizer, model)
+            st.success(f"🔤 மொழிபெயர்ப்பு: {english_text}")
 
         with st.spinner("படம் உருவாக்கப்படுகிறது..."):
-            image = img_pipe(english_text).images[0]
-        st.image(image, caption="🎨 Generated Image", use_column_width=True)
+            image = image_pipe(english_text).images[0]
+            st.image(image, caption="உருவாக்கப்பட்ட படம்", use_column_width=True)
+
+        with st.spinner("உரையை உருவாக்கப்படுகிறது..."):
+            story = textgen(english_text, max_length=100, num_return_sequences=1)[0]["generated_text"]
+            st.text_area("✍️ உருவாக்கப்பட்ட உரை", value=story, height=200)
 
 if __name__ == "__main__":
     main()