Update app.py

app.py

@@ -1,61 +1,122 @@
 import gradio as gr
+from diffusers import StableDiffusionPipeline
+import torch
+from PIL import Image
+import os
 
-def generate_recommendations(poster_type, style, message):
-    """
-    Génère des recommandations pour la création d'affiches
-    """
-    # Dictionnaire des layouts recommandés
-    layouts = {
-        "Événement": "Format vertical, titre en haut, date et lieu mis en évidence",
-        "Présentation": "Structure claire, points clés hiérarchisés",
-        "Infographie": "Organisation en sections, utilisation de graphiques",
-        "Autre": "Layout flexible selon vos besoins"
-    }
-
-    # Récupérer le layout recommandé
-    layout = layouts.get(poster_type, layouts["Autre"])
-
-    # Construire la réponse
-    response = f"""
+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 :
 
-📌 Type : {poster_type}
-🎯 Style demandé : {style}
-💭 Message : {message}
+📌 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
 
-✨ Mise en page suggérée :
-{layout}
+💡 Points d'attention :
+• Contraste et lisibilité
+• Accessibilité des éléments
+• Harmonie des couleurs
+• Impact visuel maximal
 
-💡 Conseils :
-• Utilisez des contrastes adaptés
-• Choisissez des polices lisibles
-• Vérifiez l'accessibilité des couleurs
-• Gardez une hiérarchie visuelle claire
-"""
-    return response
+♿ Vérifications :
+• Test de lisibilité
+• Contraste suffisant
+• Taille des textes adaptée
+• Éléments bien espacés
+    """
+    
+    return generated_image, recommendations
 
-# Création de l'interface
-iface = gr.Interface(
-    fn=generate_recommendations,
+# Interface Gradio
+demo = gr.Interface(
+    fn=process_request,
     inputs=[
         gr.Dropdown(
-            choices=["Événement", "Présentation", "Infographie", "Autre"],
+            choices=["Événement", "Présentation", "Infographie", "Publicité"],
             label="Type d'affiche"
         ),
         gr.Textbox(
             label="Style souhaité",
-            placeholder="Ex: moderne, minimaliste, coloré..."
+            placeholder="Ex: moderne, artistique, minimaliste..."
         ),
         gr.Textbox(
             label="Message principal",
-            placeholder="Ex: Annonce d'événement, présentation..."
+            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.Textbox(label="Recommandations"),
-    title="🎨 Assistant Equity",
-    description="Créez des affiches équitables et accessibles"
+    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"]
+    ]
 )
 
-# Lancement de l'application
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()