File size: 2,156 Bytes
ad75a51 9e37e52 ad75a51 9e37e52 ad75a51 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | import base64
import io
import json
import os
import random
import re
from typing import Any, Dict, List
import inflect
import numpy as np
import torch
import torchaudio
import uroman as ur
from audiotokenizer import AudioTokenizerForLocal
from outetts.wav_tokenizer.decoder import WavTokenizer
from transformers import AutoModelForCausalLM, AutoTokenizer
class EndpointHandler:
def __init__(self, path=""):
# Preload all the elements you are going to need at inference.
# pseudo:
# self.model= load_model(path)
hf_path = "aybdee/igbo-SpeechSynthesis"
wav_tokenizer_config_path = "/data/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml"
wav_tokenizer_model_path = "/data/wavtokenizer_large_speech_320_24k.ckpt"
audio_tokenizer = AudioTokenizerForLocal(
hf_path, wav_tokenizer_model_path, wav_tokenizer_config_path
)
self.model = AutoModelForCausalLM.from_pretrained(
hf_path, torch_dtype="auto"
).to(audio_tokenizer.device)
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
data args:
inputs (:obj: `str` | `PIL.Image` | `np.array`)
kwargs
Return:
A :obj:`list` | `dict`: will be serialized and returned
"""
prompt = audio_tokenizer.create_prompt(data, "igbo", "igbo_male2")
input_ids = audio_tokenizer.tokenize_prompt(prompt)
output = self.model.generate(
input_ids=input_ids,
temperature=0.1,
repetition_penalty=1.1,
num_beams=4,
max_length=8192,
)
codes = audio_tokenizer.get_codes(output)
# converts the codes to audio
audio = audio_tokenizer.get_audio(codes)
buffer = io.BytesIO()
torchaudio.save(buffer, audio, sample_rate=sample_rate, format="wav")
buffer.seek(0) # go to beginning
# Base64 encode the buffer content
audio_bytes = buffer.read()
audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
return {"audio": audio_base64}
|