Update prompt_enhancer.py

prompt_enhancer.py

@@ -1,182 +1,63 @@
-from typing import Dict, List, Optional
-import logging
+from typing import Dict, Any, List
 import re
 
-logger = logging.getLogger(__name__)
-
 class PromptEnhancer:
     def __init__(self):
-        self.context_keywords = {
-            # Éléments de design
-            "moderne": "modern clean professional design",
-            "vintage": "vintage retro classic design",
-            "minimaliste": "minimalist clean simple design",
-            "luxe": "luxury elegant premium design",
-            
-            # Types d'ambiance
-            "professionnel": "professional corporate business-like",
-            "créatif": "creative artistic innovative",
-            "dynamique": "dynamic energetic vibrant",
-            "élégant": "elegant sophisticated refined",
-            
-            # Éléments visuels
-            "logo": "prominent logo design professional branding",
-            "texte": "clear readable text typography",
-            "image": "main visual focal point image",
-            "photo": "photographic image realistic",
-            
-            # Caractéristiques techniques
-            "haute qualité": "high quality professional grade",
-            "détaillé": "highly detailed intricate",
-            "net": "sharp crisp clear",
-            "flou": "soft focus gentle blur",
-            
-            # Styles spécifiques
-            "3D": "three dimensional depth realistic",
-            "plat": "flat design 2D clean",
-            "graphique": "graphic design vector-style",
-            "illustré": "illustrated hand-drawn artistic"
-        }
-        
-        self.composition_patterns = {
-            # Structure de l'affiche
-            "haut": "top aligned composition with {element}",
-            "bas": "bottom aligned composition with {element}",
-            "centre": "centered composition with {element}",
-            "gauche": "left aligned composition with {element}",
-            "droite": "right aligned composition with {element}",
-            
-            # Relations spatiales
-            "au-dessus": "{element1} positioned above {element2}",
-            "en-dessous": "{element1} positioned below {element2}",
-            "à côté": "{element1} next to {element2}",
-            "autour": "{element1} surrounding {element2}"
-        }
-        
-        self.emphasis_patterns = {
-            "important": "(({})),",  # Double emphase
-            "normal": "({}),",       # Emphase simple
-            "subtil": "[{}],"        # Emphase légère
-        }
-        
-        self.common_improvements = {
-            "faire": "create",
-            "mettre": "place",
-            "avec": "featuring",
-            "contenant": "containing",
-            "il y a": "featuring",
-            "je veux": "",
-            "je souhaite": "",
-            "il faut": ""
+        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_prompt(self, user_input: str, style_context: Dict) -> str:
-        """Améliore le prompt utilisateur avec un contexte enrichi"""
-        enhanced_prompt = user_input.lower()
-        logger.debug(f"Prompt initial: {enhanced_prompt}")
-
-        # Nettoyage initial
-        enhanced_prompt = self._clean_prompt(enhanced_prompt)
-        
-        # Détection et amélioration du contexte
-        enhanced_prompt = self._add_context(enhanced_prompt)
+    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 éléments de style
-        enhanced_prompt = self._add_style_elements(enhanced_prompt, style_context)
+        # 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)
         
-        # Structure et emphase
-        enhanced_prompt = self._structure_prompt(enhanced_prompt)
+        # Construction du prompt final
+        enhanced_prompt = f"{cleaned_prompt}, {quality_string}, {composition}, {mood}"
         
-        # Ajout des qualificatifs techniques
-        enhanced_prompt = self._add_technical_qualifiers(enhanced_prompt)
-        
-        logger.debug(f"Prompt amélioré: {enhanced_prompt}")
-        return enhanced_prompt
+        # Optimisation finale
+        return self._optimize_prompt(enhanced_prompt)
 
     def _clean_prompt(self, prompt: str) -> str:
         """Nettoie et normalise le prompt"""
-        # Remplace les expressions communes par leurs versions optimisées
-        for old, new in self.common_improvements.items():
-            prompt = prompt.replace(old, new)
+        # Suppression des espaces multiples
+        cleaned = re.sub(r'\s+', ' ', prompt.strip())
         
-        # Supprime les espaces multiples
-        prompt = " ".join(prompt.split())
+        # 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 prompt
+        return cleaned.strip()
 
-    def _add_context(self, prompt: str) -> str:
-        """Ajoute du contexte basé sur les mots-clés détectés"""
-        enhanced_parts = []
+    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]
         
-        for word in words:
-            if word in self.context_keywords:
-                enhanced_parts.append(self.context_keywords[word])
-            else:
-                enhanced_parts.append(word)
+        # Réorganisation pour mettre les termes importants en premier
+        important_terms = []
+        regular_terms = []
         
-        return " ".join(enhanced_parts)
-
-    def _add_style_elements(self, prompt: str, style_context: Dict) -> str:
-        """Ajoute les éléments de style au prompt"""
-        style_elements = [
-            style_context.get("prompt_prefix", ""),
-            prompt,
-            style_context.get("layout", ""),
-            style_context.get("ambiance", ""),
-            style_context.get("palette", ""),
-            "professional poster design",
-            "high quality"
-        ]
-        
-        return ", ".join(filter(None, style_elements))
-
-    def _structure_prompt(self, prompt: str) -> str:
-        """Structure le prompt avec une emphase appropriée"""
-        # Identifie les éléments clés
-        main_elements = self._identify_main_elements(prompt)
-        
-        structured_parts = []
-        for element, importance in main_elements.items():
-            pattern = self.emphasis_patterns.get(importance, "{}") 
-            structured_parts.append(pattern.format(element))
-        
-        return " ".join(structured_parts)
-
-    def _identify_main_elements(self, prompt: str) -> Dict[str, str]:
-        """Identifie les éléments principaux et leur importance"""
-        elements = {}
-        
-        # Analyse basique des éléments clés
-        words = prompt.split()
         for word in words:
-            if len(word) > 3:  # Ignore les mots très courts
-                if word in self.context_keywords:
-                    elements[word] = "important"
-                else:
-                    elements[word] = "normal"
-                    
-        return elements
-
-    def _add_technical_qualifiers(self, prompt: str) -> str:
-        """Ajoute des qualificatifs techniques pour améliorer la qualité"""
-        technical_qualifiers = [
-            "professional quality",
-            "highly detailed",
-            "masterful composition",
-            "perfect lighting",
-            "sharp focus",
-            "8k resolution"
-        ]
+            if word.lower() in ["masterpiece", "best quality", "professional"]:
+                important_terms.append(word)
+            else:
+                regular_terms.append(word)
         
-        return f"{prompt}, {', '.join(technical_qualifiers)}"
-
-    def analyze_prompt_effectiveness(self, prompt: str) -> Dict:
-        """Analyse l'efficacité du prompt"""
-        return {
-            "length": len(prompt),
-            "key_elements": len(self._identify_main_elements(prompt)),
-            "has_context": any(keyword in prompt for keyword in self.context_keywords),
-            "has_composition": any(pattern in prompt for pattern in self.composition_patterns),
-            "technical_quality": len([q for q in ["detailed", "quality", "professional"] if q in prompt])
-        }
+        return ", ".join(important_terms + regular_terms)