switch to partialState

text_to_video.py

@@ -1,10 +1,11 @@
 import torch
 
-from transformers.tools.base import Tool, get_default_device
+from transformers.tools.base import Tool
 from transformers.utils import is_accelerate_available
 
 from diffusers import DiffusionPipeline
-
+if is_accelerate_available():
+    from accelerate import PartialState
 
 TEXT_TO_VIDEO_DESCRIPTION = (
     "This is a tool that creates a video according to a text description. It takes an input named `prompt` which "
@@ -31,7 +32,7 @@ class TextToVideoTool(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.default_checkpoint, variant="fp16"