Update app.py

app.py

@@ -1,177 +1,63 @@
-import re
 import gradio as gr
-import torch
-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,
-)
-
-analysis_pipe = pipeline(
-    "text2text-generation",
-    model=flan_model,
-    tokenizer=flan_tokenizer,
-    device=DEVICE,
-    max_new_tokens=256,
-    do_sample=True,
-    temperature=1.0,
-)
-
-# Higher temp for more variety in suggestions
-suggestion_pipe = pipeline(
-    "text2text-generation",
-    model=flan_model,
-    tokenizer=flan_tokenizer,
-    device=DEVICE,
-    max_new_tokens=256,
-    do_sample=True,
-    temperature=1.3,
-)
-
-expansion_pipe = pipeline(
-    "text2text-generation",
-    model=flan_model,
-    tokenizer=flan_tokenizer,
-    device=DEVICE,
-    max_new_tokens=128,
-    do_sample=False,
-)
-
-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()
-
-    # BLIP caption
-    caption_res = caption_pipe(image, max_new_tokens=64)
-    raw_caption = caption_res[0]["generated_text"].strip()
-    desc = raw_caption if len(raw_caption.split()) >= 3 else expansion_pipe(f"Expand into a detailed description: {raw_caption}")[0]["generated_text"].strip()
-
-    # Category
-    cat_prompt = f"Description: {desc}\n\nProvide a concise category label for this ad (e.g. 'Food', 'Fitness'):"
-    cat_out = category_pipe(cat_prompt)[0]["generated_text"].splitlines()[0].strip()
-
-    # Five-sentence analysis
-    ana_prompt = (
-        f"Description: {desc}\n\n"
-        "Write exactly five sentences explaining what this ad communicates and its emotional impact."
+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:"
     )
-    ana_raw = analysis_pipe(ana_prompt)[0]["generated_text"].strip()
-    sentences = re.split(r'(?<=[.!?])\s+', ana_raw)
-    analysis = " ".join(sentences[:5]).strip()
-
-    # **KEY CHANGE**: Use analysis in suggestion prompt!
-    sug_prompt = (
-        f"Ad description: {desc}\n"
-        f"Ad analysis: {analysis}\n\n"
-        "Based on this, suggest five unique and specific improvements for this ad. Each suggestion should be one clear sentence starting with '- ' and focus on a different aspect, like message, visuals, call-to-action, color, clarity, layout, targeting, or emotional impact. Do NOT repeat suggestions."
+    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,
     )
-    sug_raw = suggestion_pipe(sug_prompt)[0]["generated_text"].strip()
-    all_sugs = [line.strip() for line in sug_raw.splitlines() if line.strip().startswith("-")]
-    unique_sugs = []
-    seen = set()
-    for line in all_sugs:
-        line_clean = line.lower().strip().rstrip(".")
-        if line_clean not in seen and len(line_clean) > 4:
-            unique_sugs.append(line)
-            seen.add(line_clean)
-        if len(unique_sugs) == 5:
+    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:
             break
-
-    # Add defaults if needed
+    # If not enough, pad with generic
     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."
+        "- 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."
     ]
     for d in defaults:
-        d_clean = d.lower().strip().rstrip(".")
-        if len(unique_sugs) < 5 and d_clean not in seen:
-            unique_sugs.append(d)
-            seen.add(d_clean)
-    suggestions = "\n".join(unique_sugs[: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.
+        if len(suggestions) < 5 and d not in suggestions:
+            suggestions.append(d)
+    return "\n".join(suggestions[:5])
 
-            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
+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)
 
 if __name__ == "__main__":
-    demo = main()
     demo.launch()