Update app.py

app.py

@@ -1,44 +1,77 @@
-import gradio as gr
-from core import translate_tamil_to_english, generate_image, generate_creative_text
-
-def full_pipeline(tamil_text, reference_text=None):
-    if not tamil_text.strip():
-        return "Please enter Tamil text.", None, None, None, None, None, None, None, None
-
-    # Step 1: Translate Tamil to English
-    en_text, t_time, rouge_score = translate_tamil_to_english(tamil_text, reference_text)
-
-    # Step 2: Generate Image from English translation
-    image_path, img_time = generate_image(en_text)
-
-    # Step 3: Generate creative text from English translation
-    gen_text, gen_time, tokens, rep = generate_creative_text(en_text)
-
-    return en_text, t_time, rouge_score, image_path, img_time, gen_text, gen_time, tokens, rep
-
-# Gradio Interface
-demo = gr.Interface(
-    fn=full_pipeline,
-    inputs=[
-        gr.Textbox(label="✍️ Enter Tamil Text", lines=5, placeholder="எனது கனவு வீட்டை வர்ணிக்க..."),
-        gr.Textbox(label="📘 (Optional) Reference English Translation", lines=2)
-    ],
-    outputs=[
-        gr.Textbox(label="📝 English Translation"),
-        gr.Number(label="⏱️ Translation Time (s)"),
-        gr.Number(label="📊 ROUGE-L Score"),
-        gr.Image(label="🎨 Generated Image (256x256)"),
-        gr.Number(label="🖼️ Image Generation Time (s)"),
-        gr.Textbox(label="💡 Creative English Text"),
-        gr.Number(label="🕒 Text Generation Time (s)"),
-        gr.Number(label="🔢 Number of Tokens"),
-        gr.Number(label="♻️ Repetition Rate")
-    ],
-    title="🌐 Tamil to English Translator + Image & Text Generator",
-    description="Translate Tamil to English using MBart50 → Generate AI Image using StabilityAI → Generate Creative Text using GPT-2",
-    theme="soft",
-    allow_flagging="never"
-)
-
-if __name__ == "__main__":
-    demo.launch()
+import os
+import torch
+from transformers import MBartForConditionalGeneration, MBart50Tokenizer, AutoTokenizer, AutoModelForCausalLM
+from diffusers import StableDiffusionPipeline
+from PIL import Image
+import tempfile
+import time
+import streamlit as st
+
+# Use CPU (Hugging Face Spaces free tier)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load translation model
+translator_tokenizer = MBart50Tokenizer.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+translator_model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt").to(device)
+translator_tokenizer.src_lang = "ta_IN"
+
+# Load text generation model
+gen_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+gen_model = AutoModelForCausalLM.from_pretrained("gpt2").to(device)
+
+# Load image generation model
+pipe = StableDiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-1-base",
+    torch_dtype=torch.float32,
+    safety_checker=None
+).to(device)
+
+def translate_tamil_to_english(text):
+    inputs = translator_tokenizer(text, return_tensors="pt").to(device)
+    output = translator_model.generate(
+        **inputs,
+        forced_bos_token_id=translator_tokenizer.lang_code_to_id["en_XX"]
+    )
+    translated = translator_tokenizer.batch_decode(output, skip_special_tokens=True)[0]
+    return translated
+
+def generate_creative_text(prompt, max_length=100):
+    input_ids = gen_tokenizer.encode(prompt, return_tensors="pt").to(device)
+    output = gen_model.generate(
+        input_ids, max_length=max_length, do_sample=True, top_k=50, temperature=0.9
+    )
+    return gen_tokenizer.decode(output[0], skip_special_tokens=True)
+
+def generate_image(prompt):
+    image = pipe(prompt).images[0]
+    temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+    image.save(temp_file.name)
+    return temp_file.name
+
+# Streamlit UI
+st.set_page_config(page_title="Tamil → English + AI", layout="centered")
+st.title("🌐 Tamil to English + AI Image Generator")
+
+tamil_input = st.text_area("✍️ Enter Tamil Text", height=150)
+
+if st.button("🚀 Generate"):
+    if not tamil_input.strip():
+        st.warning("Please enter Tamil text.")
+    else:
+        with st.spinner("Translating..."):
+            translated = translate_tamil_to_english(tamil_input)
+            st.success("✅ Translated!")
+            st.markdown(f"**English:** `{translated}`")
+
+        with st.spinner("Generating creative text..."):
+            creative_text = generate_creative_text(translated)
+            st.success("✅ Creative text generated!")
+            st.markdown(f"**Creative Prompt:** `{creative_text}`")
+
+        with st.spinner("Generating image..."):
+            image_path = generate_image(translated)
+            st.success("✅ Image generated!")
+            st.image(Image.open(image_path), caption="🖼️ AI Generated Image", use_column_width=True)
+
+st.markdown("---")
+st.markdown("🔧 Powered by MBart, GPT2 & Stable Diffusion - Deployed on Hugging Face 🤗")