Update app.py

app.py

@@ -1,25 +1,44 @@
 import os
 import gradio as gr
-from transformers import pipeline
+import torch
+from PIL import Image
+from transformers import pipeline, AutoProcessor, AutoModelForVision2Seq, AutoTokenizer, AutoModelForCausalLM
 
-HF_TOKEN = os.environ.get("HF_TOKEN")
+# --- SETUP TOKEN ---
+HF_TOKEN = os.getenv("HF_TOKEN")  # Set in env or Secrets on Spaces
 
-# 1. Image Captioning (fast lightweight)
-captioner = pipeline(
-    "image-to-text",
-    model="google/paligemma-3b-pt-224",
-    token=HF_TOKEN
+# --- DEVICE ---
+DEVICE = 0 if torch.cuda.is_available() else -1
+
+# --- BLIP: Captioning ---
+processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-large", token=HF_TOKEN)
+blip_model = AutoModelForVision2Seq.from_pretrained("Salesforce/blip-image-captioning-large", token=HF_TOKEN)
+caption_pipe = pipeline(
+    task="image-to-text",
+    model=blip_model,
+    tokenizer=processor.tokenizer,
+    image_processor=processor.image_processor,
+    device=DEVICE,
+    token=HF_TOKEN,
 )
 
-# 2. Text Feedback/Analysis (fast lightweight)
-reviewer = pipeline(
+# --- GEMMA: Text Generation ---
+# Swap this to your preferred Gemma model ID, e.g. "google/gemma-2b-it"
+GEMMA_MODEL = "google/gemma-2b-it"
+
+gemma_tokenizer = AutoTokenizer.from_pretrained(GEMMA_MODEL, token=HF_TOKEN)
+gemma_model = AutoModelForCausalLM.from_pretrained(GEMMA_MODEL, token=HF_TOKEN)
+gemma_pipe = pipeline(
     "text-generation",
-    model="google/gemma-1.1-2b-it",
-    token=HF_TOKEN
+    model=gemma_model,
+    tokenizer=gemma_tokenizer,
+    device=DEVICE,
+    max_new_tokens=384,
+    do_sample=False,
+    token=HF_TOKEN,
 )
 
 def get_recommendations():
-    # Returns list of 10 example ad image URLs
     return [
         "https://i.imgur.com/InC88PP.jpeg",
         "https://i.imgur.com/7BHfv4T.png",
@@ -33,43 +52,69 @@ def get_recommendations():
         "https://i.imgur.com/Xj92Cjv.jpeg",
     ]
 
-def process(image):
+def process(image: Image):
     if image is None:
         return "", "", "", get_recommendations()
-    
-    # 1. Caption/ad description
-    cap = captioner(image)[0]["generated_text"].strip()
-    
-    # 2. Build feedback prompt
-    prompt = (
-        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 ways to improve this ad. Each suggestion should be a short, practical sentence."
+
+    # 1. BLIP: Caption
+    caption_res = caption_pipe(image, max_new_tokens=64)
+    description = caption_res[0]["generated_text"].strip()
+
+    # 2. GEMMA: Category
+    prompt_cat = f"This is an ad image. Description: {description}\n\nProvide a concise category label for this ad (e.g. Food, Fitness, Technology):"
+    cat_out = gemma_pipe(prompt_cat)[0]['generated_text'].splitlines()[0].strip()
+
+    # 3. GEMMA: Five-sentence analysis
+    prompt_ana = (
+        f"This is an ad image. Description: {description}\n\n"
+        "Write exactly five sentences explaining what this ad communicates and its emotional impact."
     )
-    # 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 line in lines:
-        if "category" in line.lower():
-            cat = line.split(":", 1)[-1].strip()
-        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()
+    ana_raw = gemma_pipe(prompt_ana)[0]['generated_text'].strip()
+    # Get only first five sentences.
+    import re
+    sentences = re.split(r'(?<=[.!?])\s+', ana_raw)
+    analysis = " ".join(sentences[:5])
+
+    # 4. GEMMA: Five suggestions (bullets, unique)
+    prompt_sug = (
+        f"This is an ad image. Description: {description}\n\n"
+        "Suggest five unique, practical improvements for this ad. Each must address a different aspect (message, visuals, call-to-action, targeting, layout, or design). "
+        "Each suggestion must be one sentence and start with '- '. Do NOT repeat suggestions."
+    )
+    sug_raw = gemma_pipe(prompt_sug)[0]['generated_text']
+    bullets = []
+    seen = set()
+    for line in sug_raw.splitlines():
+        if line.startswith("-"):
+            suggestion = line.strip()
+            if suggestion and suggestion not in seen:
+                bullets.append(suggestion)
+                seen.add(suggestion)
+        elif line.strip():
+            suggestion = "- " + line.strip()
+            if suggestion and suggestion not in seen:
+                bullets.append(suggestion)
+                seen.add(suggestion)
+        if len(bullets) == 5:
+            break
+    # Defaults if not enough bullets
+    defaults = [
+        "- Make the main headline more eye-catching.",
+        "- Add a clear and visible call-to-action button.",
+        "- Use contrasting colors for better readability.",
+        "- Highlight the unique selling point of the product.",
+        "- Simplify the design to reduce clutter."
+    ]
+    for default in defaults:
+        if len(bullets) < 5 and default not in seen:
+            bullets.append(default)
+    suggestions = "\n".join(bullets[:5])
+
+    return cat_out, analysis, suggestions, get_recommendations()
 
 def main():
-    with gr.Blocks(title="Smart Ad Analyzer (Fast Edition)") as demo:
-        gr.Markdown("## 📢 Smart Ad Analyzer (Fast Edition)")
+    with gr.Blocks(title="Smart Ad Analyzer (BLIP+Gemma)") as demo:
+        gr.Markdown("## 📢 Smart Ad Analyzer (BLIP + Gemma)")
         gr.Markdown(
             """
             Upload your ad image below and instantly get expert feedback.