Update utils/planner.py

utils/planner.py

@@ -24,7 +24,7 @@ processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base
 blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
 
 # ----------------------------
-# 🧠 Load CLIP Tokenizer (for optional diagnostics)
+# 🧠 Load CLIP Tokenizer (for token check)
 # ----------------------------
 tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
 
@@ -36,6 +36,8 @@ def generate_blip_caption(image: Image.Image) -> str:
         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)
+        # Clean duplicate tokens
+        caption = " ".join(dict.fromkeys(caption.split()))
         print(f"🖼️ BLIP Caption: {caption}")
         return caption
     except Exception as e:
@@ -43,11 +45,11 @@ def generate_blip_caption(image: Image.Image) -> str:
         return "a product image"
 
 # ----------------------------
-# 🧠 GPT Scene Planning
+# 🧠 GPT Scene Planning with Caption + Visual Style
 # ----------------------------
 SCENE_SYSTEM_INSTRUCTIONS = """
 You are a scene planning assistant for an AI image generation system.
-Your job is to take a caption from a product image and a user prompt, then return a structured JSON with:
+Your job is to take a caption from a product image, a visual style hint, and a user prompt, then return a structured JSON with:
 - scene (environment, setting)
 - subject (main_actor)
 - objects (main_product or items)
@@ -59,7 +61,13 @@ Respond ONLY in raw JSON format. Do NOT include explanations.
 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}"
+        visual_hint = caption if "shoe" in caption or "product" in caption else "low-top product photo on white background"
+        
+        merged_prompt = (
+            f"Image Caption: {caption}\n"
+            f"Image Visual Style: {visual_hint}\n"
+            f"User Prompt: {prompt}"
+        )
 
         response = client.chat.completions.create(
             model="gpt-4o-mini-2024-07-18",
@@ -73,10 +81,15 @@ def extract_scene_plan(prompt: str, image: Image.Image) -> dict:
         content = response.choices[0].message.content
         print("🧠 Scene Plan (Raw):", content)
 
-        # Optional logging
+        # Logging
         os.makedirs("logs", exist_ok=True)
         with open("logs/scene_plans.jsonl", "a") as f:
-            f.write(json.dumps({"caption": caption, "prompt": prompt, "scene_plan": content}) + "\n")
+            f.write(json.dumps({
+                "caption": caption,
+                "visual_hint": visual_hint,
+                "prompt": prompt,
+                "scene_plan": content
+            }) + "\n")
 
         return json.loads(content)
 
@@ -91,48 +104,42 @@ def extract_scene_plan(prompt: str, image: Image.Image) -> dict:
         }
 
 # ----------------------------
-# ✨ GPT-Powered Prompt Variations (77-tokens safe)
+# ✨ Enriched Prompt Generation (GPT, 77-token safe)
 # ----------------------------
 ENRICHED_PROMPT_INSTRUCTIONS = """
 You are a prompt engineer for an AI image generation model.
-Given a structured scene plan and product prompt, generate a visually descriptive enriched prompt that:
-
+Given a structured scene plan and a user prompt, generate a single natural-language enriched prompt that:
 1. Describes the subject, product, setting, and layout clearly
-2. Stays strictly under 77 tokens (CLIP limit for SDXL)
-3. Is natural, realistic, and suitable for Stable Diffusion XL
-4. Does NOT include quotes, explanations, or bullet points — just the enriched prompt
-
-Return only the prompt as a string.
+2. Uses natural, photo-realistic language
+3. Stays strictly under 77 tokens (CLIP token limit)
+Return ONLY the enriched prompt string. No explanations.
 """
 
 def generate_prompt_variations_from_scene(scene_plan: dict, base_prompt: str, n: int = 3) -> list:
     prompts = []
     for _ in range(n):
         try:
-            user_input = f"Scene Plan:\n{json.dumps(scene_plan)}\n\nOriginal User Prompt:\n{base_prompt}"
+            user_input = f"Scene Plan:\n{json.dumps(scene_plan)}\n\nUser Prompt:\n{base_prompt}"
             response = client.chat.completions.create(
                 model="gpt-4o-mini-2024-07-18",
                 messages=[
                     {"role": "system", "content": ENRICHED_PROMPT_INSTRUCTIONS},
                     {"role": "user", "content": user_input}
                 ],
-                temperature=0.5,
+                temperature=0.4,
                 max_tokens=100
             )
             enriched = response.choices[0].message.content.strip()
-
-            # Optional: check token count for debug
             token_count = len(tokenizer(enriched)["input_ids"])
             print(f"📝 Enriched Prompt ({token_count} tokens): {enriched}")
-
             prompts.append(enriched)
         except Exception as e:
-            print("⚠️ Prompt variation fallback:", e)
+            print("⚠️ Prompt fallback:", e)
             prompts.append(base_prompt)
     return prompts
 
 # ----------------------------
-# ❌ Generate Negative Prompt
+# ❌ Negative Prompt Generator
 # ----------------------------
 NEGATIVE_SYSTEM_PROMPT = """
 You are a prompt engineer. Given a structured scene plan, generate a short negative prompt
@@ -152,9 +159,9 @@ def generate_negative_prompt_from_scene(scene_plan: dict) -> str:
             temperature=0.2,
             max_tokens=100
         )
-        negative = response.choices[0].message.content.strip()
-        return negative
+        return response.choices[0].message.content.strip()
     except Exception as e:
         print("❌ Negative Prompt Error:", e)
         return "blurry, distorted, low quality, deformed, watermark"
 
+