Use audiocraft based implementation

handler.py

@@ -1,13 +1,15 @@
+from audiocraft.data.audio_utils import convert_audio
+from audiocraft.data.audio import audio_write
+from audiocraft.models.encodec import InterleaveStereoCompressionModel
+from audiocraft.models import MusicGen, MultiBandDiffusion
+
 from typing import Dict, List, Any
-from transformers import AutoProcessor, MusicgenForConditionalGeneration
 import torch
 import logging
 
 class EndpointHandler:
     def __init__(self, path=""):
-        # load model and processor from path
-        self.processor = AutoProcessor.from_pretrained(path)
-        self.model = MusicgenForConditionalGeneration.from_pretrained(path, torch_dtype=torch.float16).to("cuda")
+        self.model = MusicGen.get_pretrained("musicgen-medium")
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
         """
@@ -21,19 +23,13 @@ class EndpointHandler:
         parameters = data.pop("parameters", None)
         self.model.set_generation_params(**parameters)
 
-        # 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:
             with torch.autocast("cuda"):
                 outputs = self.model.generate(**inputs)
         else:
             with torch.autocast("cuda"):
-                outputs = self.model.generate(**inputs,)
+                outputs = self.model.generate(**inputs)
 
         # postprocess the prediction
         prediction = outputs[0].cpu().numpy().tolist()