change to AudioGen

handler.py

@@ -1,13 +1,19 @@
 from typing import Dict, List, Any
-from transformers import AutoProcessor, MusicgenForConditionalGeneration
-import torch
+# from transformers import AutoProcessor, MusicgenForConditionalGeneration
+# import torch
+
+# import torchaudio
+from audiocraft.models import AudioGen
+from audiocraft.data.audio import audio_write
 
 class EndpointHandler:
     def __init__(self, path=""):
         # load model and processor from path
-        path = "jamesdon/audiogen-medium-endpoint"
-        self.processor = AutoProcessor.from_pretrained(path)
-        self.model = MusicgenForConditionalGeneration.from_pretrained(path).to("cuda")
+        # path = "jamesdon/audiogen-medium-endpoint"
+        # self.processor = AutoProcessor.from_pretrained(path)
+        # self.model = MusicgenForConditionalGeneration.from_pretrained(path).to("cuda")
+        self.model = AudioGen.get_pretrained(path)
+        
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
         """
@@ -16,22 +22,11 @@ class EndpointHandler:
                 The payload with the text prompt and generation parameters.
         """
         # process input
-        inputs = data.pop("inputs", data)
-        parameters = data.pop("parameters", None)
-
-        # preprocess
-        inputs = self.processor(
-            text=[inputs],
-            padding=True,
-            return_tensors="pt",).to("cuda")
-
-        # pass inputs with all kwargs in data
-        if parameters is not None:
-            outputs = self.model.generate(**inputs, **parameters)
-        else:
-            outputs = self.model.generate(**inputs)
-
-        # postprocess the prediction
+        inputs = data.pop("inputs", data) # list of string
+        duration = data.pop("duration", 5) # seconds to generate
+        
+        self.model.set_generation_params(duration=duration)
+        outputs = self.model.generate(inputs)
         prediction = outputs[0].cpu().numpy()
 
         return [{"generated_audio": prediction}]

requirements.txt

@@ -1,3 +1,3 @@
 transformers==4.31.0
 accelerate>=0.20.3
-# audiocraft
+audiocraft