| from audiocraft.data.audio_utils import convert_audio |
| from audiocraft.data.audio import audio_write |
| from audiocraft.models import MusicGen |
|
|
| from typing import Dict, List, Any |
| import logging |
|
|
| class EndpointHandler: |
| def __init__(self, path=""): |
| self.model = MusicGen.get_pretrained("musicgen-medium") |
|
|
| def __call__(self, data: Dict[str, Any]) -> Dict[str, str]: |
| """ |
| Args: |
| data (:dict:): |
| The payload with the text prompt and generation parameters. |
| """ |
| |
| logging.info(f"data: {data}") |
| inputs = data.pop("inputs", data) |
| parameters = data.pop("parameters", None) |
| self.model.set_generation_params(**parameters) |
|
|
| outputs = self.model.generate(**inputs) |
|
|
| |
| prediction = outputs[0].cpu().numpy().tolist() |
|
|
| return [{"generated_audio": prediction}] |
|
|
