src/predict.py

"""
This file contains the Predictor class, which is used to run predictions on the
Whisper model. It is based on the Predictor class from the original Whisper
repository, with some modifications to make it work with the RP platform.
"""

from concurrent.futures import ThreadPoolExecutor
import numpy as np

from runpod.serverless.utils import rp_cuda
import boto3
import random
random.seed(0)
from glob import glob
import subprocess

import io

import numpy as np
np.random.seed(0)
import subprocess
import se_extractor

import yaml
from munch import Munch
import uuid
import shutil
from openai import OpenAI


import time
import os
import phonemizer
import torch
torch.manual_seed(0)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
from torch import nn
import torch.nn.functional as F
import torchaudio
import librosa
from nltk.tokenize import word_tokenize
import nltk
from Modules.diffusion.sampler import DiffusionSampler, ADPM2Sampler, KarrasSchedule
nltk.download('punkt')
from models import *
from utils import *
import soundfile as sf
from tortoise.utils.text import split_and_recombine_text
from resemble_enhance.enhancer.inference import denoise, enhance
from text_utils import TextCleaner
from pydantic import BaseModel, HttpUrl
from api import BaseSpeakerTTS, ToneColorConverter

class Predictor:
    def __init__(self):
        self.model = None
        self.sampler = None
        self.to_mel = None
        self.global_phonemizer = None
        self.model_params = None
        self.textclenaer = None
        self.mean = 0
        self.std = 0
        self.device = 'cuda:0'

        self.ckpt_base = 'checkpoints/base_speakers/EN'
        self.ckpt_converter = 'checkpoints/converter'
        self.base_speaker_tts = None
        self.tone_color_converter = None
        self.output_dir = 'outputs'
        self.processed_dir = 'processed'
        os.makedirs(self.processed_dir, exist_ok=True)
        os.makedirs(self.output_dir, exist_ok=True)
        self.s3_client = boto3.client('s3',aws_access_key_id=os.getenv('AWS_ACCESS_KEY'), aws_secret_access_key=os.getenv('AWS_SECRET_KEY'))
        print(os.getenv("AWS_ACCESS_KEY"))
        print(os.getenv("AWS_SECRET_KEY"))


    def setup(self):
        self.global_phonemizer = phonemizer.backend.EspeakBackend(language='en-us', preserve_punctuation=True,  with_stress=True)
        self.textclenaer = TextCleaner()
        self.to_mel = torchaudio.transforms.MelSpectrogram(
            n_mels=80, n_fft=2048, win_length=1200, hop_length=300)
        self.mean, self.std = -4, 4

        config = yaml.safe_load(open("Configs/hg.yml"))
        print(config)

        ASR_config = config.get('ASR_config', False)
        ASR_path = config.get('ASR_path', False)
        text_aligner = load_ASR_models(ASR_path, ASR_config)

        F0_path = config.get('F0_path', False)
        pitch_extractor = load_F0_models(F0_path)

        from Utils.PLBERT.util import load_plbert
        BERT_path = config.get('PLBERT_dir', False)
        plbert = load_plbert(BERT_path)

        self.model_params = recursive_munch(config['model_params'])
        self.model = build_model(self.model_params, text_aligner, pitch_extractor, plbert)
        _ = [self.model[key].eval() for key in self.model]
        _ = [self.model[key].to(self.device) for key in self.model]

        params_whole = torch.load("Models/epochs_2nd_00020.pth", map_location='cpu')
        params = params_whole['net']

        for key in self.model:
            if key in params:
                print('%s loaded' % key)
                try:
                    self.model[key].load_state_dict(params[key])
                except:
                    from collections import OrderedDict
                    state_dict = params[key]
                    new_state_dict = OrderedDict()
                    for k, v in state_dict.items():
                        name = k[7:] # remove `module.`
                        new_state_dict[name] = v
                    # load params
                    self.model[key].load_state_dict(new_state_dict, strict=False)
        #             except:
        #                 _load(params[key], model[key])
        _ = [self.model[key].eval() for key in self.model]
        self.sampler = DiffusionSampler(
            self.model.diffusion.diffusion,
            sampler=ADPM2Sampler(),
            sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0), # empirical parameters
            clamp=False
        )


    def predict(self,s3_url,passage,method_type='voice_clone'):
        output_dir = 'processed'
        gen_id = str(uuid.uuid4())
        os.makedirs(output_dir,exist_ok=True)
        raw_dir = os.path.join(output_dir,gen_id,'raw')
        segments_dir = os.path.join(output_dir,gen_id,'segments')
        results_dir = os.path.join(output_dir,gen_id,'results')
        openvoice_dir = os.path.join(output_dir,gen_id,'openvoice')
        os.makedirs(raw_dir)
        os.makedirs(segments_dir)
        os.makedirs(results_dir)
        
        
        s3_key = s3_url.split('/')[-1]
        bucket_name = 'demovidelyuseruploads'
        local_file_path = os.path.join(raw_dir,s3_key)
        self.download_file_from_s3(self.s3_client,bucket_name,s3_key,local_file_path)
        
        se_extractor.generate_voice_segments(local_file_path,segments_dir,vad=True)

        if method_type == 'voice_clone':
            #voice_clone with styletts2
            model,sampler = self.model,self.sampler
            processed_seg_dir = os.path.join(segments_dir,s3_key.split('.')[0],'wavs')
            result = self.process_audio_file(processed_seg_dir,passage,model,sampler)
            final_output = os.path.join(results_dir,f"{gen_id}-voice-clone-1.wav")
            sf.write(final_output,result,24000)
            
           
            mp3_final_output_1 = str(final_output).replace('wav','mp3')
            self.convert_wav_to_mp3(final_output,mp3_final_output_1)
            print(mp3_final_output_1)
            self.upload_file_to_s3(mp3_final_output_1,'demovidelyusergenerations',f"{gen_id}-voice-clone-1.mp3")
            return {"voice_clone_1":f"https://demovidelyusergenerations.s3.amazonaws.com/{gen_id}-voice-clone-1.mp3"}
        

    def _fn(self,path, solver, nfe, tau):
        if path is None:
            return None, None

        solver = solver.lower()
        nfe = int(nfe)
        lambd = 0.9

        dwav, sr = torchaudio.load(path)
        dwav = dwav.mean(dim=0)

        wav1, new_sr = enhance(dwav, sr, self.device, nfe=nfe, solver=solver, lambd=lambd, tau=tau)

        wav1 = wav1.cpu().numpy()

        return (new_sr, wav1)

    def _fn_denoise(self,path, solver, nfe, tau):
        if path is None:
            return None
        print(torch.cuda.is_available())
        print("Going to denoise")
        solver = solver.lower()
        nfe = int(nfe)
        lambd = 0.9

        dwav, sr = torchaudio.load(path)
        dwav = dwav.mean(dim=0)

        wav1, new_sr = denoise(dwav, sr, self.device)

        wav1 = wav1.cpu().numpy()
        print("Done noising")

        return (new_sr, wav1)

    def LFinference(self,model,sampler,text, s_prev, ref_s, alpha = 0.3, beta = 0.7, t = 0.7, diffusion_steps=5, embedding_scale=1):
        text = text.strip()
        ps = self.global_phonemizer.phonemize([text])
        ps = word_tokenize(ps[0])
        ps = ' '.join(ps)
        ps = ps.replace('``', '"')
        ps = ps.replace("''", '"')

        tokens = self.textclenaer(ps)
        tokens.insert(0, 0)
        tokens = torch.LongTensor(tokens).to(self.device).unsqueeze(0)
        
        with torch.no_grad():
            input_lengths = torch.LongTensor([tokens.shape[-1]]).to(self.device)
            text_mask = self.length_to_mask(input_lengths).to(self.device)

            t_en = model.text_encoder(tokens, input_lengths, text_mask)
            bert_dur = model.bert(tokens, attention_mask=(~text_mask).int())
            d_en = model.bert_encoder(bert_dur).transpose(-1, -2) 

            s_pred = sampler(noise = torch.randn((1, 256)).unsqueeze(1).to(self.device), 
                                            embedding=bert_dur,
                                            embedding_scale=embedding_scale,
                                                features=ref_s, # reference from the same speaker as the embedding
                                                num_steps=diffusion_steps).squeeze(1)
            
            if s_prev is not None:
                # convex combination of previous and current style
                s_pred = t * s_prev + (1 - t) * s_pred
            
            s = s_pred[:, 128:]
            ref = s_pred[:, :128]
            
            ref = alpha * ref + (1 - alpha)  * ref_s[:, :128]
            s = beta * s + (1 - beta)  * ref_s[:, 128:]

            s_pred = torch.cat([ref, s], dim=-1)

            d = model.predictor.text_encoder(d_en, 
                                            s, input_lengths, text_mask)

            x, _ = model.predictor.lstm(d)
            duration = model.predictor.duration_proj(x)

            duration = torch.sigmoid(duration).sum(axis=-1)
            pred_dur = torch.round(duration.squeeze()).clamp(min=1)


            pred_aln_trg = torch.zeros(input_lengths, int(pred_dur.sum().data))
            c_frame = 0
            for i in range(pred_aln_trg.size(0)):
                pred_aln_trg[i, c_frame:c_frame + int(pred_dur[i].data)] = 1
                c_frame += int(pred_dur[i].data)

            # encode prosody
            en = (d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(self.device))
            if self.model_params.decoder.type == "hifigan":
                asr_new = torch.zeros_like(en)
                asr_new[:, :, 0] = en[:, :, 0]
                asr_new[:, :, 1:] = en[:, :, 0:-1]
                en = asr_new

            F0_pred, N_pred = model.predictor.F0Ntrain(en, s)

            asr = (t_en @ pred_aln_trg.unsqueeze(0).to(self.device))
            if self.model_params.decoder.type == "hifigan":
                asr_new = torch.zeros_like(asr)
                asr_new[:, :, 0] = asr[:, :, 0]
                asr_new[:, :, 1:] = asr[:, :, 0:-1]
                asr = asr_new

            out = model.decoder(asr, 
                                    F0_pred, N_pred, ref.squeeze().unsqueeze(0))
        
            
        return out.squeeze().cpu().numpy()[..., :-100], s_pred #

    def length_to_mask(self,lengths):
        mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
        mask = torch.gt(mask+1, lengths.unsqueeze(1))
        return mask

    def preprocess(self,wave):
        wave_tensor = torch.from_numpy(wave).float()
        mel_tensor = self.to_mel(wave_tensor)
        mel_tensor = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - self.mean) / self.std
        return mel_tensor

    def compute_style(self,path,model):
        wave, sr = librosa.load(path, sr=24000)
        audio, index = librosa.effects.trim(wave, top_db=30)
        if sr != 24000:
            audio = librosa.resample(audio, sr, 24000)
        mel_tensor = self.preprocess(audio).to(self.device)

        with torch.no_grad():
            ref_s = model.style_encoder(mel_tensor.unsqueeze(1))
            ref_p = model.predictor_encoder(mel_tensor.unsqueeze(1))

        return torch.cat([ref_s, ref_p], dim=1)

    def process_audio_file(self,file_dir,passage,model,sampler):
        print(file_dir)
        audio_segs = glob(f'{file_dir}/*.wav')
        print(audio_segs)
        if len(audio_segs) >= 1:
            s_ref = self.compute_style(audio_segs[0], model)
        else:
            raise NotImplementedError('No audio segments found!') 
        sentences = split_and_recombine_text(passage)
        wavs = []
        s_prev = None
        for text in sentences:
            if text.strip() == "": continue
            text += '.'
            wav, s_prev = self.LFinference(model,sampler,text, 
                                    s_prev, 
                                    s_ref, 
                                    alpha = 0, 
                                    beta = 0.3,  # make it more suitable for the text
                                    t = 0.7, 
                                    diffusion_steps=10, embedding_scale=1)
            wavs.append(wav)

        audio_arrays = []
        for wav_file in wavs:
            audio_arrays.append(wav_file)
        concatenated_audio = np.concatenate(audio_arrays)
        return concatenated_audio

    def download_file_from_s3(self,s3_client,bucket_name, s3_key, local_file_path):
        try:
            s3_client.download_file(bucket_name, s3_key, local_file_path)
            print(f"File downloaded successfully: {local_file_path}")
        except Exception as e:
            print(f"Error downloading file: {e}")


    def convert_wav_to_mp3(self,wav_file, mp3_file):
        command = ['ffmpeg', '-i', wav_file, '-q:a', '0', '-map', 'a', mp3_file]
        subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)


    def upload_file_to_s3(self,file_name, bucket, object_name=None, content_type="audio/mpeg"):

        if object_name is None:
            object_name = file_name

        try:
            with open(file_name, 'rb') as file_data:
                self.s3_client.put_object(Bucket=bucket, Key=object_name, Body=file_data, ContentType=content_type)
            print("File uploaded successfully")
            return True
        except NoCredentialsError:
            print("Error: No AWS credentials found")
            return False
        except Exception as e:
            print(f"Error uploading file: {e}")
            return False