app.py

import gradio as gr
import os
from PIL import Image, ImageEnhance, ImageFilter
import requests
import io
import gc
import json
import logging
from dotenv import load_dotenv
from typing import Tuple, Optional, Dict, Any, List
import numpy as np
from dataclasses import dataclass

# Configuration avancée du logging
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("EquityArtisan")

# Chargement des variables d'environnement
load_dotenv()

@dataclass
class RenderingParams:
    """Paramètres de rendu avancés pour le contrôle fin de la génération"""
    steps: int  # Nombre d'étapes d'inférence
    cfg_scale: float  # Échelle de guidance
    width: int  # Largeur
    height: int  # Hauteur
    sampler: str  # Méthode d'échantillonnage
    seed: Optional[int]  # Seed pour la reproductibilité

@dataclass
class StyleConfig:
    """Configuration complète d'un style artistique"""
    prompt_prefix: str
    negative_prompt: str
    quality_modifiers: List[str]
    detail_modifiers: List[str]
    composition_rules: List[str]

# Styles artistiques enrichis
ART_STYLES = {
    "Ultra Photoréaliste": StyleConfig(
        prompt_prefix="masterpiece, ultra realistic photograph, octane render, unreal engine 5 rendering, ray tracing, global illumination, subsurface scattering, volumetric lighting, 8k uhd, photorealistic",
        negative_prompt="cartoon, illustration, painting, drawing, sketch, anime, low quality, blurry, noise, grain, deformed, ugly, disfigured",
        quality_modifiers=["hyperdetailed", "photographic", "award winning photography"],
        detail_modifiers=["intricate details", "sharp focus", "professional lighting"],
        composition_rules=["rule of thirds", "golden ratio", "perfect composition"]
    ),
    "Cinématique": StyleConfig(
        prompt_prefix="cinematic scene, anamorphic lens, movie still, dramatic lighting, high production value, professional cinematography, ARRI camera",
        negative_prompt="amateur, low budget, poor lighting, webcam, phone camera",
        quality_modifiers=["film grain", "color grading", "bokeh effect"],
        detail_modifiers=["depth of field", "motion blur", "lens flare"],
        composition_rules=["wide aspect ratio", "leading lines", "frame within frame"]
    ),
    # ... [Autres styles similaires] ...
}

# Paramètres de qualité avancés
QUALITY_PRESETS = {
    "Standard": RenderingParams(
        steps=30,
        cfg_scale=7.5,
        width=768,
        height=768,
        sampler="Euler a",
        seed=None
    ),
    "Haute Qualité": RenderingParams(
        steps=40,
        cfg_scale=8.5,
        width=1024,
        height=1024,
        sampler="DPM++ 2M Karras",
        seed=None
    ),
    "Ultra Qualité": RenderingParams(
        steps=50,
        cfg_scale=9.0,
        width=1536,
        height=1536,
        sampler="DPM++ SDE Karras",
        seed=None
    )
}

# Enrichissement des paramètres de composition
COMPOSITION_PARAMS = {
    "Layouts": {
        "Dynamique": {
            "description": "Composition dynamique avec lignes de force",
            "modifiers": ["dynamic composition", "leading lines", "visual flow"],
            "weights": {"foreground": 1.2, "background": 0.8}
        },
        "Équilibré": {
            "description": "Composition harmonieuse et équilibrée",
            "modifiers": ["balanced composition", "symmetrical layout", "visual harmony"],
            "weights": {"foreground": 1.0, "background": 1.0}
        }
        # ... [Autres layouts] ...
    },
    "Éclairages": {
        "Studio": {
            "description": "Éclairage professionnel type studio",
            "modifiers": ["professional studio lighting", "three-point lighting", "soft boxes"],
            "technical_params": {"contrast": 1.2, "highlights": 1.1}
        },
        "Naturel": {
            "description": "Éclairage naturel avec lumière ambiante",
            "modifiers": ["natural lighting", "golden hour", "ambient occlusion"],
            "technical_params": {"contrast": 0.9, "shadows": 0.8}
        }
        # ... [Autres éclairages] ...
    }
}

class AdvancedImageGenerator:
    def __init__(self):
        self.API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
        self.token = os.getenv('HUGGINGFACE_TOKEN')
        if not self.token:
            raise EnvironmentError("HUGGINGFACE_TOKEN non configuré")
        self.headers = {"Authorization": f"Bearer {self.token}"}
        
        # Initialisation du cache pour les paramètres optimaux
        self.optimization_cache = {}
        logger.info("AdvancedImageGenerator initialisé avec paramètres enrichis")

    def _optimize_prompt(self, base_prompt: str, style_config: StyleConfig, params: Dict[str, Any]) -> str:
        """Optimisation avancée du prompt avec analyse sémantique"""
        components = [
            style_config.prompt_prefix,
            base_prompt
        ]

        # Ajout des modificateurs de qualité selon le niveau demandé
        quality_level = params.get("quality", 50)
        if quality_level > 40:
            components.extend(style_config.quality_modifiers)
        
        # Ajout des règles de composition
        components.extend(style_config.composition_rules)

        # Optimisation contextuelle basée sur le sujet
        if "portrait" in base_prompt.lower():
            components.extend([
                "perfect facial features",
                "detailed skin texture",
                "professional portrait lighting"
            ])
        elif "paysage" in base_prompt.lower():
            components.extend([
                "atmospheric perspective",
                "perfect weather conditions",
                "golden hour lighting"
            ])

        return ", ".join(components)

    def _apply_advanced_processing(self, image: Image.Image, params: Dict[str, Any]) -> Image.Image:
        """Application de traitements avancés post-génération"""
        try:
            # Amélioration des détails
            if params.get("detail_enhancement", True):
                image = image.filter(ImageFilter.DETAIL)
            
            # Optimisation du contraste local
            if params.get("local_contrast", True):
                enhancer = ImageEnhance.Contrast(image)
                image = enhancer.enhance(1.2)
            
            # Ajustement des couleurs
            if params.get("color_optimization", True):
                enhancer = ImageEnhance.Color(image)
                image = enhancer.enhance(1.1)
            
            return image
        except Exception as e:
            logger.error(f"Erreur lors du post-traitement: {e}")
            return image

    def generate(self, params: Dict[str, Any]) -> Tuple[Optional[Image.Image], str]:
        """Génération d'image avec paramètres avancés"""
        try:
            logger.info(f"Démarrage génération avec paramètres: {json.dumps(params, indent=2)}")
            
            style_config = ART_STYLES[params["style"]]
            prompt = self._optimize_prompt(params["subject"], style_config, params)
            
            # Sélection des paramètres de rendu
            quality_preset = QUALITY_PRESETS["Ultra Qualité" if params["quality"] > 45 else "Haute Qualité"]
            
            payload = {
                "inputs": prompt,
                "parameters": {
                    "negative_prompt": style_config.negative_prompt,
                    "num_inference_steps": quality_preset.steps,
                    "guidance_scale": quality_preset.cfg_scale,
                    "width": quality_preset.width,
                    "height": quality_preset.height,
                    "sampler": quality_preset.sampler
                }
            }

            response = requests.post(
                self.API_URL,
                headers=self.headers,
                json=payload,
                timeout=60
            )

            if response.status_code != 200:
                raise Exception(f"Erreur API: {response.status_code}")

            # Traitement et optimisation de l'image générée
            image = Image.open(io.BytesIO(response.content))
            image = self._apply_advanced_processing(image, params)
            
            return image, "✨ Génération réussie avec paramètres optimaux!"

        except Exception as e:
            logger.exception("Erreur pendant la génération:")
            return None, f"⚠️ Erreur: {str(e)}"

def create_interface():
    """Création de l'interface utilisateur enrichie"""
    logger.info("Initialisation de l'interface avancée")

    css = """
    .container { max-width: 1400px; margin: auto; padding: 20px; }
    .welcome { 
        text-align: center; 
        margin: 20px 0; 
        padding: 30px;
        background: linear-gradient(135deg, #1e293b, #334155);
        border-radius: 15px;
        color: white;
    }
    .controls-group { 
        background: #2d3748; 
        padding: 20px;
        border-radius: 10px;
        margin: 15px 0;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    }
    """

    generator = AdvancedImageGenerator()

    with gr.Blocks(css=css) as app:
        gr.HTML("""
            <div class="welcome">
                <h1>🎨 Equity Artisan 4.0 Pro</h1>
                <p>Studio de Création Visuelle Avancé</p>
            </div>
        """)

        with gr.Column(elem_classes="container"):
            # Configuration principale
            with gr.Group(elem_classes="controls-group"):
                gr.Markdown("### 🎭 Style et Composition")
                with gr.Row():
                    style = gr.Dropdown(
                        choices=list(ART_STYLES.keys()),
                        value="Ultra Photoréaliste",
                        label="Style Artistique"
                    )
                    quality_preset = gr.Dropdown(
                        choices=list(QUALITY_PRESETS.keys()),
                        value="Haute Qualité",
                        label="Preset de Qualité"
                    )

            # Paramètres créatifs
            with gr.Group(elem_classes="controls-group"):
                gr.Markdown("### 🎨 Paramètres Créatifs")
                with gr.Row():
                    subject = gr.Textbox(
                        label="Description",
                        placeholder="Décrivez votre vision en détail...",
                        lines=3
                    )
                    composition = gr.Dropdown(
                        choices=list(COMPOSITION_PARAMS["Layouts"].keys()),
                        value="Dynamique",
                        label="Type de Composition"
                    )

            # Contrôles techniques
            with gr.Group(elem_classes="controls-group"):
                gr.Markdown("### ⚙️ Paramètres Techniques")
                with gr.Row():
                    quality = gr.Slider(
                        minimum=30,
                        maximum=100,
                        value=75,
                        label="Niveau de Qualité"
                    )
                    detail_level = gr.Slider(
                        minimum=1,
                        maximum=10,
                        value=8,
                        label="Niveau de Détail"
                    )

            # Génération
            with gr.Row():
                generate_btn = gr.Button("✨ Générer", variant="primary", size="lg")
                clear_btn = gr.Button("🗑️ Réinitialiser", variant="secondary")

            # Résultats
            image_output = gr.Image(
                label="Aperçu",
                elem_classes="result-image"
            )
            status = gr.Textbox(
                label="Statut",
                interactive=False
            )

        # Logique de génération
        def generate_image(*args):
            params = {
                "style": args[0],
                "quality_preset": args[1],
                "subject": args[2],
                "composition": args[3],
                "quality": args[4],
                "detail_level": args[5]
            }
            return generator.generate(params)

        # Événements
        generate_btn.click(
            generate_image,
            inputs=[
                style,
                quality_preset,
                subject,
                composition,
                quality,
                detail_level
            ],
            outputs=[image_output, status]
        )

        clear_btn.click(
            lambda: (None, "🔄 Interface réinitialisée"),
            outputs=[image_output, status]
        )

    logger.info("Interface créée avec succès")
    return app

if __name__ == "__main__":
    app = create_interface()
    app.launch()