Update app.py

app.py

@@ -1,35 +1,50 @@
 import streamlit as st
 import torch
 from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, GPT2LMHeadModel
 from diffusers import StableDiffusionPipeline
 from rouge_score import rouge_scorer
 from PIL import Image
-import clip
 import tempfile
 import os
-import math
 import time
+import torch.nn.functional as F
+import clip  # from OpenAI CLIP repo
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
 
-# Device setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# 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 MBart 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"
+)
 translator_tokenizer.src_lang = "ta_IN"
 
-# GPT-2 for creative text
-gen_model = AutoModelForCausalLM.from_pretrained("gpt2").to(device)
+# Load GPT-2
+gen_model = GPT2LMHeadModel.from_pretrained("gpt2").to(device)
+gen_model.eval()
 gen_tokenizer = AutoTokenizer.from_pretrained("gpt2")
 
-# Stable Diffusion v1.4
-pipe = StableDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-1-4",
-    torch_dtype=torch.float32,
-    use_auth_token=os.getenv("HF_TOKEN")  # set this on Hugging Face Spaces
-).to(device)
-pipe.safety_checker = None  # Optional
+# Try loading SD-2.1, fallback to lightweight
+try:
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-2-1",
+        torch_dtype=torch.float32,
+        use_auth_token=os.getenv("HF_TOKEN")
+    ).to(device)
+    pipe.safety_checker = None
+    model_loaded = "stabilityai/stable-diffusion-2-1"
+except Exception as e:
+    st.warning("⚠️ SD-2.1 failed. Using lightweight fallback model.")
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "OFA-Sys/small-stable-diffusion-v0",
+        torch_dtype=torch.float32
+    ).to(device)
+    pipe.safety_checker = None
+    model_loaded = "OFA-Sys/small-stable-diffusion-v0"
 
 # Load CLIP for image-text similarity
 clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)
@@ -53,88 +68,89 @@ def translate_tamil_to_english(text, reference=None):
 
     return translated, duration, rouge_l
 
-# Text generation with repetition & perplexity
+# Creative text generator with evaluation
 def generate_creative_text(prompt, max_length=100):
     start = time.time()
     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
+        input_ids,
+        max_length=max_length,
+        do_sample=True,
+        top_k=50,
+        temperature=0.9
     )
     text = gen_tokenizer.decode(output[0], skip_special_tokens=True)
     duration = round(time.time() - start, 2)
 
     tokens = text.split()
-    repetition_rate = sum(t1 == t2 for t1, t2 in zip(tokens, tokens[1:])) / len(tokens)
+    rep_rate = sum(t1 == t2 for t1, t2 in zip(tokens, tokens[1:])) / len(tokens) if len(tokens) > 1 else 0
 
-    # Perplexity
+    # Calculate perplexity
     with torch.no_grad():
+        input_ids = gen_tokenizer.encode(text, return_tensors="pt").to(device)
         outputs = gen_model(input_ids, labels=input_ids)
         loss = outputs.loss
-        perplexity = math.exp(loss.item())
+        perplexity = torch.exp(loss).item()
 
-    return text, duration, len(tokens), round(repetition_rate, 4), round(perplexity, 3)
+    return text, duration, len(tokens), round(rep_rate, 4), round(perplexity, 4)
 
-# Image generation + CLIP similarity
+# Generate image and CLIP similarity
 def generate_image(prompt):
     try:
         start = time.time()
         result = pipe(prompt)
         image = result.images[0].resize((256, 256))
-        duration = round(time.time() - start, 2)
-
-        # Save image
         tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
         image.save(tmp_file.name)
+        duration = round(time.time() - start, 2)
 
-        # CLIP similarity
-        image_input = clip_preprocess(image).unsqueeze(0).to(device)
-        text_input = clip.tokenize(prompt).to(device)
+        # Compute CLIP similarity
+        image_input = clip_preprocess(Image.open(tmp_file.name)).unsqueeze(0).to(device)
+        text_input = clip.tokenize([prompt]).to(device)
         with torch.no_grad():
             image_features = clip_model.encode_image(image_input)
             text_features = clip_model.encode_text(text_input)
-            similarity = torch.cosine_similarity(image_features, text_features).item()
+            similarity = F.cosine_similarity(image_features, text_features).item()
 
         return tmp_file.name, duration, round(similarity, 4)
     except Exception as e:
-        return None, 0, f"Image generation failed: {str(e)}"
+        return None, None, f"Image generation failed: {str(e)}"
 
 # Streamlit UI
 st.set_page_config(page_title="Tamil → English + AI Art", layout="centered")
-st.title("🧠 Tamil → English + 🎨 Creative Text + AI Image")
+st.title("🧠 Tamil → English + 🎨 Creative Text + 🖼️ AI Image")
 
-tamil_input = st.text_area("✍️ Enter Tamil text here", height=150)
+tamil_input = st.text_area("✍️ Enter Tamil text", height=150)
 reference_input = st.text_input("📘 Optional: Reference English translation for ROUGE")
 
 if st.button("🚀 Generate Output"):
     if not tamil_input.strip():
         st.warning("Please enter Tamil text.")
     else:
-        with st.spinner("🔄 Translating Tamil to English..."):
+        with st.spinner("🔄 Translating..."):
             english_text, t_time, rouge_l = translate_tamil_to_english(tamil_input, reference_input)
 
-        st.success(f"✅ Translated in {t_time} seconds")
+        st.success(f"✅ Translated in {t_time}s")
         st.markdown(f"**📝 English Translation:** `{english_text}`")
         if rouge_l is not None:
-            st.markdown(f"📊 **ROUGE-L Score:** `{rouge_l}`")
+            st.markdown(f"📊 ROUGE-L Score: `{rouge_l}`")
 
         with st.spinner("🖼️ Generating image..."):
-            image_path, img_time, similarity = generate_image(english_text)
+            image_path, img_time, clip_score = generate_image(english_text)
 
-        if isinstance(similarity, float):
-            st.success(f"🖼️ Image generated in {img_time} seconds")
+        if image_path:
+            st.success(f"🖼️ Image generated in {img_time}s using `{model_loaded}`")
             st.image(Image.open(image_path), caption="AI-Generated Image", use_column_width=True)
-            st.markdown(f"🎯 **CLIP Text-Image Similarity:** `{similarity}`")
+            st.markdown(f"🔍 **CLIP Text-Image Similarity:** `{clip_score}`")
         else:
-            st.error(similarity)
+            st.error(clip_score)
 
         with st.spinner("💡 Generating creative text..."):
-            creative, c_time, tokens, rep_rate, perplexity = generate_creative_text(english_text)
+            creative, c_time, tokens, rep_rate, ppl = generate_creative_text(english_text)
 
-        st.success(f"✨ Creative text generated in {c_time} seconds")
+        st.success(f"✨ Creative text in {c_time}s")
         st.markdown(f"**🧠 Creative Output:** `{creative}`")
-        st.markdown(f"📌 Tokens: `{tokens}`, 🔁 Repetition Rate: `{rep_rate}`")
-        st.markdown(f"📉 Perplexity: `{perplexity}`")
+        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 on Hugging Face")
+st.caption("Built by Sureshkumar R | MBart + GPT-2 + Stable Diffusion + CLIP")