| import torch | |
| from parler_tts import ParlerTTSForConditionalGeneration | |
| from transformers import AutoTokenizer, set_seed | |
| import soundfile as sf | |
| class EndpointHandler: | |
| def __init__(self, path=""): | |
| self.device = "cuda:0" if torch.cuda.is_available() else "cpu" | |
| self.model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-mini-expresso").to(self.device) | |
| self.tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler-tts-mini-expresso") | |
| def __call__(self, data: Any): | |
| inputs = data["inputs"] | |
| prompt = inputs["prompt"] | |
| description = inputs["description"] | |
| input_ids = self.tokenizer(description, return_tensors="pt").input_ids.to(self.device) | |
| prompt_input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device) | |
| set_seed(42) | |
| try: | |
| generation = self.model.generate(input_ids=input_ids, prompt_input_ids=prompt_input_ids) | |
| audio_arr = generation.cpu().numpy().squeeze() | |
| return audio_arr | |
| except Exception as e: | |
| logger.error(str(e)) | |
| del inputs | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| return {"error": str(e)} |