Update app.py

app.py

@@ -10,6 +10,8 @@ import logging
 import time
 from typing import Optional, Tuple
 import warnings
+import requests
+import json
 warnings.filterwarnings("ignore")
 
 logging.basicConfig(level=logging.INFO)
@@ -23,6 +25,33 @@ MAX_IMAGE_SIZE = 1024
 RATE_LIMIT_DELAY = 3
 API_RETRY_COUNT = 3
 
+IMAGE_MODELS = {
+    "Freepik Gemini 2.5 Flash": {
+        "name": "Freepik Gemini 2.5 Flash Image Preview",
+        "api": "freepik",
+        "url": "https://api.freepik.com/v1/ai/gemini-2-5-flash-image-preview",
+        "description": "🍌 Best for competition - Real Gemini 2.5 Flash generation"
+    },
+    "OpenAI DALL-E 3": {
+        "name": "DALL-E 3",
+        "api": "openai",
+        "url": "https://api.openai.com/v1/images/generations",
+        "description": "🎨 High-quality creative image generation"
+    },
+    "Stable Diffusion XL": {
+        "name": "Stable Diffusion XL",
+        "api": "stabilityai",
+        "url": "https://api.stability.ai/v1/generation/stable-diffusion-xl-1024-v1-0/text-to-image",
+        "description": "🚀 Open-source high-resolution generation"
+    },
+    "Gemini Analysis + Demo": {
+        "name": "Gemini Analysis + Visual Demo",
+        "api": "gemini",
+        "url": None,
+        "description": "📊 Analysis + visual overlay (fallback mode)"
+    }
+}
+
 if GEMINI_API_KEY:
     genai.configure(api_key=GEMINI_API_KEY)
     logger.info("Gemini API configured")
@@ -77,28 +106,34 @@ class NanoBananaApp:
     def _apply_rate_limiting(self):
         time.sleep(RATE_LIMIT_DELAY)
         
-    def _generate_with_freepik(self, image, prompt, style, editing_mode):
-        if not FREEPIK_API_KEY:
-            return None, "Freepik API key not available"
+    def _generate_with_model(self, image, prompt, style, editing_mode, model_name, api_keys):
+        model_info = IMAGE_MODELS.get(model_name)
+        if not model_info:
+            return None, f"Unknown model: {model_name}"
+            
+        api_type = model_info["api"]
+        
+        if api_type == "freepik":
+            return self._generate_with_freepik(image, prompt, style, editing_mode, api_keys.get("freepik"))
+        elif api_type == "openai":
+            return self._generate_with_dalle(image, prompt, style, editing_mode, api_keys.get("openai"))
+        elif api_type == "stabilityai":
+            return self._generate_with_stable_diffusion(image, prompt, style, editing_mode, api_keys.get("stabilityai"))
+        elif api_type == "gemini":
+            return self._generate_with_gemini_demo(image, prompt, style, editing_mode, api_keys.get("gemini"))
+        else:
+            return None, f"Unsupported API type: {api_type}"
+    
+    def _generate_with_freepik(self, image, prompt, style, editing_mode, api_key):
+        if not api_key:
+            return None, "Freepik API key not provided"
             
         try:
             buffered = io.BytesIO()
             image.save(buffered, format='PNG')
             image_b64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
             
-            style_modifiers = {
-                "realistic": "photorealistic, high-quality construction, professional architecture",
-                "futuristic": "futuristic, high-tech, modern glass and steel, sci-fi architecture",
-                "artistic": "artistic, creative design, unique architecture, colorful and innovative"
-            }
-            
-            mode_descriptions = {
-                "complete": "Complete this unfinished construction",
-                "edit": "Edit and transform this construction image",
-                "blend": "Blend and reimagine this construction"
-            }
-            
-            full_prompt = f"{mode_descriptions.get(editing_mode, 'Transform')} {prompt}. Style: {style_modifiers.get(style, '')}. Make it look professional and realistic."
+            full_prompt = self._build_enhanced_prompt(prompt, style, editing_mode)
             
             url = "https://api.freepik.com/v1/ai/gemini-2-5-flash-image-preview"
             payload = {
@@ -107,7 +142,7 @@ class NanoBananaApp:
                 "webhook_url": None
             }
             headers = {
-                "x-freepik-api-key": FREEPIK_API_KEY,
+                "x-freepik-api-key": api_key,
                 "Content-Type": "application/json"
             }
             
@@ -119,13 +154,123 @@ class NanoBananaApp:
                     img_response = requests.get(result['image_url'], timeout=30)
                     if img_response.status_code == 200:
                         generated_image = Image.open(io.BytesIO(img_response.content))
-                        return generated_image, "Generated with Freepik Gemini 2.5 Flash"
+                        return generated_image, "🍌 Generated with Freepik Gemini 2.5 Flash"
                         
             return None, f"Freepik API error: {response.status_code}"
             
         except Exception as e:
             logger.error(f"Freepik generation failed: {e}")
             return None, f"Freepik error: {str(e)}"
+    
+    def _generate_with_dalle(self, image, prompt, style, editing_mode, api_key):
+        if not api_key:
+            return None, "OpenAI API key not provided"
+            
+        try:
+            full_prompt = self._build_enhanced_prompt(prompt, style, editing_mode)
+            
+            headers = {
+                "Authorization": f"Bearer {api_key}",
+                "Content-Type": "application/json"
+            }
+            
+            payload = {
+                "model": "dall-e-3",
+                "prompt": full_prompt,
+                "n": 1,
+                "size": "1024x1024"
+            }
+            
+            response = requests.post(
+                "https://api.openai.com/v1/images/generations",
+                headers=headers,
+                json=payload,
+                timeout=60
+            )
+            
+            if response.status_code == 200:
+                result = response.json()
+                if result.get('data') and len(result['data']) > 0:
+                    image_url = result['data'][0]['url']
+                    img_response = requests.get(image_url, timeout=30)
+                    if img_response.status_code == 200:
+                        generated_image = Image.open(io.BytesIO(img_response.content))
+                        return generated_image, "🎨 Generated with DALL-E 3"
+                        
+            return None, f"DALL-E API error: {response.status_code}"
+            
+        except Exception as e:
+            logger.error(f"DALL-E generation failed: {e}")
+            return None, f"DALL-E error: {str(e)}"
+    
+    def _generate_with_stable_diffusion(self, image, prompt, style, editing_mode, api_key):
+        if not api_key:
+            return None, "Stability AI API key not provided"
+            
+        try:
+            full_prompt = self._build_enhanced_prompt(prompt, style, editing_mode)
+            
+            headers = {
+                "Authorization": f"Bearer {api_key}",
+                "Content-Type": "application/json"
+            }
+            
+            payload = {
+                "text_prompts": [{"text": full_prompt}],
+                "cfg_scale": 7,
+                "height": 1024,
+                "width": 1024,
+                "samples": 1,
+                "steps": 30
+            }
+            
+            response = requests.post(
+                "https://api.stability.ai/v1/generation/stable-diffusion-xl-1024-v1-0/text-to-image",
+                headers=headers,
+                json=payload,
+                timeout=60
+            )
+            
+            if response.status_code == 200:
+                result = response.json()
+                if result.get('artifacts') and len(result['artifacts']) > 0:
+                    image_b64 = result['artifacts'][0]['base64']
+                    image_data = base64.b64decode(image_b64)
+                    generated_image = Image.open(io.BytesIO(image_data))
+                    return generated_image, "🚀 Generated with Stable Diffusion XL"
+                    
+            return None, f"Stability AI error: {response.status_code}"
+            
+        except Exception as e:
+            logger.error(f"Stable Diffusion generation failed: {e}")
+            return None, f"Stable Diffusion error: {str(e)}"
+    
+    def _generate_with_gemini_demo(self, image, prompt, style, editing_mode, api_key):
+        if not api_key:
+            try:
+                genai.configure(api_key=api_key)
+            except:
+                return None, "Gemini API key not provided or invalid"
+        
+        return self._fallback_to_gemini_demo(image, prompt, style, editing_mode)
+    
+    def _build_enhanced_prompt(self, prompt, style, editing_mode):
+        style_modifiers = {
+            "realistic": "photorealistic, high-quality construction, professional architecture",
+            "futuristic": "futuristic, high-tech, modern glass and steel, sci-fi architecture",
+            "artistic": "artistic, creative design, unique architecture, colorful and innovative"
+        }
+        
+        mode_descriptions = {
+            "complete": "Complete this unfinished construction",
+            "edit": "Edit and transform this construction image",
+            "blend": "Blend and reimagine this construction"
+        }
+        
+        base_desc = mode_descriptions.get(editing_mode, 'Transform')
+        style_desc = style_modifiers.get(style, '')
+        
+        return f"{base_desc} {prompt}. Style: {style_desc}. Make it look professional and realistic. Architecture, construction, building."
             
     def load_yolo_optional(self):
         if not yolo_available:
@@ -151,29 +296,36 @@ class NanoBananaApp:
         except Exception as e:
             return image, f"Detection failed: {str(e)}"
     
-    def nano_banana_edit(self, image, prompt, style="realistic", editing_mode="complete"):
-        if not self.gemini_model:
-            if not GEMINI_API_KEY:
-                return image, "🔑 API key required for Nano Banana. Add GEMINI_API_KEY to use this feature."
-            return image, "Gemini Nano Banana not available"
-        
+    def nano_banana_edit_with_model(self, image, prompt, style="realistic", editing_mode="complete", model_name="Freepik Gemini 2.5 Flash", api_keys=None):
         if not prompt.strip():
             return image, "Please provide a transformation prompt"
             
-        logger.info(f"Processing: {editing_mode} mode, {style} style, prompt: {prompt[:50]}...")
+        if api_keys is None:
+            api_keys = {}
+            
+        logger.info(f"Processing with {model_name}: {editing_mode} mode, {style} style, prompt: {prompt[:50]}...")
             
         try:
             image = self._resize_image_if_needed(image)
             self._apply_rate_limiting()
             
-            # Try Freepik Gemini 2.5 Flash first (best quality)
-            if FREEPIK_API_KEY:
-                freepik_result, freepik_msg = self._generate_with_freepik(image, prompt, style, editing_mode)
-                if freepik_result is not None:
-                    return freepik_result, f"🍌 Nano Banana (Freepik): {freepik_msg}"
-                logger.warning(f"Freepik failed: {freepik_msg}, falling back to Gemini analysis")
-            
-            # Fallback to Gemini analysis + demo overlay
+            result, message = self._generate_with_model(image, prompt, style, editing_mode, model_name, api_keys)
+            if result is not None:
+                return result, message
+            else:
+                logger.warning(f"{model_name} failed: {message}, falling back to demo mode")
+                return self._fallback_to_gemini_demo(image, prompt, style, editing_mode)
+                
+        except Exception as e:
+            logger.error(f"Processing failed: {e}")
+            return self._fallback_to_gemini_demo(image, prompt, style, editing_mode)
+    
+    def _fallback_to_gemini_demo(self, image, prompt, style, editing_mode):
+        if not self.gemini_model:
+            demo_result = self._create_nano_banana_demo(image, f"Demo: {editing_mode} mode with {style} style", style, editing_mode)
+            return demo_result, f"🎭 Demo Mode: {editing_mode} mode with {style} style"
+        
+        try:
             if editing_mode == "complete":
                 base_prompt = self._get_completion_prompt(style)
                 analysis_prompt = f"Analyze this construction image and describe how to {base_prompt.lower()} User request: {prompt}. Provide detailed description of the completed construction."
@@ -189,7 +341,7 @@ class NanoBananaApp:
                     buffered = io.BytesIO()
                     image.save(buffered, format='PNG', quality=85)
                     image_bytes = buffered.getvalue()
-                    
+
                     if len(image_bytes) > 10 * 1024 * 1024:
                         return image, "Image too large. Please use a smaller image."
                     
@@ -316,6 +468,56 @@ class NanoBananaApp:
 
 app = NanoBananaApp()
 
+def update_model_description(model_name):
+    return IMAGE_MODELS.get(model_name, {}).get("description", "Model description not available")
+
+def process_nano_banana_with_settings(image, prompt, style, editing_mode, model_name, freepik_key, openai_key, stability_key, gemini_key, enable_detection, enable_voice):
+    if not image:
+        return None, None, None, None, "📷 Please upload an image to get started", None
+    
+    if not prompt or not prompt.strip():
+        return image, image, image, None, "💭 Please provide a transformation prompt", None
+    
+    user_api_keys = {
+        "freepik": freepik_key or FREEPIK_API_KEY,
+        "openai": openai_key,
+        "stabilityai": stability_key, 
+        "gemini": gemini_key or GEMINI_API_KEY
+    }
+    
+    try:
+        detection_result = image
+        detection_msg = "Detection disabled"
+        
+        if enable_detection:
+            detection_result, detection_msg = app.detect_structures_optional(image)
+        
+        processed_image, process_msg = app.nano_banana_edit_with_model(
+            image, prompt, style, editing_mode, model_name, user_api_keys
+        )
+        
+        if processed_image == image and "API key" in process_msg:
+            return image, detection_result, image, None, f"🔑 {process_msg}", None
+        
+        comparison = app.create_comparison(image, processed_image)
+        
+        audio = None
+        voice_msg = ""
+        if enable_voice:
+            if processed_image != image:
+                voice_text = f"Using {model_name}, the AI processed this construction image with {editing_mode} mode and {style} style. Request: {prompt}"
+                audio = app.generate_voice_optional(voice_text)
+                voice_msg = "🔊 Voice generated" if audio else "🔇 Voice unavailable"
+            else:
+                voice_msg = "🔇 Voice skipped (no changes)"
+        
+        status = f"🍌 {process_msg}\n📍 Detection: {detection_msg}\n🎵 Voice: {voice_msg}"
+        return image, detection_result, processed_image, comparison, status, audio
+        
+    except Exception as e:
+        logger.error(f"Processing error: {e}")
+        return image, image, image, None, f"❌ Error: {str(e)}", None
+
 def process_nano_banana(image, prompt, style, editing_mode, enable_detection, enable_voice):
     if not image:
         return None, None, None, None, "📷 Please upload an image to get started", None
@@ -430,15 +632,47 @@ with gr.Blocks(title="🍌 Nano Banana - Dynamic Image Creation", theme=gr.theme
                     value="edit",
                     label="Nano Banana Mode",
                     info="Complete: Finish construction • Edit: Modify image • Blend: Fuse elements"
-                )
-                
-                style_selector = gr.Radio(
-                    choices=["realistic", "futuristic", "artistic"],
-                    value="realistic",
+            )
+            
+            style_selector = gr.Radio(
+                choices=["realistic", "futuristic", "artistic"],
+                value="realistic",
                     label="Style",
                     info="Choose the aesthetic approach"
                 )
             
+            with gr.Group():
+                gr.Markdown("### 🤖 AI Model & API Settings")
+                model_selector = gr.Dropdown(
+                    choices=list(IMAGE_MODELS.keys()),
+                    value="Freepik Gemini 2.5 Flash",
+                    label="Image Generation Model",
+                    info="Choose your preferred AI model"
+                )
+                model_description = gr.HTML(value=IMAGE_MODELS["Freepik Gemini 2.5 Flash"]["description"])
+                
+                with gr.Accordion("🔑 API Keys (Optional - Use Your Own)", open=False):
+                    freepik_key = gr.Textbox(
+                        label="Freepik API Key",
+                        placeholder="Enter your Freepik API key for real Gemini 2.5 Flash generation",
+                        type="password"
+                    )
+                    openai_key = gr.Textbox(
+                        label="OpenAI API Key", 
+                        placeholder="Enter your OpenAI API key for DALL-E 3",
+                        type="password"
+                    )
+                    stability_key = gr.Textbox(
+                        label="Stability AI API Key",
+                        placeholder="Enter your Stability AI key for Stable Diffusion XL", 
+                        type="password"
+                    )
+                    gemini_key = gr.Textbox(
+                        label="Gemini API Key",
+                        placeholder="Enter your Gemini API key for analysis mode",
+                        type="password"
+                    )
+            
             with gr.Group():
                 gr.Markdown("### ⚙️ Optional Features")
                 enable_detection = gr.Checkbox(
@@ -475,16 +709,16 @@ with gr.Blocks(title="🍌 Nano Banana - Dynamic Image Creation", theme=gr.theme
     with gr.Row():
         gr.Examples(
             examples=[
-                ["samples/building_001.jpg", "Complete this modern building with glass facade", "realistic", "complete", False, False],
-                ["samples/bridge_049.jpg", "Transform into futuristic suspension bridge", "futuristic", "edit", True, False],
-                ["samples/road_088.jpg", "Complete as smart highway with LED lights", "futuristic", "blend", False, True],
-                ["samples/construction_019.jpg", "Add artistic elements and colorful design", "artistic", "edit", False, False],
-                ["samples/infrastructure_015.jpg", "Complete with sustainable green technology", "realistic", "complete", True, False],
-                ["samples/residential_004.jpg", "Transform into eco-friendly smart home", "futuristic", "blend", False, False],
-                ["samples/commercial_010.jpg", "Add modern commercial design elements", "realistic", "edit", False, False],
-                ["samples/construction_111.jpg", "Complete with artistic architectural details", "artistic", "complete", False, True]
+                ["samples/building_001.jpg", "Complete this modern building with glass facade", "realistic", "complete", "Freepik Gemini 2.5 Flash", "", "", "", "", False, False],
+                ["samples/bridge_049.jpg", "Transform into futuristic suspension bridge", "futuristic", "edit", "OpenAI DALL-E 3", "", "", "", "", True, False],
+                ["samples/road_088.jpg", "Complete as smart highway with LED lights", "futuristic", "blend", "Stable Diffusion XL", "", "", "", "", False, True],
+                ["samples/construction_019.jpg", "Add artistic elements and colorful design", "artistic", "edit", "Freepik Gemini 2.5 Flash", "", "", "", "", False, False],
+                ["samples/infrastructure_015.jpg", "Complete with sustainable green technology", "realistic", "complete", "Gemini Analysis + Demo", "", "", "", "", True, False],
+                ["samples/residential_004.jpg", "Transform into eco-friendly smart home", "futuristic", "blend", "OpenAI DALL-E 3", "", "", "", "", False, False],
+                ["samples/commercial_010.jpg", "Add modern commercial design elements", "realistic", "edit", "Stable Diffusion XL", "", "", "", "", False, False],
+                ["samples/construction_111.jpg", "Complete with artistic architectural details", "artistic", "complete", "Freepik Gemini 2.5 Flash", "", "", "", "", False, True]
             ],
-            inputs=[image_input, prompt_input, style_selector, editing_mode, enable_detection, enable_voice],
+            inputs=[image_input, prompt_input, style_selector, editing_mode, model_selector, freepik_key, openai_key, stability_key, gemini_key, enable_detection, enable_voice],
             label="🎯 Try These Examples"
         )
     
@@ -498,9 +732,15 @@ with gr.Blocks(title="🍌 Nano Banana - Dynamic Image Creation", theme=gr.theme
     - **🔄 Smart Fallbacks**: Multiple processing modes for reliability
     """)
     
+    model_selector.change(
+        fn=update_model_description,
+        inputs=[model_selector],
+        outputs=[model_description]
+    )
+    
     process_btn.click(
-        fn=process_nano_banana,
-        inputs=[image_input, prompt_input, style_selector, editing_mode, enable_detection, enable_voice],
+        fn=process_nano_banana_with_settings,
+        inputs=[image_input, prompt_input, style_selector, editing_mode, model_selector, freepik_key, openai_key, stability_key, gemini_key, enable_detection, enable_voice],
         outputs=[original_output, detection_output, result_output, comparison_output, status_output, audio_output]
     )