Update app.py

app.py

@@ -35,8 +35,10 @@ from smolagents import tool
 # image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
 # Import tool from Hub
 image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+from PIL import Image
+
 @tool
-def image_generator(arg1: str):
+def image_generator(arg1: str) -> Image.Image:
     """
     A tool that generates an image based on a text prompt.
     
@@ -50,16 +52,28 @@ def image_generator(arg1: str):
         result = image_generation_tool(prompt=arg1)
 
         # 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}"
+        if isinstance(result, Image.Image):
+            return result
+
+        # If result is a URL string
+        elif isinstance(result, str) and result.startswith("http"):
+            response = requests.get(result)
+            image = Image.open(BytesIO(response.content))
+            return image
+
+        # If result is a dictionary with an image URL
         elif isinstance(result, dict) and "image_url" in result:
-            return f"Generated image URL: {result['image_url']}"
-        else:
-            return f"Image generated: {result}"
+            image_url = result["image_url"]
+            response = requests.get(image_url)
+            image = Image.open(BytesIO(response.content))
+            return image
+
+        return Image.new("RGB", (256, 256), color="gray")  # fallback dummy image
+    
     except Exception as e:
-        return f"Error generating image: {str(e)}"
+        # Show a red dummy image with error in metadata or raise it
+        return Image.new("RGB", (256, 256), color="red")
+