Update app.py

app.py

@@ -1,15 +1,25 @@
+# app.py
+
 import re
 import gradio as gr
-from transformers import pipeline
-
-# 1) Image-to-text: ChatDOC/OCRFlux-3B for rich description
-image_to_text = pipeline(
-    "image-to-text",
-    model="ChatDOC/OCRFlux-3B"
+from PIL import Image
+from transformers import (
+    Blip2Processor,
+    Blip2ForConditionalGeneration,
+    pipeline
 )
 
+# 1) BLIP-2 for richer image captions
+processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
+model     = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b")
+
+def generate_caption(image: Image) -> str:
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model.generate(**inputs)
+    return processor.decode(outputs[0], skip_special_tokens=True)
+
 # 2) Helper to build Flan-T5-small text pipelines (temperature=1.0)
-def make_pipeline(model_name, max_tokens):
+def make_pipeline(model_name: str, max_tokens: int):
     return pipeline(
         "text2text-generation",
         model=model_name,
@@ -41,19 +51,14 @@ def get_recommendations():
         "https://i.imgur.com/Xj92Cjv.jpeg",
     ]
 
-# Step A: Use OCRFlux to generate a detailed caption
-def generate_caption(image):
-    result = image_to_text(image)
-    return result[0]["generated_text"].strip()
-
 # Step B: Flan interprets caption into concise category
-def generate_category(caption):
+def generate_category(caption: str) -> str:
     prompt = f"Caption: {caption}\nProvide a concise category label for this ad."
     raw = category_generator(prompt)[0]["generated_text"].strip()
     return raw.splitlines()[0]
 
 # Step C: Flan produces exactly five-sentence analysis
-def generate_analysis(caption):
+def generate_analysis(caption: str) -> str:
     prompt = (
         f"Caption: {caption}\n"
         "In exactly five sentences, explain what this ad communicates and its emotional impact."
@@ -63,7 +68,7 @@ def generate_analysis(caption):
     return " ".join(sentences[:5])
 
 # Step D: Flan suggests five actionable bullet-point improvements
-def generate_suggestions(caption):
+def generate_suggestions(caption: str) -> str:
     prompt = (
         f"Caption: {caption}\n"
         "Suggest five distinct improvements as bullet points. Each line must start with '- '."
@@ -79,7 +84,7 @@ def generate_suggestions(caption):
     return "\n".join(lines[:5])
 
 # Orchestrator: process image through all steps
-def process(image):
+def process(image: Image):
     caption     = generate_caption(image)
     category    = generate_category(caption)
     analysis    = generate_analysis(caption)