skroed commited on
Commit
6eb9b90
1 Parent(s): a48e070

Upload CustomHandler and requirements.

Browse files
Files changed (2) hide show
  1. handler.py +38 -0
  2. requirements.txt +2 -0
handler.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, Dict
2
+
3
+ import torch
4
+ from transformers import AutoModel, AutoProcessor
5
+
6
+
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]:
16
+ """
17
+ Args:
18
+ data (:dict:):
19
+ The payload with the text prompt and generation parameters.
20
+ """
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],
28
+ return_tensors="pt",
29
+ voice_preset=parameters.get("voice_preset", None),
30
+ ).to("cuda")
31
+
32
+ with torch.autocast("cuda"):
33
+ outputs = self.model.generate(**inputs)
34
+
35
+ # postprocess the prediction
36
+ prediction = outputs.cpu().numpy().tolist()
37
+
38
+ return [{"generated_audio": prediction}]
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ transformers==4.34.1
2
+ accelerate>=0.23.0