Update app.py

app.py

@@ -1,52 +1,38 @@
 import streamlit as st
 import torch
+import torch.nn.functional as F
+import os
+import time
+import tempfile
+from PIL import Image
 from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, CLIPProcessor, CLIPModel
 from diffusers import StableDiffusionPipeline
 from rouge_score import rouge_scorer
-from PIL import Image
-import tempfile
-import os
-import time
-from transformers import CLIPProcessor, CLIPModel
-import torch.nn.functional as F
 
-# Set device
+# --- Device Setup ---
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load translation model
-translator_model = MBartForConditionalGeneration.from_pretrained(
-    "facebook/mbart-large-50-many-to-many-mmt"
-).to(device)
-translator_tokenizer = MBart50TokenizerFast.from_pretrained(
-    "facebook/mbart-large-50-many-to-many-mmt"
-)
+# --- Load Models ---
+translator_model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt").to(device)
+translator_tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
 translator_tokenizer.src_lang = "ta_IN"
 
-# Load GPT-2 for creative text
 gen_tokenizer = AutoTokenizer.from_pretrained("gpt2")
 gen_model = AutoModelForCausalLM.from_pretrained("gpt2").to(device)
 gen_model.eval()
 
-# Load Stable Diffusion 1.5
-pipe = StableDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-1-5",
-    torch_dtype=torch.float32,
-).to(device)
-pipe.safety_checker = None  # Optional: disable safety filter
+pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-1-5").to(device)
+pipe.safety_checker = None
 
-# Load CLIP model
 clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
 clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
-# --- Translation ---
+# --- Functions ---
 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"]
-    )
+    outputs = translator_model.generate(**inputs, forced_bos_token_id=translator_tokenizer.lang_code_to_id["en_XX"])
     translated = translator_tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
     duration = round(time.time() - start, 2)
 
@@ -58,7 +44,6 @@ def translate_tamil_to_english(text, reference=None):
 
     return translated, duration, rouge_l
 
-# --- GPT-2 Creative Generation ---
 def generate_creative_text(prompt, max_length=100):
     start = time.time()
     input_ids = gen_tokenizer.encode(prompt, return_tensors="pt").to(device)
@@ -67,9 +52,8 @@ def generate_creative_text(prompt, max_length=100):
     duration = round(time.time() - start, 2)
 
     tokens = text.split()
-    repetition_rate = sum(t1 == t2 for t1, t2 in zip(tokens, tokens[1:])) / len(tokens)
+    repetition_rate = sum(t1 == t2 for t1, t2 in zip(tokens, tokens[1:])) / len(tokens) if len(tokens) > 1 else 0
 
-    # Perplexity
     with torch.no_grad():
         input_ids = gen_tokenizer.encode(text, return_tensors="pt").to(device)
         outputs = gen_model(input_ids, labels=input_ids)
@@ -78,7 +62,6 @@ def generate_creative_text(prompt, max_length=100):
 
     return text, duration, len(tokens), round(repetition_rate, 4), round(perplexity, 4)
 
-# --- Stable Diffusion Image Generation ---
 def generate_image(prompt):
     try:
         start = time.time()
@@ -91,7 +74,6 @@ def generate_image(prompt):
     except Exception as e:
         return None, 0, f"Image generation failed: {str(e)}"
 
-# --- CLIP Similarity ---
 def evaluate_clip_similarity(text, image):
     inputs = clip_processor(text=[text], images=image, return_tensors="pt", padding=True).to(device)
     with torch.no_grad():
@@ -112,7 +94,7 @@ if st.button("🚀 Generate Output"):
     if not tamil_input.strip():
         st.warning("Please enter Tamil text.")
     else:
-        with st.spinner("🔄 Translating..."):
+        with st.spinner("🔄 Translating Tamil to English..."):
             english_text, t_time, rouge_l = translate_tamil_to_english(tamil_input, reference_input)
 
         st.success(f"✅ Translated in {t_time}s")
@@ -120,7 +102,7 @@ if st.button("🚀 Generate Output"):
         if rouge_l is not None:
             st.markdown(f"📊 ROUGE-L Score: `{rouge_l}`")
 
-        with st.spinner("🎨 Generating image..."):
+        with st.spinner("🖼️ Generating image from text..."):
             image_path, img_time, image_obj = generate_image(english_text)
 
         if isinstance(image_obj, Image.Image):
@@ -129,16 +111,16 @@ if st.button("🚀 Generate Output"):
 
             with st.spinner("🔎 Evaluating CLIP similarity..."):
                 clip_score = evaluate_clip_similarity(english_text, image_obj)
-                st.markdown(f"�� CLIP Text-Image Similarity: `{clip_score}`")
+                st.markdown(f"🔍 **CLIP Text-Image Similarity:** `{clip_score}`")
         else:
             st.error(image_obj)
 
         with st.spinner("💡 Generating creative text..."):
             creative, c_time, tokens, rep_rate, ppl = generate_creative_text(english_text)
 
-        st.success(f"✨ Creative text generated in {c_time}s")
+        st.success(f"✨ Creative text in {c_time}s")
         st.markdown(f"**🧠 Creative Output:** `{creative}`")
         st.markdown(f"📌 Tokens: `{tokens}`, 🔁 Repetition Rate: `{rep_rate}`, 📉 Perplexity: `{ppl}`")
 
 st.markdown("---")
-st.caption("Built by Sureshkumar R using MBart, GPT-2, Stable Diffusion 1.5, and CLIP (Open Source)")
+st.caption("Built by Sureshkumar R | MBart + GPT-2 + Stable Diffusion + CLIP")