Update app.py

app.py

@@ -12,240 +12,206 @@ import numpy as np
 import cv2
 from skimage import exposure
 import torch
+import asyncio
 
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 load_dotenv()
 
-def load_art_styles():
-    return {
-        "Styles Traditionnels": {
-            "Renaissance": {"prompt": "renaissance masterpiece, anatomical precision, detailed texture, chiaroscuro lighting", "params": {"resolution": (4096, 4096), "detail_level": 0.95}},
-            "Impressionnisme": {"prompt": "impressionist style painting, visible brushstrokes, natural light", "params": {"resolution": (3072, 3072), "noise_level": 0.3}},
-            "Surréalisme": {"prompt": "surrealist dreamlike scene, symbolic elements, subconscious imagery", "params": {"randomization": 0.4}},
-            "Cubisme": {"prompt": "cubist style, geometric forms, multiple perspectives", "params": {"geometric_strength": 0.8}}
+STYLES = {
+    "Styles Traditionnels": {
+        "Renaissance": {
+            "prompt": "renaissance masterpiece, anatomical precision, detailed texture, chiaroscuro lighting, oil painting technique, museum quality",
+            "negative_prompt": "modern, abstract, simple, digital",
+            "params": {"resolution": 4096, "detail_level": 0.95}
         },
-        "Rendus Numériques": {
-            "Synthwave": {"prompt": "synthwave aesthetic, neon grid, retro-futuristic, vibrant", "params": {"saturation": 1.8, "neon": True}},
-            "Cyberpunk": {"prompt": "cyberpunk style, neon-lit, high-tech, volumetric lighting", "params": {"volumetric": True, "neon": True}},
-            "Sci-Fi": {"prompt": "sci-fi environment, futuristic technology, advanced architecture", "params": {"tech_level": 0.9}}
+        "Impressionnisme": {
+            "prompt": "impressionist style, visible brushstrokes, natural light, plein air painting",
+            "negative_prompt": "sharp, digital, modern",
+            "params": {"resolution": 3072, "noise_level": 0.3}
         },
-        "Photographie": {
-            "HDR": {"prompt": "HDR photography, extreme dynamic range, detailed shadows and highlights", "params": {"hdr_strength": 1.5}},
-            "Macro": {"prompt": "macro photography, extreme close-up, fine details", "params": {"detail_scale": 0.5}},
-            "Portrait": {"prompt": "professional portrait photography, studio lighting, bokeh", "params": {"bokeh": 0.7}},
-            "Vintage": {"prompt": "vintage photography, aged effect, retro colors", "params": {"grain": 0.4}}
+        "Surréalisme": {
+            "prompt": "surrealist dreamlike scene, symbolic imagery, unconscious imagination",
+            "negative_prompt": "realistic, ordinary",
+            "params": {"resolution": 3072, "randomization": 0.4}
         },
-        "Illustration": {
-            "Aquarelle": {"prompt": "watercolor illustration, fluid transparency, soft edges", "params": {"transparency": 0.6}},
-            "Encre": {"prompt": "ink drawing, bold strokes, high contrast", "params": {"contrast": 1.4}},
-            "Huile": {"prompt": "oil painting, thick impasto, rich colors", "params": {"texture": 0.8}}
-        },
-        "Photoréalisme": {
-            "Nature_Morte": {"prompt": "photorealistic still life, extreme detail, perfect lighting", "params": {"detail_level": 0.98}},
-            "Paysage": {"prompt": "photorealistic landscape, natural lighting, atmospheric", "params": {"atmosphere": 0.7}}
-        },
-        "Fantasy": {
-            "Fantasy": {"prompt": "fantasy art, magical atmosphere, mythical elements", "params": {"magic_effect": 0.8}},
-            "Dark_Fantasy": {"prompt": "dark fantasy, gothic elements, mysterious atmosphere", "params": {"darkness": 0.7}},
-            "Steampunk": {"prompt": "steampunk style, brass and copper, mechanical elements", "params": {"mechanical": 0.9}}
-        },
-        "Abstrait": {
-            "Holographique": {"prompt": "holographic effect, iridescent colors, light refraction", "params": {"iridescence": 0.8}},
-            "Fractal": {"prompt": "fractal art, recursive patterns, mathematical beauty", "params": {"complexity": 0.9}}
+        "Cubisme": {
+            "prompt": "cubist style, geometric forms, multiple perspectives, abstract interpretation",
+            "negative_prompt": "realistic, photographic",
+            "params": {"resolution": 2048, "geometric_strength": 0.8}
+        }
+    },
+    "Rendus Numériques": {
+        "Synthwave": {
+            "prompt": "synthwave aesthetic, neon grid, retro-futuristic, dramatic lighting, vibrant colors",
+            "negative_prompt": "natural, realistic, muted colors",
+            "params": {"saturation": 1.8, "neon_effect": True}
         },
-        "Graphisme": {
-            "Flat": {"prompt": "flat design, minimal shapes, solid colors", "params": {"simplification": 0.8}},
-            "Material": {"prompt": "material design, subtle shadows, layered elements", "params": {"layers": 0.6}},
-            "Isométrique": {"prompt": "isometric design, geometric precision, clean lines", "params": {"precision": 0.9}}
+        "Cyberpunk": {
+            "prompt": "cyberpunk style, neon-lit cityscape, high-tech, volumetric fog, futuristic",
+            "negative_prompt": "natural, vintage, rural",
+            "params": {"volumetric": True, "neon_intensity": 1.5}
+        }
+    },
+    "Photographie": {
+        "HDR": {
+            "prompt": "HDR photography, extreme dynamic range, rich details in shadows and highlights",
+            "negative_prompt": "flat lighting, low contrast",
+            "params": {"hdr_strength": 1.5, "resolution": 4096}
         },
-        "Gaming": {
-            "Pixel_Art": {"prompt": "pixel art style, retro gaming aesthetic, limited palette", "params": {"pixelation": 0.7}},
-            "Cel_Shading": {"prompt": "cel shaded style, anime-like, bold outlines", "params": {"outline": 0.8}}
+        "Portrait Studio": {
+            "prompt": "professional studio portrait, perfect lighting, bokeh effect, sharp focus",
+            "negative_prompt": "blurry, noisy",
+            "params": {"bokeh": 0.7, "sharpness": 1.4}
         }
     }
+}
 
 class ImageProcessor:
     def __init__(self):
-        self.effects = {
-            "neon": self._apply_neon,
-            "bokeh": self._apply_bokeh,
-            "grain": self._apply_grain,
-            "hdr": self._apply_hdr,
-            "pixelation": self._apply_pixelation
-        }
-
-    def process_image(self, image: Image.Image, style_params: Dict) -> Image.Image:
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    def process_image(self, image: Image.Image, params: Dict) -> Image.Image:
         try:
+            # Conversion en array numpy
             img_array = np.array(image)
             
-            for effect, value in style_params.items():
-                if effect in self.effects and value:
-                    img_array = self.effects[effect](img_array, value)
-                    
+            # Application des effets selon les paramètres
+            if params.get("neon_effect"):
+                img_array = self._apply_neon_effect(img_array)
+                
+            if params.get("volumetric"):
+                img_array = self._apply_volumetric(img_array)
+                
+            if params.get("bokeh"):
+                img_array = self._apply_bokeh(img_array, params["bokeh"])
+                
+            if params.get("hdr_strength"):
+                img_array = self._apply_hdr(img_array, params["hdr_strength"])
+            
             return Image.fromarray(img_array)
+            
         except Exception as e:
             logger.error(f"Erreur traitement: {str(e)}")
             return image
-
-    def _apply_neon(self, image: np.ndarray, strength: float) -> np.ndarray:
+            
+    def _apply_neon_effect(self, image: np.ndarray) -> np.ndarray:
         hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
-        hsv[..., 1] = np.clip(hsv[..., 1] * strength, 0, 255)
-        glow = cv2.GaussianBlur(cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB), (0, 0), 15)
-        return cv2.addWeighted(image, 1, glow, 0.5, 0)
-
+        hsv[..., 1] = np.clip(hsv[..., 1] * 1.5, 0, 255)  # Boost saturation
+        rgb = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
+        # Ajout de glow
+        blur = cv2.GaussianBlur(rgb, (0, 0), 15)
+        return cv2.addWeighted(rgb, 1, blur, 0.5, 0)
+        
+    def _apply_volumetric(self, image: np.ndarray) -> np.ndarray:
+        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+        fog = cv2.GaussianBlur(gray, (0, 0), 20)
+        return cv2.addWeighted(image, 1, cv2.cvtColor(fog, cv2.COLOR_GRAY2RGB), 0.2, 0)
+        
     def _apply_bokeh(self, image: np.ndarray, strength: float) -> np.ndarray:
         blur = cv2.GaussianBlur(image, (0, 0), int(30 * strength))
         mask = np.random.random(image.shape[:2]) > 0.5
         result = image.copy()
         result[mask] = blur[mask]
         return result
-
-    def _apply_grain(self, image: np.ndarray, strength: float) -> np.ndarray:
-        noise = np.random.normal(0, strength * 50, image.shape).astype(np.uint8)
-        return np.clip(image + noise, 0, 255)
-
+        
     def _apply_hdr(self, image: np.ndarray, strength: float) -> np.ndarray:
         return exposure.adjust_gamma(image, 1.0 / strength)
 
-    def _apply_pixelation(self, image: np.ndarray, strength: float) -> np.ndarray:
-        h, w = image.shape[:2]
-        size = int(max(h, w) * (1 - strength))
-        small = cv2.resize(image, (size, size), interpolation=cv2.INTER_LINEAR)
-        return cv2.resize(small, (w, h), interpolation=cv2.INTER_NEAREST)
-
 class ImageGenerator:
     def __init__(self):
         self.processor = ImageProcessor()
         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')}"}
-        self.styles = load_art_styles()
-
-    async def generate(self, prompt: str, style_category: str, style_name: str) -> Tuple[Optional[Image.Image], str]:
+        token = os.getenv('HUGGINGFACE_TOKEN')
+        if not token:
+            raise ValueError("HUGGINGFACE_TOKEN non trouvé dans les variables d'environnement!")
+        self.headers = {"Authorization": f"Bearer {token}"}
+        
+    def generate(self, prompt: str, style_category: str, style_name: str) -> Tuple[Optional[Image.Image], str]:
         try:
-            style_info = self.styles[style_category][style_name]
+            # Récupération du style
+            style_info = STYLES[style_category][style_name]
+            
+            # Construction du prompt final
             final_prompt = f"{prompt}, {style_info['prompt']}"
             
+            # Paramètres de génération
+            params = {
+                "inputs": final_prompt,
+                "negative_prompt": style_info["negative_prompt"],
+                "num_inference_steps": 50,
+                "guidance_scale": 7.5,
+                "width": style_info["params"].get("resolution", 1024),
+                "height": style_info["params"].get("resolution", 1024),
+            }
+            
+            # Appel API
             response = requests.post(
                 self.api_url,
                 headers=self.headers,
-                json={"inputs": final_prompt},
+                json=params,
                 timeout=30
             )
             
             if response.status_code != 200:
                 return None, f"Erreur API: {response.status_code}"
-
+                
+            # Traitement de l'image
             image = Image.open(io.BytesIO(response.content))
-            processed = self.processor.process_image(image, style_info['params'])
+            processed = self.processor.process_image(image, style_info["params"])
             
             return processed, "✨ Génération réussie!"
+            
         except Exception as e:
-            logger.error(f"Erreur: {str(e)}")
+            logger.error(f"Erreur génération: {str(e)}")
             return None, f"Erreur: {str(e)}"
         finally:
             gc.collect()
-            def create_interface():
-    generator = ImageGenerator()
 
+def create_interface():
+    generator = ImageGenerator()
+    
     with gr.Blocks() as demo:
         gr.HTML("""
             <div style='text-align: center; margin-bottom: 1rem'>
                 <h1>🎨 Equity Art Engine</h1>
-                <p>Générateur d'Images Avancé avec Styles Artistiques</p>
+                <p>Générateur d'Images avec Styles Artistiques</p>
             </div>
         """)
         
         with gr.Row():
-            with gr.Column(scale=1):
-                prompt = gr.Textbox(
-                    label="Description",
-                    placeholder="Décrivez votre image..."
-                )
+            with gr.Column():
+                prompt = gr.Textbox(label="Description", placeholder="Décrivez votre image...")
                 
                 style_category = gr.Dropdown(
-                    choices=list(generator.styles.keys()),
-                    label="Catégorie de Style",
-                    value="Styles Traditionnels"
+                    choices=list(STYLES.keys()),
+                    label="Catégorie de Style"
                 )
                 
-                style_name = gr.Dropdown(
-                    label="Style Spécifique"
-                )
+                style_name = gr.Dropdown(label="Style Spécifique")
                 
+                # Mise à jour dynamique des styles
                 def update_styles(category):
-                    return gr.Dropdown.update(
-                        choices=list(generator.styles[category].keys()) if category else []
-                    )
-                
-                style_category.change(
-                    update_styles,
-                    inputs=style_category,
-                    outputs=style_name
-                )
-                
-                advanced_params = gr.Checkbox(
-                    label="Paramètres avancés",
-                    value=False
-                )
-                
-                with gr.Column(visible=False) as advanced_options:
-                    quality = gr.Slider(
-                        minimum=1,
-                        maximum=10,
-                        value=7,
-                        step=1,
-                        label="Qualité"
-                    )
+                    return gr.Dropdown.update(choices=list(STYLES[category].keys()) if category else [])
                     
-                    seed = gr.Number(
-                        label="Seed (-1 pour aléatoire)",
-                        value=-1
-                    )
-                
-                def toggle_advanced(show):
-                    return gr.update(visible=show)
-                
-                advanced_params.change(
-                    toggle_advanced,
-                    inputs=advanced_params,
-                    outputs=advanced_options
-                )
+                style_category.change(update_styles, inputs=style_category, outputs=style_name)
                 
-                generate_btn = gr.Button("✨ Générer", variant="primary")
+                generate_btn = gr.Button("✨ Générer")
                 
-            with gr.Column(scale=2):
-                output_image = gr.Image(label="Image Générée")
-                status = gr.Textbox(label="Status")
+            with gr.Column():
+                image_output = gr.Image(label="Résultat")
+                status_output = gr.Textbox(label="Status")
         
-        def generate_image(prompt, category, style, use_advanced, quality, seed):
+        def generate(prompt, category, style):
             if not all([prompt, category, style]):
                 return None, "⚠️ Veuillez remplir tous les champs requis"
             
-            params = {
-                "quality": quality if use_advanced else 7,
-                "seed": seed if use_advanced and seed != -1 else None
-            }
-            
-            image, status_msg = generator.generate(
-                prompt=prompt,
-                style_category=category,
-                style_name=style
-            )
-            
-            return image, status_msg
+            return generator.generate(prompt, category, style)
         
         generate_btn.click(
-            generate_image,
-            inputs=[
-                prompt,
-                style_category,
-                style_name,
-                advanced_params,
-                quality,
-                seed
-            ],
-            outputs=[output_image, status]
+            generate,
+            inputs=[prompt, style_category, style_name],
+            outputs=[image_output, status_output]
         )
     
     return demo