Spaces:
Running
Running
| # coding=utf-8 | |
| # from .step0 import Step0 | |
| from .step1 import Step1 | |
| # from .step2 import Step2 | |
| from .step3 import Step3 | |
| from .step4 import Step4 | |
| from .step02 import Step02 | |
| import logging | |
| import re | |
| import os | |
| class Pipeline: | |
| def __init__(self, step0_model_dir, step1_mode, step2_model_dir, step2_mode, step3_mode): | |
| # self.step0 = Step0(step0_model_dir) | |
| self.step02 = Step02(step0_model_dir, step2_mode) | |
| self.step1 = Step1(step1_mode) | |
| # self.step2 = Step2(step2_model_dir, step2_mode) | |
| self.step3 = Step3(model_type=step3_mode) | |
| self.step4 = Step4() | |
| self.log = logging.getLogger(self.__class__.__name__) | |
| self.log.setLevel(logging.INFO) | |
| def run(self, video_input, output_dir, mode='s4', postp_mode='rep', prompt='', negative_prompt='', duration=10, seed=42): | |
| step0_resp = self.step02.run_step0(video_input) | |
| step0_resp_list = re.findall(r'(Step\d:.*?)(?=Step\d:|$)', step0_resp, re.DOTALL) | |
| step_infos = [step_info.strip().split("\n")[0] for step_info in step0_resp_list] | |
| step3_temp_dir = os.path.join(output_dir, "remove_vo") | |
| step_results = {"temp_final_audio_path": None, "temp_final_video_path": None} | |
| for step_info in step_infos: | |
| self.log.info(f"Start to {step_info}") | |
| if step_info == 'Step1: Generate audio from video.': | |
| step1_audio_path, step1_video_path = self.step1.run(video_input, output_dir, prompt, negative_prompt, duration=duration, seed=seed) | |
| step_results["step1_audio_path"] = step1_audio_path | |
| step_results["step1_video_path"] = step1_video_path | |
| elif step_info == 'Step2: Given a video and its generated audio, determine whether the audio contains voice-over.': | |
| is_vo = self.step02.run_step2(str(step_results["step1_video_path"])) | |
| step_results["is_vo"] = is_vo | |
| if not step_results["is_vo"]: # not voice-over | |
| step_results["temp_final_audio_path"] = step_results["step1_audio_path"] | |
| step_results["temp_final_video_path"] = step_results["step1_video_path"] | |
| return step_results | |
| elif step_info == 'Step3: Remove voice-over from audio.': | |
| step3_audio_path = self.step3.run(input_audio_path=step_results["step1_audio_path"], | |
| temp_store_dir=step3_temp_dir, | |
| output_dir=output_dir) | |
| step_results["step3_audio_path"] = step3_audio_path | |
| if mode == 's3': | |
| step_results["temp_final_audio_path"] = step_results["step3_audio_path"] | |
| return step_results | |
| elif step_info == 'Step4: Determine whether the audio is silent.': | |
| is_silent = self.step4.run(step_results["step3_audio_path"]) | |
| step_results["is_silent"] = is_silent | |
| else: | |
| self.log.error(f"Step-by-Step Error !!!!!!!!!") | |
| return step_results | |
| if not step_results["is_silent"]: # not silent | |
| step_results["temp_final_audio_path"] = step_results["step3_audio_path"] | |
| else: | |
| self.log.info(f"Start to post process, use mode: {postp_mode}") | |
| if postp_mode == "rm": | |
| step_results["temp_final_audio_path"] = None | |
| elif postp_mode == "rep": | |
| step_results["temp_final_audio_path"] = step_results["step1_audio_path"] | |
| step_results["temp_final_video_path"] = step_results["step1_video_path"] | |
| elif postp_mode == "neg": | |
| neg_audio_path, neg_video_path = self.step1.run(video_input, output_dir, prompt, negative_prompt='human voice', duration=duration, seed=seed, is_postp=True) | |
| step_results["temp_final_audio_path"] = neg_audio_path | |
| step_results["temp_final_video_path"] = neg_video_path | |
| else: | |
| self.log.error(f"Error postp_mode: {postp_mode}") | |
| self.log.info(f"After post-processing, audio is {step_results['temp_final_audio_path']} and video is {step_results['temp_final_video_path']}") | |
| self.log.info(f"Finish Post-Process successfully.\n") | |
| return step_results | |
| def run_for_gradio(self, video_input, output_dir, mode='s4', postp_mode='rep', prompt='', negative_prompt='', duration=10, seed=42): | |
| step_results = {"temp_final_audio_path": None, | |
| "temp_final_video_path": None, | |
| 'log': ''} | |
| step0_resp = self.step02.run_step0(video_input) | |
| step0_resp_list = re.findall(r'(Step\d:.*?)(?=Step\d:|$)', step0_resp, re.DOTALL) | |
| step_infos = [step_info.strip().split("\n")[0] for step_info in step0_resp_list] | |
| step3_temp_dir = os.path.join(output_dir, "remove_vo") | |
| for step_info in step_infos: | |
| self.log.info(f"Start to {step_info}") | |
| step_results['log'] = f"Start to {step_info}" | |
| yield step_results | |
| if step_info == 'Step1: Generate audio from video.': | |
| step1_audio_path, step1_video_path = self.step1.run(video_input, output_dir, prompt, negative_prompt, duration=duration, seed=seed) | |
| step_results["step1_audio_path"] = step1_audio_path | |
| step_results["step1_video_path"] = step1_video_path | |
| step_results['log'] = "Step1 completed." | |
| yield step_results | |
| elif step_info == 'Step2: Given a video and its generated audio, determine whether the audio contains voice-over.': | |
| is_vo = self.step02.run_step2(str(step_results["step1_video_path"])) | |
| step_results["is_vo"] = is_vo | |
| step_results['log'] = f"Step2 completed. Contain voice-over? {'Yes' if is_vo else 'No'}" | |
| yield step_results | |
| if not step_results["is_vo"]: # not voice-over | |
| step_results["temp_final_audio_path"] = step_results["step1_audio_path"] | |
| step_results["temp_final_video_path"] = step_results["step1_video_path"] | |
| step_results['log'] = "Finish step-by-step v2a." | |
| yield step_results | |
| elif step_info == 'Step3: Remove voice-over from audio.': | |
| step3_audio_path = self.step3.run(input_audio_path=step_results["step1_audio_path"], | |
| temp_store_dir=step3_temp_dir, | |
| output_dir=output_dir) | |
| step_results["step3_audio_path"] = step3_audio_path | |
| step_results['log'] = f"Step3 completed." | |
| yield step_results | |
| if mode == 's3': | |
| step_results["temp_final_audio_path"] = step_results["step3_audio_path"] | |
| step_results['log'] = "Finish step-by-step v2a." | |
| yield step_results | |
| elif step_info == 'Step4: Determine whether the audio is silent.': | |
| is_silent = self.step4.run(step_results["step3_audio_path"]) | |
| step_results["is_silent"] = is_silent | |
| step_results['log'] = f"Step4 completed. Silent? {'Yes' if is_silent else 'No'}" | |
| yield step_results | |
| else: | |
| self.log.error(f"Step-by-Step Error !!!!!!!!!") | |
| step_results['log'] = f"Step-by-Step Error !!!!!!!!!" | |
| yield step_results | |
| step_results['log'] = "Finish step-by-step v2a." | |
| yield step_results | |
| if not step_results["is_silent"]: # not silent | |
| step_results["temp_final_audio_path"] = step_results["step3_audio_path"] | |
| step_results['log'] = "Finish step-by-step v2a." | |
| yield step_results | |
| else: | |
| step_results['log'] = f"Post-processing with mode: {postp_mode}" | |
| yield step_results | |
| self.log.info(f"Start to post process, use mode: {postp_mode}") | |
| if postp_mode == "rm": | |
| step_results["temp_final_audio_path"] = None | |
| elif postp_mode == "rep": | |
| step_results["temp_final_audio_path"] = step_results["step1_audio_path"] | |
| step_results["temp_final_video_path"] = step_results["step1_video_path"] | |
| elif postp_mode == "neg": | |
| neg_audio_path, neg_video_path = self.step1.run(video_input, output_dir, prompt, negative_prompt='human voice', duration=duration, seed=seed, is_postp=True) | |
| step_results["temp_final_audio_path"] = neg_audio_path | |
| step_results["temp_final_video_path"] = neg_video_path | |
| else: | |
| self.log.error(f"Error postp_mode: {postp_mode}") | |
| self.log.info(f"After post-processing, audio is {step_results['temp_final_audio_path']} and video is {step_results['temp_final_video_path']}") | |
| self.log.info(f"Finish Post-Process successfully.\n") | |
| step_results['log'] = f"Post-processing completed." | |
| yield step_results | |
| step_results['log'] = "Finish step-by-step v2a." | |
| yield step_results | |