m-ric HF staff commited on
Commit
1f99294
·
verified ·
1 Parent(s): 660e224

Update text_to_video.py

Browse files
Files changed (1) hide show
  1. text_to_video.py +5 -4
text_to_video.py CHANGED
@@ -1,5 +1,5 @@
1
  import torch
2
-
3
  from transformers.tools.base import Tool
4
  from transformers.utils import is_accelerate_available
5
 
@@ -17,9 +17,10 @@ TEXT_TO_VIDEO_DESCRIPTION = (
17
  class TextToVideoTool(Tool):
18
  default_checkpoint = "damo-vilab/text-to-video-ms-1.7b"
19
  description = TEXT_TO_VIDEO_DESCRIPTION
20
- inputs = ['text']
21
- outputs = ['video']
22
-
 
23
  def __init__(self, device=None, **hub_kwargs) -> None:
24
  if not is_accelerate_available():
25
  raise ImportError("Accelerate should be installed in order to use tools.")
 
1
  import torch
2
+ from typing import Any
3
  from transformers.tools.base import Tool
4
  from transformers.utils import is_accelerate_available
5
 
 
17
  class TextToVideoTool(Tool):
18
  default_checkpoint = "damo-vilab/text-to-video-ms-1.7b"
19
  description = TEXT_TO_VIDEO_DESCRIPTION
20
+ name = "video_generator"
21
+ inputs = {"prompt": str}
22
+ output_type = Any
23
+
24
  def __init__(self, device=None, **hub_kwargs) -> None:
25
  if not is_accelerate_available():
26
  raise ImportError("Accelerate should be installed in order to use tools.")