Update app.py

app.py

@@ -1,249 +1,117 @@
 import gradio as gr
-import torch
-from diffusers import StableDiffusionPipeline
-import gc
 import os
 from PIL import Image
-import numpy as np
-from dataclasses import dataclass
-from typing import Optional, Dict, Any
+import requests
+import io
+import gc
 import json
-import time
+from typing import Tuple, Optional, Dict, Any
+import logging
+from dotenv import load_dotenv
 
-@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.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
+# Configuration du logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-        # Memory cleanup before model load
-        gc.collect()
-        torch.cuda.empty_cache() if torch.cuda.is_available() else None
-        
-        try:
-            self.model = StableDiffusionPipeline.from_pretrained(
-                "CompVis/stable-diffusion-v1-4",
-                torch_dtype=torch.float32,
-                safety_checker=None,
-                requires_safety_checker=False
-            )
-            
-            # Memory optimizations
-            self.model.enable_attention_slicing()
-            self.model.enable_vae_slicing()
-            
-            # Move to CPU - system doesn't have adequate GPU
-            self.model = self.model.to("cpu")
-            
-        except Exception as e:
-            print(f"Error initializing model: {str(e)}")
-            raise
+# Chargement des variables d'environnement
+load_dotenv()
 
-    def cleanup(self):
-        """Memory cleanup after generation"""
-        gc.collect()
-        torch.cuda.empty_cache() if torch.cuda.is_available() else None
+# Styles artistiques
+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"
+    }
+}
 
-    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
+# Configuration de l'API
+API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
+
+def generate_image(params: Dict[str, Any]) -> Tuple[Optional[Image.Image], str]:
+    """Génère une image via l'API Hugging Face"""
+    try:
+        headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACE_TOKEN')}"}
         
-        # 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"]
+        style = ART_STYLES[params["style"]]
+        prompt = f"{style['prompt_prefix']}, {params['subject']}"
         
-        # 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"
+        # Configuration de la requête
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "negative_prompt": style["negative_prompt"],
+                "num_inference_steps": 30,
+                "guidance_scale": 7.5,
+                "width": 768,
+                "height": 768
+            }
         }
-
-    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}"
-            
-            # 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:
-                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:
-            self.update_performance_stats(0, success=False, error=str(e))
-            raise RuntimeError(f"Generation error: {str(e)}")
         
-        finally:
-            self.cleanup()
-
-class GenerartInterface:
-    def __init__(self):
-        self.system = GenerartSystem()
+        response = requests.post(API_URL, headers=headers, json=payload)
         
-    def create_interface(self):
-        """Create the Gradio interface"""
-        with gr.Blocks(theme=gr.themes.Soft()) as demo:
-            # Header
-            gr.Markdown("# 🎨 Generart Beta")
+        if response.status_code == 200:
+            image = Image.open(io.BytesIO(response.content))
+            return image, "✨ Création réussie!"
+        else:
+            return None, f"⚠️ Erreur {response.status_code}: {response.text}"
             
-            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(self.system.styles.keys()),
-                        value="realistic",
-                        label="Style Artistique"
-                    )
-                    
-                    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())
+    except Exception as e:
+        logger.error(f"Erreur: {str(e)}")
+        return None, f"⚠️ Erreur: {str(e)}"
+
+def create_interface():
+    """Crée l'interface Gradio"""
+    with gr.Blocks() as app:
+        gr.HTML("""
+            <h1 style='text-align: center'>🎨 Generart</h1>
+            <p style='text-align: center'>Créez des affiches artistiques avec l'IA</p>
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                style = gr.Dropdown(
+                    choices=list(ART_STYLES.keys()),
+                    value="Neo Vintage",
+                    label="Style artistique"
+                )
                 
-                # 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
+                subject = gr.Textbox(
+                    label="Description",
+                    placeholder="Décrivez votre vision..."
                 )
                 
-                image = self.system.generate_image(params)
-                return [image, self.system.get_system_stats()]
-            
-            generate_btn.click(
-                fn=generate,
-                inputs=[prompt, style, steps, guidance, quality, seed],
-                outputs=[image_output, stats_output]
-            )
-            
-        return demo
+                generate_btn = gr.Button("✨ Générer")
+                
+            with gr.Column():
+                image_output = gr.Image(label="Résultat")
+                status = gr.Textbox(label="Statut")
+        
+        def on_generate(style_val, subject_val):
+            return generate_image({
+                "style": style_val,
+                "subject": subject_val
+            })
+        
+        generate_btn.click(
+            fn=on_generate,
+            inputs=[style, subject],
+            outputs=[image_output, status]
+        )
+    
+    return app
 
-# Create and launch the interface
 if __name__ == "__main__":
-    interface = GenerartInterface()
-    demo = interface.create_interface()
-    demo.launch()
+    app = create_interface()
+    app.launch()