Update app.py

app.py

@@ -1,38 +1,50 @@
+# app.py
 import streamlit as st
 import torch
 import openai
 import os
 import time
+import requests
 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 transformers import (
+    MBartForConditionalGeneration,
+    MBart50TokenizerFast,
+    AutoTokenizer,
+    AutoModelForCausalLM,
+)
 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
+openai.api_key = os.getenv("OPENAI_API_KEY")  # Make sure this is set in your environment
 
-# Load MBart
+# Load MBart model
 translator_model = MBartForConditionalGeneration.from_pretrained(
-    "facebook/mbart-large-50-many-to-many-mmt"
-).to(device)
+    "facebook/mbart-large-50-many-to-many-mmt",
+    device_map="auto",
+    low_cpu_mem_usage=True
+)
 translator_tokenizer = MBart50TokenizerFast.from_pretrained(
     "facebook/mbart-large-50-many-to-many-mmt"
 )
 translator_tokenizer.src_lang = "ta_IN"
 
-# GPT-2
-gen_model = GPT2LMHeadModel.from_pretrained("gpt2").to(device)
+# Load GPT-2 model
+gen_model = AutoModelForCausalLM.from_pretrained(
+    "gpt2",
+    device_map="auto",
+    low_cpu_mem_usage=True
+)
 gen_model.eval()
 gen_tokenizer = AutoTokenizer.from_pretrained("gpt2")
 
-# CLIP
+# Load CLIP model
 clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)
 
-# ---- Translation ----
+# ---- Translation Function ----
 def translate_tamil_to_english(text, reference=None):
     start = time.time()
     inputs = translator_tokenizer(text, return_tensors="pt").to(device)
@@ -51,7 +63,7 @@ def translate_tamil_to_english(text, reference=None):
 
     return translated, duration, rouge_l
 
-# ---- Creative Text ----
+# ---- Creative Text Generation ----
 def generate_creative_text(prompt, max_length=100):
     start = time.time()
     input_ids = gen_tokenizer.encode(prompt, return_tensors="pt").to(device)
@@ -88,10 +100,9 @@ def generate_image(prompt):
             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
+        # Save to temporary file
         tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
         image_data.save(tmp_file.name)
         duration = round(time.time() - start, 2)
@@ -108,7 +119,7 @@ def generate_image(prompt):
     except Exception as e:
         return None, None, f"Image generation failed: {str(e)}"
 
-# ---- UI ----
+# ---- Streamlit UI ----
 st.set_page_config(page_title="Tamil → English + AI Art", layout="centered")
 st.title("🧠 Tamil → English + 🎨 Creative Text + 🖼️ AI Image")