image_processor.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
import numpy as np
import cv2
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
ART_STYLES = {
    "Ultra Réaliste": {
        "prompt_prefix": "ultra realistic photograph, stunning photorealistic quality, unreal engine 5 quality, cinema quality, masterpiece, perfect composition, award winning photography, professional lighting, 8k UHD",
        "negative_prompt": "artificial, digital art, illustration, painting, drawing, artistic, cartoon, anime, unreal, fake, low quality, blurry, soft, deformed",
        "quality_boost": 1.2
    },
    "Photoréaliste": {
        "prompt_prefix": "hyperrealistic studio photograph, extremely detailed, professional photography, perfect lighting, high-end camera, 8k uhd",
        "negative_prompt": "artificial, illustration, painting, animated, cartoon, artistic",
        "quality_boost": 1.1
    },
    "Art Moderne": {
        "prompt_prefix": "modern art style, professional design, contemporary aesthetic, trending artwork, perfect composition",
        "negative_prompt": "old style, vintage, traditional, amateur, low quality",
        "quality_boost": 1.0
    },
    "Minimaliste": {
        "prompt_prefix": "minimalist design, clean composition, elegant simplicity, refined aesthetic",
        "negative_prompt": "complex, cluttered, busy, ornate, detailed",
        "quality_boost": 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 _enhance_image(self, image: Image.Image, params: Dict[str, Any]) -> Image.Image:
        """Amélioration avancée de la qualité d'image"""
        try:
            # Conversion en CV2 pour traitement avancé
            cv2_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
            
            # Débruitage adaptatif
            if params.get("quality", 35) > 40:
                cv2_image = cv2.fastNlMeansDenoisingColored(cv2_image, None, 10, 10, 7, 21)
            
            # Amélioration des détails
            kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
            cv2_image = cv2.filter2D(cv2_image, -1, kernel)
            
            # Reconversion en PIL
            image = Image.fromarray(cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB))
            
            # Ajustements fins avec PIL
            enhancers = [
                (ImageEnhance.Sharpness, params.get("detail_level", 7) / 5),
                (ImageEnhance.Contrast, params.get("contrast", 5) / 5),
                (ImageEnhance.Color, params.get("saturation", 5) / 5)
            ]
            
            for enhancer_class, factor in enhancers:
                if factor != 1.0:
                    image = enhancer_class(image).enhance(factor)
            
            return image
        except Exception as e:
            logger.error(f"Erreur traitement image: {str(e)}")
            return image

    def _build_prompt(self, params: Dict[str, Any]) -> str:
        """Construction optimisée du prompt"""
        try:
            style_info = ART_STYLES.get(params["style"], ART_STYLES["Art Moderne"])
            
            # 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 style
            prompt = f"{base_prompt}, {style_info['prompt_prefix']}"
            
            return prompt
        except Exception as e:
            logger.error(f"Erreur prompt: {str(e)}")
            return params['subject']

    def generate(self, params: Dict[str, Any]) -> Tuple[Optional[Image.Image], str]:
        try:
            if 'Bearer None' in self.headers['Authorization']:
                return None, "⚠️ Erreur: Token Hugging Face non configuré"

            # Optimisation des paramètres
            style_info = ART_STYLES.get(params["style"], ART_STYLES["Art Moderne"])
            quality_boost = style_info.get("quality_boost", 1.0)

            # Préparation de la requête
            prompt = self._build_prompt(params)
            payload = {
                "inputs": prompt,
                "parameters": {
                    "negative_prompt": style_info["negative_prompt"],
                    "num_inference_steps": min(int(40 * quality_boost), 50),
                    "guidance_scale": min(8.0 * 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))
                # Application des améliorations de qualité
                enhanced_image = self._enhance_image(image, 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()

def create_interface():
    generator = ImageGenerator()

    with gr.Blocks(css="style.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():
            # Contrôles principaux
            with gr.Group():
                gr.Markdown("### 📐 Format et Style")
                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="Art Moderne",
                        label="Style artistique"
                    )

            # Description
            with gr.Group():
                gr.Markdown("### 📝 Description")
                subject = gr.Textbox(
                    label="Description",
                    placeholder="Décrivez votre vision...",
                    lines=3
                )
                title = gr.Textbox(
                    label="Titre (optionnel)",
                    placeholder="Titre à inclure..."
                )

            # Paramètres avancés
            with gr.Group():
                gr.Markdown("### ⚙️ Paramètres")
                with gr.Row():
                    quality = gr.Slider(
                        minimum=30,
                        maximum=50,
                        value=35,
                        label="Qualité"
                    )
                    detail_level = gr.Slider(
                        minimum=1,
                        maximum=10,
                        value=7,
                        step=1,
                        label="Niveau de Détail"
                    )
                    creativity = gr.Slider(
                        minimum=5,
                        maximum=15,
                        value=7.5,
                        label="Créativité"
                    )

            # 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],
                "creativity": args[7]
            }
            return generator.generate(params)

        generate_btn.click(
            generate,
            inputs=[format_size, orientation, style, subject, title,
                   quality, detail_level, creativity],
            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()