Update app.py

app.py

@@ -1,20 +1,21 @@
 import streamlit as st
 import torch
-from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
-from transformers import AutoTokenizer, AutoModelForCausalLM, GPT2LMHeadModel
-from diffusers import StableDiffusionPipeline
-from rouge_score import rouge_scorer
-from PIL import Image
-import tempfile
+import openai
 import os
 import time
-import torch.nn.functional as F
+from PIL import Image
+import tempfile
 import clip  # from OpenAI CLIP repo
+import torch.nn.functional as F
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+from transformers import AutoTokenizer, AutoModelForCausalLM, GPT2LMHeadModel
+from rouge_score import rouge_scorer
 from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
+openai.api_key = os.getenv("OPENAI_API_KEY")  # Set this from env
 
-# Load MBart model
+# Load MBart
 translator_model = MBartForConditionalGeneration.from_pretrained(
     "facebook/mbart-large-50-many-to-many-mmt"
 ).to(device)
@@ -23,33 +24,15 @@ translator_tokenizer = MBart50TokenizerFast.from_pretrained(
 )
 translator_tokenizer.src_lang = "ta_IN"
 
-# Load GPT-2
+# GPT-2
 gen_model = GPT2LMHeadModel.from_pretrained("gpt2").to(device)
 gen_model.eval()
 gen_tokenizer = AutoTokenizer.from_pretrained("gpt2")
 
-# 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
 clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)
 
-# Translation function
+# ---- Translation ----
 def translate_tamil_to_english(text, reference=None):
     start = time.time()
     inputs = translator_tokenizer(text, return_tensors="pt").to(device)
@@ -68,7 +51,7 @@ def translate_tamil_to_english(text, reference=None):
 
     return translated, duration, rouge_l
 
-# Creative text generator with evaluation
+# ---- Creative Text ----
 def generate_creative_text(prompt, max_length=100):
     start = time.time()
     input_ids = gen_tokenizer.encode(prompt, return_tensors="pt").to(device)
@@ -85,7 +68,6 @@ def generate_creative_text(prompt, max_length=100):
     tokens = text.split()
     rep_rate = sum(t1 == t2 for t1, t2 in zip(tokens, tokens[1:])) / len(tokens) if len(tokens) > 1 else 0
 
-    # 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)
@@ -94,18 +76,28 @@ def generate_creative_text(prompt, max_length=100):
 
     return text, duration, len(tokens), round(rep_rate, 4), round(perplexity, 4)
 
-# Generate image and CLIP similarity
+# ---- Image Generation using DALL·E 3 ----
 def generate_image(prompt):
     try:
         start = time.time()
-        result = pipe(prompt)
-        image = result.images[0].resize((256, 256))
+        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(tempfile.NamedTemporaryFile(delete=False, suffix=".png"))
+        image_data = Image.open(requests.get(image_url, stream=True).raw).resize((256, 256))
+
+        # Save locally
         tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
-        image.save(tmp_file.name)
+        image_data.save(tmp_file.name)
         duration = round(time.time() - start, 2)
 
-        # Compute CLIP similarity
-        image_input = clip_preprocess(Image.open(tmp_file.name)).unsqueeze(0).to(device)
+        # CLIP similarity
+        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)
@@ -116,7 +108,7 @@ def generate_image(prompt):
     except Exception as e:
         return None, None, f"Image generation failed: {str(e)}"
 
-# Streamlit UI
+# ---- UI ----
 st.set_page_config(page_title="Tamil → English + AI Art", layout="centered")
 st.title("🧠 Tamil → English + 🎨 Creative Text + 🖼️ AI Image")
 
@@ -139,7 +131,7 @@ if st.button("🚀 Generate Output"):
             image_path, img_time, clip_score = generate_image(english_text)
 
         if image_path:
-            st.success(f"🖼️ Image generated in {img_time}s using `{model_loaded}`")
+            st.success(f"🖼️ Image generated in {img_time}s using OpenAI DALL·E 3")
             st.image(Image.open(image_path), caption="AI-Generated Image", use_column_width=True)
             st.markdown(f"🔍 **CLIP Text-Image Similarity:** `{clip_score}`")
         else:
@@ -153,4 +145,4 @@ if st.button("🚀 Generate Output"):
         st.markdown(f"📌 Tokens: `{tokens}`, 🔁 Repetition Rate: `{rep_rate}`, 📉 Perplexity: `{ppl}`")
 
 st.markdown("---")
-st.caption("Built by Sureshkumar R | MBart + GPT-2 + Stable Diffusion")
+st.caption("Built by Sureshkumar R | MBart + GPT-2 + OpenAI DALL·E 3")