Update image_transformation.py

image_transformation.py

@@ -1,5 +1,5 @@
+from PIL import Image
 import torch
-
 from transformers.tools.base import Tool
 from transformers.utils import (
     is_accelerate_available,
@@ -20,8 +20,8 @@ IMAGE_TRANSFORMATION_DESCRIPTION = (
 class ImageTransformationTool(Tool):
     default_stable_diffusion_checkpoint = "timbrooks/instruct-pix2pix"
     description = IMAGE_TRANSFORMATION_DESCRIPTION
-    inputs = ['image', 'text']
-    outputs = ['image']
+    inputs = {'image': Image.Image, 'prompt': str}
+    output_type = Image.Image
 
     def __init__(self, device=None, controlnet=None, stable_diffusion=None, **hub_kwargs) -> None:
         if not is_accelerate_available():