Update utils/planner.py

utils/planner.py

@@ -1,22 +1,44 @@
-# utils/planner.py
-
 import os
 import json
 from dotenv import load_dotenv
 from openai import OpenAI
+from PIL import Image
+import torch
+from transformers import BlipProcessor, BlipForConditionalGeneration
 
 # ----------------------------
-# 🔐 Load Environment & Client
+# 🔐 Load Environment & GPT Client
 # ----------------------------
 load_dotenv()
 client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
 
 # ----------------------------
-# 🧠 Scene Plan Extractor
+# 🧠 Load BLIP Captioning Model (once globally)
+# ----------------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
+
+# ----------------------------
+# 📸 Generate Caption from Uploaded Product Image
+# ----------------------------
+def generate_blip_caption(image: Image.Image) -> str:
+    try:
+        inputs = processor(images=image, return_tensors="pt").to(device)
+        out = blip_model.generate(**inputs, max_length=50)
+        caption = processor.decode(out[0], skip_special_tokens=True)
+        print(f"🖼️ BLIP Caption: {caption}")
+        return caption
+    except Exception as e:
+        print("❌ BLIP Captioning Error:", e)
+        return "a product image"
+
+# ----------------------------
+# 🧠 Scene Plan Extractor (GPT-4o)
 # ----------------------------
 SCENE_SYSTEM_INSTRUCTIONS = """
 You are a scene planning assistant for an AI image generation system.
-Your job is to take the user's prompt and return a structured JSON with:
+Your job is to take a caption from a product image and a user prompt, then return a structured JSON with:
 - scene (environment, setting)
 - subject (main actor)
 - objects (main product or items)
@@ -25,13 +47,16 @@ Your job is to take the user's prompt and return a structured JSON with:
 Respond ONLY in raw JSON format. Do NOT include explanations.
 """
 
-def extract_scene_plan(prompt: str) -> dict:
+def extract_scene_plan(prompt: str, image: Image.Image) -> dict:
     try:
+        caption = generate_blip_caption(image)
+        merged_prompt = f"Image Caption: {caption}\nUser Prompt: {prompt}"
+
         response = client.chat.completions.create(
             model="gpt-4o-mini-2024-07-18",
             messages=[
                 {"role": "system", "content": SCENE_SYSTEM_INSTRUCTIONS},
-                {"role": "user", "content": prompt}
+                {"role": "user", "content": merged_prompt}
             ],
             temperature=0.3,
             max_tokens=500
@@ -51,7 +76,7 @@ def extract_scene_plan(prompt: str) -> dict:
         }
 
 # ----------------------------
-# 🧠 Prompt Variation Generator
+# 🧠 Prompt Variation Generator (GPT-4o)
 # ----------------------------
 def generate_prompt_variations_from_scene(scene_plan: dict, base_prompt: str, n: int = 3) -> list:
     try: