Update app.py

app.py

@@ -1,161 +1,187 @@
 from smolagents import CodeAgent, HfApiModel, load_tool, tool
 import yaml
 import logging
+import traceback
 from tools.final_answer import FinalAnswerTool
 from Gradio_UI import GradioUI
-import traceback
 
-# Set up detailed logging
-logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+# Set up comprehensive logging
+logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
-# Custom image generation tool with better error handling
+# Create a diagnostic image generation tool
 @tool
-def enhanced_image_generator(prompt: str) -> str:
-    """Generate an image from a text prompt with enhanced error handling.
+def diagnostic_image_generator(prompt: str) -> str:
+    """Generate an image with comprehensive debugging and validation.
     Args:
         prompt: A detailed text description of the image to generate
     """
+    logger.info(f"=== DIAGNOSTIC IMAGE GENERATION START ===")
+    logger.info(f"Input prompt: {prompt}")
+    
     try:
-        logger.info(f"Attempting to generate image with prompt: {prompt}")
+        # Check if we have a base tool
+        if not hasattr(diagnostic_image_generator, '_base_tool'):
+            logger.error("No base image generation tool attached")
+            return "Error: No image generation tool available"
         
-        # Try to use the loaded image generation tool
-        if hasattr(enhanced_image_generator, '_base_tool'):
-            result = enhanced_image_generator._base_tool(prompt)
-            logger.info(f"Image generation result type: {type(result)}")
-            logger.info(f"Image generation result: {result}")
-            return result
-        else:
-            logger.error("Base image generation tool not available")
-            return "Error: Image generation tool not properly loaded"
+        base_tool = diagnostic_image_generator._base_tool
+        logger.info(f"Base tool type: {type(base_tool)}")
+        logger.info(f"Base tool: {base_tool}")
+        
+        # Call the base tool
+        logger.info("Calling base image generation tool...")
+        result = base_tool(prompt)
+        
+        # Analyze the result
+        logger.info(f"Raw result type: {type(result)}")
+        logger.info(f"Raw result: {result}")
+        
+        # Check if it's an AgentImage
+        if hasattr(result, '__class__') and 'AgentImage' in str(type(result)):
+            logger.info("Result is an AgentImage")
+            logger.info(f"AgentImage attributes: {dir(result)}")
             
+            # Try to get image properties
+            try:
+                if hasattr(result, 'size'):
+                    logger.info(f"Image size: {result.size}")
+                if hasattr(result, 'mode'):
+                    logger.info(f"Image mode: {result.mode}")
+                if hasattr(result, 'width'):
+                    logger.info(f"Image width: {result.width}")
+                if hasattr(result, 'height'):
+                    logger.info(f"Image height: {result.height}")
+                if hasattr(result, 'format'):
+                    logger.info(f"Image format: {result.format}")
+                if hasattr(result, 'show'):
+                    logger.info("Image has show method")
+                if hasattr(result, 'save'):
+                    logger.info("Image has save method")
+                    
+            except Exception as e:
+                logger.error(f"Error checking image properties: {e}")
+        
+        # Try to validate the image
+        if result and hasattr(result, 'size'):
+            width, height = result.size if hasattr(result, 'size') else (0, 0)
+            if width > 0 and height > 0:
+                logger.info(f"✅ Valid image generated: {width}x{height}")
+                return result
+            else:
+                logger.error(f"❌ Invalid image size: {width}x{height}")
+                return "Error: Generated image has invalid size"
+        
+        logger.info(f"=== DIAGNOSTIC IMAGE GENERATION END ===")
+        return result
+        
     except Exception as e:
-        logger.error(f"Image generation failed: {str(e)}")
-        logger.error(f"Traceback: {traceback.format_exc()}")
-        return f"Error generating image: {str(e)}"
+        logger.error(f"Image generation failed with exception: {str(e)}")
+        logger.error(f"Full traceback: {traceback.format_exc()}")
+        return f"Error: {str(e)}"
 
-# Initialize final answer tool
+# Initialize components
 final_answer = FinalAnswerTool()
 
-# Model configuration with fallback
-def create_model():
-    try:
-        model = HfApiModel(
-            max_tokens=1024,  # Reduced for better performance
-            temperature=0.7,
-            model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
-            custom_role_conversions=None,
-        )
-        logger.info("Primary model initialized successfully")
-        return model
-    except Exception as e:
-        logger.warning(f"Primary model failed: {e}")
-        try:
-            fallback_model = HfApiModel(
-                max_tokens=1024,
-                temperature=0.7,
-                model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
-                custom_role_conversions=None,
-            )
-            logger.info("Fallback model initialized successfully")
-            return fallback_model
-        except Exception as e2:
-            logger.error(f"Fallback model also failed: {e2}")
-            raise
-
-model = create_model()
+# Create model
+model = HfApiModel(
+    max_tokens=1024,
+    temperature=0.7,
+    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
+    custom_role_conversions=None,
+)
 
-# Load image generation tool with multiple fallbacks
-def load_image_tool():
-    tools_to_try = [
-        ("agents-course/text-to-image", "Primary image tool"),
-        ("multimodalart/stable-diffusion-xl", "SDXL fallback"),
-        ("runwayml/stable-diffusion-v1-5", "SD 1.5 fallback"),
-        ("stabilityai/stable-diffusion-2-1", "SD 2.1 fallback")
-    ]
+# Try to load image generation tool with detailed diagnostics
+logger.info("=== LOADING IMAGE GENERATION TOOL ===")
+try:
+    # First, let's try the primary tool
+    logger.info("Loading agents-course/text-to-image...")
+    base_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+    logger.info(f"✅ Tool loaded successfully: {type(base_tool)}")
     
-    for tool_name, description in tools_to_try:
-        try:
-            logger.info(f"Attempting to load {description}: {tool_name}")
-            tool = load_tool(tool_name, trust_remote_code=True)
-            logger.info(f"Successfully loaded {description}")
-            
-            # Test the tool
-            test_result = tool("test image")
-            logger.info(f"Tool test result type: {type(test_result)}")
-            
-            return tool
-            
-        except Exception as e:
-            logger.warning(f"Failed to load {description}: {e}")
-            continue
+    # Test the tool directly
+    logger.info("Testing tool directly...")
+    test_result = base_tool("a simple red circle")
+    logger.info(f"Direct test result type: {type(test_result)}")
+    logger.info(f"Direct test result: {test_result}")
     
-    logger.error("All image generation tools failed to load")
-    return None
-
-# Load the image generation tool
-base_image_tool = load_image_tool()
-
-# Attach the base tool to our enhanced version
-if base_image_tool:
-    enhanced_image_generator._base_tool = base_image_tool
-    logger.info("Enhanced image generator configured with base tool")
-else:
-    logger.error("No image generation tool available")
-
-# Load prompt templates with fallback
-def load_prompts():
-    try:
-        with open("prompts.yaml", 'r') as stream:
-            prompts = yaml.safe_load(stream)
-            logger.info("Loaded prompts from prompts.yaml")
-            return prompts
-    except Exception as e:
-        logger.warning(f"Failed to load prompts.yaml: {e}")
-        # Fallback prompts optimized for image generation
-        return {
-            "system": """You are an AI agent specialized in generating images from text descriptions. 
-            When a user requests an image, use the enhanced_image_generator tool with a detailed, descriptive prompt.
-            Always provide clear, vivid descriptions for better image generation results.
-            If image generation fails, explain the issue and suggest alternative approaches.""",
-            "user": "Generate an image based on this description: {input}"
-        }
+    # Check if test result is valid
+    if hasattr(test_result, 'size'):
+        logger.info(f"Test image size: {test_result.size}")
+        if test_result.size == (0, 0):
+            logger.warning("⚠️ Test image has size 0x0 - tool may not be working properly")
+        else:
+            logger.info("✅ Test image has valid size")
+    
+    # Attach to diagnostic tool
+    diagnostic_image_generator._base_tool = base_tool
+    image_tool_available = True
+    
+except Exception as e:
+    logger.error(f"❌ Failed to load image generation tool: {e}")
+    logger.error(f"Traceback: {traceback.format_exc()}")
+    image_tool_available = False
 
-prompt_templates = load_prompts()
+# Load prompts
+try:
+    with open("prompts.yaml", 'r') as stream:
+        prompt_templates = yaml.safe_load(stream)
+except:
+    prompt_templates = {
+        "system": "You are an AI assistant that can generate images. Use the diagnostic_image_generator tool to create images from text descriptions.",
+        "user": "{input}"
+    }
 
-# Create tools list
+# Create agent
 tools_list = [final_answer]
-if base_image_tool:
-    tools_list.append(enhanced_image_generator)
-    logger.info("Enhanced image generator added to tools")
+if image_tool_available:
+    tools_list.append(diagnostic_image_generator)
+    logger.info("✅ Diagnostic image generator added to agent")
+else:
+    logger.error("❌ No image generation tool available")
 
-# Create agent with comprehensive configuration
 agent = CodeAgent(
     model=model,
     tools=tools_list,
-    max_steps=4,  # Reduced for efficiency
-    verbosity_level=2,  # High verbosity for debugging
+    max_steps=3,
+    verbosity_level=2,
     grammar=None,
     planning_interval=None,
-    name="DebugImageAgent",
-    description="AI agent for image generation with enhanced debugging and error handling",
+    name="DiagnosticImageAgent",
+    description="AI agent with comprehensive image generation diagnostics",
     prompt_templates=prompt_templates
 )
 
-# Launch function with comprehensive error handling
-def launch_agent():
+# Create a simple test function
+def test_image_generation():
+    """Test image generation directly"""
+    logger.info("=== DIRECT IMAGE GENERATION TEST ===")
+    if image_tool_available:
+        try:
+            result = diagnostic_image_generator("a red apple on a white background")
+            logger.info(f"Direct test completed. Result: {result}")
+        except Exception as e:
+            logger.error(f"Direct test failed: {e}")
+    else:
+        logger.error("Cannot test - no image tool available")
+
+# Launch with diagnostics
+def launch_with_diagnostics():
     try:
-        logger.info("Starting image generation agent...")
-        logger.info(f"Available tools: {[tool.__name__ if hasattr(tool, '__name__') else str(tool) for tool in tools_list]}")
+        logger.info("=== LAUNCHING DIAGNOSTIC AGENT ===")
+        logger.info(f"Tools available: {len(tools_list)}")
+        logger.info(f"Image tool available: {image_tool_available}")
+        
+        # Run a quick test
+        test_image_generation()
         
-        # Launch with minimal parameters to avoid conflicts
+        # Launch the UI
+        logger.info("Starting Gradio UI...")
         GradioUI(agent).launch()
         
     except Exception as e:
-        logger.error(f"Failed to launch agent: {e}")
-        logger.error(f"Full traceback: {traceback.format_exc()}")
-        print(f"\nERROR: {e}")
-        print("Please check the logs above for detailed error information.")
+        logger.error(f"Launch failed: {e}")
+        logger.error(f"Traceback: {traceback.format_exc()}")
 
 if __name__ == "__main__":
-    launch_agent()
+    launch_with_diagnostics()