Update app.py

app.py

@@ -1,77 +1,18 @@
-import re
+import os
 import gradio as gr
-import torch
+from huggingface_hub import InferenceClient
 from PIL import Image
-from transformers import (
-    pipeline,
-    AutoProcessor,
-    AutoModelForVision2Seq,
-    AutoTokenizer,
-    AutoModelForSeq2SeqLM,
-)
-
-# Auto-detect CPU/GPU
-DEVICE = 0 if torch.cuda.is_available() else -1
-
-# Load BLIP captioning model
-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",
-    model=blip_model,
-    tokenizer=processor.tokenizer,
-    image_processor=processor.image_processor,
-    device=DEVICE,
-)
-
-# Load Flan-T5 for text-to-text
-FLAN_MODEL = "google/flan-t5-large"
-flan_tokenizer = AutoTokenizer.from_pretrained(FLAN_MODEL)
-flan_model = AutoModelForSeq2SeqLM.from_pretrained(FLAN_MODEL)
-
-category_pipe = pipeline(
-    "text2text-generation",
-    model=flan_model,
-    tokenizer=flan_tokenizer,
-    device=DEVICE,
-    max_new_tokens=32,
-    do_sample=True,
-    temperature=1.0,
-)
+import tempfile
 
-analysis_pipe = pipeline(
-    "text2text-generation",
-    model=flan_model,
-    tokenizer=flan_tokenizer,
-    device=DEVICE,
-    max_new_tokens=256,
-    do_sample=True,
-    temperature=1.0,
-)
-
-suggestion_pipe = pipeline(
-    "text2text-generation",
-    model=flan_model,
-    tokenizer=flan_tokenizer,
-    device=DEVICE,
-    max_new_tokens=256,
-    do_sample=True,
-    temperature=1.6,     # Higher temperature for more variety
-    top_p=0.95,
-)
-
-expansion_pipe = pipeline(
-    "text2text-generation",
-    model=flan_model,
-    tokenizer=flan_tokenizer,
-    device=DEVICE,
-    max_new_tokens=128,
-    do_sample=True,
-    temperature=1.0,
+# ---- Gemma-3 setup ----
+client = InferenceClient(
+    model="google/gemma-3-4b-it",
+    api_key=os.environ.get("HF_TOKEN", None),
+    provider="featherless-ai",  # or "huggingface_hub"
 )
 
 def get_recommendations():
-    # Returns list of 10 example ad image URLs
+    # As before: returns list of 10 example ad image URLs
     return [
         "https://i.imgur.com/InC88PP.jpeg",
         "https://i.imgur.com/7BHfv4T.png",
@@ -85,88 +26,82 @@ def get_recommendations():
         "https://i.imgur.com/Xj92Cjv.jpeg",
     ]
 
-def process(image: Image):
-    if image is None:
-        return "", "", "", get_recommendations()
-
-    # 1. BLIP caption
-    caption_res = caption_pipe(image, max_new_tokens=64)
-    raw_caption = caption_res[0]["generated_text"].strip()
-
-    # 1a. Expand caption if too short
-    if len(raw_caption.split()) < 3:
-        exp = expansion_pipe(f"Expand into a detailed description: {raw_caption}")
-        desc = exp[0]["generated_text"].strip()
-    else:
-        desc = raw_caption
-
-    # 2. Category
-    cat_prompt = (
-        f"Description: {desc}\n\n"
-        "Provide a concise category label for this ad (e.g. 'Food', 'Fitness'):"
-    )
-    cat_out = category_pipe(cat_prompt)[0]["generated_text"].splitlines()[0].strip()
-
-    # 3. Five-sentence analysis
-    ana_prompt = (
-        f"Description: {desc}\n\n"
-        "Write exactly five sentences explaining what this ad communicates and its emotional impact."
+def gemma_image_analysis(image: Image):
+    # Upload PIL image to Hugging Face
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+        image.save(tmp, format="PNG")
+        image_url = client.upload(tmp.name)
+    prompt = (
+        "You are an expert ad analyst. "
+        "Please give a short category for this ad, a detailed analysis of its message, visuals, and emotional impact in five sentences, "
+        "and five unique, actionable improvement suggestions (as bullet points), each addressing a different aspect (visuals, message, call-to-action, targeting, or layout). "
+        "Output should have clear sections: 'Category', 'Analysis', and 'Suggestions'."
     )
-    ana_raw = analysis_pipe(ana_prompt)[0]["generated_text"].strip()
-    sentences = re.split(r'(?<=[.!?])\s+', ana_raw)
-    analysis = " ".join(sentences[:5])
-
-    # 4. Five improvement suggestions (model + fallback defaults, no repeats)
-    sug_prompt = (
-        f"Ad description: {desc}\n"
-        f"Ad analysis: {analysis}\n\n"
-        "Suggest five ways to improve this ad. Write each suggestion as one practical sentence starting with '- '."
-    )
-    sug_raw = suggestion_pipe(sug_prompt)[0]["generated_text"].strip()
-    all_sugs = [line.strip() for line in sug_raw.splitlines() if line.strip().startswith("-")]
-
-    # Filter exact duplicates, keep order, allow model output first
-    unique_sugs = []
-    seen = set()
-    for s in all_sugs:
-        norm = s.lower().strip(".:; ")
-        if norm not in seen and len(norm) > 4:
-            unique_sugs.append(s)
-            seen.add(norm)
-        if len(unique_sugs) == 5:
-            break
-    # Add default suggestions only if needed
-    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."
+    messages = [
+        {
+            "role": "system",
+            "content": [{"type": "text", "text": "You are a helpful assistant."}]
+        },
+        {
+            "role": "user",
+            "content": [
+                {"type": "image_url", "image_url": {"url": image_url}},
+                {"type": "text", "text": prompt}
+            ]
+        }
     ]
-    for d in defaults:
-        norm = d.lower().strip(".:; ")
-        if len(unique_sugs) < 5 and norm not in seen:
-            unique_sugs.append(d)
-            seen.add(norm)
-    suggestions = "\n".join(unique_sugs[:5])
-
-    return cat_out, analysis, suggestions, get_recommendations()
+    # API call to Gemma
+    result = client.chat.completions.create(
+        model="google/gemma-3-4b-it",
+        messages=messages,
+        max_tokens=500,
+    )
+    return result.choices[0].message["content"]
 
+def process(image):
+    if image is None:
+        return "", "", "", get_recommendations()
+    # Use Gemma to get all outputs in one string
+    full_output = gemma_image_analysis(image)
+    # Parse Gemma's response (very basic, you can make fancier with regex etc)
+    # Try to split by headings if present
+    category = ""
+    analysis = ""
+    suggestions = ""
+    lines = full_output.splitlines()
+    section = None
+    for line in lines:
+        l = line.strip()
+        if l.lower().startswith("category"):
+            section = "cat"
+            category = ""
+        elif l.lower().startswith("analysis"):
+            section = "ana"
+            analysis = ""
+        elif l.lower().startswith("suggestion"):
+            section = "sug"
+            suggestions = ""
+        elif section == "cat":
+            category += l + "\n"
+        elif section == "ana":
+            analysis += l + "\n"
+        elif section == "sug":
+            suggestions += l + "\n"
+    category = category.strip()
+    analysis = analysis.strip()
+    suggestions = suggestions.strip()
+    # If parsing failed, put everything in analysis
+    if not (category or analysis or suggestions):
+        analysis = full_output.strip()
+    return category, analysis, suggestions, get_recommendations()
+
+# ---- Gradio UI ----
 def main():
-    with gr.Blocks(title="Smart Ad Analyzer") as demo:
-        gr.Markdown("## 📢 Smart Ad Analyzer")
+    with gr.Blocks(title="Smart Ad Analyzer (Gemma-powered)") as demo:
+        gr.Markdown("## 📢 Smart Ad Analyzer (Gemma-3 Edition)")
         gr.Markdown(
-            """
-            **Upload your ad image below and instantly get expert feedback.**
-
-            This AI tool will analyze your ad and provide:
-            - 📂 **Category** — What type of ad is this?
-            - 📊 **In-depth Analysis** — Five detailed sentences covering message, visuals, emotional impact, and more.
-            - 🚀 **Improvement Suggestions** — Five actionable, unique ways to make your ad better.
-            - 📸 **Inspiration Gallery** — See other effective ads for ideas.
-
-            Perfect for marketers, founders, designers, and anyone looking to boost ad performance with actionable insights!
-            """
+            "**Upload your ad image below and instantly get expert feedback.**<br>"
+            "Category, analysis, improvement suggestions—and example ads for inspiration."
         )
         with gr.Row():
             inp = gr.Image(type='pil', label='Upload Ad Image')
@@ -181,7 +116,7 @@ def main():
             inputs=[inp],
             outputs=[cat_out, ana_out, sug_out, gallery],
         )
-        gr.Markdown('Made by Simon Thalmay')
+        gr.Markdown('Made by Simon Thalmay • Powered by google/gemma-3-4b-it')
     return demo
 
 if __name__ == "__main__":