prompt_enhancer.py

from typing import Dict, Any, List
import re

class PromptEnhancer:
    def __init__(self):
        self.quality_terms = {
            "Ultra Réaliste": [
                "masterpiece", "best quality", "ultra realistic",
                "photorealistic", "8k uhd", "high resolution",
                "detailed", "sharp focus", "professional photography"
            ],
            "Artistique Pro": [
                "masterpiece", "best quality", "professional",
                "detailed", "artistic", "perfect composition",
                "award winning", "trending on artstation"
            ]
        }

    def enhance(self, prompt: str, style: str, composition: str, mood: str) -> str:
        """Améliore le prompt en ajoutant des termes de qualité et de style appropriés"""
        # Nettoyage et normalisation du prompt
        cleaned_prompt = self._clean_prompt(prompt)
        
        # Ajout des termes de qualité spécifiques au style
        quality_terms = self.quality_terms.get(style, self.quality_terms["Ultra Réaliste"])
        quality_string = ", ".join(quality_terms)
        
        # Construction du prompt final
        enhanced_prompt = f"{cleaned_prompt}, {quality_string}, {composition}, {mood}"
        
        # Optimisation finale
        return self._optimize_prompt(enhanced_prompt)

    def _clean_prompt(self, prompt: str) -> str:
        """Nettoie et normalise le prompt"""
        # Suppression des espaces multiples
        cleaned = re.sub(r'\s+', ' ', prompt.strip())
        
        # Suppression des termes de qualité redondants
        redundant_terms = ["high quality", "good quality", "best quality", "hq"]
        for term in redundant_terms:
            cleaned = re.sub(rf'\b{term}\b', '', cleaned, flags=re.IGNORECASE)
        
        return cleaned.strip()

    def _optimize_prompt(self, prompt: str) -> str:
        """Optimisation finale du prompt"""
        # Limitation de la longueur
        words = prompt.split()
        if len(words) > 77:  # SDXL peut gérer jusqu'à 77 tokens
            words = words[:77]
        
        # Réorganisation pour mettre les termes importants en premier
        important_terms = []
        regular_terms = []
        
        for word in words:
            if word.lower() in ["masterpiece", "best quality", "professional"]:
                important_terms.append(word)
            else:
                regular_terms.append(word)
        
        return ", ".join(important_terms + regular_terms)