updated handler to return audio base64 strings

handler.py

@@ -3,6 +3,7 @@ from scipy.io import wavfile
 from transformers import AutoProcessor, MusicgenForConditionalGeneration
 import torch
 import io
+import base64
 
 def create_params(params, fr):
     # default
@@ -37,7 +38,7 @@ class EndpointHandler:
         self.model = MusicgenForConditionalGeneration.from_pretrained(path, torch_dtype=torch.float16)
         self.model.to('cuda')
 
-    def __call__(self, data: Dict[str, Any]) -> bytes:
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
         Args:
             data (:dict:):
@@ -60,7 +61,7 @@ class EndpointHandler:
         with torch.cuda.amp.autocast():
             outputs = self.model.generate(**inputs, **params)
 
-        pred = outputs[0].cpu().numpy().tolist()
+        pred = outputs[0].cpu().numpy()
         sr = 32000
 
         try:
@@ -76,7 +77,11 @@ class EndpointHandler:
         # Convert BytesIO to bytes
         wav_data = wav_buffer.getvalue()
 
-        return  wav_data
+
+        # Convert the WAV binary data to Base64
+        base64_encoded_wav = base64.b64encode(wav_data).decode('utf-8')
+
+        return [{"audio": base64_encoded_wav, "sr": sr}]
 
 
 if __name__ == "__main__":

test.py

@@ -1,17 +1,12 @@
 from handler import EndpointHandler
 
 # init handler
-my_handler = EndpointHandler(path="pbotsaris/musicgen-small")
+print('init handler')
+my_handler = EndpointHandler(path=".")
+p = {"inputs": "I am quite excited how this will turn out", "duration": 2}
 
-non_holiday_payload = {"inputs": "I am quite excited how this will turn out", "duration": 2}
-# prepare sample payload
+print('calling handler')
+pred=my_handler(p)
 
-# test the handler
-non_holiday_pred=my_handler(non_holiday_payload)
-
-# show results
 print("done")
-print(non_holiday_pred)
-
-# non_holiday_pred [{'label': 'joy', 'score': 0.9985942244529724}]
-# holiday_payload [{'label': 'happy', 'score': 1}]
+print(pred)