Update app.py

app.py

@@ -9,11 +9,11 @@ from transformers import (
     AutoTokenizer,
     AutoModelForSeq2SeqLM,
 )
-from difflib import SequenceMatcher
 
+# Auto-detect CPU/GPU
 DEVICE = 0 if torch.cuda.is_available() else -1
 
-# Load BLIP
+# 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(
@@ -24,7 +24,7 @@ caption_pipe = pipeline(
     device=DEVICE,
 )
 
-# Load Flan-T5
+# 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)
@@ -38,6 +38,7 @@ category_pipe = pipeline(
     do_sample=True,
     temperature=1.0,
 )
+
 analysis_pipe = pipeline(
     "text2text-generation",
     model=flan_model,
@@ -47,6 +48,8 @@ analysis_pipe = pipeline(
     do_sample=True,
     temperature=1.0,
 )
+
+# Set higher temperature for more variety
 suggestion_pipe = pipeline(
     "text2text-generation",
     model=flan_model,
@@ -54,8 +57,9 @@ suggestion_pipe = pipeline(
     device=DEVICE,
     max_new_tokens=256,
     do_sample=True,
-    temperature=1.05,     # Slightly more conservative than before
+    temperature=1.2,
 )
+
 expansion_pipe = pipeline(
     "text2text-generation",
     model=flan_model,
@@ -66,6 +70,7 @@ expansion_pipe = pipeline(
 )
 
 def get_recommendations():
+    # Returns list of 10 example ad image URLs
     return [
         "https://i.imgur.com/InC88PP.jpeg",
         "https://i.imgur.com/7BHfv4T.png",
@@ -79,53 +84,68 @@ def get_recommendations():
         "https://i.imgur.com/Xj92Cjv.jpeg",
     ]
 
-def unique_suggestions(suggestions):
-    """Strictly remove near-duplicates, keep order, ignore case/punct."""
-    seen = []
-    for s in suggestions:
-        norm = re.sub(r'[^a-z0-9 ]', '', s.lower())
-        if all(SequenceMatcher(None, norm, re.sub(r'[^a-z0-9 ]', '', x.lower())).ratio() < 0.91 for x in seen):
-            seen.append(s)
-    return seen
-
 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()
-    desc = raw_caption
-    if len(desc.split()) < 3:
-        exp = expansion_pipe(f"Expand into a detailed description: {desc}")
+
+    # 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\nProvide a concise category label for this ad (e.g. 'Food', 'Fitness'):"
+    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\nWrite exactly five sentences explaining what this ad communicates and its emotional impact."
+    ana_prompt = (
+        f"Description: {desc}\n\n"
+        "Write exactly five sentences explaining what this ad communicates and its emotional impact."
+    )
     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 (not forcing uniqueness in prompt)
-    sug_prompt = f"Description: {desc}\n\nSuggest five ways this ad could be improved. Each suggestion should be one sentence and start with '- '."
+    # 4. Five bullet-point suggestions (uniqueness enforced)
+    sug_prompt = (
+        f"Description: {desc}\n\n"
+        "Suggest five ways this ad could be improved. Each suggestion must be about a different aspect, such as visuals, message, call-to-action, color, clarity, layout, or audience targeting. "
+        "Each suggestion must start with '- ' and be one full sentence. Make sure each is different from the others."
+    )
     sug_raw = suggestion_pipe(sug_prompt)[0]["generated_text"].strip()
-    all_sugs = [line for line in sug_raw.splitlines() if line.strip().startswith("-")]
-    unique_sugs = unique_suggestions(all_sugs)
-
-    # Default suggestions if model outputs < 5 after filtering
+    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:
+            break
+
+    # Add non-repetitive defaults 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.",
+        "- Simplify the design to reduce clutter."
     ]
     for d in defaults:
-        if len(unique_sugs) < 5 and d not in unique_sugs:
+        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()