Update app.py

app.py

@@ -1,43 +1,45 @@
 import streamlit as st
 import torch
-import openai
-import os
-import time
-import requests
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from diffusers import StableDiffusionPipeline
+from rouge_score import rouge_scorer
 from PIL import Image
 import tempfile
-import clip
+import os
+import time
+from transformers import CLIPProcessor, CLIPModel
 import torch.nn.functional as F
-from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
-from transformers import AutoTokenizer, AutoModelForCausalLM, GPT2LMHeadModel
-from rouge_score import rouge_scorer
 
 # Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# OpenAI Key
-openai.api_key = os.getenv("OPENAI_API_KEY")
-
-# ---- Load MBart (Translation) ----
+# 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"
 )
-translator_model.to(device)
 translator_tokenizer.src_lang = "ta_IN"
 
-# ---- GPT-2 ----
-gen_model = GPT2LMHeadModel.from_pretrained("gpt2")
+# Load GPT-2 for creative text
 gen_tokenizer = AutoTokenizer.from_pretrained("gpt2")
-gen_model.to(device)
+gen_model = AutoModelForCausalLM.from_pretrained("gpt2").to(device)
 gen_model.eval()
 
-# ---- CLIP ----
-clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)
+# 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
+
+# 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 ----
+# --- Translation ---
 def translate_tamil_to_english(text, reference=None):
     start = time.time()
     inputs = translator_tokenizer(text, return_tensors="pt").to(device)
@@ -56,63 +58,52 @@ def translate_tamil_to_english(text, reference=None):
 
     return translated, duration, rouge_l
 
-# ---- Creative Text ----
+# --- 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)
-    output = gen_model.generate(
-        input_ids,
-        max_length=max_length,
-        do_sample=True,
-        top_k=50,
-        temperature=0.9
-    )
+    output = gen_model.generate(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()
-    rep_rate = sum(t1 == t2 for t1, t2 in zip(tokens, tokens[1:])) / len(tokens) if len(tokens) > 1 else 0
+    repetition_rate = sum(t1 == t2 for t1, t2 in zip(tokens, tokens[1:])) / len(tokens)
 
+    # 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 = torch.exp(loss).item()
 
-    return text, duration, len(tokens), round(rep_rate, 4), round(perplexity, 4)
+    return text, duration, len(tokens), round(repetition_rate, 4), round(perplexity, 4)
 
-# ---- Image Generation ----
+# --- Stable Diffusion Image Generation ---
 def generate_image(prompt):
     try:
         start = time.time()
-        response = openai.images.generate(
-            model="dall-e-3",
-            prompt=prompt,
-            size="512x512",
-            quality="standard",
-            n=1
-        )
-        image_url = response.data[0].url
-        image_data = Image.open(requests.get(image_url, stream=True).raw).resize((256, 256))
-
+        result = pipe(prompt)
+        image = result.images[0].resize((256, 256))
         tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
-        image_data.save(tmp_file.name)
+        image.save(tmp_file.name)
         duration = round(time.time() - start, 2)
-
-        image_input = clip_preprocess(image_data).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 = F.cosine_similarity(image_features, text_features).item()
-
-        return tmp_file.name, duration, round(similarity, 4)
+        return tmp_file.name, duration, image
     except Exception as e:
-        return None, None, f"Image generation failed: {str(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():
+        outputs = clip_model(**inputs)
+        logits_per_image = outputs.logits_per_image
+        probs = F.softmax(logits_per_image, dim=1)
+        similarity_score = logits_per_image[0][0].item()
+    return round(similarity_score, 4)
 
-# ---- UI ----
+# --- 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", height=150)
 reference_input = st.text_input("📘 Optional: Reference English translation for ROUGE")
@@ -129,22 +120,25 @@ if st.button("🚀 Generate Output"):
         if rouge_l is not None:
             st.markdown(f"📊 ROUGE-L Score: `{rouge_l}`")
 
-        with st.spinner("🖼️ Generating image..."):
-            image_path, img_time, clip_score = generate_image(english_text)
+        with st.spinner("🎨 Generating image..."):
+            image_path, img_time, image_obj = generate_image(english_text)
 
-        if image_path:
-            st.success(f"🖼️ Image generated in {img_time}s using OpenAI DALL·E 3")
+        if isinstance(image_obj, Image.Image):
+            st.success(f"🖼️ Image generated in {img_time}s")
             st.image(Image.open(image_path), caption="AI-Generated Image", use_column_width=True)
-            st.markdown(f"🔍 **CLIP Text-Image Similarity:** `{clip_score}`")
+
+            with st.spinner("🔎 Evaluating CLIP similarity..."):
+                clip_score = evaluate_clip_similarity(english_text, image_obj)
+                st.markdown(f"🔍 CLIP Text-Image Similarity: `{clip_score}`")
         else:
-            st.error(clip_score)
+            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 in {c_time}s")
+        st.success(f"✨ Creative text generated 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 | MBart + GPT-2 + OpenAI DALL·E 3")
+st.caption("Built by Sureshkumar R using MBart, GPT-2, Stable Diffusion 1.5, and CLIP (Open Source)")