Update app.py

app.py

@@ -51,12 +51,46 @@ custom_role_conversions=None,
 # Import tool from Hub
 # @tool
 # image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+# @tool
+# def generate_image_from_text(prompt: str) -> Image:
+#     """
+#     A tool that generates an image based on the given text prompt and displays it.
+#     Args:
+#         prompt: A string describing the image to be generated.
+#     """
+#     try:
+#         # Load the image generation tool
+#         image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+        
+#         # Generate the image and get the URL
+#         image_url = image_generation_tool(prompt)
+
+#         # Fetch the image from the URL
+#         response = requests.get(image_url)
+#         image = Image.open(io.BytesIO(response.content))
+
+#         # Display the image inline for Jupyter Notebook or applications with inline display support
+#         display.display(image)
+#         image.save("generated_image.png")
+#         print("Image saved as generated_image.png. Open the file to view it.")
+#         # Return the image object
+#         return image
+
+#     except Exception as e:
+#         print(f"Error generating image: {str(e)}")
+
+
+
 @tool
 def generate_image_from_text(prompt: str) -> Image:
     """
     A tool that generates an image based on the given text prompt and displays it.
+    
     Args:
         prompt: A string describing the image to be generated.
+    
+    Returns:
+        Image: The generated image that can be displayed by Gradio.
     """
     try:
         # Load the image generation tool
@@ -69,16 +103,17 @@ def generate_image_from_text(prompt: str) -> Image:
         response = requests.get(image_url)
         image = Image.open(io.BytesIO(response.content))
 
-        # Display the image inline for Jupyter Notebook or applications with inline display support
-        display.display(image)
+        # Optionally save the image to a file
         image.save("generated_image.png")
         print("Image saved as generated_image.png. Open the file to view it.")
-        # Return the image object
+
+        # Return the image object for Gradio to display
         return image
 
     except Exception as e:
         print(f"Error generating image: {str(e)}")
-        
+        return None
+
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)