Update app.py

app.py

@@ -2,6 +2,22 @@ import os
 import gradio as gr
 from transformers import pipeline
 
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# 1. Image Captioning (fast lightweight)
+captioner = pipeline(
+    "image-to-text",
+    model="google/paligemma-3b-pt-224",
+    token=HF_TOKEN
+)
+
+# 2. Text Feedback/Analysis (fast lightweight)
+reviewer = pipeline(
+    "text-generation",
+    model="google/gemma-1.1-2b-it",
+    token=HF_TOKEN
+)
+
 def get_recommendations():
     # Returns list of 10 example ad image URLs
     return [
@@ -17,50 +33,43 @@ def get_recommendations():
         "https://i.imgur.com/Xj92Cjv.jpeg",
     ]
 
-# BLIP for image captioning (always on CPU, runs fast)
-captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base", device="cpu")
-
-# Gemma 1B for text generation
-gemma = pipeline("text-generation", model="google/gemma-1.1-1b-it", device="cpu")
-
 def process(image):
     if image is None:
         return "", "", "", get_recommendations()
-
-    # 1. Caption image
-    cap_result = captioner(image)
-    caption = cap_result[0]['generated_text'] if cap_result else "No caption generated."
-
-    # 2. Compose a prompt for Gemma (category, analysis, suggestions)
+    
+    # 1. Caption/ad description
+    cap = captioner(image)[0]["generated_text"].strip()
+    
+    # 2. Build feedback prompt
     prompt = (
-        f"Here is a description of an ad image: {caption}\n"
-        "1. Assign a concise ad category label (e.g., 'Food', 'Fitness').\n"
+        f"Ad description: {cap}\n"
+        "Your task: \n"
+        "1. Give a concise category label for this ad (e.g., 'Food', 'Fitness').\n"
         "2. Write exactly five sentences analyzing what this ad communicates and its emotional impact.\n"
-        "3. Suggest five specific ways to improve this ad, each as a short, practical sentence.\n"
-        "Answer in three parts clearly marked as Category, Analysis, and Suggestions."
+        "3. Suggest five ways to improve this ad. Each suggestion should be a short, practical sentence."
     )
-
-    gemma_out = gemma(prompt, max_new_tokens=256)[0]['generated_text']
-    # Split results
-    lines = gemma_out.split('\n')
+    # 3. Run through Gemma
+    resp = reviewer(prompt, max_new_tokens=256)[0]["generated_text"]
+    
+    # 4. Simple parsing logic
+    lines = resp.split('\n')
     cat, analysis, suggestions = "", "", ""
-    for i, line in enumerate(lines):
+    for line in lines:
         if "category" in line.lower():
             cat = line.split(":", 1)[-1].strip()
-        elif "analysis" in line.lower():
-            analysis = "\n".join(lines[i+1:i+6])
-        elif "suggestions" in line.lower():
-            suggestions = "\n".join(lines[i+1:i+6])
-    # Fallback if Gemma output is not perfectly formatted
+        elif "analysis" in line.lower() or "sentence" in line.lower():
+            analysis += line + " "
+        elif "suggestion" in line.lower() or line.strip().startswith("-"):
+            suggestions += line + "\n"
     if not cat: cat = lines[0][:80]
     if not analysis: analysis = "\n".join(lines[1:6])
     if not suggestions: suggestions = "\n".join(lines[6:11])
-
+    
     return cat.strip(), analysis.strip(), suggestions.strip(), get_recommendations()
 
 def main():
-    with gr.Blocks(title="Smart Ad Analyzer (BLIP+Gemma Edition)") as demo:
-        gr.Markdown("## 📢 Smart Ad Analyzer (BLIP+Gemma Edition)")
+    with gr.Blocks(title="Smart Ad Analyzer (Fast Edition)") as demo:
+        gr.Markdown("## 📢 Smart Ad Analyzer (Fast Edition)")
         gr.Markdown(
             """
             Upload your ad image below and instantly get expert feedback.