app.py

import gradio as gr
import os
from PIL import Image
import requests
import io
import gc
import json
from typing import Tuple, Optional, Dict, Any
import logging
from dotenv import load_dotenv
from prompt_enhancer import PromptEnhancer

# Configuration du logging
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Chargement des variables d'environnement
load_dotenv()

# Styles artistiques complets
ART_STYLES = {
    # Styles Classiques
    "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:
    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]:
        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

        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"):
            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"
                    )
                    orientation = gr.Radio(
                        choices=["Portrait", "Paysage"],
                        value="Portrait",
                        label="Orientation"
                    )
            
            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"
                    )
                    layout = gr.Dropdown(
                        choices=list(COMPOSITION_PARAMS["Layouts"].keys()),
                        value="Centré",
                        label="Composition"
                    )

                with gr.Row():
                    ambiance = gr.Dropdown(
                        choices=list(COMPOSITION_PARAMS["Ambiances"].keys()),
                        value="Dramatique",
                        label="Ambiance"
                    )
                    palette = gr.Dropdown(
                        choices=list(COMPOSITION_PARAMS["Palette"].keys()),
                        value="Contrasté",
                        label="Palette"
                    )
            
            with gr.Group(elem_classes="controls-group"):
                gr.Markdown("### 📝 Contenu")
                subject = gr.Textbox(
                    label="Description",
                    placeholder="Décrivez votre vision..."
                )
                title = gr.Textbox(
                    label="Titre",
                    placeholder="Titre de l'affiche..."
                )
            
            with gr.Group(elem_classes="advanced-controls"):
                gr.Markdown("### 🎯 Ajustements Fins")
                with gr.Row():
                    detail_level = gr.Slider(
                        minimum=1,
                        maximum=10,
                        value=7,
                        step=1,
                        label="Niveau de Détail"
                    )
                    contrast = gr.Slider(
                        minimum=1,
                        maximum=10,
                        value=5,
                        step=1,
                        label="Contraste"
                    )
                    saturation = gr.Slider(
                        minimum=1,
                        maximum=10,
                        value=5,
                        step=1,
                        label="Saturation"
                    )
            
            with gr.Group(elem_classes="controls-group"):
                with gr.Row():
                    quality = gr.Slider(
                        minimum=30,
                        maximum=50,
                        value=35,
                        label="Qualité"
                    )
                    creativity = gr.Slider(
                        minimum=5,
                        maximum=15,
                        value=7.5,
                        label="Créativité"
                    )
            
            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")
            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]
            }
            result = generator.generate(params)
            logger.info(f"Génération terminée avec statut: {result[1]}")
            return result
        
        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]
        )
    
    logger.info("Interface créée avec succès")
    return app

if __name__ == "__main__":
    app = create_interface()
    logger.info("Démarrage de l'application")
    app.launch()