|
|
from typing import Dict, Any |
|
|
from transformers import AutoProcessor, MusicgenForConditionalGeneration |
|
|
import scipy |
|
|
import io |
|
|
|
|
|
class EndpointHandler: |
|
|
def __init__(self, path=""): |
|
|
|
|
|
self.processor = AutoProcessor.from_pretrained( |
|
|
path, |
|
|
local_files_only=True, |
|
|
trust_remote_code=True |
|
|
) |
|
|
self.model = MusicgenForConditionalGeneration.from_pretrained( |
|
|
path, |
|
|
local_files_only=True, |
|
|
trust_remote_code=True |
|
|
) |
|
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> bytes: |
|
|
text = data.get("inputs", "") |
|
|
duration = data.get("parameters", {}).get("duration", 5) |
|
|
|
|
|
inputs = self.processor( |
|
|
text=[text], |
|
|
return_tensors="pt", |
|
|
padding=True, |
|
|
truncation=True |
|
|
) |
|
|
|
|
|
audio_values = self.model.generate( |
|
|
**inputs, |
|
|
max_new_tokens=int(duration * 50) |
|
|
) |
|
|
|
|
|
sampling_rate = self.model.config.audio_encoder.sampling_rate |
|
|
with io.BytesIO() as wav_io: |
|
|
scipy.io.wavfile.write( |
|
|
wav_io, |
|
|
rate=sampling_rate, |
|
|
data=audio_values[0, 0].numpy() |
|
|
) |
|
|
return wav_io.getvalue() |
