app.py

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

# Configuration du logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Chargement des variables d'environnement
load_dotenv()

class ImageGenerator:
    """Gestionnaire de génération d'images optimisé pour les ressources limitées"""
    
    def __init__(self):
        self.API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
        self.headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACE_TOKEN')}"}
        
        # Configurations optimisées pour les ressources limitées
        self.MAX_IMAGE_SIZE = 768  # Limite la taille maximale pour la mémoire
        self.MAX_STEPS = 40  # Limite les steps pour optimiser le temps
        self.DEFAULT_TIMEOUT = 30
        
    @staticmethod
    def get_image_dimensions(format_size: str, orientation: str) -> Tuple[int, int]:
        """Calcule les dimensions optimisées selon le format"""
        base_dimensions = {
            "A4": (768, 1024),
            "A3": (768, 1024),
            "A2": (768, 1024),
            "A1": (768, 1024),
            "A0": (768, 1024)
        }
        width, height = base_dimensions.get(format_size, (768, 768))
        return (height, width) if orientation == "Paysage" else (width, height)

    def generate(self, params: Dict[str, Any]) -> Tuple[Optional[Image.Image], str]:
        """Génère une image avec gestion optimisée des ressources"""
        try:
            # Nettoyage mémoire préventif
            gc.collect()
            
            # Préparation des dimensions
            width, height = self.get_image_dimensions(params["format_size"], params["orientation"])
            
            # Construction du prompt
            prompt = self._build_prompt(params)
            
            # Configuration de la requête
            payload = {
                "inputs": prompt,
                "parameters": {
                    "negative_prompt": self._get_negative_prompt(params["style"]),
                    "num_inference_steps": min(int(35 * (params["quality"]/100)), self.MAX_STEPS),
                    "guidance_scale": min(7.5 * (params["creativity"]/10), 10.0),
                    "width": width,
                    "height": height
                }
            }
            
            logger.info(f"Génération avec prompt: {prompt}")
            
            # Requête API avec gestion d'erreur
            response = requests.post(
                self.API_URL,
                headers=self.headers,
                json=payload,
                timeout=self.DEFAULT_TIMEOUT
            )
            
            if response.status_code == 200:
                image = Image.open(io.BytesIO(response.content))
                return image, "✨ Création réussie!"
            else:
                logger.error(f"Erreur API: {response.status_code} - {response.text}")
                return None, f"⚠️ Erreur {response.status_code}: Ajustez les paramètres"
                
        except Exception as e:
            logger.error(f"Erreur de génération: {str(e)}")
            return None, f"⚠️ Erreur: {str(e)}"
        finally:
            gc.collect()  # Nettoyage final

def create_interface():
    """Crée l'interface utilisateur Gradio"""
    
    # CSS optimisé
    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; }
    """
    
    # Initialisation du générateur
    generator = ImageGenerator()
    
    # Interface Gradio optimisée
    with gr.Blocks(css=css) as app:
        # En-tête
        gr.HTML("""
            <div class="welcome">
                <h1>🎨 Equity Artisan 3.0</h1>
                <p>Assistant de création d'affiches professionnelles</p>
            </div>
        """)
        
        # Conteneur principal
        with gr.Column(elem_classes="container"):
            # Groupe Format
            with gr.Group(elem_classes="controls-group"):
                with gr.Row():
                    format_size = gr.Dropdown(
                        choices=["A4", "A3", "A2", "A1", "A0"],
                        value="A4",
                        label="Format"
                    )
                    orientation = gr.Radio(
                        choices=["Portrait", "Paysage"],
                        value="Portrait",
                        label="Orientation"
                    )
            
            # Groupe Style
            with gr.Group(elem_classes="controls-group"):
                with gr.Row():
                    style = gr.Dropdown(
                        choices=["Art Moderne", "Neo Vintage", "Pop Art", "Minimaliste"],
                        value="Neo Vintage",
                        label="Style artistique"
                    )
                    text_effect = gr.Dropdown(
                        choices=["Standard", "Néon", "3D", "Métallique", "Graffiti"],
                        value="Standard",
                        label="Effet de texte"
                    )
            
            # Groupe Contenu
            with gr.Group(elem_classes="controls-group"):
                subject = gr.Textbox(
                    label="Description",
                    placeholder="Décrivez votre vision..."
                )
                title = gr.Textbox(
                    label="Titre",
                    placeholder="Titre de l'affiche..."
                )
            
            # Groupe Qualité
            with gr.Group(elem_classes="controls-group"):
                with gr.Row():
                    quality = gr.Slider(
                        minimum=30,
                        maximum=50,
                        value=35,
                        label="Qualité"
                    )
                    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", variant="secondary")
            
            # Sortie
            image_output = gr.Image(label="Aperçu")
            status = gr.Textbox(label="Statut", interactive=False)
        
        # Événements
        def generate(*args):
            params = {
                "format_size": args[0],
                "orientation": args[1],
                "subject": args[2],
                "style": args[3],
                "text_effect": args[4],
                "title": args[5],
                "quality": args[6],
                "creativity": args[7]
            }
            return generator.generate(params)
        
        generate_btn.click(
            generate,
            inputs=[
                format_size,
                orientation,
                subject,
                style,
                text_effect,
                title,
                quality,
                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()