import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
from PIL import Image
import os

class EquityImageGenerator:
    def __init__(self):
        # Initialisation du modèle Stable Diffusion
        self.model_id = "stabilityai/stable-diffusion-2-1"
        self.pipe = StableDiffusionPipeline.from_pretrained(
            self.model_id,
            torch_dtype=torch.float16
        ).to("cuda" if torch.cuda.is_available() else "cpu")

    def generate_image(self, poster_type, style, message, theme=""):
        # Construire le prompt pour l'IA
        base_prompt = f"high quality poster design for {message}, {style} style, {poster_type} layout"
        artistic_prompt = f"{base_prompt}, professional graphic design, trending on artstation, {theme}"
        
        # Générer l'image
        image = self.pipe(
            prompt=artistic_prompt,
            negative_prompt="low quality, bad text, blurry, amateur",
            num_inference_steps=50,
            guidance_scale=7.5
        ).images[0]
        
        return image

def process_request(poster_type, style, message, audience, budget, theme):
    # Initialiser le générateur
    generator = EquityImageGenerator()
    
    # Générer l'image
    generated_image = generator.generate_image(
        poster_type=poster_type,
        style=style,
        message=message,
        theme=theme
    )
    
    # Générer les recommandations
    recommendations = f"""
🎨 Recommandations pour votre affiche :

📌 Informations de base :
• Type : {poster_type}
• Style : {style}
• Message : {message}
• Public : {audience}
• Budget : {budget}
• Thème : {theme}

✨ Style et design :
• Adaptation du style {style}
• Cohérence visuelle assurée
• Éléments graphiques optimisés

💡 Points d'attention :
• Contraste et lisibilité
• Accessibilité des éléments
• Harmonie des couleurs
• Impact visuel maximal

♿ Vérifications :
• Test de lisibilité
• Contraste suffisant
• Taille des textes adaptée
• Éléments bien espacés
    """
    
    return generated_image, recommendations

# Interface Gradio
demo = gr.Interface(
    fn=process_request,
    inputs=[
        gr.Dropdown(
            choices=["Événement", "Présentation", "Infographie", "Publicité"],
            label="Type d'affiche"
        ),
        gr.Textbox(
            label="Style souhaité",
            placeholder="Ex: moderne, artistique, minimaliste..."
        ),
        gr.Textbox(
            label="Message principal",
            placeholder="Le message clé de votre affiche"
        ),
        gr.Dropdown(
            choices=["Grand public", "Professionnel", "Jeune public", "Seniors"],
            label="Public cible"
        ),
        gr.Radio(
            choices=["Limité", "Moyen", "Élevé"],
            label="Budget"
        ),
        gr.Textbox(
            label="Thème artistique (optionnel)",
            placeholder="Ex: cyberpunk, nature, abstrait..."
        )
    ],
    outputs=[
        gr.Image(label="Affiche générée", type="pil"),
        gr.Textbox(label="Recommandations")
    ],
    title="🎨 Assistant Equity - Générateur d'affiches IA",
    description="""
    Générez des affiches professionnelles avec l'aide de l'IA.
    Spécifiez vos préférences et obtenez une affiche unique générée par intelligence artificielle.
    """,
    theme="soft",
    examples=[
        ["Événement", "moderne et artistique", "Festival de Jazz", "Grand public", "Moyen", "musical vibrant"],
        ["Publicité", "minimaliste", "Éco-responsabilité", "Professionnel", "Élevé", "nature épurée"],
        ["Infographie", "tech moderne", "Innovation digitale", "Jeune public", "Limité", "cyberpunk minimal"]
    ]
)

if __name__ == "__main__":
    demo.launch()