m-ric HF staff commited on
Commit
d8a607f
1 Parent(s): b25e62a

Update text_to_image.py

Browse files
Files changed (1) hide show
  1. text_to_image.py +2 -1
text_to_image.py CHANGED
@@ -1,7 +1,7 @@
1
  from transformers.agents.tools import Tool
2
  from transformers.utils import is_accelerate_available
3
  import torch
4
-
5
 
6
  from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
7
  if is_accelerate_available():
@@ -11,6 +11,7 @@ TEXT_TO_IMAGE_DESCRIPTION = (
11
  "This is a tool that creates an image according to a prompt."
12
  )
13
 
 
14
  class TextToImageTool(Tool):
15
  default_checkpoint = "runwayml/stable-diffusion-v1-5"
16
  description = TEXT_TO_IMAGE_DESCRIPTION
 
1
  from transformers.agents.tools import Tool
2
  from transformers.utils import is_accelerate_available
3
  import torch
4
+ import spaces
5
 
6
  from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
7
  if is_accelerate_available():
 
11
  "This is a tool that creates an image according to a prompt."
12
  )
13
 
14
+ @spaces.GPU
15
  class TextToImageTool(Tool):
16
  default_checkpoint = "runwayml/stable-diffusion-v1-5"
17
  description = TEXT_TO_IMAGE_DESCRIPTION