Update app.py

app.py

@@ -1,23 +1,29 @@
 import re
 import gradio as gr
+import torch
 from PIL import Image
-from transformers import pipeline
+from transformers import pipeline, AutoProcessor, AutoModelForVision2Seq
 
-# 1) BLIP captioner (rich COCO captions)
+# Auto-detect CPU/GPU
+device = 0 if torch.cuda.is_available() else -1
+
+# 1) BLIP captioner
+processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+model = AutoModelForVision2Seq.from_pretrained("Salesforce/blip-image-captioning-large")
 caption_pipe = pipeline(
-    task="image-to-text",
-    model="Salesforce/blip-image-captioning-large",
-    device=-1,              # force CPU
-    max_length=64,
-    do_sample=False,
+    "image-to-text",
+    model=model,
+    processor=processor,
+    device=device
 )
 
-# 2) Flan-T5 for text‑to‑text
+# 2) Flan-T5 for text-to-text
 FLAN = "google/flan-t5-large"
 category_pipe = pipeline(
     "text2text-generation",
     model=FLAN,
     tokenizer=FLAN,
+    device=device,
     max_new_tokens=32,
     do_sample=True,
     temperature=1.0,
@@ -26,6 +32,7 @@ analysis_pipe = pipeline(
     "text2text-generation",
     model=FLAN,
     tokenizer=FLAN,
+    device=device,
     max_new_tokens=256,
     do_sample=True,
     temperature=1.0,
@@ -34,15 +41,17 @@ suggestion_pipe = pipeline(
     "text2text-generation",
     model=FLAN,
     tokenizer=FLAN,
+    device=device,
     max_new_tokens=256,
     do_sample=True,
     temperature=1.0,
 )
-# 3) Caption expander when BLIP is too short
+# Expander when BLIP caption is too short
 expansion_pipe = pipeline(
     "text2text-generation",
     model=FLAN,
     tokenizer=FLAN,
+    device=device,
     max_new_tokens=128,
     do_sample=False,
 )
@@ -62,81 +71,84 @@ def get_recommendations():
         "https://i.imgur.com/Xj92Cjv.jpeg",
     ]
 
-
+# Main processing function
 def process(image: Image):
     # 1) BLIP caption
-    caption = caption_pipe(image)[0]["generated_text"].strip()
+    caption = caption_pipe(image, max_new_tokens=64, do_sample=False)[0]['generated_text'].strip()
 
-    # 1a) Expand if too short
+    # 1a) Expand caption if too short
     if len(caption.split()) < 3:
-        exp_prompt = f"Expand this into a detailed description: {caption}"
-        desc = expansion_pipe(exp_prompt)[0]["generated_text"].strip()
+        desc = expansion_pipe(f"Expand into a detailed description: {caption}")[0]['generated_text'].strip()
     else:
         desc = caption
 
-    # 2) Single‑label category
+    # 2) Ad category
     cat_prompt = (
         f"Description: {desc}\n\n"
-        "Provide a concise category label for this ad (e.g. 'Fitness', 'Food'):"
+        "Provide a concise category label for this ad (e.g. 'Food', 'Fitness'):"
     )
-    category = category_pipe(cat_prompt)[0]["generated_text"].strip().splitlines()[0]
+    category = category_pipe(cat_prompt)[0]['generated_text'].splitlines()[0].strip()
 
-    # 3) Five‑sentence analysis
+    # 3) Five-sentence analysis
     ana_prompt = (
         f"Description: {desc}\n\n"
         "Write exactly five sentences explaining what this ad communicates and its emotional impact."
     )
-    raw_ana = analysis_pipe(ana_prompt)[0]["generated_text"].strip()
-    sents = re.split(r'(?<=[.!?])\s+', raw_ana)
-    analysis = " ".join(sents[:5])
+    raw_ana = analysis_pipe(ana_prompt)[0]['generated_text'].strip()
+    sentences = re.split(r'(?<=[.!?])\s+', raw_ana)
+    analysis = " ".join(sentences[:5])
 
-    # 4) Five bullet‑point suggestions
+    # 4) Five bullet-point suggestions
     sug_prompt = (
         f"Description: {desc}\n\n"
-        "Suggest five distinct improvements for this ad. "
-        "Each suggestion must start with '- ' and be one sentence."
+        "Suggest five distinct improvements for this ad. Each must start with '- ' and be one sentence."
     )
-    raw_sug = suggestion_pipe(sug_prompt)[0]["generated_text"].strip()
-    bullets = [l for l in raw_sug.splitlines() if l.strip().startswith("-")]
+    raw_sug = suggestion_pipe(sug_prompt)[0]['generated_text'].strip()
+    bullets = [l for l in raw_sug.splitlines() if l.startswith('-')]
     if len(bullets) < 5:
-        extras = [l.strip() for l in raw_sug.splitlines() if l.strip()]
-        for ex in extras:
-            if len(bullets) >= 5: break
-            line = ex if ex.startswith("-") else "- " + ex
-            bullets.append(line)
-    suggestions = "\n".join(bullets[:5])
+        extra_lines = [l for l in raw_sug.splitlines() if l.strip()]
+        for line in extra_lines:
+            if len(bullets) >= 5:
+                break
+            bullets.append(line if line.startswith('-') else '- ' + line)
+    suggestions = '\n'.join(bullets[:5])
 
     return caption, category, analysis, suggestions, get_recommendations()
 
-with gr.Blocks() as demo:
-    gr.Markdown("## 📢 Smart Ad Analyzer")
-    gr.Markdown(
-        "Upload an image ad to get:\n"
-        "- **BLIP Caption** (debug)\n"
-        "- **Ad Category**\n"
-        "- **Five-sentence Analysis**\n"
-        "- **Five bullet-point Suggestions**\n"
-        "- **Example Ads**"
-    )
-
-    with gr.Row():
-        inp = gr.Image(type="pil", label="Upload Ad Image")
-        with gr.Column():
-            cap_out = gr.Textbox(label="🔍 BLIP Caption", 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")
-
-    gallery = gr.Gallery(label="Example Ads")
-
-    btn.click(
-        fn=process,
-        inputs=[inp],
-        outputs=[cap_out, cat_out, ana_out, sug_out, gallery],
-    )
-
-    gr.Markdown("Made by Simon Thalmay")
-
-if __name__ == "__main__":
-    demo.launch()
+# Gradio UI
+def main():
+    with gr.Blocks() as demo:
+        gr.Markdown("## 📢 Smart Ad Analyzer")
+        gr.Markdown(
+            "Upload an image ad to get:\n"
+            "- **BLIP Caption** (raw)\n"
+            "- **Ad Category**\n"
+            "- **Five-sentence Analysis**\n"
+            "- **Five bullet-point Suggestions**\n"
+            "- **Example Ads**"
+        )
+
+        with gr.Row():
+            inp = gr.Image(type='pil', label='Upload Ad Image')
+            with gr.Column():
+                cap_out = gr.Textbox(label='🔍 BLIP Caption', 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', size='sm', variant='primary')
+
+        gallery = gr.Gallery(label='Example Ads')
+
+        btn.click(
+            fn=process,
+            inputs=[inp],
+            outputs=[cap_out, cat_out, ana_out, sug_out, gallery],
+        )
+
+        gr.Markdown('Made by Simon Thalmay')
+
+    demo.launch()
+
+if __name__ == '__main__':
+    main()
+