app.py

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

# 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 enrichis avec optimisations
ART_STYLES = {
    # Styles Réalistes
    "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, professional lighting""",
        "negative_prompt": """artificial, digital art, illustration, painting, drawing, artistic, cartoon, 
        anime, unreal, fake, low quality, blurry, soft, deformed, noisy, unclear, imperfect, amateur""",
        "quality_multiplier": 1.2
    },
    "Photoréaliste": {
        "prompt_prefix": """hyperrealistic photograph, extremely detailed, studio quality, professional photography, 
        8k uhd, perfect lighting, high-end camera, professional grade lens, expert composition""",
        "negative_prompt": """artistic, painterly, abstract, cartoon, illustration, low quality, 
        amateur, imperfect, blurry, noise""",
        "quality_multiplier": 1.1
    },
    
    # Styles Artistiques
    "Expressionniste": {
        "prompt_prefix": """expressive painting style, intense emotional art, bold brushstrokes, vibrant colors, 
        van gogh inspired, artistic masterpiece, deep emotional impact, dynamic composition""",
        "negative_prompt": """realistic, subtle, photographic, clean lines, digital art, flat, unemotional, 
        generic, weak composition""",
        "quality_multiplier": 1.0
    },
    "Impressionniste": {
        "prompt_prefix": """impressionist painting style, soft light, visible brushstrokes, outdoor scene, 
        monet inspired, artistic mastery, natural lighting, atmospheric effect""",
        "negative_prompt": """sharp details, high contrast, digital, modern, artificial, harsh, 
        unrealistic colors""",
        "quality_multiplier": 1.0
    },
    
    # Styles Modernes
    "Art Moderne": {
        "prompt_prefix": """modern art style poster, professional design, contemporary aesthetic, 
        trending on artstation, perfect composition, high-end production""",
        "negative_prompt": """traditional, cluttered, busy design, vintage, amateur, low quality, 
        unprofessional""",
        "quality_multiplier": 1.1
    },
    "Minimaliste": {
        "prompt_prefix": """minimalist design poster, clean composition, elegant simplicity, 
        perfect balance, sophisticated style, high-end design""",
        "negative_prompt": """complex, detailed, ornate, busy, cluttered, chaotic, unbalanced, 
        amateur""",
        "quality_multiplier": 1.0
    },

    # Styles Spéciaux
    "Cyberpunk": {
        "prompt_prefix": """cyberpunk style, neon lights, futuristic design, high-tech aesthetic, 
        detailed machinery, holographic elements, cinematic lighting""",
        "negative_prompt": """vintage, natural, rustic, traditional, simple, flat, dull, 
        low-tech""",
        "quality_multiplier": 1.1
    },
    "Art Déco": {
        "prompt_prefix": """art deco style, geometric patterns, luxury design, 1920s aesthetic, 
        golden age glamour, sophisticated composition""",
        "negative_prompt": """modern, minimalist, casual, contemporary, simple, rustic, 
        unrefined""",
        "quality_multiplier": 1.0
    }
}

# Paramètres de composition enrichis
COMPOSITION_PARAMS = {
    "Layouts": {
        "Centré": {
            "description": "centered composition, balanced layout, harmonious arrangement",
            "weight": 1.2
        },
        "Asymétrique": {
            "description": "dynamic asymmetrical composition, creative balance, artistic flow",
            "weight": 1.1
        },
        "Grille": {
            "description": "grid-based layout, structured composition, organized design, perfect alignment",
            "weight": 1.0
        },
        "Diagonal": {
            "description": "diagonal dynamic composition, energetic flow, dramatic arrangement",
            "weight": 1.1
        },
        "Minimaliste": {
            "description": "minimal composition, lots of whitespace, elegant spacing, perfect balance",
            "weight": 1.0
        }
    },
    "Ambiances": {
        "Dramatique": {
            "description": "dramatic lighting, high contrast, intense mood, cinematic atmosphere",
            "weight": 1.2
        },
        "Doux": {
            "description": "soft lighting, gentle atmosphere, subtle mood, delicate ambiance",
            "weight": 1.0
        },
        "Vibrant": {
            "description": "vibrant colors, energetic mood, dynamic atmosphere, bold presence",
            "weight": 1.1
        },
        "Mystérieux": {
            "description": "mysterious atmosphere, moody lighting, enigmatic feel, intriguing shadows",
            "weight": 1.1
        },
        "Serein": {
            "description": "peaceful atmosphere, calm mood, tranquil setting, harmonious lighting",
            "weight": 1.0
        }
    }
}

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 _optimize_prompt(self, params: Dict[str, Any]) -> Tuple[str, str]:
        """Optimisation avancée des prompts avec gestion contextuelle"""
        style_info = ART_STYLES.get(params["style"], ART_STYLES["Art Moderne"])
        layout_info = COMPOSITION_PARAMS["Layouts"].get(params["layout"])
        ambiance_info = COMPOSITION_PARAMS["Ambiances"].get(params["ambiance"])

        # Construction du prompt principal
        base_prompt = f"{params['subject']}"
        if params.get('title'):
            base_prompt += f", with text '{params['title']}'"

        # Ajout des éléments de composition
        composition_elements = [
            style_info["prompt_prefix"],
            layout_info["description"],
            ambiance_info["description"]
        ]

        # Calcul du multiplicateur de qualité
        quality_multiplier = (
            style_info.get("quality_multiplier", 1.0) *
            layout_info.get("weight", 1.0) *
            ambiance_info.get("weight", 1.0)
        )

        # Construction du prompt final optimisé
        enhanced_prompt = f"{base_prompt}, {', '.join(composition_elements)}"

        # Construction du negative prompt optimisé
        negative_prompt = f"{style_info['negative_prompt']}, low quality, bad anatomy, worst quality, low resolution"

        return enhanced_prompt, negative_prompt, quality_multiplier

    def _enhance_image(self, image: Image.Image, params: Dict[str, Any]) -> Image.Image:
        """Post-traitement avancé des images"""
        try:
            # Amélioration de la netteté basée sur le style
            sharpness_factor = 1.2 if params["style"] in ["Ultra Réaliste", "Photoréaliste"] else 1.1
            enhancer = ImageEnhance.Sharpness(image)
            image = enhancer.enhance(sharpness_factor)

            # Ajustement du contraste selon l'ambiance
            contrast_factor = 1.2 if params["ambiance"] == "Dramatique" else 1.1
            enhancer = ImageEnhance.Contrast(image)
            image = enhancer.enhance(contrast_factor)

            return image
        except Exception as e:
            logger.warning(f"Erreur lors de l'amélioration de l'image: {str(e)}")
            return image

    def generate(self, params: Dict[str, Any]) -> Tuple[Optional[Image.Image], str]:
        """Génération d'image avec optimisations avancées"""
        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é"

            # Optimisation des prompts
            enhanced_prompt, negative_prompt, quality_multiplier = self._optimize_prompt(params)

            # Configuration des paramètres de génération
            generation_params = {
                "num_inference_steps": min(int(50 * quality_multiplier), 60),
                "guidance_scale": min(8.5 * (params["creativity"]/10), 12.0),
                "width": 1024 if params.get("quality", 35) > 40 else 768,
                "height": 1024 if params["orientation"] == "Portrait" else 768
            }

            payload = {
                "inputs": enhanced_prompt,
                "parameters": {
                    **generation_params,
                    "negative_prompt": negative_prompt
                }
            }

            logger.debug(f"Payload final: {json.dumps(payload, indent=2)}")

            # Génération de l'image
            response = requests.post(
                self.API_URL,
                headers=self.headers,
                json=payload,
                timeout=45
            )

            if response.status_code == 200:
                image = Image.open(io.BytesIO(response.content))
                # Post-traitement de l'image
                image = self._enhance_image(image, params)
                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():
    """Création de l'interface utilisateur enrichie"""
    logger.info("Création de l'interface Gradio")

    # Styles CSS personnalisés
    css = """
    .container { max-width: 1200px; margin: auto; }
    .welcome { 
        text-align: center; 
        margin: 20px 0; 
        padding: 20px; 
        background: linear-gradient(135deg, #1e293b, #334155);
        border-radius: 10px;
        color: white;
    }
    .controls-group { 
        background: #2d3748; 
        padding: 15px; 
        border-radius: 5px; 
        margin: 10px 0;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    }
    .advanced-controls { 
        background: #374151; 
        padding: 12px; 
        border-radius: 5px; 
        margin: 8px 0;
    }
    .info-tooltip {
        color: #94a3b8;
        font-size: 0.9em;
        margin-top: 4px;
    }
    """

    generator = ImageGenerator()

    with gr.Blocks(css=css) as app:
        # En-tête et présentation
        gr.HTML("""
            <div class="welcome">
                <h1>🎨 Equity Artisan 3.0</h1>
                <p>Assistant de création d'affiches professionnelles</p>
                <p style="font-size: 0.9em; opacity: 0.8;">Powered by Stable Diffusion XL</p>
            </div>
        """)

        with gr.Column(elem_classes="container"):
            # Guide d'utilisation
            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" !
                
                *💡 Pro Tip: Pour de meilleurs résultats, soyez précis dans votre description et expérimentez avec différents styles.*
            """)

            # Contrôles principaux
            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"
                    )
                    orientation = gr.Radio(
                        choices=["Portrait", "Paysage"],
                        value="Portrait",
                        label="Orientation"
                    )

            # Style et composition
            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=["Monochrome", "Contrasté", "Pastel", "Terre", "Néon"],
                        value="Contrasté",
                        label="Palette"
                    )

            # Description et contenu
            with gr.Group(elem_classes="controls-group"):
                gr.Markdown("### 📝 Contenu")
                subject = gr.Textbox(
                    label="Description",
                    placeholder="Ex: Une affiche moderne pour un festival de musique, avec des instruments colorés flottant dans l'espace",
                    lines=3
                )
                title = gr.Textbox(
                    label="Titre (optionnel)",
                    placeholder="Le titre qui apparaîtra sur l'affiche..."
                )

            # Contrôles avancés
            with gr.Group(elem_classes="advanced-controls"):
                gr.Markdown("### 🎯 Paramètres Av# Contrôles avancés
            with gr.Group(elem_classes="advanced-controls"):
                gr.Markdown("### 🎯 Paramètres Avancés")
                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"
                    )
                    contrast = 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"
                    )

            # Paramètres de génération
            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"
                    )
                    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"
                    )

            # Boutons de contrôle
            with gr.Row():
                generate_btn = gr.Button("✨ Générer", variant="primary")
                clear_btn = gr.Button("🗑️ Effacer", variant="secondary")

            # Zone de résultat
            with gr.Group(elem_classes="controls-group"):
                gr.Markdown("### 🖼️ Résultat")
                image_output = gr.Image(label="Aperçu", height=512)
                status = gr.Textbox(label="Statut", interactive=False)

            # Zone d'historique
            with gr.Group(elem_classes="controls-group"):
                gr.Markdown("### 📋 Historique des Générations")
                history = gr.Gallery(label="Générations précédentes", show_label=True, columns=4, height=200)

        # Fonction de génération
        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

        # Connexion des événements
        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()