Update image_transformation.py

image_transformation.py

@@ -1,13 +1,14 @@
 import torch
 
-from transformers.tools.base import Tool, get_default_device
+from transformers.tools.base import Tool
 from transformers.utils import (
     is_accelerate_available,
     is_vision_available,
 )
 
 from diffusers import DiffusionPipeline
-
+if is_accelerate_available():
+    from accelerate import PartialState
 
 IMAGE_TRANSFORMATION_DESCRIPTION = (
     "This is a tool that transforms an image according to a prompt. It takes two inputs: `image`, which should be "
@@ -37,7 +38,7 @@ class ImageTransformationTool(Tool):
 
     def setup(self):
         if self.device is None:
-            self.device = get_default_device()
+            self.device = PartialState().default_device
 
         self.pipeline = DiffusionPipeline.from_pretrained(self.stable_diffusion)