Update app.py

app.py

@@ -1,63 +1,146 @@
+import re
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-# Load model and tokenizer
-MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
-# Uncomment the next line to run on GPU if available
-# model = model.to("cuda")
-
-def improve_ad(ad_analysis):
-    prompt = (
-        "You are an expert ad consultant. Based on the following ad description or analysis, "
-        "give 5 unique, actionable suggestions to improve the ad. "
-        "Each suggestion must be on a new line, start with '- ', and should NOT repeat. "
-        "If possible, cover different aspects (messaging, design, call-to-action, targeting, layout, emotion, offer). "
-        "\n\nAd Analysis or Description:\n"
-        f"{ad_analysis}\n\n"
-        "Improvement Suggestions:"
+import torch
+from PIL import Image
+from transformers import (
+    pipeline,
+    AutoProcessor,
+    AutoModelForVision2Seq,
+    AutoTokenizer,
+    AutoModelForCausalLM,
+)
+
+# Auto-detect CPU/GPU
+DEVICE = 0 if torch.cuda.is_available() else -1
+
+# Load BLIP for image 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",
+    model=blip_model,
+    tokenizer=processor.tokenizer,
+    image_processor=processor.image_processor,
+    device=DEVICE,
+)
+
+# Load Mistral-7B-Instruct for text tasks
+INSTRUCT_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"
+instr_tokenizer = AutoTokenizer.from_pretrained(INSTRUCT_MODEL)
+instr_model = AutoModelForCausalLM.from_pretrained(INSTRUCT_MODEL)
+instr_pipe = pipeline(
+    "text-generation",
+    model=instr_model,
+    tokenizer=instr_tokenizer,
+    device=DEVICE,
+    max_new_tokens=256,
+    temperature=0.5,   # Lowered temp for less repetition
+    do_sample=True,
+    repetition_penalty=1.2,  # Discourage repetition
+)
+
+def get_recommendations():
+    return [
+        "https://i.imgur.com/InC88PP.jpeg",
+        "https://i.imgur.com/7BHfv4T.png",
+        "https://i.imgur.com/wp3Wzc4.jpeg",
+        "https://i.imgur.com/5e2xOA4.jpeg",
+        "https://i.imgur.com/txjRk98.jpeg",
+        "https://i.imgur.com/rQ4AYl0.jpeg",
+        "https://i.imgur.com/bDzwD04.jpeg",
+        "https://i.imgur.com/fLMngXI.jpeg",
+        "https://i.imgur.com/nYEJzxt.png",
+        "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)
+    desc = caption_res[0]["generated_text"].strip()
+
+    # 2. Category (via instruct model)
+    cat_prompt = f"<s>[INST] The following is an ad: {desc}\nWhat is the concise category of this ad (e.g. Food, Fitness, Tech, Fashion, Travel, etc)? Respond with one word. [/INST]"
+    cat_out = instr_pipe(cat_prompt)[0]['generated_text'].split("\n")[0].strip()
+
+    # 3. Analysis (via instruct model)
+    ana_prompt = (
+        f"<s>[INST] Here is an ad description: {desc}\n"
+        "Write five sentences analyzing the ad. Cover: 1) What product or service is advertised, 2) Main visual elements, 3) The target audience, 4) The marketing message, 5) The emotional appeal or persuasion technique. Each sentence should be informative and not repeat. [/INST]"
     )
-    input_ids = tokenizer(prompt, return_tensors="pt")
-    # Uncomment the next line to run on GPU if available
-    # input_ids = input_ids.to("cuda")
-    output = model.generate(
-        **input_ids,
-        max_new_tokens=256,
-        do_sample=False,
-        pad_token_id=tokenizer.eos_token_id,
+    ana_raw = instr_pipe(ana_prompt)[0]['generated_text']
+    ana_sentences = re.split(r'(?<=[.!?])\s+', ana_raw.strip())
+    analysis = " ".join(ana_sentences[:5])
+
+    # 4. Five suggestions (via instruct model, filter repetitions)
+    sug_prompt = (
+        f"<s>[INST] Given this ad analysis: {analysis}\n"
+        "Suggest five practical, unique improvements for the ad. Each should address a different aspect (message, visuals, call-to-action, targeting, layout, or design). Each suggestion must be one sentence, not repeat, and start with '- '. [/INST]"
     )
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    # Extract suggestions: Find the part after "Improvement Suggestions:"
-    if "Improvement Suggestions:" in response:
-        response = response.split("Improvement Suggestions:")[-1]
-    # Filter and clean up to get exactly 5 suggestions
-    suggestions = []
-    for line in response.splitlines():
-        line = line.strip()
-        if line.startswith("-") and line not in suggestions:
-            suggestions.append(line)
-        if len(suggestions) == 5:
+    sug_raw = instr_pipe(sug_prompt)[0]['generated_text']
+    bullets = []
+    seen = set()
+    for line in sug_raw.splitlines():
+        if line.startswith("-"):
+            suggestion = line.strip()
+            if suggestion not in seen and suggestion:
+                bullets.append(suggestion)
+                seen.add(suggestion)
+        elif line.strip():
+            suggestion = "- " + line.strip()
+            if suggestion not in seen and suggestion:
+                bullets.append(suggestion)
+                seen.add(suggestion)
+        if len(bullets) == 5:
             break
-    # If not enough, pad with generic
+    # Add generic suggestions only if needed
     defaults = [
-        "- Add a clear and attractive call-to-action.",
-        "- Improve the visual contrast and headline font size.",
-        "- Refine the main message for clarity.",
-        "- Target a more specific audience segment.",
-        "- Use more emotional or persuasive language."
+        "- 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 d in defaults:
-        if len(suggestions) < 5 and d not in suggestions:
-            suggestions.append(d)
-    return "\n".join(suggestions[:5])
-
-with gr.Blocks() as demo:
-    gr.Markdown("## 🦙 Llama-3.1 Ad Improvement Consultant\nPaste your ad analysis or description below, and get 5 unique suggestions to improve it, powered by Meta Llama 3.1 8B Instruct (runs 100 percent locally, no API required).")
-    inp = gr.Textbox(label="Ad Analysis or Description", lines=6, placeholder="Paste your ad description, image caption, or analysis here...")
-    out = gr.Textbox(label="Improvement Suggestions", lines=7)
-    btn = gr.Button("Get Suggestions")
-    btn.click(fn=improve_ad, inputs=inp, outputs=out)
+    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") as demo:
+        gr.Markdown("## 📢 Smart Ad Analyzer")
+        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!
+            """
+        )
+        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)
+                btn = gr.Button('Analyze Ad', variant='primary')
+        gallery = gr.Gallery(label='Example Ads')
+        btn.click(
+            fn=process,
+            inputs=[inp],
+            outputs=[cat_out, ana_out, sug_out, gallery],
+        )
+        gr.Markdown('Made by Simon Thalmay')
+    return demo
 
 if __name__ == "__main__":
+    demo = main()
     demo.launch()