added tests and changed handler to respond with an wav file

handler.py

@@ -1,6 +1,8 @@
 from typing import Dict, List, Any
+from scipy.io import wavfile
 from transformers import AutoProcessor, MusicgenForConditionalGeneration
 import torch
+import io
 
 def create_params(params, fr):
     # default
@@ -35,11 +37,13 @@ class EndpointHandler:
         self.model = MusicgenForConditionalGeneration.from_pretrained(path, torch_dtype=torch.float16)
         self.model.to('cuda')
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def __call__(self, data: Dict[str, Any]) -> bytes:
         """
         Args:
             data (:dict:):
                 The payload with the text prompt and generation parameters.
+
+       Returns: wav file in bytes
         """
 
         inputs = data.pop("inputs", data)
@@ -65,7 +69,14 @@ class EndpointHandler:
         except:
              sr = 32000
 
-        return [{"audio": pred, "sr":sr}]
+         # Convert the audio data to WAV format
+        wav_buffer = io.BytesIO()
+        wavfile.write(wav_buffer, sr, pred)
+
+        # Convert BytesIO to bytes
+        wav_data = wav_buffer.getvalue()
+
+        return  wav_data
 
 
 if __name__ == "__main__":

requirements.txt

@@ -1,2 +1,3 @@
 transformers==4.31.0
 accelerate>=0.20.3
+scipy==1.11.1

test.py

@@ -0,0 +1,17 @@
+from handler import EndpointHandler
+
+# init handler
+my_handler = EndpointHandler(path="pbotsaris/musicgen-small")
+
+non_holiday_payload = {"inputs": "I am quite excited how this will turn out", "duration": 2}
+# prepare sample payload
+
+# 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}]