skroed commited on
Commit
e1735f5
1 Parent(s): 288e1c4

Fix: handler.

Browse files
Files changed (1) hide show
  1. handler.py +4 -3
handler.py CHANGED
@@ -7,9 +7,9 @@ from transformers import AutoModel, AutoProcessor
7
  class EndpointHandler:
8
  def __init__(self, path=""):
9
  # load model and processor from path
10
- self.processor = AutoProcessor.from_pretrained(path)
11
  self.model = AutoModel.from_pretrained(
12
- path,
13
  ).to("cuda")
14
 
15
  def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
@@ -21,7 +21,8 @@ class EndpointHandler:
21
  # process input
22
  text = data.pop("inputs", data)
23
  parameters = data.get("parameters", None)
24
-
 
25
  # preprocess
26
  inputs = self.processor(
27
  text=[text],
 
7
  class EndpointHandler:
8
  def __init__(self, path=""):
9
  # load model and processor from path
10
+ self.processor = AutoProcessor.from_pretrained("suno/bark-small")
11
  self.model = AutoModel.from_pretrained(
12
+ "suno/bark-small",
13
  ).to("cuda")
14
 
15
  def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
 
21
  # process input
22
  text = data.pop("inputs", data)
23
  parameters = data.get("parameters", None)
24
+ if parameters is None:
25
+ parameters = {}
26
  # preprocess
27
  inputs = self.processor(
28
  text=[text],