skroed commited on
Commit
dfc87a6
1 Parent(s): 760404a

Feat: initial test.

Browse files
Files changed (2) hide show
  1. handler.py +25 -0
  2. requirements.txt +2 -0
handler.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, Dict
2
+
3
+ from audiocraft.models import MusicGen
4
+
5
+
6
+ class EndpointHandler:
7
+ def __init__(self, path=""):
8
+ # load model and processor from path
9
+ self.model = MusicGen.get_pretrained("small")
10
+
11
+ def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
12
+ """
13
+ Args:
14
+ data (:dict:):
15
+ The payload with the text prompt and generation parameters.
16
+ """
17
+ # process input
18
+ output = (
19
+ self.model.generate_unconditional(num_samples=2, progress=True)[0]
20
+ .cpu()
21
+ .numpy()
22
+ .tolist()
23
+ )
24
+
25
+ return [{"generated_audio": output}]
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ audiocraft
2
+ torch==2.0.0