Update app.py

app.py

@@ -29,6 +29,35 @@ def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return
             return f"Image generated successfully: {result}"
     except Exception as e:
         return f"Failed to generate image from prompt '{arg1}': {str(e)}"
+from smolagents import tool
+
+# Assuming this is already loaded earlier
+# image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+
+@tool
+def image_generator(arg1: str) -> str:
+    """
+    A tool that generates an image based on a text prompt.
+    
+    Args:
+        arg1: A description or prompt for the image to be generated.
+    
+    Returns:
+        A string containing the result or URL of the generated image.
+    """
+    try:
+        result = image_generation_tool(prompt=arg1)
+
+        # Handle different possible return types from the tool
+        if isinstance(result, str):
+            return f"Generated image: {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}"
+    except Exception as e:
+        return f"Error generating image from prompt '{arg1}': {str(e)}"
+
 
 @tool
 def get_current_time_in_timezone(timezone: str) -> str: