Update app.py

app.py

@@ -1,60 +1,52 @@
-# app.py
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer, pipeline
 from diffusers import StableDiffusionPipeline
 import torch
 
-@st.cache_resource
-def load_models():
-    # Load translation model
-    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 text generation model
-    textgen = pipeline("text-generation", model="gpt2")
-
-    # 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
+@st.cache_resource(show_spinner=False)
+def load_all_models():
+    tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")
+    model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
+    text_gen = pipeline("text-generation", model="gpt2", device=-1)
+    img_pipe = StableDiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-2-base", torch_dtype=torch.float32
+    ).to("cpu")
+    return tokenizer, 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)
+    generated = model.generate(
+        **inputs,
+        forced_bos_token_id=tokenizer.get_lang_id("en"),
+        max_length=128
     )
-    if torch.cuda.is_available():
-        image_pipe = image_pipe.to("cuda")
-
-    return translator_tokenizer, translator_model, textgen, image_pipe
-
-def translate(text, tokenizer, model):
-    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)
+    return tokenizer.decode(generated[0], skip_special_tokens=True)
 
 def main():
-    st.set_page_config(page_title="Tamil Text to English + Image + Text", layout="centered")
-    st.title("🌐 தமிழ் உரையை ஆங்கிலம் + படம் + உரை ஆக மாற்றவும்")
+    st.set_page_config(page_title="Tamil→EN→Image/Text", layout="centered")
+    st.title("🌐 தமிழ் → English → Creative Text & Image")
 
-    tamil_input = st.text_area("தமிழ் உரையை உள்ளிடவும்", height=150)
+    tokenizer, model, text_gen, img_pipe = load_all_models()
+    tamil_text = st.text_area("தமிழ் உரை:", height=150)
 
-    if st.button("உரையை மொழிபெயர்த்து உருவாக்கவும்"):
-        if not tamil_input.strip():
-            st.warning("உரையை உள்ளிடவும்.")
+    if st.button("உருவாக்கு"):
+        if not tamil_text.strip():
+            st.warning("தயவு செய்து உரையை உள்ளிடவும்.")
             return
 
-        with st.spinner("மாடல்கள் ஏற்றப்படுகிறது..."):
-            tokenizer, model, textgen, image_pipe = load_models()
-
-        with st.spinner("மொழிபெயர்ப்பு நடைபெறுகிறது..."):
-            english_text = translate(tamil_input, tokenizer, model)
-            st.success(f"🔤 மொழிபெயர்ப்பு: {english_text}")
+        with st.spinner("Translating…"):
+            eng = translate_tamil(tamil_text, tokenizer, model)
+        st.success(f"🔤 Translation: {eng}")
 
-        with st.spinner("படம் உருவாக்கப்படுகிறது..."):
-            image = image_pipe(english_text).images[0]
-            st.image(image, caption="உருவாக்கப்பட்ட படம்", use_column_width=True)
+        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"]
+        st.info("📝 Creative Text:")
+        st.write(creative)
 
-        with st.spinner("உரையை உருவாக்கப்படுகிறது..."):
-            story = textgen(english_text, max_length=100, num_return_sequences=1)[0]["generated_text"]
-            st.text_area("✍️ உருவாக்கப்பட்ட உரை", value=story, height=200)
+        with st.spinner("Generating image…"):
+            img = img_pipe(eng, num_inference_steps=25, guidance_scale=7.5).images[0]
+        st.image(img, caption="🎨 Generated Image")
 
 if __name__ == "__main__":
     main()