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
import base64
from pydub.utils import mediainfo
import tempfile



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

from pydub import AudioSegment
from urllib.parse import urlparse


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'

        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"))
        self.bucket_name = os.getenv('S3_BUCKET_NAME')
        print(self.bucket_name)


    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
        )
        self.base_speaker_tts = BaseSpeakerTTS(f'{self.ckpt_base}/config.json', device=self.device)
        self.base_speaker_tts.load_ckpt(f'{self.ckpt_base}/checkpoint.pth')
        self.tone_color_converter = ToneColorConverter(f'{self.ckpt_converter}/config.json', device=self.device)
        self.tone_color_converter.load_ckpt(f'{self.ckpt_converter}/checkpoint.pth')


    def createvoice(self,s3_url,audio_base_64,cut_audio,process_audio):
        file_bytes = None
        bucket_name = self.bucket_name
        if s3_url:
            parsed_url = urlparse(s3_url)
            s3_key = parsed_url.path.lstrip('/')
            local_filename = f"{uuid.uuid4()}"
            self.download_file_from_s3(self.s3_client,bucket_name, s3_key, local_filename)
            with open(local_filename, 'rb') as file:
                file_bytes = file.read()
            os.remove(local_filename) 
        elif audio_base_64:
            file_bytes = base64.b64decode(audio_base_64)
        else:
            raise ValueError("Either s3_url or audio_base_64 must be provided.")

        file_buffer = io.BytesIO(file_bytes)
        header = file_buffer.read(12)
        print(header)

        file_format = 'mp3'  # Default format
        if b'WAVE' in header:
            file_format = 'wav'
        elif b'OggS' in header:
            file_format = 'ogg'

        unique_filename = f"{uuid.uuid4()}"
        local_filename = f"{unique_filename}.{file_format}"
        with open(local_filename, 'wb') as file_out:
            file_out.write(file_bytes)

        wav_filename = local_filename
        if file_format != "wav":
            wav_filename = f"{unique_filename}.wav"
            subprocess.run(["ffmpeg", "-i", local_filename, wav_filename])
            os.remove(local_filename)
        print(wav_filename)

        model,sampler = self.model,self.sampler
        result = self.process_audio_file(wav_filename,'this is a sample test to quickly run this model and resample',model,sampler)
        # if cut_audio > 0, means it was set
        if cut_audio > 0:
            #need to cut
            se_extractor.extract_segments_to_cut_audio(cut_audio,wav_filename)
            
        file_url = f"https://{bucket_name}.s3.amazonaws.com/{wav_filename}"

        if process_audio:
            (new_sr, wav1) = self._fn(wav_filename,"Midpoint",32,0.5)
            print('Denoised')
            buffer = io.BytesIO()
            sf.write(buffer, wav1, new_sr, format='WAV')
            print(new_sr)
            buffer.seek(0)
        else:
            wav1, sr = librosa.load(wav_filename, sr=None)
            buffer = io.BytesIO()
            sf.write(buffer, wav1, sr, format='WAV')
            buffer.seek(0)

        print("uploading")
        content_type = "audio/wav"
        try:
            self.s3_client.put_object(Bucket=bucket_name, Key=wav_filename, Body=buffer, ContentType=content_type)
            print("uploaded")
        except Exception as e:
            print(f"Error uploading to S3: {e}")
            return {"error": str(e)}

        os.remove(wav_filename)
        return {"url": file_url}


    def predict(self,s3_url,passage,process_audio,output_extension,run_type='styletts2'):
        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 = self.bucket_name
        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)
        #voice_clone with styletts2
        if run_type == 'styletts2':
            model,sampler = self.model,self.sampler
            result = self.process_audio_file(local_file_path,passage,model,sampler)
            generated_output = os.path.join(results_dir,f"{gen_id}-voice-clone-1.wav")

            sf.write(generated_output,result,24000)
            if process_audio:
                (new_sr, wav1) = self._fn(generated_output,"Midpoint",32,0.5)
                sf.write(generated_output,wav1,new_sr)
            final_output = str(generated_output).replace('wav',output_extension)
            self.convert_wav_to_output_extension(generated_output,final_output)

        if run_type == 'openvoice':
            s_ref = self.compute_style(local_file_path, self.model)
            base_speaker_tts,tone_color_converter = self.base_speaker_tts,self.tone_color_converter
            reference_speaker = local_file_path
            target_se, audio_name = se_extractor.get_se(reference_speaker, tone_color_converter, target_dir=openvoice_dir, vad=False)
            src_path = os.path.join(results_dir,f"{gen_id}-tmp.wav")
            openvoice_output = os.path.join(results_dir,f"{gen_id}-voice-clone-2.wav")
            base_speaker_tts.tts(passage,src_path,speaker='default',language='English',speed=1.0)

            source_se = torch.load(f'{self.ckpt_base}/en_default_se.pth').to(self.device)
            tone_color_converter.convert(audio_src_path=src_path,src_se=source_se,tgt_se=target_se,output_path=openvoice_output,message='')
            if process_audio:
                (new_sr, wav1) = self._fn(openvoice_output,"Midpoint",32,0.5)
                sf.write(openvoice_output,wav1,new_sr)
            final_output = str(openvoice_output).replace('wav',output_extension)
            self.convert_wav_to_output_extension(openvoice_output,final_output)
        
        self.upload_file_to_s3(final_output,bucket_name,f"{gen_id}-voice-clone.{output_extension}")
        shutil.rmtree(os.path.join(output_dir,gen_id))
        return {"voice_clone":f"https://{bucket_name}.s3.amazonaws.com/{gen_id}-voice-clone.{output_extension}"
                }
            
    


    def predict_with_emotions(self,s3_url,passage,process_audio,output_extension):
        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 = self.bucket_name
        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)
        #voice_clone with styletts2
        model,sampler = self.model,self.sampler
        result = self.process_audio_file(local_file_path,passage,model,sampler)

        base_speaker_tts,tone_color_converter = self.base_speaker_tts,self.tone_color_converter
        reference_speaker = local_file_path
        target_se, audio_name = se_extractor.get_se(reference_speaker, tone_color_converter, target_dir=openvoice_dir, vad=False)
        src_path = os.path.join(results_dir,f"{gen_id}-tmp.wav")
        openvoice_output = os.path.join(results_dir,f"{gen_id}-voice-clone-emotions.wav")
        print("extracting emotions from openai")
        base_speaker_tts.tts(passage,src_path,speaker='default',language='English',speed=1.0,use_emotions=True)
        source_se, audio_name = se_extractor.get_se(src_path, tone_color_converter, target_dir=openvoice_dir, vad=False)
        tone_color_converter.convert(audio_src_path=src_path,src_se=source_se,tgt_se=target_se,output_path=openvoice_output,message='')
        if process_audio:
            (new_sr, wav1) = self._fn(openvoice_output,"Midpoint",32,0.5)
            sf.write(openvoice_output,wav1,new_sr)
        final_ouput = str(openvoice_output).replace('wav',output_extension)
        self.convert_wav_to_output_extension(openvoice_output,final_ouput)
        self.upload_file_to_s3(final_ouput,bucket_name,f"{gen_id}-voice-clone-emotions.{output_extension}")
        shutil.rmtree(os.path.join(output_dir,gen_id))
        return {"voice_clone_emotions":f"https://{bucket_name}.s3.amazonaws.com/{gen_id}-voice-clone-emotions.{output_extension}",
                }
        

    def predict_with_multi_lang(self,s3_url,passage,process_audio,output_extension):
        print("In multi lang voice cloning")
        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 = self.bucket_name
        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)
        model,sampler = self.model,self.sampler
        result = self.process_audio_file(local_file_path,'this is a sample test to quickly run this model and resample',model,sampler)

        _,tone_color_converter = self.base_speaker_tts,self.tone_color_converter
        reference_speaker = local_file_path
        target_se, audio_name = se_extractor.get_se(reference_speaker, tone_color_converter, target_dir=openvoice_dir, vad=False)
        src_path = 'openai_source_output.mp3'
        source_se, audio_name = se_extractor.get_se(src_path, tone_color_converter, vad=True)

        
        client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
        response = client.audio.speech.create(
            model="tts-1",
            voice="fable",
            input=passage
        )

        openai_multi_lang_path = os.path.join(results_dir,f"{gen_id}-openai-gen.wav")
        response.stream_to_file(openai_multi_lang_path)
        multi_lang_with_voice_clone_path = os.path.join(results_dir,f"{gen_id}-voice-clone-multi-lang.wav")

        source_se, audio_name = se_extractor.get_se(src_path, tone_color_converter, vad=True)
        self.tone_color_converter.convert(audio_src_path=openai_multi_lang_path, src_se=source_se, tgt_se=target_se, output_path=multi_lang_with_voice_clone_path,message='')

        final_output = str(multi_lang_with_voice_clone_path).replace('wav',output_extension)
        self.convert_wav_to_output_extension(multi_lang_with_voice_clone_path,final_output)
        print(final_output)
        self.upload_file_to_s3(final_output,bucket_name,f"{gen_id}-voice-clone-multi-lang.{output_extension}")
        shutil.rmtree(os.path.join(output_dir,gen_id))
        return {"voice_clone_with_multi_lang":f"https://{bucket_name}.s3.amazonaws.com/{gen_id}-voice-clone-multi-lang.{output_extension}"
        }


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

        solver = solver.lower()
        nfe = int(nfe)
        lambd = 0.1 # lets remove denoise

        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,local_file_path,passage,model,sampler):
        print(local_file_path)
        s_ref = self.compute_style(local_file_path, model)
        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,  
                                    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_output_extension(self,wav_file, output_file):
        command = ['ffmpeg', '-i', wav_file, '-q:a', '0', '-map', 'a', output_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