Update app.py

app.py

@@ -1,74 +1,44 @@
-# app.py
-import app_ui
-import os
-import time
-import tempfile
-import torch
-from transformers import MBartForConditionalGeneration, MBart50TokenizerFast, AutoTokenizer, AutoModelForCausalLM
-from diffusers import DiffusionPipeline
-from PIL import Image
-from rouge_score import rouge_scorer
-
-# Device setup
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Hugging Face Token (required for image pipeline)
-hf_token = os.getenv("HF_TOKEN", "your_token_here")  # Replace with your token or set as environment variable
-
-# Initialize translator (Tamil to English)
-translator_tokenizer = MBart50TokenizerFast.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"
-
-# Initialize text generator
-gen_tokenizer = AutoTokenizer.from_pretrained("gpt2")
-gen_model = AutoModelForCausalLM.from_pretrained("gpt2").to(device)
-
-# Initialize Stable Diffusion image pipeline
-pipe = DiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-2-1",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    use_auth_token=hf_token
-).to(device)
-pipe.safety_checker = None  # Optional: disable safety checks
-
-def translate_tamil_to_english(text, reference=None):
-    start = time.time()
-    inputs = translator_tokenizer(text, return_tensors="pt").to(device)
-    outputs = translator_model.generate(
-        **inputs,
-        forced_bos_token_id=translator_tokenizer.lang_code_to_id["en_XX"]
-    )
-    translation = translator_tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
-    duration = round(time.time() - start, 2)
-
-    rouge_l = None
-    if reference:
-        scorer = rouge_scorer.RougeScorer(['rougeL'], use_stemmer=True)
-        scores = scorer.score(reference.lower(), translation.lower())
-        rouge_l = round(scores['rougeL'].fmeasure, 4)
-
-    return translation, duration, rouge_l
-
-def generate_image(prompt):
-    try:
-        start = time.time()
-        out = pipe(prompt)
-        img = out.images[0].resize((256, 256))
-        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
-        img.save(tmp.name)
-        return tmp.name, round(time.time() - start, 2)
-    except Exception as e:
-        return None, f"Image generation failed: {e}"
-
-def generate_creative_text(prompt, max_length=100):
-    start = time.time()
-    input_ids = gen_tokenizer.encode(prompt, return_tensors="pt").to(device)
-    out = gen_model.generate(
-        input_ids, max_length=max_length, do_sample=True, top_k=50, temperature=0.9
-    )
-    text = gen_tokenizer.decode(out[0], skip_special_tokens=True)
-    duration = round(time.time() - start, 2)
-    tokens = text.split()
-    repetition = sum(t1 == t2 for t1, t2 in zip(tokens, tokens[1:])) / len(tokens)
-    return text, duration, len(tokens), round(repetition, 4)
+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()