Update app.py

app.py

@@ -1,53 +1,58 @@
-# app.py
 import streamlit as st
-from transformers import T5Tokenizer, T5ForConditionalGeneration
+from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer, pipeline
 from diffusers import StableDiffusionPipeline
 import torch
 
-# Load translation model
 @st.cache_resource
-def load_translation_model():
-    model_name = "nithalis/tamil_translation_t5"
-    tokenizer = T5Tokenizer.from_pretrained(model_name)
-    model = T5ForConditionalGeneration.from_pretrained(model_name)
-    return tokenizer, model
-
-# Load faster image model
-@st.cache_resource
-def load_diffusion_model():
-    pipe = StableDiffusionPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-2-base",
-        torch_dtype=torch.float32
+def load_models():
+    # Translation model
+    model_name = "alirezamsh/small100"
+    tokenizer = M2M100Tokenizer.from_pretrained(model_name)
+    model = M2M100ForConditionalGeneration.from_pretrained(model_name)
+
+    # Creative text generation (GPT-2)
+    text_generator = pipeline("text-generation", model="gpt2-medium")
+
+    # Stable diffusion image generation (CPU-friendly)
+    image_pipe = StableDiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-2-base", torch_dtype=torch.float32
     ).to("cpu")
-    return pipe
 
-# 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():
-        output = model.generate(**inputs, max_length=128)
-    return tokenizer.decode(output[0], skip_special_tokens=True)
+    return tokenizer, model, text_generator, image_pipe
 
-# Main Streamlit app
-def main():
-    st.set_page_config(page_title="தமிழ்→Image", layout="centered")
-    st.title("🌐 தமிழ் → ஆங்கிலம் → படம்")
+def translate_tamil_to_english(text, tokenizer, model):
+    tokenizer.src_lang = "ta"
+    inputs = tokenizer(text, return_tensors="pt", truncation=True)
+    generated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.get_lang_id("en"))
+    return tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
 
-    user_input = st.text_area("உங்கள் உரையை தமிழ் எழுத்துப்படியாக உள்ளிடவும்:", height=150)
+def generate_creative_text(prompt, generator):
+    response = generator(f"Describe creatively: {prompt}", max_length=60, num_return_sequences=1)
+    return response[0]["generated_text"]
 
-    if st.button("உருவாக்கு"):
+def main():
+    st.set_page_config(page_title="Tamil to English → Image & Text", layout="centered")
+    st.title("🪔 தமிழ் உரையை படமாக மாற்றவும்")
+    
+    user_input = st.text_area("தமிழ் உரை உள்ளிடவும்:", height=150)
+    if st.button("உருவாக்கவும்"):
         if not user_input.strip():
-            st.warning("உரை பதிவுசெய்யவும்")
+            st.warning("தயவுசெய்து ஒரு உரை உள்ளிடவும்.")
             return
-        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()
-            img = pipe(translation, num_inference_steps=25, guidance_scale=7.5).images[0]
-            st.image(img, caption="உருவாக்கப்பட்ட படம்", use_column_width=True)
+
+        with st.spinner("மொழிபெயர்ப்பு நடக்கிறது..."):
+            tokenizer, model, textgen, pipe = load_models()
+            english_text = translate_tamil_to_english(user_input, tokenizer, model)
+            st.success(f"Translated to English: {english_text}")
+
+        with st.spinner("சிறந்த உரையை உருவாக்குகிறது..."):
+            creative = generate_creative_text(english_text, textgen)
+            st.info("✍️ Creative Text:")
+            st.write(creative)
+
+        with st.spinner("படம் உருவாக்கப்படுகிறது..."):
+            image = pipe(english_text, num_inference_steps=25).images[0]
+            st.image(image, caption="🎨 Generated Image", use_column_width=True)
 
 if __name__ == "__main__":
     main()