Update app.py

app.py

@@ -36,7 +36,8 @@ from smolagents import tool
 # Import tool from Hub
 image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
 @tool
-def image_generator(arg1: str) -> str:
+@tool
+def image_generator(arg1: str):
     """
     A tool that generates an image based on a text prompt.
     
@@ -44,20 +45,23 @@ def image_generator(arg1: str) -> str:
         arg1: A description or prompt for the image to be generated.
     
     Returns:
-        A string containing the result or URL of the generated image.
+        A PIL.Image object which Gradio can display.
     """
     try:
         result = image_generation_tool(prompt=arg1)
 
-        # Handle different possible return types from the tool
-        if isinstance(result, str):
-            return f"Generated image: {result}"
+        # If result is a PIL image, return it directly
+        if hasattr(result, 'show') and callable(result.show):
+            return result  # Gradio will display this automatically
+        elif isinstance(result, str):
+            return f"Generated image URL: {result}"
         elif isinstance(result, dict) and "image_url" in result:
             return f"Generated image URL: {result['image_url']}"
         else:
-            return f"Image generated successfully: {result}"
+            return f"Image generated: {result}"
     except Exception as e:
-        return f"Error generating image from prompt '{arg1}': {str(e)}"
+        return f"Error generating image: {str(e)}"
+
 
 
 @tool