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
import numpy as np
import cv2

# Configuration du logging
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('equity_space.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)
load_dotenv()

# 1. STYLES EQUITY (basés sur les 8 documents)
ART_STYLES = {
    "Renaissance Technologique": {
        "prompt_prefix": """ultra-detailed technical drawing, da vinci style engineering precision, 
        anatomical accuracy, golden ratio proportions, scientific illustration quality, 
        masterful light and shadow, highest level of detail, museum quality artwork""",
        "negative_prompt": "simple sketch, imprecise, cartoon, abstract, low detail",
        "quality_boost": 1.5,
        "style_params": {
            "detail_level": "maximum",
            "technical_precision": True,
            "historical_accuracy": True
        }
    },
    "Innovation Moderne": {
        "prompt_prefix": """cutting-edge technological design, futuristic engineering visual, 
        advanced mechanical detail, precise technical schematic, innovative concept art, 
        professional industrial visualization, high-tech aesthetic""",
        "negative_prompt": "vintage, retro, simplified, abstract",
        "quality_boost": 1.4,
        "style_params": {
            "tech_detail": "ultra",
            "innovation_focus": True
        }
    },
    "Photographie Analytique": {
        "prompt_prefix": """professional analytical photography, ansel adams zone system, 
        ultra sharp detail, perfect exposure, museum grade quality, technical perfection, 
        masterful composition, extreme clarity""",
        "negative_prompt": "blurry, artistic, painterly, low quality",
        "quality_boost": 1.5,
        "style_params": {
            "photographic_precision": True,
            "detail_preservation": "maximum"
        }
    },
    "Vision Futuriste": {
        "prompt_prefix": """advanced technological concept, future engineering visualization, 
        innovative architectural design, sci-fi technical precision, ultra modern aesthetic, 
        professional technical illustration""",
        "negative_prompt": "historical, vintage, traditional, hand-drawn",
        "quality_boost": 1.4,
        "style_params": {
            "future_tech": True,
            "concept_clarity": "high"
        }
    }
}

# 2. OPTIMISATION DE LA QUALITÉ
class ImageEnhancer:
    def __init__(self):
        self.enhancement_pipeline = {
            "technical_detail": self._enhance_technical_detail,
            "color_refinement": self._refine_colors,
            "sharpness_boost": self._boost_sharpness,
            "contrast_optimization": self._optimize_contrast
        }

    def _enhance_technical_detail(self, image: np.ndarray) -> np.ndarray:
        kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
        return cv2.filter2D(image, -1, kernel)

    def _refine_colors(self, image: np.ndarray) -> np.ndarray:
        lab = cv2.cvtColor(image, cv2.COLOR_BGR2Lab)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
        l = clahe.apply(l)
        lab = cv2.merge((l,a,b))
        return cv2.cvtColor(lab, cv2.COLOR_Lab2BGR)

    def _boost_sharpness(self, image: np.ndarray) -> np.ndarray:
        blurred = cv2.GaussianBlur(image, (0, 0), 3)
        return cv2.addWeighted(image, 1.5, blurred, -0.5, 0)

    def _optimize_contrast(self, image: np.ndarray) -> np.ndarray:
        lab = cv2.cvtColor(image, cv2.COLOR_BGR2Lab)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        l = clahe.apply(l)
        lab = cv2.merge((l,a,b))
        return cv2.cvtColor(lab, cv2.COLOR_Lab2BGR)

    def enhance_image(self, image: Image.Image, style_params: Dict) -> Image.Image:
        cv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        
        for enhancement, should_apply in style_params.items():
            if should_apply and enhancement in self.enhancement_pipeline:
                cv_image = self.enhancement_pipeline[enhancement](cv_image)
        
        return Image.fromarray(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))

# 3. GÉNÉRATEUR D'IMAGES PRINCIPAL
class ImageGenerator:
    def __init__(self):
        self.API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
        self.enhancer = ImageEnhancer()
        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 generate(self, params: Dict[str, Any]) -> Tuple[Optional[Image.Image], str]:
        try:
            style_info = ART_STYLES.get(params["style"])
            if not style_info:
                return None, "⚠️ Style non trouvé"

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

            payload = {
                "inputs": prompt,
                "parameters": {
                    "negative_prompt": style_info["negative_prompt"],
                    "num_inference_steps": int(50 * style_info["quality_boost"]),
                    "guidance_scale": min(9.0 * style_info["quality_boost"], 12.0),
                    "width": 1024 if params.get("quality", 35) > 40 else 768,
                    "height": 1024 if params["orientation"] == "Portrait" else 768
                }
            }

            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))
                enhanced_image = self.enhancer.enhance_image(
                    image, 
                    style_info["style_params"]
                )
                return enhanced_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 génération:")
            return None, error_msg
        finally:
            gc.collect()

# 4. INTERFACE UTILISATEUR
def create_interface():
    generator = ImageGenerator()

    # Style CSS personnalisé
    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);
    }
    """

    with gr.Blocks(css=css) as app:
        gr.HTML("""
            <div class="welcome">
                <h1>🎨 Equity Artisan 3.0</h1>
                <p>Assistant de Création d'Images Professionnelles</p>
            </div>
        """)

        with gr.Column(elem_classes="container"):
            # Format et Style
            with gr.Group(elem_classes="controls-group"):
                with gr.Row():
                    format_size = gr.Dropdown(
                        choices=["A4", "A3", "A2", "A1"],
                        value="A4",
                        label="Format"
                    )
                    orientation = gr.Radio(
                        choices=["Portrait", "Paysage"],
                        value="Portrait",
                        label="Orientation"
                    )
                    style = gr.Dropdown(
                        choices=list(ART_STYLES.keys()),
                        value="Renaissance Technologique",
                        label="Style Artistique"
                    )

            # Description
            with gr.Group(elem_classes="controls-group"):
                subject = gr.Textbox(
                    label="Description",
                    placeholder="Décrivez votre vision technique ou artistique...",
                    lines=3
                )
                title = gr.Textbox(
                    label="Titre (optionnel)",
                    placeholder="Titre à inclure dans l'image..."
                )

            # Paramètres avancés
            with gr.Group(elem_classes="controls-group"):
                with gr.Row():
                    quality = gr.Slider(
                        minimum=30,
                        maximum=50,
                        value=40,
                        label="Qualité"
                    )
                    detail_level = gr.Slider(
                        minimum=1,
                        maximum=10,
                        value=8,
                        step=1,
                        label="Niveau de Détail"
                    )

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

            # Résultat
            image_output = gr.Image(label="Résultat")
            status = gr.Textbox(label="Status", interactive=False)

        def generate(*args):
            params = {
                "format_size": args[0],
                "orientation": args[1],
                "style": args[2],
                "subject": args[3],
                "title": args[4],
                "quality": args[5],
                "detail_level": args[6]
            }
            return generator.generate(params)

        generate_btn.click(
            generate,
            inputs=[
                format_size,
                orientation,
                style,
                subject,
                title,
                quality,
                detail_level
            ],
            outputs=[image_output, status]
        )

        clear_btn.click(
            lambda: (None, "🗑️ Image effacée"),
            outputs=[image_output, status]
        )

    return app

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