Update app.py

app.py

@@ -5,9 +5,9 @@ import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
 
-
 from Gradio_UI import GradioUI
-
+from PIL import Image
+import gradio as gr
 
 # Below is an example of a tool that does nothing. Amaze us with your creativity !
 #@tool
@@ -38,12 +38,23 @@ def animal_image_generation_tool(animal_name: str) -> str:
 
         # generates an image based on the description of an animal.
         promt=f"Create a realistic image of a {animal_name}"
-        image_url = image_generation_tool(promt)
-                
-        return f"Image {animal_name} successfully generated: {image_url}"
+        image = image_generation_tool(promt)  # This returns a PIL image
+
+        # Save the image temporarily to display it in Gradio
+        image_path = f"{animal_name}_generated_image.jpg"
+        image.save(image_path)
+
+        return image_path # Return the path to the saved image       
     except Exception as e:
         return f"Error during image generation: {str(e)}"
 
+# Function to display the image in Gradio
+def display_image(animal_name: str):
+    # Call the tool to generate the image
+    image_path = animal_image_generation_tool(animal_name)
+    if image_path.startswith("Error"):
+        return image_path  # Return the error message if something went wrong
+    return image_path  # Return the path to the image
 
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
@@ -97,5 +108,19 @@ agent = CodeAgent(
     prompt_templates=prompt_templates
 )
 
+# Gradio Interface
+def gradio_interface(animal_name: str):
+    return display_image(animal_name)
+
+
+# Launch Gradio UI
+iface = gr.Interface(
+    fn=gradio_interface,
+    inputs="text",
+    outputs="image",
+    title="Animal Image Generator",
+    description="Generate an image of an animal by its name.",
+)
+
 
 GradioUI(agent).launch()