|
|
import os |
|
|
import sys |
|
|
import numpy as np |
|
|
import torch |
|
|
import soundfile as sf |
|
|
import spaces |
|
|
from config import models_path, results_path, sample_path, BASE_DIR |
|
|
|
|
|
encoder = None |
|
|
synthesizer = None |
|
|
vocoder = None |
|
|
|
|
|
def load_models(): |
|
|
global encoder, synthesizer, vocoder |
|
|
|
|
|
try: |
|
|
sys.path.append(os.path.join(BASE_DIR, 'pmt2')) |
|
|
|
|
|
from encoder import inference as encoder_module |
|
|
from synthesizer.inference import Synthesizer |
|
|
from parallel_wavegan.utils import load_model as vocoder_hifigan |
|
|
|
|
|
global encoder |
|
|
encoder = encoder_module |
|
|
|
|
|
print("Loading encoder model...") |
|
|
encoder.load_model(os.path.join(models_path, 'encoder.pt')) |
|
|
|
|
|
print("Loading synthesizer model...") |
|
|
synthesizer = Synthesizer(os.path.join(models_path, 'synthesizer.pt')) |
|
|
|
|
|
print("Loading HiFiGAN vocoder...") |
|
|
vocoder = vocoder_hifigan(os.path.join(models_path, 'vocoder_HiFiGAN.pkl')) |
|
|
vocoder.remove_weight_norm() |
|
|
vocoder = vocoder.eval().to('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
|
|
|
return True |
|
|
except Exception as e: |
|
|
import traceback |
|
|
print(f"Error loading models: {traceback.format_exc()}") |
|
|
return False |
|
|
|
|
|
@spaces.GPU(duration=120) |
|
|
def generate_speech(text, reference_audio=None): |
|
|
if not text or text.strip() == "": |
|
|
return None |
|
|
|
|
|
try: |
|
|
if reference_audio is None: |
|
|
ref_wav_path = sample_path |
|
|
else: |
|
|
ref_wav_path = os.path.join(results_path, "reference_audio.wav") |
|
|
sf.write(ref_wav_path, reference_audio[1], reference_audio[0]) |
|
|
|
|
|
print(f"Using reference audio: {ref_wav_path}") |
|
|
|
|
|
wav = synthesizer.load_preprocess_wav(ref_wav_path) |
|
|
|
|
|
encoder_wav = encoder.preprocess_wav(wav) |
|
|
embed, partial_embeds, _ = encoder.embed_utterance(encoder_wav, return_partials=True) |
|
|
|
|
|
texts = [text] |
|
|
embeds = [embed] * len(texts) |
|
|
specs = synthesizer.synthesize_spectrograms(texts, embeds) |
|
|
spec = np.concatenate(specs, axis=1) |
|
|
|
|
|
x = torch.from_numpy(spec.T).to('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
|
|
|
with torch.no_grad(): |
|
|
wav = vocoder.inference(x) |
|
|
|
|
|
wav = wav.cpu().numpy() |
|
|
wav = wav / np.abs(wav).max() * 0.97 |
|
|
|
|
|
output_filename = f"generated_{hash(text) % 10000}.wav" |
|
|
output_path = os.path.join(results_path, output_filename) |
|
|
sf.write(output_path, wav, synthesizer.sample_rate) |
|
|
|
|
|
return output_path |
|
|
|
|
|
except Exception as e: |
|
|
import traceback |
|
|
error_details = traceback.format_exc() |
|
|
print(f"Error generating speech: {error_details}") |
|
|
return None |
|
|
|
