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
    "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:
    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:
        try:
            enhancer = PromptEnhancer()
            style_info = ART_STYLES.get(params["style"], ART_STYLES["Art Moderne"])
            
            # Détection du type de contenu
            content_type = self._detect_content_type(params["subject"])
            
            # Adaptation du style selon le contenu
            style_adjustments = {
                "anime_manga": {
                    "prompt_prefix": "high quality anime artwork, detailed anime illustration, studio anime quality, professional anime art",
                    "negative_prompt": "low quality, simple, poorly drawn, blurry"
                },
                "art_digital": {
                    "prompt_prefix": "professional digital artwork, detailed digital illustration, modern digital art",
                    "negative_prompt": "traditional art, poor quality, amateur"
                },
                "art_traditionnel": {
                    "prompt_prefix": "traditional art masterpiece, fine art painting, artistic technique, professional artwork",
                    "negative_prompt": "digital art, 3d render, poor quality"
                },
                "photo_realiste": {
                    "prompt_prefix": "ultra realistic photograph, professional photography, high end photo, studio quality",
                    "negative_prompt": "drawing, painting, artificial, digital art"
                },
                "graphisme": {
                    "prompt_prefix": "professional graphic design, clean modern design, commercial quality artwork",
                    "negative_prompt": "amateur, messy, unrefined"
                },
                "fantastique": {
                    "prompt_prefix": "fantasy art masterpiece, mythical artwork, magical atmosphere, epic fantasy illustration",
                    "negative_prompt": "realistic, mundane, ordinary"
                },
                "sci_fi": {
                    "prompt_prefix": "futuristic sci-fi artwork, high tech aesthetic, advanced technology, science fiction art",
                    "negative_prompt": "vintage, retro, traditional"
                },
                "art_abstrait": {
                    "prompt_prefix": "abstract art composition, non-representational artwork, artistic expression",
                    "negative_prompt": "realistic, figurative, literal"
                },
                "art_pop": {
                    "prompt_prefix": "pop art style, bold colors, graphic art, contemporary pop culture",
                    "negative_prompt": "classical, traditional, subtle"
                },
                "art_conceptuel": {
                    "prompt_prefix": "concept art, professional design, detailed visualization, production quality",
                    "negative_prompt": "amateur, unrefined, sketch"
                }
            }

            # Ajustement du style en fonction du contenu détecté
            if content_type in style_adjustments:
                style_info = style_adjustments[content_type]
            
            # Construction du contexte
            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']]
            }
            
            base_prompt = f"{params['subject']}"
            if params.get('title'):
                base_prompt += f", with text '{params['title']}'"

            enhanced_prompt = enhancer.enhance_prompt(base_prompt, style_context)
            return enhanced_prompt

        except Exception as e:
            logger.error(f"Erreur dans la construction du prompt: {str(e)}")
            return f"{style_info['prompt_prefix']}, {params['subject']}"

    def _detect_content_type(self, subject: str) -> str:
        """Détecte le type de contenu demandé pour adapter le style"""
        content_types = {
            "anime_manga": ["manga", "anime", "dragon ball", "naruto", "one piece", "pokemon"],
            "art_digital": ["pixel art", "digital", "vectoriel", "3d", "game art"],
            "art_traditionnel": ["peinture", "aquarelle", "huile", "acrylique", "dessin"],
            "photo_realiste": ["photo", "portrait", "paysage", "architecture", "produit"],
            "graphisme": ["logo", "affiche", "flyer", "branding", "typographie"],
            "fantastique": ["fantasy", "dragon", "magie", "heroic fantasy", "mythologie"],
            "sci_fi": ["science fiction", "futuriste", "cyberpunk", "space", "robot"],
            "art_abstrait": ["abstrait", "géométrique", "non-figuratif", "minimaliste"],
            "art_pop": ["pop art", "comics", "bande dessinée", "cartoon", "graffiti"],
            "art_conceptuel": ["concept art", "character design", "environment design", "matte painting"]
        }

        subject_lower = subject.lower()
        detected_types = []
        
        for content_type, keywords in content_types.items():
            if any(keyword in subject_lower for keyword in keywords):
                detected_types.append(content_type)
        
        if len(detected_types) > 1:
            if "photo_realiste" in detected_types and any(t in detected_types for t in ["art_digital", "anime_manga"]):
                return "photo_realiste"
        
        return detected_types[0] if detected_types else "general"

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"
                    )
            
            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(
                    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"
                )
                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"
                    )
                    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"
                    )
            
            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"
                    )
            
            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(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']:
                logger.error("Token Hugging Face manquant ou invalide")
                return None, "⚠️ Erreur: Token Hugging Face non configuré"

            prompt = self._build_prompt(params)
            logger.info(f"Prompt généré: {prompt}")
            
            content_type = self._detect_content_type(params["subject"])
            logger.info(f"Type de contenu détecté: {content_type}")

            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 pour l'API: {json.dumps(payload, indent=2)}")

            try:
                response = requests.post(
                    self.API_URL,
                    headers=self.headers,
                    json=payload,
                    timeout=30
                )
                
                logger.info(f"Statut de la réponse API: {response.status_code}")
                if response.status_code != 200:
                    logger.error(f"Contenu de la réponse en erreur: {response.text}")
                
                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 requests.exceptions.Timeout:
                error_msg = "⚠️ Erreur: Le serveur met trop de temps à répondre"
                logger.error(error_msg)
                return None, error_msg
            except requests.exceptions.RequestException as e:
                error_msg = f"⚠️ Erreur de connexion: {str(e)}"
                logger.error(error_msg)
                return None, error_msg

        except Exception as e:
            error_msg = f"⚠️ Erreur inattendue: {str(e)}"
            logger.exception("Erreur détaillée pendant la génération:")
            return None, error_msg
        finally:
            gc.collect()
        
        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()