Create tts.py

tts.py

@@ -0,0 +1,29 @@
+import torch
+import torchaudio
+from fairseq.checkpoint_utils import load_model_ensemble_and_task
+from fairseq.models.text_to_speech import CodeHiFiGANVocoder
+from fairseq.models.text_to_speech.hub_interface import TTSHubInterface
+
+class TortoiseTTS:
+    def __init__(self, device='cpu'):
+        self.device = device
+        self.models, self.cfg, self.task = load_model_ensemble_and_task(
+            ["checkpoints/tts_model.pt"],
+            arg_overrides={"vocoder": "hifigan", "fp16": False}
+        )
+        self.model = self.models[0].to(device)
+        self.vocoder = CodeHiFiGANVocoder("checkpoints/vocoder.pt", model_cfg=self.cfg).to(device)
+        TTSHubInterface.update_cfg_with_data_cfg(self.cfg, self.task.data_cfg)
+
+    def text_to_speech(self, text, speaker='russian'):
+        sample = TTSHubInterface.get_model_input(self.task, text)
+        sample['net_input']['src_tokens'] = sample['net_input']['src_tokens'].to(self.device)
+        
+        with torch.no_grad():
+            wav, rate = TTSHubInterface.get_prediction(
+                self.model, 
+                self.vocoder, 
+                sample, 
+                speaker=speaker
+            )
+        return wav.cpu().numpy(), rate