Update text_to_image.py

text_to_image.py

@@ -1,4 +1,3 @@
-from PIL import Image
 from transformers.agents.tools import Tool
 from transformers.utils import is_accelerate_available
 import torch
@@ -8,7 +7,7 @@ if is_accelerate_available():
     from accelerate import PartialState
 
 TEXT_TO_IMAGE_DESCRIPTION = (
-    "This is a tool that creates an image according to a prompt, which is a text description."
+    "This is a tool that creates an image according to a prompt."
 )
 
 
@@ -16,8 +15,8 @@ class TextToImageTool(Tool):
     default_checkpoint = "runwayml/stable-diffusion-v1-5"
     description = TEXT_TO_IMAGE_DESCRIPTION
     name = "image_generator"
-    inputs = {"prompt": {"type": str, "description": "contains the image description and outputs an image"}}
-    output_type = Image.Image
+    inputs = {"prompt": {"type": "text", "description": "the image description"}}
+    output_type = "image"
 
     def __init__(self, device=None, **hub_kwargs) -> None:
         if not is_accelerate_available():