| 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=""): |
| |
| |
| |
| 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) |
|
|
| |
| audio = audio_tokenizer.get_audio(codes) |
| buffer = io.BytesIO() |
| torchaudio.save(buffer, audio, sample_rate=sample_rate, format="wav") |
| buffer.seek(0) |
|
|
| |
| audio_bytes = buffer.read() |
| audio_base64 = base64.b64encode(audio_bytes).decode("utf-8") |
| return {"audio": audio_base64} |
|
|
