import PartialState

text_to_image.py

@@ -3,7 +3,8 @@ from transformers.utils import is_accelerate_available
 import torch
 
 from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
-
+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. It takes an input named `prompt` which "
@@ -18,9 +19,7 @@ class TextToImageTool(Tool):
     outputs = ['image']
 
     def __init__(self, device=None, **hub_kwargs) -> None:
-        if is_accelerate_available():
-            from accelerate import PartialState
-        else:
+        if not is_accelerate_available():
             raise ImportError("Accelerate should be installed in order to use tools.")
 
         super().__init__()