Update app.py

app.py

@@ -3,42 +3,44 @@ import torch
 import openai
 import os
 import time
+import requests
 from PIL import Image
 import tempfile
-import clip  # from OpenAI CLIP repo
+import clip
 import torch.nn.functional as F
-import requests
 from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, GPT2LMHeadModel
 from rouge_score import rouge_scorer
-from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
 
 # 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
+# ---- Load MBart (Translation) ----
 translator_model = MBartForConditionalGeneration.from_pretrained(
-    "facebook/mbart-large-50-many-to-many-mmt",
-    device_map="auto",
-    low_cpu_mem_usage=True
+    "facebook/mbart-large-50-many-to-many-mmt"
+)
+translator_tokenizer = MBart50TokenizerFast.from_pretrained(
+    "facebook/mbart-large-50-many-to-many-mmt"
 )
-translator_tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+translator_model.to(device)
 translator_tokenizer.src_lang = "ta_IN"
 
-# Load GPT-2
-gen_model = AutoModelForCausalLM.from_pretrained("gpt2", device_map="auto", low_cpu_mem_usage=True)
-gen_model.eval()
+# ---- GPT-2 ----
+gen_model = GPT2LMHeadModel.from_pretrained("gpt2")
 gen_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+gen_model.to(device)
+gen_model.eval()
 
-# Load CLIP
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# ---- CLIP ----
 clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)
 
 # ---- Translation ----
 def translate_tamil_to_english(text, reference=None):
     start = time.time()
-    inputs = translator_tokenizer(text, return_tensors="pt")
-    inputs = {k: v.to(translator_model.device) for k, v in inputs.items()}
+    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"]
@@ -57,9 +59,7 @@ def translate_tamil_to_english(text, reference=None):
 # ---- Creative Text ----
 def generate_creative_text(prompt, max_length=100):
     start = time.time()
-    input_ids = gen_tokenizer.encode(prompt, return_tensors="pt")
-    input_ids = input_ids.to(gen_model.device)
-
+    input_ids = gen_tokenizer.encode(prompt, return_tensors="pt").to(device)
     output = gen_model.generate(
         input_ids,
         max_length=max_length,
@@ -74,7 +74,7 @@ def generate_creative_text(prompt, max_length=100):
     rep_rate = sum(t1 == t2 for t1, t2 in zip(tokens, tokens[1:])) / len(tokens) if len(tokens) > 1 else 0
 
     with torch.no_grad():
-        input_ids = gen_tokenizer.encode(text, return_tensors="pt").to(gen_model.device)
+        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()
@@ -99,7 +99,6 @@ def generate_image(prompt):
         image_data.save(tmp_file.name)
         duration = round(time.time() - start, 2)
 
-        # CLIP similarity
         image_input = clip_preprocess(image_data).unsqueeze(0).to(device)
         text_input = clip.tokenize([prompt]).to(device)
         with torch.no_grad():
@@ -111,7 +110,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")