Update app.py

app.py

@@ -127,7 +127,7 @@ def enhance_prompt_with_gemini(prompt: str, style: str) -> str:
         Return only the enhanced prompt without explanation.
         """
         
-        model = genai.GenerativeModel('gemini-2.0-flash-exp')
+        model = genai.GenerativeModel('gemini-2.5-pro')
         response = model.generate_content(enhancement_prompt)
         enhanced = response.text.strip()
         
@@ -172,9 +172,9 @@ def generate_marketing_image(prompt: str, style: str = "realistic") -> str:
                 # Initialize the genai SDK client  
                 client = genai_sdk.Client(api_key=GOOGLE_API_KEY)
                 
-                # Generate image using Imagen 3 via SDK
+                # Generate image using Imagen 4.0 (best balance of quality and speed)
                 result = client.models.generate_images(
-                    model="imagen-3.0-generate-002",
+                    model="imagen-4.0-generate-preview-06-06",
                     prompt=enhanced_prompt,
                     config={
                         "number_of_images": 1,
@@ -201,7 +201,8 @@ def generate_marketing_image(prompt: str, style: str = "realistic") -> str:
                             "prompt": prompt,
                             "enhanced_prompt": enhanced_prompt,
                             "style": style,
-                            "generation_method": "google-genai-sdk",
+                            "generation_method": "imagen-4.0",
+                            "model_name": "imagen-4.0-generate-preview-06-06",
                             "real_ai_generation": True
                         }
                         
@@ -320,8 +321,8 @@ def analyze_marketing_image_with_gemini(image_url: str, prompt: str, review_guid
             logger.error("Could not load image for analysis")
             return _fallback_image_analysis(prompt, review_guidelines)
         
-        # Generate analysis using Gemini Vision
-        model = genai.GenerativeModel('gemini-2.0-flash-exp')
+        # Generate analysis using Gemini 2.5 Pro with enhanced multimodal understanding
+        model = genai.GenerativeModel('gemini-2.5-pro')
         response = model.generate_content([analysis_prompt, image])
         analysis_text = response.text
         
@@ -396,7 +397,8 @@ def _parse_gemini_analysis(analysis_text: str, original_prompt: str) -> Dict[str
             "language_issues": language_issues,
             "effectiveness_issues": effectiveness_issues,
             "recommendations": recommendations[:5],  # Limit to top 5
-            "analysis_method": "gemini_vision",
+            "analysis_method": "gemini-2.5-pro-vision",
+            "model_name": "gemini-2.5-pro",
             "original_prompt": original_prompt
         }
         
@@ -647,14 +649,17 @@ def process_generated_image_and_results(api_response_str: str) -> Tuple[Image.Im
             review_method = metadata.get('review_method', 'unknown')
             
             generation_info = ""
-            if generation_method == "google-genai-sdk":
-                generation_info = "🎨 **Generated with**: Google Imagen3 SDK (Real AI)\n"
+            if generation_method == "imagen-4.0":
+                model_name = metadata.get('model_name', 'imagen-4.0-generate-preview-06-06')
+                generation_info = f"🎨 **Generated with**: {model_name} (Real AI)\n"
+            elif generation_method == "google-genai-sdk":
+                generation_info = "🎨 **Generated with**: Google Imagen 4.0 (Real AI)\n"
             elif generation_method == "placeholder":
                 generation_info = "🎨 **Generated with**: Placeholder (Fallback)\n"
             
             review_method_info = ""
             if review_method == "gemini_vision":
-                review_method_info = "🔍 **Reviewed with**: Gemini Vision (AI Analysis)\n"
+                review_method_info = "🔍 **Reviewed with**: Gemini 2.5 Pro Vision (AI Analysis)\n"
             elif review_method == "fallback_text":
                 review_method_info = "🔍 **Reviewed with**: Text Analysis (Fallback)\n"
             
@@ -852,9 +857,70 @@ with gr.Blocks(title="Marketing Image Generator MCP", theme=gr.themes.Soft()) as
     </div>
     """)
 
+def test_imagen4_models():
+    """Test if Imagen 4.0 models are accessible"""
+    if not GEMINI_AVAILABLE or not GOOGLE_API_KEY:
+        logger.warning("❌ Cannot test Imagen 4.0 - Google AI not configured")
+        return []
+    
+    imagen4_models = [
+        "imagen-4.0-generate-preview-06-06",
+        "imagen-4.0-fast-generate-preview-06-06", 
+        "imagen-4.0-ultra-generate-preview-06-06"
+    ]
+    
+    logger.info("🧪 Testing Imagen 4.0 model access...")
+    working_models = []
+    
+    for model_name in imagen4_models:
+        try:
+            logger.info(f"Testing {model_name}...")
+            
+            client = genai_sdk.Client(api_key=GOOGLE_API_KEY)
+            result = client.models.generate_images(
+                model=model_name,
+                prompt="A simple red circle",
+                config={
+                    "number_of_images": 1,
+                    "output_mime_type": "image/png"
+                }
+            )
+            
+            if result and hasattr(result, 'generated_images') and len(result.generated_images) > 0:
+                working_models.append(model_name)
+                logger.info(f"✅ {model_name}: ACCESSIBLE")
+            else:
+                logger.warning(f"⚠️ {model_name}: No image returned")
+                
+        except Exception as e:
+            error_msg = str(e)
+            if "404" in error_msg or "not found" in error_msg.lower():
+                logger.warning(f"⚠️ {model_name}: Model not found/available")
+            elif "403" in error_msg or "permission" in error_msg.lower():
+                logger.warning(f"⚠️ {model_name}: Permission denied")
+            else:
+                logger.error(f"❌ {model_name}: {error_msg}")
+    
+    logger.info("=" * 50)
+    logger.info(f"🎉 ACCESSIBLE IMAGEN 4.0 MODELS: {len(working_models)}")
+    for model in working_models:
+        logger.info(f"   ✅ {model}")
+    
+    if not working_models:
+        logger.info("   ⚠️ No Imagen 4.0 models accessible - continuing with Imagen 3.0")
+    
+    logger.info("=" * 50)
+    return working_models
+
 if __name__ == "__main__":
     logger.info("🚀 Starting Marketing Image Generator with MCP Server")
     logger.info(f"🔑 Google AI: {'✅ Configured' if GOOGLE_API_KEY else '❌ No API Key'}")
     logger.info("🔗 MCP Server will be available at /gradio_api/mcp/sse")
     
+    # Test Imagen 4.0 model access on startup
+    if GOOGLE_API_KEY:
+        accessible_imagen4_models = test_imagen4_models()
+        if accessible_imagen4_models:
+            logger.info(f"💡 UPGRADE OPPORTUNITY: You can use {accessible_imagen4_models[0]} for better quality!")
+    
     demo.launch(mcp_server=True)