deliverable api

src/predict.py

@@ -6,6 +6,11 @@ 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
@@ -52,6 +57,9 @@ from text_utils import TextCleaner
 from pydantic import BaseModel, HttpUrl
 from api import BaseSpeakerTTS, ToneColorConverter
 
+from pydub import AudioSegment
+
+
 class Predictor:
     def __init__(self):
         self.model = None
@@ -135,7 +143,70 @@ class Predictor:
         self.tone_color_converter.load_ckpt(f'{self.ckpt_converter}/checkpoint.pth')
 
 
-    def predict(self,s3_url,passage,method_type='voice_clone'):
+    def createvoice(self,audio_base_64,cut_audio,process_audio):
+        file_bytes = base64.b64decode(audio_base_64)
+        file_buffer = io.BytesIO(file_bytes)
+
+        header = file_buffer.read(12)
+        print(header)
+        file_format = None
+        bucket_name = 'demovidelyuseruploads'
+        if b'WAVE' in header:
+            file_format = 'wav'
+        elif header.startswith((b'\xff\xfb', b'\xff\xf3', b'\xff\xe3', b'\xff\xfa')):
+            file_format = 'mp3'
+        else:
+            file_format = 'unknown'
+        if file_format == 'unknown':
+            return {'error':'unrecognized file format, encode audio file as base64 str'}
+
+        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 == "mp3":
+            wav_filename = f"{unique_filename}.wav"
+            subprocess.run(["ffmpeg", "-i", local_filename, wav_filename])
+            os.remove(local_filename)
+        print(wav_filename)
+
+        # 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,method_type='voice_clone'):
         output_dir = 'processed'
         gen_id = str(uuid.uuid4())
         os.makedirs(output_dir,exist_ok=True)
@@ -159,29 +230,32 @@ class Predictor:
             #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)
+            result = self.process_audio_file(local_file_path,passage,model,sampler)
             final_output = os.path.join(results_dir,f"{gen_id}-voice-clone-1.wav")
+
             sf.write(final_output,result,24000)
+            if process_audio:
+                (new_sr, wav1) = self._fn(final_output,"Midpoint",32,0.5)
+                sf.write(final_output,wav1,new_sr)
 
             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}-2.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='')
-            (new_sr, wav1) = self._fn(openvoice_output,"Midpoint",32,0.5)
-            denoised_openvoice_output = os.path.join(results_dir,f"{gen_id}-voice-clone-2.wav")
-            sf.write(denoised_openvoice_output,wav1,new_sr)
-
-            
+            if process_audio:
+                (new_sr, wav1) = self._fn(openvoice_output,"Midpoint",32,0.5)
+                sf.write(openvoice_output,wav1,new_sr)
+           
            
             mp3_final_output_1 = str(final_output).replace('wav','mp3')
-            mp3_final_output_2 = str(denoised_openvoice_output).replace('wav','mp3')
+            mp3_final_output_2 = str(openvoice_output).replace('wav','mp3')
             self.convert_wav_to_mp3(final_output,mp3_final_output_1)
-            self.convert_wav_to_mp3(denoised_openvoice_output,mp3_final_output_2)
+            self.convert_wav_to_mp3(openvoice_output,mp3_final_output_2)
             print(mp3_final_output_1)
             print(mp3_final_output_2)
             
@@ -200,19 +274,19 @@ class Predictor:
             base_speaker_tts.tts(passage,src_path,speaker='default',language='English',speed=1.0,use_emotions=True)
             source_se = torch.load(f'{self.ckpt_base}/en_style_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='')
-            (new_sr, wav1) = self._fn(openvoice_output,"Midpoint",32,0.5)
-            denoised_openvoice_output = os.path.join(results_dir,f"{gen_id}-with-emotions.wav")
-            sf.write(denoised_openvoice_output,wav1,new_sr)
+            if process_audio:
+                (new_sr, wav1) = self._fn(openvoice_output,"Midpoint",32,0.5)
+                sf.write(openvoice_output,wav1,new_sr)
             
-            mp3_final_output_1 = str(denoised_openvoice_output).replace('wav','mp3')
-            self.convert_wav_to_mp3(denoised_openvoice_output,mp3_final_output_1)
+            mp3_final_output_1 = str(openvoice_output).replace('wav','mp3')
+            self.convert_wav_to_mp3(openvoice_output,mp3_final_output_1)
             print(mp3_final_output_1)
             self.upload_file_to_s3(mp3_final_output_1,'demovidelyusergenerations',f"{gen_id}-voice-with-emotions.mp3")
             shutil.rmtree(os.path.join(output_dir,gen_id))
             return {"voice_clone_with_emotions":f"https://demovidelyusergenerations.s3.amazonaws.com/{gen_id}-voice-with-emotions.mp3"
                     }
         
-        if method_type == 'voice_clone_multi_lang':
+        if method_type == 'voice_clone_with_multi_lang':
             #voice clone with multi-lingugal
             _,tone_color_converter = self.base_speaker_tts,self.tone_color_converter
             reference_speaker = local_file_path
@@ -236,9 +310,10 @@ class Predictor:
             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='')
 
             mp3_final_output_1 = str(multi_lang_with_voice_clone_path).replace('wav','mp3')
-            convert_wav_to_mp3(multi_lang_with_voice_clone_path,mp3_final_output_1)
+            self.convert_wav_to_mp3(multi_lang_with_voice_clone_path,mp3_final_output_1)
             print(mp3_final_output_1)
-            upload_file_to_s3(mp3_final_output_1,'demovidelyusergenerations',f"{gen_id}-voice-clone-multi-lang.mp3")
+            self.upload_file_to_s3(mp3_final_output_1,'demovidelyusergenerations',f"{gen_id}-voice-clone-multi-lang.mp3")
+            shutil.rmtree(os.path.join(output_dir,gen_id))
             return {"voice_clone_with_emotions":f"https://demovidelyusergenerations.s3.amazonaws.com/{gen_id}-voice-clone-multi-lang.mp3"
             }
     
@@ -249,7 +324,7 @@ class Predictor:
 
         solver = solver.lower()
         nfe = int(nfe)
-        lambd = 0.9
+        lambd = 0.1 # lets remove denoise
 
         dwav, sr = torchaudio.load(path)
         dwav = dwav.mean(dim=0)
@@ -380,14 +455,9 @@ class Predictor:
 
         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!') 
+    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
@@ -398,7 +468,7 @@ class Predictor:
                                     s_prev, 
                                     s_ref, 
                                     alpha = 0, 
-                                    beta = 0.1,  
+                                    beta = 0.3,  
                                     t = 0.7, 
                                     diffusion_steps=10, embedding_scale=1)
             wavs.append(wav)
@@ -437,4 +507,5 @@ class Predictor:
             return False
         except Exception as e:
             print(f"Error uploading file: {e}")
-            return False
+            return False
+    

src/processed/4d651a78-ccbd-4f66-96b1-0e0ede048d77/raw/631a27e2-8466-463e-a6ca-a2afd468c5a3.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c25a5bee8b60933b09cc779d942fa5c219f437e455bf64b08c2623f1c833ccfe
+size 322856

src/processed/69b28271-7198-4307-8501-e3969bbebef4/raw/631a27e2-8466-463e-a6ca-a2afd468c5a3.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c25a5bee8b60933b09cc779d942fa5c219f437e455bf64b08c2623f1c833ccfe
+size 322856

src/rp_handler.py

@@ -21,13 +21,36 @@ MODEL.setup()
 @rp_debugger.FunctionTimer
 def run_voice_clone_job(job):
     job_input = job['input']
-    method_type = job_input['method_type']
-    assert method_type in ["create_voice","voice_clone","voice_clone_with_emotions","voice_clone_with_multi_lang"]
-    s3_url = job_input['s3_url']
-    passage = job_input['passage']
-    processed_urls = MODEL.predict(s3_url,passage,method_type)
-
-    return processed_urls
+    method_type = job_input.get('method_type')
+    
+    if method_type not in ["create_voice","voice_clone","voice_clone_with_emotions","voice_clone_with_multi_lang"]:
+        return {"error":"Please set method_type: available options, create_voice, voice_clone, voice_clone_with_emotions,voice_clone_with_multi_lang"}
+    
+    if method_type == "create_voice":
+        audio_base64 = job_input.get('audio_base64')
+        if audio_base64 is None:
+            return {"error":"Needs audio file as base64"}
+        cut_audio = job_input.get('cut_audio')
+        process_audio = job_input.get('process_audio')
+        print(process_audio)
+        if process_audio is None:
+            process_audio = False
+        if cut_audio is None:
+            cut_audio = 0
+        
+        processed_urls = MODEL.createvoice(audio_base64,cut_audio,process_audio)
+        return processed_urls
+    else:
+        s3_url = job_input.get('s3_url')
+        passage = job_input.get('passage')
+        process_audio = job_input.get('process_audio')
+        print(process_audio)
+        if process_audio is None:
+            process_audio = False
+
+        result = MODEL.predict(s3_url,passage,process_audio,method_type)
+
+        return result
 
 
 runpod.serverless.start({"handler": run_voice_clone_job})

src/rp_schema.py

@@ -14,5 +14,20 @@ INPUT_VALIDATIONS = {
         'required': False,
         'default': 'None'
     },
-
+    'audio_base64': {
+        'type': str,
+        'required': False,
+        'default': 'None'
+    },
+    'cut_audio': {
+        'type': int,
+        'required': False,
+        'default': 0
+    },
+    'process_audio': {
+        'type': bool,
+        'required': False,
+        'default': False
+    }
+   
 }

src/se_extractor.py

@@ -10,7 +10,7 @@ from whisper_timestamped.transcribe import get_audio_tensor, get_vad_segments
 model_size = "medium"
 # Run on GPU with FP16
 model = None
-def split_audio_whisper(audio_path, target_dir='processed'):
+def split_audio_whisper(audio_path, target_dir='processed',needs_offset=True):
     global model
     if model is None:
         model = WhisperModel(model_size, device="cuda", compute_type="float16")
@@ -62,13 +62,18 @@ def split_audio_whisper(audio_path, target_dir='processed'):
             output_file = os.path.join(wavs_folder, fname)
             audio_seg.export(output_file, format='wav')
 
+        offset = 0.0
+        if needs_offset:
+            offset = 0.08
         if k < len(segments) - 1:
-            start_time = max(0, segments[k+1].start - 0.08)
+            start_time = max(0, segments[k+1].start - offset)
 
         s_ind = s_ind + 1
     return wavs_folder
 
 
+
+
 def split_audio_vad(audio_path, target_dir, split_seconds=10.0):
     SAMPLE_RATE = 16000
     audio_vad = get_audio_tensor(audio_path)
@@ -155,3 +160,24 @@ def generate_voice_segments(audio_path, target_dir='processed', vad=True):
 def load_model():
     model = WhisperModel(model_size, device="cpu", compute_type="int8")
 
+
+def extract_segments_to_cut_audio(max_duration,audio_path,target_dir='processed'):
+    global model
+    if model is None:
+        model = WhisperModel(model_size, device="cuda", compute_type="float16")
+    audio = AudioSegment.from_file(audio_path)
+    max_len = len(audio)
+    
+    segments, info = model.transcribe(audio_path, beam_size=5, word_timestamps=True)
+    segments = list(segments)
+    start_time = 0.0
+    end_time = max_len 
+    for segment in segments:
+        print(segment.end)
+        if segment.end > max_duration:
+            end_time = segment.end * 1000
+            break
+    max_duration_audio = audio[start_time:end_time]
+    max_duration_audio.export(audio_path,format='wav')
+
+