Update app.py

app.py

@@ -2,40 +2,31 @@ import os
 import gradio as gr
 import torch
 from PIL import Image
-from transformers import pipeline, AutoProcessor, AutoModelForVision2Seq, AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline, AutoProcessor, AutoModelForVision2Seq
 
-# --- SETUP TOKEN ---
-HF_TOKEN = os.getenv("HF_TOKEN")  # Set in env or Secrets on Spaces
+# Use HF_TOKEN from environment for private models if needed (can add below if your Gemma is gated)
+HF_TOKEN = os.environ.get("HF_TOKEN")
 
-# --- DEVICE ---
+# Auto-detect 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)
+# Load BLIP for captioning
+processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+blip_model = AutoModelForVision2Seq.from_pretrained("Salesforce/blip-image-captioning-large")
 caption_pipe = pipeline(
-    task="image-to-text",
+    "image-to-text",
     model=blip_model,
     tokenizer=processor.tokenizer,
     image_processor=processor.image_processor,
     device=DEVICE,
-    token=HF_TOKEN,
 )
 
-# --- 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)
+# Load Gemma for text generation (pick your Gemma checkpoint here)
 gemma_pipe = pipeline(
     "text-generation",
-    model=gemma_model,
-    tokenizer=gemma_tokenizer,
+    model="google/gemma-2b-it",  # Change this to any working Gemma instruct model!
     device=DEVICE,
-    max_new_tokens=384,
-    do_sample=False,
-    token=HF_TOKEN,
+    # token=HF_TOKEN  # Uncomment if your Gemma model requires a token
 )
 
 def get_recommendations():
@@ -52,81 +43,61 @@ def get_recommendations():
         "https://i.imgur.com/Xj92Cjv.jpeg",
     ]
 
+def clean_output(text):
+    # Remove prompt echoes if any (Gemma sometimes echoes)
+    if "Description:" in text:
+        text = text.split("Description:", 1)[-1]
+    if "Category:" in text:
+        text = text.split("Category:", 1)[-1]
+    return text.strip()
+
 def process(image: Image):
     if image is None:
         return "", "", "", get_recommendations()
 
-    # 1. BLIP: Caption
+    # 1. BLIP captioning
     caption_res = caption_pipe(image, max_new_tokens=64)
-    description = caption_res[0]["generated_text"].strip()
+    desc = 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()
+    # 2. Gemma: Category
+    cat_prompt = f"Classify the following ad in one or two words. Description: {desc}"
+    cat_out = gemma_pipe(cat_prompt, max_new_tokens=16)[0]['generated_text'].strip()
+    cat_out = clean_output(cat_out)
 
-    # 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. Gemma: Analysis (5 sentences)
+    ana_prompt = (
+        f"Describe in exactly five sentences what this ad communicates and its emotional impact. Description: {desc}"
     )
-    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])
+    ana_out = gemma_pipe(ana_prompt, max_new_tokens=120)[0]['generated_text'].strip()
+    ana_out = clean_output(ana_out)
 
-    # 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."
+    # 4. Gemma: Suggestions (5 bullets)
+    sug_prompt = (
+        f"Suggest five practical improvements for this ad. Each suggestion must be unique, address a different aspect (message, visuals, call to action, targeting, or layout), start with '- ', and be one sentence. Description: {desc}"
     )
-    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])
+    sug_out = gemma_pipe(sug_prompt, max_new_tokens=120)[0]['generated_text'].strip()
+    sug_out = clean_output(sug_out)
+    # Keep only lines that start with '-'
+    sug_lines = [line for line in sug_out.splitlines() if line.strip().startswith('-')]
+    suggestions = "\n".join(sug_lines[:5]) if sug_lines else sug_out
 
-    return cat_out, analysis, suggestions, get_recommendations()
+    return cat_out, ana_out, suggestions, get_recommendations()
 
 def main():
-    with gr.Blocks(title="Smart Ad Analyzer (BLIP+Gemma)") as demo:
+    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.
+            Upload your ad image below and instantly get expert feedback.  
             Category, analysis, improvement suggestions—and example ads for inspiration.
             """
         )
         with gr.Row():
             inp = gr.Image(type='pil', label='Upload Ad Image')
             with gr.Column():
-                cat_out = gr.Textbox(label='🗂️ Ad Category', interactive=False)
-                ana_out = gr.Textbox(label='📊 Ad Analysis', lines=5, interactive=False)
-                sug_out = gr.Textbox(label='🛠️ Improvement Suggestions', lines=5, interactive=False)
+                cat_out = gr.Textbox(label='Ad Category', interactive=False)
+                ana_out = gr.Textbox(label='Ad Analysis', lines=5, interactive=False)
+                sug_out = gr.Textbox(label='Improvement Suggestions', lines=5, interactive=False)
                 btn = gr.Button('Analyze Ad', variant='primary')
         gallery = gr.Gallery(label='Example Ads')
         btn.click(