Update app.py to try to render the image

app.py

@@ -5,12 +5,13 @@ import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
 from PIL import Image
+from io import BytesIO
 
 from Gradio_UI import GradioUI
 
 # Below is an example of a tool that does nothing. Amaze us with your creativity !
 @tool
-def generate_image_by_description(image_description:str)-> Image.Image: #it's import to specify the return type
+def generate_image_by_description(image_description:str)-> BytesIO: #it's import to specify the return type
     """A tool that generates an image based on the image_description variable
     Args:
         image_description: the description of the image to be generated
@@ -21,9 +22,17 @@ def generate_image_by_description(image_description:str)-> Image.Image: #it's im
 
     image = image_generation_tool(image_description)
 
-    # image.show()
+    # Normalize
+    if isinstance(image, list):
+        image = image[0]
+    elif isinstance(image, dict) and "images" in image:
+        image = image["images"][0]
 
-    return image
+    buffer = BytesIO()
+    image.save(buffer, format="PNG")
+    buffer.seek(0)
+
+    return buffer
 
 @tool
 def display_image(image:Image.Image)-> None: #it's import to specify the return type
@@ -68,7 +77,7 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, generate_image_by_description, display_image], ## add your tools here (don't remove final answer)
+    tools=[final_answer, generate_image_by_description], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,