prompt_optimizer.py

"""
Intelligent Prompt Optimization System
Enhances user prompts for better image generation accuracy
"""

import re
from typing import Dict, Tuple
from config import STYLE_PRESETS, BASE_NEGATIVE_PROMPT


class PromptOptimizer:
    """
    Analyzes and enhances user prompts to improve generation quality
    """
    
    def __init__(self):
        self.quality_tokens = [
            "highly detailed",
            "sharp focus",
            "professional",
            "4k"
        ]
        
        self.lighting_terms = {
            'photo': "soft cinematic lighting",
            'art': "dramatic lighting",
            'anime': "vibrant lighting",
            'default': "good lighting"
        }
    
    def detect_intent(self, prompt: str) -> str:
        """
        Detect the user's intended style from their prompt
        
        Args:
            prompt: User's input prompt
            
        Returns:
            Detected style category ('photo', 'art', 'anime', etc.)
        """
        prompt_lower = prompt.lower()
        
        # Check for explicit style keywords
        if any(word in prompt_lower for word in ['photograph', 'photo', 'camera', 'lens', 'dslr']):
            return 'photo'
        elif any(word in prompt_lower for word in ['anime', 'manga', 'kawaii']):
            return 'anime'
        elif any(word in prompt_lower for word in ['painting', 'art', 'illustration', 'drawing']):
            return 'art'
        elif any(word in prompt_lower for word in ['3d', 'render', 'cgi']):
            return '3d'
        
        return 'default'
    
    def clean_prompt(self, prompt: str) -> str:
        """
        Clean and normalize the user's prompt
        
        Args:
            prompt: Raw user input
            
        Returns:
            Cleaned prompt
        """
        # Remove excessive whitespace
        prompt = re.sub(r'\s+', ' ', prompt)
        
        # Remove special characters that might confuse the model
        prompt = re.sub(r'[^\w\s,.-]', '', prompt)
        
        # Capitalize first letter
        prompt = prompt.strip()
        if prompt:
            prompt = prompt[0].upper() + prompt[1:]
        
        return prompt
    
    def enhance_prompt(
        self, 
        prompt: str, 
        style: str = "None",
        add_quality_tokens: bool = True
    ) -> Tuple[str, str]:
        """
        Enhance user prompt with style-specific improvements
        
        Args:
            prompt: User's input prompt
            style: Selected style preset
            add_quality_tokens: Whether to add quality enhancement tokens
            
        Returns:
            Tuple of (enhanced_prompt, negative_prompt)
        """
        # Clean the input
        prompt = self.clean_prompt(prompt)
        
        if not prompt:
            prompt = "a beautiful scene"
        
        # Get style configuration
        style_config = STYLE_PRESETS.get(style, STYLE_PRESETS["None"])
        
        # Build enhanced prompt
        enhanced_parts = [prompt]
        
        # Add style-specific suffix
        if style_config["prompt_suffix"]:
            enhanced_parts.append(style_config["prompt_suffix"])
        
        # Add quality tokens if enabled and not already present
        if add_quality_tokens and style == "None":
            # Only add generic quality tokens if no specific style is selected
            quality_str = ", ".join(self.quality_tokens)
            enhanced_parts.append(quality_str)
        
        enhanced_prompt = ", ".join(enhanced_parts)
        
        # Build negative prompt
        negative_prompt = style_config.get("negative_prompt", BASE_NEGATIVE_PROMPT)
        
        return enhanced_prompt, negative_prompt
    
    def structure_prompt(self, prompt: str) -> str:
        """
        Convert free-form prompt into structured format
        [Subject] [Action] [Environment] [Style] [Lighting] [Quality]
        
        Args:
            prompt: User's input prompt
            
        Returns:
            Structured prompt
        """
        # This is a simple implementation
        # For production, you might use NLP to extract components
        
        parts = prompt.split(',')
        if len(parts) == 1:
            # Single phrase - assume it's the subject
            return prompt
        
        # Already somewhat structured, return as-is
        return prompt
    
    def get_prompt_info(self, prompt: str, style: str) -> Dict[str, str]:
        """
        Get detailed information about how the prompt will be processed
        
        Args:
            prompt: User's input prompt
            style: Selected style preset
            
        Returns:
            Dictionary with prompt processing details
        """
        enhanced, negative = self.enhance_prompt(prompt, style)
        intent = self.detect_intent(prompt)
        
        return {
            "original": prompt,
            "enhanced": enhanced,
            "negative": negative,
            "detected_intent": intent,
            "style": style
        }


# Test function
if __name__ == "__main__":
    optimizer = PromptOptimizer()
    
    # Test prompts
    test_prompts = [
        "a cat sitting on a windowsill",
        "photorealistic portrait of a warrior",
        "anime girl with blue hair",
        "3d render of a futuristic city"
    ]
    
    print("=== Prompt Optimizer Test ===\n")
    
    for prompt in test_prompts:
        print(f"Original: {prompt}")
        enhanced, negative = optimizer.enhance_prompt(prompt, "Photorealistic")
        print(f"Enhanced: {enhanced}")
        print(f"Negative: {negative}")
        print(f"Detected Intent: {optimizer.detect_intent(prompt)}")
        print("-" * 80)
        print()