Update app.py

app.py

@@ -1,403 +1,185 @@
 import gradio as gr
-import os
+from typing import Dict, Any, Tuple, Optional
 from PIL import Image
-import requests
-import io
-import gc
-import json
-from typing import Tuple, Optional, Dict, Any
+import torch
+from diffusers import StableDiffusionXLPipeline, DPMSolverMultistepScheduler
 import logging
-from dotenv import load_dotenv
+from dataclasses import dataclass
 from prompt_enhancer import PromptEnhancer
 
-# Configuration du logging
-logging.basicConfig(
-    level=logging.DEBUG,
-    format='%(asctime)s - %(levelname)s - %(message)s'
-)
-logger = logging.getLogger(__name__)
+@dataclass
+class GenerationConfig:
+    width: int = 1024
+    height: int = 1024
+    num_inference_steps: int = 50
+    guidance_scale: float = 7.5
+    high_noise_frac: float = 0.8
+    negative_prompt: str = ""
 
-# Chargement des variables d'environnement
-load_dotenv()
-
-# Styles artistiques complets
-ART_STYLES = {
-    # Styles Classiques
-    "Ultra Réaliste": {
-        "prompt_prefix": "ultra realistic photograph, stunning photorealistic quality, unreal engine 5 quality, octane render, ray tracing, volumetric lighting, subsurface scattering, 8k UHD, cinema quality, masterpiece, perfect composition, award winning photography",
-        "negative_prompt": "artificial, digital art, illustration, painting, drawing, artistic, cartoon, anime, unreal, fake, low quality, blurry, soft, deformed"
-    },
-    "Photoréaliste": {
-        "prompt_prefix": "hyperrealistic photograph, extremely detailed, studio quality, professional photography, 8k uhd",
-        "negative_prompt": "artistic, painterly, abstract, cartoon, illustration, low quality"
-    },
-    "Expressionniste": {
-        "prompt_prefix": "expressive painting style, intense emotional art, bold brushstrokes, vibrant colors, van gogh inspired",
-        "negative_prompt": "realistic, subtle, photographic, clean lines, digital art"
-    },
-    "Impressionniste": {
-        "prompt_prefix": "impressionist painting style, soft light, visible brushstrokes, outdoor scene, monet inspired",
-        "negative_prompt": "sharp details, high contrast, digital, modern"
-    },
-    "Art Abstrait": {
-        "prompt_prefix": "abstract art, geometric shapes, non-representational, kandinsky style, pure artistic expression",
-        "negative_prompt": "realistic, figurative, photographic, literal"
-    },
-    "Surréaliste": {
-        "prompt_prefix": "surrealist art, dreamlike imagery, symbolic elements, dali inspired, metaphysical art",
-        "negative_prompt": "realistic, conventional, ordinary, literal"
-    },
-    "Art Moderne": {
-        "prompt_prefix": "modern art style poster, professional design, contemporary aesthetic",
-        "negative_prompt": "traditional, cluttered, busy design, vintage"
-    },
-    "Pop Art": {
-        "prompt_prefix": "pop art style poster, bold colors, repeated patterns, screen print effect, warhol inspired",
-        "negative_prompt": "subtle, realistic, traditional, painterly"
-    },
-    "Minimaliste": {
-        "prompt_prefix": "minimalist design poster, clean composition, elegant simplicity",
-        "negative_prompt": "complex, detailed, ornate, busy, cluttered"
-    },
-    "Cubiste": {
-        "prompt_prefix": "cubist art style, geometric fragmentation, multiple perspectives, picasso inspired",
-        "negative_prompt": "realistic, single perspective, traditional, photographic"
-    },
-    "Futuriste": {
-        "prompt_prefix": "futuristic art style, dynamic movement, technological elements, speed and motion",
-        "negative_prompt": "static, traditional, classical, historical"
-    },
-    "Neo Vintage": {
-        "prompt_prefix": "vintage style advertising poster, retro design, classic aesthetic",
-        "negative_prompt": "modern, digital, contemporary style"
-    },
-    "Cyberpunk": {
-        "prompt_prefix": "cyberpunk style poster, neon lights, futuristic design, high-tech aesthetic",
-        "negative_prompt": "vintage, natural, rustic, traditional"
-    },
-    "Japonais": {
-        "prompt_prefix": "japanese art style poster, ukiyo-e inspired design, traditional japanese aesthetic",
-        "negative_prompt": "western, modern, photographic"
-    },
-    "Art Déco": {
-        "prompt_prefix": "art deco style poster, geometric patterns, luxury design, 1920s aesthetic",
-        "negative_prompt": "modern, minimalist, casual, contemporary"
-    },
-    "Symboliste": {
-        "prompt_prefix": "symbolic art, decorative patterns, gold elements, mystical atmosphere, klimt inspired",
-        "negative_prompt": "realistic, simple, plain, literal"
-    }
-}
-
-# Paramètres de composition enrichis
-COMPOSITION_PARAMS = {
-    "Layouts": {
-        "Centré": "centered composition, balanced layout, harmonious arrangement",
-        "Asymétrique": "dynamic asymmetrical composition, creative balance",
-        "Grille": "grid-based layout, structured composition, organized design",
-        "Diagonal": "diagonal dynamic composition, energetic flow",
-        "Minimaliste": "minimal composition, lots of whitespace, elegant spacing"
-    },
-    "Ambiances": {
-        "Dramatique": "dramatic lighting, high contrast, intense mood",
-        "Doux": "soft lighting, gentle atmosphere, subtle mood",
-        "Vibrant": "vibrant colors, energetic mood, dynamic atmosphere",
-        "Mystérieux": "mysterious atmosphere, moody lighting, enigmatic feel",
-        "Serein": "peaceful atmosphere, calm mood, tranquil setting"
-    },
-    "Palette": {
-        "Monochrome": "monochromatic color scheme, sophisticated tones",
-        "Contrasté": "high contrast color palette, bold color combinations",
-        "Pastel": "soft pastel color palette, gentle colors",
-        "Terre": "earthy color palette, natural tones",
-        "Néon": "neon color palette, vibrant glowing colors"
-    }
-}
-
-class ImageGenerator:
+class EnhancedImageGenerator:
     def __init__(self):
-        self.API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
-        token = os.getenv('HUGGINGFACE_TOKEN')
-        if not token:
-            logger.error("HUGGINGFACE_TOKEN non trouvé!")
-        self.headers = {"Authorization": f"Bearer {token}"}
-        logger.info("ImageGenerator initialisé")
-
-    def _build_prompt(self, params: Dict[str, Any]) -> str:
-        """Construction de prompt améliorée avec le PromptEnhancer"""
-        try:
-            # Initialisation du PromptEnhancer
-            enhancer = PromptEnhancer()
-
-            # Récupération du style
-            style_info = ART_STYLES.get(params["style"], ART_STYLES["Art Moderne"])
-
-            # Construction du contexte de style
-            style_context = {
-                "prompt_prefix": style_info['prompt_prefix'],
-                "layout": COMPOSITION_PARAMS['Layouts'][params['layout']],
-                "ambiance": COMPOSITION_PARAMS['Ambiances'][params['ambiance']],
-                "palette": COMPOSITION_PARAMS['Palette'][params['palette']]
-            }
-
-            # Préparation du prompt initial
-            base_prompt = f"{params['subject']}"
-            if params.get('title'):
-                base_prompt += f", with text '{params['title']}'"
-
-            # Amélioration du prompt
-            enhanced_prompt = enhancer.enhance_prompt(base_prompt, style_context)
-
-            # Analyse de l'efficacité du prompt
-            prompt_analysis = enhancer.analyze_prompt_effectiveness(enhanced_prompt)
-            logger.debug(f"Analyse du prompt: {json.dumps(prompt_analysis, indent=2)}")
-
-            return enhanced_prompt
-
-        except Exception as e:
-            logger.error(f"Erreur dans la construction du prompt: {str(e)}")
-            # Fallback sur le prompt basique en cas d'erreur
-            return f"{style_info['prompt_prefix']}, {params['subject']}"
-
-    def generate(self, params: Dict[str, Any]) -> Tuple[Optional[Image.Image], str]:
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.pipeline = self._initialize_pipeline()
+        self.prompt_enhancer = PromptEnhancer()
+        
+    def _initialize_pipeline(self) -> StableDiffusionXLPipeline:
+        pipeline = StableDiffusionXLPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-xl-base-1.0",
+            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
+            use_safetensors=True,
+            variant="fp16" if self.device == "cuda" else None
+        )
+        
+        # Optimisations cruciales pour la qualité
+        pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
+            pipeline.scheduler.config,
+            algorithm_type="dpmsolver++",
+            use_karras_sigmas=True
+        )
+        
+        if self.device == "cuda":
+            pipeline.enable_xformers_memory_efficient_attention()
+            pipeline.enable_model_cpu_offload()
+        
+        return pipeline
+
+    def _enhance_prompt(self, base_prompt: str, style_params: Dict[str, Any]) -> str:
+        """Amélioration avancée des prompts avec analyse contextuelle"""
+        enhanced = self.prompt_enhancer.enhance(
+            base_prompt,
+            style=style_params["style"],
+            composition=style_params["composition"],
+            mood=style_params["mood"]
+        )
+        
+        # Ajout d'optimisations de qualité spécifiques
+        quality_terms = [
+            "masterpiece",
+            "highly detailed",
+            "professional",
+            "award winning",
+            "stunning",
+            f"resolution {style_params.get('resolution', '8k')}",
+            "perfect composition"
+        ]
+        
+        return f"{enhanced}, {', '.join(quality_terms)}"
+
+    def generate(
+        self,
+        params: Dict[str, Any],
+        config: GenerationConfig
+    ) -> Tuple[Optional[Image.Image], str]:
         try:
-            logger.info(f"Début de génération avec paramètres: {json.dumps(params, indent=2)}")
-
-            if 'Bearer None' in self.headers['Authorization']:
-                return None, "⚠️ Erreur: Token Hugging Face non configuré"
-
-            prompt = self._build_prompt(params)
-
-            payload = {
-                "inputs": prompt,
-                "parameters": {
-                    "negative_prompt": ART_STYLES[params["style"]]["negative_prompt"],
-                    "num_inference_steps": min(int(35 * (params["quality"]/100)), 40),
-                    "guidance_scale": min(7.5 * (params["creativity"]/10), 10.0),
-                    "width": 768,
-                    "height": 768 if params["orientation"] == "Portrait" else 512
-                }
-            }
-
-            logger.debug(f"Payload: {json.dumps(payload, indent=2)}")
-
-            response = requests.post(
-                self.API_URL,
-                headers=self.headers,
-                json=payload,
-                timeout=30
-            )
-
-            if response.status_code == 200:
-                image = Image.open(io.BytesIO(response.content))
-                return image, "✨ Création réussie!"
-            else:
-                error_msg = f"⚠️ Erreur API {response.status_code}: {response.text}"
-                logger.error(error_msg)
-                return None, error_msg
-
+            # Optimisation du prompt
+            enhanced_prompt = self._enhance_prompt(params["prompt"], params)
+            
+            # Configuration avancée de la génération
+            with torch.inference_mode():
+                image = self.pipeline(
+                    prompt=enhanced_prompt,
+                    negative_prompt=config.negative_prompt,
+                    width=config.width,
+                    height=config.height,
+                    num_inference_steps=config.num_inference_steps,
+                    guidance_scale=config.guidance_scale,
+                    high_noise_frac=config.high_noise_frac,
+                ).images[0]
+
+            return image, "Génération réussie!"
         except Exception as e:
-            error_msg = f"⚠️ Erreur: {str(e)}"
-            logger.exception("Erreur pendant la génération:")
-            return None, error_msg
-        finally:
-            gc.collect()
-
-def create_interface():
-    logger.info("Création de l'interface Gradio")
-
-    css = """
-    .container { max-width: 1200px; margin: auto; }
-    .welcome { text-align: center; margin: 20px 0; padding: 20px; background: #1e293b; border-radius: 10px; }
-    .controls-group { background: #2d3748; padding: 15px; border-radius: 5px; margin: 10px 0; }
-    .advanced-controls { background: #374151; padding: 12px; border-radius: 5px; margin: 8px 0; }
-    """
-
-    generator = ImageGenerator()
-
-    with gr.Blocks(css=css) as app:
-        gr.HTML("""
-            <div class="welcome">
-                <h1>🎨 Equity Artisan 3.0</h1>
-                <p>Assistant de création d'affiches professionnelles</p>
-            </div>
-        """)
-
-        with gr.Column(elem_classes="container"):
-            # Zone d'introduction et guide
-            gr.Markdown("""
-                ### 🎯 Guide d'utilisation
-                1. Choisissez le format et l'orientation de votre affiche
-                2. Sélectionnez un style artistique et une composition
-                3. Décrivez votre vision dans "Description"
-                4. Ajustez les paramètres fins selon vos besoins
-                5. Cliquez sur "Générer" !
-            """)
-
-            with gr.Group(elem_classes="controls-group"):
-                gr.Markdown("### 📐 Format et Orientation")
-                with gr.Row():
-                    format_size = gr.Dropdown(
-                        choices=["A4", "A3", "A2", "A1", "A0"],
-                        value="A4",
-                        label="Format",
-                        info="Choisissez la taille de votre affiche. A4 est le plus petit, A0 le plus grand."
-                    )
-                    orientation = gr.Radio(
-                        choices=["Portrait", "Paysage"],
-                        value="Portrait",
-                        label="Orientation",
-                        info="Portrait (vertical) ou Paysage (horizontal)"
-                    )
-
-            with gr.Group(elem_classes="controls-group"):
-                gr.Markdown("### 🎨 Style et Composition")
-                with gr.Row():
-                    style = gr.Dropdown(
-                        choices=list(ART_STYLES.keys()),
-                        value="Art Moderne",
-                        label="Style artistique",
-                        info="Le style global de votre affiche. Chaque style a ses propres caractéristiques uniques."
-                    )
-                    layout = gr.Dropdown(
-                        choices=list(COMPOSITION_PARAMS["Layouts"].keys()),
-                        value="Centré",
-                        label="Composition",
-                        info="Comment les éléments seront organisés dans l'affiche"
-                    )
-
-                with gr.Row():
-                    ambiance = gr.Dropdown(
-                        choices=list(COMPOSITION_PARAMS["Ambiances"].keys()),
-                        value="Dramatique",
-                        label="Ambiance",
-                        info="L'atmosphère générale de votre affiche"
-                    )
-                    palette = gr.Dropdown(
-                        choices=list(COMPOSITION_PARAMS["Palette"].keys()),
-                        value="Contrasté",
-                        label="Palette",
-                        info="Les types de couleurs utilisées"
-                    )
+            logging.error(f"Erreur lors de la génération: {str(e)}")
+            return None, f"Erreur: {str(e)}"
+
+def create_enhanced_interface():
+    # Styles enrichis avec paramètres optimisés
+    ENHANCED_STYLES = {
+        "Ultra Réaliste": {
+            "prompt_enhancement": "ultra photorealistic, 8k UHD, hyperdetailed",
+            "quality_boost": 1.2,
+            "steps_multiplier": 1.3
+        },
+        "Artistique Pro": {
+            "prompt_enhancement": "professional artistic composition, perfect lighting",
+            "quality_boost": 1.1,
+            "steps_multiplier": 1.2
+        }
+        # ... autres styles
+    }
 
-            with gr.Group(elem_classes="controls-group"):
-                gr.Markdown("""### 📝 Contenu
-                *Conseils pour la description : soyez précis sur ce que vous souhaitez voir dans l'affiche*""")
-                subject = gr.Textbox(
+    with gr.Blocks(theme=gr.themes.Soft()) as interface:
+        with gr.Row():
+            with gr.Column(scale=2):
+                # Interface utilisateur améliorée
+                prompt_input = gr.Textbox(
                     label="Description",
-                    placeholder="Ex: Une affiche moderne pour un festival de musique, avec des instruments colorés flottant dans l'espace",
-                    info="Décrivez en détail ce que vous souhaitez voir dans votre affiche"
+                    placeholder="Décrivez votre vision en détail...",
+                    lines=3
                 )
-                title = gr.Textbox(
-                    label="Titre",
-                    placeholder="Le titre qui apparaîtra sur l'affiche...",
-                    info="Laissez vide si vous ne voulez pas de titre sur l'affiche"
-                )
-
-            with gr.Group(elem_classes="advanced-controls"):
-                gr.Markdown("""### 🎯 Ajustements Fins
-                *Ces paramètres permettent d'affiner le résultat final*""")
+                
                 with gr.Row():
-                    detail_level = gr.Slider(
-                        minimum=1,
-                        maximum=10,
-                        value=7,
-                        step=1,
-                        label="Niveau de Détail",
-                        info="Plus la valeur est élevée, plus l'image sera détaillée"
+                    style_selector = gr.Dropdown(
+                        choices=list(ENHANCED_STYLES.keys()),
+                        label="Style",
+                        value="Ultra Réaliste"
                     )
-                    contrast = gr.Slider(
+                    quality_slider = gr.Slider(
                         minimum=1,
-                        maximum=10,
-                        value=5,
-                        step=1,
-                        label="Contraste",
-                        info="Influence l'intensité des couleurs et la différence entre les zones claires et sombres"
-                    )
-                    saturation = gr.Slider(
-                        minimum=1,
-                        maximum=10,
-                        value=5,
-                        step=1,
-                        label="Saturation",
-                        info="Contrôle la vivacité des couleurs"
+                        maximum=5,
+                        value=4,
+                        label="Niveau de Qualité",
+                        info="Impact la finesse des détails"
                     )
 
-            with gr.Group(elem_classes="controls-group"):
-                gr.Markdown("### ⚙️ Paramètres de Génération")
                 with gr.Row():
-                    quality = gr.Slider(
-                        minimum=30,
-                        maximum=50,
-                        value=35,
-                        label="Qualité",
-                        info="Influence la qualité finale de l'image. Une valeur plus élevée prend plus de temps"
+                    resolution_selector = gr.Radio(
+                        choices=["4K", "8K"],
+                        value="8K",
+                        label="Résolution"
                     )
-                    creativity = gr.Slider(
-                        minimum=5,
-                        maximum=15,
-                        value=7.5,
-                        label="Créativité",
-                        info="Plus la valeur est élevée, plus l'IA prendra de libertés créatives"
+                    
+                    aspect_ratio = gr.Radio(
+                        choices=["1:1", "16:9", "9:16"],
+                        value="1:1",
+                        label="Format"
                     )
 
-            with gr.Row():
-                generate_btn = gr.Button("✨ Générer", variant="primary")
-                clear_btn = gr.Button("🗑️ Effacer", variant="secondary")
-
-            image_output = gr.Image(label="Aperçu")
-            status = gr.Textbox(label="Statut", interactive=False)
-
-        def generate(*args):
-            logger.info("Démarrage d'une nouvelle génération")
+            with gr.Column(scale=3):
+                output_image = gr.Image(label="Résultat")
+                status = gr.Textbox(label="Status")
+
+        generate_btn = gr.Button("Générer", variant="primary")
+
+        # Logique de génération optimisée
+        def generate_optimized(prompt, style, quality, resolution, aspect):
+            generator = EnhancedImageGenerator()
+            
+            # Configuration adaptative
+            config = GenerationConfig(
+                num_inference_steps=int(50 * ENHANCED_STYLES[style]["steps_multiplier"]),
+                guidance_scale=7.5 * ENHANCED_STYLES[style]["quality_boost"],
+                width=3840 if resolution == "4K" else 7680,
+                height=2160 if resolution == "4K" else 4320
+            )
+            
             params = {
-                "format_size": args[0],
-                "orientation": args[1],
-                "style": args[2],
-                "layout": args[3],
-                "ambiance": args[4],
-                "palette": args[5],
-                "subject": args[6],
-                "title": args[7],
-                "detail_level": args[8],
-                "contrast": args[9],
-                "saturation": args[10],
-                "quality": args[11],
-                "creativity": args[12]
+                "prompt": prompt,
+                "style": style,
+                "resolution": resolution,
+                "quality_level": quality
             }
-            result = generator.generate(params)
-            logger.info(f"Génération terminée avec statut: {result[1]}")
-            return result
+            
+            return generator.generate(params, config)
 
         generate_btn.click(
-            generate,
-            inputs=[
-                format_size,
-                orientation,
-                style,
-                layout,
-                ambiance,
-                palette,
-                subject,
-                title,
-                detail_level,
-                contrast,
-                saturation,
-                quality,
-                creativity
-            ],
-            outputs=[image_output, status]
-        )
-
-        clear_btn.click(
-            lambda: (None, "🗑️ Image effacée"),
-            outputs=[image_output, status]
+            generate_optimized,
+            inputs=[prompt_input, style_selector, quality_slider, 
+                   resolution_selector, aspect_ratio],
+            outputs=[output_image, status]
         )
 
-    logger.info("Interface créée avec succès")
-    return app
+    return interface
 
 if __name__ == "__main__":
-    app = create_interface()
-    logger.info("Démarrage de l'application")
-    app.launch()
+    app = create_enhanced_interface()
+    app.launch()