Update utils.py

utils.py

@@ -98,12 +98,13 @@ def clean_memory() -> None:
         logger.warning(f"Memory cleanup failed: {e}")
 
 
-def apply_flux_rules(prompt: str) -> str:
+def apply_flux_rules(prompt: str, analysis_metadata: Optional[Dict[str, Any]] = None) -> str:
     """
     Apply Flux optimization rules to a prompt
     
     Args:
         prompt: Raw prompt text
+        analysis_metadata: Optional metadata from image analysis including camera suggestions
         
     Returns:
         Optimized prompt following Flux rules
@@ -116,30 +117,26 @@ def apply_flux_rules(prompt: str) -> str:
     for pattern in FLUX_RULES["remove_patterns"]:
         cleaned_prompt = re.sub(pattern, '', cleaned_prompt, flags=re.IGNORECASE)
     
-    # Detect image type and add appropriate camera configuration
-    prompt_lower = cleaned_prompt.lower()
-    camera_config = ""
+    # Extract description part only (remove CAMERA_SETUP section if present)
+    description_part = _extract_description_only(cleaned_prompt)
     
-    if any(word in prompt_lower for word in ['portrait', 'person', 'man', 'woman', 'face']):
-        camera_config = FLUX_RULES["camera_configs"]["portrait"]
-    elif any(word in prompt_lower for word in ['landscape', 'mountain', 'nature', 'outdoor']):
-        camera_config = FLUX_RULES["camera_configs"]["landscape"]
-    elif any(word in prompt_lower for word in ['street', 'urban', 'city']):
-        camera_config = FLUX_RULES["camera_configs"]["street"]
+    # Check if BAGEL provided intelligent camera setup
+    camera_config = ""
+    if analysis_metadata and analysis_metadata.get("has_camera_suggestion") and analysis_metadata.get("camera_setup"):
+        # Use BAGEL's intelligent camera suggestion
+        bagel_camera = analysis_metadata["camera_setup"]
+        camera_config = f", {bagel_camera}" if not bagel_camera.startswith(",") else bagel_camera
+        logger.info(f"Using BAGEL camera suggestion: {bagel_camera}")
     else:
-        camera_config = FLUX_RULES["camera_configs"]["default"]
+        # Fallback to static rules if BAGEL didn't suggest camera
+        camera_config = _get_fallback_camera_config(description_part.lower())
+        logger.info("Using fallback camera configuration")
     
     # Add lighting enhancements if not present
-    if 'lighting' not in prompt_lower:
-        if 'dramatic' in prompt_lower:
-            cleaned_prompt += FLUX_RULES["lighting_enhancements"]["dramatic"]
-        elif 'portrait' in prompt_lower:
-            cleaned_prompt += FLUX_RULES["lighting_enhancements"]["portrait"]
-        else:
-            cleaned_prompt += FLUX_RULES["lighting_enhancements"]["default"]
+    lighting_enhancement = _get_lighting_enhancement(description_part.lower())
     
-    # Build final prompt
-    final_prompt = cleaned_prompt + camera_config
+    # Build final prompt: Description + Camera + Lighting
+    final_prompt = description_part + camera_config + lighting_enhancement
     
     # Clean up formatting
     final_prompt = _clean_prompt_formatting(final_prompt)
@@ -147,6 +144,51 @@ def apply_flux_rules(prompt: str) -> str:
     return final_prompt
 
 
+def _extract_description_only(prompt: str) -> str:
+    """Extract only the description part, removing camera setup sections"""
+    # Remove CAMERA_SETUP section if present
+    if "CAMERA_SETUP:" in prompt:
+        parts = prompt.split("CAMERA_SETUP:")
+        description = parts[0].strip()
+    elif "2. CAMERA_SETUP" in prompt:
+        parts = prompt.split("2. CAMERA_SETUP")
+        description = parts[0].strip()
+    else:
+        description = prompt
+    
+    # Remove "DESCRIPTION:" label if present
+    if description.startswith("DESCRIPTION:"):
+        description = description.replace("DESCRIPTION:", "").strip()
+    elif description.startswith("1. DESCRIPTION:"):
+        description = description.replace("1. DESCRIPTION:", "").strip()
+    
+    return description.strip()
+
+
+def _get_fallback_camera_config(prompt_lower: str) -> str:
+    """Get fallback camera configuration when BAGEL doesn't suggest one"""
+    if any(word in prompt_lower for word in ['portrait', 'person', 'man', 'woman', 'face']):
+        return FLUX_RULES["camera_configs"]["portrait"]
+    elif any(word in prompt_lower for word in ['landscape', 'mountain', 'nature', 'outdoor']):
+        return FLUX_RULES["camera_configs"]["landscape"]
+    elif any(word in prompt_lower for word in ['street', 'urban', 'city']):
+        return FLUX_RULES["camera_configs"]["street"]
+    else:
+        return FLUX_RULES["camera_configs"]["default"]
+
+
+def _get_lighting_enhancement(prompt_lower: str) -> str:
+    """Determine appropriate lighting enhancement"""
+    if 'lighting' not in prompt_lower:
+        if 'dramatic' in prompt_lower:
+            return FLUX_RULES["lighting_enhancements"]["dramatic"]
+        elif 'portrait' in prompt_lower:
+            return FLUX_RULES["lighting_enhancements"]["portrait"]
+        else:
+            return FLUX_RULES["lighting_enhancements"]["default"]
+    return ""  # No enhancement needed if lighting already mentioned
+
+
 def _clean_prompt_formatting(prompt: str) -> str:
     """Clean up prompt formatting"""
     if not prompt:
@@ -187,9 +229,17 @@ def calculate_prompt_score(prompt: str, analysis_data: Optional[Dict[str, Any]]
     detail_score = min(10, len(prompt.split(',')) * 2)  # Reward detail
     breakdown["prompt_quality"] = length_score + detail_score
     
-    # Technical details score (0-25 points)
+    # Technical details score (0-25 points) - Enhanced for BAGEL camera suggestions
+    tech_score = 0
     tech_keywords = ['shot on', 'lens', 'photography', 'lighting', 'camera']
-    tech_score = sum(5 for keyword in tech_keywords if keyword in prompt.lower())
+    for keyword in tech_keywords:
+        if keyword in prompt.lower():
+            tech_score += 5
+    
+    # Bonus points for BAGEL camera suggestions
+    if analysis_data and analysis_data.get("has_camera_suggestion"):
+        tech_score += 5  # Bonus for intelligent camera selection
+        
     breakdown["technical_details"] = min(25, tech_score)
     
     # Artistic value score (0-25 points)
@@ -199,10 +249,16 @@ def calculate_prompt_score(prompt: str, analysis_data: Optional[Dict[str, Any]]
     
     # Flux optimization score (0-20 points)
     flux_score = 0
-    if any(camera in prompt for camera in FLUX_RULES["camera_configs"].values()):
+    
+    # Check for camera configuration (either BAGEL or fallback)
+    if any(camera in prompt for camera in FLUX_RULES["camera_configs"].values()) or \
+       (analysis_data and analysis_data.get("has_camera_suggestion")):
         flux_score += 10
+        
+    # Check for lighting configuration  
     if any(lighting in prompt for lighting in FLUX_RULES["lighting_enhancements"].values()):
         flux_score += 10
+        
     breakdown["flux_optimization"] = flux_score
     
     # Calculate total