|
from typing import Dict, List, Any |
|
from transformers import AutoProcessor, AutoModel |
|
import scipy.io.wavfile |
|
|
|
class EndpointHandler: |
|
def __init__(self, path=""): |
|
self.processor = AutoProcessor.from_pretrained("suno/bark") |
|
self.model = AutoModel.from_pretrained("suno/bark") |
|
|
|
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: |
|
try: |
|
text_prompt = data.get("inputs") |
|
if not text_prompt: |
|
raise ValueError("Missing required 'inputs' field in request data.") |
|
|
|
inputs = self.processor(text=[text_prompt], return_tensors="pt") |
|
speech_values = self.model.generate(**inputs, do_sample=True) |
|
|
|
|
|
audio_data = speech_values[0].numpy() |
|
sampling_rate = 22050 |
|
|
|
|
|
audio_bytes = audio_data.tobytes() |
|
return {"audio": audio_bytes, "sampling_rate": sampling_rate} |
|
|
|
except Exception as e: |
|
return {"error": str(e)} |
|
|
