working handler

call_endpoint.py

@@ -0,0 +1,22 @@
+import requests
+
+# Define the URL of the FastAPI server
+url = "http://localhost:8000/generate_audio"
+
+# Define the text for which you want to generate audio
+text = "lo-fi music with a soothing melody"
+
+# Define the headers for the request
+headers = {"Content-Type": "application/json"}
+
+# Make a POST request to the endpoint with the text data in the request body and the specified header
+response = requests.post(url, json={"text": text}, headers=headers)
+
+# Check if the request was successful
+if response.status_code == 200:
+    # Save the audio file
+    with open("generated_audio.wav", "wb") as f:
+        f.write(response.content)
+    print("Audio file saved successfully.")
+else:
+    print("Error:", response.text)

handler.py

@@ -0,0 +1,45 @@
+from typing import  Dict, List, Any
+from transformers import pipeline
+import soundfile as sf
+import torch
+import logging
+import io
+
+logger = logging.getLogger(__name__)
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        # load the optimized model
+        # create inference pipeline
+        self.pipeline = pipeline("text-to-audio", "facebook/musicgen-stereo-large", device="mps", torch_dtype=torch.float16)
+
+    def generate_audio(self, text: str):
+        # Here you can implement your audio generation logic
+        # For demonstration purposes, let's use your existing code
+        logger.info("Generating audio for text: %s", text)
+        try:
+            music = self.pipeline(text, forward_params={"max_new_tokens": 256})
+            return music["audio"][0].T, music["sampling_rate"]
+        except Exception as e:
+            logger.error("Error generating audio for text: %s", text, exc_info=True)
+            raise e
+
+    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+        """
+        Args:
+            data (:obj:):
+                includes the input data and the parameters for the inference.
+        Return:
+            A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing :
+                - "label": A string representing what the label/class is. There can be multiple labels.
+                - "score": A score between 0 and 1 describing how confident the model is for this label/class.
+        """
+        input = data.pop("input", data)
+
+        audio_data, sampling_rate = self.generate_audio(input)
+
+        
+        with io.BytesIO() as buffer:
+            sf.write(buffer, audio_data, sampling_rate, format="WAV")
+            buffer.seek(0)
+            return buffer.getvalue()

test.py

@@ -0,0 +1,14 @@
+from handler import EndpointHandler
+
+# init handler
+my_handler = EndpointHandler(path=".")
+
+# prepare sample payload
+payload = {"input": "Lowfi hiphop with deep bass"}
+
+# test the handler
+pred=my_handler(payload)
+
+
+with open("generated_audio.wav", "wb") as f:
+    f.write(pred)