add new parameter

src/predict.py

@@ -204,7 +204,7 @@ class Predictor:
         return {"url": file_url}
 
 
-    def predict(self,s3_url,passage,process_audio):
+    def predict(self,s3_url,passage,process_audio,run_type='styletts2'):
         output_dir = 'processed'
         gen_id = str(uuid.uuid4())
         os.makedirs(output_dir,exist_ok=True)
@@ -222,41 +222,38 @@ class Predictor:
         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)
-        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}-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)
-        
-        
-        mp3_final_output_1 = str(final_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(openvoice_output,mp3_final_output_2)
-        print(mp3_final_output_1)
-        print(mp3_final_output_2)
+        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")
+
+            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)
+            mp3_final_output = str(final_output).replace('wav','mp3')
+            self.convert_wav_to_mp3(final_output,mp3_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)
+            mp3_final_output = str(openvoice_output).replace('wav','mp3')
+            self.convert_wav_to_mp3(openvoice_output,mp3_final_output)
         
-        self.upload_file_to_s3(mp3_final_output_1,'demovidelyusergenerations',f"{gen_id}-voice-clone-1.mp3")
-        self.upload_file_to_s3(mp3_final_output_2,'demovidelyusergenerations',f"{gen_id}-voice-clone-2.mp3")
+        self.upload_file_to_s3(mp3_final_output,'demovidelyusergenerations',f"{gen_id}-voice-clone.mp3")
         shutil.rmtree(os.path.join(output_dir,gen_id))
-        return {"voice_clone_1":f"https://demovidelyusergenerations.s3.amazonaws.com/{gen_id}-voice-clone-1.mp3",
-                "voice_clone_2":f"https://demovidelyusergenerations.s3.amazonaws.com/{gen_id}-voice-clone-2.mp3"
+        return {"voice_clone":f"https://demovidelyusergenerations.s3.amazonaws.com/{gen_id}-voice-clone.mp3"
                 }
             
     

src/rp_handler.py

@@ -50,7 +50,11 @@ def run_voice_clone_job(job):
             process_audio = False
 
         if method_type == 'voice_clone':
-            result = MODEL.predict(s3_url,passage,process_audio)
+            run_type = job_input.get('run_type')
+            if run_type is not None:
+                result = MODEL.predict(s3_url,passage,process_audio,run_type)
+            else:
+                result = MODEL.predict(s3_url,passage,process_audio)
         if method_type == 'voice_clone_with_emotions':
             result = MODEL.predict_with_emotions(s3_url,passage,process_audio)
         if method_type == 'voice_clone_with_multi_lang':

src/rp_schema.py

@@ -28,6 +28,10 @@ INPUT_VALIDATIONS = {
         'type': bool,
         'required': False,
         'default': False
+    },
+    'run_type': {
+        'type': str,
+        'required': False,
+        'default': False
     }
-   
 }