output_extension support

src/predict.py

@@ -204,7 +204,7 @@ class Predictor:
         return {"url": file_url}
 
 
-    def predict(self,s3_url,passage,process_audio,run_type='styletts2'):
+    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)
@@ -225,14 +225,14 @@ class Predictor:
         if run_type == 'styletts2':
             model,sampler = self.model,self.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")
+            generated_output = os.path.join(results_dir,f"{gen_id}-voice-clone-1.wav")
 
-            sf.write(final_output,result,24000)
+            sf.write(generated_output,result,24000)
             if process_audio:
-                (new_sr, wav1) = self._fn(final_output,"Midpoint",32,0.5)
-                sf.write(final_output,wav1,new_sr)
-            mp3_final_output = str(final_output).replace('wav','mp3')
-            self.convert_wav_to_mp3(final_output,mp3_final_output)
+                (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)
@@ -248,18 +248,18 @@ class Predictor:
             if process_audio:
                 (new_sr, wav1) = self._fn(openvoice_output,"Midpoint",32,0.5)
                 sf.write(openvoice_output,wav1,new_sr)
-            mp3_final_output = str(openvoice_output).replace('wav','mp3')
-            self.convert_wav_to_mp3(openvoice_output,mp3_final_output)
+            final_output = str(openvoice_output).replace('wav',output_extension)
+            self.convert_wav_to_output_extension(openvoice_output,final_output)
         
-        self.upload_file_to_s3(mp3_final_output,'demovidelyusergenerations',f"{gen_id}-voice-clone.mp3")
+        self.upload_file_to_s3(final_output,'demovidelyusergenerations',f"{gen_id}-voice-clone.{output_extension}")
         shutil.rmtree(os.path.join(output_dir,gen_id))
-        return {"voice_clone":f"https://demovidelyusergenerations.s3.amazonaws.com/{gen_id}-voice-clone.mp3"
+        return {"voice_clone":f"https://demovidelyusergenerations.s3.amazonaws.com/{gen_id}-voice-clone.{output_extension}"
                 }
             
     
 
 
-    def predict_with_emotions(self,s3_url,passage,process_audio):
+    def predict_with_emotions(self,s3_url,passage,output_extension,process_audio):
         output_dir = 'processed'
         gen_id = str(uuid.uuid4())
         os.makedirs(output_dir,exist_ok=True)
@@ -292,15 +292,15 @@ class Predictor:
         if process_audio:
             (new_sr, wav1) = self._fn(openvoice_output,"Midpoint",32,0.5)
             sf.write(openvoice_output,wav1,new_sr)
-        mp3_final_output_2 = str(openvoice_output).replace('wav','mp3')
-        self.convert_wav_to_mp3(openvoice_output,mp3_final_output_2)
-        self.upload_file_to_s3(mp3_final_output_2,'demovidelyusergenerations',f"{gen_id}-voice-clone-emotions.mp3")
+        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,'demovidelyusergenerations',f"{gen_id}-voice-clone-emotions.{output_extension}")
         shutil.rmtree(os.path.join(output_dir,gen_id))
-        return {"voice_clone_emotions":f"https://demovidelyusergenerations.s3.amazonaws.com/{gen_id}-voice-clone-emotions.mp3",
+        return {"voice_clone_emotions":f"https://demovidelyusergenerations.s3.amazonaws.com/{gen_id}-voice-clone-emotions.{output_extension}",
                 }
         
 
-    def predict_with_multi_lang(self,s3_url,passage,process_audio):
+    def predict_with_multi_lang(self,s3_url,passage,output_extension,process_audio):
         print("In multi lang voice cloning")
         output_dir = 'processed'
         gen_id = str(uuid.uuid4())
@@ -342,12 +342,12 @@ class Predictor:
         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='')
 
-        mp3_final_output_1 = str(multi_lang_with_voice_clone_path).replace('wav','mp3')
-        self.convert_wav_to_mp3(multi_lang_with_voice_clone_path,mp3_final_output_1)
-        print(mp3_final_output_1)
-        self.upload_file_to_s3(mp3_final_output_1,'demovidelyusergenerations',f"{gen_id}-voice-clone-multi-lang.mp3")
+        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,'demovidelyusergenerations',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://demovidelyusergenerations.s3.amazonaws.com/{gen_id}-voice-clone-multi-lang.mp3"
+        return {"voice_clone_with_multi_lang":f"https://demovidelyusergenerations.s3.amazonaws.com/{gen_id}-voice-clone-multi-lang.{output_extension}"
         }
 
 
@@ -520,8 +520,8 @@ class Predictor:
             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]
+    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)
 
 

src/rp_handler.py

@@ -45,20 +45,24 @@ def run_voice_clone_job(job):
         s3_url = job_input.get('s3_url')
         passage = job_input.get('passage')
         process_audio = job_input.get('process_audio')
+        output_extension = job_input.get('output_extension')
         print(process_audio)
+        if output_extension not in ["mp3","ogg"]:
+            return {"error" : "only supports mp3 and ogg as output_extension"}
+        print(output_extension)
         if process_audio is None:
             process_audio = False
 
         if method_type == 'voice_clone':
             run_type = job_input.get('run_type')
             if run_type is not None:
-                result = MODEL.predict(s3_url,passage,process_audio,run_type)
+                result = MODEL.predict(s3_url,passage,process_audio,output_extension,run_type)
             else:
-                result = MODEL.predict(s3_url,passage,process_audio)
+                result = MODEL.predict(s3_url,passage,process_audio,output_extension)
         if method_type == 'voice_clone_with_emotions':
-            result = MODEL.predict_with_emotions(s3_url,passage,process_audio)
+            result = MODEL.predict_with_emotions(s3_url,passage,process_audio,output_extension)
         if method_type == 'voice_clone_with_multi_lang':
-            result = MODEL.predict_with_multi_lang(s3_url,passage,process_audio)
+            result = MODEL.predict_with_multi_lang(s3_url,passage,process_audio,output_extension)
 
         return result
 

src/rp_schema.py

@@ -33,5 +33,10 @@ INPUT_VALIDATIONS = {
         'type': str,
         'required': False,
         'default': False
+    },
+    'output_extension': {
+        'type': str,
+        'required': False,
+        'default': 'ogg'
     }
 }