Update text_to_image.py

text_to_image.py

@@ -1,3 +1,4 @@
+from PIL import Image
 from transformers.tools.base import Tool
 from transformers.utils import is_accelerate_available
 import torch
@@ -15,8 +16,9 @@ TEXT_TO_IMAGE_DESCRIPTION = (
 class TextToImageTool(Tool):
     default_checkpoint = "runwayml/stable-diffusion-v1-5"
     description = TEXT_TO_IMAGE_DESCRIPTION
-    inputs = ['text']
-    outputs = ['image']
+    name = "image_generator"
+    inputs = {"prompt": str}
+    output_type = Image.Image
 
     def __init__(self, device=None, **hub_kwargs) -> None:
         if not is_accelerate_available():