Update app.py

app.py

@@ -1,347 +1,249 @@
 import gradio as gr
+import torch
+from diffusers import StableDiffusionPipeline
+import gc
 import os
 from PIL import Image
-import requests
-import io
-import gc
+import numpy as np
+from dataclasses import dataclass
+from typing import Optional, Dict, Any
 import json
-from typing import Tuple, Optional, Dict, Any
-import logging
-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 étendus
-ART_STYLES = {
-    "Art Moderne": {
-        "prompt_prefix": "modern art style poster, professional design",
-        "negative_prompt": "traditional, photorealistic, cluttered, busy design"
-    },
-    "Neo Vintage": {
-        "prompt_prefix": "vintage style advertising poster, retro design",
-        "negative_prompt": "modern, digital, contemporary style"
-    },
-    "Pop Art": {
-        "prompt_prefix": "pop art style poster, bold design",
-        "negative_prompt": "subtle, realistic, traditional art"
-    },
-    "Minimaliste": {
-        "prompt_prefix": "minimalist design poster, clean composition",
-        "negative_prompt": "complex, detailed, ornate, busy"
-    },
-    "Cyberpunk": {
-        "prompt_prefix": "cyberpunk style poster, neon lights, futuristic design",
-        "negative_prompt": "vintage, natural, rustic, traditional"
-    },
-    "Aquarelle": {
-        "prompt_prefix": "watercolor art style poster, fluid artistic design",
-        "negative_prompt": "digital, sharp, photorealistic"
-    },
-    "Art Déco": {
-        "prompt_prefix": "art deco style poster, geometric patterns, luxury design",
-        "negative_prompt": "modern, minimalist, casual"
-    },
-    "Japonais": {
-        "prompt_prefix": "japanese art style poster, ukiyo-e inspired design",
-        "negative_prompt": "western, modern, photographic"
-    }
-}
-
-# Paramètres de composition
-COMPOSITION_PARAMS = {
-    "Layouts": {
-        "Centré": "centered composition, balanced layout",
-        "Asymétrique": "dynamic asymmetrical composition",
-        "Grille": "grid-based layout, structured composition",
-        "Diagonal": "diagonal dynamic composition",
-        "Minimaliste": "minimal composition, lots of whitespace"
-    },
-    "Ambiances": {
-        "Dramatique": "dramatic lighting, high contrast",
-        "Doux": "soft lighting, gentle atmosphere",
-        "Vibrant": "vibrant colors, energetic mood",
-        "Mystérieux": "mysterious atmosphere, moody lighting",
-        "Serein": "peaceful atmosphere, calm mood"
-    },
-    "Palette": {
-        "Monochrome": "monochromatic color scheme",
-        "Contrasté": "high contrast color palette",
-        "Pastel": "soft pastel color palette",
-        "Terre": "earthy color palette",
-        "Néon": "neon color palette"
-    }
-}
-
-class ImageGenerator:
+import time
+
+@dataclass
+class GenerationParams:
+    prompt: str
+    style: str = "realistic"
+    steps: int = 20
+    guidance_scale: float = 7.0
+    seed: int = -1
+    quality: str = "balanced"
+    
+class GenerartSystem:
     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 _build_prompt(self, params: Dict[str, Any]) -> str:
-        """Construction de prompt améliorée"""
-        style_info = ART_STYLES.get(params["style"], ART_STYLES["Neo Vintage"])
-        prompt = f"{style_info['prompt_prefix']}, {params['subject']}"
-
-        # Ajout des paramètres de composition
-        if params.get("layout"):
-            prompt += f", {COMPOSITION_PARAMS['Layouts'][params['layout']]}"
-        if params.get("ambiance"):
-            prompt += f", {COMPOSITION_PARAMS['Ambiances'][params['ambiance']]}"
-        if params.get("palette"):
-            prompt += f", {COMPOSITION_PARAMS['Palette'][params['palette']]}"
-
-        # Ajout des ajustements fins
-        if params.get("detail_level"):
-            detail_strength = params["detail_level"]
-            prompt += f", {'highly detailed' if detail_strength > 7 else 'moderately detailed'}"
-
-        if params.get("contrast"):
-            contrast_strength = params["contrast"]
-            prompt += f", {'high contrast' if contrast_strength > 7 else 'balanced contrast'}"
-
-        if params.get("saturation"):
-            saturation_strength = params["saturation"]
-            prompt += f", {'vibrant colors' if saturation_strength > 7 else 'subtle colors'}"
-
-        if params.get("title"):
-            prompt += f", with text saying '{params['title']}'"
-
-        logger.debug(f"Prompt final: {prompt}")
-        return prompt
-
-    def generate(self, params: Dict[str, Any]) -> Tuple[Optional[Image.Image], str]:
+        self.model = None
+        self.styles = {
+            "realistic": {
+                "prompt_prefix": "professional photography, highly detailed, photorealistic quality",
+                "negative_prompt": "cartoon, anime, illustration, painting, drawing, blurry, low quality",
+                "params": {"guidance_scale": 7.5, "steps": 20}
+            },
+            "artistic": {
+                "prompt_prefix": "artistic painting, impressionist style, vibrant colors",
+                "negative_prompt": "photorealistic, digital art, 3d render, low quality",
+                "params": {"guidance_scale": 6.5, "steps": 25}
+            },
+            "modern": {
+                "prompt_prefix": "modern art, contemporary style, abstract qualities",
+                "negative_prompt": "traditional, classic, photorealistic, low quality",
+                "params": {"guidance_scale": 8.0, "steps": 15}
+            }
+        }
+        self.quality_presets = {
+            "speed": {"steps_multiplier": 0.8},
+            "balanced": {"steps_multiplier": 1.0},
+            "quality": {"steps_multiplier": 1.2}
+        }
+        self.performance_stats = {
+            "total_generations": 0,
+            "average_time": 0,
+            "success_rate": 100,
+            "last_error": None
+        }
+        
+    def initialize_model(self):
+        """Initialize the model with memory optimizations"""
+        if self.model is not None:
+            return
+
+        # Memory cleanup before model load
+        gc.collect()
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
+        
         try:
-            logger.info(f"Début de génération avec paramètres: {json.dumps(params, indent=2)}")
+            self.model = StableDiffusionPipeline.from_pretrained(
+                "CompVis/stable-diffusion-v1-4",
+                torch_dtype=torch.float32,
+                safety_checker=None,
+                requires_safety_checker=False
+            )
             
-            if 'Bearer None' in self.headers['Authorization']:
-                return None, "⚠️ Erreur: Token Hugging Face non configuré"
-
-            prompt = self._build_prompt(params)
+            # Memory optimizations
+            self.model.enable_attention_slicing()
+            self.model.enable_vae_slicing()
             
-            # Configuration de base
-            payload = {
-                "inputs": prompt,
-                "parameters": {
-                    "negative_prompt": ART_STYLES[params["style"]]["negative_prompt"],
-                    "num_inference_steps": min(int(35 * (params["quality"]/100)), 40),
-                    "guidance_scale": min(7.5 * (params["creativity"]/10), 10.0),
-                    "width": 768,
-                    "height": 768 if params["orientation"] == "Portrait" else 512
-                }
-            }
+            # Move to CPU - system doesn't have adequate GPU
+            self.model = self.model.to("cpu")
             
-            logger.debug(f"Payload: {json.dumps(payload, indent=2)}")
+        except Exception as e:
+            print(f"Error initializing model: {str(e)}")
+            raise
 
-            response = requests.post(
-                self.API_URL,
-                headers=self.headers,
-                json=payload,
-                timeout=30
-            )
+    def cleanup(self):
+        """Memory cleanup after generation"""
+        gc.collect()
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
+
+    def update_performance_stats(self, generation_time: float, success: bool = True, error: Optional[str] = None):
+        """Update system performance statistics"""
+        self.performance_stats["total_generations"] += 1
+        
+        # Update average time
+        prev_avg = self.performance_stats["average_time"]
+        self.performance_stats["average_time"] = (prev_avg * (self.performance_stats["total_generations"] - 1) + 
+                                                generation_time) / self.performance_stats["total_generations"]
+        
+        # Update success rate
+        if not success:
+            self.performance_stats["success_rate"] = (self.performance_stats["success_rate"] * 
+                                                    (self.performance_stats["total_generations"] - 1) + 
+                                                    0) / self.performance_stats["total_generations"]
+            self.performance_stats["last_error"] = error
+            
+    def get_system_stats(self):
+        """Get current system statistics"""
+        return {
+            "total_generations": self.performance_stats["total_generations"],
+            "average_time": round(self.performance_stats["average_time"], 2),
+            "success_rate": round(self.performance_stats["success_rate"], 1),
+            "memory_usage": f"{torch.cuda.memory_allocated()/1024**2:.1f}MB" if torch.cuda.is_available() 
+                          else "CPU Mode"
+        }
+
+    def generate_image(self, params: GenerationParams) -> Image.Image:
+        """Generate image with given parameters"""
+        try:
+            # Initialize model if needed
+            if self.model is None:
+                self.initialize_model()
+                
+            # Prepare generation parameters
+            style_config = self.styles[params.style]
+            quality_config = self.quality_presets[params.quality]
+            
+            # Construct final prompt
+            full_prompt = f"{style_config['prompt_prefix']}, {params.prompt}"
             
-            if response.status_code == 200:
-                image = Image.open(io.BytesIO(response.content))
-                return image, "✨ Création réussie!"
+            # Calculate final steps
+            final_steps = int(min(25, params.steps * quality_config["steps_multiplier"]))
+            
+            # Set random seed if needed
+            if params.seed == -1:
+                generator = None
             else:
-                error_msg = f"⚠️ Erreur API {response.status_code}: {response.text}"
-                logger.error(error_msg)
-                return None, error_msg
-
+                generator = torch.manual_seed(params.seed)
+                
+            start_time = time.time()
+            
+            # Generate image
+            with torch.no_grad():
+                image = self.model(
+                    prompt=full_prompt,
+                    negative_prompt=style_config["negative_prompt"],
+                    num_inference_steps=final_steps,
+                    guidance_scale=params.guidance_scale,
+                    generator=generator,
+                    width=512,
+                    height=512
+                ).images[0]
+                
+            generation_time = time.time() - start_time
+            self.update_performance_stats(generation_time, success=True)
+            
+            return image
+            
         except Exception as e:
-            error_msg = f"⚠️ Erreur: {str(e)}"
-            logger.exception("Erreur pendant la génération:")
-            return None, error_msg
+            self.update_performance_stats(0, success=False, error=str(e))
+            raise RuntimeError(f"Generation error: {str(e)}")
+        
         finally:
-            gc.collect()
+            self.cleanup()
 
-def create_interface():
-    logger.info("Création de l'interface Gradio")
-    
-    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; }
-    .advanced-controls { background: #374151; padding: 12px; border-radius: 5px; margin: 8px 0; }
-    """
-    
-    generator = ImageGenerator()
-    
-    with gr.Blocks(css=css) as app:
-        gr.HTML("""
-            <div class="welcome">
-                <h1>🎨 Equity Artisan 3.0</h1>
-                <p>Assistant de création d'affiches professionnelles</p>
-            </div>
-        """)
+class GenerartInterface:
+    def __init__(self):
+        self.system = GenerartSystem()
         
-        with gr.Column(elem_classes="container"):
-            # Format et Orientation
-            with gr.Group(elem_classes="controls-group"):
-                gr.Markdown("### 📐 Format et Orientation")
-                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"
-                    )
+    def create_interface(self):
+        """Create the Gradio interface"""
+        with gr.Blocks(theme=gr.themes.Soft()) as demo:
+            # Header
+            gr.Markdown("# 🎨 Generart Beta")
             
-            # Style et Composition
-            with gr.Group(elem_classes="controls-group"):
-                gr.Markdown("### 🎨 Style et Composition")
-                with gr.Row():
+            with gr.Row():
+                # Left column - Controls
+                with gr.Column(scale=1):
+                    prompt = gr.Textbox(label="Description", placeholder="Décrivez l'image souhaitée...")
+                    
                     style = gr.Dropdown(
-                        choices=list(ART_STYLES.keys()),
-                        value="Neo Vintage",
-                        label="Style artistique"
-                    )
-                    layout = gr.Dropdown(
-                        choices=list(COMPOSITION_PARAMS["Layouts"].keys()),
-                        value="Centré",
-                        label="Composition"
+                        choices=list(self.system.styles.keys()),
+                        value="realistic",
+                        label="Style Artistique"
                     )
-
-                with gr.Row():
-                    ambiance = gr.Dropdown(
-                        choices=list(COMPOSITION_PARAMS["Ambiances"].keys()),
-                        value="Dramatique",
-                        label="Ambiance"
-                    )
-                    palette = gr.Dropdown(
-                        choices=list(COMPOSITION_PARAMS["Palette"].keys()),
-                        value="Contrasté",
-                        label="Palette"
-                    )
-            
-            # Contenu
-            with gr.Group(elem_classes="controls-group"):
-                gr.Markdown("### 📝 Contenu")
-                subject = gr.Textbox(
-                    label="Description",
-                    placeholder="Décrivez votre vision..."
+                    
+                    with gr.Group():
+                        steps = gr.Slider(
+                            minimum=15,
+                            maximum=25,
+                            value=20,
+                            step=1,
+                            label="Nombre d'étapes"
+                        )
+                        
+                        guidance = gr.Slider(
+                            minimum=6.0,
+                            maximum=8.0,
+                            value=7.0,
+                            step=0.1,
+                            label="Guide Scale"
+                        )
+                        
+                        quality = gr.Dropdown(
+                            choices=list(self.system.quality_presets.keys()),
+                            value="balanced",
+                            label="Qualité"
+                        )
+                        
+                        seed = gr.Number(
+                            value=-1,
+                            label="Seed (-1 pour aléatoire)",
+                            precision=0
+                        )
+                        
+                    generate_btn = gr.Button("Générer", variant="primary")
+                    
+                    # System Stats
+                    with gr.Group():
+                        gr.Markdown("### 📊 Statistiques Système")
+                        stats_output = gr.JSON(value=self.system.get_system_stats())
+                
+                # Right column - Output
+                with gr.Column(scale=1):
+                    image_output = gr.Image(label="Image Générée", type="pil")
+                    
+            # Generation Event
+            def generate(prompt, style, steps, guidance_scale, quality, seed):
+                params = GenerationParams(
+                    prompt=prompt,
+                    style=style,
+                    steps=steps,
+                    guidance_scale=guidance_scale,
+                    quality=quality,
+                    seed=seed
                 )
-                title = gr.Textbox(
-                    label="Titre",
-                    placeholder="Titre de l'affiche..."
-                )
-            
-            # Ajustements fins
-            with gr.Group(elem_classes="advanced-controls"):
-                gr.Markdown("### 🎯 Ajustements Fins")
-                with gr.Row():
-                    detail_level = gr.Slider(
-                        minimum=1,
-                        maximum=10,
-                        value=7,
-                        step=1,
-                        label="Niveau de Détail"
-                    )
-                    contrast = gr.Slider(
-                        minimum=1,
-                        maximum=10,
-                        value=5,
-                        step=1,
-                        label="Contraste"
-                    )
-                    saturation = gr.Slider(
-                        minimum=1,
-                        maximum=10,
-                        value=5,
-                        step=1,
-                        label="Saturation"
-                    )
-            
-            # Paramètres de génération
-            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é"
-                    )
+                
+                image = self.system.generate_image(params)
+                return [image, self.system.get_system_stats()]
             
-            # Boutons
-            with gr.Row():
-                generate_btn = gr.Button("✨ Générer", variant="primary")
-                clear_btn = gr.Button("🗑️ Effacer", variant="secondary")
+            generate_btn.click(
+                fn=generate,
+                inputs=[prompt, style, steps, guidance, quality, seed],
+                outputs=[image_output, stats_output]
+            )
             
-            # Sortie
-            image_output = gr.Image(label="Aperçu")
-            status = gr.Textbox(label="Statut", interactive=False)
-        
-        def generate(*args):
-            logger.info("Démarrage d'une nouvelle génération")
-            params = {
-                "format_size": args[0],
-                "orientation": args[1],
-                "style": args[2],
-                "layout": args[3],
-                "ambiance": args[4],
-                "palette": args[5],
-                "subject": args[6],
-                "title": args[7],
-                "detail_level": args[8],
-                "contrast": args[9],
-                "saturation": args[10],
-                "quality": args[11],
-                "creativity": args[12]
-            }
-            result = generator.generate(params)
-            logger.info(f"Génération terminée avec statut: {result[1]}")
-            return result
-        
-        generate_btn.click(
-            generate,
-            inputs=[
-                format_size,
-                orientation,
-                style,
-                layout,
-                ambiance,
-                palette,
-                subject,
-                title,
-                detail_level,
-                contrast,
-                saturation,
-                quality,
-                creativity
-            ],
-            outputs=[image_output, status]
-        )
-        
-        clear_btn.click(
-            lambda: (None, "🗑️ Image effacée"),
-            outputs=[image_output, status]
-        )
-    
-    logger.info("Interface créée avec succès")
-    return app
+        return demo
 
+# Create and launch the interface
 if __name__ == "__main__":
-    app = create_interface()
-    logger.info("Démarrage de l'application")
-    app.launch()
+    interface = GenerartInterface()
+    demo = interface.create_interface()
+    demo.launch()