Spaces:
Running
Running
File size: 2,105 Bytes
f5790af |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio
import wget
import json
import os
TTS_FOLDER = "./TTSModel"
TTS_MODEL_NAME = "vits"
TTS_MODEL_CONFIG = "config.json"
TTS_MODEL_WEIGHTS = "pytorch_model.bin"
TTS_VOCAB = "vocab.json"
TTS_CONFIG_URL = "https://huggingface.co/kakao-enterprise/vits-vctk/resolve/main/config.json"
TTS_MODEL_WEIGHTS_URL = "https://huggingface.co/kakao-enterprise/vits-vctk/resolve/main/pytorch_model.bin"
TTS_VOCAB_URL = "https://huggingface.co/kakao-enterprise/vits-vctk/resolve/main/vocab.json"
TTS_FILES_URLS = [
(TTS_CONFIG_URL, TTS_MODEL_CONFIG),
(TTS_MODEL_WEIGHTS_URL, TTS_MODEL_WEIGHTS),
(TTS_VOCAB_URL, TTS_VOCAB),
]
def ensure_tts_files_exist():
os.makedirs(TTS_FOLDER, exist_ok=True)
for url, filename in TTS_FILES_URLS:
filepath = os.path.join(TTS_FOLDER, filename)
if not os.path.exists(filepath):
wget.download(url, out=filepath)
class VITS(nn.Module):
def __init__(self, spec_channels, segment_size, num_speakers, num_languages, num_symbols):
super().__init__()
self.spec_channels = spec_channels
self.segment_size = segment_size
self.num_speakers = num_speakers
self.num_languages = num_languages
self.num_symbols = num_symbols
self.embedding = nn.Embedding(num_symbols, 192)
self.decoder = Generator(spec_channels)
def forward(self, text):
x = self.embedding(text)
audio = self.decoder(x)
return audio
class Generator(nn.Module):
def __init__(self, spec_channels):
super().__init__()
self.spec_channels = spec_channels
self.initial_conv = nn.ConvTranspose2d(192, spec_channels, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
self.final_conv = nn.Conv2d(spec_channels, 1, kernel_size=(7, 7), padding=(3, 3))
def forward(self, encoder_outputs):
x = encoder_outputs.unsqueeze(2)
x = self.initial_conv(x)
x = self.final_conv(x)
return x.squeeze(1) |