Update app.py

app.py

@@ -23,7 +23,7 @@ import gradio as gr
 
 # tool which i creating. A tool for generating an animal image
 @tool
-def animal_image_generation_tool(animal_name: str) -> str:
+def animal_image_generation_tool(animal_name: str) -> Image.Image:
     """Generating an image of an animal by its name.
     Args:
         animal_name: A string representing the name of the animal to search for (e.g., 'cat').
@@ -40,21 +40,17 @@ def animal_image_generation_tool(animal_name: str) -> str:
         promt=f"Create a realistic image of a {animal_name}"
         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       
+        return image # 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
+    image = animal_image_generation_tool(animal_name)
+    if isinstance(image, str):  # If an error message is returned
+        return image  # Return the error message
+    return image  # Return the PIL image object
 
 @tool
 def get_current_time_in_timezone(timezone: str) -> str: